Why Low-Precision Transformer Training Fails: An Analysis on Flash Attention

ArXiv: 2510.04212

🎯 Pitch

Training GPT-2 with flash attention in BF16, a standard recipe, catastrophically explodes due to a hidden vicious cycle: when multiple pre-softmax scores share the exact maximum value, systemically biased rounding errors in the attention output align with structurally similar low-rank representations to inject a coherent, destructive gradient that balloons weight norms and derails training.

1. Executive Summary

This paper analyzes a notorious training failure — catastrophic loss explosion — that occurs when training GPT-2 models with flash attention in BF16 precision on the OpenWebText dataset. The authors provide the first mechanistic explanation, revealing that the failure arises from two intertwined phenomena: the emergence of similar low-rank representations within the attention mechanism (structurally similar (PK)[T]⊤X[T] matrices across tokens and training steps) and the compounding effect of biased rounding errors in BF16 arithmetic (systematically negative errors in the ¯PV product when multiple attention probabilities equal exactly 1). These factors create a vicious cycle where biased rounding errors act as coefficients for the low-rank representations, accumulating as a biased gradient update that increases weight spectral norms and ultimately derails training — establishing that low-precision flash attention fails only when repeated maxima in pre-softmax scores trigger biased rounding that exploits structural coherence in the gradient error direction. A minimal modification to the safe softmax that dynamically adjusts normalization when multiple identical maxima are detected stabilizes training for both GPT-2S and GPT-2M across AdamW and Muon optimizers, confirming the diagnosis.

2. Context and Motivation

The Core Problem: Low-Precision Flash Attention Training Is Unreliable

The specific gap this paper addresses is deceptively narrow but practically devastating: training GPT-2-size transformers with flash attention in BF16 precision fails catastrophically for reasons no one has mechanistically explained. The failure manifests as a sudden loss explosion after several thousand training steps (Figure 2, Section 3.1), turning a previously converging model into numerical garbage within a few iterations. This is not a hypothetical edge case — it is a concrete, reproducible bug report that has sat unresolved in open-source repositories for over two years (nanoGPT Issue 303, 2023; nanoGPT Issue 524, 2024; flash-attention Issue 337, 2024).

The authors frame this as a knowledge gap, not a bug fix. Existing stabilization techniques — QK normalization, QK-clipping, Gated Attention — are deployed empirically in production models (Kimi-Team, 2025; Qwen-Team, 2025) without a causal understanding of why they work. The field knows that certain interventions prevent loss spikes, but not what numerical process they interrupt. This paper's explicit goal is to fill that gap: provide a complete causal chain from BF16 addition to training collapse, enabling principled rather than ad-hoc solutions.

Why This Problem Matters: Practical and Theoretical Significance

The practical importance is straightforward: flash attention is indispensable for training long-context models, and BF16 is the default precision for large-scale training. Flash attention reduces the memory complexity of attention from $O(N^2)$ to $O(N)$ with respect to sequence length (Dao et al., 2022), making it possible to train transformers on contexts of 8K, 32K, or 128K tokens — the backbone of modern LLMs. BF16, with its FP32-matching dynamic range, eliminates the gradient underflow problems that plague FP16 training (Kalamkar et al., 2019; Wang & Kanwar, 2019) and has become the industry standard for pretraining models from GPT-3 to LLaMA to Chinchilla. The combination of these two technologies — flash attention for memory efficiency, BF16 for compute efficiency — is essentially the default training recipe for modern transformers. When this combination fails silently after thousands of steps, the cost is not just the wasted compute of a single run; it is the uncertainty that any long training run might be accumulating invisible numerical damage before an abrupt collapse.

Beyond the immediate failure case, the paper positions this as an instance of a broader problem: the gap between the deep learning community's empirical stabilization practices and its understanding of numerical failure modes. The authors note (Section 1, Section 5 Discussion) that similar instabilities are observed in production-scale models (Kimi-Team, 2025; Qwen-Team, 2025), often linked empirically to phenomena like large spectral norms of weights (Yang et al., 2023; Rybakov et al., 2024) and attention sinks (Xiao et al., 2023). These observations have spawned a patchwork of fixes — normalization layers, clipping thresholds, architectural modifications — but without a mechanistic explanation, each new failure case requires starting from scratch. The paper's analytical workflow (Section 5, Appendix F) — isolate the error source, identify accumulation mechanisms, trace to root arithmetic cause — is positioned as a diagnostic blueprint that generalizes beyond this specific case to other architectures, scales, and low-precision formats like FP8.

The theoretical significance centers on what the paper reveals about how low-precision arithmetic interacts with the structure of gradient updates. The key finding (Section 3.3.1) — that the gradient error is not random noise but a biased accumulation along structurally similar low-rank directions — challenges the implicit assumption that rounding errors in deep learning are benign, zero-mean perturbations. The paper demonstrates that when (a) the attention mechanism produces representations with consistent low-rank structure across tokens and training steps, and (b) the arithmetic introduces biased rounding errors, these errors compound rather than cancel. This is not an obvious property of either transformers or BF16 individually; it is an emergent failure mode that only manifests under their interaction. Understanding this interaction provides a principled lens for reasoning about why certain architectural choices (QK normalization, Gated Attention) work: they disrupt the structural coherence that transforms random rounding errors into systematic bias.

Prior Approaches and Where They Fall Short

Stabilization techniques in practice. The paper catalogs four categories of empirical interventions (Section 1, Section 5 Discussion):

1. QK normalization (Henry et al., 2020; Qwen-Team, 2025): normalizes query and key vectors before computing attention scores. Empirically prevents loss spikes. The paper's analysis (Section 5 Discussion) provides a mechanistic explanation: normalization disrupts the low-rank structure of (PK)[T]⊤X[T] matrices, preventing rounding errors from accumulating coherently.

2. QK-clipping (Kimi-Team, 2025): clips extreme values in query-key dot products. This reduces the likelihood of producing attention probabilities of exactly 1, which the paper identifies as the trigger for biased rounding. However, the effectiveness of clipping depends on choosing the right threshold — too aggressive and it degrades model quality; too permissive and it fails to prevent instability.

3. Gated Attention (Qiu et al., 2025; Qwen-Team, 2025): introduces non-linear gating mechanisms that modify the attention computation. The authors suggest this disrupts the structural similarity of error matrices, preventing systematic accumulation.

4. Reverting to higher precision: using FP32 for attention computations or for the entire model. This works but eliminates the efficiency gains — memory and speed — that motivated low-precision training in the first place.

Where these fall short. All four categories share a fundamental limitation: they are empirical patches applied without causal understanding. They were discovered through trial and error, validated by "it works on my training run," and deployed because the cost of not deploying them (a crashed training run) outweighs the cost of deploying them (minor architectural overhead). This leaves three problems:

• No principled guidance for hyperparameters. For QK-clipping, what threshold should be used? For QK normalization, should it be applied to all layers or only problematic ones? Without understanding the underlying mechanism, these choices are made by grid search or heuristics.

• No guarantee of transfer to new settings. A stabilization technique that works for GPT-2 at 125M parameters may or may not work for a 7B model in FP8. Without knowing why the technique works, one cannot predict whether it will generalize.

