Low-Precision Attention May Have Biased Rounding Errors

低精度Attention可能存在有偏的舍入误差 苏剑林 · 2025-10-27 · Machine Learning approximation analysis optimization attention

Some time ago I came across a paper on arXiv, Why Low-Precision Transformer Training Fails: An Analysis on Flash Attention, whose experimental observations closely match some of the phenomena we saw while training Kimi K2 — for instance, in both cases the problem starts to appear from the second Attention layer onward. The paper attributes this to an inherent biased error in low-precision Attention, which is a rather unexpected angle of analysis, so I read through it with some interest.

However, I found the paper's presentation somewhat hard to follow — admittedly partly because I myself am not very familiar with low-precision arithmetic. In any case, after repeatedly asking the authors for clarification, I finally managed to understand the paper, and I'm recording my understanding here for reference.

Summary of the Conclusion

It's worth pointing out that although the paper's title calls out "Flash Attention," according to the paper's own description, the same problem still occurs even when the block size is set as large as the training length. So Flash Attention's block-wise computation is not the cause of the issue, which means we can simplify our analysis by working with a naive low-precision Attention implementation instead. more

For simplicity, we'll only analyze single-head Attention. Let $\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}\in\mathbb{R}^{n\times d}$, and denote $\boldsymbol{S} = \boldsymbol{Q}\boldsymbol{K}^{\top}$, where the bold $\boldsymbol{1}$ denotes the all-ones matrix of shape $n\times 1$, and $\boldsymbol{S}_{\max}$ denotes the $n\times 1$ matrix obtained by taking the row-wise maximum of $\boldsymbol{S}$. Then

\begin{equation}\boldsymbol{O} = \frac{\exp(\boldsymbol{S})\boldsymbol{V}}{\exp(\boldsymbol{S})\boldsymbol{1}} = \frac{\exp(\boldsymbol{S} - \boldsymbol{S}_{\max})\boldsymbol{V}}{\exp(\boldsymbol{S}- \boldsymbol{S}_{\max})\boldsymbol{1}}\end{equation}

We write $\bar{\boldsymbol{P}} = \exp(\boldsymbol{S} - \boldsymbol{S}_{\max})$, so the key computation in Attention is the matrix multiplication $\bar{\boldsymbol{P}}\boldsymbol{V}$, which is generally carried out in BF16 precision. The paper's conclusion is: under low-precision computation, this step $\bar{\boldsymbol{P}}\boldsymbol{V}$ has a biased rounding error. That is, over the long run on average, the expectation of the difference between the low-precision computed $\bar{\boldsymbol{P}}\boldsymbol{V}$ and its exact value is not zero.

As a result, biases across different training steps can keep accumulating, potentially leading to MaxLogit explosion, loss spikes, and other issues, up to and including training collapse. Strictly speaking, of course, this is only one possible mechanism behind MaxLogit explosion and similar issues — not necessarily the whole story — but even so, it's well worth studying and thinking through.

Round-to-Even

To understand the paper's conclusion, let's first review some basic facts about rounding errors. The reason I'm writing this section is, as mentioned at the outset, that I myself am not familiar with low-precision arithmetic — so this section is entirely for my own benefit, to fill in the fundamentals. Readers already familiar with this can skip ahead.

We know that the common rounding method is "round half up": in decimal, for a positive number with one decimal digit that we want to drop, 0–4 rounds down to 0, producing an error of $0,-0.1,-0.2,-0.3,-0.4$; 5–9 rounds up to 10, producing an error of $0.5,0.4,0.3,0.2,0.1$. Have you noticed that the average of these errors isn't 0, but 0.05? That is, "round half up," on average, tends to inflate the original number, producing a positive bias.

Of course, the relative bias shrinks as more digits are dropped — for example, if a number with 2 decimal digits has both digits dropped, the average error becomes 0.005. But regardless, this positive bias from "round half up" always exists, just varying in magnitude. The root of the bias lies at the midpoint: for instance, 0.51 and 0.49 round up/down respectively, and their errors exactly cancel out, but for 0.50, whichever direction we choose to round it — up or down — there's no other number to cancel its error against.

To eliminate this bias, IEEE 754 introduced the "round-to-even" rule: for the midpoint case, rounding should go toward whichever neighbor is even. For example, 2.5 rounds down to 2 when the last digit is dropped, but 3.5 rounds up to 4. This way, "5" has an equal chance of producing an error of $\pm 5$ in either direction, so the average error becomes zero, eliminating the bias.

Back to computing. As we know, computers use binary, which has only 0 and 1, so 1 plays the role that "5" plays in decimal. The bias in binary "round half up" is even more vivid, because the last bit can only be 0 or 1: if it's 0, naturally nothing needs to change, but if it's 1, "round half up" is triggered and we carry the 1. So when a binary number's last bit is dropped under "round half up," the result must be greater than or equal to the original number — which is exactly why "round-to-even" is also needed here to eliminate the bias.

