Activation-Keyed Momentum: An Anisotropic Momentum Update via the Delta Rule¶

会议: NeurIPS 2026
arXiv: 2608.19491
代码: https://github.com/euijinh/AK-Momentum
领域: 优化/理论
关键词: 动量更新、Delta法则、各向异性、隐式预条件、大语言模型预训练

一句话总结¶

本文提出基于 Delta 法则的激活键控动量（AK-Momentum），将线性层输入激活视作键、反传误差视作值，构建输入密度自适应的非各向同性动量遗忘机制，在无需矩阵求逆与持久显存的前提下实现隐式输入端曲率校正，在大模型预训练中节省高达 46.6% 的达到目标验证损失步数。

研究背景与动机¶

现代深度神经网络优化器（如 SGD-momentum、AdamW、Shampoo 与 Muon）普遍依赖一阶动量缓冲区来平滑随机梯度噪声并提供搜索惯性。然而，这些优化器底层的一阶动量无一例外均沿用指数滑动平均（EMA）机制，以固定的标量衰减率对所有参数方向一视同仁地遗忘历史梯度。这种对所有子空间施加均匀遗忘的假设与深度网络训练中高度各向异性（anisotropic）的特征几何产生了本质冲突：在特征空间中，输入的协方差谱通常由少数高能量的主方向主导，伴随大量极少被激发的平坦尾部方向。

统一标量衰减率在特征谱的两端均会引发严重的次优性：在高频出现的输入方向上，局部损失曲率随权重更新剧烈漂移，最新梯度携带关键的新信息，但统一 EMA 仍保留了过多陈旧梯度的惯性；而在极少激发的罕见方向上，本应通过长期累积平均来抑制方差，但其历史有效信息却在下一批罕见样本到来前被机械衰减殆尽。尽管 K-FAC、Shampoo、SOAP 以及 Muon 等预条件或正交化算法致力于校正曲率各向异性，但它们均是将额外变换外挂在标准 EMA 动量之外，从未触碰动量累积本身的更新逻辑。

更深层的结构性矛盾在于：线性层的梯度天然分解为输入激活与反向传播误差的外积，二者自然构成关联记忆中的“键（Key）”与“值（Value）”，指明了当前修正所对应的输入子空间与输出更新量。标准 EMA 将梯度打平为全矩阵机械相加，直接丢失了该键值配对结构以及方向被查询的频率信息。本文的核心 idea 是：将线性层输入激活作为键、输出误差作为值，引入经典联想记忆的 Delta 法则更新动量缓冲区，使得频繁查询的方向快速刷新陈旧记忆、罕见方向长期保留，在不求逆矩阵和零持久显存开销下实现隐式输入端自适应自然梯度预条件。

方法详解¶

整体框架¶

AK-Momentum 将一阶动量缓冲区重新构想为一个在线线性联想记忆矩阵，其目标是建立从输入特征空间到输出误差信号的自适应映射。在反向传播计算线性层梯度时，算法直接利用前向保存的输入激活计算归一化键向量，并结合反传输出误差，按正则化 Delta 规则对动量矩阵执行自适应擦除与写入操作。

整体数据流从前向传播捕获激活开始，经过反向传播计算误差与外积梯度，随后通过逐层特征二阶矩收缩动量矩阵在当前子空间上的旧分量，并写入最新误差响应，最后无缝交付给基座优化器（如 AdamW 的二阶矩缩放或 SGD 步长更新）。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["前向传播输入激活 X 与反传误差 Δ"] --> B["键向量归一化与二阶矩估计<br/>计算单位范数键并捕获方向密度"]
    B --> C["Delta法则各向异性记忆传输<br/>对齐当前键的方向收缩擦除旧动量"]
    C --> D["正则化维纳预测子稳态更新<br/>写入最新误差构造隐式输入端预条件"]
    D --> E["最适配算子收缩调度与零持久显存落地<br/>双收缩路径执行并传递至基座优化器"]

关键设计¶

1. 键向量归一化与二阶矩估计：解耦输入尺度与宽度以保障 μP 超参无缝迁移

