---
license: apache-2.0  
base_model:  
- microsoft/MiniLM-L6-v2  
tags:  
- transformers  
- sentence-transformers  
- sentence-similarity  
- feature-extraction  
- text-embeddings-inference  
- information-retrieval  
- knowledge-distillation  
language:
- en
---
<div style="display: flex; justify-content: center;">  
    <div style="display: flex; align-items: center; gap: 10px;">  
        <img src="logo.webp" alt="MongoDB Logo" style="height: 36px; width: auto;">  
        <span style="font-size: 32px; font-weight: bold">MongoDB/mdbr-leaf-ir</span>  
    </div>  
</div>

## Introduction

`mdbr-leaf-ir` is a compact high-performance text embedding model specifically designed for **information retrieval (IR)** tasks. 

Enabling even greater efficiency, `mdbr-leaf-ir` supports [flexible asymmetric architectures](#asymmetric-retrieval-setup) and is robust to [vector quantization](#vector-quantization) and [MRL truncation](#mrl).

If you are looking to perform other tasks such as classification, clustering, semantic sentence similarity, summarization, please check out our [`mdbr-leaf-mt`](https://huggingface.co/MongoDB/mdbr-leaf-mt) model.

Note: this model has been developed by MongoDB Research and is not part of MongoDB's commercial offerings.</span>

## Technical Report

A technical report detailing our proposed `LEAF` training procedure is [available here (TBD)](http://FILL_HERE_ARXIV_LINK).

## Highlights  

* **State-of-the-Art Performance**: `mdbr-leaf-ir` achieves new state-of-the-art results for compact embedding models, ranking <span style="color:red">#TBD</span> on the public BEIR benchmark leaderboard for models <30M parameters with an average nDCG@10 score of <span style="color:red">[TBD HERE]</span>.
* **Flexible Architecture Support**: `mdbr-leaf-ir` supports asymmetric retrieval architectures enabling even greater retrieval results. [See below](#asymmetric-retrieval-setup) for more information.
* **MRL and quantization support**: embedding vectors generated by `mdbr-leaf-ir` compress well when truncated (MRL) and/or are stored using more efficient types like `int8` and `binary`.  [See below](#mrl) for more information.

  
<!-- ## Performance  
### Benchmark Results  

* Values are nDCG@10
* Scores exclude CQADupstack and MSMARCO; full BEIR results are available on the [public leaderboard](https://huggingface.co/spaces/mteb/leaderboard).
* Scores in bold highlight when our model outperforms comparisons in either standard or asymmetric mode; we also highlight cases when comparisons outperform our model in standard mode. Blue are scores when asymmetric outperforms standard. 
* `BM25` scores are obtained with `(k₁=0.9, b=0.4)`.

| Model | Size | arg. | fiqa | nfc | scid. | scif. | quora | covid | nq | fever | c-fever | dbp. | hotpot | avg. |  
|-------|------|------|------|-----|-------|-------|--------|-------|----|----- |---------|------|--------|------|  
| **`mdbr-leaf-ir` (asym.)** | 23M | **<span style="color:blue">58.5</span>** | **<span style="color:blue">42.1</span>** | **36.1** | <span style="color:blue">20.4</span> | **69.9** | <span style="color:blue">86.2</span> | **<span style="color:blue">83.7</span>** | **<span style="color:blue">61.4</span>** | **<span style="color:blue">86.4</span>** | **<span style="color:blue">37.4</span>** | **<span style="color:blue">44.8</span>** | **<span style="color:blue">69.0</span>** | **<span style="color:blue">58.0</span>** |  
| **`mdbr-leaf-ir`** | 23M | **56.7** | **38.1** | **36.2** | 19.5 | **70.0** | 71.0 | **83.0** | **58.2** | **85.4** | **32.4** | 43.7 | 68.2 | **55.2** |  
| **Comparisons** | | | | | | | | | | | | | | |  
| `snowflake-arctic-embed-xs` | 23M | 52.1 | 34.5 | 30.9 | 18.4 | 64.5 | 86.6 | 79.4 | 54.8 | 83.4 | 29.9 | 40.2 | 65.3 | 53.3 |  
| `MiniLM-L6-v2` | 23M | 50.2 | 36.9 | 31.6 | **21.6** | 64.5 | **87.6** | 47.2 | 43.9 | 51.9 | 20.3 | 32.3 | 46.5 | 44.5 |  
| `BM25` | -- | 40.8 | 23.8 | 31.8 | 15.0 | 67.6 | 78.7 | 58.9 | 30.5 | 63.8 | 16.2 | 31.9 | 62.9 | 43.5 |  
| `SPLADE v2` | 110M | 47.9 | 33.6 | 33.4 | 15.8 | 69.3 | 83.8 | 71.0 | 52.1 | 78.6 | 23.5 | 43.5 | **68.4** | 51.7 |  
| `ColBERT v2` | 110M | 46.3 | 35.6 | 33.8 | 15.4 | 69.3 | 85.2 | 73.8 | 56.2 | 78.5 | 17.6 | **44.6** | 66.7 | 51.9 |  
-->

## Quickstart  
  
### Sentence Transformers  
  
```python  
from sentence_transformers import SentenceTransformer  
  
# Load the model  
model = SentenceTransformer("MongoDB/mdbr-leaf-ir")  
  
# Example queries and documents  
queries = [
    "What is machine learning?",  
    "How does neural network training work?"  
]  
  
documents = [  
    "Machine learning is a subset of artificial intelligence that focuses on algorithms that can learn from data.",  
    "Neural networks are trained through backpropagation, adjusting weights to minimize prediction errors."  
]  
  
# Encode queries and documents  
query_embeddings = model.encode(queries, prompt_name="query")  
document_embeddings = model.encode(documents)  
  
# Compute similarity scores  
scores = model.similarity(query_embeddings, document_embeddings)  

# Print results
for i, query in enumerate(queries):
    print(f"Query: {query}")
    for j, doc in enumerate(documents):
        print(f" Similarity: {scores[i, j]:.4f} | Document {j}: {doc[:80]}...")

# Query: What is machine learning?
#  Similarity: 0.6908 | Document 0: Machine learning is a subset of ...
#  Similarity: 0.4598 | Document 1: Neural networks are trained ...
# 
# Query: How does neural network training work?
#  Similarity: 0.4432 | Document 0: Machine learning is a subset of ...
#  Similarity: 0.5794 | Document 1: Neural networks are trained ...
```
  
### Transformers Usage  

<span style="color:red">CHECK THAT safe_open WORKS WITH URLS; link to code in repo</span>

<!-- ```python  
from safetensors import safe_open
from transformers import AutoModel, AutoTokenizer

# Load the model
tokenizer = AutoTokenizer.from_pretrained(MODEL)
model = AutoModel.from_pretrained(MODEL)

tensors = {}
with safe_open(MODEL + "/2_Dense/model.safetensors", framework="pt") as f:
    for k in f.keys():
        tensors[k] = f.get_tensor(k)

W_out = torch.nn.Linear(in_features=384, out_features=768, bias=True)
W_out.load_state_dict({
    "weight": tensors["linear.weight"], 
    "bias": tensors["linear.bias"]
})

_ = model.eval()
_ = W_out.eval()

# Example queries and documents  
queries = [
    "What is machine learning?",  
    "How does neural network training work?"  
]  
  
documents = [  
    "Machine learning is a subset of artificial intelligence that focuses on algorithms that can learn from data.",  
    "Neural networks are trained through backpropagation, adjusting weights to minimize prediction errors."  
]

# Tokenize
QUERY_PREFIX = 'Represent this sentence for searching relevant passages: '
queries_with_prefix = [QUERY_PREFIX + query for query in queries]

query_tokens = tokenizer(queries_with_prefix, padding=True, truncation=True, return_tensors='pt', max_length=512)
document_tokens =  tokenizer(documents, padding=True, truncation=True, return_tensors='pt', max_length=512)

# Perform Inference
with torch.inference_mode():
    y_queries = model(**query_tokens).last_hidden_state
    y_docs = model(**document_tokens).last_hidden_state

    # perform pooling
    y_queries = y_queries * query_tokens.attention_mask.unsqueeze(-1)
    y_queries_pooled = y_queries.sum(dim=1) / query_tokens.attention_mask.sum(dim=1, keepdim=True)

    y_docs = y_docs * document_tokens.attention_mask.unsqueeze(-1)
    y_docs_pooled = y_docs.sum(dim=1) / document_tokens.attention_mask.sum(dim=1, keepdim=True)

    # map to desired output dimension
    y_queries_out = W_out(y_queries_pooled)
    y_docs_out = W_out(y_docs_pooled)

    # normalize and return
    query_embeddings = F.normalize(y_queries_out, dim=-1)
    document_embeddings = F.normalize(y_docs_out, dim=-1)

similarities = query_embeddings @ document_embeddings.T
print(f"Similarities:\n{similarities}")
# Similarities:
#  tensor([[0.6908, 0.4598],
#          [0.4432, 0.5794]])
```   -->
  
### Asymmetric Retrieval Setup  
  
`mdbr-leaf-ir` is *aligned* to [`snowflake-arctic-embed-m-v1.5`](https://huggingface.co/Snowflake/snowflake-arctic-embed-m-v1.5), the model it has been distilled from, making the asymmetric system below possible: 
  
```python  
# Use a larger model for document encoding (one-time, at index time)  
doc_model = SentenceTransformer("Snowflake/snowflake-arctic-embed-m-v1.5")  
document_embeddings = doc_model.encode(documents)  
  
# Use mdbr-leaf-ir for query encoding (real-time, low latency)  
query_model = SentenceTransformer("MongoDB/mdbr-leaf-ir")  
query_embeddings = query_model.encode(queries, prompt_name="query")  
  
# Compute similarities  
scores = query_model.similarity(query_embeddings, document_embeddings)  
```
Retrieval results from asymmetric mode are usually superior to the [standard mode above](#sentence-transformers).

### MRL

Embeddings have been trained via [MRL](https://arxiv.org/abs/2205.13147) and can be truncated for more efficient storage:
```python
from torch.nn import functional as F

query_embeds = model.encode(queries, prompt_name="query", convert_to_tensor=True)
doc_embeds = model.encode(documents, convert_to_tensor=True)

# Truncate and normalize according to MRL
query_embeds = F.normalize(query_embeds[:, :256], dim=-1)
doc_embeds = F.normalize(doc_embeds[:, :256], dim=-1)

similarities = model.similarity(query_embeds, doc_embeds)

print('After MRL:')
print(f"* Embeddings dimension: {query_embeds.shape[1]}")
print(f"* Similarities:\n\t{similarities}")

# After MRL:
# * Embeddings dimension: 256
# * Similarities:
# 	tensor([[0.7202, 0.5006],
#           [0.4744, 0.6083]])
```

### Vector Quantization
Vector quantization, for example to `int8` or `binary`, can be performed as follows:

**Note**: For vector quantization to types other than binary, we suggest performing a calibration to determine the optimal ranges, [see here](https://sbert.net/examples/sentence_transformer/applications/embedding-quantization/README.html#scalar-int8-quantization). 
Good initial values, according to the [teacher model's documentation](https://huggingface.co/Snowflake/snowflake-arctic-embed-m-v1.5#compressing-to-128-bytes), are:
* `int8`: -0.3 and +0.3
* `int4`: -0.18 and +0.18 
```python
from sentence_transformers.quantization import quantize_embeddings
import torch

query_embeds = model.encode(queries, prompt_name="query")
doc_embeds = model.encode(documents)

# Quantize embeddings to int8 using -0.3 and +0.3 as calibration ranges
ranges = torch.tensor([[-0.3], [+0.3]]).expand(2, query_embeds.shape[1]).cpu().numpy()
query_embeds = quantize_embeddings(query_embeds, "int8", ranges=ranges)
doc_embeds = quantize_embeddings(doc_embeds, "int8", ranges=ranges)

# Calculate similarities; cast to int64 to avoid under/overflow
similarities = query_embeds.astype(int) @ doc_embeds.astype(int).T

print('After quantization:')
print(f"* Embeddings type: {query_embeds.dtype}")
print(f"* Similarities:\n{similarities}")

# After quantization:
# * Embeddings type: int8
# * Similarities:
#  [[119073  78877]
#   [ 76174  99127]]
```

## Evaluation

Please refer to this <span style="color:red">TBD</span> script to replicate results.
The checkpoint used to produce the scores presented in the paper [is here](https://huggingface.co/MongoDB/mdbr-leaf-ir/commit/ea98995e96beac21b820aa8ad9afaa6fd29b243d).

## Citation  
  
If you use this model in your work, please cite:  
  
```bibtex  
@article{mdb_leaf,  
  title         = {LEAF: Lightweight Embedding Alignment Knowledge Distillation Framework},  
  author        = {Robin Vujanic and Thomas Rueckstiess},  
  year          = {2025}
  eprint        = {TBD},
  archiveprefix = {arXiv},
  primaryclass  = {FILL HERE},
  url           = {FILL HERE}
}  
```  
  
## License  
  
This model is released under Apache 2.0 <span style="color:red">(TBD)</span> License.  
  
## Contact  
  
For questions or issues, please open an issue or pull request. You can also contact the MongoDB ML research team at robin.vujanic@mongodb.com.