PaperlayerView exact revision on arXiv

Source-linked AI summary

MDN: Parallelizing Stepwise Momentum for Delta Linear Attention

Yulong Huang, Xiang Liu, Hongxiang Huang, Xiaopeng Lin, Zunchang Liu, Xiaowen Chu, Zeke Xie, Bojun Cheng

arXiv:2605.05838v1cs.LGcs.NE

TL;DR

Linear attention models using naive SGD struggle with fine-grained historical details, while stepwise momentum is difficult to parallelize efficiently. MDN addresses this with a causal chunkwise-parallel algorithm and constrained gating, outperforming existing linear-attention baselines across downstream tasks.

  • Problem

    Existing linear-attention mechanisms use naive SGD updates and struggle with fine-grained historical details, while stepwise momentum remains difficult to parallelize efficiently.

  • Method

    MDN combines a causal chunkwise-parallel stepwise-momentum algorithm with dynamical-systems-guided constrained gating for stable linear attention.

  • Results

    Momentum DeltaNet outperforms existing linear-attention baselines across a range of downstream tasks while achieving comparable training efficiency to Mamba2 and KDA.

  • Takeaways & Limitations

    MDN supports a practical balance between expressive dynamics and efficiency for linear attention with stepwise momentum.

  • Takeaways & Limitations

    Unconstrained second-order coefficients can cause catastrophic numerical failures, so MDN relies on gating constraints to avoid unstable sign-flipping dynamics.

Abstract

from arXiv · show

Linear Attention (LA) offers a promising paradigm for scaling large language models (LLMs) to long sequences by avoiding the quadratic complexity of self-attention. Recent LA models such as Mamba2 and GDN interpret linear recurrences as closed-form online stochastic gradient descent (SGD), but naive SGD updates suffer from rapid information decay and suboptimal convergence in optimization. While momentum-based optimizers provide a natural remedy, they pose challenges in simultaneously achieving training efficiency and effectiveness. To address this, we develop a chunkwise parallel algorithm for LA with a stepwise momentum rule by geometrically reordering the update coefficients. Further, from a dynamical systems perspective, we analyze the momentum-based recurrence as a second-order system that introduces complex conjugate eigenvalues. This analysis guides the design of stable gating constraints. The resulting model, Momentum DeltaNet (MDN), leverages Triton kernels to achieve comparable training throughput with competitive linear models such as Mamba2 and KDA. Extensive experiments on the 400M and 1.3B parameter models demonstrate consistent performance improvements over strong baselines, including Transformers, Mamba2 and GDN, across diverse downstream evaluation benchmarks. Code: https://github.com/HuuYuLong/MomentumDeltaNet .

1. Introduction

Linear Attention addresses Transformers’ quadratic sequence-length scaling but struggles with fine-grained historical retrieval. MDN introduces causally exact, chunkwise-parallel stepwise momentum with dynamical-systems-guided gating to improve efficiency and stability.

  • Motivation: Linear Attention avoids Transformers’ quadratic O(L^2) scaling with sequence length but remains limited in long-context scalability and fine-grained historical retrieval.These limitations are reflected in constrained context-retrieval capability.
  • Motivation: Momentum accumulates gradient information in an auxiliary state, attenuating noise, smoothing updates, stabilizing optimization, and extending accessible historical information in Linear Attention.This motivation follows from interpreting Linear Attention recurrences as online optimization.
  • Problem: Prior blockwise momentum improves hardware utilization but breaks strict temporal causality, causing training–inference mismatch and weaker intra-block dependency modeling.Larger blocks exacerbate the mismatch by weakening dependencies within each block.
  • Method: MDN proposes a chunkwise-parallel algorithm for stepwise momentum that decouples recursive coefficients geometrically, preserving strict causality during efficient parallel computation.Stepwise momentum uses block size 1, maintaining exact consistency between parallel training and decoding.
  • Method: MDN models momentum as a second-order dynamical system, exposing complex eigenvalues that guide constrained gating mechanisms.The dynamical-systems formulation is used to design stable gating constraints.

2. Notation and Preliminaries