• No path to more fundamental solutions. If the root cause is unknown, improvements are limited to tuning existing heuristics. A mechanistic understanding enables targeted, minimal interventions (like the paper's dynamic softmax modification) that address the root cause directly rather than working around it.

Prior work on low-precision training stability. The paper situates itself within a broader literature on numerical stability in low-precision training (Appendix A). Key threads include:

• Gradient scaling (Micikevicius et al., 2017; Zhao et al., 2021): prevents underflow in FP16 by scaling the loss, which shifts gradient magnitudes into representable ranges. This addresses dynamic range issues (exponent) but not precision issues (significand bits). The failure case in this paper occurs in BF16, which has the same dynamic range as FP32 — underflow is not the problem.

• Per-tensor scaling for FP8 (Perez et al., 2023; Peng et al., 2023; Balanc¸a et al., 2024): dynamically adjusts scaling factors to keep values within FP8's narrow representable range. Again, this targets dynamic range rather than the rounding bias the paper identifies.

• Optimizer modifications (Molybog et al., 2023; Huang et al., 2025; Wortsman et al., 2023): detect and mitigate gradient spikes through momentum reset, spike-aware clipping, or hybrid optimizers. These treat the symptoms (large gradients) rather than the cause (biased accumulation in weight updates). The paper's contribution is orthogonal: it explains why those large gradients appear in the first place.

• Stochastic rounding (Ben Ali et al., 2024): instead of deterministic round-to-nearest, randomly rounds with probability proportional to the truncated bits. This eliminates systematic bias in rounding error but requires hardware support that is not universally available. The paper's dynamic softmax solution achieves a similar effect — preventing systematic bias — without changing the rounding mode.

The specific gap: no mechanistic explanation for flash attention failure. None of the prior work addresses the phenomenon this paper investigates: the catastrophic failure that occurs specifically when flash attention and BF16 are combined. The nanoGPT issues that the paper cites (Issue 303, 2023; Issue 524, 2024; Issue 554, 2024) document the failure but offer only empirical workarounds (use FP32, use standard attention). Lee et al. (2024) report that roughly 10% of GPT-2 pretraining runs diverge under pure BF16 versus 0% under TF32, but do not explain why. Golden et al. (2024) study flash attention stability more broadly, but their analysis focuses on general numerical properties rather than this specific failure mode.

The absence of explanation is not just an academic gap — it has practical consequences:

"The absence of a clear causal chain from numerical error to loss explosion has left the community reliant on ad-hoc patches rather than principled solutions, hindering progress in robust low-precision training." (Section 3.1)

How This Paper Positions Itself

The paper's positioning is distinctive: it is not proposing a new stabilization technique as its primary contribution. The dynamic softmax modification (Section 4) is presented as a validation experiment — proof that the diagnosed mechanism is correct — rather than as a recommended production fix. This is a subtle but important distinction. Most papers on training stability propose a technique and evaluate its effectiveness; this paper proposes an explanation and uses a minimal intervention to test it.

As a diagnostic methodology. The paper explicitly positions its analytical workflow — isolate error source → identify accumulation mechanism → trace to root arithmetic cause — as a generalizable framework (Section 5, Appendix F). This is not just a description of what the authors did; it is a prescription for how other researchers should approach similar failures in other settings (FP8 training, larger models, different architectures). The paper's value proposition includes both the specific explanation for this failure case and the methodology that produced it.

As a mechanistic explanation of empirical phenomena. The paper connects its findings to several well-known but poorly understood observations in transformer training:

• Attention sinks (Xiao et al., 2023): tokens that attract disproportionately high attention scores. The paper provides a numerical mechanism: by creating attention probabilities of exactly 1, attention sinks are precisely the conditions that trigger biased rounding in the ¯PV product. This transforms attention sinks from an architectural curiosity into a direct causal factor in training instability.

• Growth of weight spectral norms (Yang et al., 2023; Rybakov et al., 2024): the observation that unstable training runs exhibit anomalously large spectral norms in specific layers and heads. The paper explains this as the accumulation of biased low-rank gradient updates — each training step adds a small error in the same structural direction, causing the norm to grow monotonically.

• Effectiveness of QK normalization and Gated Attention: The paper's explanation — that these techniques work by disrupting structural similarity in error matrices — is a post-hoc mechanistic interpretation, not experimentally validated in this paper. But it provides a unified framework for understanding why such diverse architectural interventions converge on the same empirical outcome (stable training).

Relationship to the broader low-precision landscape. The paper explicitly acknowledges that industry practice is moving toward FP8 for compute-bound operations while retaining BF16 for memory-bound operations like attention (Liu et al., 2024; Qwen-Team, 2025). This raises the stakes of the analysis: if attention is the precision bottleneck in mixed-precision training, understanding its failure modes becomes increasingly important as other parts of the model move to even lower precision. The paper does not claim its specific mechanism (biased rounding in ¯PV) directly transfers to FP8 attention — FP8 has different rounding behavior and a narrower dynamic range — but the analytical approach is positioned as transferable.

What the paper is not. It is important to clarify the boundaries of the paper's claims. The paper does not:

• Claim that all low-precision training failures have the same root cause. The analysis is specific to this failure case (GPT-2 + BF16 + flash attention on OpenWebText).
• Claim that the dynamic softmax modification is the optimal or only solution. It is a minimal validation experiment. Other interventions (QK normalization, stochastic rounding) might address the same root cause through different mechanisms.
• Provide a theoretical guarantee that the modified flash attention prevents all future instabilities. The experiments in Figure 7 run to 600K steps (GPT-2S) and 100K steps (GPT-2M), which is substantial but not a proof of asymptotic stability.
• Analyze the interaction of this mechanism with distributed training (the experiments use DDP on 4 GPUs, but the analysis does not explore whether synchronization amplifies or mitigates the error).

Why This Failure Case Is a Principled Choice

The authors made several deliberate choices in selecting this failure case that strengthen the paper's contribution:

Reproducible and deterministic. By recording and reusing the exact sequence of data batches from an initial failing run (Section 3.1), the authors eliminate data randomness as a confounding factor. This is critical for causal analysis: if the failure depends on which data batches are seen, it becomes much harder to isolate the numerical mechanism. The deterministic reproduction means that every experiment in the paper processes identical data, making comparisons between configurations (e.g., high-precision vs. low-precision δ computation) directly causal rather than correlational.

Small enough to analyze, large enough to matter. GPT-2S (12 layers, 768 embedding dimension, 12 heads) is a model scale where per-head diagnostics are tractable — the authors can visualize individual (PK)[T]⊤X[T] matrices (Figure 4) and trace a specific BF16 addition (Section 3.3.2) — but the architecture is representative of the transformer family used in production. The extension to GPT-2M and the Llama-3.1-8B analysis (Appendix D) provide preliminary evidence that the identified mechanisms (low-rank structure in error matrices, multiple attention maxima) exist in larger and more modern models.

Isolated failure to a single layer and head. One of the paper's most striking findings is that the failure originates in a specific attention head (head 8 in layer 2, Figure 3) and that computing only that head's output in FP32 stabilizes the entire model. This extreme locality is methodologically valuable: it means the analysis can focus on a single 64-dimensional attention head rather than the entire 12-layer model, dramatically reducing the search space for root causes. It also suggests that the failure is threshold-based — most of the model operates benignly in BF16, but one head crosses a numerical boundary that triggers catastrophic propagation.

In summary, the paper addresses a concrete, long-standing, and practically significant failure at the intersection of two critical technologies (flash attention and BF16 training). It positions its contribution not as a new stabilization trick but as the first complete causal explanation of why the failure occurs, validated by a minimal intervention, and accompanied by a generalizable diagnostic methodology. This fills a gap that empirical fixes have papered over but never explained.

3. Technical Approach

3.1 Reader Orientation

This paper is fundamentally a mechanistic analysis — not a proposal for a new training algorithm or architecture, but a forensic investigation into why an existing, widely-used combination of technologies (GPT-2 + flash attention + BF16 precision) catastrophically fails. The "system" being analyzed is the standard training pipeline for a transformer language model using flash attention in mixed precision, and the core idea is that the failure is not a random numerical artifact but a deterministic consequence of how BF16 rounding errors interact with structurally coherent representations that emerge naturally in the attention mechanism. The paper solves the problem of explaining this failure by tracing a complete causal chain from individual BF16 addition operations, through biased accumulation in the attention output, to corrupted weight gradients that increase spectral norms and trigger loss explosion — and validates this explanation by showing that a minimal, targeted modification to safe softmax (dynamically adjusting normalization when multiple identical attention maxima are detected) prevents the failure entirely.

3.2 Big-Picture Architecture (Diagram in Words)

The analysis traces the failure mechanism through four connected stages, each building on the previous one:

1. Failure isolation experiments (Section 3.2): A series of targeted ablations that systematically narrow down the source of instability by selectively replacing low-precision operations with high-precision (FP32) counterparts. These experiments identify that the failure originates specifically in the computation of δ = rowsum(dO ◦ Olp) within the backward pass of flash attention in layer 2, head 8 of the GPT-2 model.

2. Gradient error decomposition (Section 3.3.1): Mathematical analysis of how errors in δ propagate into the gradient of the query projection matrix WQ. This reveals that the gradient error is a sum of rank-1 matrices (PK)[T]⊤X[T] weighted by the scalar error (δlp - δhp)[T], and that these rank-1 matrices are structurally similar across tokens and training steps — creating a pathway for errors to accumulate rather than cancel.

3. Rounding error root cause analysis (Section 3.3.2): Investigation of why the scalar coefficients (δlp - δhp)[T] are systematically biased positive. This traces the bias to the ¯PV product in the forward pass, where attention probabilities of exactly 1 (occurring at repeated row-maxima in pre-softmax scores) combine with predominantly negative values in V to trigger a specific BF16 rounding behavior: significand overflow requiring a right shift and round-up, introducing systematically negative error in ¯O that propagates into positive error in δ.

4. Validation via minimal intervention (Section 4): A modified safe softmax that dynamically adjusts the normalization factor when multiple identical maxima are detected, ensuring that ¯P contains only values strictly less than 1. This prevents the biased rounding condition and stabilizes training across model sizes and optimizers, confirming the causal mechanism.

3.3 Roadmap for the Deep Dive

• First, the experimental setup and failure reproduction (Section 3.4.1): I explain the exact model, data, optimizer, and precision configuration used to reproduce the failure deterministically — this is the "crime scene" that constrains all subsequent analysis.
• Second, the isolation experiments (Section 3.4.2): I walk through the sequence of targeted ablations that narrow the failure to a specific computation (δlp) in a specific attention head (head 8, layer 2), establishing the minimal scope for the root cause analysis.
• Third, the gradient error decomposition (Section 3.4.3): I derive how errors in δ propagate into dWQ, explain why the low-rank structure of (PK)[T]⊤X[T] matters, and show evidence that these matrices are similar across tokens and steps — this is the "amplification mechanism" that transforms small rounding errors into systematic weight corruption.
• Fourth, the rounding error root cause (Section 3.4.4): I trace the positive bias in (δlp - δhp)[T] back through Olp - Ohp to the ¯PV product, analyze the specific BF16 addition that introduces negative error when ¯P[T, t] = 1 and V[t, i] < 0, and explain the bit-level mechanism of significand overflow, right shift, and round-up — this is the "trigger" that initiates the failure cascade.
• Fifth, the validation modification (Section 3.4.5): I describe the dynamic softmax adjustment, explain why it must be conditional and dynamic rather than a fixed offset, and show how it prevents ¯P elements from reaching exactly 1.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a mechanistic analysis paper whose core idea is that low-precision flash attention training fails because (1) the attention mechanism produces structurally similar low-rank representations that create a coherent error direction in weight gradients, and (2) BF16 rounding introduces a systematic bias when multiple attention probabilities equal exactly 1, providing the scalar coefficients that drive error accumulation along that direction. The interaction of these two factors — neither sufficient alone — creates a vicious cycle that increases weight spectral norms and causes loss explosion.

3.4.1 Experimental Setup and Deterministic Failure Reproduction

The investigation begins with a precisely specified failure case that the authors make fully reproducible. The model is a GPT-2 architecture with 12 layers, 12 attention heads per layer, an embedding dimension of 768, and a context length of 1024 tokens. This is the standard "GPT-2 Small" configuration, placing it at roughly 125M parameters — large enough to exhibit the failure but small enough for detailed per-head analysis. The model is pre-trained on the OpenWebText dataset (Gokaslan et al., 2019), an open-source recreation of the WebText corpus used in the original GPT-2 paper (Radford et al., 2019).

Training configuration. The authors use the AdamW optimizer with β1 = 0.9, β2 = 0.95, and zero weight decay. The learning rate follows a cosine schedule with a 2000-step linear warmup to a peak of 1 × 10^{-3}, decaying to 1 × 10^{-5}. Global gradient clipping is applied with a maximum norm of 1.0. Training runs on 4 NVIDIA A100 (80GB) GPUs using PyTorch's Distributed Data Parallel (DDP) module. Each GPU processes a micro-batch of 32 sequences, and gradient accumulation over 4 steps yields an effective global batch size of 524,288 tokens per optimization step. The precision strategy is automatic mixed precision: the forward pass (including flash attention) uses BF16, while the backward pass uses FP32 for gradient accumulation. This is the standard mixed-precision recipe: weights are stored in FP32, casts to BF16 for forward computations, and gradients are accumulated in FP32 before being applied to the FP32 master weights.

Deterministic reproduction. A critical methodological choice is that the authors record and reuse the exact sequence of data batches from an initial run that led to the failure. This means that every experiment in the paper processes identical data in identical order, eliminating data randomness as a confounding variable. If the failure were caused by a particular rare data pattern, it would still be reproducible, but the authors could not distinguish data effects from numerical effects. By fixing the data stream, any change in training behavior between configurations (e.g., computing δ with Olp versus Ohp) is guaranteed to be caused by the numerical change, not by seeing different data. The failure manifests as a sudden loss explosion after several thousand training steps, as shown in Figure 2: the high-precision (FP32 flash attention) configuration converges smoothly, while the low-precision (BF16) configuration diverges catastrophically around step 6,600–6,700. This is the same failure pattern reported in community bug reports (nanoGPT Issue 303, 2023; nanoGPT Issue 524, 2024).

Why this configuration matters. The combination of BF16 forward pass with FP32 backward pass is deliberate and representative of standard practice. In pure BF16 training, both forward and backward passes use BF16, which can cause gradient underflow even with loss scaling. The mixed-precision approach — FP32 master weights, BF16 forward, FP32 backward — is the recommended PyTorch recipe and is used in most production training pipelines. The fact that the failure occurs in this standard configuration, not in an exotic or aggressive low-precision setup, underscores the practical significance: this is a bug that affects default training behavior, not an edge case from pushing precision too far.

3.4.2 Isolation Experiments: Narrowing Down the Failure Source

The authors conduct a systematic sequence of ablation experiments to pinpoint which specific operation in flash attention is responsible for the failure. Each experiment modifies one aspect of the computation — replacing it with a high-precision or alternative implementation — and checks whether training stabilizes. The logic is: if changing operation X fixes the failure, then X (or something downstream of X) is part of the causal chain.

Tiling is not the source of failure. The first hypothesis tested is whether the block-wise processing (tiling) in flash attention — which partitions Q, K, V into smaller blocks processed iteratively — introduces errors that accumulate across blocks. To test this, the authors set the block size equal to the sequence length (1024), effectively disabling tiling so that flash attention computes with full matrices in one pass. Training still fails identically. This rules out tiling as a causal factor and, importantly, simplifies all subsequent analysis: the authors can use a non-tiled version of flash attention, which eliminates the complexity of block-wise interactions and focuses attention on the core arithmetic operations (softmax, matrix multiply, normalization).

Failure originates in a single layer. The authors monitor the spectral norms of weight matrices across all layers during training (Figure 11 in Appendix). The spectral norm — the largest singular value of a matrix — is a standard diagnostic for training instability: anomalously large spectral norms indicate that weight updates are pushing the model into a regime where activations grow unboundedly (Yang et al., 2023; Rybakov et al., 2024). The monitoring reveals an anomalous spike specifically in layer 2's attention weights. Two targeted experiments confirm this localization:

• Using flash attention only in layer 2 (with standard attention in all other layers) reproduces the training failure.
• Replacing flash attention with standard attention in layer 2 (while retaining flash attention in all other layers) restores training stability.

These results are striking: the failure is not a diffuse numerical degradation across the model but a localized pathology in a single transformer layer. This dramatically narrows the analysis scope — the authors need to understand what goes wrong in layer 2's attention, not in the entire 12-layer stack.

Failure is linked to the computation of δ. In the backward pass of flash attention (Algorithm 2), an intermediate term δ = rowsum(dO ◦ O) is computed for efficiency. This term appears in the gradient of the attention scores: dS = P ◦ (dP - δ), which then propagates into gradients for Q, K, and V. There is a mathematically equivalent alternative formulation: δ = rowsum(dP ◦ P), where dP = dO V^⊤. The key difference is that the standard formulation uses O (the attention output from the forward pass), while the alternative uses P (the attention probabilities) and dP. The authors find that replacing δ = rowsum(dO ◦ O) with δ = rowsum(dP ◦ P) restores training stability. This is a critical finding because the two are mathematically identical in exact arithmetic — the only difference is that O was computed in BF16 during the forward pass. This experiment demonstrates that numerical errors in the low-precision O are the source of instability, since using an alternative formulation that avoids Olp eliminates the failure.

Numerical errors in O are the source of failure. Building on the δ finding, the authors isolate the error to the low-precision output matrix Olp with two additional experiments:

• During the backward pass, instead of using Olp from the forward pass to compute δ, the authors recompute O as PV in FP32 within the backward pass. This change stabilizes training.
• During the forward pass, the authors compute O in high precision (FP32) — that is, Ohp instead of Olp — while keeping all other forward operations in BF16. The backward pass then uses δhp = rowsum(dO ◦ Ohp). This also stabilizes training.

These experiments establish the claim that the paper returns to throughout the analysis:

Claim 1. Low-precision δlp = rowsum(dO ◦ Olp) causes training failure.

The phrase "low-precision δlp" specifically means the δ vector computed using the BF16 output Olp from the forward pass, as opposed to the high-precision δhp computed using Ohp (which would be the result if O were computed in FP32). The error that matters is δlp - δhp, the difference between what the backward pass actually uses and what it should use in exact arithmetic.

Failure is localized to specific attention heads. Layer 2 has 12 attention heads. To further narrow the analysis, the authors track the spectral norm of each head's query projection matrix WQ. Figure 3 shows that heads 1, 7, 8, 9, 11, and 12 have elevated spectral norms, with head 8 showing the largest (6.27, compared to values around 2–4 for the other heads). A targeted experiment confirms the causal role: selectively computing O in high precision for these six outlier heads (while keeping all other heads in BF16) is sufficient to restore training stability. Since head 8 has the largest spectral norm, the authors focus all subsequent analysis on layer 2, head 8. This is an extreme form of dimensionality reduction — instead of analyzing a 12-layer × 12-head = 144-dimensional space of possible failure points, the investigation narrows to a single 64-dimensional attention head (embedding dimension 768 ÷ 12 heads = 64 dimensions per head).

Why this isolation strategy is methodologically elegant. The sequence of experiments follows a principle of minimum sufficient intervention: at each step, the authors change the smallest possible component of the system to see if stability is restored. If changing X fixes the failure, the root cause must involve X. By starting with broad changes (disable tiling) and progressively narrowing (single layer → single computation → single head), the authors converge on the minimal causal unit — δlp for head 8 of layer 2 — without assuming which part of the system is responsible. This bottom-up approach is what distinguishes the analysis from prior empirical work that observed instability but could not locate its origin.

3.4.3 Cause 1: Similar Low-Rank Matrices Bias Weight Updates

With the failure isolated to δlp in a specific attention head, the authors now analyze how errors in δ corrupt the training dynamics. The key insight is that the gradient error is not random — it has a consistent low-rank structure across tokens and training steps, which causes errors to accumulate rather than cancel out.

Derivation of the gradient error. The analysis begins by expressing how errors in δ propagate into the gradient of the query projection matrix WQ. The gradient flow in attention is: the loss gradient with respect to attention scores is dS = α P ◦ (dP - δ), where α is the attention scaling factor (typically 1/√d), P is the attention probability matrix, dP = dO V^⊤ is the gradient of the loss with respect to P, and δ controls the row-wise normalization. The gradient with respect to the query matrix Q is then dQ = dS K. The gradient with respect to WQ is the outer product of the input features X and dQ: dWQ = dQ^⊤ X.

The difference between high-precision and low-precision gradients therefore depends on the difference between δhp and δlp. The authors derive this step by step:

First, the difference in score gradients:

$dQ_{hp} - dQ_{lp} = \alpha \cdot \text{diag}(\delta_{lp} - \delta_{hp})(PK)$

where diag(δlp - δhp) is a diagonal matrix whose diagonal entries are the elements of the vector difference δlp - δhp, and PK is the matrix product of the attention probabilities and the key matrix. The diag(v) multiplication means that each row T of PK is scaled by the scalar (δlp - δhp)[T].

What this equation means operationally: take the high-precision query gradient and subtract the low-precision query gradient. The difference equals α times the matrix PK with each row multiplied by the corresponding error in δ. If (δlp - δhp)[T] is positive for some token position T, then row T of PK is over-represented in the low-precision gradient (i.e., dQlp is pushed in a direction that dQhp would not take).

Next, expanding to the weight gradient:

$dWQ_{hp} - dWQ_{lp} = \alpha \sum_{T=1}^{N} (\delta_{lp} - \delta_{hp})[T] \cdot (PK)[T]^{\top} X[T]$

where N is the sequence length, (PK)[T] is the T-th row of the PK matrix (a vector of length d_head = 64), X[T] is the T-th row of the input features (a vector of length d_model = 768), and (PK)[T]⊤ X[T] is a rank-1 matrix of size d_head × d_model.

What this equation computes: the total error in the weight gradient dWQ is a weighted sum of N rank-1 matrices, where the weight for token position T is the scalar error (δlp - δhp)[T], and the rank-1 matrix is the outer product of that token's (PK) row and input feature row. Each term (PK)[T]⊤ X[T] is a matrix whose rank is at most 1 — it can be written as the outer product of two vectors, meaning it has only one non-zero singular value and represents a single direction in the weight space.

Why this form matters: this decomposition separates the gradient error into two factors: (a) scalar coefficients (δlp - δhp)[T] that depend on numerical precision, and (b) rank-1 direction matrices (PK)[T]⊤ X[T] that depend on the model's representations. If the scalar coefficients are zero-mean and independent across tokens and training steps, the sum across many tokens will tend to cancel — rounding errors in opposite directions will average out. But if the scalar coefficients are systematically biased (consistently positive or negative) AND the direction matrices are structurally similar (pointing in roughly the same direction in weight space), then the errors accumulate rather than cancel. Each training step adds a small perturbation in the same direction, and over hundreds of steps, this coherent accumulation can substantially modify the weight matrix.

Evidence for similar low-rank structure. The authors present Figure 4 as evidence that (PK)[T]⊤ X[T] matrices are structurally similar across tokens and training steps. Panels (a) and (d) show the PK matrices for head 8 at two different training steps (6610 and 6619) and two different batch indices (190 and 209). Visual inspection reveals similar column patterns — certain features (columns of PK, corresponding to head dimensions) have consistently large or small values across tokens. Panels (b) and (e) show the input features X, which also exhibit structural consistency. Panels (c) and (f) show the rank-1 matrices (PK)[T]⊤ X[T] for specific tokens (token 50 at step 6610, token 718 at step 6619), with similar columns highlighted at input feature indices 546 and 678.

The critical claim is not that the matrices are identical — they are not — but that they share similar low-rank structure, meaning that their dominant singular vectors point in similar directions. This means that when the scalar coefficients (δlp - δhp)[T] are consistently positive (as shown in Figure 5a), the low-precision gradient dWQlp receives repeated nudges in roughly the same direction in weight space, while the high-precision gradient dWQhp (which uses accurate δhp) does not.

The authors formalize this by approximating the sum:

$dWQ_{hp} - dWQ_{lp} \approx \alpha \left(\sum_{T=1}^{N} (\delta_{lp} - \delta_{hp})[T]\right) \cdot R$

where R denotes the common low-rank structure emerging across different tokens and training steps. This approximation says: instead of adding up N slightly different rank-1 matrices, treat them as effectively the same direction R, and the total error is proportional to the sum of the scalar coefficients.

Evidence for positive bias in scalar coefficients. Figure 5a tracks the cumulative sum of (δlp - δhp)[T] over a sequence of training steps from 6580 to 6680 — the period leading up to the loss explosion. The cumulative sum is consistently positive and increasing, indicating that across many tokens and training steps, δlp systematically exceeds δhp. A zero-mean random error would produce a cumulative sum that fluctuates around zero; a consistently increasing sum means the errors are biased in one direction. This bias is what prevents cancellation: each training step adds a positive multiple of R to the weight update, and these accumulate over hundreds of steps.

Consequences of biased accumulation. The accumulation of biased updates in direction R has two observable effects that the paper documents (and that match empirical observations in the broader literature):

1. Growth of weight spectral norms (Yang et al., 2023; Rybakov et al., 2024): The spectral norm of WQ for head 8 increases disproportionately compared to other heads (Figure 3, Figure 11), because the accumulated error α · Σ(δlp-δhp)[T] · R increases the magnitude of the weight matrix along the direction R. Since R is low-rank, this growth is concentrated — it doesn't uniformly scale all singular values, but specifically amplifies the direction corresponding to R.

2. Growth of activations: As the weight spectral norm increases, the forward pass produces increasingly large activations (the query vectors Q = X WQ grow in magnitude). This is a positive feedback loop: larger activations push the softmax into a sharper regime (more extreme attention probabilities), which increases the likelihood of attention probabilities of exactly 1, which triggers more biased rounding, which further biases weight updates.

Claim 2. Weight update in low-precision training is biased by (δlp - δhp)[T]R, which arises from the structurally similar matrices (denoted as R) across tokens and training steps, and its positively-biased coefficient (δlp - δhp)[T]. This bias accumulates error, preventing cancellation, increasing weight spectral norm and activation, and leading to loss explosion.

Why this finding challenges implicit assumptions. The standard assumption in mixed-precision training is that rounding errors are approximately zero-mean and independent across operations, so they average out over large batches and many training steps (Micikevicius et al., 2017). This paper demonstrates a counterexample where this assumption fails: the rounding errors are not zero-mean (they have a positive bias in δ) and the gradient error directions are not independent (they share low-rank structure). The interaction of these two violations — bias in the coefficients, coherence in the directions — creates a failure mode that neither factor alone would produce. Random zero-mean errors along a coherent direction would cancel (positive and negative alternate). Biased errors along random directions would be harmless (each step updates a different direction, so no single weight grows abnormally). But biased errors along a coherent direction compound systematically.

3.4.4 Cause 2: Biased Rounding Error Leads to Positive (δlp - δhp)[T]

Having established that the failure requires both structural similarity in gradient errors AND a positive bias in (δlp - δhp)[T], the authors now trace the origin of that bias to its root cause: a specific BF16 rounding behavior in the forward pass computation of ¯O = ¯P V.

Locating the error in Olp - Ohp. Recall that δ = rowsum(dO ◦ O). The error in δ for token position T is:

$\delta_{lp}[T] - \delta_{hp}[T] = \sum_{i} dO[T, i] \cdot (O_{lp}[T, i] - O_{hp}[T, i])$

where the sum runs over feature dimensions i (0 to 63 for a single head). This decomposes the δ error into contributions from each feature dimension, weighted by the upstream gradient dO[T, i]. The authors focus on token position T = 718 where the error is positive and large.

Figure 5(b) and (c) reveal a strong sign correlation for specific feature dimensions (20 and 29): both the gradient dO[T, i] and the output error O_{lp}[T, i] - O_{hp}[T, i] are consistently negative. Since the product of two negative numbers is positive, these dimensions contribute positively to the δ error. The fact that the output error is systematically negative (Olp[T, i] < Ohp[T, i]) means the low-precision O is biased toward more negative values than the high-precision O. The rest of the analysis traces where this negative bias in O comes from.

Isolating the error to ¯PV. The attention output O is computed from an intermediate unnormalized output ¯O via a safe softmax and normalization:

$\bar{P} = \exp(S - \text{rowmax}(S))$ $\bar{O} = \bar{P} V$ $O = \bar{O} / \text{rowsum}(\bar{P})$

Further experiments show that computing only ¯O = ¯P V in FP32 is sufficient to stabilize training. This narrows the root cause to the matrix multiplication ¯PV — the normalization step (O = ¯O / rowsum(¯P)) is not the source of the bias. The analysis therefore focuses on the element-wise computation:

$\bar{O}_{lp}[T, i] - \bar{O}_{hp}[T, i] = (\bar{P}_{lp}[T, :] V[:, i])_{lp} - (\bar{P}_{hp}[T, :] V[:, i])_{hp}$

where the subscript (·)lp on the parentheses means "compute the dot product of ¯P[T, :] and V[:, i] in FP32, then round the final result to BF16," while (·)hp means "compute entirely in FP32 without rounding."

How the error accumulates across token positions. To see when during the summation the error arises, the authors plot the cumulative error as the sum over token positions progresses:

$\bar{O}_{\text{error}}(t) = \left(\sum_{t'=1}^{t} \bar{P}[T, t'] V[t', i]\right)_{lp} - \left(\sum_{t'=1}^{t} \bar{P}[T, t'] V[t', i]\right)_{hp}$

Figure 6(b) and (c) show this cumulative error for feature dimension i = 20. The error accumulates in significant negative steps — sharp downward jumps in the cumulative error curve — at specific token positions t. Figure 6(c), which provides a zoomed-in view, reveals that these negative jumps occur at token positions where the corresponding attention probability ¯P[T, t] is exactly 1.

Why ¯P[T, t] = 1 is the critical condition. The attention probability ¯P is computed from the pre-softmax scores S as ¯P = exp(S - rowmax(S)). If S[T, t] is the maximum value in row T, then S[T, t] - rowmax(S[T, :]) = 0, and exp(0) = 1. If there are multiple identical maxima in the same row, then multiple positions will have ¯P[T, t] = 1. This is exactly the condition that occurs repeatedly in the failing head.

When ¯P[T, t] = 1, the product ¯P[T, t] V[t, i] simplifies to just V[t, i] — the attention probability doesn't scale the value, it just selects it. Figure 6(a) shows that for feature dimension i = 20, the values V[:, 20] are predominantly negative. So the accumulation Σ ¯P[T, t'] V[t', i] involves adding together many negative BF16 numbers with values coming directly from V (when ¯P[T, t] = 1).

The BF16 addition mechanism that produces negative bias. The core numerical mechanism is the rounding behavior when adding two negative BF16 numbers with the same exponent. In floating-point arithmetic, adding two numbers with the same sign can cause the significand (the fractional part) to overflow the available precision. The BF16 format has 1 sign bit, 8 exponent bits, and 7 fraction (significand) bits, with an implicit leading 1 (so the significand is effectively 8 bits: 1.xxxxxxx). When two BF16 numbers are added:

1. Their exponents are aligned (the smaller number's significand is right-shifted until exponents match).
2. The significands are added.
3. If the sum overflows (e.g., 1.xxxxxxx + 1.yyyyyyy = 1z.zzzzzzz — a leading bit beyond the implicit 1), the significand is right-shifted by one bit and the exponent is incremented to re-normalize.
4. The right-shifted significand must be rounded to fit the 7 fraction bits. The standard rounding mode is "round to nearest, ties to even."

The rounding bias arises because the summation is performed in FP32 (which has 23 fraction bits), so the accumulator holds a much more precise intermediate sum than BF16 can represent. When a new BF16 value (V[t, i], negative) is added to this FP32 accumulator, the lower-order bits already present in the accumulator (from previous small additions) can activate the sticky bit — a flag indicating that non-zero bits were shifted past the rounding boundary during exponent alignment. The sticky bit forces a round-up when the final result is converted back to BF16, because "round to nearest" with a sticky bit treats the exact value as slightly above the halfway point.

Concrete bit-level example. The authors provide a worked example of the specific BF16 addition that causes the large negative error jump in Figure 6(c). The FP32 representations are:

11000000000110100000111000101110  (-2.4071154594421387)
11000000000100110000000000000000  (-2.296875)

Both numbers have the same exponent (the 8-bit exponent field is 10000000, which with bias 127 gives exponent 1, so the values are in the range [-4, -2)). The addition of their significands:

(-1.00110100000111000101110) + (-1.00100110000000000000000) = -10.01011010000111000101110

The result -10.0101101... has overflowed — the leading 10 means the significand needs to be shifted right by one bit, and the exponent incremented:

Significand: 10.01011010000111000101110 → 1.001011010000111000101110 (after right shift)
Exponent: 10000000 → 10000001

The exact FP32 result is 1 10000001 00101101000011100010111, which represents -4.703990459442139. To convert this to BF16, only 7 fraction bits are kept: the fraction is 0010110 (7 bits), and the next bit (the rounding bit) is 1. Because the rounding bit is 1 and there are additional non-zero bits after it (00011100010111 — these are the lower-order bits that were shifted out, and their presence activates the sticky bit logic), the round-to-nearest rule rounds up, adding 1 to the last fraction bit:

BF16 fraction: 0010110 + 1 = 0010111

The final BF16 result is 1100000010010111, which represents -4.71875. The true sum was -4.703990459442139, so the rounding introduced an error of -0.014759540557861328 — the BF16 result is more negative than the exact sum by about 0.015.

Why this creates systematic negative bias. The systematic nature of the bias comes from the interaction of several factors:

• ¯P[T, t] = 1 occurs repeatedly when there are multiple identical row maxima. Each occurrence adds V[t, i] (a negative BF16 number) to the accumulator.
• V[:, i] is predominantly negative for the problematic feature dimensions (Figure 6a). So the accumulator is summing negative numbers.
• Each addition of two negative numbers with similar magnitude risks significand overflow, which triggers the right-shift-and-round-up procedure analyzed above.
• The round-up makes the BF16 result more negative than the exact sum (rounding -4.704 to -4.719 is a negative error). Since the numbers being added are negative, "rounding up" in the significand means "rounding toward more negative" — the magnitude increases.
• The positive error from rounding down is smaller because when the rounding bit is 0, the result is rounded down (toward zero, i.e., less negative), but the magnitude of this positive error is typically smaller than the negative error from rounding up because the values being added are usually small.

This asymmetry — negative rounding errors dominate positive ones — creates the systematic negative bias in ¯Olp relative to ¯Ohp. That negative bias propagates into Olp - Ohp (the low-precision output is more negative), which (when multiplied by negative dO values) produces positive contributions to δlp - δhp, which accumulate as biased weight updates in direction R.

Claim 3. With more than one ¯P[T, t] = 1 and negative V[t, i], the addition in ¯PV can cause overflow of significand. This necessitates a right shift and a round-up of significand due to sticky bit, introducing negative rounding error in O and leading to positive (δlp - δhp)[T].

Why this mechanism is specific to the combination of flash attention and BF16. Several factors must align for this failure to occur:

• Multiple identical row maxima in pre-softmax scores. This produces ¯P[T, t] = 1 at multiple token positions. Standard training with FP32 precision also has this, but the rounding error is negligible at 23 fraction bits. In BF16 with 7 fraction bits, the rounding error is large enough to matter.
• Predominantly same-sign values in V. If V[:, i] had a mix of positive and negative values, the rounding errors from positive and negative additions might partially cancel. The consistent negativity of V[:, 20] (Figure 6a) means all the rounding errors push in the same direction.
• The flash attention forward pass uses BF16 for the ¯PV product. If this product were computed in FP32 (as the isolation experiment showed stabilizes training), the rounding errors would be at FP32 precision and would be ~6 orders of magnitude smaller — too small to accumulate to a meaningful bias over the number of training steps before the failure.

This is the paper's complete causal chain: repeated attention maxima → ¯P values of exactly 1 → biased BF16 rounding in ¯PV when V values share a sign → systematically negative error in ¯O → propagates to negative error in Olp → multiplied by negative dO values yields positive δlp - δhp → these positive scalars weight structurally similar rank-1 matrices (PK)[T]⊤X[T] → biased weight updates accumulate in direction R → weight spectral norms increase → activations grow → more extreme softmax → more repeated maxima (positive feedback) → loss explosion.

3.4.5 Validation: Mitigating Bias via Dynamic Softmax Adjustment

The final component of the paper is not a proposed production fix but a validation experiment: a minimal modification to flash attention that, according to the diagnosed mechanism, should prevent the failure. If the modification works, it confirms the causal analysis; if it doesn't, the analysis is wrong or incomplete.

The intervention strategy. The root cause is ¯P[T, t] = 1 occurring at multiple token positions due to repeated identical maxima in the pre-softmax scores. The intervention must prevent ¯P from containing any value of exactly 1 while preserving the mathematical correctness of attention in exact arithmetic. The softmax function has a shift-invariance property: softmax(z) = softmax(z - c) for any constant c. Standard flash attention exploits this by subtracting the row maximum rm = rowmax(S) before computing the exponential, ensuring numerical stability (the largest exponent is 0). The authors' modification exploits the same property differently: when multiple identical maxima are detected, they subtract a slightly larger value so that S - m at the maximum positions is strictly negative, making ¯P = exp(S - m) < 1 at all positions.

The modified safe softmax. The modification is applied within the tiled forward pass of flash attention (Algorithm 3, magenta-highlighted lines):

rm = rowmax(S), rs = rowsum(rm - S ≤ ε)
m' = where(rm > 0 ∧ rs > 1, β·rm, rm)
m = where(rm < 0 ∧ rs > 1, 0, m')
¯P = exp(S - m)

where ε is a small numerical tolerance (e.g., 10^{-3}) to account for floating-point imprecision in detecting identical values, and β > 1 (the authors use β ∈ [2, 8]).

What each line does:

1. rm = rowmax(S): computes the row-wise maximum of the pre-softmax scores, as in standard safe softmax.

2. rs = rowsum(rm - S ≤ ε): counts how many entries in each row are within ε of the maximum. If rs > 1, there are multiple (effectively) identical maxima in that row.

3. m' = where(rm > 0 ∧ rs > 1, β·rm, rm): for rows where the maximum is positive AND there are multiple maxima, the normalization factor m is set to β·rm instead of rm. This means the maximum value in the exponent becomes rm - β·rm = -(β-1)·rm, which is strictly negative (since rm > 0 and β > 1).

4. m = where(rm < 0 ∧ rs > 1, 0, m'): for rows where the maximum is negative AND there are multiple maxima, m is set to 0. The maximum exponent becomes rm - 0 = rm, which is strictly negative (since rm < 0). If rm is negative and unique (rs = 1), m stays at rm and the behavior is identical to standard softmax.

5. ¯P = exp(S - m): the attention probabilities are computed as before, but now with the guarantee that max(S - m) < 0 whenever multiple maxima were detected, so max(¯P) < 1.

Why this must be dynamic and conditional, not a fixed offset. Appendix C explains why simpler alternatives fail:

• A fixed offset (always subtracting a small constant from the row maximum) would cause ¯P values to be consistently rounded in one direction during BF16 conversion, introducing a new systematic rounding error. The dynamic adjustment is applied only when the specific condition (multiple maxima) that triggers the original biased rounding is detected.

• An unconditional adjustment (always using β·rm when rm > 0) would cause underflow when rm is large. If rm is a single, very large positive value, then exp(-(β-1)·rm) could underflow to zero, making the normalization factor zero and causing a division-by-zero error when computing O. The conditional application preserves numerical stability in all cases where the original softmax is stable.

• The choice of β ∈ [2, 8]: smaller values (e.g., β = 1.001) risk having exp(-(β-1)·rm) round back to exactly 1 in BF16 (if (β-1)·rm is very small), defeating the purpose. Larger values (e.g., β = 100) risk underflow for moderate rm. The range [2, 8] balances these concerns: the exponent is sufficiently negative that exp(-(β-1)·rm) is reliably less than 1 even after BF16 rounding, but not so negative that it underflows to zero.

Why the rm < 0 case sets m = 0. When the row maximum is negative and repeated, setting m = γ·rm with γ ∈ (0, 1) would make the maximum exponent (1-γ)·rm. But if γ is close to 1, (1-γ)·rm is close to zero (since rm is negative), and exp((1-γ)·rm) could round to exactly 1 in BF16 — reintroducing the very condition we're trying to avoid. Setting m = 0 guarantees that the maximum exponent is rm (which is negative), so exp(rm) < 1. Appendix C confirms that this is a robust choice.

Mathematical equivalence in exact arithmetic. The modification is mathematically equivalent to standard attention in exact arithmetic because it uses the shift-invariance property of softmax: softmax(z) = softmax(z - c). Standard flash attention uses c = rm; the modified version uses c = β·rm or c = 0 when multiple maxima are detected. In exact arithmetic, both produce identical attention outputs because softmax is invariant to constant row-wise shifts. The only difference is in finite-precision arithmetic, where the modified version prevents ¯P from taking the value 1, thereby avoiding the biased rounding path.

Validation results (Section 4, Figure 7). The modified flash attention is evaluated on:

• GPT-2S with AdamW (Figure 7a): training for 600K steps. The original flash attention diverges (loss explodes), while the stabilized version converges smoothly, matching the behavior of high-precision flash attention from Figure 2.
• GPT-2S with Muon optimizer (Figure 7b): Muon (Jordan et al., 2024) is a recently proposed optimizer that uses matrix orthogonalization for hidden layer weights. The stabilized flash attention also prevents loss explosion with Muon, demonstrating that the fix is not optimizer-specific.
• GPT-2M with AdamW (Figure 7c): a larger model variant trained for 100K steps. The stabilized version prevents failure at this larger scale, providing preliminary evidence of generalization beyond GPT-2S.

What this validation proves and what it doesn't. The experiments demonstrate that the specific mechanism diagnosed — ¯P values of exactly 1 triggering biased BF16 rounding — is sufficient to cause the failure, because preventing that mechanism (while changing nothing else) restores stability. However, the paper does not claim that this modification is the optimal or only solution, nor that it addresses all possible low-precision instabilities. It is a minimal validation of the causal analysis, not a recommended production fix. The authors explicitly position QK normalization and Gated Attention as alternative interventions that address the same root cause through different mechanisms (Section 5 Discussion): they disrupt the low-rank structure R rather than preventing the biased rounding trigger, but the end result — preventing coherent accumulation — is the same.

Summary of Design Choices and Their Justifications

• Deterministic data replay over random data loading: eliminates data randomness as a confounding variable, enabling causal attribution of stability changes to numerical modifications.
• Isolation by minimum sufficient intervention over analyzing the full model: progressively replacing operations with FP32 equivalents narrows the failure to a single computation (δlp) in a single head, making root cause analysis tractable.
• Gradient error decomposition into scalar coefficients and rank-1 directions over treating the error as an opaque tensor: separates the "how much" (bias in δlp - δhp) from the "in what direction" (structure of (PK)[T]⊤X[T]), enabling diagnosis of why errors accumulate rather than cancel.
• Bit-level analysis of BF16 addition over statistical characterization of rounding error: the specific mechanism (significand overflow → right shift → round-up with sticky bit) explains why the bias is negative rather than just observing that it exists.
• Dynamic, conditional softmax adjustment over fixed offset or unconditional modification: prevents the trigger condition (¯P[T, t] = 1) without introducing new numerical instabilities (underflow from large exponents, new systematic rounding errors from fixed offsets).
• Validation by minimal intervention over proposing a new stabilization technique: the modification tests the diagnosis rather than claiming to be the best solution; its simplicity (a few lines in the softmax) matches the specificity of the diagnosed mechanism.

4. Key Insights and Innovations

Innovation 1: The Failure Is an Emergent Interaction, Not a Single Bug

The paper’s most fundamental conceptual move is reframing the training failure from “a numerical bug in flash attention” to an emergent interaction between two independently benign properties: the low-rank structure of attention representations and the rounding behavior of BF16 arithmetic. Neither factor alone causes failure — BF16 training with standard attention is stable, and flash attention in FP32 is stable — but their combination creates a positive feedback loop that drives loss explosion.

Prior work treated low-precision training failures as problems of insufficient precision (gradients underflowing in FP16, values falling outside FP8’s dynamic range). The solutions — loss scaling (Micikevicius et al., 2017), per-tensor scaling (Perez et al., 2023), stochastic rounding (Ben Ali et al., 2024) — all aim to keep values within representable range or eliminate bias in individual operations. This paper demonstrates that the BF16 + flash attention failure is qualitatively different: the precision is not too low in an absolute sense (most of the model trains fine in BF16), but a specific pattern of values (multiple ¯P = 1 entries) triggers a specific arithmetic pathology (biased round-up under significand overflow) that exploits a specific structural property of the gradients (low-rank coherence). The failure is not that BF16 is “not enough bits” — it’s that those 7 fraction bits, under these precise conditions, produce errors that all point in the same direction in weight space.

This is a conceptual shift with implications beyond this paper. It suggests that other mysterious instabilities — in FP8 training, in larger models, in different architectures — may similarly require emergent interaction explanations rather than simply “use more bits.” The paper’s analytical workflow (isolate error source → identify accumulation mechanism → trace to root arithmetic cause, Appendix F) operationalizes this reframing as a diagnostic methodology. Rather than treating numerical stability as a binary property (stable/unstable), the paper treats it as a causal chain that can be decomposed and analyzed component by component. This is a fundamentally more powerful approach than the empirical trial-and-error that produced QK normalization and gradient clipping — it enables reasoning about why an intervention works, which in turn enables predicting whether it will transfer to new settings.

The strength of this reframing is evidenced by how cleanly it explains previously puzzling observations. The localization of failure to a single head (head 8, layer 2) makes sense under the interaction view: that head happened to develop the combination of low-rank (PK)[T]⊤X[T] structure and attention-score distributions that produce multiple identical maxima. The effectiveness of QK normalization (disrupting the low-rank structure) and QK-clipping (reducing the probability of ¯P = 1) are unified as interventions on different nodes of the same causal chain. The observation that the failure is deterministic and reproducible (given fixed data) but only occurs after thousands of steps reflects the slow accumulation of biased weight updates along a coherent low-rank direction — a process that requires both time and structural consistency to build to catastrophic levels.

Where this innovation sits on the incremental-to-fundamental spectrum: fundamental reframing. It changes what “solving” a numerical instability means — from finding an empirical patch to understanding a causal mechanism — and provides a template for how to do that understanding.

Innovation 2: The Diagnostic Methodology as a Transferable Contribution

The paper’s second distinctive contribution is the analytical workflow itself, positioned explicitly as a generalizable framework rather than a one-off investigation. The three-step methodology — (1) isolate the error source through systematic high-precision substitution, (2) identify accumulation mechanisms by decomposing gradient errors into coefficient and direction components, (3) trace to root arithmetic cause via bit-level analysis of the specific floating-point operation — is presented (Section 5, Appendix F) as a blueprint for diagnosing numerical instabilities in other architectures, scales, and precision formats.

This matters because the field currently lacks such a methodology. The dominant approach to numerical instability in deep learning is empirical: observe a loss spike, try stabilization techniques (normalization, clipping, different optimizers, different precision), and keep what works. This produces a growing catalog of useful tricks but no understanding of why they work or when they will fail. The consequence is that each new model scale, architecture, or precision format requires rediscovering stability through trial and error — an increasingly expensive proposition as training runs cost millions of dollars.

The paper’s analytical workflow is a different kind of contribution than a new method or architecture. It is a process innovation — a way of thinking about and investigating numerical failures that is more systematic and generalizable than the current ad-hoc approach. The paper validates this methodology not by claiming it produces the optimal solution (the dynamic softmax modification is explicitly a validation experiment, not a recommended fix) but by demonstrating that it produces a mechanistically complete explanation: a causal chain from individual BF16 additions to loss explosion where each link is experimentally verified.

Several design choices in the paper’s investigation are non-obvious and instructive as components of the methodology:

• Deterministic data replay (Section 3.1): by fixing the data stream, the authors eliminate an entire class of confounding variables. If the failure were data-dependent (triggered by a rare batch), any intervention that changed the data order could spuriously appear to “fix” the problem. The replay ensures that changes in training behavior are causally attributable to the numerical modification.

• Localization by minimum sufficient intervention (Section 3.2): rather than analyzing the full model, the authors progressively narrow the failure scope by replacing operations with high-precision equivalents. Each replacement that stabilizes training identifies a node in the causal chain. This converges on the minimal causal unit (δlp in head 8, layer 2) without assuming which part of the system is responsible.

• Error decomposition into coefficient × direction (Section 3.3.1): the mathematical decomposition of dWQhp - dWQlp = α Σ(δlp - δhp)[T] · (PK)[T]⊤X[T] separates the gradient error into a scalar weighting term and a rank-1 direction term. This separation is what enables the diagnosis that both bias in the coefficients AND structural similarity in the directions are necessary for catastrophic accumulation. Without this decomposition, the error would appear as an opaque tensor difference, and the role of low-rank structure would be invisible.

• Bit-level root cause analysis (Section 3.3.2): the paper does not stop at “rounding errors are biased” — it traces the bias to the specific mechanism of significand overflow, right shift, and round-up with sticky bit in BF16 addition. This level of detail is what enables the targeted intervention (preventing ¯P = 1) rather than a coarser fix (e.g., “use more precision everywhere”).

The significance of this innovation lies not in its novelty — debugging by bisection and high-precision substitution is standard practice in numerical analysis — but in its systematic application to deep learning training dynamics, where the complexity of the system (millions of interacting parameters, stochastic optimization, distributed computation) has historically made such analysis seem intractable. The paper demonstrates that, at least for this failure mode, the complexity can be managed by aggressive localization (to a single head) and structural decomposition (separating coefficients from directions). This opens the door to similar analyses of other instabilities.

Where this sits on the spectrum: incremental as a debugging technique, fundamental as a contribution to deep learning practice. Each individual step (high-precision substitution, error decomposition) is standard, but their integration into a coherent workflow that produces a complete mechanistic explanation — and the paper’s explicit framing of this workflow as a transferable contribution — is a meaningful advance in how the field can approach numerical stability.

Innovation 3: The Low-Rank Coherence Mechanism as an Explanation for Empirical Phenomena

The paper’s third distinctive contribution is the identification of structural similarity in gradient error directions as the amplification mechanism that transforms small per-step rounding errors into catastrophic weight corruption. This is not just a property of this specific failure — it provides a unifying mechanistic explanation for several independently observed but poorly understood phenomena in transformer training.

The key insight is that (PK)[T]⊤X[T] matrices — the rank-1 components of the weight gradient error — are not random and independent across tokens and training steps. They share a similar low-rank structure (Figure 4), which means that when the scalar coefficients (δlp - δhp)[T] are systematically biased (rather than zero-mean), the per-step errors accumulate along a consistent direction R rather than canceling. This is what the paper formalizes as dWQhp - dWQlp ≈ α · Σ(δlp-δhp)[T] · R.

The significance of this finding extends beyond the immediate failure case. It provides a mechanistic explanation for three empirical observations that the broader training stability literature has documented but not explained:

1. Growth of weight spectral norms (Yang et al., 2023; Rybakov et al., 2024): the observation that unstable training runs exhibit anomalously large spectral norms. Under the paper’s mechanism, this growth is a direct consequence of accumulating biased updates along R — each step adds a small multiple of R to WQ, and since these multiples are consistently positive, the weight matrix’s norm along R grows monotonically. The spectral norm (the largest singular value) increases because the updates are concentrated in a low-rank subspace rather than distributed isotropically.

2. Attention sinks (Xiao et al., 2023): the phenomenon where certain tokens (often the first token or punctuation) attract disproportionately high attention scores. The paper provides a numerical mechanism: attention sinks create the condition of multiple ¯P = 1 entries (because the sink token and other tokens both achieve maximum attention scores), which triggers the biased rounding that initiates the failure cascade. This transforms attention sinks from an architectural curiosity into a direct causal factor in training instability — they are the conditions under which the failure mechanism activates.

3. Effectiveness of QK normalization and Gated Attention (Henry et al., 2020; Qiu et al., 2025): these architectural modifications empirically stabilize training. The paper’s mechanism explains why they work: they disrupt the structural similarity of (PK)[T]⊤X[T] matrices. QK normalization constrains the magnitude and direction of query and key vectors, preventing the emergence of consistent low-rank structure. Gated Attention introduces non-linearity that breaks the coherence of error directions. In both cases, the interventions don’t eliminate rounding errors — they eliminate the structural pathway that allows rounding errors to accumulate systematically. Rounding errors still occur, but they point in random directions and cancel out over training steps.

What makes this innovation more than a restatement of the paper’s mechanism is its explanatory unification. Prior to this work, attention sinks, spectral norm growth, QK normalization, and training instability were related only by empirical correlation — “models with attention sinks sometimes diverge, and QK normalization sometimes helps.” The paper provides a causal framework that connects these phenomena: attention sinks → repeated maxima → biased rounding → coherent gradient errors → spectral norm growth → instability, with QK normalization intervening at the structural coherence step. This transforms a collection of observations into a coherent causal narrative.

The paper also provides preliminary evidence that this structural similarity is not unique to GPT-2. Appendix D shows similar patterns of coherent (PK)[T]⊤X[T] structure in Llama-3.1-8B (Figure 8) and multiple attention maxima (Figure 9), suggesting that the conditions for this failure mode exist in larger, more modern models — even if those models don’t fail in BF16 because they employ stabilization techniques or train at scales where other factors dominate.

Where this sits on the spectrum: fundamental conceptual contribution. The low-rank coherence mechanism is not an incremental refinement of existing understanding — it is a genuinely new explanation for why rounding errors in deep learning can be catastrophic rather than benign, and it unifies several previously disparate empirical observations under a single causal framework. This is the kind of insight that enables principled reasoning about numerical stability rather than empirical guesswork.

Innovation 4: The Validation-by-Minimal-Intervention Paradigm

The paper’s fourth distinctive contribution is methodological: the use of a targeted, minimal modification not as a proposed solution but as a validation experiment for a mechanistic hypothesis. The dynamic softmax adjustment (Section 4) is explicitly not positioned as “the fix for low-precision flash attention” — it is a test of the causal chain: if the mechanism we diagnosed is correct, then preventing ¯P[T, t] = 1 should prevent the failure, while changing nothing else about the computation.

This approach is subtly but importantly different from how stabilization techniques are typically developed and evaluated in deep learning. The standard approach is:

1. Observe instability.
2. Propose a technique that seems plausible (e.g., QK normalization, gradient clipping).
3. Show that it prevents the instability in experiments.
4. Conclude that the technique “solves” the problem.

The problem with this approach is that it doesn’t distinguish between interventions that address the root cause and interventions that address a downstream symptom. Gradient clipping, for example, prevents loss explosion by capping gradient magnitudes — but it doesn’t explain why gradients were exploding in the first place. A model trained with gradient clipping might still have biased rounding errors and low-rank accumulation; it just can’t express them catastrophically because the optimizer bounds the updates.

The paper’s validation-by-minimal-intervention approach is different:

1. Diagnose the complete causal chain (Sections 3.2–3.3.2).
2. Design an intervention that targets the root cause — the earliest node in the causal chain where a change would break the feedback loop.
3. Show that this minimal, targeted change prevents the failure.
4. Interpret success as confirmation of the mechanism, not as a claim that the intervention is the optimal or only solution.

The dynamic softmax modification is an ideal validation experiment because it is surgical: it changes only the normalization constant in the softmax when specific conditions are met, preserving mathematical equivalence in exact arithmetic, and targeting precisely the condition (¯P = 1) that the mechanism identifies as the trigger. Its success (Figure 7) — preventing loss explosion across two model sizes and two optimizers — is strong evidence that the diagnosed mechanism is correct, because the intervention is too narrowly targeted to fix the problem through any other pathway.

This paradigm has implications for how the field should approach numerical stability research. Rather than racing to propose new stabilization techniques, the paper demonstrates the value of first achieving mechanistic understanding, then using that understanding to design targeted interventions, and interpreting those interventions as hypothesis tests rather than solutions. This is closer to how physics or biology investigate failures — formulate a causal model, design an experiment that would confirm or refute it, and interpret the result as evidence about the model — than to the empirical trial-and-error that currently dominates deep learning stability work.

The paper is transparent about the limitations of this approach (Section 5): the dynamic softmax modification is not a recommended production fix (other interventions like QK normalization may be more practical or general), and the analysis is specific to BF16 + flash attention + GPT-2. But the paradigm — diagnose first, validate with a minimal intervention, interpret as confirmation — is positioned as a template for investigating other instabilities.

Where this sits on the spectrum: incremental as an individual experiment, fundamental as a research paradigm. The dynamic softmax modification itself is a small code change (a few lines in Algorithm 3), but the paper’s use of it as a validation instrument rather than a solution represents a meaningful shift in how the field could approach stability research — prioritizing mechanistic understanding over empirical patches, and using minimal interventions as diagnostic tools rather than as performance claims.

5. Experimental Analysis

Evaluation Methodology

• Dataset. All experiments use the OpenWebText corpus (Gokaslan et al., 2019), an open-source recreation of the WebText dataset used in GPT-2 pretraining. For deterministic reproducibility, the authors record and reuse the exact sequence of data batches from an initial failing run, ensuring every experiment processes identical data in identical order and isolating failure from data-related randomness (Section 3.1).

• Base model(s). The primary model is GPT-2 Small (GPT-2S): 12 layers, 12 attention heads, embedding dimension 768, context length 1024 (approximately 125M parameters). This scale is chosen because it is large enough to exhibit the failure reproducibly but small enough to enable detailed per-head diagnostic analysis — individual attention head computations can be traced, visualized, and selectively modified. A larger GPT-2 Medium (GPT-2M) variant is used for generalization testing in Section 4. For cross-model evidence, the authors also analyze Llama-3.1-8B in Appendix D, though this analysis is observational (visualizing structural similarity of gradient components and presence of multiple attention maxima) rather than a training stability experiment.

• Metrics. Two primary metrics are tracked:

  • Validation loss (reported in Figures 2, 7, and Figure 13 in Appendix): the standard cross-entropy loss on held-out data, used as the primary indicator of training stability versus catastrophic divergence.
  • Spectral norm of weight matrices (reported in Figures 3 and 11): the largest singular value of the query projection matrix WQ for individual attention heads, used as a diagnostic for detecting anomalous weight growth that precedes loss explosion, following the methodology of Yang et al. (2023) and Rybakov et al. (2024).
  • For the Llama-3.1-8B analysis in Appendix D: cosine similarity between columns of (PK)[T]⊤X[T] matrices from different token positions, used to quantify structural similarity of gradient error directions.

• Baselines. The paper does not compare against alternative stabilization techniques as baselines in the traditional sense — it is a mechanistic analysis, not a method comparison. The key "baselines" are:

  • High-precision flash attention (FP32 forward pass): the stable configuration shown in Figure 2, serving as the reference for what correct training behavior looks like.
  • Standard BF16 mixed-precision flash attention: the failing configuration, serving as the "untreated" condition against which modifications are evaluated.
  • Various component-level high-precision substitutions (e.g., computing only ¯PV in FP32, computing only O in FP32, replacing rowsum(dO ◦ O) with rowsum(dP ◦ P)) used as diagnostic experiments rather than competing methods.
  • Muon optimizer (Jordan et al., 2024): used in Section 4 to demonstrate that the stabilized flash attention works across optimizer choices, not just AdamW.

• Generation budget / compute accounting. The paper does not involve test-time compute budgets or generation-based metrics since it analyzes a pretraining failure, not inference-time strategies. Compute is measured in training steps (iterations of the optimizer). The primary experiments run to 20,000 steps for failure analysis (Section 3), 600,000 steps for GPT-2S stability validation, and 100,000 steps for GPT-2M validation (Section 4). The precision configuration is automatic mixed precision: forward pass in BF16 (including flash attention), backward pass gradient accumulation in FP32, with an FP32 master copy of weights — this is the standard PyTorch recipe and is representative of industrial training practice.

• Cross-validation / statistical protocol. There is no cross-validation or statistical significance testing in the traditional sense, which is appropriate for a mechanistic analysis paper. The key reproducibility measure is deterministic data replay (Section 3.1): by recording and reusing the exact sequence of data batches, the authors eliminate data sampling as a source of variance. This means that any difference between two experimental configurations (e.g., high-precision δ versus low-precision δ) is causally attributable to the numerical modification rather than to data randomness. The paper reports results from single training runs under each configuration, which is appropriate given the deterministic setup — multiple random seeds would only add variance from data order, which has been eliminated by the replay design. However, the paper does not report whether the failure occurs consistently across multiple independent initial data orders (e.g., different random shuffles of OpenWebText), which would strengthen the claim that the mechanism is general rather than dataset-order-specific.

Main Quantitative Results

Since this is a mechanistic analysis paper rather than a method-comparison paper, the "quantitative results" are primarily diagnostic measurements that establish the causal chain, plus validation experiments showing that the identified mechanism is sufficient to cause the failure. I organize these by the stages of the investigation.

Failure Reproduction and Isolation

The failure is reproducible and deterministic. Figure 2 shows the validation loss curve for the low-precision (BF16) configuration: training proceeds normally for approximately 6,600 steps, then the loss explodes catastrophically — rising from a converged value around 3–4 to above 10 within a few hundred steps. The high-precision (FP32) configuration trains stably throughout. The paper corroborates this with community reports: nanoGPT Issue 303 (2023), Issue 524 (2024), and Issue 554 (2024) all document the same failure pattern, with loss curves shown in Appendix Figure 10 from two independent runs exhibiting identical sudden-divergence behavior.

The failure is localized to a single layer and computation. The diagnostic experiments in Section 3.2 produce a series of binary outcomes (failure vs. stability) that narrow the source:

• Disabling tiling (setting block size equal to sequence length): failure persists. This rules out the tiling strategy as causal.
• Using flash attention only in layer 2 (standard attention elsewhere): failure reproduces.
• Replacing flash attention with standard attention in layer 2 (flash attention elsewhere): training stabilizes. This localizes the failure origin to layer 2.
• Replacing δ = rowsum(dO ◦ O) with δ = rowsum(dP ◦ P) (mathematically equivalent in exact arithmetic): training stabilizes. This identifies the low-precision computation of δ as the causal node.
• Recomputing O = PV in FP32 during backward pass: training stabilizes. This isolates the error to Olp.
• Computing O in FP32 during forward pass (all other operations in BF16): training stabilizes. This confirms that the error in Olp — specifically, the difference between BF16 Olp and what Ohp would be — is the direct cause.

The paper does not report quantitative metrics for each of these diagnostic experiments (e.g., "validation loss at step X with intervention Y"), which is a minor limitation — the outcomes are binary (explodes vs. doesn't) at the scale of Figure 2's loss curve.

The failure is localized to specific attention heads. Figure 3 shows the spectral norm of WQ for each of the 12 attention heads in layer 1 (note: this is layer 1 in the figure, while the failure originates in layer 2 per the text — there may be a discrepancy in the paper's reporting, or the figure shows layer 1 as a representative example while layer 2 is the actual failure locus). The spectral norms are: head 1 (4.34), head 2 (2.68), head 3 (1.90), head 4 (2.67), head 5 (2.04), head 6 (2.66), head 7 (4.12), head 8 (6.27), head 9 (5.32), head 10 (1.74), head 11 (4.64), head 12 (5.39). Head 8 has the largest spectral norm, and the paper reports that "selectively computing the output O in high precision for these outlier heads (1, 7, 8, 9, 11, and 12), which is sufficient to restore training stability." The exact number of heads requiring FP32 O computation is not stated as a threshold — it's possible that fixing only head 8 would suffice, or that a subset of the six outlier heads is necessary, but the paper does not perform this ablation.

Quantitative Evidence for Low-Rank Gradient Error Accumulation

The cumulative sum of (δlp - δhp)[T] is consistently positive. Figure 5(a) tracks Σ(δlp - δhp)[T] accumulated over training steps 6580 to 6680 — the period leading up to the failure. The cumulative sum starts near zero, grows approximately linearly to ~0.00038 by step 6600, then accelerates to ~0.00040 by step 6680. The sum is monotonically increasing (no significant negative fluctuations), confirming a systematic positive bias rather than zero-mean random noise. The paper does not report the absolute magnitude of individual (δlp - δhp)[T] values or their variance, only the cumulative trajectory over ~100 steps.

The gradient error direction matrices are structurally similar. Figure 4 provides visual evidence that (PK)[T]⊤X[T] matrices share structural features across tokens and training steps. Panels (c) and (f) highlight specific columns (input features 546 and 678) that exhibit similar patterns for token 50 at step 6610 and token 718 at step 6619. The paper does not provide a quantitative similarity metric (e.g., cosine similarity, CCA, subspace overlap) for these matrices in the GPT-2 analysis. However, Appendix D does report cosine similarity scores for Llama-3.1-8B: for column 2598, the cosine similarity between (PK)[T]⊤X[T] at token 3 ("brown") and token 4 ("fox") is 0.999994; for column 2050, the similarity is 0.999994; for column 3899, it is 0.999993. These near-1.0 cosine similarities provide quantitative evidence that the structural similarity observed qualitatively in GPT-2 (Figure 4) also exists in larger, more modern models, though the Llama analysis is observational (not linked to a training failure).

Quantitative Evidence for Biased Rounding in ¯PV

The output error Olp - Ohp is systematically negative for specific features. Figure 5(b) and (c) show that for token position T = 718, the values of O[T, 20] and O[T, 29] (specific feature dimensions) differ between low-precision and high-precision: the error Olp[T, i] - Ohp[T, i] is approximately -0.012 to -0.014 for these features. The upstream gradient dO[T, i] at these same dimensions is approximately -1.5 × 10^{-6}. The product dO[T, i] × (Olp - Ohp)[T, i] is therefore positive for both dimensions, contributing positively to the δ error. The paper reports this for two specific feature dimensions but states that "this strong sign correlation is also observed across other tokens," though no aggregate statistics are provided.

The ¯PV error accumulates in large negative steps at positions where ¯P[T, t] = 1. Figures 6(b) and (c) show the cumulative error ¯Oerror(t) for feature dimension i = 20 as the sum over token positions t progresses. The cumulative error begins near zero, experiences several sharp downward jumps of magnitude -0.014 to -0.015 at specific token positions (around tokens 640–680 in the detailed view of Figure 6c), and ends at approximately -0.12 after all 1024 tokens. The sharp downward jumps are shown to coincide with positions where ¯P[T, t] = 1 (the attention probability is exactly 1). Figure 6(a) provides the context for why these jumps are negative: for feature dimension i = 20, the values of V[:, 20] are "predominantly negative," ranging from approximately -6 to +2 but with the vast majority of tokens having negative values in the range [-6, 0]. The paper presents one specific bit-level example (Section 3.3.2) of a BF16 addition that produces a negative error of -0.014759540557861328, and states that "when such rounding events occur systematically across many additions in the ¯PV product, the errors accumulate."

How many tokens are affected? The paper does not aggregate how many token positions per sequence have ¯P[T, t] = 1 or how many feature dimensions of V are predominantly negative. Figure 13 in Appendix E shows the "number of activated SFA condition" (i.e., number of rows with multiple identical maxima) plotted alongside the validation loss curve. The number of multiple-maxima rows begins increasing before the loss explodes: at step 7000, the number of activated conditions spikes sharply to the range of 2 × 10^{9} (the paper uses this unit; it likely refers to a cumulative count or a batch-aggregated total — the exact meaning is not fully explained), roughly coinciding with the loss divergence. This temporal correlation suggests that the emergence of repeated maxima is a leading indicator of impending failure, consistent with the mechanism where repeated maxima trigger the biased rounding that drives the loss explosion.

Validation via Stabilized Flash Attention

The modified softmax prevents loss explosion across configurations. Figure 7 presents the main validation results:

• GPT-2S + AdamW (Figure 7a): Training for 600K steps (substantially longer than the ~7K steps at which the original failure occurs). The stabilized flash attention (SFA, β = 2) converges smoothly, with validation loss decreasing from approximately 10 to below 3, following a trajectory nearly identical to what high-precision flash attention would produce (compare to Figure 2's stable FP32 curve). The original low-precision flash attention diverges catastrophically around step 7K.

• GPT-2S + Muon (Figure 7b): The Muon optimizer (Jordan et al., 2024) is a recently proposed alternative to AdamW that uses matrix orthogonalization. The stabilized flash attention also prevents loss explosion for 600K steps, with validation loss decreasing from approximately 10 to approximately 3. The original flash attention diverges. This demonstrates that the fix is not optimizer-specific — the failure mechanism is numerical, not optimizer-dependent.

• GPT-2M + AdamW (Figure 7c): The larger model is trained for 100K steps. The stabilized flash attention prevents loss explosion, with validation loss decreasing from approximately 10 to between 3 and 4. The original flash attention diverges. This provides preliminary evidence of generalization to larger model scales, though 100K steps is shorter than the 600K run for GPT-2S.

What β value was used? The paper states β = 2 for the experiments in Figure 7 and notes that β ∈ [2, 8] is the recommended range. No sweep over β values is reported — there is no evidence about sensitivity to this hyperparameter. A β that is too small (e.g., 1.001) risks ¯P rounding back to 1 in BF16; a β that is too large risks underflow. The choice of β = 2 appears to be a reasonable default but is not empirically justified in the paper.

Ablation Studies and Robustness Checks

Since this is a mechanistic analysis rather than a method paper, the "ablations" are primarily the diagnostic isolation experiments described in Section 3.2, which I have already covered in detail above. These serve the same function as ablations — systematically removing or modifying components to identify which are necessary for the failure. I summarize them here with a focus on the logical structure of the evidence they provide:

• Tiling ablation: Setting block size equal to sequence length (disabling tiling) does not prevent failure. This demonstrates that the failure is not caused by block-wise accumulation of errors in the tiled flash attention algorithm — it occurs even when attention is computed with full matrices in a single pass. This is a valuable negative result because tiling was a plausible hypothesis (error accumulating across blocks in the online softmax).

• Layer-wise high-precision substitution: Using flash attention only in layer 2 reproduces the failure; replacing flash attention with standard attention in layer 2 prevents it. This demonstrates that the failure is localized to layer 2, not a diffuse numerical problem across all layers. The paper does not explain why layer 2 specifically — what is it about this layer's representations or attention patterns that makes it vulnerable? This is a gap in the analysis.

• Head-wise high-precision substitution: Computing O in FP32 for heads 1, 7, 8, 9, 11, and 12 prevents failure. This demonstrates that the failure is localized to these specific attention heads. The paper does not report whether computing FP32 O for only head 8 (the one with the largest spectral norm) would suffice, which would further narrow the causal unit. This ablation is missing.

• δ computation substitution: Replacing δ = rowsum(dO ◦ O) with δ = rowsum(dP ◦ P) prevents failure. This is a critical experiment because the two formulations are mathematically identical — the only difference is whether O is recomputed from P and V in the backward pass or reused from the forward pass. The fact that this change fixes the failure isolates the error to the stored Olp from the forward pass, not to any structural property of how δ is used in subsequent gradient computations.

• ¯PV high-precision substitution: Computing ¯O = ¯P V in FP32 (while keeping everything else in BF16) prevents failure. This narrows the root cause to the matrix multiplication in the unnormalized output computation, ruling out the subsequent normalization (O = ¯O / rowsum(¯P)) as the source of bias.

• Beyond the isolation experiments, the paper includes several additional robustness checks and supporting analyses:

• ¯P = 1 mechanism validation via multiple-maxima frequency tracking (Appendix E, Figure 13): The frequency of rows with multiple identical maxima (the "activated SFA condition") is plotted alongside the validation loss. The number of multiple-maxima rows increases sharply before the loss explodes, providing temporal evidence that the trigger condition (¯P = 1 at multiple positions) emerges as a precursor to failure, not as a consequence of it. However, correlation is not causation — it's possible that some third factor causes both the increased multiple maxima and the loss explosion. The causal role of ¯P = 1 is established by the intervention in Section 4, not by this correlation.

• Cross-model evidence in Llama-3.1-8B (Appendix D): This is not a training experiment but an observational analysis of a pretrained Llama-3.1-8B checkpoint. Figure 8 shows that (PK)[T]⊤X[T] matrices for tokens 3 ("brown") and 4 ("fox") in layer 1, head 13 share structurally similar columns with cosine similarities > 0.999. Figure 9 shows that multiple attention maxima (values of 1.0 in the attention probability distribution) occur in various layers and heads of Llama-3.1-8B — for example, layer 1 head 6 query 44, layer 16 head 27 query 11, layer 26 head 1 queries 45 and 82, layer 29 head 10 query 96, and layer 30 head 24 query 49 all exhibit attention distributions where the maximum probability is 1.0 and shared across multiple key positions (typically including the first token, consistent with attention sink behavior). This demonstrates that the preconditions for the failure — structural similarity of gradient error directions and multiple identical attention maxima — exist in models beyond GPT-2. However, the paper does not claim that Llama-3.1-8B actually fails in BF16 training (it is provided as a pretrained model, and its training configuration is not specified). The Llama analysis is suggestive but not conclusive — it shows that the conditions for failure exist, not that the failure would occur under BF16 training of Llama-3.1-8B.

• Optimizer robustness (Figure 7): The stabilized flash attention works with both AdamW and Muon optimizers. This demonstrates that the failure is not caused by an interaction with AdamW's specific update rule (e.g., momentum accumulation of biased gradients). The paper does not test other common optimizers (e.g., SGD, Lion, Adafactor), so the claim of optimizer-independence is supported for two optimizers but not exhaustively tested.

• Model scale robustness (Figure 7c): The stabilized flash attention works for GPT-2M (larger than GPT-2S). The paper does not specify the exact parameter count or architecture of GPT-2M, and trains it for only 100K steps compared to 600K for GPT-2S, so the evidence for scale generalization is preliminary.

Critical Assessment

Does the evidence support the central claim that biased BF16 rounding in ¯PV causes the failure?

The evidence for this claim is strong but circumstantial, built from a chain of experiments where each link is tested independently:

1. The isolation experiments (Section 3.2) convincingly narrow the failure to Olp → δlp in a specific attention head. Every alternative explanation tested (tiling, other layers, other computations) is ruled out. The experiments are clean binary comparisons — does changing X fix the failure? — and the answers are unambiguous.

2. The gradient error analysis (Section 3.3.1) provides strong evidence that the error accumulates along a coherent low-rank direction rather than canceling. The cumulative sum in Figure 5(a) shows monotonic growth, which is exactly what the theoretical decomposition predicts if both bias and structural similarity are present. However, the quantitative measurement of structural similarity is weak for the GPT-2 case — Figure 4 is a qualitative visualization, not a quantitative comparison. The Llama-3.1-8B cosine similarities (>0.999) are compelling quantitative evidence that such structure exists in transformers, but they are from a different model and are not linked to a training failure.

3. The rounding error analysis (Section 3.3.2) provides a plausible mechanism for why Olp - Ohp is negative and why δlp - δhp is positive. The bit-level worked example is convincing for the specific addition analyzed, and the visual correlation in Figure 6 between ¯P[T, t] = 1 and negative error jumps is suggestive. However, the paper does not provide aggregate statistics on how often the biased rounding condition occurs per training step, how many feature dimensions are affected, or what fraction of the total δlp - δhp error is attributable to this mechanism versus other rounding errors. The analysis is a existence proof (this mechanism can produce the observed bias) rather than a quantitative accounting (this mechanism explains X% of the total error). It is possible that other, unidentified rounding errors also contribute to the bias, and the ¯P = 1 mechanism is just one component.

4. The validation experiment (Section 4) is the strongest evidence: a minimal, targeted intervention that addresses only the diagnosed mechanism (preventing ¯P = 1) and that stabilizes training across configurations. If the diagnosed mechanism were wrong or incomplete, this intervention would not be expected to work — or at minimum, it would only partially mitigate the failure. The fact that it completely prevents loss explosion for 600K steps (Figure 7) is strong evidence that the ¯P = 1 → biased rounding → coherent accumulation chain is the dominant failure pathway.

Overall assessment: The evidence is sufficient to support the mechanistic claim, but the quantitative accounting is incomplete. We know that this mechanism causes the failure (because preventing it fixes the problem), but we don't know how much of the total numerical error it accounts for or whether other, secondary mechanisms also contribute.

Does the evidence support the claim that this is the first mechanistic explanation?

The paper's claim to be "the first mechanistic explanation" (Section 1, Section 3.1, Section 5) is supported by negative evidence — the authors cite several community bug reports (nanoGPT Issues 303, 524, 554; flash-attention Issue 337) and prior work (Lee et al., 2024; Golden et al., 2024) that document the failure but do not provide a causal mechanism. The paper's explanation is empirically grounded (each link in the chain is tested) and mechanistically complete (from individual BF16 addition to loss explosion). Whether it is literally the first such explanation depends on unpublished work or industry internal knowledge, but within the public literature, the paper's claim appears valid.

The more important question is whether the explanation is correct and complete, not whether it is first. On correctness: the validation experiment provides strong evidence. On completeness: there are gaps (discussed below), but the explanation covers the major causal nodes — trigger (¯P = 1), bias mechanism (significand overflow in BF16 addition), amplification (low-rank gradient error accumulation), and observable consequences (spectral norm growth, loss explosion).

Genuine weaknesses in the experimental evidence

1. No aggregate statistics on the biased rounding mechanism. The paper analyzes one specific token position (T = 718), one feature dimension (i = 20), and one specific BF16 addition in detail. It shows that ¯P[T, t] = 1 correlates with negative error jumps in the cumulative error plot (Figure 6c). But it does not report: how many tokens per sequence have ¯P = 1? How many feature dimensions of V are predominantly negative? What fraction of the total δ error is accounted for by this mechanism? Without these aggregates, we cannot assess whether the analyzed addition is representative or cherry-picked — it is possible that most additions do not exhibit biased rounding, and the paper has selected the one that does.

2. No quantitative measurement of structural similarity in GPT-2. The evidence for Claim 2 (low-rank coherence of gradient errors) in GPT-2 is Figure 4, which is a qualitative heatmap visualization. The paper does not report cosine similarities, subspace angles, or any quantitative metric for how similar (PK)[T]⊤X[T] matrices are across tokens and steps in GPT-2. The Llama-3.1-8B analysis (Appendix D) provides quantitative cosine similarities (>0.999), which is compelling, but this is a different model and a forward-pass analysis of a pretrained checkpoint, not an analysis of training dynamics. The structural similarity claim is central to the mechanism (without it, biased errors would cancel), yet the quantitative evidence for it in the actual failing model is thin.

3. Single deterministic data order. The paper's use of a fixed data replay is methodologically elegant for causal attribution — it eliminates data randomness as a confound. However, it also means the entire analysis is based on one specific sequence of training data. The paper does not report whether the failure occurs for other data orders (different random shuffles of OpenWebText) or for other random seeds. If the failure only occurs for certain data orders, the mechanism might be triggered by specific rare data patterns rather than being a general property of BF16 + flash attention. The community reports (nanoGPT issues) suggest the failure is common, but the paper's controlled experiments are all within a single data trajectory.

4. Missing ablation on which heads must be in FP32. The paper reports that selectively computing O in FP32 for six outlier heads (1, 7, 8, 9, 11, 12) stabilizes training, but it does not ablate whether fixing only head 8 (the one with the largest spectral norm) would suffice. If fixing only head 8 works, the mechanism is even more localized than reported. If fixing all six is necessary, there may be cross-head interactions that the paper does not analyze.

5. No sweep over β in the validation experiment. The stabilized flash attention uses β = 2. The paper mentions that β ∈ [2, 8] is the recommended range, but no experiments test sensitivity to this hyperparameter. A β that is too small might not prevent ¯P from rounding to 1; a β that is too large might cause underflow. Without a sensitivity analysis, the practical guidance for choosing β is incomplete.

6. Limited scale generalization. The GPT-2M experiment (Figure 7c) provides preliminary evidence of generalization, but 100K training steps is only 1/6 of the GPT-2S run (600K steps). The paper does not specify the exact architecture of GPT-2M. The Llama-3.1-8B analysis is observational, not a training experiment — we don't know whether Llama-3.1-8B would fail under the same conditions or whether its training recipe (which may include QK normalization, different precision settings, or other stabilization) would prevent the failure.

7. No comparison to existing stabilization techniques. The paper's dynamic softmax modification is presented as a validation experiment, not a proposed solution. However, the paper does not experimentally compare it to existing stabilization techniques (QK normalization, QK-clipping, Gated Attention) to demonstrate that it addresses the same root cause. The claim in Section 5 Discussion that these techniques "disrupt the underlying structure" of (PK)[T]⊤X[T] is a mechanistic interpretation, not an experimentally tested hypothesis. An experiment showing that QK normalization, like the dynamic softmax, prevents the loss explosion and disrupts the low-rank structure would strengthen the unification claim.

8. No analysis of distributed training effects. The experiments use DDP on 4 GPUs, but the analysis does not explore whether gradient synchronization across GPUs amplifies or mitigates the error accumulation. If the biased gradient updates are consistent across GPUs (because they all process data from the same distribution and the mechanism is data-independent), synchronization might amplify the bias. If they are inconsistent, synchronization might provide some averaging benefit. The paper does not address this.

Missing experiments that would strengthen the paper

• Quantitative decomposition of the δ error: measure what fraction of the total δlp - δhp error is attributable to the ¯P = 1 mechanism (by comparing the actual error to a counterfactual where only ¯P < 1 additions occur). This would quantify how dominant this mechanism is versus other rounding errors.

• Sensitivity to V value distribution: if the mechanism depends on V[:, i] being predominantly negative, experiments that artificially balance the signs of V (e.g., by adding a constant offset) should reduce or eliminate the bias. This would test the causal role of V's sign distribution more directly.

• Sweep over β values in the stabilized flash attention to characterize the sensitivity and provide practical guidance.

• Comparison of the dynamic softmax to QK normalization in the same experimental setup, measuring both training stability and the structural similarity of (PK)[T]⊤X[T] under each intervention.

• Multiple data orders to assess whether the failure and the proposed mechanism generalize beyond the specific recorded data trajectory.

• Analysis of why layer 2, head 8 specifically develops the problematic combination of low-rank structure and multiple attention maxima, while other layers and heads do not. This would require analyzing the representations and attention patterns across layers during early training.

Conditions under which the claims hold

The paper's central claim — that biased BF16 rounding in ¯PV, amplified by low-rank gradient error coherence, causes the training failure — is supported for this specific configuration: GPT-2S, BF16 mixed precision, flash attention, OpenWebText, AdamW, the specific hyperparameters listed in Section 3.1, and the specific data order of the recorded trajectory. The paper provides evidence of generalization to GPT-2M (100K steps), Muon optimizer, and suggests cross-model relevance via the Llama-3.1-8B analysis, but these extensions are preliminary.

The claim that the mechanism explains all "similar" failures in the community (nanoGPT issues, Kimi-Team, Qwen-Team reports) is plausible but unverified. The paper does not reproduce those specific failure cases or demonstrate that the same mechanism operates in them. The community failures may have different root causes that produce similar symptoms (loss explosion). The paper's analytical workflow is positioned as a general diagnostic tool, but its application to other failure cases remains future work.

The claim that the dynamic softmax modification is a "practical solution" (abstract) is qualified by the paper's own discussion: it is a validation experiment, and the authors suggest that QK normalization and Gated Attention may be more practical interventions. The dynamic softmax has not been tested at scales beyond GPT-2M or on production training pipelines, so its practicality is unproven.

6. Limitations and Trade-offs

Limitation 1: The Analysis Is Based on a Single Deterministic Data Trajectory

The assumption or constraint. The paper’s entire causal investigation — from failure isolation to bit-level rounding analysis — is conducted on one specific sequence of training data recorded from an initial failing run and replayed identically across all experiments. The authors explicitly state: "For deterministic reproducibility, 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" (Section 3.1). This means every diagnostic experiment, every visualization (Figures 4–6), and every quantitative measurement (Figure 5a) reflects the model’s behavior on a single data trajectory.

The consequence. The paper cannot distinguish between a failure mechanism that is universal to BF16 + flash attention (occurring regardless of data order) and one that is triggered by specific rare data patterns encountered in this particular trajectory. If the low-rank structure of (PK)[T]⊤X[T] or the emergence of multiple attention maxima depends on the specific sequence of training batches — for instance, a cluster of unusual documents that push the attention mechanism into a pathological regime — then the diagnosed mechanism might not reproduce under a different random shuffle of the same dataset. A practitioner training GPT-2 on OpenWebText with a different data loader seed might never encounter the failure, or might encounter it through a different causal pathway that the paper’s analysis does not cover. Conversely, the failure might be more common than the paper implies if it is triggered by many different data patterns.

What evidence exists in the paper. The paper provides no experiments with alternative data orders. All results in Sections 3–4 are from the single recorded trajectory. The community reports cited (nanoGPT Issue 303, 2023; Issue 524, 2024; Issue 554, 2024) document similar loss explosions in independent runs, but those runs use different data loader implementations, potentially different dataset versions, and different hyperparameters — they are not controlled replications of this paper’s exact setup. The loss curves from two independent runs shown in Appendix Figure 10 exhibit the same qualitative failure pattern (sudden divergence), but the paper cannot verify whether the exact mechanism — biased rounding at ¯P = 1 in head 8 of layer 2 — is operative in those runs. The deterministic replay is a methodological strength for causal attribution within this trajectory but a limitation for claims of generality across trajectories.

Mitigation status. Not addressed. The paper does not discuss the dependence on data order as a limitation, nor does it suggest experiments with alternative data trajectories. The authors present the deterministic replay as an unqualified methodological improvement without acknowledging the tradeoff in generality. Future work could test whether the failure occurs under different data shuffles, measure the variance in failure timing or mechanism across data orders, and quantify whether the trigger condition (¯P = 1 at multiple positions) is data-dependent or an inevitable consequence of the model’s learning dynamics.

Limitation 2: No Aggregate Statistics on the Prevalence of the Biased Rounding Mechanism

The assumption or constraint. The paper’s root-cause analysis (Section 3.3.2) traces the failure to biased BF16 rounding in the ¯PV product, occurring when ¯P[T, t] = 1 at multiple token positions and V[:, i] is predominantly negative. The evidence for this mechanism consists of: (a) a detailed bit-level analysis of one specific BF16 addition for token position T = 718 and feature dimension i = 20, quantifying a rounding error of -0.01476 (Section 3.3.2); (b) a cumulative error plot for i = 20 (Figure 6b–c) showing sharp negative jumps correlated with ¯P[T, t] = 1 positions; and (c) a visualization showing that V[:, 20] is predominantly negative (Figure 6a). The paper implicitly assumes that this analyzed case is representative — that similar biased rounding events occur systematically across many tokens and feature dimensions, and that they dominate the total error in δlp.

The consequence. Without aggregate statistics, a reader cannot assess how much of the total numerical error this mechanism accounts for. Several key quantities are unreported: (1) What fraction of token positions per sequence have ¯P[T, t] = 1 at multiple key positions? (2) For what fraction of feature dimensions is V[:, i] predominantly negative, creating the conditions for systematic round-up bias? (3) What fraction of the total δlp - δhp error is attributable to the ¯P = 1 mechanism versus other rounding errors (e.g., from operations where ¯P < 1, from the normalization step, from other matrix multiplications)? It is possible that the analyzed addition is the largest or most dramatic example, selected for pedagogical clarity, while most rounding errors in the ¯PV product are small, zero-mean, and harmless — and some other, unidentified mechanism is the true dominant contributor to the δ bias. The paper demonstrates sufficiency (preventing ¯P = 1 fixes the failure) but not necessity (the failure cannot occur through any other mechanism), because the validation experiment in Section 4 prevents all instances of ¯P = 1, not just the ones that would cause biased rounding in negatively-signed V columns. If the dynamic softmax also incidentally changes other numerical properties (e.g., the distribution of ¯P values, the conditioning of the softmax), the stabilization might be partially attributable to those side effects rather than solely to preventing biased rounding.

What evidence exists in the paper. The paper provides one detailed arithmetic example and one feature dimension’s cumulative error plot. Figure 13 (Appendix E) shows the temporal correlation between "number of activated SFA condition" (rows with multiple maxima) and the loss explosion, but this is a trigger frequency metric, not a measurement of how much error those triggers actually produce. The paper does not report: per-step histograms of ¯P[T, t] values, per-dimension statistics of V sign distributions, or variance decomposition of δlp - δhp into components attributable to different operations. The analysis in Section 3.3.2 is a case study, not a population-level characterization.

Mitigation status. Not addressed. The paper does not discuss the representativeness of the analyzed addition, does not provide aggregate statistics, and does not measure what fraction of the total δ error is caused by ¯P = 1 additions with negatively-signed V. The validation experiment (Section 4) demonstrates that the diagnosed mechanism is a sufficient cause of the failure, but without quantitative accounting, the paper cannot claim it is the primary or only cause. A practitioner trying to assess whether other stabilization techniques (e.g., stochastic rounding, which eliminates bias in all operations) are necessary or whether the dynamic softmax alone suffices would need this decomposition.

Limitation 3: The Analysis Is Specific to One Model Scale and Architecture, With Minimal Evidence of Generalization

The assumption or constraint. All causal experiments (Sections 3.2–3.3.2) are conducted on GPT-2 Small — a 12-layer, 12-head, 768-dimensional model with approximately 125M parameters — trained on OpenWebText with a context length of 1024. The paper provides two forms of generalization evidence: (a) a GPT-2 Medium variant trained for 100K steps with the stabilized flash attention (Figure 7c), and (b) an observational analysis of a pretrained Llama-3.1-8B checkpoint (Appendix D) showing that structurally similar (PK)[T]⊤X[T] matrices and multiple attention maxima exist in a larger, more modern model. The authors acknowledge this scope limitation: "The generalizability of our findings to other architectures, larger scales, or different low-precision formats like FP8 requires further investigation" (Section 5, Limitations).

The consequence. The paper cannot claim that this specific causal mechanism — ¯P = 1 → significand overflow → biased round-up → low-rank gradient accumulation → spectral norm growth → loss explosion — is the explanation for training instability in large-scale production models (Kimi-Team, 2025; Qwen-Team, 2025) or in FP8 training pipelines. Several factors could differ at scale:

• Model width: GPT-2S has 64-dimensional attention heads. Larger models typically have 128-dimensional heads, which changes the dimensionality of (PK)[T]⊤X[T] and may affect the likelihood of emergent low-rank structure.
• Number of layers: GPT-2S has 12 layers. Deeper models (e.g., 32, 64, 80 layers) may exhibit different patterns of where spectral norms grow — the paper found the failure localized to layer 2, but this could be depth-dependent.
• Training data scale: OpenWebText is a relatively small corpus by modern standards (tens of GB). Models trained on trillions of tokens may encounter the failure condition (¯P = 1 at multiple positions) at different frequencies or in different layers.
• Precision format: The paper’s mechanism depends on BF16’s 7-bit fraction and round-to-nearest behavior. FP8 has different rounding characteristics (only 2–3 fraction bits in E5M2 format, or 7 fraction bits but narrower dynamic range in E4M3), and FP8 training pipelines typically use per-tensor scaling (Perez et al., 2023) that could alter the magnitude of values entering the ¯PV product.
• Architectural differences: Llama-3.1-8B uses rotary position embeddings (RoPE), SwiGLU activations, and RMSNorm — all of which differ from GPT-2’s architecture and could change the structural properties of attention representations.

What evidence exists in the paper. The generalization evidence is preliminary:

• GPT-2M (Figure 7c): Trained for only 100K steps versus 600K for GPT-2S. The exact parameter count and architecture of GPT-2M are not specified (standard GPT-2 Medium is 345M parameters with 24 layers, 16 heads, 1024 embedding dimension — but the paper does not confirm this). A 100K-step run demonstrates that the stabilized flash attention prevents failure at this larger scale, but it does not demonstrate that the same mechanism causes the failure (the paper does not localize the failure to a specific layer and head in GPT-2M, does not show that multiple ¯P = 1 occurrences trigger it, and does not analyze the low-rank structure of gradients).
• Llama-3.1-8B (Appendix D): This is an observational analysis of a pretrained checkpoint, not a training experiment. The paper shows that (PK)[T]⊤X[T] matrices exhibit similar columns (cosine similarity >0.999) across tokens, and that multiple attention maxima occur in various layers and heads. This demonstrates that the preconditions for the failure exist in a large modern model — structural similarity in gradient directions, and the trigger condition of ¯P = 1 — but it does not demonstrate that Llama-3.1-8B would actually fail under BF16 flash attention training. In fact, Llama-3.1-8B was presumably trained successfully, implying that either (a) its training recipe included stabilization techniques that prevented the failure, (b) the mechanism manifests differently at scale (e.g., lower learning rates, different data distributions), or (c) the preconditions exist but some other factor prevents them from combining into catastrophic accumulation. The paper does not investigate which of these explanations holds.

Mitigation status. Partially addressed through acknowledgment (Section 5 Limitations) and the preliminary generalization experiments (Figure 7c, Appendix D). The paper explicitly flags this as future work: "Future work could extend this analysis to FP8 training, larger models, and different architectures" (Section 5). However, the paper does not discuss why the mechanism might or might not transfer — what specific properties of larger models (wider heads, more layers, different data distributions) would amplify or suppress the failure. A practitioner deciding whether to apply the dynamic softmax modification to a 7B or 70B model would need to extrapolate from GPT-2 scale with minimal guidance.

Limitation 4: The Validation Experiment Does Not Distinguish Between Preventing the Trigger and Preventing the Amplification

The assumption or constraint. The paper’s central validation is that modifying the softmax to prevent ¯P = 1 (by dynamically adjusting the normalization factor when multiple maxima are detected) stabilizes training across configurations (Section 4, Figure 7). The interpretation is: preventing ¯P = 1 eliminates the biased rounding trigger, which prevents the vicious cycle of error accumulation → spectral norm growth → sharper softmax → more repeated maxima → more biased rounding. The authors do not experimentally test whether the dynamic softmax also changes other properties of the attention computation that might independently contribute to stability.

The consequence. The dynamic softmax modification might stabilize training through an unintended mechanism rather than (or in addition to) preventing biased rounding. Specifically:

• The modification changes the distribution of ¯P values: when multiple maxima are detected, the softmax temperature effectively decreases (the exponent differences become larger because m = β·rm instead of m = rm). This sharpens the attention distribution — making high-attention positions even higher and low-attention positions even lower. A sharper attention distribution might reduce the effective rank of PK, which could alter the structural similarity of (PK)[T]⊤X[T] matrices (the amplification mechanism) rather than only preventing the trigger (¯P = 1).
• The modification is applied conditionally: the where(rs > 1, ...) check means that rows with a single maximum are processed identically to standard flash attention, while rows with multiple maxima use a different normalization. This introduces an asymmetry in how different rows are computed — rows with concentrated attention (one dominant position) use standard softmax, while rows with diffuse attention (multiple maxima) use a sharpened softmax. This could affect the gradient flow through the attention mechanism in ways that are independent of the biased rounding mechanism. For example, sharpening the softmax in multi-maximum rows might reduce the gradient signal through those rows (because the softmax gradient is small when probabilities are near 0 or 1), which could stabilize training by reducing the magnitude of dS and therefore dWQ — a mechanism that has nothing to do with BF16 rounding.

The paper’s validation therefore demonstrates that the combination of (a) preventing ¯P = 1 and (b) sharpening the softmax when multiple maxima occur stabilizes training. It does not demonstrate that (a) alone would suffice. A minimal test of the mechanism would be: prevent ¯P from being exactly 1 (e.g., by clamping ¯P to max(¯P, 1 - ε)) without changing the softmax temperature, and verify that this also stabilizes training. Without such a test, the causal attribution to the biased rounding mechanism specifically — rather than to the side effects of the dynamic softmax — is plausible but not uniquely identified.

What evidence exists in the paper. The paper provides no ablation that isolates the effect of preventing ¯P = 1 from the effect of changing the softmax temperature. The discussion in Appendix C (Design Considerations) explains why a naive fixed offset is insufficient ("would cause the computed values of ¯P to be consistently rounded in one direction during BF16 conversion, introducing a fixed error"), which motivates the dynamic approach, but this does not address whether the dynamic approach works through the hypothesized mechanism or through a different pathway. The paper also explains why an unconditional adjustment (always using β·rm) would risk underflow, motivating the conditional application — but again, this is about designing a safe intervention, not about verifying that the intervention works for the diagnosed reason.

Mitigation status. Not addressed. The paper does not discuss the possibility that the dynamic softmax stabilizes training through a mechanism other than preventing biased rounding, nor does it propose control experiments to rule out alternative explanations. A practitioner seeking the simplest possible fix might wonder: could I just clamp ¯P to < 1 everywhere without the conditional logic? Or reduce the learning rate? Or use gradient clipping? The paper does not isolate the necessary and sufficient components of its intervention, making it difficult to determine the minimal change required to prevent the failure.

Limitation 5: Upstream Cost of the Dynamic Softmax Modification Is Not Characterized

The assumption or constraint. The stabilized flash attention (Algorithm 3) introduces additional operations into the inner loop of the forward pass: computing rs = rowsum(rm - S ≤ ε) (a reduction over the key dimension to count repeated maxima), two where conditional operations, and a potential adjustment to the normalization factor m. These operations are executed for every query block and every key/value block — meaning they scale with the number of tiles in the tiled flash attention algorithm. The paper does not measure or report the computational overhead of these additional operations relative to the standard flash attention forward pass.

The consequence. The paper presents the dynamic softmax as a "practical solution" (abstract) and a "minimal modification" (Section 4), but without throughput or memory measurements, a practitioner cannot assess whether the modification is actually practical for production training. The overhead could be non-trivial for several reasons:

• The rowsum(rm - S ≤ ε) operation is a broadcast comparison followed by a reduction, which may not fuse efficiently with the matrix multiplications and exponentials that dominate the flash attention kernel. Flash attention’s performance comes from careful kernel fusion — introducing a reduction that is not part of the standard attention computation could require additional synchronization points or register pressure.
• The conditional logic (where(rm > 0 ∧ rs > 1, β·rm, rm)) introduces branching in the kernel, which could reduce GPU utilization if different rows within a thread block take different paths.
• The frequency of the condition being activated matters: if rs > 1 is rare (occurring in only a few rows per sequence), the overhead of checking the condition (which must be done for every row) might be small. But Figure 13 (Appendix E) shows that the number of activated conditions spikes to ~2×10⁹ (cumulative) around the failure point — if this represents a high per-step frequency in the failing regime, the modification might have higher overhead precisely when it is most needed.

Beyond throughput, the modification also changes the numerical behavior of attention in exact arithmetic whenever multiple maxima are detected, because it changes the normalization constant m. While this is mathematically equivalent to standard attention in exact arithmetic (softmax is shift-invariant), it means that the output ¯P is different from what standard flash attention would compute — the values are strictly less than 1 instead of possibly being exactly 1. This could affect downstream computations (e.g., the value aggregation in ¯PV, the normalization O = ¯O/rowsum(¯P)) in ways that change the model’s learning dynamics even in the absence of rounding errors. The paper’s validation shows that training converges (Figure 7), but it does not compare the final model quality (e.g., validation perplexity at convergence, downstream task performance) between the stabilized flash attention and high-precision (FP32) flash attention. A practitioner would want to know: does the modification change the solution the model converges to, or only prevent the failure without affecting the final model?

What evidence exists in the paper. None. The paper provides no throughput measurements, no profiling of the modified kernel, no comparison of wall-clock training time, and no comparison of final validation loss between stabilized and high-precision flash attention at convergence. Figure 7 shows that the stabilized version converges (loss decreases), but the curves for the original BF16 configuration are truncated by the loss explosion, so it’s impossible to compare final converged loss between the stabilized and stable (FP32) configurations. If the stabilized version converges to a slightly higher loss than FP32 attention, the modification might trade training stability for model quality — a tradeoff the paper does not characterize.

Mitigation status. Not addressed. The paper does not discuss computational overhead, kernel fusion compatibility, or numerical quality impacts as limitations. The "Limitations" paragraph in Section 5 mentions that "our proposed mitigation is tailored to the specific rounding error identified and may not address other sources of numerical instability" but does not discuss the practical cost of deploying the mitigation. This is a significant gap for a paper that positions its modification as enabling "practical" low-precision training.

Limitation 6: The Low-Rank Coherence Mechanism Is Postulated but Not Experimentally Manipulated

The assumption or constraint. Claim 2 (Section 3.3.1) asserts that the training failure requires both biased scalar coefficients (δlp - δhp)[T] and structurally similar rank-1 matrices (PK)[T]⊤X[T] — the former provides the systematic push, the latter provides the consistent direction that allows errors to accumulate rather than cancel. The evidence for structural similarity is Figure 4 (qualitative heatmap visualization in GPT-2) and Appendix D (cosine similarities >0.999 in Llama-3.1-8B). The paper does not experimentally manipulate the structural similarity — for instance, by adding noise to break the coherence, or by amplifying the similarity to accelerate the failure — to test whether this factor is truly necessary for the failure.

The consequence. Without experimental manipulation, the claim that low-rank coherence is a necessary condition for the failure rests on correlational evidence (similar-looking heatmaps, high cosine similarity) and the logical argument that errors would cancel without coherent structure. But several alternative hypotheses are not ruled out:

• The bias alone might be sufficient: if Σ(δlp - δhp)[T] is large enough (consistently positive across many steps), the weight update error α · Σ(δlp-δhp)[T] · R could be harmful even if R varies somewhat across steps — as long as R doesn’t vary so much that the dot product R_step1 · R_step2 is negative on average. The paper doesn’t measure the cross-step similarity of R, only cross-token similarity within a step.
• The low-rank structure might be a consequence, not a cause: as the weight spectral norm grows (due to some other instability), the attention mechanism might naturally produce more structured (PK)[T]⊤X[T] matrices — the low-rankness could be a symptom of the developing pathology rather than a necessary condition for it.
• The failure might occur through a different accumulation mechanism entirely: for example, the biased δ error might directly increase the magnitude of dS and therefore dQ, without requiring structural similarity — each step simply adds a large error in whatever direction (PK)[T]⊤X[T] happens to point, and the accumulation is driven by the magnitude of the error rather than its direction.

The paper’s discussion (Section 5) proposes that QK normalization and Gated Attention work by disrupting the low-rank structure of (PK)[T]⊤X[T], but this is a post-hoc interpretation, not an experimentally tested mechanism. The paper does not measure whether QK normalization actually reduces the structural similarity of these matrices, or whether the dynamic softmax (which targets the trigger, not the structure) also incidentally affects the structure. Without manipulation experiments, the role of low-rank coherence remains a plausible and elegant hypothesis rather than an empirically verified necessary condition.

What evidence exists in the paper. The evidence is observational: Figure 4 shows similar-looking heatmaps; Appendix D shows high cosine similarities in Llama-3.1-8B; Figure 5a shows monotonic accumulation of (δlp - δhp)[T], which is consistent with coherent error directions but does not rule out alternative explanations (e.g., the accumulation is driven by bias magnitude alone, and the observed structural similarity is incidental). The paper decomposes the gradient error mathematically (Equation 2) into a sum of rank-1 terms weighted by scalars, but this decomposition is an algebraic identity — it doesn’t itself demonstrate that structural similarity is necessary for accumulation.

Mitigation status. Partially addressed through the Llama-3.1-8B analysis (Appendix D), which provides quantitative similarity measurements (cosine similarity) rather than only qualitative visualizations — strengthening the evidence that this structural property exists beyond GPT-2. But the paper does not experimentally test necessity: it does not show that models without structural similarity are immune to the failure, does not show that artificially disrupting the similarity (while keeping the rounding bias) prevents the failure, and does not show that amplifying the similarity (while keeping the rounding bias constant) accelerates the failure. The mechanistic interpretation of QK normalization and Gated Attention as structural disruptors is offered as an explanation for their empirical effectiveness (Section 5 Discussion) but is not verified in this paper’s experimental setup — the authors did not train models with QK normalization and measure the structural similarity of their gradient components. A practitioner evaluating whether to deploy QK normalization, Gated Attention, or the dynamic softmax would want to know whether these interventions are interchangeable (addressing the same root cause through different mechanisms) or complementary (addressing different nodes in the causal chain), and the paper does not provide this guidance.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper shifts the conversation around low-precision training stability from empirical patch management to mechanistic causal understanding. Before this work, the field's approach to numerical instabilities in transformer training followed a well-worn pattern: observe a loss spike, try stabilization techniques (QK normalization, gradient clipping, higher precision for suspect operations), keep what works, and move on. The nanoGPT issues that motivated this paper (#303, #524, #554) exemplify this — two years of community discussion produced workarounds (use FP32, use standard attention) but no explanation of why BF16 + flash attention fails. This paper demonstrates that such failures are not mysterious acts of numerical chaos but deterministic consequences of specific interactions between attention's representational structure and BF16's rounding behavior — and that these interactions can be diagnosed systematically, validated with minimal interventions, and explained through a complete causal chain from individual bit operations to loss explosion.

This is a methodological shift rather than a paradigm shift in architecture or training algorithms. The paper does not propose a new way to train transformers; it proposes a new way to think about training failures. Its lasting contribution may be less the specific mechanism it identifies (biased rounding at ¯P = 1, amplified by low-rank gradient coherence) and more the analytical workflow it demonstrates: isolate the error source through systematic high-precision substitution → decompose gradient errors into coefficient and direction components to identify accumulation mechanisms → trace to root arithmetic cause through bit-level analysis → validate with a minimal, targeted intervention. This workflow (codified in Appendix F as a "blueprint") transforms numerical debugging from an art practiced by a few experts into a structured methodology that can be taught, replicated, and applied to new failure cases.

The paper resolves several previously puzzling contradictions in the stability literature:

• Why does BF16 training sometimes fail and sometimes not? The paper shows that failure requires a specific alignment of factors: multiple identical row maxima in pre-softmax scores (producing ¯P = 1), predominantly same-sign values in the corresponding V columns (creating conditions for systematic round-up), AND structurally similar (PK)[T]⊤X[T] matrices (providing a coherent error accumulation direction). Lee et al. (2024) reported that roughly 10% of GPT-2 BF16 runs diverge — this paper's mechanism explains why it's stochastic: the model must develop the specific combination of low-rank structure and attention-score distributions that trigger biased rounding, which depends on the random seed, data order, and learning dynamics. The 90% of runs that don't diverge likely don't develop this specific combination, even though they still use BF16 and flash attention.

• Why do stabilization techniques like QK normalization and Gated Attention work? The paper provides a mechanistic interpretation (Section 5 Discussion) that unifies these diverse interventions: they disrupt the structural coherence of (PK)[T]⊤X[T] matrices, preventing rounding errors from accumulating along a consistent low-rank direction. QK normalization constrains the magnitude and direction of query/key vectors, reducing the emergence of structured representations. Gated Attention introduces non-linearity that breaks coherence. QK-clipping reduces the probability of ¯P = 1 by bounding extreme attention scores. These techniques were developed independently through trial and error; the paper's framework explains why they converge on stable training — they interrupt the same causal chain at different nodes.

• Why does the failure localize to specific layers and heads? The finding that a single attention head (head 8, layer 2) is the failure origin (Section 3.2) was previously inexplicable — why would one 64-dimensional subspace out of 144 attention heads be the sole locus of catastrophic instability? The paper's mechanism explains this: that specific head happened to develop the combination of low-rank gradient structure and attention-score distributions (multiple repeated maxima) that triggers the biased rounding → coherent accumulation → spectral norm growth feedback loop. Other heads might have one factor but not the other, or might have both but at sub-critical magnitudes.

This work redirects research attention in several ways:

• Away from "more bits" as the default solution. The paper's mechanism shows that the failure is not caused by BF16 having insufficient precision in general — most of the model trains fine in BF16 — but by a specific numerical pathology triggered under specific conditions. This suggests that the path to robust low-precision training is not simply FP32 everywhere (which eliminates the efficiency gains) but targeted interventions that prevent the pathological conditions while preserving the benefits of low precision.

• Toward verifier/accumulation decomposition as a diagnostic lens. The mathematical decomposition of gradient errors into scalar coefficients and rank-1 direction matrices (Equation 2: dWQ_hp - dWQ_lp = α Σ(δ_lp - δ_hp)[T] · (PK)[T]⊤X[T]) separates "how much error" from "in what direction." This decomposition is what enables the diagnosis that both bias in the coefficients AND structural similarity in the directions are necessary for catastrophic accumulation. Future investigations of other numerical instabilities can adopt this lens: when training fails, decompose the gradient error, check whether the scalar coefficients are biased, and check whether the direction matrices are coherent across tokens and steps.

• Toward preventive monitoring rather than reactive patching. Figure 13 (Appendix E) shows that the number of rows with multiple identical maxima begins increasing before the loss explodes. This suggests that the trigger condition (¯P = 1 at multiple positions) is a leading indicator of impending failure — it could be monitored during training and used to trigger interventions (dynamic softmax adjustment, precision increase for affected heads, learning rate reduction) before catastrophic divergence occurs. The paper's mechanism provides a principled basis for such monitoring: rather than watching for generic symptoms (large gradients, loss spikes), watch for the specific preconditions that the causal chain identifies.

The paper also shifts the perceived tractability of numerical stability research. Before this work, analyzing why a 125M-parameter model training on billions of tokens fails after thousands of steps might have seemed hopelessly complex — too many interacting components, too much stochasticity, too many potential failure points. The paper demonstrates that with aggressive localization (to a single head), deterministic data replay, and structural decomposition of gradient errors, the complexity is manageable. This lowers the barrier to entry for other researchers to conduct similar mechanistic investigations, potentially accelerating progress on other long-standing stability puzzles (FP8 training divergence, loss spikes in large-scale models, instability in mixture-of-experts architectures).

Follow-Up Research This Work Enables

1. Quantitative decomposition of the δ error to measure the biased rounding mechanism's contribution. The paper identifies the biased rounding mechanism (¯P = 1 + negative V → significand overflow → round-up) as the root cause and validates that preventing ¯P = 1 stabilizes training, but does not measure what fraction of the total δ_lp - δ_hp error is attributable to this mechanism versus other rounding errors. A direct follow-up would instrument the flash attention backward pass to track, for each training step: (a) the total δ_lp - δ_hp error magnitude, (b) the number of token positions where ¯P[T, t] = 1 and V[t, i] shares sign with the error direction, (c) the error contribution from those positions versus from positions where ¯P < 1, and (d) the distribution of per-addition rounding errors across the ¯PV summation. If the ¯P = 1 mechanism accounts for >80% of the total δ error, the case is strong that it is the dominant failure pathway. If it accounts for <20%, then other rounding error sources (e.g., from the normalization step, from the exponentiation in softmax) are equally or more important, and the dynamic softmax may be stabilizing training partially through side effects rather than solely through preventing biased rounding. This experiment would use the same deterministic replay setup as the paper but with additional logging; it would not require training new models from scratch, only instrumenting the existing failing trajectory.

2. Direct experimental manipulation of low-rank gradient coherence to test necessity. Claim 2 (Section 3.3.1) asserts that structural similarity of (PK)[T]⊤X[T] matrices is necessary for the failure — without it, biased coefficients would produce weight updates in random directions that cancel out. But this claim is based on observational evidence (Figure 4 heatmaps, Appendix D cosine similarities) and logical argument, not experimental manipulation. A strong follow-up would test necessity by artificially decorrelating the error directions: during training of the failing configuration, add independent Gaussian noise to the PK or X matrices (or their product) in the backward pass, scaled to preserve the per-step gradient magnitude but randomize the direction of (PK)[T]⊤X[T] across tokens and steps. If the failure is prevented (or delayed substantially) under decorrelated directions despite identical biased coefficients, the necessity of low-rank coherence is confirmed. A complementary experiment would amplify coherence: replace (PK)[T]⊤X[T] with its rank-1 approximation (SVD truncated to the dominant singular vector) during the backward pass, making all error directions perfectly aligned. If the failure accelerates (occurs in fewer training steps), the sufficiency of coherence (combined with biased coefficients) is confirmed. These experiments would modify only the gradient computation for WQ in the affected head, not the forward pass, providing a clean test of Claim 2 independent of the trigger mechanism.

3. Replication of the full causal analysis on an FP8 training failure. The paper's analytical workflow (Appendix F) is positioned as a general diagnostic methodology, but it has only been demonstrated on one specific failure (GPT-2 + BF16 + flash attention). A critical test of the methodology's generality would be to apply it to a known FP8 training failure — for example, the divergence reported by Lee et al. (2024) when training transformers in FP8 without per-tensor scaling, or the instability observed by Fishman et al. (2024) in trillion-token FP8 training runs with SwiGLU activations. The FP8 format (E4M3 or E5M2) has different rounding characteristics than BF16: E4M3 has 4 exponent bits and 3 fraction bits (narrower dynamic range, similar precision to BF16's 7 fraction bits but different rounding thresholds), while E5M2 has 5 exponent bits and 2 fraction bits (wider dynamic range but very low precision). The biased rounding mechanism in this paper (significand overflow → right shift → round-up with sticky bit) may or may not transfer — FP8 addition could exhibit different bias patterns, or the dynamic range limitations could introduce underflow/overflow as dominant error sources before precision-based bias becomes relevant. A successful replication would: (1) identify whether FP8 training failures localize to specific layers/heads (as BF16 does) or are more diffuse; (2) determine whether the root cause is biased rounding (precision), dynamic range (exponent), or an interaction; (3) validate with a targeted intervention. A negative result — the FP8 failure has a completely different mechanism — would still be valuable, as it would map the boundary conditions of the paper's specific mechanism and demonstrate that "low-precision instability" is not a single phenomenon.

4. Training stability monitoring using multiple-maxima frequency as an early warning signal. Figure 13 (Appendix E) shows that the frequency of rows with multiple identical maxima (the "activated SFA condition") increases sharply before the loss explosion. This suggests that rs > 1 (from the dynamic softmax condition check) could serve as a leading indicator for impending training instability. A practical follow-up would instrument a training run to track, at each optimizer step: (a) the number of attention rows with rs > 1 (multiple maxima within tolerance ε), (b) the spectral norms of WQ per head, (c) the maximum |(δ_lp - δ_hp)[T]| value, and (d) the loss on a held-out validation set. The experiment would run multiple training configurations (different random seeds, different learning rates, different model scales) and measure whether the multiple-maxima frequency consistently rises before other indicators (spectral norm growth, loss increase). If the temporal precedence is robust, this provides a cheap, online diagnostic — the rs computation is already part of the dynamic softmax and could be logged without overhead. A practical system could trigger automatic interventions (switching to FP32 for affected heads, reducing learning rate, applying the dynamic softmax only when rs > 1 becomes frequent) when the frequency exceeds a threshold. The experiment would need to establish: (1) the false-positive rate (how often rs > 1 rises without subsequent failure), (2) the lead time (how many steps before loss explosion the indicator rises), and (3) whether the indicator generalizes to models and datasets beyond GPT-2/OpenWebText.

5. Systematic comparison of stabilization techniques through the lens of the causal chain. The paper's Discussion (Section 5) proposes that QK normalization, QK-clipping, Gated Attention, and the dynamic softmax all address the same root cause but at different nodes in the causal chain: the dynamic softmax prevents the trigger (¯P = 1), QK-clipping reduces the probability of the trigger, and QK normalization / Gated Attention disrupt the amplification mechanism (low-rank coherence of (PK)[T]⊤X[T]). This is a testable hypothesis that the paper does not experimentally verify. A follow-up study would train GPT-2 (or a larger model) with each stabilization technique under the paper's deterministic replay setup and measure, for the problematic head (head 8, layer 2): (a) the frequency of ¯P = 1 occurrences, (b) the structural similarity (cosine similarity or subspace angle) of (PK)[T]⊤X[T] across tokens and steps, (c) the bias in (δ_lp - δ_hp)[T], and (d) the growth of WQ spectral norm. The predictions are: dynamic softmax reduces (a) but not (b); QK normalization reduces (b) but not necessarily (a); QK-clipping reduces (a); Gated Attention reduces (b). If these predictions hold, the causal model is validated and practitioners can choose interventions based on which node is most problematic for their specific setup. If the predictions don't hold — for example, if QK normalization also reduces ¯P = 1 frequency, or if the dynamic softmax also reduces structural similarity — then the interventions have overlapping effects that complicate the causal interpretation. The experiment would also measure whether combining interventions (e.g., dynamic softmax + QK normalization) provides benefits beyond either alone, or whether they are redundant.

6. Cross-architecture stress test: does the mechanism generalize to vision transformers, state-space models, or encoder-decoder architectures? The paper's analysis is specific to decoder-only transformer architectures (GPT-2, and observationally Llama-3.1-8B) using standard scaled dot-product attention. Several architectural variants could be stress-tested to map the boundary conditions of the mechanism:

• Vision Transformers (ViT): Do image patches produce similar low-rank (PK)[T]⊤X[T] structure and multiple attention maxima? If not, ViT training in BF16 + flash attention might be inherently stable, suggesting the mechanism is specific to language modeling's token distributions.
• Cross-attention in encoder-decoder models: The mechanism depends on ¯P = 1 at multiple positions in self-attention. In cross-attention, where queries come from the decoder and keys/values from the encoder, the attention patterns may be fundamentally different — less prone to attention sinks or repeated maxima.
• State-space models (Mamba, etc.): These replace attention entirely with structured state-space recurrences, eliminating the ¯PV product that is the locus of biased rounding. If SSM training in BF16 is stable without special interventions, it supports the paper's mechanism as attention-specific.
• Linear attention approximations: Methods that avoid computing the full softmax (e.g., Performer, Linformer) would not produce ¯P = 1 through the same pathway. Testing whether they exhibit similar BF16 instabilities would clarify whether the mechanism depends specifically on the softmax's argmax behavior or on more general properties of attention. Each stress test would replicate the paper's deterministic replay setup (training a small model to failure, localizing the error source, analyzing gradient coherence) on the target architecture. A negative result for some architecture would not invalidate the paper's mechanism for GPT-2 but would define its scope, which is valuable for practitioners deciding whether to apply the dynamic softmax or related interventions to non-standard architectures.

Practical Applications and Downstream Use Cases

1. Debugging and stabilizing custom low-precision training pipelines. The paper's analytical workflow gives practitioners a concrete procedure for diagnosing numerical instabilities in their own training setups. When a training run diverges in BF16 (or FP8), rather than guessing which operation to upcast or which stabilization technique to apply, an engineer can: (1) use deterministic data replay to make the failure reproducible, (2) systematically replace operations with FP32 equivalents (starting from the full attention module and narrowing to specific computations like δ = rowsum(dO ◦ O) and ¯PV), (3) analyze the gradient error for bias and structural coherence by instrumenting the backward pass, and (4) target the root cause with a minimal intervention. The paper's finding that computing only ¯PV in FP32 is sufficient to stabilize training (Section 3.2) provides a concrete, low-overhead fix for the specific failure: in frameworks that support mixed-precision at the operation level, practitioners can keep most of flash attention in BF16 while upcasting only the ¯PV multiplication to FP32, avoiding the precision cost of full FP32 attention or the implementation complexity of the dynamic softmax. This is immediately actionable for anyone encountering the nanoGPT-style loss explosion — it requires no architectural changes, no hyperparameter tuning, and can be implemented as a one-line modification to the flash attention kernel (cast ¯P and V to FP32 before the matmul, cast the result back to BF16).

2. Training-time stability monitoring in large-scale production runs. The correlation between multiple-maxima frequency and impending loss explosion (Figure 13, Appendix E) enables a practical early warning system for training instability. During large-scale pretraining (where a single run costs millions of dollars and early detection of divergence can save substantial compute), practitioners can: (1) log the number of attention rows with rs > 1 (the same condition checked by the dynamic softmax) at each optimizer step, computed as a lightweight byproduct of the forward pass, (2) track this metric alongside standard diagnostics (gradient norm, weight spectral norm, validation loss), and (3) trigger automatic interventions — such as temporarily increasing precision for affected layers, reducing the learning rate, or activating the dynamic softmax — when the frequency exceeds a threshold or its rate of change accelerates. The paper's Figure 13 shows the frequency rising before loss divergence (~step 7,000), providing lead time for intervention. The cost of monitoring is negligible compared to the cost of a failed training run: the rs computation is a row-wise reduction over the key dimension that can be fused into the existing flash attention kernel or computed as a separate lightweight pass. The specific threshold and lead time would need to be calibrated per model and dataset, but the paper provides the conceptual foundation — watch for the trigger condition, not just the symptoms.

3. Guiding precision allocation in mixed-precision training recipes. The paper's finding that only a few attention heads (1, 7, 8, 9, 11, 12 in layer 2) require high-precision O computation to stabilize training (Section 3.2) suggests a precision-heterogeneous training strategy: rather than using uniform precision across all operations (e.g., "everything in BF16" or "attention in FP32"), allocate higher precision only to the small fraction of computations that are numerically vulnerable. In the GPT-2S failure case, computing O in FP32 for 6 out of 144 attention heads (4.2% of heads) or computing only ¯PV in FP32 for the entire model is sufficient to prevent divergence. This is far cheaper than full FP32 attention and preserves most of BF16's memory and speed benefits. A practitioner training a model at scale could: (1) run a short diagnostic phase at the start of training (or use a smaller proxy model) to identify which layers and heads develop elevated spectral norms or multiple attention maxima, (2) selectively upcast those components to higher precision for the remainder of training, and (3) dynamically adjust the precision allocation if other heads become vulnerable later in training. The paper's mechanism provides the theoretical justification for this strategy — the vulnerability is localized because the failure requires a specific alignment of structural and numerical factors that only occurs in a subset of heads — and the isolation experiments provide a template for how to identify the vulnerable components.

4. Informing hardware and library design for low-precision training. The paper's bit-level analysis of BF16 rounding bias (Section 3.3.2) identifies a specific failure mode — significand overflow during addition of same-sign numbers with identical exponents, combined with sticky bit activation forcing round-up — that could be addressed at the hardware or numerical library level. Two concrete implications:

• Stochastic rounding for accumulation operations: The biased rounding occurs because deterministic round-to-nearest with sticky bit systematically rounds up when two same-sign numbers overflow the significand. Stochastic rounding (Ben Ali et al., 2024) — which rounds probabilistically based on the truncated bits rather than deterministically — would eliminate this systematic bias: an addition that would deterministically round up (due to rounding bit = 1 and sticky bit set) would instead round up with probability proportional to the truncated value and round down otherwise. Over many additions, the expected error is zero, preventing the systematic negative bias in ¯O. If hardware supporting stochastic rounding in matrix multiplication accumulators becomes available, the paper's mechanism would be naturally mitigated without any software changes.
• Mixed-precision accumulators for reductions: The paper shows that the ¯PV product is the critical operation — the accumulator (the running sum over token positions) needs higher precision than the operands to avoid biased rounding. A library-level optimization could compute ¯PV using an FP32 accumulator internally (as PyTorch's matmul in BF16 already does for the dot product) but round the final per-element result to BF16. The paper's mechanism suggests that even this single rounding at the end (rather than at each addition) would drastically reduce the error, because the biased round-up only occurs when the accumulator's significand overflows relative to BF16's 7-bit fraction — a much rarer event after summing many small contributions than after each individual addition.

When to Prefer This Method

(This section is not applicable. The paper does not position the dynamic softmax modification as a recommended production technique to be chosen over named alternatives — it is explicitly a validation experiment. The paper's Discussion (Section 5) notes that QK normalization and Gated Attention are existing techniques that address the same root cause through different mechanisms, but does not provide experimental comparisons or decision rules for choosing among them. The dynamic softmax has not been benchmarked against QK normalization on throughput, final model quality, or stability at scale. A forced "Prefer X when Y" matrix would be fabricating tradeoffs the paper does not establish.)