PaperlayerView exact revision on arXiv

Source-linked AI summary

Gated DeltaNet-2: Decoupling Erase and Write in Linear Attention

Ali Hatamizadeh, Yejin Choi, Jan Kautz

arXiv:2605.22791v1cs.AI

TL;DR

Fixed-size recurrent attention compresses long histories into limited memory, while existing delta-rule gates tie erasing old associations to writing new values. Gated DeltaNet-2 separates these decisions channel-wise and achieves the strongest overall results across language modeling, commonsense reasoning, and retrieval, especially long-context retrieval.

  • Problem

    Fixed-size recurrent attention compresses associations into finite memory, while KDA and Gated DeltaNet tie key-side erasing and value-side writing to one scalar gate.

  • Method

    Gated DeltaNet-2 uses channel-wise erase and write gates, preserves channel-wise decay, recovers KDA and Gated DeltaNet as special cases, and supports efficient chunkwise training.

  • Results

    Gated DeltaNet-2 achieves the strongest overall results across language modeling, commonsense reasoning, and retrieval, with clearest gains on long-context multi-key retrieval.

  • Takeaways & Limitations

    Decoupling erase and write improves the recurrent and hybrid frontier across language modeling, commonsense reasoning, synthetic retrieval, and real-world recall.

  • Takeaways & Limitations

    Fixed-size recurrent states retain old associations until later updates overwrite them indirectly, limiting exact retrieval in long contexts.

Abstract

from arXiv · show

Linear attention replaces the unbounded cache of softmax attention with a fixed-size recurrent state, reducing sequence mixing to linear time and decoding to constant memory. The hard part is not just what to forget, but how to edit this compressed memory without scrambling existing associations. Delta-rule models subtract the current read before writing a new value, and Kimi Delta Attention (KDA) sharpens forgetting with channel-wise decay. But the active edit still uses a single scalar gate to control two different things: how much old content to erase on the key side and how much new content to commit on the value side. We introduce Gated DeltaNet-2, which generalizes both Gated DeltaNet and KDA by inheriting adaptive forgetting and channel-wise decay while addressing their shared limitation, the scalar tie between erasing and writing. Gated Delta Rule-2 separates these roles with a channel-wise erase gate b_t and a channel-wise write gate w_t, reducing to KDA when both gates collapse to the same scalar and to Gated DeltaNet when the decay also collapses. We derive a fast-weight update view, a chunkwise WY algorithm with channel-wise decay absorbed into asymmetric erase factors, and a gate-aware backward pass that preserves efficient parallel training. At 1.3B parameters trained on 100B FineWeb-Edu tokens, Gated DeltaNet-2 achieves the strongest overall results among Mamba-2, Gated DeltaNet, KDA, and Mamba-3 variants across language modeling, commonsense reasoning, and retrieval. Its advantage is most pronounced on long-context RULER needle-in-a-haystack benchmarks, where it improves the evaluated multi-key retrieval setting and remains strong in both recurrent and hybrid settings. Code is available at https://github.com/NVlabs/GatedDeltaNet-2.

1. Introduction

Gated DeltaNet-2 addresses the finite-memory and editing limitations of linear recurrent attention by decoupling key-side erasing from value-side writing while preserving channel-wise decay. The resulting recurrence retains efficient chunkwise training and improves long-context retrieval, especially in multi-key settings.

  • Transformer self-attention offers direct historical access and parallel training, but its quadratic sequence-length cost impedes long-context training and high-throughput inference.
  • Linear recurrent attention uses a fixed-size key-value state for linear-time sequence mixing, but compressed memory makes exact retrieval difficult over long contexts.
  • Gated DeltaNet-2 decouples key-side erase gates from value-side write gates, removing the scalar tie between which old coordinates are removed and which incoming coordinates are committed.
  • Absorbing cumulative channel-wise decay into rank-one erase factors yields a compact WY recurrence with the chunkwise structure required by efficient delta-rule kernels.
  • Gated DeltaNet-2 shows its clearest gains on long-context RULER retrieval, particularly the evaluated multi-key case, and maintains the strongest overall retrieval profile in recurrent and hybrid settings.

