PaperlayerView exact revision on arXiv

Source-linked AI summary

The Query Knows What to Forget: A Second Erase Direction for Linear Attention

Dhruman Gupta, Aritra Das, Debayan Gupta

arXiv:2608.13668v1cs.LG

TL;DR

Linear attention’s key-derived erase step cannot directly remove query-visible interference in the fixed-size state. QED adds a query-derived, key-orthogonal erase direction and improves long-context retrieval, extending useful context from approximately 8K to 16K tokens at comparable retrieval accuracy.

  • Problem

    Key-derived erasure cannot directly modify query-visible content in the fixed-size state, creating a mismatch that limits retrieval at long context.

  • Method

    QED augments GDN-2 with a gated, query-derived erase correction projected orthogonal to the key while preserving its rank-one update.

  • Results

    QED improves long-context retrieval beyond the 2048-token training window and FDA recall, extending useful context from approximately 8K to 16K tokens at comparable retrieval accuracy.

  • Takeaways & Limitations

    The experiments support using the query to define QED’s added erase direction, while the erase gate and key-orthogonal projection are not individually established as necessary.

  • Takeaways & Limitations

    Ablation effects vary substantially across training runs, so additional runs and a training-time λmax sweep are needed to establish remaining component-level effects.

Abstract

from arXiv · show

Linear attention keeps a state of fixed size. At long context, many stored items share this state, and interference between them degrades retrieval. Gated DeltaNet-2 (GDN-2), like every delta-rule model before it, derives its erase vector from the key of the current token. However, the interference in its reads is measured through the query, and the erase step cannot reach it. We introduce the Query-derived Erase Direction (QED). QED adds a second erase direction derived from the query and orthogonal to the key. In the fast-weight view, a key-directed delta edit cannot change the key-orthogonal part of a read. It uses the editable part to cancel old-state content measured along the query. It also improves retrieval at every length past the training window, and it about doubles the usable context length on S-NIAH-1.

1 Introduction

Linear attention’s fixed-size state causes growing interference, while key-derived erasure cannot directly modify query-visible content. QED addresses this mismatch with a gated, key-orthogonal erase direction derived from the query and improves long-context retrieval.

  • Motivation: Fixed-size recurrent states make linear attention efficient but cause stored associations to interfere as context grows, reducing retrieval reliability.Linear attention uses constant memory and per-token sequence-mixing cost during autoregressive decoding.
  • Motivation: Key-derived erasure mismatches query-based reading because query-visible content can lie outside the state component that key erasure directly modifies.The query selects returned content, while the key-derived erase vector selects content removed by the update.
  • QED: QED adds a gated erase direction derived from the query, projected orthogonal to the key, to cancel old-state content inaccessible to the standard delta update.The change preserves the rank-one update and its stability constraints.
  • Results: Up to 39 percentage points: QED improves S-NIAH-1 retrieval beyond the training length across two training seeds.It also improves FDA by 3.4 and 3.7 points.
  • Ablations: Query-defined added erase directions are consistently beneficial, while gating and key-orthogonal projection have inconclusive effects across runs.These findings come from controlled ablations.

2 Background

The background traces linear attention from additive state updates to increasingly flexible forgetting and correction mechanisms. GDN-2 decouples erase and write control, but its erase direction remains derived from the current key, which is the target of the proposed change.

  • Linear attention and the delta rule: Linear attention stores a fixed-size state matrix instead of a softmax cache and updates it additively with the current key and value.The base state is S_t ∈ R^(K×V), updated by adding k_t v_t^⊤.
  • Linear attention and the delta rule: DeltaNet replaces additive updates with an error-correcting rank-one step that erases old content at the key before writing the new value.The scalar β_t controls the rank-one correction that removes the old value at the key.
  • Gated DeltaNet and KDA: Gated DeltaNet adds scalar decay, while KDA makes decay channel-wise so each key channel forgets at its own rate.GDN applies one scalar α_t to every state channel; KDA uses a vector α_t ∈ (0, 1)^K.
  • Gated DeltaNet-2 (GDN-2): GDN-2 decouples erase and write control with channel-wise gates, yet its erase vector remains b_t ⊙ k_t and therefore derives from the current key.The key-directed rank-one edit removes the state component read through the gated key before writing the new value at k_t.

3 Query-derived Erase Direction

