MDN: Parallelizing Stepwise Momentum for Delta Linear Attention
Yulong Huang, Xiang Liu, Hongxiang Huang, Xiaopeng LIN, Zunchang LIU, Xiaowen Chu, Zeke Xie, Bojun Cheng
Abstract
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.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 0ecc76f3-df65-44d0-99a3-988a071d6465Cited by top-tier papers2
- Semantic Integrity Matters: Benchmarking and Preserving High-Density Reasoning in KV Cache CompressionXiang Liu, Zhenheng Tang, Hong Chen, Peijie Dong et al.ICML 2026 · 16 citations
- DiffAdapt: Difficulty-Adaptive Reasoning for Token-Efficient LLM InferenceXiang Liu, Xuming Hu, Xiaowen Chu, Eunsol ChoiICLR 2026 · 15 citations
Builds on32
- Efficiently Modeling Long Sequences with Structured State SpacesAlbert Gu, Karan Goel, Christopher RéICLR 2022 · 3,482 citations
- Transformers are RNNs: Fast Autoregressive Transformers with Linear AttentionAngelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, François FleuretICML 2020 · 2,665 citations
- Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space DualityTri Dao, Albert GuICML 2024 · 1,407 citations
- Going deeper with Image TransformersHugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve et al.ICCV 2021 · 1,279 citations
- Test-Time Training with Self-Supervision for Generalization under Distribution ShiftsYu Sun, Xiaolong Wang, Zhuang Liu, John Miller et al.ICML 2020 · 1,220 citations
Related papers
- Preconditioned DeltaNet: Curvature-aware Sequence Modeling for Linear RecurrencesNeehal Tumma, Noel Loo, Daniela RusICML 2026 · 3 citations
- Gated Delta Networks: Improving Mamba2 with Delta RuleSonglin Yang, Jan Kautz, Ali HatamizadehICLR 2025
- Parallelizing Linear Transformers with the Delta Rule over Sequence LengthSonglin Yang, Bailin Wang, Yu Zhang, Yikang Shen et al.NeurIPS 2024 · 412 citations
- Log-Linear AttentionHan Guo, Songlin Yang, Tarushii Goel, Eric P. Xing et al.ICLR 2026 · 41 citations
- Gated KalmaNet: A Fading Memory Layer through Test-time Ridge RegressionLiangzu Peng, Aditya Chattopadhyay, Luca Zancato, Elvis Nunez et al.CVPR 2026 · 10 citations
