What should post-training optimize? A test-time scaling law perspective¶
会议: NeurIPS 2026
arXiv: 2605.10716
OpenReview: https://openreview.net/forum?id=u0f6gRsMJ0
领域: 对齐/RLHF
关键词: 后训练, 测试时缩放定律, Best-of-N, 尾部外推, 优势估计
一句话总结¶
针对后训练单样本均值奖励与部署端 Best-of-N 极值采样之间的目标不匹配以及采样预算不对称难题,提出基于高斯尾部统计外推的尾部外推优势估计器(TEA)及去偏版本 Prefix-TEA,仅用少量训练采样即可精准优化大预算推理性能。
研究背景与动机¶
现代大语言模型的标准落地流程高度依赖部署阶段的测试时计算扩展(Test-Time Compute),其中最为基础且广泛应用的策略是 Best-of-\(N\)(BoN)采样:针对给定 prompt,模型独立采样生成 \(N\) 个候选回答,经由奖励模型(RM)或验证器打分后返回得分最高的最优候选。然而,这一主流部署范式与标准的强化学习后训练目标存在本质错配。现有的 RLHF 或 GRPO 算法主要优化单次采样的期望奖励均值,但 BoN 部署的实际表现完全由奖励分布的最上层极端尾部(Upper Tail)决定。更严重的是,过往针对单样本均值的奖励微调已被广泛证实会显著压缩生成多样性,导致模式坍塌,直接摧毁了依靠多次采样探索极端优质解的随机性基石。
尽管近期已有少数关注测试时感知的后训练方法(如直接针对 pass@\(k\)、max@\(k\) 或 BoN 的训练目标),但现有算法普遍暗含一个严苛假设:训练阶段必须使用与部署阶段相当的逐 prompt 采样预算(即 matched-budget,\(m \approx N\))。在实际工业级后训练中,这一假设完全不可行。训练阶段模型必须兼顾覆盖数以万计的多样化 prompt 数据集,平摊到单个 prompt 上的并行生成预算极为受限(通常仅 \(m \ll N\));而在推理部署阶段,随着推理计算成为独立 Scaling 轴,人们往往愿意对单道难题投入巨大的搜索与重采样预算(\(N\) 可达 128、256 乃至上千)。这种“训练低预算 \(m\) vs 部署高预算 \(N\)”的预算错配(Budget Mismatch),构成了测试时感知训练落地的一大核心瓶颈。
本文敏锐地从测试时缩放定律(Test-Time Scaling Laws)的视角切入破解该矛盾:既然在给定策略下,BoN 的期望最优性能随样本量 \(N\) 的增长规律能够被奖励分布的高层极端尾部所刻画,那么便无需在训练时直接耗费 \(N\) 次采样,而是可以用少量采样(\(m \ll N\))拟合局部上尾部充分统计量,进而显式外推出面向大预算 \(N\) 的策略梯度。核心 idea:基于奖励分布极端上尾部的高斯渐近结构建立测试时目标的高精度代理函数,推导仅需 \(m \ll N\) 样本即可计算的尾部外推优势估计器 TEA 及其矩消除去偏变体 Prefix-TEA,实现以小博大的最佳样本梯度优化。
方法详解¶
整体框架¶
本文的核心目标是在训练阶段每道 prompt 仅生成 \(m\) 个样本(\(m \ll N\))的前提下,构建针对部署端 Best-of-\(N\) 目标 \(V_N(\theta; x) = \mathbb{E}[\max_{1 \leqslant i \leqslant N} R(x, y_i) \mid x]\) 的无偏或低偏差策略梯度估计。
整个优化流程包含四个连贯阶段:首先,模型对输入 prompt 采样生成 \(m\) 个候选回复并由奖励模型评分;其次,利用分位数截断(默认取前 \(\alpha = 0.25\))圈定极端高分区域,提取局部高斯尾部统计量(经验截断阈值 \(\hat{r}_m\)、截断均值 \(\hat{\mu}_m\) 与截断方差 \(\hat{\sigma}_m\));随后,依据极值理论中的尾部曲率缩放因子 \(\widetilde{c}_N\),对落入上尾部的样本赋予结合线性超额奖励与方差二次修正的加权尾部优势,同时利用 Prefix-TEA 的多尺度前缀抵消有限样本系统性偏差;最后,经正部截断与组内均值中心化后,反传策略梯度并结合 KL 散度约束更新策略参数。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入 Prompt: x"] --> B["阶段1:有限分组采样<br/>生成 m 个候选回复并计算奖励 R_i"]
B --> C["阶段2:高斯尾部代理目标构建<br/>提取经验阈值、截断均值与方差"]
C --> D["阶段3:尾部外推优势估计 (TEA / Prefix-TEA)<br/>线性超额奖励 + 方差二次曲率修正"]
D --> E["阶段4:多尺度矩消除去偏与中心化<br/>正部截断与组内中心化优势"]
E --> F["输出: 策略梯度反传更新与 KL 正则"]
关键设计¶
1. 高斯尾部代理目标构建:将极端极值估算转化为局部尾部统计量外推
标准 Best-of-\(N\) 价值函数 \(V_N(\theta; x)\) 完全受制于极端尾部分布,直接求导需在 \(N\) 阶统计量上采样。为化解训练预算不足的痛点,本文引入高斯尾部假设(Assumption 1):假定提示词 \(x\) 诱导的奖励分布在上尾部 \(2\alpha\) 分位点以上满足局部一维高斯密度。基于该结构假设,引理 1 证明了 \(V_N(\theta; x)\) 与局部前 \(\alpha\) 区域的高阶统计量之间存在指数级逼近关系: $\(|V_N(\theta; x) - (\mu_{\theta, \alpha}(x) + \widetilde{c}_N \sigma_{\theta, \alpha}(x))| \leqslant C(1 - 2\alpha)^N\)$ 其中 \(\mu_{\theta, \alpha}(x)\) 与 \(\sigma_{\theta, \alpha}(x)\) 为真实奖励超过上 \(\alpha\)-分位数 \(r_{\theta, \alpha}(x)\) 时的条件均值与条件标准差,外推系数 \(\widetilde{c}_N = (c_N - \lambda_\alpha) / \sqrt{\delta_\alpha}\) 纯粹由标准正态极值期望 \(c_N = \mathbb{E}[\max_{1\leqslant i\leqslant N} Z_i]\) 与截断正态常数解析确定。由此,庞大部署预算 \(N\) 的极值优化被解析转化为针对局部稳定统计量 \(\mu_{\theta, \alpha}(x) + \widetilde{c}_N \sigma_{\theta, \alpha}(x)\) 的可微代理目标 \(\widehat{J}_{\mathrm{TT}}(\theta)\)。
2. 尾部外推优势估计 (TEA):线形溢出项与方差曲率项的协同加权
将代理目标求导后(引理 2),其策略梯度展现出极具物理启发性的结构: $\(g_\theta(x) = \frac{1}{\alpha} \mathbb{E}_{y \sim \pi_\theta(\cdot \mid x)} \left[ \mathbb{I}\{R(x, y) \geqslant r_{\theta, \alpha}(x)\} \widetilde{R}_{\eta_{\theta, \alpha}(x)}(R(x, y)) \nabla_\theta \log \pi_\theta(y \mid x) \right]\)$ 其中尾部赋形奖励函数定义为: $\(\widetilde{R}_\eta(u) := (u - r) + \frac{\widetilde{c}_N}{2\sigma} \left( (u - \mu)^2 - (r - \mu)^2 \right)\)$ 在有限样本 \(m\) 下,通过经验分位数 \(q_m = \lceil \alpha m \rceil\) 选出前 \(q_m\) 个最高奖励集合 \(\mathcal{I}_m\),直接代入经验截断均值 \(\hat{\mu}_m\)、经验截断方差 \(\hat{\sigma}_m^2\) 与阈值 \(\hat{r}_m\)。该式内含精妙机制:指示函数 \(\mathbb{I}\) 将优化注意力完全收束到前 \(\alpha\)(如前 25%)的优质样本上,阻断低分回复对高阶潜力的干扰;第一项 \((u - \hat{r}_m)\) 奖励样本超出截断线的线性溢出值;第二项二次项则对应方差曲率修正——当目标部署预算 \(N\) 增大时,\(\widetilde{c}_N\) 相应增加,赋予高分极端离群点更高的正向梯度权重,主动拉伸奖励分布的上尾跨度。在训练工程落地中,TEA 进一步取原始优势的正部并执行组内均值中心化:\(A_i^{\mathrm{TEA}} = (A_i^{\mathrm{raw}})_+ - \frac{1}{m}\sum_{j=1}^m (A_j^{\mathrm{raw}})_+\)。
3. 多尺度矩消除去偏估计器 (Prefix-TEA):打破小样本系统性阶次偏差
定理 1 揭示出直接代入经验统计量的 TEA 虽具备 \(\mathcal{O}(m^{-1} \log N)\) 的方差界,但其梯度估计存在 \(\mathcal{O}(m^{-1} \sqrt{\log N})\) 的系统性有限样本偏差。在随机梯度上升中,非零偏差下界会限制算法收敛到真正驻点的精度。为此,本文设计了基于矩消除(Moment Cancellation)与广义折刀法(Generalized Jackknife)的 Prefix-TEA 估计器。利用渐近展开定理,通过在大小为 \(m\) 的单次采样组内截取 \(J\) 个不同规模的嵌套前缀子集 \((m_1, \dots, m_J)\),分别计算局部经验尾部向量,并求解代数 Vandermonde 矩消除方程得出加权系数 \((w_1, \dots, w_J)\)。在主实验配置 \((k=2, J=4)\) 下,定理 2 从理论上证明 Prefix-TEA 将梯度估计偏差从 \(\mathcal{O}(m^{-1})\) 骤降至 \(\mathcal{O}(m^{-2})\),同时维持 \(\mathcal{O}(m^{-1})\) 的极佳方差阶次,从而在理论与大批量训练场景中提供了更稳健的收敛性保证。
损失函数 / 训练策略¶
在训练执行层面,所有方法统一集成至通用分组策略梯度训练框架。对采样组内的每个候选 \(y_i\),其损失函数结合策略梯度重要性采样裁剪与参考模型 \(\pi_{\mathrm{ref}}\) 的 token 级 KL 散度约束: $\(\mathcal{L}(\theta) = -\frac{1}{m} \sum_{i=1}^m A_i \log \pi_\theta(y_i \mid x) + \beta \, \mathrm{KL}\bigl(\pi_\theta(\cdot \mid x) \,\|\, \pi_{\mathrm{ref}}(\cdot \mid x)\bigr)\)$ 其中 \(A_i\) 即为 TEA 或 Prefix-TEA 产出的停止梯度标量优势。基准超参数配置为:超参数 \(\alpha = 0.25\),\(N_{\mathrm{target}} = 128\)(控制尾部外推曲率),KL 系数 \(\beta = 0.1\),学习率 \(10^{-5}\),采用余弦退火调度,总训练步数 450 步,每 25 步在独立的 128 个验证 prompt(每题采样 32 条取 top-10 均值)上进行无偏检查点早停评测。
实验关键数据¶
主实验¶
主实验在 UltraFeedback 指令遵循数据集上开展,基座策略为 Llama-3.2-1B-Instruct,训练与评测奖励模型均采用 Skywork-Reward-V2-Llama-3.1-8B。在保留测试集 core250 上,各方法统一采用 \(m=64\) 训练采样预算,评测端对每个 prompt 采样 256 条回答并计算分组 Best-of-\(N\) 收益。
| 方法 | 训练预算 | 选定步数 | bo1 | bo8 | bo32 | bo64 | bo128 | bo256 |
|---|---|---|---|---|---|---|---|---|
| GRPO | \(m=64\) | 100 | -1.129 | 6.731 | 9.803 | 11.033 | 12.149 | 13.273 |
| GRPO-Z | \(m=64\) | 450 | 0.448 | 7.690 | 10.684 | 11.999 | 13.124 | 14.207 |
| BoN-max second | \(m=64\) | 250 | -1.445 | 6.871 | 10.017 | 11.304 | 12.510 | 13.536 |
| CAT-BoN | \(m=64\) | 425 | -0.800 | 7.083 | 10.158 | 11.395 | 12.482 | 13.484 |
| BoN mean (\(2\times\)) | \(m=128\) | 275 | -1.431 | 6.579 | 9.672 | 11.009 | 12.191 | 13.289 |
| Prefix-TEA (Ours) | \(m=64\) | 450 | 0.401 | 8.031 | 11.105 | 12.403 | 13.590 | 14.548 |
| TEA (Ours) | \(m=64\) | 425 | 0.581 | 8.241 | 11.299 | 12.568 | 13.718 | 14.750 |
如上表所示,TEA 与 Prefix-TEA 在整个 Best-of-\(N\) 前沿上全线碾压基线。相较于 GRPO,TEA 在 bo128 与 bo256 上分别取得 +1.569(95% CI: [+1.305, +1.819])与 +1.477(95% CI: [+1.177, +1.767])的显著净提升;即便面对使用双倍训练算力(\(m=128\))的强基线 BoN mean \(2\times\),仅用 \(m=64\) 的 TEA 仍保持绝对胜出。
消融实验¶
消融实验重点探究训练预算 \(m \in \{16, 32, 64\}\) 到测试超额预算 \(N \in \{128, 256, 512\}\) 的跨尺度外推鲁棒性,以及外推目标参数 \(N_{\mathrm{target}}\) 的敏感度(表内均为相对同预算 GRPO 的差值)。
| 配置 / 训练预算 | 评估指标 | 相对于 GRPO 增益 | 95% 置信区间 | 平均长度变化 | 说明 |
|---|---|---|---|---|---|
| \(m=16 \to N=512\) | \(\Delta \text{bo512}\) | +0.867 | [+0.520, +1.212] | -22.9 | 极度预算错配下仍保持大胜 |
| \(m=32 \to N=512\) | \(\Delta \text{bo512}\) | +1.427 | [+0.988, +1.853] | -34.4 | 中等预算时提升显著扩大 |
| \(m=64 \to N=512\) | \(\Delta \text{bo512}\) | +1.519 | [+1.149, +1.865] | -10.5 | 默认训练预算,增益最强 |
| \(N_{\mathrm{target}}=128\)(默认) | \(\Delta \text{bo128}\) | +1.569 | [+1.307, +1.828] | -10.5 | 最优偏差-方差平衡点 |
| \(N_{\mathrm{target}}=256\) | \(\Delta \text{bo128}\) | +0.828 | [+0.588, +1.052] | -11.7 | 外推过远导致估计方差轻微增大 |
| \(N_{\mathrm{target}}=512\) | \(\Delta \text{bo128}\) | +0.308 | [+0.102, +0.513] | +0.7 | 保持正收益但增益有所衰减 |
关键发现¶
- 打破严苛预算限制:在极其严苛的 \(m=16 \ll N=512\) 极端不对称场景下,TEA 相对 GRPO 在 bo512 上依旧斩获 +0.867 的稳定净增益,证明尾部外推机制能有效利用有限局部样本推测全局极端分布。
- 与生成长度无关的真提升:生成文本长度分析证实,TEA 生成的平均长度(311.4 tokens)甚至短于 GRPO(321.9 tokens),彻底排除了 RM 偏好长回复带来的作弊虚增(Length Bias)。
- 独立大模型裁判一致裁决:使用 DeepSeek-v4 Pro 进行严格随机双盲评测,在 core250 的 bo128 采样结果中,DeepSeek-v4 Pro 判定 TEA 对 GRPO 的胜率高达 60.8%(152 胜 98 负),对 GRPO-Z 胜率达 54.4%,对最佳测试感知基线胜率达 58.0%。
- 广泛的泛化与模型鲁棒性:在切换不同策略骨干(Qwen3-4B-Instruct 上 \(\Delta\text{bo256} = +3.542\),Llama-3.1-8B LoRA 上 \(\Delta\text{bo256} = +1.102\))、不同数据集(Anthropic HH-helpful 上 \(\Delta\text{bo256} = +1.800\))以及不同打分模型(Skywork-Qwen3-8B RM 上胜率达 56.0%)时,TEA 的性能提升无一例外持续成立。
亮点与洞察¶
- 视角的范式革新:以往关于 Test-Time Scaling Laws 的研究仅将高斯尾部模型视作“评测预测工具”或用于离线拟合;本文首次将其翻转为“在线策略优化器”,将统计外推理论直接植入强化学习策略梯度的解析推导中,视角极其敏锐。
- 真实工业约束下的解耦设计:深刻指出了“单题训练成本受限”与“单题推理算力无限”这一学术界常忽略的工业核心痛点,摆脱了以往方法对 \(m=N\) 匹配采样的笨重依赖,用优雅的数学公式消解了硬件瓶颈。
- 可复用的通用即插即用优势算子:TEA 与 Prefix-TEA 完全封装为简单的优势权重分配模块,其正部截断与组内中心化算子可无缝平替任何基于 PPO、GRPO 或 ReMax 的 RL 基础设施,无需更改网络结构或引入额外 Critic。
局限与展望¶
- 连续光滑高斯尾部假设的边界:本文的理论与优势推导依赖于连续光滑奖励分布的局部高斯假设。针对诸如代码执行 pass@\(k\) 或数学答案完全匹配等离散、非光滑的 0-1 稀疏奖励,该高斯连续假设可能失效,需进一步探索重尾(如极值帕累托 Pareto 分布)或自适应非参数尾部估计。
- 树搜索与代理交互的推广:当前算法针对的是单轮独立的 Best-of-\(N\) 采样空间。如何将尾部外推思想延伸到多轮 MCTS 树搜索、自反思(Self-Refine)以及带外部工具调用的多步智能体轨迹搜索中,仍有待深入研究。
相关工作与启发¶
- vs GRPO / GRPO-Z (Shao et al., 2024):GRPO 采用组内奖励均值中心化,本质上是对准单样本平均奖励的提升;而 TEA 引入截断指示函数与二次曲率修正,专为拉伸高分极端上尾设计,梯度的真实 oracle 余弦相似度从 GRPO 的 0.176 跃升至 0.761。
- vs Chow BoN-RL (Chow et al., 2024):Chow BoN-RL 需将小样本组硬拆为选择集与校正集进行蒙特卡洛纠偏,小样本下梯度方差极大且极其脆弱;TEA 则利用尾部分布整体拟合,充分利用全组统计量,表现极为稳健。
- vs Compute Aligned Training (CAT-BoN, Ousherovitch & Tewari, 2026):CAT-BoN 通过经验累积经验分布函数(ECDF)对样本做幂次加权,在 \(m \ll N\) 时受限于经验分位数的粗糙离散阶梯;而 TEA 依托参数化极值外推,平滑刻画了超越观测样本极值的密度演化趋势。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次从测试时缩放定律推导强化学习后训练目标,将极值理论精妙转化为策略梯度估计。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 UltraFeedback 与 HH-RLHF 双数据集、多模型骨干、跨奖励模型、独立 LLM 裁判及深入的梯度对齐机制剖析。
- 写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰严谨,从引理到有限样本上界层层递进,动机阐述切中要害。
- 价值: ⭐⭐⭐⭐⭐ 直击推理算力扩展时代训练与部署预算错配的核心痛点,具有极高的工业实用与理论指导价值。