Fantastic (small) Retrievers and How to Train Them: mxbai-edge-colbert-v0 Tech Report
Rikiya Takehi, Benjamin Clavié, Sean Lee, Aamir Shakir
TL;DR
The paper presents mxbai-edge-colbert-v0, two compact ColBERT-backed encoders (17M and 32M) designed for scalable retrieval from cloud to on-device contexts. It builds a modern dense backbone from Ettin/ModernBERT, trained via a three-stage pipeline: contrastive pre-training, supervised fine-tuning with hard negatives, and Stella-style embedding distillation using a teacher like StellaV5 1.5B, with a 2-layer projection and an $ ext{L}_2 $ loss. Through extensive ablations on optimizers, projection dimensions and layers, and casing, the work demonstrates strong BEIR performance—most notably that the 17M model outperforms ColBERTv2 despite smaller embeddings—and significant gains for the 32M model from distillation and architectural choices, especially in long-context settings. The work emphasizes practical impact for on-device reranking and edge retrieval, providing a solid foundational baseline and a path for future efficient, small-scale ColBERT models.
Abstract
In this work, we introduce mxbai-edge-colbert-v0 models, at two different parameter counts: 17M and 32M. As part of our research, we conduct numerous experiments to improve retrieval and late-interaction models, which we intend to distill into smaller models as proof-of-concepts. Our ultimate aim is to support retrieval at all scales, from large-scale retrieval which lives in the cloud to models that can run locally, on any device. mxbai-edge-colbert-v0 is a model that we hope will serve as a solid foundation backbone for all future experiments, representing the first version of a long series of small proof-of-concepts. As part of the development of mxbai-edge-colbert-v0, we conducted multiple ablation studies, of which we report the results. In terms of downstream performance, mxbai-edge-colbert-v0 is a particularly capable small model, outperforming ColBERTv2 on common short-text benchmarks (BEIR) and representing a large step forward in long-context tasks, with unprecedented efficiency.
