PaperlayerView exact revision on arXiv

Source-linked AI summary

Gated Delta Networks: Improving Mamba2 with Delta Rule

Songlin Yang, Jan Kautz, Ali Hatamizadeh

arXiv:2412.06464v3cs.CLcs.LG

TL;DR

Linear Transformers remain limited by retrieval and long-context challenges, including memory collisions and insufficiently targeted memory control. This paper introduces the gated delta rule and hardware-efficient parallel training, yielding Gated DeltaNet models that consistently outperform Mamba2 and DeltaNet across diverse benchmarks.

  • Problem

    Linear Transformers face retrieval limitations from memory collisions, while existing gating and delta mechanisms each provide incomplete memory control.

  • Method

    The paper combines adaptive gating with delta updates in the gated delta rule and extends chunkwise parallelism for hardware-efficient training.

  • Results

    Gated DeltaNet consistently outperforms Mamba2 and DeltaNet across language modeling, reasoning, retrieval, length extrapolation, and long-context understanding benchmarks.

  • Takeaways & Limitations

    Hybrid architectures combining Gated DeltaNet with sliding window attention or Mamba2 further improve training efficiency and model performance.

  • Takeaways & Limitations

    Delta-rule variants retain theoretical limitations and moderate real-world performance concerns, while more expressive nonlinear recurrences can sacrifice training parallelism.

Abstract

from arXiv · show

Linear Transformers have gained attention as efficient alternatives to standard Transformers, but their performance in retrieval and long-context tasks has been limited. To address these limitations, recent work has explored two distinct mechanisms: gating for adaptive memory control and the delta update rule for precise memory modifications. We observe that these mechanisms are complementary: gating enables rapid memory erasure while the delta rule facilitates targeted updates. Building on this insight, we introduce the gated delta rule and develop a parallel training algorithm optimized for modern hardware. Our proposed architecture, Gated DeltaNet, consistently surpasses existing models like Mamba2 and DeltaNet across multiple benchmarks, including language modeling, common-sense reasoning, in-context retrieval, length extrapolation, and long-context understanding. We further enhance performance by developing hybrid architectures that combine Gated DeltaNet layers with sliding window attention or Mamba2 layers, achieving both improved training efficiency and superior task performance.

1 INTRODUCTION

Linear Transformers reduce inference memory but face retrieval failures from memory collisions and limitations in adaptive forgetting. Gated DeltaNet combines gating with delta updates and extends chunkwise parallel training, outperforming Mamba2 and DeltaNet across diverse benchmarks.

  • Self-attention scales quadratically with sequence length, creating substantial computational demands for training and inference.
  • Linear Transformers reduce inference memory by reframing kernelized attention as a linear RNN with matrix-valued states.
  • When sequence length exceeds model dimensionality, linear Transformers face inevitable key-value memory collisions that hinder exact retrieval.
  • Mamba2 uniformly decays all key-value associations, so forgetting one association also forgets unrelated associations.Its dynamic ratio αt controls state decay at each time step.
  • DeltaNet selectively replaces old key-value pairs but lacks rapid memory clearing during context switches.The paper proposes a gated delta rule that combines targeted updates with adaptive memory erasure, trained using an extended chunkwise-parallel algorithm.
  • Gated DeltaNet consistently outperforms Mamba2 and DeltaNet across language modeling, reasoning, retrieval, extrapolation, and long-context understanding benchmarks.Hybrid variants combining Gated DeltaNet with sliding window attention or Mamba2 further improve training efficiency and model performance.

2 PRELIMINARY

This section develops linear-attention formulations with decay and delta updates, emphasizing their recurrent, matrix, and chunkwise-parallel forms for efficient training. It also explains how decay forgets historical information while the delta rule selectively replaces stored associations.

  • Linear attention with decay: Vanilla linear attention underperforms Transformers in language modeling, motivating data-dependent decay to forget historical information.
  • Linear attention with decay: The decay recurrence uses a data-dependent scalar α_t ∈ (0, 1), whose cumulative products express the equivalent parallel form.The formulation is also related to state space duality and appears in several gated or recurrent architectures.
  • Chunkwise training: Chunkwise parallel training splits inputs and outputs into chunks, computing each chunk from the previous chunk’s final state and current query, key, and value blocks.This form enables hardware-efficient, linear-time training.
  • Chunkwise training: The chunkwise algorithm uses matrix multiplications and chunk-boundary decay notation, supporting tensor-core hardware optimization.Left and right arrows denote decay toward a chunk’s first and last positions, respectively; the state matrix decays over the entire chunk.
  • Delta networks: The delta update rule combines selective memory replacement with generalized Householder transitions, writing a value determined by the current value, old value, and writing strength β_t.β_t ∈ (0, 1) controls the linear combination, while extensions allow β_t ∈ (0, 2).
  • Delta networks: A hardware-efficient DeltaNet chunkwise algorithm uses cumulative products and the WY transform to express the recurrence with matrix multiplications suitable for tensor-core GPU optimization.This addresses the computational inefficiency that previously limited DeltaNet’s attention despite its associative-recall and language-modeling performance.

