为什么大模型优化器需要学会健忘：K-Maxwell动量退火如何连破NanoGPT世界纪录

训练一个大语言模型，为什么有时候记性太好，反而成了累赘？

在几乎所有深度学习的代码里，优化器都被赋予了恒定不变的记忆力。 第一步怎么留存过去的动量，到了第三千步，依然按一模一样的节奏去留存。 但在开源社区的 NanoGPT 极速竞赛里，有人发现了一个完全反直觉的现象。 只要让优化器在训练后期慢慢学会健忘，把记忆的平均跨度从 58 步缩短到 26 步，训练所需的步数就能硬生生砍掉近百步。 不改网络架构，不加任何训练数据，不调整批次大小。 在这个由八张英伟达 H100 严格卡死评测标准的赛场上，两位研究者用这个改动，一天之内连续打破了两个世界纪录。

这项比赛叫 NanoGPT 优化基准。 它的规则非常严苛。 参赛者必须把一个类 GPT-2 模型，在 FineWeb 数据集上训练到 3.28 的验证集损失。 网络结构不准动，数据量不准动，唯一的比拼维度就是算法效率：看谁能用最少的迭代步数撞线。 以往大家打破纪录，往往要在矩阵正交化、学习率调度或者精细的显存调度上下功夫。 但研究者 Jack Zheng 和 Jeffrey Cheng 盯上的，是深度学习里最基础的一个部件：动量。

训练神经网络，就像在浓雾笼罩的崇山峻岭里摸索下山的路。 计算机在每一步计算出的梯度，受到小批量数据的随机性影响，往往像醉汉一样到处乱晃。 如果只看当前这一步的坡度，模型很容易在狭窄的山谷两壁来回震荡，根本走不快。 所以几十年前，研究者就给优化器装上了动量机制。 做法很简单，就是维护一个历史梯度的指数移动平均。 过去走过的方向会形成一股惯性。 有了这股惯性，高频的噪声被抹平，模型就能顺着平缓的主山脉大步流星往前冲。

但这里埋着一个所有人都默认接受的假设。 从经典的带动量随机梯度下降，到目前大火的 Muon 优化器，动量的衰减系数通常都是一个固定的常数。 这意味着优化器在整个训练周期里，记忆的时间半衰期是死板固定的。 在训练刚开始的时候，这套机制运作得很好。 损失函数的曲面开阔，地势平缓，巨大的惯性能带着参数毫不犹豫地冲过各种坑洼。 可是到了训练后半程，情况全变了。

当模型越来越接近最优解的谷底，学习率开始退火，周围的损失曲面变得极其陡峭而狭窄。 这时候模型需要的不再是闭着眼睛往前冲。 它需要极其精准的微调，落进那个微小的盆地中央。 如果优化器手里还死死攥着 50 步之前的历史动量，那些过期的记忆就变成了有毒的负重。 想要转弯的时候，过去的惯性在生拉硬拽；想要刹车停进车位的时候，旧动量推着它直接冲出了最优区域。 很多训练后期的震荡和徘徊，根源都在于惯性太大，车头转不过来。

怎么解决这个问题？ 物理学里早就给出过答案。 在材料科学中，描述粘弹性流体松弛过程的麦克斯韦模型，如果只用单一弹簧和阻尼器，就无法刻画复杂的流动。 物理学家会并联多个不同时间尺度的单元，组成广义麦克斯韦模型。 Jack Zheng 和 Jeffrey Cheng 把这套物理直觉搬进了优化器，命名为 K-Maxwell 动量。 他们跳出单一指数移动平均缓冲池的框架，一口气并联了 6 到 8 个对数等距分布的缓冲池。 最快的缓冲池只保留最近 3 步的信息，专抓瞬时变化；最慢的缓冲池能追溯 64 步，保持长线大趋势。

更绝的一步在于动态退火。 在以往的类似尝试中，多个缓冲池的权重也是固定的。 但这次他们让缓冲池的混合权重，随着训练步数实时发生线性偏移。 在基准 Muon 优化器上，从第 1000 步开始，记忆池的加权平均寿命从 58 步逐渐滑落到 26 步。 在针对更大矩阵做了超球面投影的 MuonH 优化器上，从第 750 步开始，平均记忆寿命从 50 步压缩到 22 步。 这相当于给高速行驶的列车安装了一套自适应转向系统。 前半程，它是一辆重载列车，靠着超长记忆的巨大惯性在荒原上破风前行； 后半程，它在飞奔中卸下重甲，蜕变成一台转向极其灵敏的轻型赛车，看准谷底的针尖位置轻巧一踩，精准入弯。

结果干脆利落。 在 Muon 优化器上，达成 3.28 目标所需的步数从 3250 步降到 3160 步，净省 90 步，效率提升近 3%。 在最新的 MuonH 优化器上，步数从 3125 步进一步压到了 3065 步，再省 60 步。 在连续八个不同随机种子的严苛测试下，每一颗种子的表现都稳定超越了原先的世界纪录。

当整个行业都在把注意力押在采购更多显卡、堆叠更多数据参数的时候，这项纪录展示了算法底层的另一种美感。 我们每天用来训练海量参数的数学工具，底层依然留着大量粗糙的经验主义空白。 知道何时蓄力前行固然重要，但懂得在接近终点时果断抛掉过去的惯性，才是真正的高手过弯。