topk-small-65k-wallclock-matched

topk-small-65k-wallclock-matched

A 1,813.94M-parameter 12L/d=1024 GPT with a flat learned-dictionary TopK SAE bottleneck at mlp_in (dictionary 65,536, k=48). This is the learned-dictionary control for CayleySAE / ParityTransformer — same chassis, same insertion point, same sparsity; the dictionary is learned instead of parity-hashed.

What "wall-clock-matched" means here — read before quoting

This model was trained for 9.713 hours on 8×A100, the wall-clock that PT-Small-2L consumed for its 20B-token run.

The hardware is not identical. PT-Small-2L trained at ~572k tok/s on its host; this host runs the identical script at ~800k tok/s — 1.40× faster. So 9.713 h here buys roughly 40% more effective compute than the reference model ever received. That asymmetry is deliberate: the control was given every opportunity to win. Do not describe this as a same-hardware wall-clock match.

At the measured 122,818 tok/s this budget yielded 4.294B tokens (2,730 iters at 1,572,864 tok/iter — the same tok/iter as PT-Small-2L, so iteration counts are directly comparable).

Checkpoints

ckpt_iter002000.pt sits only 89 iters into an 819-iter warmdown (LR still ~1.07e-2). It is a lower bound on what a genuinely ~7.1 h-budgeted run would reach, because the end-of-run anneal is exactly what it lacks. Use it as a trajectory point, not as a "7.5 h arm".

Training

wf=0.3, not PT's 0.9, because at 4.29B tokens against D* ≈ 20 × 1.81B ≈ 36B this run sits at ~0.12× Chinchilla saturation, where a super-Chinchilla warmdown fraction would handicap it. Each arm gets the schedule appropriate to its budget.

Dictionary width is PT-Small-2L's unfolded feature count (32,768 × 2 polarity faces), since polarity faces are the units the interpretability claims are made in.

Efficiency

On identical hardware, PT-Small-2L trains at ~800k tok/s and this model at 122,818 tok/s — 6.5× slower for the same feature count and sparsity, because a learned dictionary materializes 2·d·dict_size parameters per layer where the parity construction adds none.

Provenance

Training script train_topk_small_65k_wallclock_matched.sh is included in this repo. Full experimental record: mh/reports/45-t6c-wallclock-matched-topk65k.md. wandb: h76qtfjv