用优化器思想改进Transformer，提升训练效果和泛化能力。
Momentum Streams for Optimizer-Inspired Transformers
- 将Transformer层残差更新类比为优化器步，引入动量机制。
- 三重动量Transformer在相同算力下验证损失最低，优于原版Transformer。
- 动量带来更平坦的极小值，减少遗忘，适合长序列任务。
预归一化Transformer层的残差更新可解释为作用于代理令牌能量的一阶优化器一步，其中注意力与MLP子层充当梯度代理。基于此观察，我们构建了一系列优化器启发的Transformer（三重动量、Adam/AdamW、Muon、SOAP），并在匹配计算量下进行对比。主预训练实验中，三重动量TMMFormer达到最低验证损失，优于原始Transformer及先前架构变体。受控消融与支持性理论表明，性能提升主要源于动量而非预条件化。进一步显示，TMMFormer及其他动量型设计能收敛至比原始Transformer更平坦的极小值，从而减少遗忘并提升泛化能力。
原文摘要 · Abstract (English)
The residual update of a pre-norm Transformer layer admits an interpretation as one step of a first-order optimizer acting on a surrogate token energy, wherein the attention and MLP sublayers function as gradient oracles. Based on this observation, we build a family of optimizer-inspired Transformers (triple-momentum, Adam/AdamW, Muon, SOAP) and compare them under matched compute. In our main pretraining experiment, the triple-momentum TMMFormer achieves the lowest validation loss, outperforming the vanilla Transformer and prior architectural variants. A controlled ablation and supporting theory show that momentum, not preconditioning, is the main source of the gain. We further show that TMMFormer and other momentum-based designs reach flatter minima than the vanilla Transformer, which leads to less forgetting and better generalization.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。