MDN: Parallelizing Stepwise Momentum for Delta Linear Attention

📄 arXiv: 2605.05838v1 📥 PDF

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

分类: cs.LG, cs.NE

发布日期: 2026-05-07

🔗 代码/项目: GITHUB

💡 一句话要点

提出Momentum DeltaNet (MDN)，通过分块并行动量机制优化线性注意力模型

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 线性注意力 长序列建模 动量优化 并行计算 动力系统 大语言模型

📋 核心要点

  1. 现有线性注意力模型基于朴素SGD更新，导致长序列建模中信息衰减过快，且优化过程收敛性不佳。
  2. 提出MDN模型，通过几何重排系数实现步进动量的分块并行计算，并利用动力系统分析设计稳定门控。
  3. 实验证明MDN在保持高训练吞吐量的同时，在400M和1.3B参数规模下显著超越了Mamba2和GDN等基线。

📝 摘要（中文）

线性注意力（LA）通过规避自注意力的二次复杂度，为大语言模型（LLM）处理长序列提供了有效范式。近期如Mamba2和GDN等模型将线性递归解释为闭式在线随机梯度下降（SGD），但朴素的SGD更新存在信息衰减过快及优化收敛次优的问题。虽然基于动量的优化器能改善这一现状，但如何在保持训练效率的同时提升效果仍具挑战。为此，本文开发了一种基于几何重排更新系数的步进动量分块并行算法。此外，从动力系统视角出发，将动量递归分析为引入复共轭特征值的二阶系统，并据此设计了稳定的门控约束。由此构建的Momentum DeltaNet (MDN) 利用Triton内核实现了与Mamba2等模型相当的训练吞吐量。在400M和1.3B参数规模下的实验表明，MDN在多项下游基准测试中均优于Transformer、Mamba2及GDN等强基线模型。

🔬 方法详解

问题定义：线性注意力模型常被建模为在线SGD过程，但这种一阶递归结构在处理长序列时，容易导致梯度信息迅速衰减，且缺乏动量机制导致优化收敛速度和精度受限，难以在长程依赖建模中达到最优表现。

核心思路：引入动量机制以增强递归更新的稳定性与记忆能力。通过将动量更新转化为二阶动力系统，利用几何重排技术解决动量计算在并行化过程中的依赖性问题，从而在保持线性复杂度的同时实现高效训练。

技术框架：MDN采用分块并行算法，将长序列划分为多个块（chunk），在块内利用Triton内核进行高效并行计算，块间通过状态传递实现递归。整体架构结合了动量更新规则与门控机制，确保模型在长序列上的信息流稳定性。

关键创新：核心创新在于将动量递归视为二阶系统，通过分析复共轭特征值，设计了能够抑制发散并保持长程记忆的稳定门控约束。此外，通过几何重排系数，成功将原本串行的动量更新转化为可并行计算的算子。

关键设计：采用了基于Triton的定制化算子实现，确保了计算效率；引入了针对二阶动力系统设计的门控参数，通过约束特征值分布，平衡了模型对近期信息的敏感度与远期信息的保持能力，从而提升了整体优化性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

MDN在400M和1.3B参数规模下进行了广泛评估。实验结果显示，MDN在多项下游任务中表现优于Mamba2、GDN及标准Transformer。在保持与Mamba2相当的训练吞吐量的同时，MDN通过引入动量机制，在长序列建模的准确率和收敛速度上均取得了显著的性能提升。

🎯 应用场景

MDN适用于需要处理超长上下文的大语言模型场景，如长文档分析、长视频理解、复杂代码库推理及大规模时序数据预测。其高效的并行计算特性使其在资源受限的训练环境下具有显著优势，为构建高性能、长记忆的轻量化LLM提供了新的技术路径。

📄 摘要（原文）

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 .