用教师模型指导小模型,实现高效文本检索与多任务学习。
LEAF: Knowledge Distillation of Text Embedding Models with Teacher-Aligned Representations
- 通过对齐教师模型表示,训练轻量级嵌入模型。
- 在BEIR和MTEB基准上达到同规模模型最佳性能,支持异构编码架构。
- 无需标注数据或难例,小批量训练即可,适合资源有限场景。
我们提出LEAF(轻量嵌入对齐框架),一种面向文本嵌入模型的知识蒸馏框架。其核心特点是蒸馏后的叶子模型与教师模型对齐。在信息检索场景中,可实现灵活的异构架构:文档由大教师模型编码,查询则由小型叶子模型处理。实验表明,当教师模型具备最大相对长度(MRL)和量化鲁棒性时,叶子模型能自动继承这些特性,无需显式训练。为验证框架能力,我们发布了leaf-ir,一个2300万参数的信息检索专用嵌入模型,基于LEAF训练,在BEIR基准上排名榜首(公开榜单第一),同规模模型中表现最优;异构模式下检索性能进一步提升。该方法不限于检索任务,我们还构建了多任务leaf-mt模型,在MTEB v2(英语)榜单上同样位列第一。LEAF适用于黑盒模型,无需标注或难例样本,支持小批量训练,对数据集和训练基础设施要求低。所有模型均以Apache 2.0许可证开源。
原文摘要 · Abstract (English)
We present LEAF ("Lightweight Embedding Alignment Framework"), a knowledge distillation framework for text embedding models. A key distinguishing feature is that our distilled leaf models are aligned to their teacher. In the context of information retrieval, this allows for flexible asymmetric architectures where documents are encoded with the larger teacher model, while queries can be served with the smaller leaf models. We also show that leaf models automatically inherit MRL and robustness to output quantization whenever these properties are present in the teacher model, without explicitly training for them. To demonstrate the capability of our framework we publish leaf-ir, a 23M parameters information retrieval oriented text embedding model trained using LEAF, which sets a new state-of-the-art (SOTA) on BEIR, ranking #1 on the public leaderboard for this benchmark and for models of its size. When run in asymmetric mode, its retrieval performance is further increased. Our scheme is however not restricted to the information retrieval setting, and we demonstrate its wider applicability by synthesizing the multi-task leaf-mt model. This also sets a new SOTA, ranking #1 on the public MTEB v2 (English) leaderboard for its size. LEAF is applicable to black-box models and in contrast to other embedding model training frameworks, it does not require judgments nor hard negatives, and training can be conducted using small batch sizes. Thus, dataset and training infrastructure requirements for our framework are modest. We make our models publicly available under a permissive Apache 2.0 license.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。