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

1. Same backbone limitations (vocab=6400, limited training data)
2. Only trained on T2Reranking — domain transfer to other reranking tasks untested
3. MoE version not yet trained

🙏 Acknowledgements

• Base: MiniMind by @jingyaogong
• Route: Qwen3-Reranker
• Data: C-MTEB/T2Reranking
• Eval: MTEB

📄 License

Apache 2.0