JaydeepR/clm-8b-repro-heads

CLM-8B reproduction heads

An independent reproduction of CLM-8B ("Contrastive Language Models", Kwok et al., Stanford and NVIDIA Research, 2026): two small projection heads on top of a frozen Qwen3-8B. Given a situation and a list of options, the model scores every option: the situation and each option go through the encoder once, and no text is generated.

Question / context -> Qwen3-8B (frozen, last token) -> State head  -> 512-d ┐
Option             -> Qwen3-8B (frozen, last token) -> Action head -> 512-d ┴ cosine x learned scale -> softmax

Each head is 4096 -> 1536 -> 1536 (LayerNorm) -> 512 with GELU, 18.9M trained parameters in total. Training code, logs and the full write-up are at github.com/jaydeepraijada/clm-experiments. This is not an official release of the CLM authors; their heads are at Contrastive-LM/CLM-v0.1-8B and jackyk02/qwen3_8b_head.

Files

File	Stage	Notes
clm-8b-repro-post.pt	final, post-trained (step 400)	the head to use
stages/mid_best.pt	mid-trained (step 1,344)	best tricky-Q&A model before agent data
stages/pre_best.pt	pre-trained (step 30,000)
stages/post_v2_best.pt	post-trained, ablation	agent actions embedded as bare decisions; see below
results/*.json		same-code comparisons against the authors' heads
logs/*.jsonl		evaluation logs during training

Checkpoints use the CLM format (state_head, action_head, logit_scale, cfg), so they load in the authors' server:

pip install git+https://github.com/Contrastive-LM/CLM.git
vllm serve Qwen/Qwen3-8B --runner pooling --max-model-len 2048 --port 8090   # the encoder, revision b968826d
clm-serve --ckpt clm-8b-repro-post.pt

They are torch.save pickles: load them only from a source you trust (torch.load(..., weights_only=False)). A head only works with the encoder and pooling it was trained against: Qwen3-8B at revision b968826d9c46dd6066d109eabc6255188de91218, last-token pooling, L2-normalised.

Results

Same evaluation code and held-out data for every head. "Authors" are the CLM authors' released heads.

Test	Published	Authors' head	This model
Tricky Q&A top-1 after pre-training	52.1%	51.7%	50.7%
Tricky Q&A top-1 after mid-training	69.2%	69.2%	68.5%
Tricky Q&A top-1 after post-training	68.5%	68.6%	69.0%
General Q&A top-1 (1,024-way)		83.6%	82.7%
DeepSWE best-of-4, seed 1234	31/38	31/38	30/38
DeepSWE best-of-4, mean of 5 seeds		31.0/38	30.8/38
T-Rex median decision latency in the 60 FPS game loop (L40S)	16.5 ms	17.0 ms	17.1 ms
T-Rex courses survived, safety shield on	5/5	5/5	3/5
T-Rex agreement with the game's planner, shield on	65.8%	65.3%	2.7%

• Tricky Q&A: 191,161 held-out questions (Nemotron DQA shard_00 chunks 28-29), each with its right answer and up to 10 look-alike wrong answers. Frozen Qwen3-8B with plain cosine and no heads scores 14.4%.
• DeepSWE: the head is fine-tuned on DeepSWE rollouts with the authors' script, then picks one of 4 candidate solutions for each of 38 held-out tasks (random picking solves 28, a perfect picker 34).

Limitations

• Game-style decisions (T-Rex) are much worse than the authors'. Our pre- and mid-trained heads play like the authors' mid-trained head; the authors' post-training teaches the skill and ours does not. Post-training data recipe is the likely cause: the authors' 8B post-training data and renderer are not public, and we used a reconstruction (build_agent_data.py). Embedding bare actions (post_v2) helped only slightly.
• Hand-written prompts are unreliable. In a small informal check on 8 hand-written decisions (ticket routing, document questions, next agent step), this model picked right 1 of 8 and the authors' head 2 of 8. Use it on tasks close to its training data, or fine-tune it (as for DeepSWE), and measure before relying on it.
• English text only.

Training

All on Modal, about $82 in total.

Stage	Data	Steps	Batch	LR	Time
Pre	62.1M Q&A pairs (InfoNCE, in-batch negatives, same-document pairs masked)	30,000	16,384	0.00653	2.4 h, 1x L40S
Mid	shard_00 chunks 0-27 with up to 10 hard negatives per question	1,344	16,384	0.001	57 min
Post	60% agent steps (914,106 pairs from 39 ADP v2 configs) + 40% replay of hard-negative Q&A	400	16,384	0.0003	7 min

AdamW, no weight decay, gradient clip 1.0, learned temperature capped at scale 100. The agent pairs (at most 6 random steps per recording, states cut to their last 2,047 tokens in the Qwen chat template) were embedded with Qwen3-8B on H100s, about 1.56B tokens.

Licence and data

The heads are released under Apache 2.0, like the CLM code. The training pipeline is adapted from VizuaraAI/system-one-from-scratch (Apache 2.0). They were trained on vectors derived from Nemotron-CC Diverse QA (NVIDIA Open Data License) (via Contrastive-LM/CLM-v0.1-Pretrain-Nemotron and jackyk02/nemotron_dqa_qwen3_8b_neg_embeddings) and from the Agent Data Protocol collection (neulab/adp-v2), whose subsets carry their own licences. Check those terms for your use.

Citation

Please cite the original work:

@misc{kwok2026contrastivelanguagemodels,
  title={Contrastive Language Models},
  author={Jacky Kwok and Hangoo Kang and Tarun Suresh and Jon Saad-Falcon and Marco Pavone and Christopher Ré and Azalia Mirhoseini},
  year={2026},
  note={Notion Blog},
  url={https://contrastive-lm.notion.site}
}