BF16 Addition

Next, let's review the BF16 format. BF16 represents a floating-point number using 16 bits: 1 sign bit, 7 mantissa bits, and 8 exponent bits. The 8-bit exponent gives it the same representable range as FP32 (1 sign bit, 23 mantissa bits, 8 exponent bits), which is exactly why it has become the dominant floating-point format for LLM training today.

BF16 keeps a relatively large number of exponent bits, at the cost of fewer mantissa bits — hence its lower representable precision. To mitigate the accumulated error caused by low precision, BF16 arithmetic adopts the strategy of "accumulate in FP32": that is, when BF16 numbers are summed, they are first converted to FP32, added together in FP32 space to get an FP32 result, and only then converted back to BF16.

Now let's consider two BF16 numbers with the same sign and the same exponent being added. Why choose the case of equal exponents for analysis? Because we want to estimate the error, and equal exponents mean the two numbers are of the same order of magnitude — which is precisely when addition is most likely to produce the largest error. For example, if two numbers being added differ by a factor of 100, then even if we simply return the larger one, the error would be no more than 1%. So the largest errors tend to occur when numbers of the same order of magnitude are added.

When two BF16 numbers with the same sign and exponent are added, a carry inevitably occurs — for example, "1.0000001 + 1.0000100 = 10.0000101 = 1.00000101 × 10," where the exponent must be incremented by 1, and the last bit "1" must be dropped to convert back to BF16 format. As discussed in the previous section, if the last bit is dropped via "round half up," a positive bias results. But as we already know, scientists long ago discovered this bias, which is precisely why "round-to-even" was introduced to eliminate it.

Two Large, One Small

So far, then, everything is within controlled and expected bounds — no bias has appeared yet. However, as is often the case, the unexpected happens.

Now let's consider adding three numbers with the same sign, with the following characteristic: two of them have the same, large exponent, while the third is very small. For instance, take the previous section's example "1.0000001 + 1.0000100" and add "0.0000000001" to it: "1.0000001 + 1.0000100 + 0.0000000001 = 10.0000101001 = 1.00000101001 × 10."

Originally, with just the two numbers added, the result was "1.00000101 × 10," and dropping the last bit would trigger "round-to-even," giving "1.0000010 × 10." But now, with the extra tiny number added, the mantissa bits to be dropped when converting to BF16 become "1001," which is larger than the midpoint, so the round-up rule is triggered instead, giving "1.0000011 × 10." So from the perspective of the original two-number sum, the appearance of the third, tiny number has broken the "round-to-even" rule, and the positive bias reappears!

Of course, the conditions for this to occur seem quite stringent. First, the three numbers need the same sign; second, we need a "two large, one small" configuration, where the two large numbers happen to trigger a carry, and the small number is small enough to only affect FP32's mantissa (i.e., bits 9 through 23 of the mantissa). In this scenario, the small number is so tiny that dropping it alone would cause almost no error — but its mere presence happens to disrupt the "round-to-even" behavior of the two large numbers, thereby introducing a one-sided bias.

Tailor-Made

Given how stringent these conditions seem, can they really occur in practice? Under normal circumstances, this really isn't easy — but for Attention, it's as if this bug were "tailor-made" for it!

Let's take some row and column (i.e., some element) of $\bar{\boldsymbol{P}}\boldsymbol{V}$, which can be written as

\begin{equation}\sum_{i=1}^n \bar{p}_i v_i \label{eq:sum-pi-vi}\end{equation}

