qcz/muon-retriever-finetuning

Comparing Muon, NorMuon and AdamW for fine-tuning a dense retriever

Final checkpoints (step 3,907) of every run in the blog post Comparing Muon, NorMuon and AdamW for Fine-tuning a Dense Retriever. Code, configs and all result tables: github.com/qcznlp/muon-retriever-finetuning.

Every model starts from lightonai/DenseOn-unsupervised and is fine-tuned for one epoch on 500K queries from lightonai/embeddings-fine-tuning (one positive and seven hard negatives per query, InfoNCE at temperature 0.02, batch 128). The runs differ only in the optimizer and its learning rates. They are research checkpoints for comparing optimizers.

Load a model

Each run is a Sentence Transformers model in its own folder:

from huggingface_hub import snapshot_download
from sentence_transformers import SentenceTransformer

name = "muon-5e-4-seed42"
path = snapshot_download("qcz/muon-retriever-finetuning", allow_patterns=f"{name}/*")
model = SentenceTransformer(f"{path}/{name}")
queries = model.encode(["what does the muon optimizer do"], prompt_name="query")
documents = model.encode(["Muon orthogonalizes the momentum of each hidden weight matrix."], prompt_name="document")

Runs

Muon and NorMuon update the 88 hidden weight matrices; the token embeddings and norm weights are trained by AdamW exactly like the best AdamW run (fourth column). Val loss = contrastive loss on 4,096 held-out queries, which selected the learning rates. BEIR = mean nDCG@10 (x100) over 14 decontaminated BEIR tasks.

Folder	Optimizer	Hidden-matrix LR	Embeddings and norms	Seed	Val loss	BEIR
adamw-1e-6-seed42	AdamW	1e-6	same as hidden	42	0.3428	56.50
adamw-3e-6-seed42	AdamW	3e-6	same as hidden	42	0.2662	58.34
adamw-1e-5-seed42	AdamW	1e-5	same as hidden	42	0.2148	58.86
adamw-3e-5-seed42	AdamW	3e-5	same as hidden	42	0.1857	58.93
adamw-5e-5-seed42	AdamW	5e-5	same as hidden	42	0.1799	59.18
adamw-1e-4-seed42	AdamW	1e-4	same as hidden	42	0.1909	58.93
adamw-5e-5-seed2027	AdamW	5e-5	same as hidden	2027	0.1844	59.13
adamw-5e-5-seed3407	AdamW	5e-5	same as hidden	3407	0.1805	59.00
muon-1e-4-seed42	Muon	1e-4	5e-5	42	0.2054	59.04
muon-2e-4-seed42	Muon	2e-4	5e-5	42	0.1859	59.16
muon-3e-4-seed42	Muon	3e-4	5e-5	42	0.1782	59.21
muon-5e-4-seed42	Muon	5e-4	5e-5	42	0.1776	59.08
muon-1e-3-seed42	Muon	1e-3	5e-5	42	0.1978	58.92
muon-3e-3-seed42	Muon	3e-3	5e-5	42	0.3472	55.69
muon-5e-4-seed2027	Muon	5e-4	5e-5	2027	0.1812	59.15
muon-5e-4-seed3407	Muon	5e-4	5e-5	3407	0.1757	59.37
normuon-1e-4-seed42	NorMuon	1e-4	5e-5	42	0.2074	58.98
normuon-2e-4-seed42	NorMuon	2e-4	5e-5	42	0.1865	59.18
normuon-3e-4-seed42	NorMuon	3e-4	5e-5	42	0.1790	59.11
normuon-5e-4-seed42	NorMuon	5e-4	5e-5	42	0.1763	59.06
normuon-1e-3-seed42	NorMuon	1e-3	5e-5	42	0.1920	59.02
normuon-3e-3-seed42	NorMuon	3e-3	5e-5	42	0.3290	56.03
normuon-5e-4-seed2027	NorMuon	5e-4	5e-5	2027	0.1802	59.23
normuon-5e-4-seed3407	NorMuon	5e-4	5e-5	3407	0.1752	59.15

Each folder also contains trainer_state.json with the logged training loss.