📊 statistics

Correcting Stochastic Update Bias in Preconditioned Language Model Optimizers

本文通过提出一种结合交叉拟合预条件与方差校正逆运算的单批次框架，识别并修正了预条件语言模型优化器中的两种有限样本偏差——梯度预条件器耦合与非线性逆偏差，从而降低了预训练损失，并提升了 AdamW、Sophia 和 Shampoo 等模型的性能。

原作者： Nikhil Nayak, Julia White, Urchade Zaratiana, Kelton Zhang, Henrijs Princis, Dhruv Atreja, Henry Fawcett, Matthew Thomas, George Hurn-Maloney, Ash Lewis

发布于 2026-05-21

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Nikhil Nayak, Julia White, Urchade Zaratiana, Kelton Zhang, Henrijs Princis, Dhruv Atreja, Henry Fawcett, Matthew Thomas, George Hurn-Maloney, Ash Lewis

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

想象一下，你正在尝试教一个机器人写诗。为此，你向它展示成千上万首诗的示例，并要求它调整其“大脑”（即其参数）以变得更好。机器人使用一种称为优化器（optimizer）的特殊工具，来确定该如何微调其“大脑”。

大多数现代优化器都是“聪明”的。它们不仅查看机器人犯下的错误，还会观察错误的形状，以此决定迈出多大的步伐。这被称为预条件（preconditioning）。

然而，本文的作者发现，这些聪明的优化器存在一个隐藏的缺陷：由于它们试图基于小样本且有噪声的数据做出决策，因此存在轻微的偏差。这就像只透过窗户看五秒钟就试图判断天气一样。

以下是问题所在及其解决方案，通过简单的类比进行解释。

两个隐藏缺陷

该论文指出了优化器在处理小批量数据时，在数学上出错的两种具体方式：

1. “自利”偏差（耦合偏差）

类比：想象一个学生正在参加考试。为了给自己的作业打分，他们使用刚刚考过的那同一张试卷来计算分数。自然地，他们可能会因为问题和答案在脑海中完美匹配，而倾向于给自己打分过严或过宽。
现实：在优化器中，机器人使用完全相同的小组数据来计算“移动方向”（梯度）和“步长大小”（预条件子）。由于它们来自同一来源，它们在统计上是“耦合”的。这会在更新规则中产生微妙的扭曲，导致机器人的步伐略微偏离目标。

2. “非线性”偏差（逆偏差）

类比：想象你试图猜测汽车的平均速度。你知道驾驶一英里的平均时间是 1 分钟。你可能会想：“好吧，所以平均速度是 60 英里/小时。”但数学并非如此简单！如果你先对时间取平均值，然后再计算速度，得到的答案与先计算每次行程的速度然后再取平均值是不同的。这是因为“速度”是“时间”的倒数，而对倒数进行数学运算是棘手的。
现实：优化器需要除以一个数字（预条件子）来确定步长。即使机器人平均而言能完美地估计预条件子，取倒数（将其翻转以进行除法）这一行为也会引入系统性误差。这就像试图通过先平均一群人的鞋码来猜测他们的平均身高；数学关系因此发生了扭曲。

解决方案：“双重检查”系统

作者提出了一种巧妙的修复方法，无需机器人重新读取数据或显著降低速度。他们称之为单批次偏差校正（Single-Batch Bias Correction）。

他们同时应用了两个主要技巧：

1. 交叉拟合（“独立裁判”）

工作原理：机器人不再使用同一组数据来计算“方向”和“步长”，而是将当前的数据批次分成两个独立的组。
- A 组计算方向。
- B 组计算步长。
结果：通过保持裁判的独立性，你消除了“自利”偏差。步长不再受该特定组数据具体方向的影响。

2. 方差校正（“噪声计”）

工作原理：机器人观察“步长”估计值在数据微小子组之间的变化程度。如果估计值不稳定（高方差），机器人就知道由于上述非线性逆问题，数学计算可能存在偏差。
结果：机器人使用一个统计公式（基于“Delta 方法”）来减去预期的误差。这就像一位机械师知道，在潮湿天气下，某种特定工具往往会高出 2% 的读数，因此他们会自动从读数中减去 2% 以获得真实值。

尝试后的结果如何？