3 GATED DELTA NETWORKS

Gated DeltaNet combines adaptive state decay with targeted delta updates, pairing complementary memory-management and key-value association mechanisms. The section develops its formulation, analyzes its online-learning interpretation, evaluates retrieval behavior, and describes hardware-efficient and hybrid architectures.

  • Formulation: The gated delta rule combines adaptive gating for memory management with delta updates for effective key-value association learning.The data-dependent gate controls state decay, while the delta structure supports targeted updates.
  • Formulation: The online-learning analysis treats recurrent state updates as closed-form solutions and explains why retention becomes problematic when states saturate.The framework compares regularized objectives that preserve prior state information with mechanisms designed to address saturation.
  • Retrieval analysis: In S-NIAH-1, DeltaNet achieves near-perfect performance across sequence lengths, while Mamba2 degrades beyond 2K sequences and Gated DeltaNet degrades less severely.The setting tests long-term retention with repeated synthetic context.
  • Retrieval analysis: In S-NIAH-2/3, Mamba2 and Gated DeltaNet maintain better performance than DeltaNet because gating filters irrelevant information and reduces memory collision.These tasks use real-world essay context and fixed state size, testing efficient memory management.
  • Retrieval analysis: In S-NIAH-3, Gated DeltaNet outperforms Mamba2 on UUID memorization, supporting the delta rule’s stronger complex-pattern memorization ability.Mamba2’s performance drops quickly when values change from numbers to UUIDs.
  • Architecture and training: Gated DeltaNet adapts DeltaNet’s chunkwise algorithm to enable hardware-efficient training and combines with sliding-window attention or Mamba2 in hybrid models.The token mixer replaces self-attention with gated delta rule mixing within a Llama-style macro architecture.

4 EXPERIMENTS

Experiments compare Gated DeltaNet and hybrid variants with recent recurrent, Transformer, and linear-attention baselines under matched training conditions. Across reasoning, retrieval, long-context, extrapolation, and throughput evaluations, Gated DeltaNet generally improves performance while retaining throughput close to DeltaNet.

  • Experimental setup: All models use 1.3B parameters, 100B FineWeb-Edu tokens, identical training conditions, and a 4K training length; hybrids use a 2K sliding window.The setup uses AdamW, the Llama2 tokenizer, and a 0.5M-token batch size.
  • Common-sense reasoning: Gated DeltaNet consistently outperforms linear-model baselines, including Mamba2 and DeltaNet, on language modeling and zero-shot common-sense reasoning at 400M and 1.3B scales.The hybrid variant further enhances performance.
  • In-context retrieval: Gated DeltaNet outperforms DeltaNet and Mamba2 on real-world retrieval, while hybrid models combining recurrence and attention outperform pure attention models.The improvement over recurrent baselines is smaller than in synthetic retrieval settings, partly because repetition errors dominate these tasks.
  • Length extrapolation: Gated DeltaNet achieves the lowest overall perplexity among RNN models across six long-sequence benchmarks and shows relatively robust extrapolation up to 20K tokens.Hybrid models improve further by using attention for local context modeling, reducing memory-management demands on recurrent components.
  • Long-context understanding: Gated DeltaNet shows consistent advantages on LongBench, especially for single-document QA, few-shot in-context learning, and code tasks.These gains correspond to retrieval, in-context learning, and state-tracking capabilities across the evaluated tasks.
  • Throughput: Gated DeltaNet achieves essentially the same throughput as DeltaNet and is only slightly slower than Mamba2, while H1 and H2 outperform standalone Gated DeltaNet.Mamba2 reaches 2-3K tokens/sec; the hybrid throughput advantage is especially compelling for H1 across sequence lengths.

5 RELATED WORK

