Technical Summary: Don't Read Everything: A Curvature-Conditioned Query for Linear Attention

1. Problem Statement

Linear attention mechanisms have emerged as a solution to the quadratic computational cost ($O(T^2)$) and memory overhead of standard softmax attention in Transformers. By replacing the softmax inner product with a kernelized dot product, linear attention reframes the read operation as a linear Recurrent Neural Network (RNN) with a matrix-valued state, achieving linear time complexity ($O(T)$) and constant cache ($O(1)$).

However, despite recent advancements in write-side mechanisms (e.g., data-dependent gating in GLA, Mamba, and RWKV; delta-rule updates in Gated DeltaNet), linear attention models consistently underperform standard softmax attention on in-context retrieval and long-context tasks. The authors argue that this performance gap stems not only from how information is written into memory but, crucially, from how it is read.

In standard softmax attention, the log-sum-exp normalizer creates a competitive environment where the most distinctive key dominates the output, effectively suppressing matches that lie in high-density regions of the memory. Linear attention discards this normalizer to gain efficiency. Consequently, every past key contributes additively to the output forever. This causes "useful targets" (the "needles") to be diluted by the "bulk" of stored vectors (the "haystack"), particularly when the query overlaps with directions where memory is dense. Existing remedies focus on the write side (deciding what enters the state) but cannot react to which directions a future query will read against.

2. Methodology: Curvature-Conditioned Query (CCQ)

The paper proposes Curvature-Conditioned Query (CCQ), a lightweight, read-time correction mechanism inspired by the local geometry of the missing softmax normalizer.

Theoretical Derivation

The authors observe that the softmax competition term, $A_t(q_t) = \log \sum_{j \le t} \exp(k_j^\top q_t)$, is a function of the query $q_t$ alone once the keys are fixed. They perform a second-order Taylor expansion of this log-partition function around the isotropic-attention point ($q=0$):

$$A_t(q) \approx \text{const} + \mu_t^\top q + \frac{1}{2} q^\top (\bar{C}_t - \mu_t \mu_t^\top) q$$

Here, the Hessian of the expansion corresponds to the running sample covariance of the keys, $\Sigma_t = \text{Cov}_{\le t}[k]$. This covariance matrix captures the directions of high density in the memory.

The Mechanism

CCQ utilizes this geometric insight to construct a "cleaned" query before it interacts with the linear attention state. Instead of reading the state with the raw query $q_t$, the model applies a linear contraction:

$$q_t^{\text{clean}} = (I - \lambda_t \Sigma_t) q_t = q_t - \lambda_t \Sigma_t q_t$$

• $\Sigma_t$: The running key covariance, maintained via the same recurrent or chunkwise mechanism as the backbone state (requiring one additional $d_k \times d_k$ matrix and one $d_k$-vector per head).
• $\lambda_t$: A scalar gate conditioned on the query ($\lambda_t = \sigma(W_\lambda q_t + b_\lambda)$), controlling the aggressiveness of the contraction.
• Effect: The operator contracts the query along the high-variance (high-density) directions of the memory. This penalizes keys that align with the bulk of the stored vectors, effectively "pulling" the query away from crowded directions to favor isolated targets.

Integration

CCQ modifies only the read step. It is fully composable with any linear-attention backbone (e.g., GLA, Gated DeltaNet) because it does not alter the write rule or the state recurrence $S_t = \text{Update}(S_{t-1}; k_t, v_t, \dots)$. The final output is computed as $o_t = S_t q_t^{\text{clean}}$.

3. Key Contributions

1. Identification of the Read-Step Bottleneck: The paper identifies the read step as a distinct, under-explored lever for closing the performance gap between softmax and linear attention, complementing existing write-side remedies.
2. CCQ Mechanism: The proposal of a curvature-aware linear contraction derived from a single Taylor expansion of the softmax log-partition. It is additive to any linear-attention write rule.
3. Empirical Validation: Extensive experiments showing that CCQ improves perplexity, zero-shot downstream accuracy, needle-in-a-haystack (S-NIAH) retrieval, length extrapolation, and LongBench performance when attached to GLA and Gated DeltaNet backbones.

4. Experimental Results

The authors evaluated CCQ on 500M and 1.3B parameter models trained on the FineWeb-Edu subset.

• Language Modeling & Downstream Accuracy:
  • At 1.3B, CCQ-GLA achieved the best WikiText perplexity (16.92) and LAMBADA accuracy (47.82%) among recurrent models, outperforming its backbone by significant margins.
  • CCQ-Gated DeltaNet achieved the best 7-task average accuracy (56.24%) among recurrent models.
• Synthetic Retrieval (S-NIAH):
  • CCQ significantly rescued backbones on length extrapolation tasks. For instance, at 1.3B, CCQ-GLA improved S-NIAH-1 (pass-key) accuracy at 4K context by +49.0% over the GLA backbone.
  • While standard Transformers collapsed to 0% accuracy at 8K context (due to 4K training limits), recurrent CCQ variants retained substantial accuracy.
• Length Extrapolation:
  • On six long-context corpora (GovReport, QMSum, etc.), CCQ variants consistently achieved lower perplexity than their backbones across 4K to 20K contexts.
  • At 20K, CCQ variants beat their respective backbones by 1–2.5 perplexity points.
• Long-Context Understanding (LongBench):
  • CCQ-GLA achieved the best average score (15.77) among recurrent models on the 14-task LongBench, outperforming Mamba2 by +2.64 points.
  • Both CCQ variants showed marked improvements on tasks requiring long-output generation (e.g., GovReport, MultiNews) and few-shot reasoning.

5. Significance and Claims

The paper claims that the performance gap in linear attention is not merely a capacity issue but a geometric one: the lack of competition in the read step allows memory density to drown out specific targets.

• Geometric Insight: The authors demonstrate that the "missing" competition in linear attention can be approximated by the local curvature (Hessian) of the softmax log-partition, which coincides with the key covariance.
• Memory-Aware Correction: Unlike input-only query projections, CCQ is "memory-aware" because it uses the running statistics of the stored keys to condition the query.
• Modularity: The method is presented as a drop-in replacement for the read step, compatible with any linear-attention backbone without modifying the write dynamics.
• Modest Scope: The authors explicitly note limitations, stating that results are currently limited to 500M/1.3B scales, two specific backbones (GLA, Gated DeltaNet), and specific benchmarks. They do not claim CCQ solves all linear attention issues or that it generalizes to 7B+ scales or state-space models like Mamba without further investigation.

In summary, the paper argues that by "not reading everything" (i.e., contracting the query along dense memory directions), linear attention can recover much of the retrieval and long-context capability lost by discarding the softmax normalizer, at a minimal computational cost.