Compute-Aligned Training: Optimizing for Test-Time Inference¶
会议: NeurIPS 2026
arXiv: 2604.24957
OpenReview: https://openreview.net/forum?id=HUYwkuQIwd
代码: https://github.com/adamoosya/compute-aligned-training
领域: Reasoning
关键词: 测试时计算扩展, 策略对齐训练, 多数投票, Best-of-N, 强化学习
一句话总结¶
将测试时推理策略形式化为作用在基础策略上的可微算子，推导出对齐Pass@N、多数投票及Best-of-N边际效用的梯度加权准则，使模型在不增加训练算力的前提下大幅提升测试时扩展性能。
研究背景与动机¶
随着大语言模型从训练时参数扩展转向测试时计算（Test-Time Compute, TTC）扩展，以思维链采样、Pass@N、多数投票（Majority Vote）以及Best-of-N验证为代表的推理搜索策略已成为复杂推理任务的核心。然而，现有的后训练范式（如监督微调 SFT 与标准强化学习 RL）从根本上建立在优化“单个生成样本的期望奖励或对数似然”之上。这种单点优化的训练目标隐式迫使模型将概率质量收敛在少数“安全且高似然”的输出路径上，导致策略分布过早坍缩，丧失了搜索所需的探索多样性。
这种优化目标与实际部署模式的割裂，在测试时扩展中引发了显著的效率浪费与性能瓶颈。在 Pass@N 场景中，若模型生成正确答案的基础概率已达 50%，将其硬推到 90% 对最终 Pass@100 的边际收益微乎其微，但标准交叉熵仍会分配巨量梯度去拟合这一已被解决的样本；在多数投票场景中，正确答案只需以 35% 的微弱优势压制 10% 的对手即可胜出，但标准训练却将其视为高误差而强行向确定性拟合，耗费有限的模型容量。现有研究要么局限于特定策略（如针对 Pass@N 的 DCO），要么依赖复杂的变分近似，缺乏统一的数学视角来解释并解决这一全局性失配。
本文的核心思考在于：测试时推理策略本质上不是打补丁式的启发式后处理，而是以基础生成策略和搜索预算为输入、输出最终答案的确定性映射算子。核心 idea：将测试时推理策略建模为作用在基础策略上的函数算子，建立对齐测试时诱导分布的目标函数，通过对角梯度分解导出无需额外采样的样本边际效用加权系数，实现无额外训练开销的策略与推理协同优化。
方法详解¶
整体框架¶
计算对齐训练（Compute Aligned Training, CAT）的核心在于将部署环境下的推理策略视为概率单纯形上的映射算子 \(\mathcal{T}\)。给定输入提示词 \(x\)，基础策略 \(\pi_\theta(\cdot|x)\) 产生候选输出分布；推理策略（如 Pass@N、多数投票、Best-of-N）结合超参数 \(\phi\)（如采样预算 \(N\)、共识阈值 \(k\)），将基础策略转化为一个等效的系统级部署分布 \(\tilde{\pi}_\theta(y|x) = \mathcal{T}(\pi_\theta(\cdot|x), \phi)(y)\)。
CAT 的训练目标是在该等效分布 \(\tilde{\pi}_\theta\) 上最大化真实答案的似然（SFT 范式）或期望效用（RL 范式）。由于直接对整个输出空间计算完整 Jacobian 矩阵会导致计算量爆炸，CAT 引入了对角梯度分解（Diagonal Gradient Decomposition），将完整梯度解构为“标准梯度乘以标量边际效用权重”与“非对角残差向量”，并在理论上证明了在比例衰减假设下非对角残差与主梯度保持安全一致，从而在零额外计算开销下完成精准的测试时对齐。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入数据与基础策略预测<br/>p = π_θ(y*|x)"] --> B["算子形式化与等效分布构建<br/>π~_θ = T(π_θ, ϕ)"]
B --> C["对角梯度分解与残差有界保证<br/>∇ℒ_CAT = w · ∇ℒ_CE + ε_vec"]
C --> D["边际效用动态梯度重加权<br/>Pass@N / MajVote / BoN 缩放系数"]
D --> E["策略参数更新与测试时搜索扩展<br/>保持探索方差与决策边界锐度"]关键设计¶
1. 算子形式化与等效测试时分布：将推理搜索抽象为策略变换 传统训练直接对单样本输出似然 \(\log \pi_\theta(y^*|x)\) 施加惩罚，忽略了下游集成或筛选对成功概率的重塑。CAT 将任意给定超参数 \(\phi\)（如采样预算 \(N\)）的测试时策略抽象为算子 \(\mathcal{T}: \Delta^{|\mathcal{Y}|-1} \times \Phi \to \Delta^{|\mathcal{Y}|-1}\)，导出等效系统策略 \(\tilde{\pi}_\theta(y|x) = \mathcal{T}(\pi_\theta(\cdot|x), \phi)(y)\)。在 SFT 设定下，训练损失定义为该诱导分布下的负对数似然： $\(\mathcal{L}_{\text{CAT}}(\theta) = -\mathbb{E}_{(x, y^*) \sim \mathcal{D}} [\log \tilde{\pi}_\theta(y^*|x)]\)$ 对于 Pass@N 策略，正确答案出现至少一次即为成功，未命中则并入虚拟失败状态，其有效成功概率严格等于 \(\tilde{p} = 1 - (1-p)^N\)；对于多数投票，成功概率则为正确答案在 \(N\) 次采样中命中次数达到或超过共识阈值 \(k\) 的累积二项分布概率 \(\tilde{p}_{\text{maj}} = \sum_{i=k}^N \binom{N}{i} p^i (1-p)^{N-i}\)。
2. 对角梯度分解与安全有界保证：摆脱全量雅可比的零开销标量更新 直接根据链式法则对 \(\mathcal{L}_{\text{CAT}}\) 求导需遍历指数级输出空间 \(\mathcal{Y}\) 中所有候选输出 \(y'\)，计算 \(\frac{\partial \tilde{p}}{\partial \pi_\theta(y'|x)}\)，这在自回归生成中计算不可行。CAT 提出了对角梯度分解定理，严格分离出对角主梯度与非对角残差 \(\epsilon_{\text{vec}}\)： $\(\nabla_\theta \mathcal{L}_{\text{CAT}} = \left( \frac{p}{\tilde{p}} \frac{\partial \tilde{p}}{\partial p} \right) \cdot \nabla_\theta \mathcal{L}_{\text{CE}} + \epsilon_{\text{vec}}\)$ 为了保证舍弃 \(\epsilon_{\text{vec}}\) 后的标量代理梯度 \(w(p, \phi) \cdot \nabla_\theta \mathcal{L}_{\text{CE}}\) 依然是有效的下降方向，CAT 引入了比例衰减假设（Proportional Decay Assumption）：即当正确答案概率 \(p\) 提升时，错误竞争路径按其当前比例同步失权 \(\nabla_\theta \pi_\theta(y'|x) \approx -\frac{\pi_\theta(y'|x)}{1-p} \nabla_\theta p\)。在这一假设下，CAT 证明了只要推理策略满足竞争性（Competitive，即增加错误答案概率不会增加系统成功率）和对抗对称性（Adversarial Symmetry，即对最强竞争对手的敏感度不大于对正确答案的敏感度），残差向量与对角梯度的内积非负 \(\langle g_{\text{diag}}, \epsilon_{\text{vec}} \rangle \geq 0\)。标量加权梯度构成真实梯度的安全下界估计，从数学上根除了方向偏移与策略崩溃风险。
3. 多范式闭式缩放因子推导：跨 SFT 与 RL 的边际效用精准分配 针对不同推理策略与后训练范式，CAT 分别导出了无需额外推理开销的闭式缩放因子 \(w(p, \phi)\)。对于 SFT 下的 Pass@N，因子为 \(w_{\text{pass, SFT}} = \frac{N p (1-p)^{N-1}}{1-(1-p)^N}\)，随着 \(p \to 1\) 单调衰减至 0，对易样本自适应梯度退火，将模型容量倾斜给难样本；对于多数投票，SFT 因子为 \(w_{\text{maj, SFT}} = \frac{k \binom{N}{k} p^k (1-p)^{N-k}}{\sum_{i=k}^N \binom{N}{i} p^i (1-p)^{N-i}}\)，将梯度聚焦于决策临界边界 \(k\) 附近的样本。 在强化学习设定下，优化目标从对数概率变为系统期望奖励 \(J_{\text{CAT}}(\theta) = \mathbb{E}_{y \sim \tilde{\pi}_\theta}[R(y|x)]\)，其策略梯度缩放因子转化为纯导数形式 \(w(p, \phi) = \frac{\partial \tilde{p}}{\partial p}\)。在 Best-of-N RL 中，通过 rollout 批次中劣于当前候选样本的经验分位数 \(\hat{P}_{<y}\)，动态估计加权系数 \(w_{\text{bon, RL}} = N (\hat{P}_{<y})^{N-1}\)，形成赢家通吃的强化机制，促使策略生成具有极端方差与高奖励峰值的候选解。
损失函数 / 训练策略¶
在实现过程中，为防止初期策略概率过低导致加权系数数值不稳定，模型先在标准交叉熵目标下进行 1-2 个 epoch 的充分预热（Warmup），使基础策略进入具备基础推理轨迹的概率区间；随后切换至 CAT 目标。在 RL 范式中，CAT 直接将缩放系数乘以 GRPO 的归一化优势项（Normalized Advantages），对每个 rollout 轨迹进行自适应重加权，全程不引入任何额外的价值网络或 Monte Carlo 树搜索。
实验关键数据¶
主实验¶
论文在 MATH 基准上基于 Mistral-7B-Instruct-v0.2 进行了微调与强化学习实验，涵盖 Pass@N、Majority Vote 以及蛋白质语言模型 ProtGPT2 上的 Best-of-N 序列生成。以下汇总 SFT 与 RL 模式下不同推理预算的代表性评估结果。
| 训练范式与模型配置 | 测试指标 @1 | @4 / @8 | @16 | @32 | @64 |
|---|---|---|---|---|---|
| SFT Baseline (Cross-Entropy) | 15.8% (1.6) | 29.6% (@4) | 46.3% (2.2) | 53.7% (2.2) | 59.8% (2.2) |
| CAT Pass@N=4 (SFT) | 15.2% (1.6) | 30.5% (@4) | 49.2% (2.2) | 57.5% (2.2) | 63.8% (2.1) |
| CAT Pass@N=16 (SFT) | 14.5% (1.6) | 30.1% (@4) | 50.9% (2.2) | 59.8% (2.2) | 66.2% (2.1) |
| CAT Pass@N=64 (SFT) | 13.5% (1.5) | 28.5% (@4) | 49.4% (2.2) | 59.3% (2.2) | 67.6% (2.1) |
| SFT Baseline (MajVote, CoT) | 21.9% (1.8) | 22.9% (@4) | 23.6% (1.9) | 23.8% (1.9) | 23.9% (1.9) |
| CAT MajVote N=16 (SFT) | 22.4% (1.9) | 24.3% (@4) | 25.9% (2.0) | 26.1% (2.0) | 26.2% (2.0) |
| Standard RL (GRPO) | 8.3% (1.2) | 24.1% (@8) | 30.3% (2.1) | 35.8% (2.1) | - |
| CAT Pass@16 (RL) | 8.3% (1.2) | 25.1% (@8) | 32.5% (2.1) | 40.0% (2.2) | - |
| Standard RL (MajVote) | 11.1% (1.4) | 14.4% (@4) | 19.2% (1.8) | - | - |
| CAT Maj@4 (RL) | 10.9% (1.4) | 15.0% (@4) | 23.0% (1.9) | - | - |
消融实验¶
为了验证对抗竞争项对决策边缘稳定性的影响，论文在 MATH 子集上对比了标准 SFT、仅考虑对角近似的 CAT Majority Vote 以及引入逐 token 竞争边际最大化的混合对比损失（Pairwise Contrastive, M=4 局部对抗 token）。
| 模型与优化策略 | Pass@1 精度 | Majority Vote 精度 (N=64) | 聚合增益 (Aggregation Gain) |
|---|---|---|---|
| Standard SFT | 12.90% | 16.20% | +3.30% |
| CAT Majority Vote (对角标量加权) | 13.03% | 17.60% | +4.57% |
| Pairwise Contrastive (结合对抗边缘惩罚) | 12.47% | 18.20% | +5.72% |
关键发现¶
- 搜索税（Search Tax）现象明显：针对大预算（如 \(N=64\)）对齐训练的模型在单次生成准确率（Pass@1）上承受了约 1%~2.3% 的跌幅，但该折损换来了更宽阔的探索空间，并在高计算预算下实现了高达 7.8% 的绝对增益。
- 方差膨胀与偏差权衡：Majority Vote 目标中，随着目标预算 \(N\) 增加，难以收敛样本的梯度方差按 \(O(N^2)\) 放大。实验观察到 \(N=16\) 呈现最佳性能甜蜜点，而过大的 \(N\) 会因梯度方差过高导致增益略微回落。
- 模态无关的泛化能力：在 ProtGPT2 的非凸疏水性蛋白质生成任务中，标准 RL 陷入了局部最优平台（奖励仅 8.25~11.40），而通过 Best-of-N CAT 训练的模型自发形成了双峰分布，成功跨越能量洼地，期望最大奖励直接攀升至 12.16（接近理论上限）。
亮点与洞察¶
- 算子视角的范式转换：将后处理推理策略（Pass@N、多数投票）从单纯的测试期外挂转化为可微分的策略算子，建立了连接训练目标与搜索机制的通用理论基座。
- 数学优美的有界性与下界证明：在严格的几何分析下证明了对抗对称策略的非对角残差与主梯度方向共线偏正，使单标量加权不仅具备高计算效率，更成为保证优化稳定性的保守安全梯度。
- 零额外算力适配工业落地：不同于复杂的树搜索或多阶段蒸馏方案，CAT 仅修改反向传播时的梯度标量加权或优势项，无需额外显存或采样开销，可直接插入任何 SFT 与 RL 管线。
局限与展望¶
- 高阶搜索策略的形式化难度：当前方法主要适配单阶段采样过滤策略（如 Pass@N、多数投票、BoN），对于多轮交互式智能体或带有回溯机制的蒙特卡洛树搜索（MCTS），算子 \(\mathcal{T}\) 难以推导解析形式。
- 大预算下的梯度方差膨胀：Majority Vote 与高阶 BoN 在极端难样本上的方差放大严重，未来需要引入原理级方差缩减（Variance Reduction）与自适应梯度裁剪技术。
- 对角近似在长程依赖下的衰减：在极端多峰、存在强竞争伪解的长链推理中，忽略非对角交叉协方差可能低估对抗 token 的破坏力，结合轻量对比辅助损失是进一步强化的可行路径。
相关工作与启发¶
- vs DCO (Direct Coverage Optimization): DCO 首次发现了标准交叉熵在 Pass@N 上的过度拟合缺陷并提出了覆盖优化，但其仅局限于 SFT 和 Pass@N 指标。CAT 构成了 DCO 的泛化理论超集，成功扩展至多数投票、RL 范式以及非 NLP 领域。
- vs BoNBoN Alignment: BoNBoN 旨在通过蒸馏使单个样本逼近 BoN 采样分布；而 CAT 则是在明确部署时使用多采样推理的前提下，对基座模型进行针对性塑形，使其充当更优的集成候选池生成器。
- vs RL 推理时间优化 (Tang et al., 2025): Tang 等人通过直接构建基于重采样的梯度估计器来优化 Pass@k，但存在采样估计方差大、计算负担重的问题。CAT 依赖闭式导数与分位数近似，避免了高昂的训练采样。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次从算子可微性与对角梯度分解视角系统构建统一测试时对齐框架。
- 实验充分度: ⭐⭐⭐⭐⭐ 跨越 SFT、RL（GRPO）以及蛋白质生成三大维度，理论推导与实验数据严密自洽。
- 写作质量: ⭐⭐⭐⭐⭐ 逻辑结构严谨，数学推导与直觉解释层次分明。
- 价值: ⭐⭐⭐⭐⭐ 为大模型测试时计算扩展提供了兼具理论深度与极高工程可行性的训练新准则。