llm-semantic-router / mmbert-embed-32k-2d-matryoshka

mmBERT-Embed-32K-2D-Matryoshka

A multilingual embedding model with 32K context window and 2D Matryoshka support for flexible efficiency-quality tradeoffs.

Model Highlights

Key Results

What is 2D Matryoshka?

This model supports two dimensions of flexibility :

1. Dimension Reduction (Matryoshka): Truncate embeddings to smaller dimensions with minimal quality loss
2. Layer Reduction (Adaptive): Use intermediate layer outputs for faster inference

Usage

Basic Usage (Sentence Transformers)

from sentence_transformers import SentenceTransformer

# Load model
model = SentenceTransformer("llm-semantic-router/mmbert-embed-32k-2d-matryoshka")

# Encode sentences
sentences = [
    "This is a test sentence.",
    "这是一个测试句子。",
    "Dies ist ein Testsatz.",
]
embeddings = model.encode(sentences)
print(embeddings.shape)  # (3, 768)

Matryoshka Dimension Reduction

import torch.nn.functional as F

# Encode with full dimensions
embeddings = model.encode(sentences, convert_to_tensor=True)

# Truncate to smaller dimension (e.g., 256)
embeddings_256d = embeddings[:, :256]
embeddings_256d = F.normalize(embeddings_256d, p=2, dim=1)

# Or truncate to 64 dimensions for maximum compression
embeddings_64d = embeddings[:, :64]
embeddings_64d = F.normalize(embeddings_64d, p=2, dim=1)

Long Context (up to 32K tokens)

# For long documents, set max_seq_length
model.max_seq_length = 8192  # or up to 32768

long_document = "..." * 10000  # Very long text
embedding = model.encode(long_document)

Layer Reduction (Advanced)

For latency-critical applications, you can extract embeddings from intermediate layers:

from transformers import AutoModel, AutoTokenizer
import torch
import torch.nn.functional as F

model = AutoModel.from_pretrained(
    "llm-semantic-router/mmbert-embed-32k-2d-matryoshka",
    trust_remote_code=True,
    output_hidden_states=True
)
tokenizer = AutoTokenizer.from_pretrained("llm-semantic-router/mmbert-embed-32k-2d-matryoshka")

# Encode
inputs = tokenizer(sentences, padding=True, truncation=True, return_tensors="pt")
with torch.no_grad():
    outputs = model(**inputs)
    
    # Use layer 6 for 3.3× speedup (56% quality)
    hidden = outputs.hidden_states[6]
    hidden = model.final_norm(hidden)
    
    # Mean pooling
    mask = inputs["attention_mask"].unsqueeze(-1).float()
    pooled = (hidden * mask).sum(1) / mask.sum(1)
    embeddings = F.normalize(pooled, p=2, dim=1)

Evaluation Results

MTEB Benchmark (24 tasks)

STS Benchmark

2D Matryoshka Quality Matrix (STS)

Long-Context Retrieval (4K tokens)

Throughput (AMD MI300X)

Latency Comparison vs BGE-M3 and Qwen3-Embedding-0.6B

mmBERT-Embed is significantly faster due to:

1. Flash Attention 2 - BGE-M3 lacks FA2 (O(n) vs O(n²) attention)
2. Encoder architecture - Qwen3 uses decoder with causal masking
3. Smaller model - 307M vs 569M/600M params

Batch Size = 1

Batch Size = 8

Key insight : The FA2 advantage grows with sequence length and batch size:

• At short sequences (512), BGE-M3 is faster (no FA2 overhead)
• At 2K+ tokens, mmBERT pulls ahead significantly
• At 4K batch=8: mmBERT is 3.1× faster than BGE-M3

Benchmarked on AMD MI300X, bf16 precision.

Training

Data

Trained on BAAI/bge-m3-data (73GB, 279 JSONL files) with:

• Multilingual triplets (query, positive, negative)
• Diverse domains and languages

Configuration

• Base Model : llm-semantic-router/mmbert-32k-yarn
• Loss : Matryoshka2dLoss (combines AdaptiveLayerLoss + MatryoshkaLoss)
• Matryoshka Dimensions : [768, 512, 256, 128, 64]
• Epochs : 1
• Batch Size : 16 (effective 32 with gradient accumulation)
• Learning Rate : 2e-5
• Max Sequence Length : 32,768
• Hardware : AMD Instinct MI300X

Use Cases

When to Use mmBERT-Embed

1. Multilingual RAG for 1800+ languages (especially low-resource languages not covered by Qwen3 or BGE-M3)
2. Long-document retrieval where chunking loses cross-section relationships
3. Edge deployment where 307M params matters vs 600M+
4. Flexible inference where you need to trade quality for speed/storage at runtime

When to Use Alternatives

• Maximum quality on major languages : Qwen3-Embed-8B
• Production stability : BGE-M3 (more battle-tested)
• Very short texts only : Smaller models may suffice

Limitations

• Layer reduction quality (56% at 6L) is lower than full model; use for latency-critical applications where moderate quality is acceptable
• MTEB mean (61.4) is slightly below BGE-M3 (64.5) but with 4× longer context and 2D flexibility
• Optimized for retrieval tasks; may need fine-tuning for other downstream tasks

Citation

@misc{mmbert-embed-2d-matryoshka,
  title={mmBERT-Embed: Multilingual Embedding Model with 2D Matryoshka Training},
  author={vLLM Semantic Router Team},
  year={2025},
  url={https://huggingface.co/llm-semantic-router/mmbert-embed-32k-2d-matryoshka}
}

License

Apache 2.0