📊 statistics

Training for the Model You Return: Improving Optimization for Iterate-Averaged Language Models

本文提出了 PACE，一种针对 AdamW 的优化器封装器，它将训练视为一个旨在最小化迭代平均语言模型误差的最优控制问题，并展示了其在监督微调和预训练任务中的理论收敛保证及经验性改进。

原作者： Kwok Chun Au, Adam Block

发布于 2026-06-25

📖 1 分钟阅读☕ 轻松阅读

原作者： Kwok Chun Au, Adam Block

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下你正在训练一个巨大的数字大脑（语言模型）去学习一项新技能，比如写诗或编程。你通过向它展示数千个示例，并让它逐步调整其内部的“权重”（它的知识）。

通常情况下，当你完成训练后，你会取出最后一个版本的脑力并说：“这就是最终成品。”但许多现代研究人员发现了一个秘密：大脑在训练过程中经历的所有版本的平均值，往往比仅靠最后一个版本更聪明、更稳定。这就像是在说：“我所有练习测试的平均分，比我在最后一天参加的那次考试成绩更能预示我的最终考试水平。”

问题在于，训练过程并不知道它最终会被这个“平均值”所评判。它只想尽可能快地冲向终点，这可能会让路径变得摇摆不定，从而使最终的平均值不够完美。

这篇论文介绍了一种名为 PACE（用于高效优化的回拉平均控制，Pullback Averaging Control for Efficient Optimization）的新方法来解决这个问题。以下是它的工作原理，使用简单的类比进行说明：

1. 问题所在：摇晃的钢丝绳

把训练语言模型想象成在钢丝绳上行走，走向一个特定的目的地（完美的答案）。

标准训练 (AdamW)： 你向前迈步。有时走得太远，有时又左右摇晃。当你到达终点时，你可能稍微偏离了中心。
“平均值”技巧： 我们不只是看你最终停在哪里，而是决定观察你在行走过程中所站过的每一个位置的平均值。
问题在于： 如果你知道自己将根据你的平均位置被评判，你应该以不同的方式行走。你不应该只是冲向终点；你应该尝试在整个过程中保持路径平稳且居中。但标准的训练算法并不知道这一点；它们只是在盲目冲刺。

2. 解决方案：“幽灵向导”（最优控制）

作者提出了一个根本性的问题：如果我们知道自己将被基于“平均值”进行评判，我们应该如何改变行走方式，才能让这个平均值达到最优？

他们使用了高级数学（具体来说是用于引导火箭和机器人的“最优控制理论”）来解决这个问题。他们想象了一个简化的世界，目标是一个完美的谷底，而训练中的“噪声”就像是把行路人吹离航线的风。

他们计算出了完美的策略：不要只是向前走；偶尔回头看看你的“幽灵向导”（你过去所处位置的平均值），并轻轻地将自己拉回到它身边。

3. 现实世界的工具：PACE

完美的数学解法对于在真实的、大规模计算机上运行来说过于复杂。因此，作者创建了一个实用的、轻量级的版本，称为 PACE。

它是如何工作的： PACE 是一个位于标准训练工具 (AdamW) 之上的“包装器”。
机制： 每隔几步，PACE 就会检查模型当前的位置与它过去步骤的“平均值”之间的差异。
回拉： 如果模型偏离平均值太远，PACE 会轻轻地将其拉回。这就像一名教练牵着一只不断跑离路径的小狗的绳子；教练并没有停止小狗，而是轻轻地将它拉回队伍的中心。
智能推力： 这种拉力的强度并非随机的。它是经过“裁剪”（限制）的，因此不会猛烈地拽动模型，并且会根据模型在当前步骤中的信心程度进行调整。

4. 结果：更平稳的旅程

论文在几种不同的语言模型（参数量从 10 亿到 20 亿不等）上测试了 PACE。

类比： 想象两名跑步者。一名跑步时狂野奔跑，在通往终点的路上左右横跳（标准训练）。另一名跑步者则平稳奔跑，不断修正路径以保持居中（PACE）。
结果： “平稳”的跑步者（PACE）最终获得的“平均位置”始终优于“狂野”的跑步者。
关键发现： 即使在允许标准方法在结束时降低学习率（一种常用的稳定技巧）的情况下，PACE 的表现也优于标准方法。PACE 在整个训练过程中都实现了这种稳定性，而不仅仅是在结束阶段。

总结

简而言之，这篇论文认为，如果你计划使用训练步骤的平均值作为你的最终模型，那么你应该以不同的方式进行训练。你不应该仅仅瞄准终点；你应该努力让你的整个旅程都保持居中。PACE 是一个简单的、全新的工具，它能轻轻地将模型拉回到它自己的历史轨迹中，从而产生一个更聪明、更稳定的最终产品。

该论文并不声称：

它并不声称这适用于医疗诊断或临床用途。
它并不声称这适用于参数量超过 20 亿的模型（他们只测试到了 20 亿规模）。
它并不声称这能取代所有其他方法，而是通过结合平均化处理，对标准的“AdamW”方法进行了改进。

技术摘要：针对返回模型的训练优化

问题陈述

现代语言模型（LM）流水线通常返回的是一个平均模型——通常是训练迭代过程的指数移动平均（EMA）——而非最终的迭代点本身。虽然迭代平均已知可以稳定优化并提高泛化能力，但一个根本性的问题仍然存在：鉴于最终输出将是一个迭代平均值，那么如何修改训练过程本身，以优化这一特定估计器的性能？

先前的工作通过观察到平均化允许使用更激进的学习率，在隐含层面解决了这个问题。然而，本文采取了一种直接的方法：它试图设计一种优化算法，通过显式控制训练轨迹来最小化返回的迭代平均估计器的误差。