QED addresses query-aligned interference that key-directed erasure cannot remove in fixed-size linear-attention states. It adds a gated, query-derived erase direction orthogonal to the key, enabling correction of the current read through key–query overlap.

  • Motivation: Longer contexts pack more key–value pairs into the same fixed-size state, increasing query-aligned interference and degrading retrieval.The read is a mixture containing the target plus interfering stored pairs selected by the query.
  • QED mechanism: QED adds a second erase direction derived from the query and projected orthogonally to the key, allowing the layer to remove content selected by what it reads.The added direction is d_t = P⊥_k_t(b_t ⊙ q_t), with learned per-head strength λ_h ≥ 0.
  • Fast-weight view: Key-directed delta edits cannot change state responses along directions orthogonal to the current key, leaving part of the query-read interference unreachable.This inaccessible content is the old-state component selected by the query outside the key direction.
  • QED mechanism: QED samples old-state content along the gated query-derived direction, writes its negative sample along the key, and reaches the query through key–query overlap ρ_t.The gate scales the correction channel-wise, preventing erasure on every read and preserving content needed by later reads.

4 Experimental setup

The experiments compare GDN-2-based model arms that isolate the gate, projection, and correction direction, using controlled training and long-context retrieval evaluations. They also test QED recurrence strength at inference time with fixed weights.

  • Model and ablations: The baseline is a 24-layer decoder whose blocks contain one GDN-2 layer and one SwiGLU layer with RMSNorm.Each arm removes one mechanism component to identify which part drives performance.
  • Model and ablations: QED-rand replaces the query with a fixed random direction, while QED-nogate removes the gate but retains the query direction and projection.The arms isolate whether the query-derived direction and gate are necessary.
  • Training: 340M parameters, 24 layers, hidden size 1024, 6 heads, head dimension 128, tied embeddings, and 15B training tokens define every model.Training uses FineWeb-Edu, 2048-token sequences, AdamW at learning rate 3 × 10−4, BF16 with FP32 master weights, and seeds 0 and 1.
  • Evaluations: 500 samples for each length from 1K to 32K tokens are used on RULER S-NIAH-1, S-NIAH-2, S-NIAH-3, plus SWDE, FDA, and SQuAD-completion.The latter three are Based recall tasks.
  • Inference-time strength test: For each final QED checkpoint, inference scales the learned strength by s ∈ {−1, 0, 0.5, 1, 1.5, 2} during prefill and decoding with fixed weights.The same deterministic S-NIAH-1 examples are compared against the trained setting s = 1 using the exact McNemar test.

5 Results

QED leaves validation loss unchanged while improving S-NIAH-1 accuracy at every context length beyond the 2048-token training window. Its correction-strength effects vary across checkpoints, and recall-task gains are specific to FDA.

  • Validation loss: QED does not change validation loss, consistent with targeting retrieval rather than perplexity.The full ten-run spread is 0.0034 nats, with at most 0.0013 nats within each seed.
  • S-NIAH-1: QED increases S-NIAH-1 accuracy at every length above the 2048-token training length, in both seeds.The 4K–32K evaluations are extrapolation beyond the training window, and QED remains above GDN-2 at every such length.
  • Correction-strength ablation: At 16K, one QED checkpoint reaches 44.8% with reversed correction, 55.0% with correction removed, 65.6% at trained strength, and 73.0% when doubled.For the other checkpoint, accuracy stays within 2.4 points of the trained setting without a consistent directional trend.
  • Recall tasks: Every arm with the projection increases FDA above baseline in both seeds, whereas SQuAD-completion is unchanged and SWDE shows only a small seed-0 gain.These short-document recall tasks measure recall rather than extrapolation; FDA comparisons use checkpoints trained for the same 15B-token schedule.

6 Discussion

QED improves long-context retrieval beyond the 2048-token training window without measurable changes in validation loss or SQuAD-completion. Ablations consistently support the query-derived erase direction, while evidence for the erase gate, projection, and inference-time correction strength remains inconclusive.

  • Retrieval gains: QED improves S-NIAH-1 accuracy at every evaluated length beyond the 2048-token training window across both training seeds and also improves FDA recall.Validation loss and SQuAD-completion do not measurably change, suggesting the gains concern access to compressed recurrent memory rather than general language-modeling quality.
  • Ablations: QED outperforms the random-direction control at every extrapolation length in both training seeds.This is the most consistent ablation evidence for defining the added erase direction from the query.
  • Ablations: The erase gate and key-orthogonal projection provide inconsistent benefits across runs, so neither sub-component is individually established.The experiments support the complete QED update and its query-derived direction, but not definitive component-level conclusions for the gate or projection.
  • Projection analysis: In 1.18M token-head samples from two unprojected models, the observed transition eigenvalue remained within [−0.011, 0.985], and counterfactual projection changed it by less than 10−3.Although the unprojected eigenvalue can in principle lie in [−0.25, 1.25], learned corrections showed negligible alignment with the key at this scale.
  • Limitations: Inference-time correction-strength changes differed across training runs, and additional runs plus a training-time λmax sweep are needed to establish the remaining effects.Accuracy increased with strength in one run but changed little in the other; the test therefore does not determine QED-trained models’ reliance on the correction.

