arXiv:2605.05838cs.LGcs.NE2026-05被引 2

提出并行化步进动量的Delta线性注意力,提升长序列模型训练效率与稳定性。

MDN: Parallelizing Stepwise Momentum for Delta Linear Attention

论文配图:MDN: Parallelizing Stepwise Momentum for Delta Linear Attention
图 1 · 摘自论文原文
  • 通过几何重排更新系数,实现步进动量的分块并行计算
  • 在400M和1.3B参数模型上超越Transformer、Mamba2等基线性能
  • 从动力系统视角设计稳定门控约束,兼顾训练效率与收敛性

线性注意力(LA)通过避免自注意力的二次复杂度,为扩展大语言模型处理长序列提供了前景。近期模型如Mamba2和GDN将线性递归解释为闭式在线随机梯度下降(SGD),但朴素的SGD更新存在信息快速衰减和优化收敛不佳的问题。虽然基于动量的优化器可缓解此问题,但难以同时保证训练效率与有效性。为此,本文提出一种基于几何重排更新系数的分块并行算法,实现步进动量规则下的线性注意力。进一步从动力系统角度分析,将动量递归视为二阶系统,引入复共轭特征值,并据此设计稳定门控约束。所提出的模型Momentum DeltaNet(MDN)利用Triton内核,在训练吞吐量上与Mamba2和KDA等先进线性模型相当。在400M和1.3B参数模型上的大量实验表明,MDN在多个下游评估基准上持续优于包括Transformer、Mamba2和GDN在内的强基线。代码已开源:https://github.com/HuuYuLong/MomentumDeltaNet。

原文摘要 · Abstract (English)

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 .

线性注意力动量优化并行计算长序列建模

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。