该团队在三种用于训练大型语言模型（如驱动聊天机器人的模型）的“智能”优化器上测试了这一修复方案：

AdamW：标准的、主力型优化器。
Sophia：一种试图猜测问题“曲率”（例如判断山坡是陡峭还是平坦）的优化器。
Shampoo：一种非常复杂的优化器，利用矩阵数学来处理海量数据。

结果：

预训练（从头学习）：当从零开始训练模型时，偏差校正效果非常好。它显著降低了所有三种优化器的误差率（损失）。这就像机器人因为不再被自身有噪声的计算所误导，从而更快、更准确地学会了写诗。
指令微调（对已有智能模型进行微调）：当他们尝试教一个已经训练好的模型学习新任务时，结果是“中性至积极”的。校正没有造成损害，在某些情况下甚至提供了微小的提升，但收益较小。这是有道理的，因为模型已经稳定，因此“噪声”不再是主要问题。

核心结论

该论文认为，我们一直将这些优化器视为完美的数学机器，但实际上，由于它们处理的数据有限，它们正在犯下微小但一致的错误。

通过简单地拆分数据以保持方向和步长的独立性，并测量噪声以减去逆运算误差，作者创建了一个“偏差校正层”。这一层使训练过程更加高效，允许语言模型学得更好、更快，尤其是在它们从零开始的时候。这是一个微小的统计调整，却对这些 AI 模型的学习效果产生了惊人的巨大影响。

技术摘要：修正预条件语言模型优化器中的随机更新偏差

1. 问题陈述

预条件优化器（例如 AdamW、Sophia、Shampoo）是训练大型语言模型（LLM）的核心。这些方法从迷你批次（minibatches）中估计下降方向和几何项（预条件器），以逼近理想的总体更新：
$u^\star(\theta) = P(\theta)^{-1}\nabla f(\theta)$
在实践中，优化器使用迷你批次估计值 $\hat{g}$ 和 $\hat{P}$ 应用随机更新：
$\hat{u} = \hat{P}^{-1}\hat{g}$

本文指出，即使 $\hat{g}$ 和 $\hat{P}$ 是总体梯度和预条件器的无偏估计量，乘积 $\hat{P}^{-1}\hat{g}$ 通常也是理想更新的有偏估计量。这种有限样本偏差源于两个不同的来源：

梯度–预条件器耦合偏差：当 $\hat{g}$ 和 $\hat{P}$ 从同一个迷你批次估计时，它们在统计上是相关的。这种耦合引入了一个交叉矩项，使其偏离目标总体向量场。
逆预条件器偏差：矩阵求逆（或逆根运算）是一种非线性变换。根据詹森不等式（Jensen's inequality），即使 $\hat{P}$ 对 $P$ 是无偏的，逆的期望 $\mathbb{E}[\hat{P}^{-1}]$ 也不等于期望的逆 $(\mathbb{E}[\hat{P}])^{-1}$ 。当预条件器估计值噪声大、激进或为矩阵值时，这种偏差尤为显著。

这些偏差并非实现错误，而是将非线性自适应规则应用于有限随机估计时的固有后果。

2. 方法论

作者提出了一种单批次偏差校正框架，无需重新访问相同的参数值（这在神经网络训练中是不切实际的），即可解决上述两种偏差来源。该方法在固定的参数步长 $\theta_t$ 处运行，通过将当前训练批次划分为微批次组（microbatch groups）来实现。

2.1 交叉拟合预条件（解决耦合问题）

为了消除梯度–预条件器耦合，该方法从独立的微批次组中估计分子（梯度）和分母（预条件器）。

设 $A$ 和 $B$ 为从同一训练分布中采样的两个独立组。
梯度从 $A$ 估计： $\hat{g}_A$ 。
预条件器从 $B$ 估计： $\hat{P}_B$ 。
更新变为 $\hat{u}_{cf} = \hat{P}_B^{-1}\hat{g}_A$ 。
由于独立性， $\mathbb{E}[\hat{P}_B^{-1}\hat{g}_A] = \mathbb{E}[\hat{P}_B^{-1}]\mathbb{E}[\hat{g}_A]$ ，从而有效地消除了耦合偏差项。

2.2 方差校正求逆（解决非线性问题）

