低精度注意力崩溃的根源在共享的QK通道,而非具体错误源。
One QK Channel, Many Sources: Guarding Low-Precision Attention Collapse
- 发现崩溃源于QK通道的谱漂移,修复该通道即可稳定训练。
- 仅修复QK通道可阻止崩溃,即使其他错误源仍存在。
- 提出无需针对每个错误源单独修复,适合高并发低精度部署。
bfloat16下的Transformer可在数千步内正常训练,随后突然崩溃。不同低精度误差会引发相同失败,难以判断是否需各自修复或统一阻断。我们将一个复现的GPT-2级崩溃定位至streaming-softmax累加器,使用fp32累加可修复。通过可控扰动实验,发现即使错误位于注意力之外,仍会触发相同的查询-键(QK)谱漂移。仅修正QK通道即可维持训练稳定,而其他修复无效。该现象在多个架构与规模下均成立,并在另一GPU架构上复现。因果探测显示,每次更新都脱离当前QK权重前三个主奇异方向:查询投影的最大奇异值保持11.1,而移除等量能量后其升至237,表明QK通道是早期漂移驱动者而非被动追踪者。其触发依赖跨步的时间符号一致性,非累积偏差。QK-Guard通过一个惰性控制器,在注意力对数饱和时启用无参数的QK归一化,可有效遏制所有测试中的漂移,且在60,000步内表现优于持续启用的归一化,而同一触发点的非QK操作则失效。结果支持在共享的QK位置干预,而非对每个故障源分别修复。
原文摘要 · Abstract (English)
A bfloat16 transformer can train normally for many steps and then collapse abruptly. Distinct low-precision errors can trigger the same failure, leaving unclear whether each source needs its own repair or one shared route can be blocked. We isolate a reproduced GPT-2-class collapse to the streaming-softmax accumulator, where fp32 accumulation repairs it, and use the fault as an assay for moving controlled errors across sources. Errors placed outside attention still drive the same query-key (QK) spectral runaway, while correcting only QK keeps training stable with the source fault active. This source-channel dissociation shows that fault source is not failure channel. It holds across the tested architectures and scales and reproduces on a second GPU architecture. A causal probe projects each update off the current QK weights' leading three singular directions: the query projection's largest singular value stays at 11.1, whereas removing equal energy elsewhere leaves it at 237. The QK channel therefore drives the early runaway rather than merely tracking it. Entry depends on temporal sign-coherence across steps, not aggregate deviation. QK-Guard closes the channel with a dormant controller that switches on parameter-free QK normalization when attention-logit saturation begins. It contains every tested runaway and matches always-on QK normalization over 60k steps, while non-QK actions at the same trigger fail. The results support intervention at the shared QK locus rather than separate repair at each fault source.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。