面向优化器启发式 Transformer 的动量流
机器学习
2026-05-26 v1 人工智能
计算与语言
摘要
预规范 Transformer 层的残差更新可被解释为对 surrogate token 能量执行的一阶优化器的一步，其中注意力和 MLP 子层充当梯度 oracle。基于这一观察，我们构建了一系列优化器启发式 Transformer（三动量、Adam/AdamW、Muon、SOAP），并在匹配的计算资源下进行比较。在我们的主要预训练实验中，triple-momentum TMMFormer 取得了最低的验证损失，超越了 vanilla Transformer 和 prior 架构变体。受控的消融实验和支持理论表明，动量而非 preconditioning 是主要收益来源。我们进一步表明，TMMFormer 和其他动量-based 设计比 vanilla Transformer 达到更平坦的最小值，这导致更少的遗忘和更好的泛化。
关键词
引用
@article{arxiv.2605.24425,
title = {Momentum Streams for Optimizer-Inspired Transformers},
author = {Jingchu Gai and Nai-Chieh Huang and Jiayun Wu},
journal= {arXiv preprint arXiv:2605.24425},
year = {2026}
}