8-layer distillation from BAAI/bge-m3 with2.5x speedup
This is an embedding model distilled from BAAI/bge-m3 on a combination of public and proprietary datasets. It is a 8-layer model --instead of 24 layers) in 366m-parameter size and achieves 2.5x speedup with little-to-no loss in retrieval performance.
Motivation
We are a team that have developed some of the real use cases of semantic search and RAG, and no other models apart from BAAI/bge-m3 have proved to be useful in a variety of domains and use cases, especially in multimodal settings. However, it's extra large and prohibitively expensive to serve for large user groups with a low latency and/or index large volumes of data. That's why we wanted the same retrieval performance in a smaller model size and with higher speed. We composed a large and diverse dataset of 10m texts and applied a knowledge distillation technique that reduced the number of layers from 24 to 8. The results were surprisingly promising --we achieved a Spearman Cosine score of 0.965 and MSE of 0.006 in the test subset, which can be even taken to be within numerical error ranges. We couldn't observe a considerable degredation in our qualitative tests, either. Finally, we measured a 2.5x throughput increase (454 texts / sec instead of 175 texts / sec, measured on a T4 Colab GPU).
Future Work
Even though our training dataset was composed of diverse texts in Turkish, the model retained a considerable performance in other languages as well --we measured a Spearman Cosine score of 0.938 in a collection 10k texts in English, for example. This performance retention motivated us to work on the second version of this distillation model trained on a larger and multilingual dataset as well as an even smaller distillation. Stay tuned for these updates, and feel free to reach out to us for collaboration options.
1from sentence_transformers import SentenceTransformer
23# Download from the 🤗 Hub4model = SentenceTransformer("altaidevorg/bge-m3-distill-8l")5# Run inference6sentences =[7'That is a happy person',8'That is a happy dog',9'That is a very happy person',10]11embeddings = model.encode(sentences)12print(embeddings.shape)13# [3, 1024]1415# Get the similarity scores for the embeddings16similarities = model.similarity(embeddings, embeddings)17print(similarities.shape)18# [3, 3]
Approximate statistics based on the first 1000 samples:
sentence
label
type
string
list
details
min: 5 tokens
mean: 55.78 tokens
max: 468 tokens
size: 1024 elements
Citation
BibTeX
Sentence Transformers
bibtex
1@inproceedings{reimers-2019-sentence-bert,
2 title = "Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks",
3 author = "Reimers, Nils and Gurevych, Iryna",
4 booktitle = "Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing",
5 month = "11",
6 year = "2019",
7 publisher = "Association for Computational Linguistics",
8 url = "https://arxiv.org/abs/1908.10084",
9}
MSELoss
bibtex
1@inproceedings{reimers-2020-multilingual-sentence-bert,
2 title = "Making Monolingual Sentence Embeddings Multilingual using Knowledge Distillation",
3 author = "Reimers, Nils and Gurevych, Iryna",
4 booktitle = "Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing",
5 month = "11",
6 year = "2020",
7 publisher = "Association for Computational Linguistics",
8 url = "https://arxiv.org/abs/2004.09813",
9}
bge-m3
bibtex
1@misc{bge-m3,
2 title={BGE M3-Embedding: Multi-Lingual, Multi-Functionality, Multi-Granularity Text Embeddings Through Self-Knowledge Distillation},
3 author={Jianlv Chen and Shitao Xiao and Peitian Zhang and Kun Luo and Defu Lian and Zheng Liu},
4 year={2024},
5 eprint={2402.03216},
6 archivePrefix={arXiv},
7 primaryClass={cs.CL}
8}