2. Preliminary

Linear attention stores token associations in a fixed-size recurrent matrix and supports parallel or chunkwise computation, but its basic additive update cannot explicitly remove stale associations. Successive decay and delta-rule variants add forgetting and active editing, while KDA retains a scalar active gate that jointly controls erasure and writing.

  • Linear attention: Linear attention stores a fixed-size matrix state, reads it with queries, and replaces tokenwise recurrence with a parallel causal matrix form.The state size is independent of sequence length.
  • Linear attention limitation: Every outer product is added to the state and none is removed, so old associations persist until later writes overwrite them indirectly through superposition.
  • Chunkwise computation: Chunkwise computation keeps recurrence only across chunks while expressing within-chunk interactions as dense matrix products, preserving linear sequence complexity for fixed chunk size C.This schedule maps well to tensor cores.
  • Prior mechanisms: Mamba-2 adds data-dependent scalar decay, DeltaNet performs an active read-subtract-write edit, and Gated DeltaNet combines global forgetting with selected-association editing.In DeltaNet, β_t = 1 overwrites the association at key k_t, while β_t = 0 leaves it unchanged.
  • KDA limitation: KDA makes decay channel-wise and retains DeltaNet’s efficient WY-based chunkwise algorithm, but its scalar active gate β_t still jointly controls erasure and writing.This remaining tie motivates Gated DeltaNet-2.

3. Gated DeltaNet-2

Gated DeltaNet-2 decouples channel-wise erasure and writing through independent gates while retaining adaptive forgetting and channel-wise decay. Its asymmetric delta update preserves efficient chunkwise computation and parallel training, and recovers KDA and Gated DeltaNet as tied special cases.

  • Gated Delta Rule-2: Gated Delta Rule-2 uses channel-wise erase gate b_t for key-side removal and write gate w_t for value-side insertion, eliminating KDA’s scalar tie.The erase gate selects key coordinates used to read old content, while the write gate selects value coordinates being inserted.
  • Gated Delta Rule-2: Gated Delta Rule-2 recovers KDA when both gates equal scalar β_t and recovers Gated DeltaNet when channel-wise decay further becomes scalar.These models are tied subspaces of the more general update.
  • Fast-weight update: The fast-weight interpretation decays the state, reads old content through a gated erase direction, and writes a correction toward a gated value target.Unlike correlation-writing Mamba variants, the delta family writes a residual between the target value and currently read memory.
  • Chunkwise algorithm: Channel-wise decay is absorbed into asymmetric erase factors, leaving fixed-size chunk updates composed of triangular solves and dense matrix multiplication.The erase gate enters the erase-side auxiliary, while the write gate enters the write-side auxiliary; fused Triton kernels implement the computation.
  • Backward pass: The backward pass retains KDA’s matrix structure but explicitly carries distinct erase- and write-gate factors through vector-Jacobian products.The scalar-gated shortcut fails because the gates act diagonally on different key and value channels.

4. Experiments

Experiments compare matched 1.3B recurrent and hybrid models trained on 100B FineWeb-Edu tokens across language modeling, reasoning, retrieval, ablations, and throughput. Gated DeltaNet-2 achieves the best average results while retaining efficient scaling with sequence length.

  • Setup: All models use 1.3B parameters, are trained on 100B FineWeb-Edu tokens, and are evaluated in recurrent-only and hybrid forms.The hybrid form pairs the recurrent token mixer with sliding-window attention.
  • Language modeling and common-sense reasoning: Gated DeltaNet-2 achieves the best average in both recurrent and hybrid settings across language modeling and common-sense reasoning.Because recurrent state size is matched, the passage attributes the gain to a stronger update rule rather than larger memory; the trend persists with SWA.
  • In-context retrieval on synthetic data: Gated DeltaNet-2 is strongest on synthetic RULER retrieval tasks where memory editing matters most.Table 3 reports S-NIAH and MK-NIAH under fixed-state memory, testing retention, interference control, high-entropy value storage, and multi-key discrimination.
  • Gate structure and erase range ablations: Both scalarized gate variants trail full Gated DeltaNet-2 in the channel-structure ablations.The ablations remove channel-wise gate variation while keeping the original projections and parameter count fixed.
  • Throughput comparison: 38.0 to 36.1 Kt/s: Gated DeltaNet-2’s hybrid training throughput drops only mildly as sequence length grows on a single H100.The small gap relative to KDA reflects the added channel-wise erase and write gates, indicating a modest constant efficiency cost.

