🤖 machine learning

Compute Aligned Training: Optimizing for Test Time Inference

本文提出了“计算对齐训练”这一新颖框架，该框架通过推导新的损失函数，将大型语言模型的训练目标与推理阶段的测试策略相一致，从而相较于标准的监督微调（SFT）和强化学习（RL）方法显著提升了性能扩展能力。

原作者： Adam Ousherovitch, Ambuj Tewari

发布于 2026-04-29

📖 1 分钟阅读☕ 轻松阅读

原作者： Adam Ousherovitch, Ambuj Tewari

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

想象你正在为一种非常特定的期末考试培训学生。

旧方法（标准训练）：
传统上，当我们教授人工智能模型（例如那些撰写文章或解决数学问题的模型）时，我们就像一位严格的老师，单独批改每一份作业。如果学生答对了问题，我们就说“做得好！”；如果答错了，我们就说“再试一次”。我们想要学生每次都完美地答对第一个答案。

问题是，这并不能让他们为真正的考试做好准备。在现实世界中，我们不仅仅是向模型索取一个答案。我们要求它生成许多不同的答案（就像要求学生写 10 篇不同的文章草稿），然后我们挑选最好的一个，或者投票选出最常见的那个。

如果你训练学生完美地写出第一稿，他们可能会变得过于自信，从而停止探索新想法。他们可能会陷入一种“安全”的思维模式。当你后来要求他们生成 10 份草稿时，他们可能只会写出 10 份略有不同的同一份“安全”答案，而这些答案中没有一个可能是你所需的精彩、创造性的解决方案。

新方法（计算对齐训练 - CAT）：
这篇论文的作者提出了一种名为**计算对齐训练（Compute Aligned Training, CAT）*的新训练方法。他们不是根据单次作业来给学生打分，而是根据学生在实际*考试形式中的表现来评分。

以下是其工作原理，通过几个类比来说明：

1. “通过@N"类比（彩票票）

想象考试允许你为一道题购买N 张彩票（假设是 64 张）。只要你的 64 张彩票中任何一张是中奖号码，你就赢了。

标准训练： 如果学生用一张彩票就有 50% 的胜率，老师会不断逼迫他们达到 90% 或 99%。但在 64 张彩票的抽奖中，一旦你有了 50% 的胜率，购买更多彩票就足以保证获胜了。老师浪费精力试图让一个“好”学生在简单问题上变得“完美”，却忽略了学生只有 1% 胜率的那些难题。
CAT 训练： 老师意识到：“嘿，这个学生用 64 张彩票很可能已经会赢了。我不再这么严格地批改这个简单问题了。”相反，他们将所有精力集中在学生目前失败的难题上。他们教导学生分散机会，确保 64 张彩票中至少有一张是中奖票，而不是试图让某一张特定的彩票变得完美。

2. “多数投票”类比（选举）

想象考试要求模型生成 10 个答案，然后全班投票选出最受欢迎的那个。

标准训练： 老师试图让学生的第一个答案成为绝对最受欢迎的，即使它已经以压倒性优势获胜。这就像一位已经赢得 90% 选票的候选人；老师却不断告诉他要更努力地竞选，这是浪费时间。
CAT 训练： 老师关注“临界点”。如果学生的答案目前以微弱劣势落败，老师会给予巨大的提升，帮助他们冲过终点线。如果学生已经轻松获胜，老师就不再给他们额外的分数。这迫使模型将精力集中在“战场”问题上，因为在这些问题上，一点点额外的努力就能扭转选票。

3. “最佳 N 选”类比（才艺表演）

想象模型生成 10 首歌曲，而你只保留其中最好的一首。

标准训练： 老师试图让平均水平的歌曲听起来不错。这导致音乐安全、无聊、平庸，虽然从不糟糕，但也从不惊艳。
CAT 训练： 老师告诉模型：“没关系，哪怕你的 10 首歌里有 9 首很糟糕，只要第 10 首是杰作就行。”这鼓励模型去冒险，尝试奇怪、高方差的想法。它学会了产生多样化的输出，因为它知道“搜索”过程（挑选最好的那个）会过滤掉失败者并保留获胜者。

他们实际上做了什么？

研究人员通过三种具体方式测试了这一想法：

