cayley-131k-3L-mlp_in-v2
205M-parameter GPT with a three-level CayleySAE sparsity bottleneck inserted
at mlp_in in every transformer block. 131,072 leaf features per layer.
Best val loss: 3.1309 (FineWeb-Edu, iter 1000 of 1907).
Successor to
aemack-org/cayley-10b
(3L, 65k leaves, 3.173 val) and the previous
markhenry/cayley-131k-3L-mlp_in-v1 (3L, 131k leaves, 3.1514 val — the
v4-fullrun snapshot). This is an improvement of
~0.020 nats on the
prior 131k-3L high score.
Architecture
| Parameter | Value |
|---|
| n_layer | 12 |
| n_head | 8 |
| n_embd | 1024 |
| block_size | 1024 |
| vocab_size | 50304 |
| bias | False |
| norm | RMSNorm (affine) |
| MLP | GELU, 4× expansion |
| tokenizer | GPT-2 (tiktoken) |
| dtype | bfloat16 |
| sparsity_mode | cayley |
| cayley_locations | mlp_in |
| cayley_levels | 10,16,0; 15,32,256; 17,64,256 |
| cayley_per_parent_budget | True |
| cayley_score_standardize | True |
CayleySAE hierarchy — three levels, [L, k, delta] per row, where m = 2**L
is the features at that level, k is top-k per token, and delta is the
per-parent budget at child levels:
- L0: 1024 features, k=16
- L1: 32,768 features, k=32, per-parent budget 256
- L2: 131,072 features, k=64, per-parent budget 256
Total 164,864 features per layer, 112 active per token (~1470× overcomplete
at the L2 leaf dictionary).
Training provenance
This checkpoint is the product of a two-stage recipe: a fresh end-to-end
pretraining run (v6-fullrun-floor) followed by a landing-LR fork probe
that reloaded v6's best checkpoint with a fresh optimizer and a new warmdown
schedule (v12).
Stage 1 — cayley-131k-3L-mlp_in-v6-fullrun-floor (fresh init, the parent)
| Parameter | Value |
|---|
| init | fresh (random) |
| optimizer | Muon (hidden 2D) + AdamW (embeddings, biases) |
| peak LR (muon / adamw) | 1.2e-2 / 1.2e-2 |
| min LR | 1.5e-4 |
| lr_schedule | linear_warmdown, warmdown_frac=0.9 |
| warmup_iters | 200 |
| batch_size | 24 |
| grad_accum_steps | 64 |
| max_iters | 9537 |
| tokens seen | ~15.0B |
| dataset | FineWeb-Edu-100B (first 15B tokens) |
| best_val_loss | 3.1347 (iter ~9000) |
| wall clock | ~4h51m on 16× RTX 5090 |
The v6-fullrun-floor recipe raised the minimum LR from v1's (v4-fullrun's)
5e-5 to 1.5e-4 after observing that the last ~130 iters of v4 at LR below
2.77e-4 produced no productive work — those tokens were wasted below the
useful LR band. Combined with a 25% token-budget extension (12B → 15B), this
shifted ~1500 additional iterations into the productive LR range.
Stage 2 — cayley-131k-3L-mlp_in-v12-landing-1e3-ga128-6b (fork-probe, this checkpoint)
The v6 ckpt_best.pt was loaded, its optimizer state was discarded, and a
fresh 1907-iter warmdown schedule was run from a new (lower) peak LR. Token
budget was doubled vs the typical fork-probe template (3B → 6B) and gradient
accumulation was doubled (64 → 128), so each iter covers 2× tokens while the
iter grid matches the earlier fork probes (v9/v10/v11) — enabling direct
eval-by-eval comparison across the full sweep.
| Parameter | Value |
|---|
| init | fork of v6's ckpt_best.pt (optimizer discarded) |
| optimizer | Muon (hidden 2D) + AdamW (embeddings, biases) |
| peak LR (muon / adamw) | 1e-3 / 1e-3 |
| min LR | 5e-5 |
| lr_schedule | linear_warmdown, warmdown_frac=0.8 |
| warmup_iters | 50 |
| batch_size | 24 |
| grad_accum_steps | 128 |
| tokens / iter | ~3.15M |
| max_iters | 1907 |
| tokens seen (stage 2) | ~6.0B |
| cumulative tokens | ~21.0B (15B v6 + 6B v12) |
| dataset | FineWeb-Edu-100B |
| best_val_loss | 3.1309 (iter 1000) |
| wall clock (stage 2) | ~118 min on 16× RTX 5090 |
The fork-probe sweep (v9 @ 5e-4, v10 @ 1e-3, v11 @ 2e-3, v12 @ 1e-3 + ga128)
established a noise-floor mean of ~3.135 across all fork-landing configurations;
v12 was the ga128 "settle-it" replay of the best recipe and represents the most
robust point inside that floor (per-eval std 0.0054, vs 0.0116 for v10 at ga64).
The full analysis is written up in mh/reports/20-training-tips.md §13 and §13a
in the source repo.
Purpose
Primary interpretability subject model for 131k-3L CayleySAE experiments.
Successor to
markhenry/vanilla-v5-parity
and
aemack-org/cayley-10b,
used for feature-profiling, circuit-discovery, and alignment-tax comparisons
against the vanilla baselines.