5. Related Work

Efficient sequence models replace quadratic self-attention with recurrent or linear-time token mixers that maintain fixed-size states. The field progressed from mostly data-independent structured state-space and recurrent transitions to selective dynamics, SSD, and learned decay gates.

  • Efficient Sequence Models: Efficient sequence models maintain fixed-size recurrent states to replace quadratic self-attention with recurrent or linear-time token mixing.Early structured state-space and recurrent models primarily used data-independent transitions.
  • Selective Dynamics and Gating: Mamba and Mamba-2 introduced data-dependent selective dynamics and the SSD framework, while gated linear attention improved memory control with learned decay gates.These approaches extend efficient sequence modeling beyond mostly data-independent transitions.

6. Conclusion

Gated DeltaNet-2 decouples active memory editing into channel-wise erase and write decisions. It removes the shared scalar gate while recovering Gated DeltaNet and KDA as special cases and retaining efficient chunkwise training.

  • Core mechanism: Gated DeltaNet-2 uses channel-wise erase and write gates to separate key-side removal from value-side commitment during memory updates.The erase gate b_t selects decayed-state key coordinates to read and remove, while the write gate w_t selects value coordinates to commit.
  • Core mechanism: Decoupling the gates removes the scalar β_t tie shared by Gated DeltaNet and KDA.The separate decisions eliminate the constraint that erasing and writing must be controlled by one scalar.
  • Special cases: Gated DeltaNet and KDA are recovered as special cases of Gated DeltaNet-2.The model’s generalized gate formulation contains both prior mechanisms as limiting cases.
  • Efficient training: A WY form with gate-aware kernels preserves efficient chunkwise training.The implementation maintains parallel training efficiency while supporting the decoupled erase and write gates.

A. Chunkwise derivation for Gated DeltaNet-2 … A.5. Tied-gate reductions

The appendix derives a chunkwise WY formulation for Gated DeltaNet-2 by absorbing channel-wise decay into asymmetric erase factors and sharing a triangular solve across erase- and write-side computations. It also shows that tying channel gates recovers KDA and, with tied decay, Gated DeltaNet.

  • A. Chunkwise derivation for Gated DeltaNet-2: The chunkwise derivation operates on a length-C chunk with start state S0 and per-head key- and value-dimensional variables.The state is S_r ∈ R^{d_k×d_v}, with key-side vectors in R^{d_k} and value-side vectors in R^{d_v}.
  • A. Chunkwise derivation for Gated DeltaNet-2: All exponentials, products, and ratios involving γ act elementwise over the key-channel axis.This establishes the channel-wise interpretation used throughout the chunkwise derivation.
  • A.1. Decay-normalized recurrence: After decay normalization, channel-wise decay disappears from the recurrence and is carried by the left and right factors of each rank-one edit.The normalized factors are defined as K̄ = γ^-1 ⊙ K, Ē = γ ⊙ (B ⊙ K), and Z = W ⊙ V.
  • A.1. Decay-normalized recurrence: The chunk computation forms a strictly lower-triangular causal matrix T, whose unit-diagonal inverse A is obtained by forward substitution.The construction uses T = tril(ĒK̄^⊤, −1), A = (I + T)^−1, Y = AĒ, and U = AZ.
  • A.2. Compact state formula: The compact state formula expresses normalized residual increments as R = AZ − AĒS0 = U − YS0, proving the corresponding prefix-state update.The first r rows of the relevant matrices define each prefix, and substituting the increments into the normalized recurrence yields the stated state formula.
  • A.3. Compact output formula: The output derivation defines Qγ row-wise from γr ⊙ qr and combines it with a causal score matrix.This gives the compact output construction corresponding to Eq. 24 in the main text.
  • A.4. Row recurrences: Erase-side and write-side auxiliaries solve the same lower-triangular system with different right-hand sides, allowing one WY inverse to be shared.The shared inverse supports efficient chunkwise computation for both computations.
  • A.5. Tied-gate reductions: Tying b_r and w_r to the same scalar gate recovers KDA, while additionally tying α_r to a scalar recovers Gated DeltaNet without changing the algorithm.Under these reductions, the chunkwise equations specialize accordingly; channel-wise decay otherwise prevents Ē from being a scalar row scaling of K̄.

