Record: Cosine TTT scheduling with per-layer lr — mean val_bpb=1.0970 (3 seeds)

[mrdavtan]

Summary

val_bpb=1.0970 (3-seed mean, std=0.0010). 15.4-15.8 MB artifact. 8xH100 SXM, FA2.

Seed	Steps	Pre-TTT	Post-TTT	Artifact
1337	7,101	1.1577	1.0959	15.4 MB
42	6,700	1.1588	1.0971	15.5 MB
7	6,987	1.1580	1.0979	15.8 MB

Training architecture follows the community stack. The main change from prior work is in the TTT schedule. All runs used FA2; FA3 Hopper would improve pre-TTT quality through faster training steps. The schedule is independent of the attention kernel and should apply to any architecture.

TTT scheduling

Two modifications to AdamW TTT (PR #442):

Cosine lr decay over 30 epochs. Starts at full lr to address large-scale quantization damage, progressively reduces to refine without overshooting. Flat lr must compromise between these two regimes.

Per-layer lr groups based on quantization damage. MLP output projections showed 3.4× higher relative quantization error than input projections on our trained checkpoint. TTT receives: 3× base lr for output projections, 0.5× for input projections, 1× for the rest. Ratios are model-specific.

TTT_OPTIMIZER=adamw  TTT_LR=0.0005  TTT_EPOCHS=30
TTT_COSINE=1  TTT_PERLAYER=1  TTT_FREEZE_BLOCKS=0
TTT_BATCH_SEQS=64 (per GPU, 512 total with DDP sharding)

30 epochs at ~15.5s/epoch = ~465s total. We also tested flat lr, SGD, focal loss, and KL divergence from pre-quant model. Focal loss and KL divergence did not improve over cross-entropy. Full comparison in the README.

Other findings

See PR #212 for a non-record submission documenting 25+ experiments with negative results on codebook quantization, magnitude pruning, multi-token prediction, embedding factorization, and depth recurrence.

Acknowledgments

• Training architecture: PRs Record: Int6 MLP3x + SmearGate + BigramHash + MuonWD + SWA (mean val_bpb=1.1483) #162 (raahilshah), Record: 10L Int5-MLP + BigramHash(10240) + SWA(0.4) + WD=0.04 (val_bpb=1.1428, mean 3 seeds) #180 (thwu1)
• TTT as eval-time technique: PR [record bpb=1.195] sliding window + LoRA TTT #77 (samacqua)
• Aggressive TTT: PR Non-record: 11L EMA + TTT(20ep,freeze=0) + 15-run ablation study — val_bpb=1.1213 (3-seed) #398 (felipe-parodi)
• AdamW for TTT: PR Record: 11L EMA + AdamW TTT 10ep (mean val_bpb=1.1027) #442 (sjp611)
• DDP TTT sharding: PRs Non-record: 11L EMA + TTT(20ep,freeze=0) + 15-run ablation study — val_bpb=1.1213 (3-seed) #398, Record: 11L XSA4 + Tight SWA + FA3 + Two-Phase TTT (3-seed mean val_bpb=1.1227) #417 (EthanYangTW)
• Partial RoPE, LN Scale: PR Record: 11L Partial RoPE + LN Scale + EMA + XSA4 (val_bpb: 1.1248) #315 (jfprincz)
• XSA: arXiv:2603.09078
• Muon, modded-nanogpt: foundation codebase

Reproduction

git clone https://github.com/mrdavtan/parameter-golf.git
cd parameter-golf && git checkout next-gen
pip install flash-attn --no-cache-dir --no-build-isolation
pip install zstandard sentencepiece huggingface_hub
python3 data/cached_challenge_fineweb.py --variant sp1024
bash run_competition.sh 1337

Hardware: 8xH100 SXM (RunPod), PyTorch 2.9.1+cu128, Flash Attention 2

[mrdavtan]

Updating to sliding window eval. New scores in progress. The per-layer lr groups, cosine scheduling, and ablation findings are independent of the evaluation method.

[mrdavtan]

Closing: multi-epoch TTT is invalid per the clarified rules in #402. The per-layer LR and cosine scheduling contributions remain available for legal single-pass (Case 2) TTT implementations.