方法论

理论公式化

作者将该问题构建为连续时间随机二次模型中的一个最优控制问题。

动力学： 参数轨迹 $\theta_t$ 遵循一个随机微分方程（SDE），该方程代表了在二次极小值附近的带噪梯度下降（Ornstein–Uhlenbeck 过程），并由一个控制输入 $u_t$ 进行增强。
目标： 寻找一个控制策略 $u_t$ ，在对控制干预幅度施加二次惩罚（正则化参数 $\lambda$ ）的条件下，最小化返回的时刻平均估计器 $\hat{\theta}_{EMA, T}$ 与真实最优值 $\mu^*$ 之间的均方误差。
推导： 利用线性二次高斯（LQG）控制理论和确定性等价原理，作者推导出了一个最优控制器。精确解取决于给定轨迹下最优值的后验均值，这在大型语言模型中是难以处理的。

实际近似：PACE

为了使解决方案具有可实现性，作者引入了一系列启发式近似：

估计器替换： 用当前的迭代平均值（ $\hat{\theta}_{EMA, t}$ ）代替未知的最优值后验均值。
时间近似： 假设当前时间 $t$ 远小于总训练时间 $T$ （即 $t \ll T$ ），从而简化随时间变化的系数。
离散化与预处理： 对连续控制进行离散化，并将未知的 Hessian 矩阵替换为 AdamW 中使用的对角预处理器（即二阶矩估计 $v_t$ ）。

由此产生了 PACE（用于高效优化的拉回平均控制，Pullback Averaging Control for Efficient Optimization），这是一个围绕 AdamW 构建的轻量级封装。其更新规则增加了一个“拉回”项，将实时权重向 EMA 权重方向移动：
$\theta_{k+1} = \theta_k - \eta g_k + C(\hat{\theta}_{EMA, k} - \theta_k)$
其中 $C$ 是由 Adam 预处理器导出的对角矩阵，且拉回强度被限制以确保更新仍是标准步长与拉回项的凸组合。

核心贡献

最优控制公式化： 本文提供了第一个专门用于最小化迭代平均估计器误差而非最终迭代点误差的优化算法的有理推导。
理论保证：
- 凸函数设定： 作者证明了 PACE 的一个程式化版本在具有无偏随机梯度的任何凸损失函数下，都能以标准的 $O(1/\sqrt{T})$ 速率收敛，至多受限于一个依赖于 EMA 強度的常数因子。这确保了该方法在任何情况下都不会显著差于标准 SGD。
- 二次函数设定： 在用于推导的二次设定中，受控动力学严格改善了迭代平均估计器的极限平方误差，优于不受控的动力学（标准带有 EMA 的 SGD）。在某些二次实例上，这种提升可以是任意大的。
实证验证： 该方法在以下场景进行了评估：
- 微调： 在 smol-smoltalk 数据集上的 SmolLM2-1.7B、Qwen3-1.7B 和 Gemma3-1B 模型。
- 预训练： 在 Chinchilla 最优 Token 预算下的 GPT-2 (124M) 以及 FineWeb 数据集。
- 基准测试： 将 PACE 与原生 AdamW、带有 EMA 评估的 AdamW 以及 Schedule-Free 算法进行了对比。

结果

性能： 在所有测试的模型（1B–2B 参数）和数据集上，PACE 均严格优于原生 AdamW 和带有 EMA 评估的 AdamW。
鲁棒性： 这种提升在广泛的学习率、衰减调度（恒定、余弦、Warmup-Stable-Decay）和超参数（拉回强度 $c$ 、EMA 指数 $\kappa$ ）范围内都是鲁棒的。
与 Schedule-Free 的比较： PACE 与 Schedule-Free 算法（Defazio 等人，2024）表现相当，后者同样利用迭代平均来消除对学习率衰减的需求。然而，PACE 使用了不同的更新规则。
学习率衰减： PACE 在恒定学习率下也表现良好，表明它可以在通常需要学习率衰减的阶段稳定训练。在 Token 预算对比中，PACE 的恒定学习率轨迹能像 WSD 分支一样有效地追踪损失景观的“河床”，甚至在某些情况下表现更优。

意义与局限性

意义：
本文证明了针对“返回的模型”（即平均值）而非“最终迭代点”进行优化，可以带来切实的性能增益。通过将训练轨迹视为一个控制问题，PACE 提供了一种有理论依据的方法，用于在无需复杂调度调优的情况下稳定训练并提高泛化能力。作者指出，尽管由于 EMA 因子的存在，在一般凸函数设定下的理论收敛保证略慢，但在二次设定中，PACE 提供了比标准方法更好的理论收敛性。

局限性：

理论假设： 推导依赖于一个带有加性噪声的简化二次模型。虽然在一般凸函数下证明了收敛性，但如何将动量和自适应预处理（AdamW 的核心）严谨地纳入理论分析仍是一个开放的研究方向。
规模： 实验范围限于微调阶段最高 2B 参数的模型以及预训练阶段 124M 参数的模型。其在大规模（如 70B+）下的有效性尚未得到验证。
内存开销： PACE 需要在训练期间维护一份额外的模型权重副本（EMA 状态），这引入了不可忽视的内存开销，尽管在实践中可以通过大 Batch Size 来缓解。
种子敏感性： 由于计算资源限制，主要的运行实验使用了单一随机种子，尽管多种子鲁棒性研究确认了该方法的排序稳定性。

作者总结道，PACE 代表了迈向使优化算法与现代 LM 实际部署策略（即平均模型作为最终产品）相一致的重要一步。