B. Backward derivation … B.5. Why scalar post-scaling is invalid

The appendix derives chunk-level vector-Jacobian products for the recurrent state and output, including a gate-aware WY inverse path and elementwise cumulative-decay gradients. It shows that independently gated erase and write operators must enter the accumulation products directly, making scalar post-scaling valid only for tied gates.

  • B. Backward derivation: Chunk-level vector-Jacobian products use upstream gradients dO and dS_C together with the forward chunk equations.The forward quantities are U = AZ, Y = ĀE, A = (I + T)^−1, and T = tril(ĒK̄⊤, −1).
  • B.1. Output and state paths: The output and state backward paths apply the causal mask to dA_qk and use rowsum to return a d_k-dimensional vector.The residual relation from Eq. 49 supplies the remaining backward connection.
  • B.2. Gate-aware WY inverse path: The WY inverse path computes auxiliary products and uses only the strictly lower triangular part of dT.This follows from the construction of T and the associated gate-aware accumulation equations.
  • B.4. Elementwise gates and cumulative decay: Elementwise gate gradients use normalized keys and queries, with implementation performing a reverse cumulative sum over each chunk.The vector-Jacobian products are derived for the normalized quantities before accumulation.
  • B.5. Why scalar post-scaling is invalid: In KDA, scalar β_r multiplies both value and erase right-hand sides, allowing the factor to move after the dot product.This factoring applies to the tied scalar-gate case.
  • B.5. Why scalar post-scaling is invalid: Gated DeltaNet-2 replaces β_s v_s with w_s ⊙ v_s, so a distinct diagonal write operator cannot be recovered by row or column scalar post-scaling.The independently gated erase side has the analogous b_s ⊙ k_s issue.
  • B.5. Why scalar post-scaling is invalid: Because erase and write gates act inside their respective products, they must be baked into the dot products of Eqs. 64 and 65.A scalar post-scale is correct only when the gates are tied.

C. Layer and kernel implementation … E. Experimental settings

The implementation separates channel-wise erase and write processing across chunked and recurrent kernels, with gate-aware backward computation and hardware-specific scheduling. Numerical safeguards use fp32 for decay, recurrent state, accumulators, and precision-sensitive solves, while verification compares chunkwise computation against recurrent references.

  • C.1. Layer parameterization: Independent projections produce channel-wise erase and write gates, with shapes dmodel → Hdk and dmodel → Hv dv; grouped value attention repeats key-side tensors across value-head groups.The repeated tensors are q, k, log-decay g, and b; v and w already use the value-head axis.
  • C.2. Forward kernels: Chunked execution fixes C = 64 and uses intra-chunk causal products, a WY forward substitution solve, auxiliary construction, then state recurrence and output kernels.Erase factors enter the key tile before the dot product, while state and output kernels retain KDA-compatible matrix shapes.
  • C.3. Backward kernels: The gate-aware WY backward kernel emits gradients for the erase and write gates, whose shapes are B × T × H × dk and B × T × Hv × dv.Its direct gradients include dW = dZ ⊙ V, dV = dZ ⊙ W, and dB += dĒ ⊙ γ ⊙ K.
  • C.4. Autotuning and hardware dispatch: On Hopper GPUs, autotuning restricts the fused WY backward search to two and four warps because eight warps can trigger a Triton WGMMA layout assertion; Ampere retains the full search.The restriction changes only scheduling, not the mathematical operation.
  • C.5. Recurrent decoding kernel: The recurrent decoding kernel processes tokens autoregressively, keeps state in fp32, applies exp(g_t), reads through b_t ⊙ k_t, and writes w_t ⊙ v_t along k_t.A forward-only recurrent kernel is provided for short sequence lengths.
  • C.6. Variable-length sequences: Packed variable-length batches use cumulative sequence lengths, reset recurrent state at sequence boundaries, share layouts across kernels, and remove padding before projection before restoring it afterward.The same layout covers chunk forward, chunk backward, and recurrent decoding.
  • D.1. Decay precision; D.3. State and accumulator dtypes; D.4. WY solve precision: Explicit fp32 decay computation, fp32 recurrent states and matrix accumulators, and an exposed precision flag for the triangular solve protect path-length-dependent products and precision-sensitive forward substitution.The conservative IEEE fp32 solve path is selected when required by hardware, while other matrix products may use tensor cores.
  • D.2. Query and key normalization; D.5. Initialization and output gate; D.6. Correctness checks: Queries and keys are L2-normalized per head, all linear layers use Xavier uniform initialization with gain 2−2.5 and zero biases when present, and correctness is checked against tokenwise recurrent references and autograd.The output applies RMSNorm and a SiLU gate; fp64 gradients agree to machine precision, while production fp32 and bfloat16 errors match expected accumulation and mantissa noise.

