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

Contrastive projection head on frozen Qwen3-8B embeddings — pre-trained (in-batch InfoNCE)

Training notes

Depth-2 ladder protocol (plain GELU MLP, no LayerNorm, batch 1024, constant LR 3.5e-4, 16-epoch cap, early stop on val). This file is the best-validation snapshot (0.7432 at step 846,034); the run's final eval read 0.7399 after three epochs at 0.745 +/- 0.002, so treat the head as ~0.743 +/- 0.003. At this rung the 76M / 151M / 302M 2-layer heads reach 0.7588 / 0.7553 / 0.7399 -- still descending at 302M, so the depth-2 optimum for 51.2B tokens is at or above this size (fit: N* ~ D^1.24). A 3-layer head reaches a lower loss at a quarter of the parameters; this is the best 2-layer head, the lowest-validation-loss checkpoint of the big-head ablation.

Files

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

Loading

import torch
ck = torch.load("run_bh623_w32768.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