Low-Precision Transformer Failure in Flash Attention

This October 5 2025 paper presents the first mechanistic explanation for a persistent **training instability** experienced when using **low-precision arithmetic** (specifically BF16) with the **Flash Attention** algorithm in transformer models. The paper identifies the core problem as a "catastrophic loss explosion" caused by two interacting phenomena: the emergence of **similar low-rank representations** within the attention mechanism and the accumulation of **biased rounding errors** inherent to BF16 addition during the attention output calculation. This bias leads to a systematic error in the gradient updates, causing the spectral norm of weights to increase and derailing the training process. To validate this analysis, the authors introduce a minimal modification to the softmax computation in Flash Attention that **mitigates the rounding bias** and successfully stabilizes the training, offering a practical solution to this long-standing issue.Source:https://arxiv.org/pdf/2510.04212

There are no comments yet.

Please log in to write the first comment.