原始 Delta 规则直接应用未归一化的输入特征 \(x_t\) 作为键，会导致其收缩算子的谱范数强依赖于特征绝对范数及网络隐藏层宽度 \(n\)（即 \(\|x_t x_t^\top\|_{\mathrm{op}} \sim n\)），迫使步长系数 \(\eta = \mathcal{O}(1/n)\) 且在训练各阶段极易数值发散。为消除尺度敏感性，AK-Momentum 引入逐样本单位范数键向量 \(\hat{x}_t = x_t / \|x_t\|_2\)。

对于微批次中的 \(N\) 个 token 输入矩阵 \(X \in \mathbb{R}^{N \times n}\)，行归一化得到 \(\hat{X} \in \mathbb{R}^{N \times n}\)，并定义归一化输入二阶矩 \(\hat{\Sigma} = \frac{1}{N} \hat{X}^\top \hat{X}\)。根据性质，\(\operatorname{tr}(\hat{\Sigma}) = 1\)，其特征值 \(\hat{\lambda}_i \in [0, 1]\) 严格反映输入特征的方向性分布密度（directional density），而非特征能量尺度。这使得算子谱范数恒等于 1，满足宽度无关的稳定条件 \(\eta < \beta\)，保证了最大更新参数化（\(\mu\text{P}\)）下 \(\eta\) 与学习率超参数在不同模型尺寸间的零样本直接迁移。

2. Delta法则各向异性记忆传输：按输入激活访问频率自适应调整有效记忆时域

针对标准 EMA 标量衰减无法适应各向异性特征分布的问题，AK-Momentum 构造了状态相关的非各向同性动量递推更新式：

\[M_t = \beta M_{t-1} + \eta (\delta_t - M_{t-1} \hat{x}_t) \hat{x}_t^\top = M_{t-1} (\beta I_n - \eta \hat{x}_t \hat{x}_t^\top) + \eta \delta_t \hat{x}_t^\top\]

右乘算子 \(\beta I_n - \eta \hat{x}_t \hat{x}_t^\top\) 构成了数据驱动的各向异性遗忘因子。在期望意义下，沿着输入协方差矩阵的特征基 \(u_i\)（对应特征值 \(\hat{\lambda}_i\)），历史动量的期望传输衰减率转变为 \(\tilde{\beta}_i = \beta - \eta \hat{\lambda}_i\)。由此导出各子方向上的有效记忆时域（memory horizon）：

\[\tau_i = \frac{1}{1 - \tilde{\beta}_i} = \frac{1}{(1 - \beta) + \eta \hat{\lambda}_i}\]

对于被高频访问的活跃特征方向（\(\hat{\lambda}_i\) 较大），\(\tau_i\) 显著缩短，动量能够以更快速度刷新、快速冲刷旧梯度；而对于几乎不被激发的长尾罕见方向（\(\hat{\lambda}_i \approx 0\)），衰减率回退至基线 EMA 水平 \(\tau_i \approx \frac{1}{1 - \beta}\)，完整保留长期历史信息并持续压制方差。在非平稳训练目标漂移下，其稳态滞后界被严格证明小于 EMA。

3. 正则化维纳预测子稳态更新：无需矩阵求逆的隐式输入端自然梯度预条件

标准 EMA 在准静态下收敛至总体梯度的无偏估计 \(\bar{G}\)（即最小化权重空间 MSE 目标 \(J_w\)）；而纯 Delta 规则对应于最小化函数空间预测误差 \(J_f(M) = \frac{1}{2} \mathbb{E}[\|\delta - M x\|_2^2]\)，其稳态解为最优线性维纳预测子 \(\bar{G} \Sigma_x^{-1}\)。AK-Momentum 结合均匀衰减 \(\beta < 1\) 引入 Tikhonov 正则化，其稳态解析定点为：

\[M^\star = \eta \bar{G} \bigl((1 - \beta) I_n + \eta \hat{\Sigma}\bigr)^{-1} = \bar{G} (\mu I_n + \hat{\Sigma})^{-1}, \quad \text{其中 } \mu := \frac{1 - \beta}{\eta}\]

