Why Low-Precision Transformer Training Fails: An Analysis on Flash Attention
Authors: Haiquan Qiu
Quanming Yao
Venue: ICLR 2026 Oral (2026)
Keywords: low-precision training
transformer
attention
Primary area: optimization
TL;DR: For the first time, we mechanistically explain why low-precision training with flash attention fails, identifying a vicious cycle of rounding errors and proposing a simple, effective fix.

Abstract: 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.

Published as a conference paper at ICLR 2026
WHY LOW-PRECISION TRANSFORMER TRAINING FAILS: AN ANALYSIS ON FLASH ATTENTION
Haiquan Qiu²
Quanming Y.0
Department of Flectronic Engineering, Tsinghua University
2 Beijing National Research Center for Information Science and Technology State Key laboratory of Space Network and Communications qbq22@mails tss tsinghua. edu. en, qyacaa@tsinghus edu. cn
ABSTRACT
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 expla- nation for long-standing and unresolved failure case where training with flash attention in low- ow-pprcissonn settings leads to catastrophic loss explosion. Our in- depth analysis reveals that the failure not 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 de- railing the training dynamics. To validate our findings, introduce minimal modification to the flash attention that mitigates the bias in rounding errors. This simple change stabilizes the training process, confirming ing a practical solution to this persistent problem. Code is available at https: tigithub.ccmm/ www.//oo//cccccceeooo /why low preci sion- traiiirinnnnnn saaininn-faals fails
analysis and offer-
INTRODUCTION
The pursuit of training ever-larger and more powerful transformer models is a relentless drive for computational efficiency (Brown al. 2020 Hoffmann et al., (022). A key strategy in this en- deavor is the adoption of low-precision numerical formats (Micikevicius et al., 2017]] Wang et al.) 2018 Kalamkar et al. 2019 Liu et al. 2024). which promise substantial reductions memory foot- print and significant boosts in training speed. In industrial practice, it is common to use BF16 for memory-bound operations like flash attention while pushing compute- bound operations like FFNs to even lower precisions such as FP8 (Liu et al.) [2024] (Qwen-Tean (025). This highlights the height- ened sensitivity of attention mechanisms to numerical precision. Despite development of stabi- lization techniques like QK normalization (Henry et al. 20200 Qwen-Team 2025). QK-clip (Kimi- Team 2025) and Gated Attention (Qiu et 2005) Qwen- -eaen-Teammaa 2025). the path to further reducing precision is often blocked by a lack of understanding of the underlying failure mechanisms.
This paper confronts this challenge by dissecting a notorious and long-standing failure issue involv- ing flash attention. By reducing the memory complexity of the attention mechanism from quadratic to linear with respect to sequence length, flash attention has become a comeestone algorithm for efficient transformer training, making it indispensable for handling the long contexts required by
modern large-scale models (Dao et al.] [0022] Dao
2024
Shah et al., 2024). This failure, arising
in low-precision settings (flash-attention Issue 337
2024
nanoGPT Issue 303 2023 nanoGPT Is-
sue 524 2024 na Issue 554 2024: Lee et al. 2024 Golden et al. 22024). presents significant
bottleneck. We focus on specific. reproducible failure case reported by the community (nanoGPT)
Issue 303 2023 nanoGPT Issue 524 2024), which has remained unresolved for over two years. Our in-depth analysis provides the first mechanistic explanation for this failure, revealing that it is not random artifact but direct consequence of two intertwined phenomena: the emergence of simi- lar low-rank representations across different training steps and tokens, and the compounding effect
Corresponding author

---