where $\bar{p}_i = \exp(s_i - \max(s_i))\leq 1$. As we know, one hallmark of Softmax Attention is its ability to "concentrate attention" — that is, attention may focus on just a handful of tokens. This shows up as: for a few tokens, $\bar{p}_i$'s corresponding $\bar{p}_i$ is close to 1, while the rest are very close to 0, though because of $\exp$, they can't be exactly 0 (unless they underflow to zero in BF16's representable range).

Then, as layers stack and training proceeds, the input $\boldsymbol{V}$ may exhibit "anisotropy," one manifestation of which is that the sign distribution along certain dimensions becomes uneven. Without loss of generality, suppose most of the entries in $v_i$ are positive (the negative case is analogous), and they're roughly of the same order of magnitude. Then the sum $\eqref{eq:sum-pi-vi}$ can be split into two parts: the main term, consisting of the handful of near-1 $\bar{p}_i$ values multiplied by $v_i$, and the remainder term, consisting of the vast majority of near-0 $\bar{p}_i$ values multiplied by $v_i$.

The paper considers a special case: the main term's corresponding few $\bar{p}_i$ values aren't merely close to 1 but exactly equal to 1 — that is, certain rows of $\boldsymbol{S}$ have multiple entries equal to $\max$ simultaneously. This special case is naturally harder to satisfy, but easier to understand: in this case, the main term's $\bar{p}_i v_i$ inherently has only BF16 precision. With this, "the timing and conditions are right," and the bug described in the previous section is triggered exactly as described:

Most terms are positive, the main terms all have BF16 precision, and their sum satisfies the carry condition; the remainder term is tiny, only affecting the very last mantissa bits of FP32, and just so happens to break "round-to-even," causing a bias; finally, because of "concentrated attention," the number of main terms isn't large, so there isn't too much carrying involved either (the more digits dropped, the smaller the bias), which keeps the bias within a noticeable range!

Put all of this together, and isn't it exactly a "custom-tailored bug" for Attention?

Killing the Remainder Term

Now that we understand the mechanism behind the problem, let's think about how to fix it. The simplest solution is to just store $\bar{\boldsymbol{P}}\boldsymbol{V}$ in FP32 rather than rounding it down to BF16 — though this does increase GPU memory usage somewhat. Let's instead try to discuss solutions that stay within BF16.

On the surface, the cause of the bias is that the tiny remainder term breaks "round-to-even," but thinking about it more deeply, the root cause is actually that "round half up" has a discontinuity right at the midpoint, and near that discontinuity, small perturbations easily introduce a bias. "Round-to-even" eliminates the bias, but it doesn't eliminate the discontinuity itself. The ideal fix is Stochastic Rounding, i.e., rounding up or down probabilistically, which avoids bias from small perturbations to the greatest extent possible.

However, apparently Stochastic Rounding is hard to implement efficiently at the hardware level, so most current hardware matrix-multiplication operators don't support it. As a result, the original paper takes a different route, confronting the issue head-on — an approach I'd call "killing the remainder term." Specifically, when a certain trigger condition is detected, we modify Attention's computation formula to

\begin{equation}\boldsymbol{O} = \frac{\exp(\boldsymbol{S})\boldsymbol{V}}{\exp(\boldsymbol{S})\boldsymbol{1}} = \frac{\exp(\boldsymbol{S} - \beta\boldsymbol{S}_{\max})\boldsymbol{V}}{\exp(\boldsymbol{S}- \beta\boldsymbol{S}_{\max})\boldsymbol{1}}\end{equation}

where $\beta > 1$. In this way, every term needs to be divided by an additional factor of $\exp((\beta-1)\boldsymbol{S}_{\max})$, which is a fairly substantial number (the paper sets $\beta \geq 2$), so the originally tiny remainder term is much more likely to underflow to zero and vanish — at which point "round-to-even" once again works properly, eliminating the bias.

So what's the detection condition? The original paper adopts a fairly simple one: the modification is triggered whenever a row of the matrix $\boldsymbol{S}$ has its maximum value achieved at least twice, meaning $\bar{p}_i$ has at least two entries equal to 1. But I suspect there's plenty of room for refinement here — this is left as a direction for future improvement. Also worth noting: since Flash Attention computes things block by block, this detection condition and the corresponding modification are also applied per block; see the code in the paper's appendix for details.

Further Thoughts

Overall, the paper offers a rather distinctive perspective for understanding phenomena like MaxLogit explosion. It can explain some things, but it doesn't cover the full picture, and it leaves a lot of open questions (and some points I'd quibble with) worth thinking about.

First, the paper's analysis of Attention bias relies on the anisotropy of $\boldsymbol{V}$, which might explain why MaxLogit explosion and similar anomalies only start appearing from the second Attention layer onward: the first Attention layer's input is the embedding, which is relatively less prone to anisotropy, whereas the input to the second and later Attention layers has already passed through a preceding Attention layer, which may inherently introduce anisotropy (see reference).

However, this doesn't explain why MaxLogit explosion only shows up in specific layers — for example, the paper's experiments show the issue only in layer 2, while K2's results show issues in layers 2 through 4. Similarly, this clearly can't explain why Muon is more prone to MaxLogit explosion than Adam (as observed in Moonlight and K2). So this is likely a combined result of architecture, optimizer, and low-precision factors together — looking at precision alone gives an incomplete picture.

There's also a question worth pondering deeply: causality. Another condition the paper identifies for Attention bias is that attention concentrates on a small number of tokens; intervening in the Attention computation under this condition successfully prevented subsequent anomalies. However, I observed a small model trained normally, and its attention wasn't as concentrated as one might expect — for instance, the average Top-1 probability was less than 0.2, and it took the cumulative probability of the Top-400 tokens to reach 0.9 (with a training length of 4096).

So is Attention bias truly the "cause" of training collapse, or its "effect"? In other words, when we observe "attention concentrating on a handful of tokens," might that actually indicate the model has already entered a collapse regime? If we only intervene at that point, might it already be "too late"? For example, even if certain metrics show that some anomaly has been prevented, is it possible that the model can no longer scale further? For now, these questions remain unanswered.

Summary

This post shared a paper analyzing bias in low-precision Attention computation, and along the way, gave me a chance to brush up on the fundamentals of low-precision arithmetic.