在特征基 \(u_i\) 上，动量相比梯度的缩放因子为 \(\phi_i = \frac{1}{\mu + \hat{\lambda}_i}\)。这恰好等价于 K-FAC 自然梯度更新中的输入协方差求逆因子 \(\mathbb{E}[x x^\top]^{-1}\)。因此，AK-Momentum 仅凭每步简单的 rank-1 外积修正，就在无需计算阻尼矩阵求逆的前提下，内生实现了输入端曲率白化，避免了显式二阶预条件算法庞大的显存开销与求逆不稳定性。

4. 最适配算子收缩调度与零持久显存落地：自适应选择收缩路径严格控制计算开销

在工程部署中，批量修正项形式为 \(\eta (G - M \hat{\Sigma})\)，其中 \(G = \frac{1}{N} \Delta^\top X\) 即为反向传播已计算出的标准权重梯度，完全免费；唯一的额外开销仅来自矩阵乘法 \(M \hat{\Sigma} = \frac{1}{N} M (\hat{X}^\top \hat{X})\)。针对层维度形状差异，算法提出了成本最优的算子分派法则：

\[M \hat{\Sigma} = \begin{cases} M (\hat{X}^\top \hat{X}), & \text{若 } n < 2m \text{（Gram 矩阵优先）} \\ (M \hat{X}^\top) \hat{X}, & \text{若 } n \ge 2m \text{（全连接中间张量优先）} \end{cases}\]

对于平方层与高瘦层（如 QKVO 及 MLP 上投影），采用 Gram 路径并让同一输入的多个共享投影复用已归一化的 Gram 矩阵；对于门控 MLP 的宽降维层（\(n = \gamma d \ge 2d\)），采用全连接路径避免物化庞大的 \(n \times n\) 矩阵。理论证明每个 Transformer 块相对于线性前向/反向的额外 FLOPs 占比严格收敛在窄带 \(\frac{6 + 4\gamma}{24 + 18\gamma} \in (22.2\%, 25.0\%]\) 内（Llama-2 架构下为 23.1%）。在前向中抓取的激活在反向完成后即刻释放，持久优化器显存开销为严格的 0 字节。

实验关键数据¶

主实验¶

论文在 FineWeb-Edu 10BT/20BT 数据集上预训练 Llama-2 架构语言模型（67M、370M 和 1B 参数规模），采用固定 μP 超参数迁移规则对比 AK-AdamW 与标准 AdamW，并在匹配验证损失水平下统计步数减少与算力收益。

规模 (种子数, Token预算)	平均验证损失差距 (nats)	最大验证损失差距 (nats)	平均步数节省比例 (%)	最大步数节省比例 (%)
67M (3 seeds, 10BT)	\(0.0936 \pm 0.0150\)	\(0.1077 \pm 0.0146\)	\(40.85 \pm 4.99\%\)	\(\mathbf{46.64 \pm 4.52\%}\)
370M (3 seeds, 10BT)	\(0.0383 \pm 0.0028\)	\(0.0501 \pm 0.0016\)	\(17.79 \pm 0.99\%\)	\(\mathbf{22.14 \pm 0.81\%}\)
1B (1 seed, 20BT)	\(0.0245\)	\(0.0361\)	\(13.47\%\)	\(\mathbf{19.37\%}\)

在相同的 Optuna 搜索调优协议下引入结构化优化器 Muon 进行对照，AK-AdamW 在 67M 上相比 Muon 取得了 \(0.120 \pm 0.008\) nats 的验证损失优势（平均步数节省 \(49.4 \pm 2.4\%\)），在 370M 上保持 \(0.109 \pm 0.003\) nats 的稳定优势。

算力与跨架构开销分析¶

下表汇报了 H200 GPU 上的实际算力开销、单步耗时以及在越过收敛交叉点后在匹配验证损失下的实际 FLOPs 与墙上钟表时间（Wall-clock）节省，并附带视觉模型 CIFAR-10 泛化实验。

