# One QK Channel, Many Sources: Guarding Low-Precision Attention Collapse

> Low-precision transformer collapse converges on a single query-key channel regardless of error source, and per-head QK normalization at that locus contains the runaway.

- **Source:** [arXiv](https://arxiv.org/abs/2608.02091)
- **Published:** 2026-08-29
- **Permalink:** https://picx.dev/p/RhCb1j
- **Whiteboard:** https://picx.dev/p/RhCb1j/image

## Summary

# Summary of "One QK Channel, Many Sources: Guarding Low-Precision Attention Collapse"

## Summary (Overview)

- **Core finding**: The paper establishes a **source-channel dissociation** in low-precision transformer training collapse: heterogeneous error sources (e.g., bf16 streaming-softmax accumulation, weight-gradient quantization) converge on a single query–key (QK) failure channel, and correcting only the QK route blocks the collapse while the original fault remains active.

- **Entry mechanism identified**: Temporal sign-coherence—the persistence of each attention head's error sign across optimizer steps—is what gates error entry into the QK channel, not aggregate deviation magnitude or per-tile scalar means.

- **Causal validation**: Projecting each optimizer step's update off the current QK weights' three leading singular directions suppresses the spectral runaway ($\hat{\sigma}(W^Q)$ held to 11.1), while removing equal energy elsewhere leaves it at 237, proving the QK channel drives the collapse rather than merely tracking it.

- **Practical remedy**: **QK-Guard**, a dormant controller that monitors attention-logit saturation (threshold 30) and switches on parameter-free per-head QK normalization at the QK locus, contains all tested runaways across architectures, scales, and fault sources, matching always-on QK-norm over 60k-step horizons.

- **Scope boundary**: The dissociation holds for temporally sign-coherent low-precision error perturbing residual-block parameter updates or entering upstream of those blocks in pre-norm transformer setups; a fault acting after the final attention module would escape.

---

## Introduction and Theoretical Foundation

### Background and Motivation

Bfloat16 transformer training is routine, yet runs can appear healthy for thousands of steps before abruptly collapsing. Prior work (Qiu and Yao, 2024) linked this pattern to **structurally aligned attention errors** whose compounding rounding contributions destabilize training, and identified query–key (QK) normalization as the disruptor of that compounding.

### The Open Question

The central unresolved question: **Do low-precision faults require repair at their sources, or do they share a blockable failure channel?** Collapse alone cannot answer this—it reveals the endpoint, not what the failure depended on.

### Theoretical Foundation

The paper builds on several established results:
- Low-precision accumulation is known to destabilize training (Wang et al., 2024), though classical analysis expected shorter tiled reductions to be safer (Dao et al., 2022)
- Four-bit attention training requires high-precision output solely for the backward $D$-term (Gao et al., 2024)
- Query–key gradients exhibit low-rank structure (Zhao et al., 2024)
- Temporally correlated gradient estimates destabilize training (Sutskever et al., 2013)

### Key Definitions

The **source-channel dissociation** is the paper's thesis: *fault source is not failure channel*. Multiple distinct error generators can feed a single dominant failure pathway (the QK channel), meaning intervention should target the channel rather than each source individually.

---

## Methodology

### 1. Fault Isolation via Controlled Reproduction

The authors reproduced GPT-2 small collapse on OpenWebText using a PyTorch re-implementation of FlashAttention-2's tiled online-softmax recurrence. The key carries are:
- **Output carry** $O$: reaches $\bar{d}Q$ and $dK$, never $dV$, through the per-row $D$-term
- **Statistics carry** (running max and exponential sum): reaches $dQ$, $dK$, and $dV$ through the saved log normalizer

### 2. Manipulable Injection Assay

The output carry yields an exactly computable injection: the bf16-minus-fp32 saved-output difference enters through the per-row $D$-term, while an fp32-accumulation forward pass keeps the residual stream clean. The error is movable along **four assay axes**:
- **Fault source**: what generates the error and where
- **Placement**: emulation sub-choice within a source
- **Dose**: scaling factor
- **Structure**: error pattern

### 3. Route Correction and Source Shift

- **Route correction**: Each step runs the same weights/batch backward through both faulty and fp32-accumulation reference, splicing the clean gradient into one named route (QK, MLP, V/O, or whole attention) while leaving others faulty
- **Source shift**: Replaces bf16 accumulation with post-reduction weight-gradient quantization (MXFP4-style: 4-bit elements, blocks of 32, shared power-of-two scale) across different module targets

### 4. Error-Structure Interventions

To test what gates entry, the paper manipulates single properties of the injected backward error:
- Sign scrambling (element-wise)
- Batch-axis permutation
- Gaussian noise replacement
- Scalar mean injection/removal
- Per-head DC bias injection/removal
- Whole-tensor vs. per-head sign randomization

### 5. Direction-Specific Removal (Causal Probe)

At each optimizer step, project the realized update of the current $W^Q$ and $W^K$ off their three leading singular directions, then rescale the remainder back to the update's original norm. Control: remove equal energy from outside that subspace.

### 6. QK-Guard Design

A dormant closed-loop controller:
- **Monitor**: attention-logit saturation (maximum logit exceeding threshold 30 at periodic readings)
- **Action**: permanently switch on parameter-free per-head QK normalization at the QK locus

---

## Empirical Validation / Results

### Fault Localization (Table 1)

| Condition | Collapsed | Step |
|---|---|---|
| bf16 O + bf16 stats | 6/6 | ×700–1400 |
| accumulator → fp32 | 0/2 | ▶4900 |
| SDPA, bf16 operands | 0/2 | ▶4900 |
| FlashAttention-2 forced | 0/2 | ▶4900 |
| Full matrix, no tiling | 0/2 | ▶4900 |

**Carry decomposition**: Both bf16/fp32 and fp32/bf16 carry combinations collapse (2/2 each), while fp32/fp32 survives (0/2). The fault localizes to the bf16 streaming-accumulator artifact, not bf16 operands or tiling itself.

### Route Correction Across Architectures (Table 2)

Cleaning only the query–key route rescues every tested architecture (plain GPT-2, RoPE variant, 350M, LLaMA-style) to healthy train loss (3.23–3.88); MLP cleaning rescues none. On LLaMA-style, 4/4 untreated collapses vs. 0/4 with per-head QK normalization.

### Source Shift Results

- Weight-gradient quantizer on QK gradient: runaway in $\hat{\sigma}(W^Q)$ to 292 and 205 (healthy ≈ 6); QK-norm suppresses to 24 and 26
- Same quantizer on MLP/value gradients: 4/4 unprotected collapses, 0/4 with QK-norm
- GEMM operand quantization (fp32 accumulation): healthy through 11900 steps—4-bit arithmetic alone is not the trigger

### Entry Condition (Table 3)

| Treatment | Collapsed | Step |
|---|---|---|
| True error | 5/5 | ×1000–2900 |
| Batch-permuted | 3/3 | ×1000–1300 |
| Sign scrambled | 0/4 | ▶7900, 19900 |
| Gaussian noise | 0/2 | ▶11900 |

Per-tile scalar mean injection is benign (0/2), while error-minus-mean still collapses (2/2). Per-head DC bias is the carrier: injecting it collapses 6/6; removing it protects only in the clean-forward surrogate.

### Causal Probe Results

- **On-target removal** (projecting off leading singular directions): all 4 treated arms hold $\hat{\sigma}(W^Q)$ between 11.1 and 15.4, none collapse through 15900 steps
- **Equal-energy of-target control**: $\hat{\sigma}(W^Q)$ runs to 237–491; model-wide arms collapse at 8100 and 15600
- **Mere energy reduction** (no direction targeting): all 4 arms collapse between 9000 and 11600 steps

### QK-Guard Performance

- **Plain GPT-2**: All 6 guarded arms fire before collapse and finish healthy; unguarded seeds end at 10.61 and 9.86 train loss
- **RoPE transfer**: 6/6 pairs fire between 6100–13400 steps; 6/6 unguarded twins collapse, 0/6 guarded arms collapse through 15900 steps
- **350M scale**: Same threshold transfers; final losses 4.04 (QK-norm) vs. 4.03 (fp32 hot-swap) vs. 8.28 (unguarded)
- **Held-out fault** (weight-gradient quantization): Both seeds fire at 4600 and 5600, finishing within +0.006 and +0.012 nats of always-on QK-norm

### Durability (60k-step horizon)

| Method | Validation Loss (3 seeds) |
|---|---|
| Always-on QK-norm | 3.13, 3.09, 3.09 |
| QK-Guard | 3.13, 3.10, 3.09 |
| Controls | collapse at 7200, 8600, 7400 |

Peak $\sigma(W^Q)$ for survivors: 46.5—far exceeding values where unguarded runs collapse, showing durability rests on loss, not spectral safety.

---

## Theoretical and Practical Implications

### Theoretical Implications

1. **Source-channel dissociation reframes failure analysis**: The finding that heterogeneous error sources converge on one QK channel means collapse diagnostics must distinguish between where errors originate and where they propagate. A single dominant attractor (QK channel) can amplify diverse low-precision artifacts.

2. **Temporal sign-coherence as the gating mechanism**: The identification of cross-step sign persistence as the entry condition separates pathwise properties from per-step statistics. This connects to classical results on temporally correlated gradient estimates destabilizing training, but localizes the mechanism to attention's QK pathway.

3. **The QK channel as causal driver**: Direction-specific removal proves the channel's dominant singular directions actively drive the runaway, not merely track it. This is a causal claim about the dynamics, not a correlational one.

4. **Boundary conditions**: The dissociation holds within a defined scope (pre-norm transformers, residual-block parameter updates or upstream entry). A fault after the final attention module would escape—this falsifiable bound strengthens the theoretical claim.

### Practical Implications

1. **Intervention locus over source repair**: The paper's central practical message: act at the QK locus rather than repairing each source separately. This is particularly important when low precision is the design point (e.g., 4-bit quantization) rather than an artifact.

2. **QK-Guard as a deployable remedy**: The dormant controller adds negligible cost (parameter-free normalization, only activated at threshold crossing) and matches always-on QK-norm quality over 60k steps. Its trigger (attention-logit saturation) is monitorable without per-step spectral decompositions.

3. **Limitations of source-local fixes**: Fp32 accumulation fixes the streaming-softmax fault but not weight-gradient quantization or external errors. The paper's evidence suggests a more general remedy at the QK channel.

4. **Cost-free containment**: No detectable quality cost (within +0.02 nats of always-on QK-norm reference), and the action locus (not reactivity alone) determines rescue success.

---

## Conclusion

### Main Takeaways

1. **Fault source is not failure channel**: Multiple low-precision error generators—streaming-softmax accumulation, weight-gradient quantization, external injections—converge on a single QK spectral runaway channel that one action (per-head QK normalization) blocks while each generator remains active.

2. **Temporal sign-coherence gates entry**: The persistence of each head's error sign across optimizer steps admits errors into the QK channel, not aggregate magnitude or per-tile means. However, ordinary feature learning is also temporally coherent, so this quantity cannot serve as a trigger.

3. **QK-Guard provides triggered containment**: By monitoring the downstream consequence (attention-logit saturation) rather than the entry condition itself, QK-Guard fires only when a runaway develops and contains it across architectures (plain, RoPE, LLaMA-style), scales (125M to 350M), and fault sources (accumulation, quantization).

### Known Boundaries

- **Scale**: Largest model tested has 350M parameters
- **Arithmetic**: Software-emulated quantizer; native 4-bit hardware remains untested
- **Horizon**: Durability measured to 60k steps; some low-precision failures appear only after prolonged training

### Future Directions

- Native four-bit hardware validation of the dissociation
- Longer-horizon durability studies
- Testing whether containment constitutes repair or merely postpones failure
- Extending the channel-level account to post-attention fault sources

The paper's closing claim encapsulates its contribution: *"one action contains faults that share no source: fault source is not failure channel."*

---

_Markdown view of https://picx.dev/p/RhCb1j, served by PicX — AI-generated visual whiteboard summaries of research papers._