What should post-training optimize? A test-time scaling law perspective

📄 arXiv: 2605.10716v1 📥 PDF

作者: Muheng Li, Jian Qian, Wenlong Mou

分类: cs.LG, stat.ML

发布日期: 2026-05-11

💡 一句话要点

提出尾部外推估计器（TEA），解决测试时大规模采样与训练时有限算力间的失配问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 后训练优化 Best-of-N采样 策略梯度 极值理论 强化学习 算力预算

📋 核心要点

  1. 现有后训练目标优化平均奖励，忽略了测试时Best-of-N策略依赖奖励分布上尾的特性，导致训练目标与部署策略存在本质失配。
  2. 论文提出尾部外推估计器（TEA），通过对有限训练采样进行上尾统计量外推，有效近似大规模采样下的策略梯度，弥补了训练与部署间的算力预算鸿沟。
  3. 在指令遵循任务中，TEA及Prefix-TEA在不同模型和奖励模型配置下，均展现出优于传统方法的Best-of-N性能，验证了该方法在资源受限场景下的有效性。

📝 摘要（中文）

大型语言模型日益依赖测试时策略：即采样N个响应并利用奖励模型选择最优者。然而，现有的后训练目标通常优化单个响应的平均奖励，这与由奖励分布上尾决定的“Best-of-N”性能存在本质失配。近期针对测试时感知的目标函数虽有所改善，但往往假设训练与部署阶段具有相同的每提示采样预算，这在实际中难以实现。本文研究了“预算失配”机制，即训练时仅能获取m（m远小于N）个采样，但目标是优化N个采样下的部署性能。基于奖励尾部的结构性假设，作者证明了Best-of-N目标的策略梯度可以通过外推上尾统计量从较小的采样组中近似得到。由此提出了一系列尾部外推估计器（TEA），包括直接估计器、TEA优势函数及基于矩抵消的Prefix-TEA。在指令遵循任务上的实验表明，该方法在多种模型、奖励模型及预算设置下均显著提升了Best-of-N性能。

🔬 方法详解

问题定义：现有后训练方法多优化单样本期望奖励，而实际部署采用Best-of-N策略，其性能由奖励分布的上尾决定。当训练阶段的采样预算m远小于部署阶段的N时，直接优化会导致严重的偏差，无法有效提升测试时性能。

核心思路：利用极值理论的思想，假设奖励分布的尾部具有特定的结构。通过对训练阶段获取的少量样本（m个）进行统计分析，外推估计出大规模采样（N个）下的奖励分布上尾特征，从而构建针对Best-of-N目标的策略梯度估计器。

技术框架：该方法包含三个核心组件：一是基于小规模采样的奖励分布建模；二是利用尾部外推技术（Tail-Extrapolation）计算梯度；三是引入Prefix-TEA进行偏差修正，通过矩抵消技术消除估计中的低阶偏差，确保在有限预算下梯度估计的稳定性。

关键创新：首次在预算失配（m << N）的设定下，将Best-of-N优化问题转化为尾部统计量的外推问题。相比于传统的强化学习目标，该方法直接针对测试时行为进行优化，且无需在训练时进行大规模采样。

关键设计：核心技术包括TEA（Tail-Extrapolated Advantage）估计器，它通过对奖励分布的尾部进行参数化建模来近似梯度；以及Prefix-TEA，通过固定阶数的去偏处理，利用矩抵消机制进一步提升了在极端分布下的估计精度。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验在多种指令遵循数据集上进行，对比了标准RLHF及现有测试时感知目标。结果显示，TEA和Prefix-TEA在不同模型规模（如Llama系列）和奖励模型下，均能显著提升Best-of-N的胜率。在m=4、N=64的预算失配场景下，该方法相较于基线方法在奖励得分上实现了显著的性能增益，验证了其在资源受限下的鲁棒性。

🎯 应用场景

该研究适用于所有依赖Best-of-N采样策略的大型语言模型部署场景，特别是在算力受限的后训练阶段。其价值在于无需增加训练时的采样开销，即可显著提升模型在推理阶段的指令遵循能力、代码生成质量及复杂逻辑推理表现，具有极高的工业落地价值。

📄 摘要（原文）

Large language models are increasingly deployed with test-time strategies: sample $N$ responses, score them with a reward model or verifier, and return the best. This deployment rule exposes a mismatch in post-training: standard objectives optimize the mean reward of a single response, whereas best-of-$N$ performance is governed by the upper tail of the reward distribution. Recent test-time-aware objectives partly address this mismatch, but typically assume that training can use the same per-prompt rollout budget as deployment, which is impractical when post-training must cover many prompts while deployment can allocate much larger per-prompt test-time compute. We study this budget-mismatch regime, where only $m\ll N$ per-prompt rollouts are available during training but the target objective is best-of-$N$ deployment. Under structural assumptions on the reward tails, we show that the policy gradient of the best-of-$N$ objective can be approximated from a much smaller rollout group by extrapolating upper-tail statistics. This yields a family of Tail-Extrapolated estimators for best-of-$N$-oriented post-training: a simple direct estimator, Tail-Extrapolated Advantage (TEA), and a fixed-order debiased Prefix-TEA estimator based on moment cancellation. Experiments on instruction-following tasks show that TEA and Prefix-TEA improve best-of-$N$ performance across different language models, reward models and datasets under various training and test-time budget settings.