MDN: 面向 Delta 线性注意力的逐步动量并行化方法 MDN: Parallelizing Stepwise Momentum for Delta Linear Attention

MDN的完整方法流程包含以下步骤。第一步是建立递推公式：从latent loss $\mathcal{L} = \frac{1}{2}\|v_t - \tilde{S}_{t-1}^\top k_t\|^2$ 出发，计算梯度 $\nabla \mathcal{L} = -k_t(v_t - \tilde{S}_{t-1}^\top k_t)^\top$，得到动量更新 $M_t = \mu_t M_{t-1} - \eta_t k_t \tilde{v}_t^\top$（其中修正值 $\tilde{v}_t = v_t - S_{t-1}^\top p_t$，$p_t = \alpha_t k_t$）和fast weight更新 $S_t = \alpha_t S_{t-1} - \beta_t M_t$。第二步是推导并行公式：将 $M_t$ 展开为 $M_t = \bar{\mu}_t M_0 - \sum_{i=1}^t \frac{\bar{\mu}_t}{\bar{\mu}_i} k_i \tilde{v}_i^\top$，代入 $S_t$ 后得到 $S_t = \bar{\alpha}_t S_0 - b_t M_0 + \sum_{i=1}^t \gamma_{t,i} k_i \tilde{v}_i^\top$，其中 $b_t = \bar{\alpha}_t c_t$，$\gamma_{t,i} = \frac{\bar{\alpha}_t}{\bar{\mu}_i}(c_t - c_{i-1})$，$c_t = \sum_{i=1}^t \frac{\beta_i \bar{\mu}_i}{\bar{\alpha}_i}$。第三步是几何重排：原求和 $\sum_{i=1}^t \sum_{j=1}^i a_i b_j$ 按行扫描可改写为列扫描 $\sum_{i=1}^t \sum_{j=i}^t a_j b_i$，从而把内积系数与外积分离，使块内可写成 $(Q K^\top) \odot \Gamma \cdot \tilde{V}$ 的形式。第四步是分块实现：每个chunk内并行计算修正值 $\tilde{V}[t] = U[t] - Y[t] S[t] + Z[t] M[t]$，其中 $U[t] = T[t] V[t]$，$Y[t] = T[t] \text{Diag}(\bar{\alpha}_{0 \to C-1}) P[t]$，$Z[t] = T[t] \text{Diag}(b_{0 \to C-1}) P[t]$，$T[t] = \text{Tril}(I + P K^\top \odot \Gamma^-)^{-1}$。第五步是系数对数域计算：为避免数值溢出，所有累积乘积 $\bar{\alpha}_t, \bar{\mu}_t, c_t$ 在log域内用cumsum和logcumsumexp算子计算，$\Gamma[t]$ 矩阵通过下三角log-broadcasting构造（$O(C^2)$ 空间、$O(\log C)$ 并行时间）。第六步是稳定门控：基于二阶系统的特征值分析（$\lambda_\pm = \frac{\alpha_t + \mu_t - \alpha_t\beta_t\eta_t\|k_t\|^2 \pm \sqrt{(\alpha_t + \mu_t - \alpha_t\beta_t\eta_t\|k_t\|^2)^2 - 4\alpha_t\mu_t}}{2}$），约束 $\beta_{\max} = \sin^2(\theta)$，$\alpha_{\max} = \cos^2(\theta)$，$\theta = \arctan(\eta_t \cdot s)$，并把 $\log \mu$ 下界clamp到-2、$\mu_t$ 限制在 $[e^{-1}, 1)$，确保所有特征值落入右半平面（实部非负），避免符号翻转。第七步是架构集成：MDN沿用GDN的主干结构，加入输出端query修正 $q_t = q_t - d k_t$（来自Comba），门控参数 $\alpha_{\log}, \mu_{\log}, \beta, \eta$ 由低秩矩阵 $W_{\alpha/\beta/\mu/\eta} \in \mathbb{R}^{d_{in} \times h}$ 计算（$h \ll d_{in}$，引入可忽略参数），最终通过Triton kernel实现，训练时启用 $\tilde{V}[t]$ 物化策略以加速反向传播。

MDN的技术新颖性体现在三个层面。第一是算法层面的geometric decoupling：传统chunkwise算法（如GDN的WY representation）只处理一阶乘积的累加，而MDN需要处理带有动量系数的双层求和。通过对下三角域 $\{1 \le j \le i \le t\}$ 做行-列扫描重排，把 $O(t^2)$ 依赖的嵌套累加变换为可分形式，这是本文最具技术深度的贡献，也是首次在二阶递推上实现严格因果的分块并行。第二是理论层面的二阶动力学视角：本文首次把动量版线性注意力写为 $2d_k d_v$ 维block状态空间 $[S_t; M_t] = A_t [S_{t-1}; M_{t-1}] + B_t$，推导出 $A_t$ 的闭式谱 $\{\alpha_t, \mu_t, \lambda_+, \lambda_-\}$，揭示了动量引入的复共轭特征值会扩展一阶系统的表达能力（一阶只有实特征值，无法建模阻尼振荡），同时也引入了符号翻转的风险。这一分析不仅给出稳定性条件 $-(1-\alpha_t)(1-\mu_t) \le \alpha_t\beta_t\eta_t\|k_t\|^2 \le (1+\alpha_t)(1+\mu_t)$，更直接指导了 $\beta \le 1 - \alpha$ 的门控约束设计。第三是工程层面的训练优化：通过物化修正值 $\tilde{V}[t]$ 而非完整状态 $S[t], M[t]$，MDN在保留chunkwise并行能力的同时降低了反向传播的内存开销，使训练吞吐量与Mamba2、KDA相当。这三个层面的结合让MDN成为首个在400M/1.3B规模上稳定训练的动量型线性注意力模型。

