LEAF: Knowledge Distillation of Text Embedding Models with Teacher-Aligned Representations
This work addresses the need for efficient and high-performing text embedding models for applications like information retrieval, offering a practical distillation method that reduces infrastructure requirements.
The paper tackles the problem of knowledge distillation for text embedding models by introducing LEAF, a framework that aligns distilled models to their teachers, enabling flexible asymmetric architectures in information retrieval. The resulting leaf-ir model with 23M parameters achieves state-of-the-art performance on the BEIR benchmark (#1 on the leaderboard for its size) and further improves in asymmetric mode, while leaf-mt also sets a new SOTA on MTEB v2 (English) for its size.
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.