Official PyTorch implementation of Gated DeltaNet-2: Decoupling Erase and Write in Linear Attention.

Ali Hatamizadeh, Yejin Choi, and Jan Kautz.

Linear attention compresses an unbounded KV cache into a fixed-size recurrent state. The hard part is not just what to forget, but how to edit this compressed memory without scrambling existing associations. Prior delta-rule models (Gated DeltaNet, Kimi Delta Attention) tie erasing and writing to a single scalar gate — even though they act on different axes of the state.

Gated DeltaNet-2 decouples these two roles:

• ✂️ Channel-wise Erase Gate b_t — selects which key-side coordinates of the decayed state are read and removed
• ✍️ Channel-wise Write Gate w_t — selects which value-side coordinates of the new content are committed
• 🌀 Channel-wise Decay — inherited from KDA for fine-grained global forgetting
• 🔁 Strict Generalization — recovers KDA when both gates collapse to the same scalar, and Gated DeltaNet when the decay also collapses
• ⚡ Hardware-efficient Training — fast-weight WY chunkwise algorithm with gate-aware backward, fused in Triton

Given an erase gate b_t ∈ [0,1]^{d_k}, a write gate w_t ∈ [0,1]^{d_v}, and channel-wise decay D_t = Diag(α_t), the recurrent state evolves as:

Compared with KDA, the right factor of the rank-one erase becomes channel-selective on the key axis, and the write term becomes channel-selective on the value axis. The two decisions no longer share a single scalar.

We train all models at 1.3B parameters on 100B tokens of FineWeb-Edu, matched in recurrent state size, and compare against Mamba-2, Gated DeltaNet, KDA, and Mamba-3 (SISO and MIMO).

Gated DeltaNet-2 achieves the best average across both recurrent-only and hybrid settings:

Gated DeltaNet-2 is strongest where memory editing matters most — particularly the interference-heavy multi-key needle-in-a-haystack settings:

Across SWDE, SQuAD, FDA, TriviaQA, NQ, and DROP, Gated DeltaNet-2 leads the recurrent and hybrid frontier:

Gated DeltaNet-2 retains near-flat scaling with sequence length on a single H100 (training, hybrid 1.3B), with only a small constant overhead over KDA for the added channel-wise gates.

Ablations confirm both gates contribute, with the erase gate b_t accounting for most of the gain — consistent with its role in selectively protecting or revising key-side associations in the recurrent state.

• 05/21/2026: 🔥 Code Release: Train your own Gated DeltaNet-2 on FineWeb-Edu
• Watch this space for more exciting updates!

Launch your training with our streamlined command:

💡 Pro Tip: Add --interactive_job --debug for interactive debugging sessions!

We train 1.3B-parameter models on 100B tokens of FineWeb-Edu with:

• AdamW, peak LR 4e-4, weight decay 0.1, gradient clip 1.0
• Cosine schedule with 1B-token warmup
• Global batch size 0.5M tokens, sequence length 4K
• Hybrid models use a 2K sliding-window attention size
• 16 heads, d_k = d_v = 128, matched recurrent state size against Mamba-2/3 baselines

Copyright © 2026, NVIDIA Corporation. All rights reserved.

Licensed under the NVIDIA Source Code License-NC. See LICENSE for details.

Built on the shoulders of giants:

• Gated DeltaNet
• Kimi Delta Attention
• Flash Linear Attention
• Samba
• LiTGPT

If you find this work useful, please consider citing:

If you find this work useful, please consider:

• Starring the repository
• Citing our paper
• Contributing to the codebase

Join us in pushing the boundaries of linear attention! 🚀