MiniLM-L6, unsupervised SimCSE on MS MARCO passages

Control model from the embeddings and RAG step of mistral-optimisation. It shows what unsupervised contrastive learning (SimCSE) brings to a raw MiniLM, compared with training on (query, passage) pairs. It is a weak retriever, not meant for real use: for that, see the supervised version minilm-msmarco-contrastive, or better, all-MiniLM-L6-v2.

How SimCSE works here

Each passage is encoded twice with dropout active. The two slightly different vectors form a positive pair, and the other passages of the batch are the negatives. No query and no label is used: the model only learns that a passage should stay close to itself under noise.

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%
This model (SimCSE, unsupervised) 3.3% 8.6% 12.1%
Supervised, same data and settings 25.6% 73.7% 87.9%
all-MiniLM-L6-v2 (reference) 42.2% 89.2% 98.5%

SimCSE triples the raw model's scores, but stays far behind supervision: making a passage close to itself does not teach which passage answers a question.

Training

  • Base: nreimers/MiniLM-L6-H384-uncased
  • 20,000 MS MARCO v2.1 train passages (the relevant passage of each query), 1 epoch, batch 64, lr 3e-5
  • InfoNCE loss with in-batch negatives, temperature 0.05, dropout as the only augmentation
  • 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-simcse")
model = AutoModel.from_pretrained("Falcot922/minilm-msmarco-simcse").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)
Downloads last month
33
Safetensors
Model size
22.7M params
Tensor type
F32
ยท
Inference Providers NEW
This model isn't deployed by any Inference Provider. ๐Ÿ™‹ Ask for provider support

Model tree for Falcot922/minilm-msmarco-simcse

Finetuned
(23)
this model

Dataset used to train Falcot922/minilm-msmarco-simcse