7 Conclusion and Further Work

QED augments the erase vector with a query-derived correction while preserving the rank-one recurrent update. In matched 340M-parameter training schedules, it improves long-context retrieval beyond training length and FDA recall without measurable changes in validation loss or SQuAD-completion, while motivating broader evaluation and stability studies.

  • Conclusion: QED adds a query-derived correction to the erase vector while preserving the rank-one recurrent update.The ablations consistently support using the query to define the added erase direction.
  • Conclusion: 340M-parameter models trained on matched 15B-token schedules consistently improve long-context retrieval beyond training length and FDA recall.These gains occur without a measurable change in validation loss or SQuAD-completion.
  • Further Work: Further work should evaluate QED across more training runs, model scales, context lengths, and retrieval tasks.The proposed scope is intended to test whether the reported improvements generalize broadly.
  • Further Work: A training-time sweep over λmax should determine useful correction strength, while larger-scale experiments should test when the projection’s stability constraint becomes active.These studies target the correction-strength choice and the operating regime of the stability constraint.

Appendix · A The strength parameter

QED’s strength is controlled by one learned, sigmoid-gated scalar per layer and head, bounded by λmax. It starts with a small correction, changes slowly in practice, and adds limited computation while preserving the GDN-2 kernel structure.

  • A The strength parameter: A learned scalar λh controls QED strength for each layer and head.The scalar is set by a sigmoid gate.
  • A The strength parameter: The sigmoid gate bounds λh using the maximum correction strength λmax.λmax is the cap applied to the correction.
  • A The strength parameter: Each θh starts at −2, initializing λh at λmax sigmoid(−2) ≈0.03.This makes the correction active but small at the start of training.
  • A The strength parameter: θh uses a learning-rate multiplier of 10 during training.The multiplier is applied to the parameter controlling the strength gate.
  • A The strength parameter: In 1B-token test runs, the mean λ moved less than 1% from its initial value at the default rate.Despite the multiplier, θ moves very slowly at the default rate.
  • A The strength parameter: At λ = 0, equation (6) reduces to equation (4), and QED changes only the erase vector.The QED kernel therefore retains the structure of the GDN-2 kernel.
  • A The strength parameter: QED adds one dot product and one subtraction for each token and head.This is the stated extra computational work of the QED kernel.

A.1 Inference-time strength test

Inference-time scaling of QED strength with fixed weights produced checkpoint-dependent effects on S-NIAH-1 accuracy. One checkpoint benefited consistently from stronger correction, while the other changed little; the sweep measures the accumulated recurrence effect and does not establish a training-time strength constraint.

  • Method: The sweep changes the full recurrence during both prefill and decoding, with fixed model weights, rather than modifying only the final read.Because every earlier update uses s, the incoming state depends on the strength and the test measures the accumulated QED recurrence effect.
  • Results: At 16K, one checkpoint’s accuracy rose from 44.8% with the correction reversed to 55.0% with it removed, 65.6% at trained strength, and 73.0% when doubled.Accuracy increased at every context length as correction strength increased for this checkpoint.
  • Results: For the other checkpoint, changing correction strength altered accuracy by no more than 2.4 points.The two checkpoints therefore did not show the same inference-time sensitivity.
  • Method: Since d_t ⊥ k_t, changing strength does not alter the nontrivial eigenvalue of the projected transition, ruling out eigenvalue instability as the cause of accuracy changes.The passage attributes the measured effect to the accumulated recurrence rather than projected-transition eigenvalue instability.
  • Limitation: The gain at s = 2 motivates stronger-correction training tests but does not show that λ_max = 0.25 was binding; a training-time λ_max sweep is required.Changing the training bound could alter both θ_h and the remaining model weights.
Loading 2608.13668v1…