Related work develops linear recurrent models through gating, delta-rule memory updates, and more expressive state transitions. These approaches improve memory control or expressiveness but involve limitations in retrieval, real-world performance, or training parallelism.

  • Gated linear RNN: Gated linear RNNs evolved from data-independent decay mechanisms toward data-dependent gating and forgetting mechanisms.This progression includes S4, S5, LRU, RWKV4/5, RetNet, HGRN1/2, Mamba1/2, RWKV6, and GSA.
  • Gated linear RNN: Modern forget gates remove dependence on the previous hidden state and rely only on input data, enabling efficient parallelism across sequence lengths.The absence of a forget gate is identified as a limitation of DeltaNet.
  • Delta rule: The delta rule has theoretical limitations and only moderate real-world performance, while nonlinear recurrence can improve expressiveness at the cost of training parallelism.Other proposed enhancements preserve parallelism through negative eigenvalues or products of Householder transition matrices.
  • Online learning: Alternative online-learning formulations use nonlinear regression or entire-history regression, but nonlinear variants require chunk-level updates or approximations.These approaches include TTT, Titans, and related nonlinear recurrence methods.
  • Hybrid models: Hybrid research also explores interleaving attention layers across layers and combining linear and softmax attention within a single layer.The cited examples include MiniMax-01 and Hybrid Mamba2-Attention.

6 CONCLUSION

The conclusion presents Gated DeltaNet as combining stronger key-value association learning with more adaptive memory clearance, supported by efficient training and hybrid extensions.

  • 6 CONCLUSION: Gated DeltaNet enables better key-value association learning than Mamba2 and more adaptive memory clearance than DeltaNet.The paper reports consistently better empirical results across various tasks.
  • 6 CONCLUSION: The authors extend Yang et al. (2024b)’s parallel algorithm to support hardware-efficient Gated DeltaNet training.Hybrid Gated DeltaNet models achieve higher training throughput and overall performance.

A EXTENDED WY REPRESENTATION FOR GATED DELTA RULE

This appendix introduces the extended WY representation for the gated delta rule and states that its validity is established by mathematical induction.

  • A EXTENDED WY REPRESENTATION FOR GATED DELTA RULE: The appendix restricts the presentation to the first chunk to reduce notation clutter.This is a presentation simplification rather than a change to the method.
  • A EXTENDED WY REPRESENTATION FOR GATED DELTA RULE: It presents the extended WY representation for the state S_t.The supplied passage identifies the representation but does not provide its full expression.
  • A EXTENDED WY REPRESENTATION FOR GATED DELTA RULE: The appendix states that the representation is proved by mathematical induction.The proof claim is given without further proof steps in the supplied passage.

B.1 EVALUATION

The evaluation covers commonsense reasoning, synthetic and real-world in-context retrieval, and long-context understanding, alongside a controlled ablation study.

  • Commonsense reasoning: Commonsense evaluation uses PIQA, HellaSwag, WinoGrande, ARC-easy, ARC-challenge, SIQA, BoolQ, Wikitext, and LAMBADA.These benchmarks follow the evaluation setup of Gu & Dao (2023).
  • In-context retrieval: In-context retrieval evaluation combines RULER’s S-NIAH-1, S-NIAH-2, and S-NIAH-3 synthetic tasks with diverse real-world datasets.The synthetic tasks test passkey, numerical-needle, and word-based-needle retrieval.
  • Long context understanding: Long-context understanding is evaluated on 14 LongBench tasks spanning narrative, scientific, multi-hop, summarization, and specialized tasks.The listed examples include Narrative QA, QasperQA, MultiField QA, HotpotQA, 2WikiMulti QA, Musique, GovReport, QMSum, and MultiNews.
  • Ablation study: The ablation study compares Gated DeltaNet block variants using average perplexity and zero-shot commonsense reasoning accuracy.All models have 400M parameters and are trained for 15B tokens on the same FineWeb-Edu subset.

B.2 ABLATION STUDY

The ablation studies identify which Gated DeltaNet components and design choices most affect performance, while also finding a superior ordering for hybrid architectures.

  • Component ablations: Short convolution and the output gate are crucial for Gated DeltaNet performance, whereas output normalization provides only marginal improvements.
  • Normalization and activation: L2 normalization is essential for optimal performance, while the choice of feature map has less influence.
  • Normalization and activation: SiLU consistently outperforms the other evaluated activation functions.
  • Architectural choices: A head dimension of 128 offers the best trade-off between performance and computational efficiency.
  • Hybrid architectures: Among the evaluated hybrid designs, ordering Mamba2, Gated DeltaNet, and SWA produces superior results.Evaluations use the lm-evaluation-harness, Llama tokenizer, and a shared FineWeb-Edu subset.
Loading 2412.06464v3…