架构 / 任务规模	单步 FLOPs 额外开销	实测单步耗时比率	匹配损失 FLOPs 净节省	匹配损失墙上时间净节省
Llama-2 67M (FineWeb-Edu)	10.0%	\(1.1526\times\) AdamW	35.1% (最高 41.1%)	32.0% (最高 38.3%)
Llama-2 370M (FineWeb-Edu)	15.5%	\(1.1775\times\) AdamW	5.6% (最高 9.4%)	4.0% (最高 7.7%)
Llama-2 1B (FineWeb-Edu)	18.1%	\(1.1844\times\) AdamW	正向收益（全程低于AdamW）	正向收益（全程低于AdamW）
ViT-Tiny (CIFAR-10, 120ep)	—	—	训练损失差距 0.079 nats	相对 AdamW 改善 25.1%
ResNet-18 (CIFAR-10, 120ep)	—	—	训练损失差距 0.041 nats	相对 AdamW 改善 11.1%

关键发现¶

各向异性白化得到实证确证：训练过程诊断显示，AK-AdamW 的动量缓冲区在整个训练期始终保持与真实总体梯度显著更高的余弦相似度，且对输出误差的预测误差明显更低；在 3,000 步之后，特征协方差矩阵的条件数显著降低、Top-1 方差占比大幅缩小、有效秩显著提高，充分证实了理论预测的隐式输入端曲率校正机制。
算力开销完全被步数收益覆盖：尽管单步因矩阵乘法产生 10%~18% 的理论计算开销，但由于达到相同验证损失所需步数削减达 20%~46%，在 67M 和 370M 规模下越过早期交叉点后，总训练 FLOPs 与端到端耗时均实现净收益。
超参数跨宽度完全零样本迁移：在 67M 规模调优的 \(\eta=0.37\) 和 \(\beta_1=0.99\) 在 370M 和 1B 上未经任何调整直接使用，两项规模下的收敛曲线从热身结束起即与 AdamW 彻底分离且永不相交。

亮点与洞察¶

将优化器动量重新抽象为关联记忆联想矩阵：突破了过去数十年视动量为平坦指数平均的惯性思维，通过输入激活与输出误差的键值外积分解，让 Delta 规则精确作用于特征空间。
无求逆实现自然梯度输入端预条件：无需维护复杂的高维二阶协方差逆矩阵，仅凭一阶动量自身的线性联想动态，在稳态下内生逼近了 Tikhonov 正则化的输入白化预条件矩阵。
极度轻量且即插即用：持久显存开销为零，代码实现只需修改动量累积几行算子，能够与二阶矩更新、权重衰减乃至外部矩阵正交化算法（如 Muon/Shampoo）直接正交叠加。

局限与展望¶

计算开销随宽度逐步逼近上限：随着模型隐藏层宽度增大，输入层维度扩展使额外计算开销从 67M 的 10.0% 提升至 1B 规模的 18.1%，在超大模型上对矩阵算子融合内核的工程优化要求更高。
缺乏输出端曲率的主动校正：AK-Momentum 仅依据前向输入激活实现了输入侧的曲率平衡，对于反向传播输出误差（output-side error）本身的协方差各向异性尚未做结构化处理，理论上需要与输出端预条件器协同。
验证范围集中在语言与分类：目前主要验证了语言预训练以及 CIFAR-10 上的视觉模型，多模态生成、强化学习及长序列场景下的动态特性仍待进一步探索。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 巧妙跳出 EMA 传统框架，首次将关联记忆 Delta 规则引入优化器一阶动量更新，打通了联想记忆与自然梯度的内在联系。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 67M 到 1B 语言模型完整预训练与 CIFAR-10 视觉基准，消融诊断深入到特征谱条件数与余弦对齐，具备 μP 严谨性。
写作质量: ⭐⭐⭐⭐⭐ 理论证明完整严密，从定点分析到非平稳跟踪理论环环相扣，实验成本核算坦诚详实。
价值: ⭐⭐⭐⭐⭐ 零持久显存占用、直接替代一阶动量的高效改进方案，具备成为下一代大规模模型训练标准基础组件的潜力。