数学问题： 他们训练人工智能模型解决数学问题。当他们使用 CAT 来为“通过@N"（生成许多答案并挑选正确的那个）做准备时，与标准方法相比，模型在解决难题方面取得了显著进步。
投票： 他们训练模型为“多数投票”生成答案。同样，当应用投票策略时，CAT 模型的表现更好。
蛋白质设计： 他们甚至将这一方法测试于一种完全不同的 AI，即设计蛋白质（生命的构建模块）的 AI。在 AI 必须从许多劣质结构中找到稀有、高质量的蛋白质结构的场景中，经过 CAT 训练的模型比标准模型更擅长找到“头奖”蛋白质。

核心结论

这篇论文认为，你训练模型的方式应该与你使用它的方式相匹配。

如果你计划使用模型生成许多选项并挑选最好的一个，你不应该训练它完美地通过第一次尝试。你应该训练它擅长创建一个选项池，以便最有可能从中找到最好的那个。

他们称之为“计算对齐训练”，因为它将训练过程与测试时使用的“计算”（额外的思考能力）对齐了。这是一种在不需更强大的计算机或更多训练时间的情况下获得更好结果的方法；你只需改变游戏规则，以匹配模型实际被使用的方式。

以下是 Ousherovitch 和 Tewari 的论文《计算对齐训练：优化测试时推理》（Compute Aligned Training: Optimizing for Test Time Inference）的详细技术总结。

1. 问题陈述：训练与推理的错位

本文探讨了现代大语言模型（LLM）开发中的一个根本性脱节：训练目标针对单样本性能进行优化，而部署却日益依赖测试时计算策略（搜索）。

现状： 标准的后训练方法，如监督微调（SFT）和强化学习（RL），旨在优化单个正确输出的似然性（例如，最小化交叉熵或最大化单个样本的期望奖励）。
冲突： 现代推理策略（例如，Pass@N、多数投票、Best-of-N）生成多个候选解并将其聚合以找到最佳答案。
- 示例（Pass@N）： 如果模型有 50% 的概率正确，将其提升至 90% 对于 Pass@100 指标（其中成功已近乎确定）而言收益递减。然而，标准训练仍继续推动达到 100% 的似然性，从而在“简单”样本上浪费容量。
- 示例（多数投票）： 模型仅需 plurality（例如 35% 的概率）即可战胜对手。标准训练推动 $p \to 1.0$ ，鼓励“压倒性”胜利而非稳健的共识，这可能导致模型在被迫作为集成系统运行时性能下降。
后果： 这种错位导致策略崩溃（多样性降低），并阻碍模型有效利用增加的测试时计算预算。

2. 方法论：计算对齐训练（CAT）

作者提出了计算对齐训练（CAT），这是一个将训练目标与测试时使用的特定推理策略对齐的框架。

核心概念化

作者将推理策略概念化为可微算子（ $T$ ），而非事后启发式方法，它将基础策略 $\pi_\theta$ 转换为有效的测试时分布 $\tilde{\pi}_\theta$ ：
$\tilde{\pi}_\theta(y|x) = T(\pi_\theta(\cdot|x), \phi)$
其中 $\phi$ 代表超参数（例如，采样预算 $N$ ）。

统一损失函数

CAT 不是最小化基础策略（ $\pi_\theta$ ）下真实标签的负对数似然，而是最小化有效系统级策略（ $\tilde{\pi}_\theta$ ）下的损失：
$L_{CAT}(\theta) = -\mathbb{E}_{(x, y^*)} [\log \tilde{\pi}_\theta(y^*|x)]$

梯度缩放机制

为了使这在无需采样完整输出空间的情况下在计算上可行，作者推导了对角梯度近似。他们表明，CAT 目标的梯度可以通过将标准梯度乘以权重因子 $w(p, \phi)$ 来近似：
$\nabla_\theta L_{CAT} \approx w(p, \phi) \cdot \nabla_\theta L_{CE}$
这里， $p$ 是正确答案的基础概率， $w$ 代表相对于测试时策略增加 $p$ 的边际效用。

推导出的关键缩放因子：

