mmbert-rerank-32k-2d-matryoshka

mmbert-rerank-32k-2d-matryoshka

A 32K context multilingual reranker with 2D Matryoshka support, enabling flexible layer and dimension reduction for efficient inference.

Model Description

This model is a cross-encoder reranker built on mmbert-32k-yarn, a ModernBERT model extended to 32K context via YaRN position interpolation and trained on 1800+ languages from the Glot500 corpus.

Key Features

2D Matryoshka Architecture

This model implements 2D Matryoshka Representation Learning, allowing you to trade quality for speed by using:

• Fewer layers: 3, 6, 11, or 22 (full)
• Fewer dimensions: 64, 128, 256, 512, or 768 (full)

This provides 20 different quality-speed configurations from a single model.

Benchmark Results

Long Document Reranking (Key Advantage)

We evaluate on synthetic long documents (~1500 tokens) where the relevant answer is placed at different positions. This tests whether models can use their full context window.

Key Finding: BGE truncates inputs at 512 tokens, causing it to completely miss relevant information beyond that limit. mmbert-rerank maintains 100% accuracy with its 32K context window.

Latency Benchmark

Measured on AMD MI300X GPU:

*BGE truncates at 512 tokens, so it doesn't actually process long docs

Multilingual Performance

Tested on 8 languages including low-resource African languages (Swahili, Yoruba, Amharic, Hausa, Igbo):

BEIR Benchmark (Short Documents)

On standard BEIR datasets with short documents (where BGE's 512-token limit is sufficient):

On short documents, BGE is slightly better due to its focused training on shorter contexts.

When to Use This Model

Usage

Installation

pip install transformers torch

Basic Usage

import torch
from transformers import AutoTokenizer

# Load model (custom class required)
import sys
sys.path.append("path/to/scripts")
from train_rerank import Matryoshka2DReranker

model = Matryoshka2DReranker.from_pretrained("vllm-sr/mmbert-rerank-32k-2d-matryoshka")
tokenizer = AutoTokenizer.from_pretrained("vllm-sr/mmbert-rerank-32k-2d-matryoshka")

model.eval()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

# Score query-passage pairs
pairs = [
    ("What is machine learning?", "Machine learning is a subset of AI that enables systems to learn from data."),
    ("What is machine learning?", "The weather is sunny today."),
]

scores = model.compute_score(pairs, tokenizer, normalize=True)
# Output: [0.9686, 0.0017]

Using Reduced Layers/Dimensions (2D Matryoshka)

# Full model (22 layers, 768 dims) - highest quality
scores = model.compute_score(pairs, tokenizer, normalize=True)

# Reduced layers (use layer 11 for good balance)
scores = model.compute_score(pairs, tokenizer, layer_idx=11, normalize=True)

# Reduced dimensions
scores = model.compute_score(pairs, tokenizer, dim_idx=256, normalize=True)

# Both reduced (fastest)
scores = model.compute_score(pairs, tokenizer, layer_idx=11, dim_idx=256, normalize=True)

Reranking Pipeline

def rerank(query: str, passages: list[str], top_k: int = 10) -> list[tuple[str, float]]:
    """Rerank passages by relevance to query."""
    pairs = [(query, p) for p in passages]
    scores = model.compute_score(pairs, tokenizer, normalize=True)
    
    # Sort by score descending
    ranked = sorted(zip(passages, scores), key=lambda x: x[1], reverse=True)
    return ranked[:top_k]

# Example
query = "How does photosynthesis work?"
passages = [
    "Photosynthesis is the process by which plants convert sunlight into energy.",
    "The stock market closed higher today.",
    "Plants use chlorophyll to absorb light during photosynthesis.",
    "Python is a popular programming language.",
]

results = rerank(query, passages, top_k=2)
# [("Photosynthesis is the process...", 0.95), ("Plants use chlorophyll...", 0.87)]

Training Details

Training Data

• Dataset: cfli/bge-m3-data
• Samples: 1,569,864 query-passage pairs
• Languages: 100+ languages from multilingual sources

Training Configuration

2D Matryoshka Configuration

Comparison with Competitors

Model Architecture

Matryoshka2DReranker
├── encoder: ModernBertModel (308M params)
│   └── 22 transformer layers
│   └── hidden_size: 768
│   └── YaRN position interpolation (32K context)
├── final_norm: LayerNorm (applied to intermediate layers)
└── layer_heads: ModuleDict
    ├── "3": {"64": Linear, "128": Linear, ..., "768": Linear}
    ├── "6": {"64": Linear, "128": Linear, ..., "768": Linear}
    ├── "11": {"64": Linear, "128": Linear, ..., "768": Linear}
    └── "22": {"64": Linear, "128": Linear, ..., "768": Linear}

Limitations

• Training data: While the model supports 32K tokens, training was primarily on shorter passages from BGE-M3 data. Performance on very long passages relies on the base model's YaRN-extended context.
• Domain-specific data: Model was trained on general web data; may require fine-tuning for specialized domains.
• Layer reduction speedup: Due to the need for all hidden states, layer reduction provides minimal speedup without early-exit modifications.

License

Apache 2.0

Citation

@misc{mmbert-rerank-2025,
  title={mmbert-rerank-32k-2d-matryoshka: A Long-Context Multilingual Reranker with 2D Matryoshka},
  author={LLM Semantic Router Team},
  year={2025},
  url={https://huggingface.co/vllm-sr/mmbert-rerank-32k-2d-matryoshka}
}

Acknowledgments

• ModernBERT for the base architecture
• BGE-M3 for training data
• Matryoshka Representation Learning for the dimensional reduction approach