MiniLM-L6, contrastive training on MS MARCO

Sentence embedding model for passage retrieval, trained from a raw MiniLM that had never learned to compare sentences. It is a learning exercise from the embeddings and RAG step of mistral-optimisation. For real use, all-MiniLM-L6-v2 (same architecture, a billion training pairs) is much better.

Results

1,000 MS MARCO v2.1 validation queries; the corpus is every passage proposed for these queries (9,984). Recall@K is the share of queries whose relevant passage is in the top K (FAISS, inner product).

Model	Recall@1	Recall@5	Recall@10
Raw MiniLM (starting point)	1.0%	3.2%	3.9%
Unsupervised SimCSE, same data	3.3%	8.6%	12.1%
This model	25.6%	73.7%	87.9%
all-MiniLM-L6-v2 (reference)	42.2%	89.2%	98.5%

Training

Base: nreimers/MiniLM-L6-H384-uncased
20,000 (query, relevant passage) pairs from MS MARCO v2.1 train, 1 epoch, batch 64, lr 3e-5
InfoNCE loss with in-batch negatives, temperature 0.05
Embedding: mean of token vectors (padding excluded), L2-normalized, max 128 tokens

Usage

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

tok = AutoTokenizer.from_pretrained("Falcot922/minilm-msmarco-contrastive")
model = AutoModel.from_pretrained("Falcot922/minilm-msmarco-contrastive").eval()

def embed(texts):
    enc = tok(texts, padding=True, truncation=True, max_length=128, return_tensors="pt")
    with torch.no_grad():
        hidden = model(**enc).last_hidden_state
    mask = enc["attention_mask"].unsqueeze(-1)
    return F.normalize((hidden * mask).sum(1) / mask.sum(1), dim=-1)

q = embed(["what is the capital of france"])
p = embed(["Paris is the capital of France.", "Bananas are yellow."])
print(q @ p.T)

Loading…