This section introduces self-attention and linear-attention formulations, emphasizing the latter’s lower-complexity chunkwise computation and extensions with decay and Delta updates.

  • Self-Attention: Self-attention requires O(L) memory per step for its KV cache and O(L2) aggregate computational complexity.The causal mask restricts each token to attend only to preceding positions.
  • Linear Attention: Linear attention removes the Softmax operator to avoid self-attention’s quadratic cost, although its fully parallel causal form remains quadratic in L.The formulation represents the recurrent state as a fast-weights matrix.
  • Chunkwise Parallelism: Chunkwise parallel linear attention achieves O(LCd + Ld2) training complexity, interpolating between fully parallel computation at C = L and recurrence at C = 1.When L ≫C, this is lower than the fully parallel O(L2d) cost.
  • Decay Mechanisms: Data-dependent decay selectively forgets historical information through scalar αt ∈(0, 1), and its cumulative products define a decay-aware causal mask.The decay formulation extends naturally to chunkwise algorithms under the SSD framework.
  • Delta Rule: Delta-based linear attention dynamically updates value representations and becomes computationally tractable through an efficient chunkwise algorithm.GDN combines the Delta rule with decay, while later work extends scalar gating to vector-valued gating.

3. Method

MDN incorporates stepwise momentum into Linear Attention through an exact chunkwise-parallel recurrence derived by coefficient reordering. Its second-order dynamical formulation enables complex-eigenvalue expressivity, while stability-aware gating and Triton-oriented implementation address numerical and efficiency challenges.

  • Stepwise Momentum Recurrence: MDN constructs recurrent and exact chunkwise-parallel updates for Linear Attention by incorporating optimizer momentum, decay, and learning-rate parameters.The method begins with the Stepwise Momentum recurrence and derives its parallel form before presenting the architecture.
  • Parallel Formulation: Reordering the shared lower-triangular index domain decouples coefficients from outer products, removing the nested-summation obstacle to parallelization.The transformation changes the traversal from row-wise to column-wise while preserving equality.
  • Coefficient Chunkwise: Chunkwise coefficient computation uses log-domain prefix operations, with O(log C) within-chunk prefix-sum complexity and O(1) log-cumsum-exp time per chunk.The fixed chunk size permits parallel log-cumsum-exp computation with O(C^2) space per chunk.
  • Second Order Dynamics and Expressivity: The momentum rule forms a second-order system admitting complex conjugate eigenvalues, enabling damped oscillatory behavior beyond first-order real-valued decay dynamics.This expands the eigenvalue space and addresses the limitation that real-valued systems cannot capture oscillatory dependencies.
  • Stability via Quadrant Constraint: MDN introduces stability-aware gating because unconstrained second-order coefficients can cause NaNs through sign-flipping modes associated with negative-real-part eigenvalues.The proposed parameterization balances expressivity and numerical robustness by constraining the relevant dynamics.
  • MDN Architecture: The gating design clamps µlog to prevent vanishing momentum and constrains αmax + βmax = 1, while using low-rank trainable matrices with h ≪ din for negligible parameter overhead.The decay function matches GDN and Mamba2, and tanh(·) + 1 ∈ (0, 2) controls the mean of the parameterized quantity.

4. Experiments

MDN improves retrieval, reasoning, and long-context performance over linear-attention baselines while preserving competitive efficiency. Ablations attribute gains to momentum, stability-aware constraints, and more active fast-weight dynamics.

  • MQAR Retrieval: MDN achieves strong MQAR retrieval accuracy across sequence lengths from 256 to 2k tokens, competing with KDA.Models are trained with sequences up to 256 tokens and 4–64 key-value pairs, then evaluated at longer contexts.
  • Downstream Results: MDN delivers the strongest average reasoning performance with competitive perplexity and consistently outperforms other linear baselines on in-context retrieval.The evaluation covers commonsense reasoning, in-context retrieval, long-context modeling, and Needle-In-A-Haystack tasks.
  • Long-context Modeling: MDN achieves the highest average LongBench score, with particularly strong improvements on Code and Summarization tasks.This evaluation uses 1.3B-parameter models and 16K-length contexts.
  • Needle-In-A-Haystack: 38.60 on MK, 35.15 on MQ, and 27.60 on MV are achieved at 8k context length in challenging multi-needle NIAH settings.MDN consistently improves accuracy across tasks, especially beyond the 4k training context length.
  • Efficiency Analysis: MDN has decoding latency nearly identical to GDN and Comba, while achieving training throughput comparable to Mamba2 and KDA despite dual-state computation.The implementation uses Triton kernels, with correction values materialized at manageable memory overhead.
  • Ablation Study: Removing output correction still leaves MDN ahead of GDN and Comba, whereas removing stability constraints causes divergence and weakening other constraints degrades performance.The momentum lower bound µlog_min = −2 provides the best overall trade-off, while a 7:1 linear/full-attention ratio improves LM average with competitive retrieval accuracy.

