Record: 0.6864 BPB — K-LoRA + Min-NLL + FlashAttention-3

[bigbag]

Summary

• val_bpb: 0.6864 (seed 42, 8xH100 SXM)
• Artifact: 15.53 MB (97.1% of 16MB limit)
• Training: 600s (7313 steps at 82ms/step)
• Eval: 588s (within 600s budget)

Three Innovations

1. K-Projection LoRA (from PR #611)

LoRA on K projections (not just Q/V) with 0.3x LR multiplier. Adapting K alongside Q/V gives more expressive per-document specialization at marginal compute cost.

2. Min-NLL Epoch Selection (from PR #611)

Track minimum average NLL per document across all TTT epochs, use best epoch's scores. Prevents late-epoch overfitting — safely run 6 epochs without any document degrading.

3. FlashAttention-3 (our addition)

flash_attn_func for causal attention + Rotary cache .clone() fix for CUDA graph compatibility. ~3% speed boost.

LoRA TTT Details

• Rank-8 Q/K/V LoRA + rank-16 LM-head LoRA
• Per-block bias tuning, per-document reset at BOS boundaries
• Batched 64 docs/GPU, Adam lr=0.01, 6 epochs, per-step cosine LR
• Per-layer LR: LM-head 2x, V 1.5x, Q 0.5x, K 0.3x, bias 3x
• Temperature rescaling T=0.98, wall-clock deadline 550s

Results

• Pre-quant BPB: 1.1624
• Post-quant BPB: 1.1755 (quant gap: 0.013)
• Post-TTT BPB: 0.6864 (TTT gain: 0.489)
• Eval breakdown: short docs 27s + long docs 340s + scoring 221s = 588s

Based On

• PR Record: Chimera TTT — K-Projection LoRA + Min-NLL (0.5601 BPB, 3-seed mean) #611 (Chimera TTT, K-LoRA + Min-NLL)
• PR Record: DeepQuant V10b — 11L INT6 + 8ep LoRA TTT (val_bpb=0.6430) #596 (DeepQuant V10b, LoRA TTT architecture)

Test plan

• Artifact under 16MB (15.53MB)
• Training under 600s (600s)
• Eval under 600s (588s)

[valerio-oai]

The min-NLL scheme leaks information. This is the same as training on the val set, and is therefore disallowed.