MDN: Parallelizing Stepwise Momentum for Delta Linear Attention
:::info Stub — Full Engineering Breakdown Coming This paper has a linked code implementation and was featured on Hugging Face Papers with 1 upvotes. A full breakdown with production viability rating, implementation notes, and honest limitations is being written. Subscribe to AI Letters → :::
| Authors | Yulong Huang et al. |
| Year | 2026 |
| HF Upvotes | 1 |
| arXiv | 2605.05838 |
| Download | |
| Code | https://github.com/HuuYuLong/MomentumDeltaNet |
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 .
Engineering Breakdown
Plain English
This paper proposes MDN, a new algorithm that adds momentum-based optimization to Linear Attention models while keeping them parallelizable and efficient. The key contribution is a geometric reordering technique that lets you apply stepwise momentum updates in chunks rather than sequentially, avoiding the information decay problem that plagues naive linear recurrence implementations without sacrificing the speed benefits of linear attention.
Key Engineering Insight
By reordering update coefficients geometrically, you can parallelize momentum updates across chunks instead of computing them step-by-step, which is the critical engineering trick that lets you get both optimization quality (momentum helps convergence) and computational efficiency (parallelization avoids sequential bottlenecks).
Why It Matters for Engineers
Linear attention models are becoming practical alternatives to transformers for long sequences because they scale linearly instead of quadratically, but they've had convergence and information retention issues. If MDN delivers production-ready momentum benefits without sacrificing speed, it directly improves model quality on long-context tasks like document processing or code understanding without increasing inference latency or memory costs.
Research Context
Recent linear attention work (Mamba2, GDN) reframed recurrences as SGD updates but suffered from poor convergence and information decay. MDN advances this by adding momentum—a well-proven optimization technique—while solving the parallelization problem that momentum naturally introduces. This bridges the gap between optimization theory and efficient implementation, potentially making linear attention competitive with transformers on both speed and quality metrics.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::