arXiv:2608.19491cs.LGcs.CL2026-08

让优化器学会按方向遗忘,提升训练效率

DeltaMomentum: A Key-Value based Anisotropic Momentum Update via Delta Rule

论文配图:DeltaMomentum: A Key-Value based Anisotropic Momentum Update via Delta Rule
图 1 · 摘自论文原文
  • 用键值结构动态调整各方向的遗忘速率
  • 在多个模型上减少46%训练步数达同样效果
  • 可直接替换现有优化器,开销仅25%以内

现代优化器通常用指数移动平均(EMA)计算动量,以固定速率遗忘所有方向。但深度网络训练中输入方向高度异质:少数方向频繁出现,多数极少被触及。现有方法在缓冲区外增加处理模块,未改变动量更新本身。本文提出DeltaMomentum,将方向感知嵌入动量更新规则。核心观察是线性层梯度可分解为输入(键)与输出误差(值)。利用键值结构,采用经典δ规则更新动量,使每个方向的遗忘率由其出现频率决定。理论证明该方法为有效动量,无需矩阵求逆即可实现输入侧曲率校正,并在固定与漂移最优解下清除陈旧方向更快。它是任意优化器动量缓冲区的即插即用替代品,系数在μP条件下跨宽度可迁移,额外计算量占门控MLP线性部分的22.2%~25.0%,无持久内存开销。在FineWeb-Edu预训练中,使用DeltaMomentum的AdamW(DeltaAdamW)在67M规模上比原AdamW少46.39±4.32%步数达到相同验证损失,在370M上减少22.12±0.80%,1B规模下增益仍持续。基于相同协议调优的Muon基线在两个语言模型尺度上均低于DeltaAdamW,且对SGD、ResNet-18和ViT-Tiny在CIFAR-10上的性能提升亦成立。训练诊断证实了预期机制,展现更优梯度追踪与更健康的输入方向。

原文摘要 · Abstract (English)

Most modern optimizers form their momentum as an exponential moving average (EMA) of past gradients, forgetting every direction at one fixed rate. However, the inputs a deep network sees during training can be highly anisotropic, with a few directions queried frequently while most are seen rarely. Recent methods address this anisotropy by wrapping extra processing around this buffer, leaving the momentum update itself unchanged. We propose DeltaMomentum, which builds direction-awareness into the momentum update rule. The main observation is that the gradient of a linear layer splits into an input that acts as a key and an output-side error that acts as a value. Exploiting the key-value structure, DeltaMomentum updates the momentum buffer by the canonical delta rule, so each direction is forgotten at a rate set by how often it appears. We prove that it is a valid momentum, that it applies the input-side curvature correction without matrix inversion, and that it clears stale directions faster than EMA under both a fixed and a drifting optimum. It is a drop-in replacement for the momentum buffer of any optimizer, its coefficient transfers across widths under $μ$P, and its extra compute stays between $22.2\%$ and $25.0\%$ of a gated-MLP block's linear cost with no persistent memory. In FineWeb-Edu pretraining, AdamW with DeltaMomentum (DeltaAdamW) reaches AdamW's validation loss in up to $46.39 \pm 4.32\%$ fewer steps at 67M and $22.12 \pm 0.80\%$ at 370M over three seeds, and the gain persists at 1B on a Chinchilla-optimal budget. A Muon baseline tuned under the same protocol sits above DeltaAdamW at both language-model scales, and the gain holds for SGD, ResNet-18, and ViT-Tiny on CIFAR-10. Training-time diagnostics confirm the predicted mechanism, better gradient tracking and healthier input directions.

优化器动量更新键值机制训练加速

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