独立分析MDN的弱点如下。第一，门控超参数数量增加：相比GDN只有 $\alpha, \beta$ 两个门控，MDN引入 $\mu, \eta$ 增加了调参复杂度，$\eta$ 的温度 $\tau = \sqrt{d_{in}/h}$ 和缩放因子 $s$ 的选择对最终性能有显著影响（实验显示 $\eta_t = 2\sigma(\cdot)$ 替代 $\tanh(\cdot) + 1$ 时LAMBADA困惑度从41.62退化到49.10）。改进方向：可借鉴Muon optimizer中的自适应缩放思想，让 $\eta_t$ 根据更新范数自动调整。第二，$\mu$ 下界clamp为固定值-2的鲁棒性问题——表4显示把clamp值改为-1.5、-1.357、-1都会导致LAMBADA困惑度退化到42.65~44.90，这说明稳定性和表达力之间存在trade-off，clamp值需要细致的网格搜索。改进方向：可考虑学习式下界或per-head动态clamp。第三，反向传播依赖修正值物化 $\tilde{V}[t]$ 显著提升了训练吞吐量，但代价是峰值显存占用（400M MDN 39.38 GB vs GDN 37.21 GB），在大batch或长序列训练中可能成为瓶颈。改进方向：可探索chunkwise重计算策略，在不显著牺牲速度的前提下按需重建 $\tilde{V}$。第四，二阶系统的特征值约束 $\beta_t \le 1 - \alpha_t$ 牺牲了一部分表达能力空间——某些理论上可表达的动力学被排除在外。改进方向：可考虑更细粒度的per-channel门控，让不同特征维度拥有不同的稳定边界。第五，MQAR的合成评估中MDN与KDA打平而非显著超越，意味着在纯粹的关联记忆任务上KDA的vector-valued gating可能更有效。改进方向：未来工作可结合MDN的动量机制与KDA的channel-wise门控。

作者明确提出的未来方向：（1）探索更复杂的优化器启发的门控策略（如Nesterov动量、Adam、Muon缩放），并将其纳入线性注意力框架；（2）进一步优化kernel以提升训练吞吐量，目标是对齐GDN/Comba的效率；（3）扩展到更大规模（7B及以上）验证MDN的扩展性。基于论文成果可延伸的方向包括：（1）将MDN应用于多模态长序列建模，如视频、语音、基因组数据，这些场景对长程依赖建模有强需求；（2）探索MDN与状态空间模型（如Mamba3）的混合架构，利用二者的互补优势——Mamba3的强序列建模能力配合MDN的动量记忆机制；（3）研究动量系数 $\mu_t, \eta_t$ 的自适应机制，例如用Muon-like的谱归一化替代当前的固定clamp和温度缩放；（4）将geometric decoupling思想推广到更高阶递推（如三阶动量、Nesterov加速梯度），为非平稳线性递推提供一个通用的并行化模板；（5）研究MDN在强化学习、online learning场景中的应用，因为其递推本质上就是一个online learning过程；（6）探索更细粒度的混合架构设计，如学习式layer placement，让模型自动决定哪些层使用MDN、哪些使用full attention或SSM。

MDN的复现评估整体较为友好。代码已开源，地址为github.com/HuuYuLong/MomentumDeltaNet，仓库中应包含完整的Triton kernel实现。实验方面，训练数据为SlimPajama的100B tokens子集（公开可获取），使用GPT-2 tokenizer（公开）；评测协议严格遵循GDN和Comba的设置，包括commonsense reasoning（LM Evaluation Harness）、in-context retrieval（prefix-linear-attention，2K input）、LongBench、Needle-In-A-Haystack（RULER），所有这些都是公开benchmark。训练配置透明：400M模型用15B tokens、0.5M batch size、AdamW、cosine schedule（peak 3e-4，warmup 0.5B，min 3e-5），1.3B模型用100B tokens、1M batch size、warmup 1B。算力需求方面，1.3B实验在H100上完成，论文中报告了完整的训练吞吐量和显存占用（400M MDN需39.38 GB，1.3B MDN-H需44.89 GB）；复现一个400M实验需要约数天H100/A100算力，1.3B则需要更长时间。复现难度属于中等偏上：（1）Triton kernel的log-cumsum-exp算子实现有一定门槛，需要深入理解chunkwise并行算法的实现细节（论文Appendix E提供了PyTorch-style伪代码可作参考）；（2）门控参数 $\alpha_{\log}, \mu_{\log}, \beta, \eta$ 的初始化（特别是 $\mu$ 的clamp下界-2）对稳定性敏感，移除约束直接导致NaN；（3）超参搜索范围较广（学习率从5e-5到2e-2），MQAR的搜索需要9个学习率×4个hidden dim组合；（4）二阶系统的特征值分析虽然论文给出了闭式解，但实践中验证稳定性约束是否正确实施需要一定的数值计算经验。总体而言，对熟悉flash-linear-attention生态和Triton编程的研究者来说，1-2周可以完成400M规模的复现。