5. Conclusion

Momentum DeltaNet scales stepwise momentum in linear attention with an efficient chunkwise parallel algorithm and constrained gating. Experiments show it outperforms existing linear-attention baselines across downstream tasks while balancing expressive dynamics and efficiency.

  • Conclusion: Momentum DeltaNet scales stepwise momentum in linear attention through an efficient chunkwise parallel algorithm.The algorithm addresses computational bottlenecks in scaling stepwise momentum.
  • Conclusion: Constrained gating mechanisms help balance expressive dynamics and efficiency.The model integrates constrained gating while resolving computational bottlenecks.
  • Conclusion: Momentum DeltaNet outperforms existing linear attention baselines across a range of downstream tasks.The conclusion attributes this finding to the reported experimental results.

Impact Statement · A. Notation

The paper’s impact is primarily methodological, advancing efficient sequence modeling and linear attention for large language models using public academic benchmarks without introducing sensitive personal-data datasets. Its notation defines tensor conventions, triangular masks, chunk indexing and state continuity, and cumulative decay products.

  • Impact Statement: The work advances efficient sequence modeling and linear attention for large language models as a primarily methodological contribution.Experiments use public academic benchmarks and do not introduce datasets containing personal or sensitive information.
  • A. Notation: Matrices use bold uppercase letters, column vectors use bold lowercase letters, and learnable parameter matrices use italic uppercase letters.The t-th row vector of Q is denoted q⊤t.
  • A. Notation: M and M− denote lower-triangular masks with and without diagonal elements, respectively.The notation distinguishes whether diagonal elements are included.
  • A. Notation: A sequence of length L is split into L/C chunks of length C, with each chunk represented as □[t] ∈ R^C×d for □ ∈ {Q, K, V, · · ·}.The r-th element of chunk t is indexed by □[t]r = □tC+r, with t ∈ [0, L/C) and r ∈ [1, C].
  • A. Notation: Chunk state matrices are re-indexed so that each chunk’s initial state is the last state of the previous chunk.This establishes continuity between adjacent chunks.
  • A. Notation: The bar notation denotes cumulative products, with ¯x_r defined through products of x_k and ratios of cumulative terms.For chunkwise notation, ¯α[t] abbreviates the case j = 1 and i = C.

B. Extended Related Work

Extended related work organizes linear attention around gating, correction-based online learning, and expanding-memory recurrences. It traces a progression from vanilla identity gating through DeltaNet variants with decay, channel-wise gating, query correction, and composed memory dynamics.

  • Linear Attention with Gating: Vanilla linear attention uses identity gating, G_t = I, within a general gated formulation that accumulates a recurrent memory state.The formulation distinguishes the accumulated state S_t from the memory gate G_t.
  • Linear Attention with Correction: Correction-based linear attention treats the recurrent state as fast-weight memory optimized through online learning, addressing saturation and input-noise sensitivity in Hebbian accumulation.The correction framework interprets the hidden state as a Fast Weight Programmer.
  • Linear Attention with Correction: The Delta rule corrects prior associations using reconstruction error, implicitly orthogonalizing memory when G_t = I and improving utilization of limited capacity.The correction term is e_t^v = v_t − S_{t−1}^⊤k_t.
  • Linear Attention with Correction: GDN adds input-dependent scalar decay, KDA adds channel-wise diagonal gating, and Comba introduces query correction on top of Delta-rule updates.GDN sets G_t = α_tI; KDA sets G_t = Diag(α_t); Comba uses e_t^q = q_t − d k_t.
  • Linear Attention with Expanding Memory: Expanding-memory approaches compose multiple gated or Delta-based recurrences: ABC couples two vanilla modules, GSA adds input-dependent gating, and MesaNet uses conjugate gradient for a test-time linear loss.MesaNet can be viewed as a dual recurrence derived from an in-context regression objective.