Pass@N（SFT 与 RL）：
- $w(p, N) = \frac{N p (1-p)^{N-1}}{1 - (1-p)^N}$
- 行为： 抑制那些成功已得到保证的简单样本（高 $p$ ）的梯度，将学习集中在“困难”样本上，因为边际收益能改善覆盖率。
多数投票（SFT 与 RL）：
- $w(p, N, k) \approx \frac{k \binom{N}{k} p^k (1-p)^{N-k}}{\sum_{i=k}^N \binom{N}{i} p^i (1-p)^{N-i}}$
- 行为： 将梯度质量集中在共识阈值附近（ $p \approx k/N$ ）。它忽略那些已经获胜或注定失败的样本，仅关注“转折点”。
Best-of-N（RL）：
- $w \propto N (P_{<y})^{N-1}$
- 行为： 创造“赢家通吃”的动态。梯度仅强化批次中奖励最高的样本，鼓励模型探索分布的高方差、高奖励尾部，而不是坍缩到安全的均值。

理论保证

论文证明，对于竞争策略（即增加对手的概率会降低成功率），对角近似是真实梯度的保守低估。这确保优化器永远不会朝错误的方向迈出一步，充当防止策略崩溃的自然稳定器。

3. 主要贡献

统一算子框架： 将测试时策略形式化为可微算子，推导出一种通用机制，使训练能够与任意推理策略（SFT 和 RL）对齐。
超越 Pass@N 的扩展： 将先前的工作（如直接覆盖率优化）推广到多数投票和Best-of-N，涵盖 SFT 和 RL 范式。
模态无关性： 展示了该框架不仅适用于 LLM，还适用于蛋白质语言模型（PLMs），证明了其在复杂、多模态奖励景观中性能的提升。
方差分析： 提供了梯度方差膨胀的理论和实证分析，提出了归一化技术（如对数加权估计量）以稳定训练。

4. 实验结果

作者在MATH 基准（使用 Mistral-7B）和蛋白质生成（使用 ProtGPT2）上评估了 CAT。

A. 语言模型（MATH 基准）

Pass@N（SFT）： 使用 CAT 针对 $N=64$ 训练的模型牺牲了单样本准确率（Pass@1 下降了约 2.3%），但相比标准 SFT，Pass@64 覆盖率实现了7.8% 的绝对提升。
多数投票（SFT）： CAT 模型在所有推理预算下均一致优于基线。在 $N=16$ 处观察到一个“甜蜜点”，此时模型在偏差和方差之间达到了最佳平衡。
强化学习（GRPO）： 标准 RL 遭受策略崩溃，迅速饱和（Pass@32 $\approx$ 35.8%）。CAT 训练的模型保留了搜索所需的方差，在 Pass@32 处达到40.0%。
预热效应： 在 SFT 预热阶段使用 CAT 显著提高了后续 RL 的性能，甚至通过丰富模型的推理路径提升了单样本（Pass@1）准确率。

B. 蛋白质生成（蛋白质语言模型）

无条件生成（死亡之谷）： 在多模态奖励景观中（局部最优在 $h \approx 0.35$ ，全局最优在 $h \approx 0.75$ ，中间被一个“聚合谷”隔开），标准 RL 被困在局部最优。Best-of-N CAT 成功“跨越”了该山谷，通过维持高分布方差找到了全局最大值。
条件生成： CAT 模型学会了生成与输入具有逆疏水性的序列，展示了比标准 RL 更优越的扩展性和与关系任务的对齐性。

5. 意义与影响

范式转变： 论文主张，训练目标的设计必须考虑下游推理过程。随着测试时计算成为扩展模型能力的主要杠杆，训练必须针对聚合结果进行优化，而非单一样本。
效率： CAT 在不增加训练计算成本的情况下实现了这些收益。它仅根据样本对测试时策略的边际效用重新加权现有梯度。
稳定性与探索： 该框架提供了一种原则性的方法，在单样本精度与分布多样性之间进行权衡，使模型能够有效利用 Best-of-N 和多数投票等搜索过程，而不会坍缩为模式寻求行为。
通用性： 通过证明其在 LLM 和蛋白质 LM 中的有效性，论文表明“计算对齐训练”是利用测试时搜索的任何生成模型的通用原则。

总之，计算对齐训练弥合了学习与搜索之间的差距，确保模型不仅擅长生成一个答案，而且在拥有计算预算进行搜索时，能够被优化为擅长找到最佳答案。