MiniMind-Rerank-Dense
A cross-encoder reranker built on the MiniMind-3 small language model (64M params), with the design fully aligned with Qwen3-Reranker. Part of the minimind-embedding project.
📌 How It Works
💡 Instead of a separate scoring head, this model reuses the language model's
lm_head: given a query and a document, it predicts whether the next token is是(yes, relevant) or否(no, irrelevant). The probability of是becomes the relevance score.
Prompt template:
<|im_start|>system
判断下面的文档是否符合查询需求,只回复是或否<|im_end|>
<|im_start|>user
查询:{query}
文档:{document}<|im_end|>
<|im_start|>assistant
<think>
</think>
The model predicts the next token after </think>: 是 (id=357) or 否 (id=1332). The softmax probability of 是 is the relevance score (0~1).
🏗️ Model Details
| Item | Value |
|---|---|
| Architecture | MiniMind-3 (causal decoder-only Transformer) |
| Parameters | 64M |
| Hidden size | 768 |
| Layers | 8 |
| Vocab size | 6400 |
| Max sequence length | 8192 |
| Scoring | pointwise yes/no (next-token prediction) |
| Target tokens | 是 (id=357) / 否 (id=1332) |
🎓 Training
Trained on C-MTEB/T2Reranking (80% train split, 20% held-out for evaluation — no data leakage).
| Config | Value |
|---|---|
| Training data | T2Reranking, 80k query-doc pairs |
| Epochs | 3 |
| Batch size | 16 |
| Learning rate | 1e-5 (cosine decay) |
| Loss | Pointwise cross-entropy |
Training loss (red) 0.79→0.35, accuracy (blue) 55%→85%, 3 epochs.
📊 Evaluation
C-MTEB/T2Reranking (MAP@10)
| Method | MAP@10 | Accuracy | Note |
|---|---|---|---|
| Zero-shot (pretrained backbone) | 0.472 | — | baseline |
| Full fine-tuning + Stage3 merge | 0.915 | 85% | ✅ +94% |
| Frozen backbone (lm_head only) | 0.650 | 63% | also works (+38%) |
The Three Bugs Story
The reranker's development hit three cascading bugs that initially made all experiments look like "training makes it worse." Each fix drove the score up:
| Bug | Symptom | Root Cause | Fix | After Fix |
|---|---|---|---|---|
| ① Label mapping reversed | All training decreased score (0.47→0.24) | cross_entropy(logits, labels): label=1 maps to index 1 (否), but intent was "relevant→是"(index 0) | target = 1 - labels | 0.47→0.65 |
| ② Data leakage | Score inflated (0.94) | T2Reranking has only dev split; train and eval used the same data | 80/20 split train/test | 0.94→0.92 (real) |
| ③ Tokenizer fragmentation | Prompt wasted tokens | <Query>/<Document> split into 5 BPE fragments | Switch to Chinese 查询:/文档: | tokens 79→56 |
Key takeaway: the 64M MiniMind backbone is fully capable of reranker fine-tuning. The earlier "catastrophic forgetting" conclusion was a label bug artifact. After fixes, full fine-tuning achieves MAP@10=0.915.
🚀 Usage
import torch
import torch.nn.functional as F
# Load with the project's modeling code
# from shared.model import MiniMindForRerank
# from shared.configs import rerank_dense_64m
# from shared.tokenizer import load_tokenizer, build_rerank_inputs
# model = MiniMindForRerank(rerank_dense_64m())
# model.load_state_dict(torch.load("rerank_768.pth"), strict=False)
# tokenizer = load_tokenizer(padding_side="left")
queries = ["天空为什么是蓝色"]
documents = ["阳光散射使天空呈蓝色。", "红烧肉的做法是..."]
# Build prompt inputs
# enc = build_rerank_inputs(tokenizer, queries, documents, max_length=512)
# Score = probability of 是
# scores = model.predict(enc["input_ids"], enc["attention_mask"])
# Higher score = more relevant
⚠️ Limitations
- Same backbone limitations (vocab=6400, limited training data)
- Only trained on T2Reranking — domain transfer to other reranking tasks untested
- MoE version not yet trained
🙏 Acknowledgements
- Base: MiniMind by @jingyaogong
- Route: Qwen3-Reranker
- Data: C-MTEB/T2Reranking
- Eval: MTEB
📄 License
Apache 2.0