MongoDB
/

mdbr-leaf-ir

@@ -1,148 +1,276 @@
----
-tags:
-- sentence-transformers
-- sentence-similarity
-- feature-extraction
-- dense
-base_model: sentence-transformers/all-MiniLM-L6-v2
-pipeline_tag: sentence-similarity
-library_name: sentence-transformers
----
-# SentenceTransformer based on sentence-transformers/all-MiniLM-L6-v2
-This is a [sentence-transformers](https://www.SBERT.net) model finetuned from [sentence-transformers/all-MiniLM-L6-v2](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2). It maps sentences & paragraphs to a 768-dimensional dense vector space and can be used for semantic textual similarity, semantic search, paraphrase mining, text classification, clustering, and more.
-## Model Details
-### Model Description
-- **Model Type:** Sentence Transformer
-- **Base model:** [sentence-transformers/all-MiniLM-L6-v2](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2) <!-- at revision c9745ed1d9f207416be6d2e6f8de32d1f16199bf -->
-- **Maximum Sequence Length:** 512 tokens
-- **Output Dimensionality:** 768 dimensions
-- **Similarity Function:** Cosine Similarity
-<!-- - **Training Dataset:** Unknown -->
-<!-- - **Language:** Unknown -->
-<!-- - **License:** Unknown -->
-### Model Sources
-- **Documentation:** [Sentence Transformers Documentation](https://sbert.net)
-- **Repository:** [Sentence Transformers on GitHub](https://github.com/UKPLab/sentence-transformers)
-- **Hugging Face:** [Sentence Transformers on Hugging Face](https://huggingface.co/models?library=sentence-transformers)
-### Full Model Architecture
-```
-SentenceTransformer(
-  (0): Transformer({'max_seq_length': 512, 'do_lower_case': False, 'architecture': 'BertModel'})
-  (1): Pooling({'word_embedding_dimension': 384, 'pooling_mode_cls_token': False, 'pooling_mode_mean_tokens': True, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False, 'pooling_mode_weightedmean_tokens': False, 'pooling_mode_lasttoken': False, 'include_prompt': True})
-  (2): Dense({'in_features': 384, 'out_features': 768, 'bias': True, 'activation_function': 'torch.nn.modules.linear.Identity'})
-  (3): Normalize()
-)
-```
-## Usage
-### Direct Usage (Sentence Transformers)
-First install the Sentence Transformers library:
-```bash
-pip install -U sentence-transformers
-```
-Then you can load this model and run inference.
-```python
-from sentence_transformers import SentenceTransformer
-# Download from the 🤗 Hub
-model = SentenceTransformer("sentence_transformers_model_id")
-# Run inference
-queries = [
-    "Which planet is known as the Red Planet?",
-]
-documents = [
-    "Venus is often called Earth's twin because of its similar size and proximity.",
-    'Mars, known for its reddish appearance, is often referred to as the Red Planet.',
-    'Saturn, famous for its rings, is sometimes mistaken for the Red Planet.',
-]
-query_embeddings = model.encode_query(queries)
-document_embeddings = model.encode_document(documents)
-print(query_embeddings.shape, document_embeddings.shape)
-# [1, 768] [3, 768]
-# Get the similarity scores for the embeddings
-similarities = model.similarity(query_embeddings, document_embeddings)
-print(similarities)
-# tensor([[0.4166, 0.6312, 0.5094]])
-```
-<!--
-### Direct Usage (Transformers)
-<details><summary>Click to see the direct usage in Transformers</summary>
-</details>
--->
-<!--
-### Downstream Usage (Sentence Transformers)
-You can finetune this model on your own dataset.
-<details><summary>Click to expand</summary>
-</details>
--->
-<!--
-### Out-of-Scope Use
-*List how the model may foreseeably be misused and address what users ought not to do with the model.*
--->
-<!--
-## Bias, Risks and Limitations
-*What are the known or foreseeable issues stemming from this model? You could also flag here known failure cases or weaknesses of the model.*
--->
-<!--
-### Recommendations
-*What are recommendations with respect to the foreseeable issues? For example, filtering explicit content.*
--->
-## Training Details
-### Framework Versions
-- Python: 3.12.7
-- Sentence Transformers: 5.1.0
-- Transformers: 4.52.4
-- PyTorch: 2.6.0+cu126
-- Accelerate: 1.2.1
-- Datasets: 3.1.0
-- Tokenizers: 0.21.0
-## Citation
-### BibTeX
-<!--
-## Glossary
-*Clearly define terms in order to be accessible across audiences.*
--->
-<!--
-## Model Card Authors
-*Lists the people who create the model card, providing recognition and accountability for the detailed work that goes into its construction.*
--->
-<!--
-## Model Card Contact
-*Provides a way for people who have updates to the Model Card, suggestions, or questions, to contact the Model Card authors.*
--->

+---
+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.png" alt="MongoDB Logo" style="height: 36px; width: auto;">
+        <span style="font-size: 32px; font-weight: bold">MongoDB/mdbr-leaf-ir</span>
+    </div>
+</div>
+**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 [`mdb-leaf-mt`](https://huggingface.co/MongoDB/mdb-leaf-mt) model.
+**Note**: this model has been developed by MongoDB Research and is not part of MongoDB's commercial offerings.
+## 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]]
+```
+## 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 [email protected].

logo.png ADDED Viewed