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

file	iter	tokens	note
ckpt_final.pt	2,730	4.294B	the deliverable; fully warmed down
ckpt_iter002000.pt	2,000	3.146B	trajectory snapshot, un-annealed

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

architecture	12 layers, d=1024, 8 heads, learned pos enc
bottleneck	TopK SAE at mlp_in, dict 65,536, k=48, unit-norm decoder
params	1,813.94M = 202.53M backbone + 12 × 134.28M dictionary
data	FineWeb-Edu sample-100BT, 4.294B tokens
optimizer	Muon 1.2e-2 / AdamW 1.2e-2 (lockstep), momentum 0.95, ns_steps 5
schedule	linear_warmdown, warmup 200, wf=0.3 (819 warmdown iters from 1,911)
batch	bs8 × ga192 (global), seq 1024 → 1,572,864 tok/iter
hardware	8×A100-SXM4-80GB, 55.5 GB/GPU peak

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