E.1. Training

The study compares Gated DeltaNet-2 with Transformer and recurrent baselines in recurrent-only and hybrid configurations. Models are matched for scale and trained under a common 1.3B-parameter, 100B-token FineWeb-Edu setup.

  • Evaluation setup: Evaluations compare Gated DeltaNet-2 with Transformer, Mamba-2, Gated DeltaNet, KDA, and Mamba-3 baselines.Each recurrent architecture is tested in recurrent-only and hybrid forms; hybrids combine the recurrent token mixer with sliding-window attention.
  • Evaluation setup: Hybrid models use sliding-window attention with a 2K window under the same residual block structure.The hybrid design follows Section 3.5 and pairs each recurrent token mixer with SWA.
  • Model matching: Mamba-3 is evaluated in both SISO and MIMO variants, with MIMO preserving the SISO main recurrent state size while adding rank-R_MIMO parameterization.This isolates the MIMO parameterization without changing the main recurrent state size.
  • Model matching: For fair recurrent comparisons, models are matched in parameter count and main recurrent state size.Gated DeltaNet, KDA, and Gated DeltaNet-2 use H=16 heads with d_k=128 and d_v=128; Mamba-2 and Mamba-3 use expansion factor 2, head dimension 64, and d_state=64.
  • Training configuration: 1.3B-parameter models are trained on 100B FineWeb-Edu tokens with AdamW, a 4 × 10^-4 peak learning rate, 0.1 weight decay, and 1.0 gradient clipping.Training uses cosine annealing, a 1B-token warm-up, a 0.5M-token global batch, 4K-token sequences, and 2K SWA windows for hybrids.

E.2. Evaluation

The evaluation covers language modeling, zero-shot commonsense reasoning, and in-context retrieval. Retrieval is tested in both synthetic needle-in-a-haystack settings and real-world datasets.

  • Language modeling and common-sense reasoning: Language modeling uses perplexity on WikiText and LAMBADA, alongside LAMBADA accuracy for zero-shot transfer.The zero-shot suite also includes PIQA, HellaSwag, WinoGrande, ARC-Easy, ARC-Challenge, OpenBookQA, Social IQa, and BoolQ.
  • In-context retrieval: Synthetic retrieval uses RULER’s Single Needle-In-A-Haystack and Multi-Key Needle-In-A-Haystack tasks.S-NIAH-1 tests passkey retrieval, S-NIAH-2 numerical needles, S-NIAH-3 word-based needles, and MK-NIAH-1 includes distractor key-value pairs.
  • In-context retrieval: Real-world retrieval covers structured relation extraction, key-value retrieval, and question answering across SWDE, FDA, SQuAD, TriviaQA, DROP, and Natural Questions.SWDE uses HTML, FDA uses PDFs, and the remaining datasets provide question-answering evaluations.
Loading 2605.22791v1…