激活键控动量：通过 Delta 规则进行各向异性动量更新

大多数现代优化器将其动量形成为过去梯度的指数移动平均值（EMA），以一个固定的速率忘记每个方向。然而，深度网络在训练过程中看到的输入可能是高度各向异性的，一些方向被频繁查询，而大多数方向很少被查询。预处理方法通过围绕该缓冲区进行额外处理来解决这种各向异性，并使动量更新本身保持不变。我们提出了激活键动量（AK-Momentum），它将方向意识构建到动量更新规则中。线性层的梯度分为充当键的输入激活和充当值的输出侧误差。根据该激活，AK-Momentum 通过规范 delta 规则更新动量缓冲区，因此每个方向都会按照其出现频率设定的速率被遗忘。我们证明它是一个有效的动量，它应用输入侧曲率校正而无需矩阵求逆，并且在固定和漂移最优条件下，它比 EMA 更快地清除陈旧方向。它是任何优化器动量缓冲区的直接替代品，其系数跨宽度传输低于 $μ$P，并且其额外计算量保持在门控 MLP 块线性成本的 $22.2\%$ 和 $25.0\%$ 之间，无需持久内存。在 FineWeb-Edu 预训练中，使用 AK-Momentum (AK-AdamW) 的 AdamW 在 67M 步数下达到 AdamW 的验证损失高达 $46.39 \pm 4.32\%$，在 370M 步数下达到 $22.12 \pm 0.80\%$，超过三个种子，并且在 Chinchilla 最佳预算下增益持续保持在 1B。在同一协议下调整的 Muon 基线在两种语言模型尺度上均高于 AK-AdamW，并且 CIFAR-10 上的 SGD、ResNet-18 和 ViT-Tiny 的增益均保持不变。训练时诊断确认了预测机制、更好的梯度跟踪和更健康的输入方向。