C. Chunkwise Parallel Derivation for Momentum Delta Rule. · D. Coefficients Chunkwise Parallelization

The paper derives a general parallel formulation of the momentum delta rule by geometrically reordering recurrence coefficients, then converts it into an algebraically equivalent chunkwise implementation. It further computes chunk coefficients in the log domain with parallel prefix operations and Triton-based log-sum-exp for numerical stability.

  • C. Chunkwise Parallel Derivation for Momentum Delta Rule.: The momentum delta recurrence is expanded into a general parallel formulation that supports chunkwise execution.The derivation begins from the recurrent momentum formulation and explicitly considers the case µ_t ≠ 0.
  • C. Chunkwise Parallel Derivation for Momentum Delta Rule.: Geometric reordering exchanges row-wise and column-wise accumulation over the same lower-triangular index domain, exposing dependencies for parallelization.The reordering traverses {(i, j) | 1 ≤ j ≤ i ≤ t} along orthogonal directions.
  • C. Chunkwise Parallel Derivation for Momentum Delta Rule.: The resulting coefficient construction transforms the correction recursion into an iteration form that can be solved in parallel.The formulation uses a strict lower-triangular matrix Γ− and collects terms to solve the correction value e_vt.
  • C. Chunkwise Parallel Derivation for Momentum Delta Rule.: Although the full parallel formulation is inefficient, each chunk remains algebraically equivalent and decomposes updates into inter-chunk and intra-chunk components.Within each chunk, coefficients are computed and the output combines inter-chunk state contributions with intra-chunk correction terms.
  • C. Chunkwise Parallel Derivation for Momentum Delta Rule.: The geometric decoupling provides a reusable template for parallelizing non-stationary linear recurrences with structured dependencies and potentially incorporating advanced optimization-inspired rules.The paper identifies possible extensions beyond MDN, including Nesterov-style updates.
  • D. Coefficients Chunkwise Parallelization: Accumulated coefficients are computed chunkwise in parallel in the log domain using lower-triangular matrices and log-domain intermediate quantities.The coefficient procedure defines Γ[t], A_log[t], and C_log[t] within each chunk.
  • D. Coefficients Chunkwise Parallelization: The within-chunk cumsum has O(log C) parallel complexity, while Triton implements log-cumsum-exp with masking, row-wise maximum subtraction, and log-domain accumulation.The implementation broadcasts vectors to a masked lower-triangular matrix to improve numerical stability.

E. Pytorch-like Pseudo Code for Recurrent and Chunkwise MDN · F. Stability Condition of Gated Momentum Dynamics

