KURE/R2 + Tanh Reparam + Parallel EMA + LoRA TTT

[machdragon]

Summary

Built on PR openai#201 (LAWA-EMA + Int6 + Overtone + MLP3x, val_bpb=1.1551). Adds four improvements targeting quantization fidelity and eval-time adaptation.

Changes (1857 lines, +329 from base)

1. KURE + R2 regularization (lines 1081-1098, 1635)

• quant_reg_loss() — kurtosis→1.8 penalty + 2×std outlier penalty
• Added to training loop using unwrapped base_model (DDP-safe)
• eps=1e-8 in kurtosis to prevent NaN when var≈0

2. Tanh reparameterization (lines 546-553)

• CastedLinear._tanh_reparam class flag; forward applies torch.tanh(w) for 2D weights ≥ 64×64
• .weight kept as nn.Parameter (safe for all ~15 callers: _init_weights, optimizer setup, TTT adapter sizing, tied embeddings, etc.)
• LAWA shadows track tanh(P) directly; materialized before export
• QAT STE still applies after tanh (correct ordering)

3. Parallel EMA tracks (lines 1569-1582, 1660-1666, 1703-1728)

• Three decay rates: 0.995, 0.999, 0.9995 (default raised from 0.995)
• Clone-then-pick via proxy eval at export time
• Safe: doesn't mutate shadows during eval passes

4. Causal LoRA TTT (lines 1107-1287, 1838-1849)

• Ported from PR [record bpb=1.195] sliding window + LoRA TTT openai/parameter-golf#77: BatchedLinearLoRA, BatchedTTTLoRA, eval_val_ttt_lora()
• Integrated into GPT.forward() (lora= kwarg), Block.forward() (q/v delta fns), CausalSelfAttention.forward() (q/v deltas)
• Per-token loss (reduction="none") when lora is set; mean loss otherwise
• Per-doc LoRA reset, score-before-train, causal attention — leakage-free
• Runs after int6 roundtrip quantization

Hyperparameters (env vars)

Variable	Default	Description
KURE_LAMBDA	0.01	Kurtosis regularization strength
R2_LAMBDA	0.01	Range regularization strength
TANH_REPARAM	1	Enable tanh weight reparameterization
LAWA_EMA_DECAY	0.999	Default EMA decay
TTT_LORA_ENABLED	1	Enable LoRA TTT at eval
TTT_LORA_RANK	8	LoRA rank
TTT_LORA_LR	0.01	LoRA learning rate
TTT_CHUNK_SIZE	256	Chunk size for TTT
TTT_EVAL_SEQ_LEN	1024	Context window for TTT
TTT_BATCH_SIZE	64	Batch size for TTT eval

Verification

• python3 -c "import ast; ast.parse(...)" — syntax OK
• quant_reg_loss, KURE_LAMBDA, R2_LAMBDA — reg wired in
• tanh_reparam, torch.tanh(w) — tanh in CastedLinear.forward
• lawa_decays, lawa_averaged — parallel tracks + safe export
• BatchedTTTLoRA, eval_val_ttt_lora — TTT ported
• 1e-8 in kurtosis — eps present
• base_model in reg loss — DDP unwrap correct
• .weight still exists on CastedLinear (inherits nn.Linear)
• Run eval → update submission.json

Test plan

• Run training on 8xH100 with default hyperparameters
• Verify KURE/R2 reg loss appears in training logs
• Verify parallel EMA proxy eval picks best decay
• Verify TTT LoRA eval completes and reports val_bpb
• Compare val_bpb against PR LAWA-EMA frontier fork (pr198 base, SWA -> LAWA val_bpb=1.1551) openai/parameter-golf#201 baseline (1.1551) and PR 11-Layer Int6 + WD=0.04 + SWA + FA3 (val_bpb: 1.1318) openai/parameter-golf#198 (1.1318)
• Verify int6 artifact size stays under 16MB

🤖 Generated with Claude Code

---

[devin-ai-integration]

✅ Devin Review: No Issues Found

Devin Review analyzed this PR and found no bugs or issues to report.