⏶18
Muon 在尾部关联记忆学习中表现优于 Adam
发表
由 
Fengzhuo Zhang 提交
Fengzhuo Zhang 提交作者: Shuche Wang, Fengzhuo Zhang, Jiaxiang Li, Cunxiao Du, Chao Du, Tianyu Pang, Zhuoran Yang, Mingyi Hong, Vincent Y. F. Tan
摘要
AI 生成总结
Muon 优化器通过有效优化关联记忆参数并在重尾数据中平衡跨类学习，在训练 LLM 方面优于 Adam。Muon 优化器在训练大型语言模型 (LLM) 时始终比 Adam 更快，但其成功的机制仍不清楚。本文通过联想记忆的视角揭示了这一机制。通过对 Muon 优化的 Transformer 组件进行消融实验，我们发现 LLM 的联想记忆参数，即 Value 和 Output (VO) 注意力权重以及前馈网络 (FFN)，是 Muon 优越性的主要贡献者。基于这种联想记忆的视角，我们进而解释了 Muon 在现实世界语料库上的优越性，这些语料库本质上是重尾的：少数类别（尾部类别）出现的频率远低于其他类别。通过两个关键特性解释了其优越性：(i) 其更新规则始终比 Adam 产生更各向同性的奇异值谱；因此，(ii) 在重尾数据上，它比 Adam 更有效地优化尾部类别。除了经验证据，我们还在类别不平衡数据下通过分析单层联想记忆模型来理论证实了这些发现。我们证明了 Muon 无论特征嵌入如何，都能始终跨类别实现均衡学习，而 Adam 则可能根据嵌入特性引起学习错误的巨大差异。总而言之，我们的经验观察和理论分析揭示了 Muon 的核心优势：其更新规则与线性联想记忆的外积结构一致，从而在重尾分布中比 Adam 更均衡、更有效地学习尾部类别。
Muon 优化器在训练大型语言模型（LLM）时比 Adam 更快，但其成功的机制仍然不清楚。本文通过联想记忆的视角揭示了这种机制。通过去除 Muon 优化的 Transformer 组件，我们发现 LLM 的联想记忆参数，即 Value 和 Output (VO) 注意力权重以及前馈网络 (FFN)，是 Muon 优越性的主要贡献者。受这种联想记忆观点的启发，我们随后解释了 Muon 在实际语料库上的优越性，这些语料库本质上是重尾的：少数类别（尾部类别）比其他类别出现的频率低得多。其优越性通过两个关键属性来解释：(i) 它的更新规则比 Adam 持续产生更具各向同性的奇异谱；结果是，(ii) 在重尾数据上，它比 Adam 更有效地优化尾部类别。除了经验证据，我们还通过分析类别不平衡数据下的单层联想记忆模型来理论上证实这些发现。我们证明了 Muon 无论特征嵌入如何，都能在类别之间实现均衡的学习，而 Adam 则可能根据嵌入属性引起学习误差的巨大差异。总而言之，我们的经验观察和理论分析揭示了 Muon 的核心优势：其更新规则与线性联想记忆的外积结构一致，从而在重尾分布中实现比 Adam 更均衡、更有效的尾部类别学习。