vanilla-small-20B

vanilla-small-20B — dense baselines (12L / d=1024, 202.53M)

Dense (sparsity_mode=none) counterparts to the cayley-small-*-mlp_in-20B ParityTransformer siblings. One 20B-token run, five checkpoints: three val-matched early stops plus the token-matched end of the run.

These are the dense comparators for the interpretability claim: if a ParityTransformer is more interpretable than a dense model that is at least as good on val, the comparison cannot be dismissed as a handicapped baseline.

Checkpoints

Files are named by iter. They were previously named by the --save-at-val threshold that fired; those names misstated the actual loss, so they were renamed (ckpt_val_3.1584.pt → ckpt_val_1900.pt, 3.1484 → 2100, 3.1330 → 2300).

Siblings, measured in the same pass on the same batches:

Paired differences vs the siblings

Negative = dense is better. Every model saw byte-identical batches, so these are exact paired comparisons, not error bars added in quadrature (paired SEM ±0.0004).

Note ckpt_val_1900 clears 2L but does not clear 3L. Pick the checkpoint your comparison requires from this table.

ckpt_final_20B.pt is also better than ckpt_best.pt by 0.0029 (110σ) despite the name: ckpt_best was selected by a noisy 4.9M-token in-run eval that happened to dip at iter 12200. Prefer ckpt_final_20B.pt for everything; ckpt_best.pt is kept only for provenance.

How these numbers were measured

evals/language_modeling/eval_val_loss.py over the entire FineWeb-Edu val split — contiguous non-overlapping windows, bf16 autocast, batch 16 × 1024, fineweb-edu-25B, seed 0.

The paired SEM is ~13× tighter because shared passage difficulty cancels. Quote the paired differences.

Two reasons the run log disagrees with this table:

Measured on this protocol, dense sits ~0.0025 nats below its in-run trigger, while the cayley siblings sit ~0.011 above their published numbers. The cayley siblings' published numbers are in-run values too, so use the measured column for any cross-model claim and do not mix the two — a ~0.01 nat protocol artifact otherwise lands inside the very gap being resolved.

One caveat we cannot yet separate: the cayley path must run bf16 through this loader (it builds the backbone fp32 but forces the SAE to bf16 + triton), so its offset could be protocol or genuine bf16 sensitivity in top-k selection. Dense is bf16-insensitive here — its fp32 and autocast runs agree to 0.0002.

Alignment tax

Equal budget (20B tokens both sides):

Token efficiency — where dense first reaches each sibling's measured val, interpolated on the in-run eval curve shifted by the measured −0.0025 protocol offset:

Both inside the <15× target. Eval spacing (100 iters = 157M tokens) caps resolution, so treat as two significant figures. Note the large chassis gives ~3.5× on the same construction (markhenry/vanilla-large-20B) — the ratio is not a constant across scale.

Corrections to earlier versions of this card / report.md

All five corrections move the same way: measured consistently, the dense baselines are stronger than the in-run numbers implied, so the alignment tax is larger. That is the conservative direction for the interpretability claim.

Architecture and training

12 layers, 8 heads, d=1024, block_size 1024, learned positional encoding, RMSNorm (affine), no bias, vocab 50304, 202.53M params — identical to the cayley-small siblings, which add buffers but no nn.Parameters.

Muon + AdamW lockstep 1.2e-2 (min 1.5e-4), linear_warmdown wf=0.9, warmup 200, tokens/iter = 32 × 1024 × 48 = 1,572,864 → 12,716 iters. FineWeb-Edu sample-100BT, GPT-2 tokenizer. 2× H200, 8h32m, ~695k tok/s. Wandb 9iezcto0.

--n-head 8 must be passed explicitly: train.py arch defaults are mode-dependent, so sparsity_mode=none at d=1024 otherwise resolves n_head=12 and 1024 % 12 != 0 asserts. The cayley siblings hit the cayley default of 8.

Also in this repo: report.md, eval_curve.csv (full in-run eval trajectory), train.log, and the training script.

Loading

import torch
ck = torch.load("ckpt_val_2100.pt", map_location="cpu", weights_only=False)
ck["iter_num"]                              # 2100
cfg, sd = ck["model_config"], ck["model"]   # GPTConfig fields, state dict

Checkpoints carry optimizer_states, so they are resumable, not inference-only.

Siblings

• markhenry/cayley-small-2L-mlp_in-20B — measured 3.1671
• markhenry/cayley-small-3L-mlp_in-20B — measured 3.1425
• markhenry/vanilla-large-20B — same construction, large chassis

---

Re-evaluation and card by Claude Code. Training run and report.md predate it; where they disagree, the corrections table above states why.