qwen3-8b-dqa-head-302m-midtrained

Contrastive projection head on frozen Qwen3-8B embeddings — mid-trained on hard negatives

Mid-training: before → after

Training notes

Warm-started from the pre-trained 302M checkpoint and mid-trained on Gemini-generated hard negatives (K<=10 per pair, weight 1.0, doc-masked bidirectional InfoNCE), shard_00 chunks 0000-0027, early-stopped on hard-negative val loss (best at epoch 3, stopped at epoch 5). Deviation from the reference recipe: batch 4096 and LR 5e-4 instead of 16384 / 1e-3 -- a 302M head with 10 negatives per row does not fit a 40 GB GPU at the larger batch; all negatives were kept so the loss definition is unchanged. Cost: general val loss rose 0.726 -> 0.745 (+0.019) while best-of-10 selection rose 46.4% -> 67.2% (+20.8 pts).

Files

• best_head.pt — torch.save dict with state_head / action_head state dicts, logit_scale, and the training cfg.

Loading

import torch
ck = torch.load("best_head.pt", map_location="cpu", weights_only=False)
cfg = ck["cfg"]  # width=32768, depth=2, layernorm=False
# build two heads with experiments/scaling_projection_head.make_head(width, gen, device, depth,
#   512, activation=cfg.get("activation","gelu"), layernorm=cfg.get("layernorm",False),
#   residual=cfg.get("residual",False)) and load_state_dict each tower.

Inputs are Qwen3-8B last-token embeddings (4096-d, L2-normalised); the question side prepends the source document. Score = cosine(state_head(q), action_head(a)) × exp(logit_scale).

Code and training logs: https://github.com/jackyk02/contrastive_learning