论文

用于 μ 子优化的物理响应和记忆模型

A Physical Response-and-Memory Model for Muon Optimization

模型训练预训练

摘要

训练 大语言模型 的成本很高。相同的计算最终能达到多低的损失取决于每个步骤的梯度如何转换为权重更新;执行此转换的规则是优化器。从 SGD 和 AdamW 到最近的 Muon,有效的更新规则大多由工程直觉形成,然后根据基准进行选择。 Muon 在应用更新之前对动量矩阵进行半正交,并不断打破公共训练基准的记录;然而,为什么半正交方向有效,以及动量在多长时间的历史中应该平均,是目前主要通过经验回答的两个问题。在这里,我们将训练过程中的权重矩阵视为具有记忆的响应介质,并为其构建物理模型,其中两个问题都找到了答案:半正交方向是输出侧安全预算下的最大耗散响应,这解释了它的工作原理;动量是介质积累的内应力;它应该平均多长时间是由该应力的松弛决定的,并且真正的介质在多个时间尺度上松弛,最简单的形式是一个快一个慢。在此基础上我们提出了Bi-Maxwell优化器。该框架进一步产生了一个可测试的结果:梯度方向在训练早期变化很快,后期变化更慢,因此最佳记忆长度应该随着训练阶段而增长;通过只读探针在 8 个独立训练轨迹上对其代理进行的逐步测量与此结果一致。单独替换内存内核(从单个时间尺度到两个时间尺度)可以在公共大型语言模型优化器基准测试中以明显更少的步骤进行目标损失训练。