Sections E–F provide Pytorch-like recurrent and chunkwise MDN pseudocode, then formulate gated momentum dynamics in state-space form and state a necessary-and-sufficient stability condition based on eigenvalue magnitudes.

  • E. Pytorch-like Pseudo Code for Recurrent and Chunkwise MDN: The recurrent MDN function accepts tensor states, gates, optimizer coefficients, optional scale and initial states, and an output-final-state flag.The pseudocode converts q, k, v, p, log_alpha, log_mu, and beta to float32 before proceeding.
  • E. Pytorch-like Pseudo Code for Recurrent and Chunkwise MDN: The chunkwise computation builds shifted momentum terms and lower-triangular gamma masks before forming attention as (p @ k.transpose(-1, -2)) * gamma_mask.It also constructs attn_inv by negating attn and updates prior positions within each chunk.
  • F. Stability Condition of Gated Momentum Dynamics: The stability analysis reformulates the coupled momentum-state and fast-weight-state recurrences as a unified discrete state-space dynamic representation.The section begins from the recursive updates for Mt and St.
  • F. Stability Condition of Gated Momentum Dynamics: Substituting Mt into the St update yields an explicit representation of the coupled dynamics.This substitution makes the dependence between the momentum and fast-weight states explicit.
  • F. Stability Condition of Gated Momentum Dynamics: The formulation keeps αt and βt input-dependent for expressivity while treating optimizer-related coefficients such as µt and ηt as fixed scalars.The resulting system is expressed compactly in block matrix form.
  • F. Stability Condition of Gated Momentum Dynamics: The block-matrix dynamics use a state transition matrix A, whose entries define the system’s transition structure.The passage introduces A before giving its specific form.
  • F. Stability Condition of Gated Momentum Dynamics: The matrix A has a closed-form spectral characterization through its spectrum.The cited passage introduces the eigenvalue-based description without showing the spectrum in the supplied excerpt.
  • F. Stability Condition of Gated Momentum Dynamics: |λ| ≤1 holds for all eigenvalues if and only if |αt| ≤1, |µt| ≤1, ∥kt∥2 = 1, and the remaining stated condition is satisfied.The supplied passage presents these requirements as necessary and sufficient for stability.

G. Additional Experiment Details

The experiments cover MQAR, language modeling at 400M and 1.3B parameters, and diverse downstream and retrieval-intensive in-context benchmarks. Training and evaluation use longer or challenging sequence settings to assess model performance.

  • MQAR Experiments Details: MQAR models train on 64–256-token sequences with 4–64 key–value pairs, then evaluate on 512–2048-token sequences with 32–512 pairs.The setup largely follows Arora et al. (2023a).
  • Language Model Experiments Details: Language models are trained from scratch at 400M and 1.3B parameters using 4k sequences, identical configurations, and AdamW.The models use 15B and 100B training tokens with batch sizes of 0.5M and 1M, respectively; the learning rate peaks at 3 × 10^-4 after warmup under cosine decay.
  • Language Model Experiments Details: Evaluation includes WikiText perplexity plus HellaSwag, LAMBADA, WinoGrande, ARC-Easy, ARC-Challenge, BoolQ, and SciQA.These downstream tasks cover commonsense reasoning and question answering and follow Yang et al. (2025).
  • Language Model Experiments Details: Retrieval-intensive in-context evaluation uses 2K input tokens and benchmarks SWDE, SQuAD, FDA, TQA, NQ, and DROP.The protocol follows prefix-linear-attention and uses minimally transformed benchmarks designed for non-instruction-tuned models.

H. Additional Experiments

This section reports complete additional ablation studies on 400M models. The experiments use a fixed training setup while modifying one variable at a time.

  • Additional ablation studies: Table 6 reports the complete results of the additional ablation studies.The reported experiments concern the 400M models.
  • Additional ablation studies: The additional ablation studies evaluate 400M models.
  • Additional ablation studies: Each experiment follows the same training setup and modifies only one variable at a time.

I. Limitations and Future Work

MDN’s scaling beyond 1.3B parameters remains unverified, and its current implementation has lower throughput than optimized first-order baselines. Future work targets system optimization, broader hybrid-attention studies, and applications to other long-sequence modalities.

  • Scaling limitations: MDN has only been evaluated at 400M and 1.3B parameters, leaving its behavior at 7B and larger scales unclear.Larger-scale experiments were precluded by limited compute resources.
  • System optimization: The additional momentum state and materialized correction values make MDN slower to train than GDN and Comba, though comparable to Mamba2 and KDA.The correction values are materialized in the chunkwise training implementation to improve backward efficiency.
  • System optimization: Future system work will pursue more optimized kernels and memory-efficient backward strategies to improve MDN’s training throughput.The current implementation leaves room for further system optimization.
  • Hybrid designs: MDN’s hybrid experiments cover only a small set of linear/full-attention ratios, motivating broader studies of layer placement, ratios, and gating parameterizations.These studies may further improve the efficiency–performance trade-off.
  • Beyond language modeling: Future applications include speech, video, time-series, and genomics, where efficient long-range dependency modeling is important.These modalities extend evaluation beyond language modeling.
Loading 2605.05838v1…