论文

受优化器启发的 Transformer 的动力流

Momentum Streams for Optimizer-Inspired Transformers

模型架构Transformer

摘要

预规范 Transformer 层的残差更新允许将其解释为作用于代理词元能量的一阶优化器的一个步骤,其中注意力和 MLP 子层充当梯度预言机。基于这一观察,我们构建了一系列受优化器启发的 Transformer(三重动量、Adam/AdamW、Muon、SOAP),并在匹配计算下对它们进行比较。在我们的主要预训练实验中,三动量 TMMFormer 实现了最低的验证损失,优于普通 Transformer 和之前的架构变体。受控消融和支持理论表明,动量而不是预处理是增益的主要来源。我们进一步表明,TMMFormer 和其他基于动量的设计比普通 Transformer 达到更平坦的最小值,这会导致更少的遗忘和更好的泛化。