低精度训练时闪注意力会导致损失爆炸，原因被首次揭穿。
Why Low-Precision Transformer Training Fails: An Analysis on Flash Attention
- 发现低精度下注意力层出现相似低秩表征，引发误差累积。
- 量化误差的偏置效应与注意力结构相互作用，导致训练崩溃。
- 仅修改闪注意力即可稳定训练，适合做高效模型训练的研究者。
追求计算效率促使人们在训练Transformer模型时采用低精度格式，但常因严重的训练不稳定性而受阻。本文首次揭示了一个长期未解难题的机制：在低精度设置下使用闪注意力（Flash Attention）会引发灾难性损失爆炸。深入分析表明，该失败并非偶然，而是由两个交织现象导致：注意力机制中出现相似的低秩表示，以及低精度算术固有的有偏舍入误差的累积效应。这两者共同形成误差积累的恶性循环，破坏权重更新，最终瓦解训练动态。为验证结论，我们对闪注意力进行了最小修改，有效缓解了舍入误差偏差。该简单改动成功稳定了训练过程，证实了分析的有效性，并提供了一种切实可行的解决方案。代码已开源：https://github.com/ucker/why-low-precision-training-fails。
原文摘要 · Abstract (English)
The pursuit of computational efficiency has driven the adoption of low-precision formats for training transformer models. However, this progress is often hindered by notorious training instabilities. This paper provides the first mechanistic explanation for a long-standing and unresolved failure case where training with flash attention in low-precision settings leads to catastrophic loss explosion. Our in-depth analysis reveals that the failure is not a random artifact but caused by two intertwined phenomena: the emergence of similar low-rank representations within the attention mechanism and the compounding effect of biased rounding errors inherent in low-precision arithmetic. We demonstrate how these factors create a vicious cycle of error accumulation that corrupts weight updates, ultimately derailing the training dynamics. To validate our findings, we introduce a minimal modification to the flash attention that mitigates the bias in rounding errors. This simple change stabilizes the training process, confirming our analysis and offering a practical solution to this persistent problem. Code is available at https://github.com/ucker/why-low-precision-training-fails.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。