Published as a conference paper at ICLR 2026
of biased rounding errors inherent in low- Now-precision arithmetic. We demonstrate how these biased rounding errors act as coefficients for the low-rank representations, causing them to accumulate as biased gradient update to the weights. This pushes the spectral norm of weights and activations to increase abnormally, ultimately overwhelming the training dynamics. To validate our analysis, introduce minimal modification flash attention that mitigates the bias in rounding errors, allowing the low-rank weight updates to cancel out during training. This experiment confirms our analysis and stabilizes the training process, offering practical solution this persistent problem.
Notations We use bold lowercase letters for vectors (e.g., ô, m, I') and bold uppercase letters for matrices (e.g.. Q, K. V). We use notation dW to denote the gradient of the loss & with respect to variable W for simplicity, c.g., dQ == /QQ. diag(v) denotes a diagonal matrix with the elements of vector v on its diagonal. o denotes the element- wise product. We use Python-style indexing, e.g., :[ denotes the i-th row of matrix M. Subscripts ip and hp distinguish between low-precision (BF16) and high-precision (FP32) computations. Binary strings are represented using a typewriter font (e.g.. 101010). The where function mimics torch. where. Unless specified otherwise, operations follow Py'Torch's broadcasting nules. Key claims highlighted in light cyan box at the end of some sections.
2 PRELIMINARY
2.1 LOW-PRECISION TRAINING
Low-precision training is a cornerstone of modern deep learning, enabling the development of in- creasingly large models by reducing memory usage and accelerating computation (Liu et al. 2024] Tseng et al. 2023). This is achieved by representing weights, activations, and gradients using nu- merical formats with fewer bits than the standard 32-bit single-precision (FP32). Mixed-precision training (Micikevicius et al., 2017) has been widely adopted in practice, which combines 16-bit formats like FP16 or 6 (BF16) for most computations with an FP32 master copy of weights to maintain accuracy. While FP16 offers higher precision, its limited dynamic range often leads to gradient underflow, requiring techniques like loss scaling. contrast, BF16, originally developed for Google's s TPUs and now widely supported, provides the same dynamic range as FP32, making it more robust against underflow and a preferred choice for training large language models (Kalamkar et al. 2019 Wang & Kanwar 2019). However, the reduced precision of BF16 can still introduce numerical errors that lead to training failure, the focus of this paper.
The 6 (BF16) format is 16-bit foating-poin representation with sign, 8 exponent, and 7 significand bits. It matches the dynamic range of 32-bit cision, making it popular choice for balancing computation and numerical range in deep learning. Adding two BF16 numbers involves aligning their exponents, adding the significands, normalizing the sum, and rounding the result to fit the 7-bit fraction. This final rounding step. typically "round to nearest, ties to even", is one of the primary sources of error. While this rounding method is designed to be unbiased for random data, sequence of operations on data with specific distribution could lead to biased rounding error. This accumulation of error in one direction is a critical factor to the training fallure observed in low-precision settings. See Appendix B for details of BF16 addition.
(FP32) but has lower pre-
2.2 FLASH ATTENTION
Flash Attention (FA) DDao et al. 2022 Dao 2024 Shah et al. 2024) is an 1/O-aware exact attention algorithm designed to overcome the memory bottleneck of standard attention. Standard attention, defined as O = soft
V. requires materializing the N x N attention score matrix = a QK leading to a memory complexity of ((22) with respect to the sequence length N. Flash attention reduces this to ((N) by partitioning the input matrices Q K. V e R into blocks and processing them iteratively. These blocks are loaded from high-bandwidth memory (HBM) into fast on-chip SRAM, minimizing costly memory transfers.
In this paper, we focus on analyzing flash attention 2. The forward pass of FA computes the output O and log-sum-exp statistics L using online softmax method (Algorithm D. It iterates through blocks of Q (outer loop) and blocks of K, V (inner loop). For each query block Q. it maintains running statistics: the maximum score m; and the normalization factor f,. In each inner loop step,
2

---

Published as a conference paper at ICLR 2026
Section 3.2
3.3.1
Section 3.3.2
Isolate Failure Source: Low precision Up in ip rowsum (dO 0gy)
Low-rank (P = R
8ip 8p is blased towards
is similar across training steps and tokens (Claim 2)
positive due to safe softmax and PV (Claim 3)
causes training failure (Claim 1)
Soction
D High and low precision gradient error (biased weight update) dw hp
Find the Root Cause
dW (p
(8i
8up)(T] [T] R
11
Section 4
Low-precision gradient in training increases spectral norm of
weights and activations
Mitigote Blas in Rounding Error: use dynamic maximum in safe softmax
Loss Exploding!!! *
Figure 1: Analysis in different sections. Our paper traces the causal chain of training failure (blue
box) in reverse to identify the root causes.
it computes unnormalized attention scores P( = exp(ss)) - m and updates an unnormalized
())
output by accumulating the product p V This accumulation is a key focus of our analysis. After iterating through all key/value blocks, the final output block O; is correctly normalized. This tiling strategy avoids materializing the full N x N score matrix. The backward pass (Algorithm 2 leverages the same tiling strategy. It first computes a key intermediate term, 6 rowsum (O which is central to our investigation. Then, recomputes attention scores on- on-the-ft the-fly to calculate the gradient of the scores, ds(") p(1) (dp! 5.), where ddp(") dO. The final gradients dQ, dK, dV are accumulated block-wise. This approach maintains I/O efficiency both passes.
RoOT CAUSES OF INSTABILITY FLASH ATTENTION
We first introduce the failure case in Section 3.1 In Section 3.2 we narrow down the source of the
numerical errors within the flash attention. Further analysis in Section 3.3 reveals that the failure stems from a combination of two factors: the emergence of low-rank representations and the ac- cumulation of biased rounding errors inherent to BF16 arithmetic. The full process of such failure from root cause to loss exploding is shown Fig. [
3.1 THE FAILURE CASE OF LOW-PRECISION FLASH ATTENTION
Our investigation targets a well-documented and per- sistent failure: the catastrophic loss explosion that oc- curs when training Generative Pre-trained Transformer 2 (GPT-2) models with flash attention in BF16 preci- sion G Issue 303 2023) nanoGPT Issuc 524 2024 nanoGPT Issue 554 2024). This failure case, re- ported for over two years, manifests as a sudden loss ex- plosion after several thousand training steps (see Fig. (0)). While empirical workarounds like reverting to standard attention or using higher precision (FP32) stabilize train- ing, they come at the cost of efficiency. This instability is not an isolated incident; the broader community has ob- served similar failures when training large language mod- els (Kimi- Team 2025] Twen-Teemm (025). These fail- ures are often empirically linked to phenomena such large spectral norms weights, large activations (Yang)
Triining Lowppecciion
Tratning Stps
Figure 2: The failure case using BF16 and flash attention results in a sudden loss explosion. while the stable config- uration converges.
et al. 2023; Rybakov et al., 2024), and attention sinks (Xiao et al. 2023), leading to suite of fixes including QK normalization (Henry et al. 2020; Qwen Team (025). OK-clipping (Kimi- Team 2025), and Gated Attention (Qiu et al.] 2022 Qwen-Teent
2025). Despite these interventions,
fundamental understanding of the root causes has remained clusive. The absence clear causal
chain from numerical error to loss explosion has left the community reliant ad-hoc patches rather
3

---

Published as a conference paper at ICLR 2026
than principled solutions, hindering progress in robust low-precision training. This paper provides the first mechanistic explanation by reproducing the failure, dissecting its root causes, and proposing practical, principled solution.
To reproduce the failure, we employ a GPT-2 architecture with 12 layers, 12 attention heads, an em- bedding dimension of 768, and context length of 1024. The model is pre-trained the Open Web- Text dataset (Gokaslan et al. 2019). For deterministic reprodueibility, we deviate from a standard random data loader by recording and reusing the exact sequence of data batches from an initial run that led to the failure. This ensures all subsequent experiments process identical data in the same order, isolating the failure from data-related randomness.
We train the model using the AdamW optimizer with 31 = 0.9, 82 = 0.95, and zero weight decay. The learning rate follows cosine schedule with 2000-iteration linear warmup a peak of 1 x 10-3 decaying to 1 x 10-5, We apply global gradient clipping with a maximum norm of 1.0. Training is conducted on 4 NVIDIA A A100 (80GB) GPUs using PyTorch's Distributed Data Parallel (DDP) module. We use automatic mixed precision with BF16 for the forward pass and FP32 for the backward pass. Each GPU processes a micro-batel size 32, and with gradient accumulation over 4 steps, the effective global batch size is 524,288 tokens per optimization step.
3.2 ISOLATING THE SOURCE OF FAILURE WITHIN FLASH ATTENTION
To pinpoint the source of failure within flash attention, we conduct a series of targeted experiments. We systematically modify the digorihmm-disabling tiling, selectively replacing flash attention with a standard implementation, and performing key computations in high precision-to narrow down the potential causes of instability. To accelerate this analysis, we monitor empirical indicators such as the spectral norm of weights to quickly identify failing configurations.
To determine if the block-wise processing in flash attention was responsible for the failure, we conduct an experiment where we disabled tiling by setting the block size equal to the sequence length. This forces the algorithm to compute with the full matrices at once. The training process still failed, leading to the same loss explosion. This finding rules out the tiling strategy as the cause of the problem. For all subsequent experiments, we therefore use this non-tiled setup to simplify the analysis and focus on the core numerical computations.
Tiling is not the Source of Failure.
Failure Originates in a Single Layer. layers (Yang al., 2023 Rybakov et al. 2024). This reveals an anomalous spike specifically within the second layer's attention (see Fig. III We confirm this finding with two targeted experiments: (1) using flash attention only in layer 2 was sufficient reproduce the training failure, and (2) replac- ing flash attention with standard attention in layer 2, while retaining it in all other layers, restored training stability. These results conclusively identify the flash attention layer as the origin of the failure. So. subsequent analysis focuses on this module to dissect the failure mechanism.
We first analyze the spectral norms of weights across all
Spectral Norm wf per Head in Layar
Failure is Linked to the Computation of 8. The backward pass of flash attention uses the term ô rowsum(dO O) E V for computational efficiency. An alternative, mathematically equivalent formulation com- putes this term as = rowsum(dP o P), where dP dOVT. We find that replacing the efficient computa- tion with this alternative formulation restored training sta- bility. This experiment demonstrates that numerical er- rors introduced when computing O in BF16 is likely the primary source of failure, as the alternative formulation which avoids training failure is equivalent to use O com- puted in FP32.
Spstre hore
Figure 3: Wº of attention head 8 has the largest spectral norm. Subsequent analysis focuses on this head.
Numerical Errors in O are the Source of Failure. Building on the finding that the computation of 8 is critical, we further isolate the source of error to the output matrix Olp in low-precision 6,p = rowsum/ Oip) We conduct two key experiments. First, instead of using the low-precision O from the forward pass to compute 5, we recompute it as PV in FP32 within the backward pass; this change stabilizes training. Second, we find that com- puting O in high precision (FP32) during the forward pass, i.e., diep rowsum(dO Ohp), while
4

---

Published as a conference paper at ICLR 2026
keeping all other operations in BF16, also restored stability. This evidence conclusively demon- strates that numerical errors introduced during the BF16 computation of O are the direct cause of the failure (refer to Claim
Failure is Localized to Specific Attention Heads. To further narrow down the source of failure, we analyze individual attention heads by tracking the spectral norm (Yang et al. 2023 Rybakov et al. 2024) of their query projection matrices (WQ), as shown in Fig. This reveals that a few heads exhibit disproportionately large spectral norms. We confirm their role by selectively comput- ing the output in high precision for these outlier heads (1, 7, 8, 9, 11, and 12), which sufficient to restore training stability. Since head 8 shows the largest spectral norm, focus our subsequent analysis on this head to dissect the precise failure mechanism.
Claim 1. Low-precision dip = rowsum (dO o Oip) causes training failure.
3.3 THE RooT CAUSES OF TRAINING FAILURE
Our investigation in this section uncovers the two interconnected root causes of the training failure. In Section 3.3.T we demonstrate how low-precision dip drives the training failure by a biased weight update, and find that bias arises from the emergence of similar low-rank representations R, whose coefficients (òip - [kp) [Z'] are biased towards positive values, causing the error to accumulate rather than cancel out. In Section 3.3.2. trace the origin of these positive coefficients (ōip (Tp to biased rounding errors inherent in the 6F16 addition within the PV product.
3.3.1 CAUSE 1. SIMILAR LOW-RANK MATRICES BIAS WEIGHT UPDATES
This section traces the training failure to biased weight update. We first analyze the difference between the high and low precision gradients calculated with Öhp and dip, respectively. Then we find that the low-rank representations and biased (ōtp 6kp) (T) lead to loss explosion.
Analysis of the Error between High and Low Precision Gradients To understand how nu- merical errors propagate into the gradients, we analyze the difference between the high-precision (hp) and low-precision (lp) gradients for the query matrix, dQ. The gradient dQ is computed from the gradient of the attention scores, dS, dQ dSK. The score gradient is given by dS = aP (dP - 5). where 5 = powsum(dO) o O) is the only term that differs between the high- and low oow-precision backward passes based Section 3.2 and is scaling factor in attention.
The difference between the high-precision and low precision query gradients can be derived as:
dQhp dQip (dShp KK
= =(aP o (dP 6np) aP (dP d1p)) K = (a.P (ōtp ōnp)) )
(I)
=a diag Ótp- 8 (PK).
In the final step, we express the row-wise scaling operation P (ötp 8np) (follow broadcasting rule) as matrix multiplication, where diag( ōkp) a diagonal matrix whose diagonal entries are the elements of the vector difference dip ōnp This formulation reveals that the gradient error directly proportional to the error in å and modulated by the term PK.
The gradient of the query projection matrix, dW is given by the outer product of the input features X and the query gradient dQ. The difference between the high-precision (hp) and low-precision (Ip) gradients for W can be expressed as:
dw hp
dw 1p (dQup dQip)) T =x==xxax (PK)) diag Sip
8kp)
T=1 (
where [T and X[T] are the T-th row vectors of their matrices. This equation shows that the total gradient error is a weighted sum of rank-1 matrices, with weights given by the error in d.
Similar Low-rank Updates of Weight Cause Training Failure In Fig. the rows of PK (panels d) and X (panels b, exhibit strong structural similarity across different training steps and token
5

---

Published as a conference paper at ICLR 2026
PK thoze Sutch trdes Trunnig 95201
(PRISOD 5CC (Head D, Bach India 191, Tha ring 6610)
(b)
(c)
intea 20s, Tralnnn 6614)
X (Batch Irdie 209, Tra Traiiing 6619)
(PRITA00 "E77EE tread Butch Index 200, Traiiigg 0013)
Featuree
(d)
(I)
Figure 4: PK, X, and (PK) X[T] at different batch indices and training steps. (c) and (f) show that (РК) X[T] for different tokens and training steps have some similar columns in input features 546 and 678.
positions. This implies that the resulting rank-1 matrices, (PK) [ to one another. For instance, Fig. (panels c, f) shows this similarity for tokens 50 and 718 at training steps 6610 and 6619, respectively. Because these rank-1 error components are structurally consistent, we can approximate the total gradient difference as
also highly similar
ip d dwo R hp
(6
[T
T=1
where R. denotes the common low-rank structure emerging across different tokens and training steps.
Eqn. (3) shows that the accumulation of the low-rank error direction governed by the scalar term Er=( (ō (6.p-6%p))(7] If this sum is biased towards non-zero values, the error across different training of -
steps will accumulate rather than cancel out. We track the cumulative over sequence of training steps (6580 to 6680) leading up to the failure, as shown in Fig. 5(a). The plot reveals that this sum consistently positive, indicating systematic bias. This bias causes the еитог in the low-rank direction R. to compound with each training step. Because R is also similar for different steps. this ultimately corrupts the weight updates, increases the spectral norm (Fig. I and activations (Yang et al.) 2023 Rybakov et al. (2244), and causes the training failure (refer to Claim 2). The following section finds the root cause of this positive bias by analyzing the weights and gradients training step 6619, point of significant positive contribution identified in Fig 5 a).
8 [T]
Weight update in low-precision training is biased by (ōtp [T]R, which arises from the structurally similar matrices (denoted as R) across tokens and training steps, and its positively-biee biased coefficient (ōip 6 (T). This bias accumulates error, preventing cancellation, increasing weight spectral norm and activation, and leading to loss explosion.
Claim 2.
3.3.2 CAUSE 2. BIASED ROUNDING ERROR LEADS TO POSITIVE (ō1p S
This section investigates the origin of the positive bias in (8,p [TT. We trace this error to the interaction between dO and the numerical discrepancy in which itself arises from biased rounding errors in BF16 addition during the PV computation.

---

Published as a conference paper at ICLR 2026
Campaltee - -2,
www.
Muw
M
(a) Positively-biased (ōip [T] (b) Large values in features 20 & (c) Large error in features 20 & 29 29
Figure Analysis of 6 = rowsum(dO O).
Ve Ved ml
Vinuitaation OfT. Ousdll
Detal Ven of Boodlt
(a) Most [ia are negative
(b)
Large error when PT,
Detailed view of (b)
Figure 6: Analysis of PV
Locate the Large Error in (ōip - [Tp)[T] We first investigate the source of the positive bias in ET= [T). As analyzed in Section 3.2 the error in ô stems from the product of the upstream gradient dO and the numerical error in the low-precision output, Oip -Opp To dissect
this, focus token position T 718 where the error component (ōtp h [T] is positive.
In Fig. b) and (c), we observe strong sign correlation between the gradient dO[T;:] and the output error O [] - Onp :] for specific feature dimensions such as 20 and 29 (also observed across other tokens). In these dimensions, both dO and the output error O lp Onpp are consistently negative. This alignment ensures their product, which contributes the error in 5, is positive. The
fact that the output error tends to be negative (Ti :] Ohp T, il) indicates that the low-precision computation of systematically biased towards more negative values. Our subsequent analysis, therefore, focuses on identifying the origin of this computational bias.
The output computed from an intermediate unnormalized output, Ô. The computation involves safe softmax followed by matrix multiplication and normalization:
P exp(S - rowmax (S)),
Ô = PV,
O - Ô rowsum(P)
Further experiments pinpoint the source of failure to the computation of the unnormalized output, O = PV. We find that computing only this product in FP32 is sufficient to stabilize training. To understand the origin this bias, we examine the difference between the low-precision and high- precision computation of single element, O i] (for feature index i 20 in our analysis):
O T, i] O [T, (P [T,:] V(; (P hp [T, V
(4)
where the inputs P and V are themselves the results of prior BF16 operations. Specifically, the subscript (-)ip here computes dot product FP32 with the final result rounded BF16, while (-) computes entirely in FP32.
To understand how the error (p [ O [T,) becomes systematically negative, we plot the cumulative emorr as the sum over token positions progresses in Fig. ((b) and (c):
P PVT,'VV ,)),, -EE P PIT. hp
errur (t)
(5)
7

---

Published as a conference paper at ICLR 2026
The figure shows that the error accumulates in significant negative steps. These steps occur at token positions where the corresponding attention probability P(T. f is exactly (also observed in other token positions). This happens when the pre-softmax score S[T, is the maximum value in its row, causing exp(S(T. max (A)) to evaluate to exp(O) = 1.
Analysis of Biased Rounding Error Furthermore, the bias arises from the interaction of these unit values with the distribution of the value matrix V. As observed in Fig. a for the problematic feature dimension i = 20, the values of V[:, are predominantly negative. When P(T, t] = 1, the product P[T, [t., i] is simply Vit, d), a negative BF16 number. A systematic error then occurs when two such negative BF16 numbers are added. In floating-poin arithmetic, adding two numbers with the same sign can cause the resulting significand to overflow (e.g., 1xxxxxx 1.yyyy -10. zzzz), requiring right shift and an exponent increment to re-normalize. The bits shifted out of the 7-bit BF16 fraction determine the rounding direction. When adding two negative numbers, the rounding operation (e.g., round-to-nearest can introduce a consistent bias
To illustrate how this rounding bias oceurs, consider the addition of two significands that cause an overflow, requiring a right shift for normalization. The bit that is shifted out (the rounding bit) determines the rounding direction. We show all possible additions of the last two 2-bit numbers, where the green bit represents the bit that would be shifted out (the rounding bit):
00 + 00 = 00
00 + 01 = 01
00 + 10 = 10
00 + 11 = 11
01 + 01 = 10
01 + 10 11 01 11 = 100 10 10 100 10 11 101 11 11 110
Because the summation is performed in FP32, the accumulation of lower-order bits from the small numbers can activate the sticky bit. Consequently, this forces a rounding up when the subsequent BF16 number is added. Therefore, the rounding bit 1 indicates that rounding up required. Because the operands P PTT,t] [t, i] are negative and have large exponent (because P(T, t make PIT, 4] [t. smaller), the error rounding up is magnified, resulting a negative error. When the rounding bit is O, the result is rounded down, introducing positive error. Furthermore, the positive error is smaller compared to the negative error from rounding up, as the values being added usually very small. This asymmetry leads to the rounding error dominated by the rounding up, resulting in negative rounding error as observed in our analysis. This systematic negative bias in the computation of O is the ultimate source of the training failure.
does not
Remark 1. When P(T,t) t] < is less than I and has mix of zero and non-zero bits in its signifi- cand's bits, its product with VII il will also have a non-zero value in those bits. When rounding to BF16, this will not introduce biased rounding error.
Analysis of Rounding Error in Fig. 6 To make this concrete, we now analyze the specific BF16 number addition that causes the large negative error jump shown Fig. ((c). Because the sum- mation is performed in FP32, the first BF16 value is added by some small values from other to- kens before the second BF16 value is added. This can activate the sticky bit, which can force round-up when adding the second BF16 value. For this example, resentations. The first operand is the accumulated sum of previous terms, which includes a BF16 (110000000001000 T0110000000000000000000000000000 the sticky bit. The second operand is another BF16 value. Their FP32 representations are:
start with the FP32 rep-
2.40625) plus small residual (~ 0.00087) that activate
11000000111000011110000000000010
11000000000011000000000000000000 (-2)
Since their exponents are identical, the addition is performed on their significands:
h1tp:/00111110000000110111111111 = 10001011011000111000000111111110
(-
The result overflows the significand's format, requiring normalization. right by one bit, and the exponent is incremented:
The significand is shifted
8

---

Published as a conference paper at ICLR 2026
wold
Stor
(a) Adam
(b) GPT-2S+Muun
(c) 4PT-2M+ W
Figure 7: Validation losses for classical flash attention and the stablized flash attention with AdamW, Muon optimizers, and GPT-2M model.
Significand: 10.0101101 + 1.00101101
Exponent: 10000000 - 10000001
which corresponds to -4. 42199999 To store this result in BF16, must be rounded to fraction bits. The significand is 1. 1.00101101. The 7-bit fraction is 0010110. Because the rounding bit is and there are nonzero bits after the rounding bit, the round-to-nearest (ties-to-even) rule rounds up (adding to the last bit of the fraction):
The exact result in FP32 is 1 110000001 01111000011110110000001111111111
BF16 Fraction: 0010110 + 1 0010111
The final BF16 result is 1100000010010111. which represents -4.71875. This rounded value is more negative than the true sum of -4. -4470399045944223333 The error introduced by this single When such rounding events occur systematically across
addition is 0.0144595540557711 many additions in the PV product, the errors accumulate, creating the negative bias in O that ultimately destabilizes training.
Claim 3. With more than one P = and negative V(t, El, the addition in PV can cause overflow of significand. This necessitates right shift and round-up up of significand due to sticky bit, introducing negative rounding error in O and leading to positive (ōip 8 (7).
4 EXPERIMENT: MITIGATE BIAS IN ROUNDING ERROR
Our analysis traces the failure to biased rounding errors in the PV computation. This occurs when multiple identical maxima in a row of the pre-softmax scores S cause corresponding elements in P to become exactly To validate findings, we modify the softmax to detect this specific condition and adjust the normalization, ensuring all elements of P are strictly less than 1. This prevents the biased rounding and restores training stability.
To prevent biased rounding, we introduce a tar- geted modification to the safe softmax compu-
rowmax(S)), (S), r, = rowsum(r.. -S56 m = where(r,. > 0 A r, > L, Bm),
tation. The core idea is to dynamically adjust
the normalization factor m only when a row of the score matrix S contains multiple identical maximum values. This adjustment ensures that
m = where(rm < 0 AAAr, > m')
P = exp(S m)
the argument to the exponential function, S m. becomes strictly negative at these maximal po- sitions, which in turn guarantees that all elements P = exp(S - m) are less than 1. A naive approach, such subtracting small fixed constant, is insufficient introduces new systematic rounding errors (see Appendix C); hence, a dynamic maximum strategy is required. Our modifica- tion is presented above.
This modification prevents elements of from becoming exactly 1. We first check whether the row maximum Fm. is attained by more than entry by evaluating rowsum (rm S S e), where is

---

Published as a conference paper at ICLR 2026
small numerical tolerance (e.g.. 10 account for praatinn--pin precision. If repeated maxima are detected and I'm > O, the normalization factor is adjusted to m frmm with B > 1. The new maximum in the exponent then becomes - (3 1)rm, which is strictly negative. If instead I' < and is repeated, we set m = O, which likewise ensures that the maximum exponent remains negative. In both cases, the adjustment guarantees that max(S m) < 0 and thus max(P) 1,
preventing the conditions that lead to biased rounding.
Crucially, this modification is mathematically equivalent to standard attention in exact arithmetic, as it leverages the shift-invariance property of the softmax function (soft softmax(2 - c)). Our method simply chooses a different row-wise constant e ensure numerical stability. In our experiments, we set 3 E [2, 8]. as smaller values risk having the result round back to 1, while larger values risk underflow. This modification is integrated into the standard flash attention tiling algorithm (line with magenta in AlgorithmB 3 without altering the backward pass.
To validate our approach, we pre-train the GPT-2S model for 600K steps in BF16 precision using our stabilized flash attention (3 = 2) with both AdamW and Muon optimizers (Jordan et al. 2024)
The AdamW configuration uses the same hyperparameters as described in Section 3.1 For the Muon setup, we employ RMS-matched Muon (Liu et al. 2025) with learning rate of 5 x 10 and weight decay of 0.01, while retaining Adam for embeddings, the language model head, and other one-dimensional parameters. We also extend our experiments to the larger GPT-2M model, training it for 100K steps with the A optimizer under the same hyperparameters as in Section 3.1 As shown in Fig. 7 the modified flash attention successfully stabilizes training for both optimizers and larger model, preventing the loss explosion seen in the original implementation.
CONCLUSION
This paper presents the first mechanistic explanation for a notorious loss explosion in low-precision flash attention training. We pinpoint the root cause to an interplay between emergent low-rank rep- resentations and biased BF16 rounding errors. A minimal, targeted modification to flash attention validates our analysis by restoring stability. Our analytical workflow provides blueprint for di- agnosing similar numerical instabilities in other architectures, scales, and low-precision formats, paving the way for more robust and efficient large-scale model training.
Discussion Our findings are consistent across various hardware (NVIDIA A100, RTX 4090, Huawei Ascend 910B) and mechanistically explain empirical observations of training instability. The growth of weight spectral norms results from the accumulation of a low-rank error matrix in the gradients. We also clarify the role of attention sinks: by attracting high attention scores, they are more likely to produce attention probabilities of 1, which triggers the biased rounding error in the PV computation. This provides a direct numerical link between the architectural behavior of sinks and the arithmetic instability that derails training. Finally, our analysis provides compelling explanation for the effectiveness of established stabilization techniques. We posit that a root cause of training failure is the structural similarity within the error matrices, (PK) X(T), which creates a pathway for rounding errors to accumulate systematically. Techniques like QK normalization and Gated Attention disrupt this underlying structure. In doing so, they ensure that even when rounding emrors are present, they lack the coherence to compound, thereby averting the instabilities observed during training.
Limitations Our analysis focuses on a specific failure case in a GPT-2 model. The generaliz- ability of our findings other architectures, larger scales, different low- wow-precision formats like FP8 requires further investigation. Additionally, our proposed mitigation is tailored to the specific rounding error identified and may not address other sources of numerical instability.
Future Work Our work provides a foundation for exploring numerical stability in low-precision training. Our analytical framework, i.e., locating the source of error, finding similar update direc- tions, and tracing back to root causes, can be applied to other architectures and precision formats. Future work could extend this analysis to FP8 training, larger models, and different architectures. Additionally, developing automated tools to detect and mitigate such numerical instabilities during training would be valuable for the community.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENT
This work is supported by Beijing Science and Technology Program NN0.22511000081250000
REFERENCES
Sami Ben Ali, Silviu-loan Filip, and Olivier Sentieys. A stochastic rounding-enabled low-precision floating-point mac for dnn training. In 2024 Design, Automation & Test in Europe Conference & Exhibition (DATE), pp. 1-6. IEEE, 2024.
Paul Balança, Sam Hosegood, Carlo Luschi, and Andrew Fitzgibbon. Scalify: scale propagation for efficient low-precision Ilm training. arXiv preprint arXiv:2407. 1 17353, 2024.
Charlie Blake, Constantin Eichenberg, Josef Dean, Lukas Balles, Luke Y Prince, Bjöm Deiseroth, Andres Felipe Cruz-Salinas, Carlo Luschi, Samuel Weinbach, and Douglas Orr. u-ppp The unit- scaled maximal update parametrization. arXiv preprint arXiv:2407.17465 2024.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240): 1-113, 2023.
Tri Dao. FlashAttention-2:: Faster attention with better parallelism and work partitioning. In Inter- national Conference on Learning Representations (ICLR), 2024.
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In Advances in Neural Information Process- ing Systems (NeurIPS), 2022.
Maxim Fishman, Brian Chmiel, Ron Banner, and Daniel Soudry. Scaling fp8 training to trillion- token llms. arXiv preprint arXiv:2409. /2517, 2024.
flash-attention Issue 337. Discussion on flash-attention issue #337: pretraining loss blows up with triton flash attention. lattps://githul com Saom/DaooAlln Bataaa//Elaaa-aattetttee
issues/337 2024. Accessed: 2025-09-07.
Aaron Gokaslan, Vanya Cohen, Ellie Pavlick, and Stefanie Tellex. Openwebtext corpus. http: 2019.
Alicia Golden, Samuel Hsia, Fei Sun, Bilge Acun, Basil Hosmer, Yejin Lee, Zachary DeVito, Jeff Johnson, GuoYeon Wei, David Brooks, et al. Is flash attention stable? arXiv preprint arXiv:2405.02883, 2024.
Xin He, Jianhua Sun, Hao Chen, and Dong Li. Campo: Campo:(Cost-a)) (Cost-Aware) performance optimization for MixxddPrrecisoon neural network training. In 2022 USENIX Annual Technical Conference (USENIX ATC 22), pp. 505-518. 2022.
Alex Henry, Prudhvi Raj Dachapally. Shubham Pawar, and Yuxuan Chen. Query-key normalization for transformers. arXiv preprint arXiv:2010.04245. 2020.
Jordan Hoffimann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, al. Train- ing compute-optimal large language models. arXiv preprint arXiv:2203. 15556, 2022.
Tianjin Huang, Ziquan Zhu, Gaojie Jin, Lu Liu, Zhangyang Wang, and Shiwei Liu. Spam: Spike- aware adam with momentum reset for stable IIm training. arXiv preprint 25X1v:2501.06842 2025.
Keller Jordan, Yuchen Jin, Vlado Boza, You Jiacheng, Franz Cesista, Laker Newhouse. and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024, URL https: 7/kel lerellerjrrdan..omm gi thub. io/postss/uont

---

Published as a conference paper at ICLR 2026
Dhiraj Kalamkar, Dheevatsa Mudigere, Naveen Mellempudi, Dipankar Das, Kunal Banerjee, Sasikanth Avancha, Dharma Teja Vooturi, Nataraj Jammalamadaka, Jianyu Huang, Hector Yuen, et al. A study of 6 for deep learning training. arXiv preprint arXiv:1905.12322. 2019.
Kimi-Team. Kimi k2: Open agentic intelligence. arXiv preprint arXiv:2507.2053: 2025.
Joonhyung Lee, Jeongin Bae, Byeongwook Kim, Se Jung Kwon, and Dongsoo Lee. To fp8 and back again: Quantifying the effects of reducing precision IIm training stability. CoRR, 2024.
Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chenggi Deng, Chenyu Zhang, Chong Ruan, et al. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19437, 2024.
Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, et al. Muon is scalable for Ilm training. arXiv preprint arXiv:2502.16982 2025.
Naveen Mellempudi, Sudarshan Srinivasan, Dipankar Das, and Bharat Kaul. Mixed precision train- ing with 8-bit floating point. arXiv preprint 1905. 12334, 2019.
Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, et al. Mixed precision training. arXiv preprint arXiv:1710.032 2017.
Paulius Micikevicius, Dusan Stosic, Neil Burgess, Marius Cornea, Pradeep Dubey, Richard Grisen- thwaite, Sangwon Ha, Alexander Heinecke, Patrick Judd, John Kamalu, al. Fp8 formats for deep learning. arXiv preprint arXiv:2209.05433 2022.
Igor Molybog, Peter Albert, Moya Chen, Zachary Vito, David Esiobu, Naman Goyal, Punit Singh Koura, Sharan Narang, Andrew Poulton, Ruan Silva, et al. A theory on adam instability in large- scale machine learning. arXiv preprint arXiv:2304.0900 098001 2023.
nanoGPT Issue 303. Discussion on nanogpt issue #303: Gradient explosion when training with /033, 2023, Accessed:
69 httt:///ggtaabb...o///aaaaattt/t 2025-09-07.
nanoGPT Issue 524. Discussion on nanogpt issue #524: Training loss becomes nan when using bfloat16. http:////.ttt..cccoco//aaattaaaa 2025-09-07.
524, 2024. Accessed:
nanoGPT Issue 554. Discussion nanogpt issue #554: Loss diverges when using 6 with flash attention. https::/ cessed: 2025-09-07.
com http://tww..uutee...aaaaatttttaa http:///e..aatttt//a//////////// 2024. Ac-
Badreddine Noune, Philip Jones, Daniel Justus, Dominic Masters, and Carlo Luschi. 8-bit numerical formats for deep neural networks. arXiv preprint arXiv:2206.02915. 2022.
Houwen Peng, Kan Wu, Yixuan Wei, Guoshuai Zhao, Yuxiang Yang. 7e Liu, Yifan Xiong, Ziyue Yang, Bolin Ni, Jingcheng Hu, et al. Fp8-Im: Training fp8 large language models. arXiv preprint arXiv:2310.188 18311,, 2023.
Sergio P Perez, Yan Zhang, James Briggs, Charlie Blake, Josh Levy-Kramer, Paul Balanca, Carlo Luschi, Stephen Barlow, and Andrew William Fitzgibbon. Training and inference of large lan- guage models using 8-bit floating point. arXiv preprint arXiv:2309. /7224, 2023.
Zihan Qiu, Zekun Wang, Bo Zheng, Zeyu Huang, Kaiyue Wen, Songlin Yang, Rui Men, Le Yu, Fei Huang. Suozhi Huang, et al. Gated attention for large language models: Non-linearity, sparsity, and attention-sink-fiee arXiv preprint arXiv:2505.06 257700088, 2025.
Qwen-Team. Qwen3 technical report, 2025. URL https: www.i.xiiiiv...c www.iii.orggabaa /tt. 09388
Alec Radford, Jeff Wu. Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
12

---

Published as a conference paper at ICLR 2026
Jack W Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXh::2112.11446 2021.
Samyam Rajbhandari, Jeff' Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Perfor- mance Computing, Networking. Storage and Analysis, pp. 1-16. IEEE, 2020.
Oleg Rybakov, Mike Chrzanowski, Peter Dykas, Jinze Xue, and Ben Lanir. Methods of improving IIm training stability, arXiv preprint TTt.:2410.166822 2024.
Jay Shah, Ganesh Bikshandi, Ying Zhang. Vijay Thakkar, Pradeep Ramani, and Tri Dao. Flashattention-3: Fast and accurate attention with asynchrony and low-precision. Advances in Neural Information Processing Systems, 37:68658-666 68685, 2024.
Hugo Touvron, Thibaut Lavril, Gautier Izacard. Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302 13971, 2023.
Albert Tseng, Tao Yu, and Youngsuk Park. arXiv:2502.20556, 2025.
Training llms with mxfp4. arXiv preprint
Naigang Wang, Jungwook Choi, Daniel Brand, Chia- Yu Chen, and Kailash Gopalakrishnan. Train- ing deep neural networks with 8-bit floating point numbers. Advances in neural information processing systems, 31, 2018.
Ruizhe Wang. Yeyun Gong, Xiao Liu, Guoshuai Zhao, Ziyue Yang. Baining Guo, Zhengjun Zha, and Peng Cheng. Optimizing large language model training using fp4 quantization. arXiv preprint arXiv:2501. 1710////1//// 2025.
Shibo Wang and Pankaj Kanwar Bfloat16: The secret to high performance on cloud -eaae-aannneeeeeerinn// August 2019.
tpus. https://eLoud.gooql....
Accessed: 2025-09-07.
Mitchell Wortsman, Tim Dettmers, Luke Zettlemoyer, Ari S. Morcos, Ali Farhadi, and Lud- wig Schmidt. Stable and low-precision training for large-scale vision-language models. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https:
//openreview. met/forum? oett://orrmmm.ssssccs.p.pp...
d=sqcASmpA2R
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309.17455 2023.
Greg Yang, Edward J Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ry- der, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. arXiv preprint arXiv:2203.03466 2022.
Greg Yang, James B Simon, and Jeremy Bernstein. A spectral condition for feature learning, arXiv preprint (tt::2310.18//.. 2023.
Ruizhe Zhao, Brian Vogel. Tanvir Ahmed, and Wayne Luk. Reducing underflow in mixed precision training by gradient scaling. In Proceedings of the Twenty-Ninth International Conference International Joint Conferences on Artificial Intelligence, pp. 2922-2928, 2928, 2021.
Jiecheng Zhou, Ding Tang, Rong Fu, Boni Hu, Haoran Xu, Yi Wang, Zhilin Pei, Zhongling Su, Liang Liu, Xingcheng Zhang, et al. Towards efficient pre-training: Exploring fp4 precision in large language models. arXiv preprint arXiv:2502.11458, 2025.
13

---

Published as a conference paper at ICLR 2026
A RELATED WORK
A.1 MIXED-PRECISION BFI6 TRAINING.
Contemporary large language model (LLM) pretraining almost universally employs mixed-precision arithmetic. Early efforts by Micikevicius al. (2017) demonstrated that FP16 training- an FP32 master copy of weights and fixed loss scaling could match FP32 accuracy for many models. However, the narrow exponent range of FP16 often causes many gradients underflow, necessitat- ing careful tuning. The bfloat 16 (BF16) format, with an 8-bit exponent and 7-bit mantissa, retains the wide dynamic range of FP32 while halving storage cost. Kalamkar et al. (2019) showed that BF16 achieves convergence parity with FP32 on large models without specialized tuning. Since then, BF16 has become the default 16-bit format for many large-scale training frameworks, with native support in PyTTrch and TensorFlow (Wang & Kanwar 2019).
could
BF16 mixed precision has enabled training of landmark L.LMs at unprecedented scales, including GPT-3 (175B) (Brown et al. 2020), Google PaLM (540B) (Chowdhery et al., [023), DeepMind Gopher (280B) (Rae et al. (021), Chinchilla (70B) (Hoffmann et al., 2002),, and Meta's LLaMA
family (7B-6BB) (Touvron et al. 2023). To handle the massive memory footprint, parallel training frameworks such as Megatron and DeepSpeed integrate BF16 training with techniques like the Zero Redundancy Optimizer (ZeRO) (Rajbhandari et al. 2020
Despite its advantages, BF16 training can still exhibit instabilities. Empirical studies show that FP16 is highly unstable even with loss scaling, whereas BF16 eliminates most precision-related tuning and faitures Wang & Kanwar (2019). However, [ee et al. (2024) report that roughly 10% of GPT-2 pretraining runs diverged under pure BF16, compared 0% under TF32. This suggests that while BF16 substantially improves stability, complementary stabilization techniques remain necessary at scale.
A.2 STABILIZING LOW-PRECISION TRAINING
Gradient Scaling, Early work by Micikevicius et al. (201)) introduced FP16 mixed precision training. where weights, activations, and gradients are stored half-precision while maintaining master FP32 copy. They also proposed loss scaling to prevent FP16 underflows. Even with loss scaling, some underflow can still occur in deep networks. To address this, Zhao al. (202T) introduced gradient scaling, which dynamically computes per-layer scaling factors to avoid both underflow and overflow.
Ultra-Low-Preeision (FP8/INT8) Training. To further reduce cost, recent works explore FP8 or INT8 precision for training and inference. However, naive FP8 training tends to diverge. Lee
et al. (2024) note that the direct application of FPS to LLM training is unstable without additional stabilization techniques. To address /erez et al. (023) propose dynamically adjusted per-tensor scaling factors for FP8 matrix multiplications. Using this scheme, they successfully train GPT- and LLaMA-style models up 70B parameters entirely in FP8. Similarly, [Peng et al.) (2023) introduce FP8-LM, a framework that progressively applies FP8 to gradients, optimizer states, and distributed communication, achieving a 39% memory reduction and 75% speedup compared to BF16. [allaçça et al. (2024) present SCALIFY, which propagates scale factors throughout the computation graph to ensure stable FPS operations without manual tuning. These approaches collectively demonstrate that careful scaling management enables FP8 or INT8 training to match BF16 performance while reducing memory and compute requirements.
Optimizer and Gradient Stabilization. Optimizer algorithms play a critical role in training sta- bility. Molybogg al. (2023) theoretically analyze Adam and show that catastrophic divergence often arises when the update direction becomes uncorrelated with the true descent direction large-scale models. To address gradient instability, HHang et al. (2025) propose SPAM (Spike-Aware Adam with Momentum Reset), which detects and mitigates rare but severe "gradient spikes" by resetting momentum and applying spike-aware clipping. In parallel, Wortsman et al. (2023) investigate loss spikes in lanounlanguage models and show that Adam W often underestimates the second moment before spikes occur. They propose hybrid Maaa-AdaFactoo optimizer that adaptively corrects
14

---

Published as a conference paper at ICLR 2026
recond-moment underestimation, outperforming gradient clipping alone. These methods highlight how optimizer modifications directly mitigate divergence in low- precisioc regimes.
Activation and Architectural Techniques. The choice of activation functions and initialization strategies also impacts stability. Fishmann et al. (0024) observe that the SwiGLU activation amplifies outliers during long FP8 training runs. They introduce Smooth-SwiGLU, a modified activation that prevents outlier amplification, enabling stable trillion-token FP8 training. In the vision-language domain, Wortsman et al.] (2023) show that "layer-scale zero" initialization and carefully designed low-precision linear layers (e.g., SwitchBack) further improve stability in int8 training.
A.3 FLASH ATTENTION
Flash Attention (Dao et al., 2022 an 1/O-aware exact attention algorithm designed to overcome the memory bottleneck of standard self-attention, which scales quadratically with sequence length N. Standard attention requires materializing the full N x N attention score matrix, leading to O(N2) memory usage that becomes prohibitive for long sequences. Flash Attention mitigates this by employing a tiling and kernel fusion strategy. The query, key, and value matrices are partitioned into smaller blocks, which iteratively loaded from slow high-bandwidth memory (HBM) into fast on-chip SRAM.
Within the SRAM, the core computation is performed. A key innovation is the of online softmax algorithm. Instead of computing the full score matrix before normalization, Flash Atten- tion maintains running statisties (the maximum score and the sum of exponentials) as it processes blocks of keys and values. This allows it to compute the correctly normalized output block-by- by-boock without ever storing the entire intermediate attention matrix. By fusing the matrix multiplications and softmax operations into single GPU kernel, it drastically reduces the number of memory read/write operations to HBM, which is the primary performance bottleneck. This 1/O-aware de- sign reduces the memory complexity from ((N2) to (NN) and significantly accelerates computa- tion. Flash Attention (Dao 2024) further optimized this approach by improving work partitioning among GPU thread blocks and warps and reducing non-matmul computation overhead, achieving near-optimal performance and making it an indispensable component for training and inferring long- context LLMs. See Algorithm 1 and Algorithm ( for pseudocode of Flash Attention 2's forward and backward passes.
Algorithm Flash Attention: Forward Pass
Require: Matrices Q,K,
K, V Nxd block sizes Be, By.
1: Divide Q into T,
N B
blocks QI,..., Qr, of size B, each, and divide K, V in to
To N B.
blocks K K....., Kr. and Vi,..., VT., of size B, x each.
2: Divide the output O E NRd into T, blocks OI...., Or, of size B, x d each, and divide the logsumexp L into T, blocks Lj...., LT, of size B, each.
3: for s 17.do do
4: Initialize (0))
(0) E B.Bxd dl0) (0) E RR m (0) = (-00) E B.B.
5: for 1 و T. do
6: Compute s(
Q B...B.
Compute m()) max max(m (m (j-1) rowmax(S!")) E RBB. púi exp(s!)) m())) + rowsum(P!))) E B.B. j0 10.1-11 p'lv.
7:
B. (pointwise). ef-) m --1)mm
mm e
Compute 0.0) d 1- m
8:
9:
end for
Compute O (To) (T))
10:
Compute L, = m (T.) + log( (T)))
11:
12:
Write O, the i-th block of O.
13:
Write L, as the i-th block of L.
end for
15:
Return the output O and the logsumexp L.
15

---

Published as a conference paper at ICLR 2026
Algorithm 2 Flash Attention: Backward Pass
Require: Matrices Q, K, V, O, dO E RNxd vector L E IR,, block sizes Bo, Br.
1: Divide Q into T
N B+
blocks QI,..., Qr, of size B, each, and divide K, V to
T. = B. blocks ,,...., Kr. and VI...., VT., of size Be x d each.
2: Divide O into T, blocks O OI,...; Or, of size B, d each, divide d0 into T, blocks dOi,..., ddO,. of size B, x d each, and divide L into T, blocks Li....,
. Lr. of size B, each. 3: Initialize dQ (0) Nxd and divide it into T, blocks dQ...... ......QQ. of size B, x each. Divide dK, dV E R Wxd in to T. blocks dKi.... dKr, and dV dV...... dVr.,. of size Be x d each.
4: Compute 5 rowsum(do o O) IRN (pointwise multiply), and divide it into T, blocks
81 81,..., δr, of size B, each.
5: for 1i T, do
6: Initialize dK, (0) Bgxd; dV, = (0)
7: for 1 < < T, do
8:
Compute s/ Q.KÝ E BRB.BB.
Compute P (1) = exp(S.j - L;) E B.B.B.
9:
Compute dV dV; (P())) d0, BB
10:
Compute dP () do,V, ER B.
11:
12:
Compute ds -P(3) o (dP! (i)
6,) E B.
Update dQ, dQ, as U E R Bddd
14:
Compute dK, + dK, + ds
dQ; E E Bxxd
end for
16:
end for
17:
Return dQ, (K, dV.
B BF16 ADDITION
The bfloat16 (Brain Floating-Point) format is 16-bit floating-point representation widely used in deep learning for its balance between computational efficiency and numerical range. consists of 1 sign bit, 8 exponent bits, and 7 fraction (or mantissa) bits. This structure gives 6 the same dynamic range as the bi-
format (FP32) but with significantly less precision.
The addition of two 6 numbers, say and 6, follows the standard procedure for peating-point arithmetic:
Exponent Alignment: The exponents of the two numbers are compared. The number with the smaller exponent has its significand (the combination of the implicit leading bit and the fraction) shifted to the right until its exponent matches the larger one. Each right shift increases the ex- ponent by one. Bits shifted past the available precision lost, which is an initial source of error.
The aligned significands are added together. The sign of the result is determined by the signs and magnitudes of the operands.
Significand Addition:
3. Normalization: The result is normalized to ensure it conforms to the 1.xxxx... 2 format. the addition resulted in an overflow (c.g., 10. 30.xxxx...). the significand is shifted right and the exponent is incremented. If it resulted in cancellation (e.g., 0.00xx...), the significand is shifted left and the exponent is decremented until the leading bit 1.
The resulting significand, which may have more than 7 fraction bits after normaliza- tion, must be rounded. The standard mode is "round to nearest, ties even". This means if the truncated part is greater than half the value of the last storable bit (LSB), the number is rounded up. If is less, is rounded down. If is exactly half, it is rounded to the nearest value with an even LSB.
Rounding:
The key source of numerical error comes from steps 1 and 4. During exponent aligmment, precision is lost from the smaller-magnitude number. After addition, the result must be rounded back to the
16

---

Published as a conference paper at ICLR 2026
7-bit fraction, which introduces another rounding error. While "round to nearest, ties even" is designed to be unbiased for random data, sequence of additions on data with a specific distribution (e.g., mostly negative numbers being added together) accumulated error consistently pushes the result in one direction. This accumulation of biased error critical factor in the training failure observed in low-precision settings.
lead to biased rounding error, where the
C DESIGN CONSIDERATIONS FOR MITIGATING BIASED ROUNDING ERROR IN FLASH ATTENTION
Use Dynamic Maximum Value Rather Than Fixed Offset A fixed offset would cause the com- puted values of P be consistently rounded one direction during BF16 conversion, introducing a fixed error. Since the elements of V often share the same sign, this fixed rounding error in P does average out to zero when computing PV. This leads biased error in the output O, which in turn creates a biased term ö, reintroducing the very failure we aim to solve.
Dynamic Maximum is Applied Conditionally Our modification is applied conditionally- when a row contains multiple identical maximum values- to avoid introducing new numerical insta- bilities. An unconditional adjustment is not better option. For example, if row has single, very large positive maximum value Im. applying our rule would mean calculating exp(S - 3.rm.. The largest term in the exponent would become -(3 1)rm. and exp(-(3 1)rm) could underflow to zero. This would cause the normalization factor to become zero, leading to a division-by-zero error when computing the output O. By applying the modification only the specific case that causes biased rounding, we preserve the numerical stability of the standard online softmax in all other scenarios.
Explanation on Dealing Negative Repeated Row Maximum We also explore alternative stabi- lization methods for the negative, repeated maximum (r.m. < 0). One approach involves setting the normalization factor m TTm for some 7 E (0, 1). This makes the new maximum value the close 1, this new maximum approaches zero. In low-precision arithmetic, exp((1 - 7) rm can round to exactly 1, reintroducing the very failure we aim to prevent. Consequently, we found that setting 7 (i.e., m 0) robust choice, as it ensures the maximum value in the exponent remains sufficiently negative.
exponent (1 7) However, we observe that if y
D SIMILAR PATTERNS DISCOVERED IN LLAMA-3.1-8D
To see if the phenomena that leads to the failure case in GPT-2 also exists in other models, we present (PK) [T) X(T) and the attention score Alama-3.1-80 which shows the similar patterns we discovered in GPT-2. This section is not meant to say Llama 3-B will definitely fail in BF16 training, but rather show that the phenomena we discovered not unique to GPT-2.
D.1 STRUCTURALLY SIMILAR (PK) [(T]
Given an input "The quick brown fox jumps over the lazy dog", we visualize [PK] for attention head 13 in layer 1 of Llama-3.1-8B. We visualize the representations and similar columns in Fig.
D.2 MULTIPLE MAXIMUM P
We also observe multiple maximum values in the attention scores of Llama-3.1-88 3.1-BB. We visualize some examples in Fig. 9 This experiment shows the attention sink (Xiao et al., 2023) behavior as well, where the first token has high attention scores.
17

---

Published as a conference paper at ICLR 2026
layer
ofvestte 177
(c)
(d)
(e)
(1)
Figure 8: Similar (PK)[T] X[T] discovered different tokens in Llama-3.1-8D1 (c)-(f) show the
similar columns.
OTHER DETAILS
E. RELATION BETWEEN MULTIPLE MAXIMA AND LOss
To investigate the connection between the multiple maxima phenomenon and training instability, plot the frequency of multiple maxima occurrences against the training loss in Fig. [3] The visualization reveals a clear temporal pattern: the number of multiple maxima begins to iccraase before the loss explodes (see 7000 step). This strong correlation suggests that the emergence of multiple maxima is leading indicator of the impending training failure.
E.2 DERIVING EQUATION
The following equations are prepared to derive equation I in the main text.
dQ dSK
dS = o o (dP 5)
(7)
Here the detailed derivation.
We begin with the goal of deriving an expression for the difference dQuip dQup-
Apply the definition of dQ. We start with the left-hand side of the equation and substitute the definition from equation 6 for both the high-precision (hp) and low-precision (Ip) terms.
dQuup dQip = (dS,, - (dS
Using the right distributive property of matrix multiplication, we can factor out the matrix K.
d dQup dQip (dSh.p dSip) K
(8)
18

---

Published as a conference paper at ICLR 2026
layert
Upportte qunrystt
Adwelee Proguityy
Key Postios
Kay Powidon
(a)
(b)
1aysm26,, haspei quaryeds
byyemos heade1, quarymi2
Mlordion Propad
0.
Key Postion
May Pan
(d)
querynds
0.4
Poskiie
Pestiie
(e)
(I)
Figure 9: Two maxima observed in Llama-3. 1-8B. The first token has maxima due to the attention sink.
Substitute the definition dS. Next, we substitute the expressions for dSap and dSi using the definition from equation
dS = aP (dP -kp))
dS aP (dP -dtp)
Plugging these into our expression equation
dQup dQip [aPP (dP ō (aP (dP -1p))) K
Simplify using properties of the Hadamard product. The Hadamard product (o) is distributive over matrix addition and subtraction, and we can factor out the common scalar and matrix P.
dQup dQip = (a ( (P (dP -kp)) P (dP 5tp)) K
[aP ((dP 8np) (dP 8tp) K
Now, we simplify the expression inside the parentheses. The dP terms cancel each other out.
dQip dQip aP o (dP - -ōhp dP + 6 K
oP (8lp 5 K
4. Convert the Hadamard product to matrix multiplication. The final step involves rewriting the row-wise scaling operation, represented by the Hadamard product with broadcasting, as a standard matrix multiplication. Let the vector difference be dip dep. The expression P (ōtp 6hp)
19

---

Published as a conference paper at ICLR 2026
Algorithm 3 Stabilized Flash Attention by Mitigating Biased Rounding Error: Forward Pass Require: Matrices Q, K, V E RRxxa, block sizes Be, B,, 3 > 1.
1: Divide Q into T,
N
blocks Qu...., Qr, of size B, each, and divide K, V to
B
Tc = B. blocks K Ki...., Kr. and Vi...., VT., of size Be x d each.
2: Divide the output O E IR into T, blocks O....., Or, of size B, x d each, and divide the logsumexp L into T, blocks Li...., Lr, of size B, each.
3: for 1 T, do
4: Initialize O = (0) B. xd E ERBxxd, ( (0) f E ER m (0)
(-00) E ERA,
5: for 1 & Te do
6: Compute S Q B.
7:
rowsum(rm S Se) SS)
8:
Im rowmax
m G)) where(Tm mm 00 r 1, 3r I'm
9:
()) where(rm OAAr, 1, m, (i3))
10:
Compute m (5) = max(m: (j-1)/ Im)
E RRB, p0) = exp(S - m E R.B.B.
11:
(pointwise), e(d) m (11-1) -m el (1i-1)
+ rowsum (Pl)) E E.B.
Compute 0.50 diag(em. -m jo (1-1) +p
13:
end for
Compute O, (i)) - (T.)
Compute L, m (To) + log( ((7)))
15:
Write O, as the i-th block of Ó.
17:
Write L; as the i-th block of L.
end for
19:
Return the output O and the logsumexp L.
means that each of matrix P is element-win multiplied by the vector Specifically, the i-th row of P is scaled by the i-th element of v.
Let's examine this at the element level. The (i, element of the resulting matrix is: (Po Pij V;
This is precisely the result of pre-multiplying the matrix P by diagonal matrix whose diagonal entries are the elements of the vector Let D = diag(v) diag(d;p 8hp). The (i, element of the product DP is:
(DP), D.
Since D a diagonal matrix, Dik is non-zero only when = i, where D = Vi- Thus, the sum simplifies to:
= D. Pij
This confirms that Po diag(v)P. Applying this identity to our expression equation 10
dQupp dQip [ax diag(d;; P) K
regroup the terms to arrive at the Finally, using the associativity of matrix multiplication, final equation.
dQup - dQip = a diag - 6kp) This completes the detailed derivation from the initial definitions to the final expression.
(10)
F SUGGESTIONS FOR HANDLING FAILURE IN LOW-PRECISION TRAINING
Our analysis provides blueprint for diagnosing numerical instabilities that may arise in other low- precision formats, such as FP8. The analytical workflow presented in this paper is possible to be generalized other settings. The key steps
20

---

Published as a conference paper at ICLR 2026
val/loss
10
8
6
4
2
Step
0
20
40
60
80
(a) Loss nanoGPT Issue #303
(b) Loss in nanoGPT Issue #524
Figure 10: Loss curves of two independent runs of GPT-2 training with flash attention in BF16 reported in the Github issue nanoGPT.
1. Isolate the Source of Error: Systematically narrow down the failure to specific operation module through targeted high-precision substitutions.
2. Identify Error Accumulation Mechanisms: Analyze the gradient error to find structural patterns, such the emergence similar update directions that allow errors accumulate rather than cancel out.
3. Trace to Root Arithmetic Cause:
Examine the low-level arithmetic the specific format to find the origin of any systematic bias.
By following this methodology, researchers and practitioners in this field move from observing a high-level training failure to understanding its root cause the underlying hardware arithmetic, enabling the development of principled solutions.
21

---

Published as a conference paper at ICLR 2026
THE USE OF LARGE LANGUAGE MODELS (LLMS)
We use LLMs for various purposes, including: polishing the language of the paper and assistant code implementation. We carefully review and edit all content generated by LLMs ensure accuracy and appropriateness.
22

---

Published as a conference paper at ICLR 2026
Spectral Norm Across Layers and Training Steps
aattn w
C amn weightk
Cattoo atn gro.
Cattn proj weight
oumbo weigggh
mmipp proj weight
Lattne, wnngwogg
ww k
tamoo art we ght
1 puprej weight
t.mppe to weight
mipp mmppcc_re cuprej weight
aattoe weighti q
2.tt cuattn we k
2.mmn ait we w
2 cuproj weight
2.mpp e_ In weight
2mip cuproj weight
Sattne_wit ent weight q
3.atnn oit we ghikk
cuattn v
3.atin. pro weight
Sumpp is weight
aumbo.o weight
4attii weig.we
4.atn.o
ghik
4.tti o cuatth we ght
4.atho.o proj weight
4.mlp. weiight
4.mip.e weigght
satiii,, atn.welgnt.
-70
watree e we k
a.attn.o o orn we ghov
Model Layers
s.attrie, puos wwiggtt
Bumboue o fo weight
Simipa weight
atthuu, with wolggnt
B.itt e with ww k
ght.v
witn cu weight
Bimpp e, fo weight
cuproj teight
Pattne, e_ with weight q
10
Taatn.c_am giikk
7.rtn with wow ght
7.athnc_ppr c_poo weight
7 cu..c_cfoo weight
7.mpp c weight
Batiic weighti
B.atn.o, giik
aattie_att we ghtvv
s.attr.o projuweight
Bumbo o weeight
Bumipuo. weight
attnic wolgggggh
B.amn.o with we ghtk
B.atn.o, ait ghtv
Baath o proj weight
Bumbp.c o to wewiight
aumipie, proj weightt
10 atttt arth wolght
10 teattoeeem w
att.o we v
10.wth c_ proj weight
10.mlple to weightt
10 cuprej weight
11 attioe wiighii
11attn custtn k
11.tiii momeeeatin with ww ght
11.ttin cupme weight
11 mimplee fe wwight
11 mimpl weigh
2000
Au8i
10000
12000
Training Steps
Figure 11: Spectral norm across layers and training steps.
23

---

Published as a conference paper at ICLR 2026
Lifent
Figure 12: Token difference visualization
Validation Loss and SFA Condition Activation During Training
Number of Astivatec SMK Concition
Number Activated SFA Condition
Validzation Loss
7-
Training Steps
Figure 13: Correlation between multiple maxima occurrences and loss curve.
24


Reviews:
[{"review_id": "qmP10k5U1J", "reviewer": "Reviewer_7ALf", "summary": "This work provides the first mechanistic explanation for why transformer training with FlashAttention in low precision (bf16) leads to unstable loss spikes. Through detailed analysis and ablation studies, the authors attribute the training instabilities to two related effects: 1. the emergence of similar low-rank matrix updates in the model's internal representations across tokens and steps, and 2. biased rounding errors in bf16 arithmetic when multiple attention scores achieve the exact same maximum. Together, these create a biased gradient direction that inflates spectral norms and thus destabilizes training. Based on their mechanistic analysis, the authors propose a minimal “dynamic softmax” patch to the FlashAttention algorithm and show that it restores stable training.", "questions": "- In Section 3.3.1, the authors observe that the outer products $(PK)[T]^T X[T]$ exhibit strong structural similarity and approximate them with a common low-rank matrix $R$. Did the authors investigate why this shared structure emerges? For example, could it reflect commonly-occuring tokens, correlated attention patterns across tokens, or a form of mode collapse within the model's internal representations?\n- In Section 3.3.2, the authors attribute the biased rounding error to rows where multiple elements of $S$ share the same maximum in bfloat16, yielding several $\\bar{P}[T, t] = 1$. It would greatly bolster the proposed mechanisms to provide an experiment measuring how frequently this condition occurs during training, and showing whether its occurrence correlates with loss spikes (e.g. in Figure 8).\n- Relatedly, in the discussion, the authors suggest that attention sinks may contribute to loss spikes by producing more attention probabilities of 1 (in bfloat16) and thus triggering the biased rounding error highlighted in Section 3.3.2. Did the authors conduct experiments testing this connection?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- The work investigates a very important problem, loss spikes during Transformer training due to low-precision using FlashAttention. Since attention and FlashAttention are the workhorse of modern AI systems, this problem is of great interest to the community at large.\n- The work is very thorough, with systematic ablations into every step of the FlashAttention algorithm to isolate the source of the unstable training. The authors provide a thorough characterization of the failure modes and, using this insight, propose a simple fix that alleviates the loss spikes.", "weaknesses": "Overall, the analysis and ablation studies are thorough, and the proposed fix is convincing. However, one key piece of the story remains somewhat unclear: the origin of the shared low-rank structure $R$ in Section 3.3.1 (see questions). An investigation into where this structure comes from (e.g. feature collapse, attention sinks?) would make the characterization more complete.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZrSMvewyzc", "reviewer": "Reviewer_8vQG", "summary": "This paper proposes both an explanation and a solution to the problem of training instability for models using bf16 flashattention. They isolate the problem to a biased rounding error in the computation of delta, which produces correlated (low-rank) updates to the weight matrices, eventually leading to divergence. The paper proposes a modification to the flashattention algorithm to resolve this issue and demonstrates empirically that this modification stabilizes training.", "questions": "Questions:\n* Can you explain how existing stability fixes (e.g., QK norm) fit into your explanation? I think this is related to my question regarding how the persistent small bias leads to the large spike in loss.\n* Related to my other questions, shouldn't we expect gradient dynamics to correct the small rounding drift if the drift leads the model to higher loss regimes?\n* Is there a reason you set m' = 0 in the case that r_m < 0? why not set it to gamma*r_m, where 1>gamma>0?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "* The paper proposes a plausible explanation for the observed training instability.\n* The paper proposes a sensible fix to this problem and demonstrates that it works.\n* To the best of my knowledge, the paper's ideas and contributions are original.\n* The ideas are mostly clearly explained in a step-by-step manner that allows the reader to follow along.\n* The problem is very significant, as low-precision flash-attention is very popular. For example, I have personally encountered issues related to training stability, as have other researchers I know.\n\nOverall, I quite like this paper.", "weaknesses": "I would separate the weaknesses into two categories.\n\nWeaknesses in the explanation for bf16 attention instability:\n* The identified cause is a small but persistent bias. However, the observed effect is a sharp, near-instantaneous spike. The paper does not explain why the sharp spike occurs or why the persistent bias causes it. This is the most significant weakness, in my opinion.\n* Minor: I believe remark 1 is slightly imprecise. For example, if $\\bar{\\mathbf{P}}[T,t] = 1/2$, its product with $\\bar{\\mathbf{P}}[t,i]$ would have last 16 bits 0. However, the intended meaning of the remark, that with high probability the product has non-zero least 16 bits, is true.\n\n\nWeaknesses in experimental results:\n* The findings isolating the source of training instability could be validated on other models, sizes, and/or training scripts.\n* The evidence that the stabilized flash attention is actually stable could be more robust. The authors could test more model sizes, train for longer, etc.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kq2YA2WWLF", "reviewer": "Reviewer_zAjM", "summary": "This paper explains why training Transformers with Flash Attention in low-precision (BF16) settings often leads to loss explosions. The authors show that the problem comes from an interaction between low-rank representations and biased rounding errors in BF16 arithmetic. They locate the issue in the backward pass and propose a simple fix to prevent the rounding bias. The proposed method stabilizes training and confirms their analysis, offering both practical and theoretical values for improving numerical stability.", "questions": "1. It would be more readable to include the definition of Flash Attention (Line 98) in a self-contained manner (e.g. in appendices). \n2. The notations, such as \"$d\\mathbf{Q}, d\\mathbf{K}, d\\mathbf{V}, d\\mathbf{O}$\", are somewhat confusing and require additional clarifications to help distinguish them from derivatives. \n3. The derivation presented in Line 246 appears a bit abrupt; providing more intermediate steps would help improve the logical flow. \n4. Claim 3 in Line 424 contains a typo (\"significant\" is misspelled as \"significand\").", "limitations": "", "rating": 4, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. This paper provides a mechanistic explanation for the long-standing issue of low-precision training failures in Flash Attention.  \n2. The authors identify the root cause as a vicious cycle between emergent low-rank representations and biased BF16 rounding errors.  \n3. The paper proposes a fix through a dynamic softmax adjustment that stabilizes training.", "weaknesses": "1. The analysis and solution are validated only on GPT-2, leaving their effectiveness on larger models unverified. For example, if the failure is located at layer 2, head 8, how can we identify the corresponding head in larger models? \n2. I wonder whether the analysis can be generalized to other numerical precisions beyond BF16 and FP32. \n3. The dynamic maximum adjustment introduces additional operations into Flash Attention's critical path. Whether do these operations incur extra computational overhead? What is the additional computation complexity?   \n4. The proposed fix is verified only on a single model. Can this strategy be applied to other scenarios as well?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "htPGMXsqHv", "reviewer": "Reviewer_RoBX", "summary": "This paper investigates a long-standing training instability in transformer models when using flash attention with BF16 precision, which manifests as catastrophic loss explosions during training. It provides a mechanistic explanation for this failure as (1) the emergence of structurally similar low-rank representations across different tokens and training steps in the attention mechanism's intermediate computations, and (2) biased rounding errors inherent in BF16 arithmetic operations. Through systematic experimental analysis on a reproducible GPT-2 failure case, they trace the error propagation, demonstrating how biased gradient updates accumulate in a specific low-rank direction rather than canceling out. The bias originates from a specific numerical condition which is detailed. To validate their analysis, the authors make a small modification to the flash attention's softmax computation, showing this stabilizes training while remaining mathematically equivalent.", "questions": "The analysis focuses on BF16 arithmetic. The deep learning community is increasingly interested in even lower precision formats like FP8 or INT8 for training. Are you aware of other outstanding (BF16, FP8, INT8) issues that would benefit from this type of analysis? If so, consider listing them in an Appendix to promote their prompt solution by early-career researchers.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "This is an excellent paper.\n\n**Clarity and Exposition:** Remarkably clear writing, particularly given the intricate technical details included. The notation is well-defined, figures are informative and well-integrated, and mathematical derivations provide appropriate detail. It reads very smoothly.\n\n**Methodical Investigative Structure:** Reads like a detective story, systematically tracing the failure from symptom to root cause, via targeted experiments. The step-by-step elimination builds a clear causal chain that is highly rigorous.\n\n**Highly Convincing Analysis:** The multi-layered evidence is persuasive. The successful flash attention code modification based on insights gained supports the preceding theoretical analysis with a long-term solution for the community.", "weaknesses": "Nothing major to speak of.\n\nThe focus of the paper is (by design) very narrow.  \n\nDetail is lacking at a couple of points. But these ommissions are minor, and given the extensive detail included, I assume this was done to meet the page limit.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)