为了减少由求逆非线性引入的偏差，该方法利用了预条件器估计组 $B$ 内部的变异性。

将组 $B$ 划分为 $m$ 个微批次 $B_1, \dots, B_m$ 。
对角情况（AdamW, Sophia）：作者应用Delta 方法校正。对于标量预条件器条目 $p$ ，逆函数 $T(p) = (p+\lambda)^{-1}$ 被展开至二阶。主导偏差项与微批次均值的方差成正比，对其进行估计并减去：
$\tilde{T}(\bar{p}) = \frac{1}{\bar{p} + \lambda} - \frac{\widehat{\text{Var}}(\bar{p})}{(\bar{p} + \lambda)^3}$
该校正逐个元素地应用于对角预条件器。
矩阵情况（Shampoo）：对于矩阵值预条件器，该方法在平均预条件器 $\bar{P}$ 的特征基中应用校正。将微批次预条件器投影到该基中，并对特征值（或逆根特征值）应用标量 Delta 方法校正，然后再重构矩阵。
留一法（Jackknife）替代方案：还描述了一种更通用的留一（LOO）Jackknife 校正，它可以抵消主导的 $O(1/m)$ 偏差项，尽管对于矩阵预条件器而言，其计算成本可能较高。

3. 主要贡献

统计分解：本文明确将随机预条件更新误差分解为两种有限样本偏差：梯度–预条件器耦合以及求逆的非线性。
校正框架：提出了一种实用的单步框架，结合交叉拟合预条件（消除耦合）和方差校正求逆（消除逆偏差），利用微批次变异性实现。
实证实例化：该框架针对三大优化器家族进行了实例化：
- AdamW：基于对角动量（使用 LOO 交叉拟合和詹森/Delta 校正）。
- Sophia：基于对角曲率。
- Shampoo：矩阵预条件（使用特征基校正）。

4. 实验结果

该框架在 Qwen2.5-0.5B 模型上进行了评估，涵盖两种模式：

预训练：在 FineWeb-Edu 序列（干净和混合质量/噪声）上从随机初始化开始。
指令微调：在 Alpaca 风格数据上进行监督微调（SFT）。

主要发现：

预训练性能：偏差校正在所有三种优化器家族中一致地降低了保留的预训练损失（交叉熵）：
- AdamW：损失降低了 0.149 nats（从 4.836 降至 4.687）。
- Sophia：损失降低了 0.070 nats（从 6.665 降至 6.595）。
- Shampoo：损失降低了 0.110 nats（从 5.792 降至 5.681）。
- 在混合质量（噪声）预训练中，校正后的 AdamW 变体在干净评估数据上显示出较小但积极的增益（0.019 nats），同时对受损片段的拟合不那么激进。
指令微调：结果“中性至积极”。
- AdamW：在重新调整学习率后显示出微小增益（例如 -0.001 nats）。
- Sophia：在小数据集上与标准优化几乎持平，在较大的 5K 数据集上略好。
- Shampoo：在仅注意力（attention-only）的 SFT 设置中效果基本中性。
消融实验：本文指出，朴素的两折划分（使用一半批次作为分母）可能因噪声增加而过于保守。为 AdamW 提出的留一（LOO）构造在消除耦合的同时保持了分母尺度，被证明对性能至关重要。

5. 意义与主张

本文将这项工作定位为一种实用机制，通过解决有限样本更新中固有的统计偏差来改进预条件优化器，而不是提出新的优化器架构或曲率近似。

理论洞察：它确立了标准自适应更新由于耦合和非线性，是总体预条件方向的有偏估计量。
实用价值：该校正对于对角优化器而言轻量级，并可扩展至矩阵优化器。它提高了样本效率和稳定性，特别是在估计值噪声大且训练运行时间长的预训练模式中。
适度范围：作者对效应大小的普遍性持谨慎态度。他们承认结果基于单一模型规模（0.5B）和单一种子运行。他们指出，在从强预训练检查点开始的短期指令微调运行中，样本划分的方差成本可能会抵消偏差减少带来的收益，从而导致中性结果。
结论：主要贡献在于证明校正随机更新中的有限样本偏差可以使预条件语言模型优化更有效，特别是在预训练阶段，而无需对优化器的内存占用或更新结构进行根本性改变。