One QK Channel, Many Sources: Tracing Low-Precision Attention Collapse
Not on the map. Why?
Abstract (arXiv)
A bfloat16 transformer can train normally, then collapse abruptly. Prior work links collapse to structured attention errors and shows QK normalization disrupts their compounding. Distinct low-precision errors trigger the same collapse, leaving unclear whether each needs a fix at its source or one shared route can be blocked instead. We isolated the fault behind a reproduced GPT-2-class collapse to the streaming-softmax accumulator, where an fp32 streaming core repairs it, and turned it into an assay for moving a controlled error across sources. Using it, we found that errors placed outside attention still drove the same QK spectral runaway, and that correcting only QK kept training stable while the fault stayed active. This is a source-channel dissociation: fault source is not failure channel. It held across tested architectures and scales, and reproduced on a second GPU architecture. As a causal probe, projecting each update off the current QK weights' leading three singular directions held the query projection's largest singular value to 11.1, whereas removing equal energy elsewhere left it at 237: the QK channel causally drives the early runaway. What lets the injected error in is temporal sign-coherence, its per-head sign persisting across steps, not aggregate deviation; once inside, the runaway shows as attention-logit saturation. QK-Guard, a dormant controller, tests this by switching on parameter-free QK normalization at the first monitored threshold crossing. On the runs designated for this test, the QK-local action prevented the failure of each matched or same-configuration unguarded run; on plain GPT-2, all 12 final train and validation losses were within 0.03 nat of same-configuration always-on QK-norm, and both methods ran 60k steps without collapse. Intervention at the QK locus therefore suffices in place of a fix at each source.
Key facts. No citing papers in the map yet (data through 14 Oct 2026). Builds on 27 of its 40 references that are in the map. First posted on arXiv 3 Aug 2026.
Builds on
27 of 40 references.
References · 27
- 2021RoFormer: Enhanced Transformer with Rotary Position Embeddingcited by 4,703
- 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,673
- 2023FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioningcited by 2,598
- 2019Root Mean Square Layer Normalizationcited by 1,480
- 2017Mixed Precision Trainingcited by 1,156
- 2024Chameleon: Mixed-Modal Early-Fusion Foundation Modelscited by 927
- 2015Deep Learning with Limited Numerical Precisioncited by 376
- 2020Query-Key Normalization for Transformerscited by 280
- 2019A Study of BFLOAT16 for Deep Learning Trainingcited by 232
- 2023Small-scale proxies for large-scale Transformer training instabilitiescited by 183
- 2018Online normalizer calculation for softmaxcited by 162
- 2023Stabilizing Transformer Training by Preventing Attention Entropy Collapsecited by 161
15 more
- 2024Scaling FP8 training to trillion-token LLMscited by 65
- 2023Oscillation-free Quantization for Low-bit Vision Transformerscited by 57
- 2023A Theory on Adam Instability in Large-Scale Machine Learningcited by 56
- 2025SPAM: Spike-Aware Adam with Momentum Reset for Stable LLM Trainingcited by 28
- 2024Methods of improving LLM training stabilitycited by 13
- 2025Why Low-Precision Transformer Training Fails: An Analysis on Flash Attentioncited by 12
- 2019Accumulation Bit-Width Scaling For Ultra-Low Precision Training Of Deep Networkscited by 11
- 2025Taming Transformer Without Using Learning Rate Warmupcited by 9
- 2025AdaGC: Enhancing LLM Pretraining Stability via Adaptive Gradient Clippingcited by 6
- 2026Pretraining large language models with MXFP4 on Native FP4 Hardwarecited by 5
- 2025Characterization and Mitigation of Training Instabilities in Microscaling Formatscited by 4
- 2025Controlling changes to attention logitscited by 3
- 2025Spectral Alignment as Predictor of Loss Explosion in Neural Network Trainingcited by 2
- 2026Rank-Aware Spectral Bounds on Attention Logits for Stable Low-Precision Trainingcited by 2
- 2026Mechanism-Driven Monitors for Preemptive Detection of LLM Training Instabilitycited by 2
Built on by
Not cited yet.