MDN: 델타 선형 어텐션을 위한 단계적 모멘텀 병렬화

MDN: Parallelizing Stepwise Momentum for Delta Linear Attention

May 7, 2026

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

cs.AI

초록

선형 주의(Linear Attention, LA)는 자기 주의의 이차 복잡성을 피함으로써 대규모 언어 모델(LLM)이 긴 시퀀스를 확장할 수 있는 유망한 패러다임을 제공한다. 최근 Mamba2 및 GDN과 같은 LA 모델은 선형 재귀를 폐쇄형 온라인 확률적 경사 하강법(SGD)으로 해석하지만, 단순한 SGD 업데이트는 빠른 정보 소멸과 최적화에서의 최적 이하 수렴을 겪는다. 모멘텀 기반 최적화기가 자연스러운 해결책을 제공하지만, 이는 훈련 효율성과 효과성을 동시에 달성하는 데 어려움을 제기한다. 이를 해결하기 위해, 우리는 업데이트 계수를 기하학적으로 재배열하여 단계별 모멘텀 규칙을 갖는 LA를 위한 청크 단위 병렬 알고리즘을 개발한다. 나아가 동역학 시스템 관점에서, 모멘텀 기반 재귀를 복소 켤레 고유값을 도입하는 2차 시스템으로 분석한다. 이 분석은 안정적인 게이팅 제약 조건의 설계를 안내한다. 결과 모델인 Momentum DeltaNet(MDN)은 Triton 커널을 활용하여 Mamba2 및 KDA와 같은 경쟁력 있는 선형 모델과 유사한 훈련 처리량을 달성한다. 400M 및 1.3B 파라미터 모델에 대한 광범위한 실험은 Transformer, Mamba2 및 GDN을 포함한 강력한 기준선 대비 다양한 다운스트림 평가 벤치마크에서 일관된 성능 향상을 입증한다. 코드: https://github.com/HuuYuLong/MomentumDeltaNet

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 .