截断差项，提升优项：基于排名分类的BoN式蒸馏

# 基于排名分类的截断坏项、提升好项蒸馏方法  
来源：https://arxiv.org/html/2608.19748  

## 截断坏项，提升好项：基于排名分类的BoN风格蒸馏  
Yaniv Romano\[0.5em\]  
电气计算机工程系\[0.2em\]  
计算机科学系\[0.3em\]  
以色列理工学院\[0.5em\]  
[email protected]  
[email protected]  

###### 摘要  
推理时选择方法（如Best-of-N）通过采样候选池并根据奖励模型选择最高排名的补全来改善生成效果。蒸馏旨在通过将原始奖励替换为池内排名，并学习一个提升高排名补全权重的策略，将该过程压缩为单一策略。然而，现有的基于排名的策略通常使用平滑的全支持重加权，因此低排名补全虽然获得较低的概率质量，但仍保留在目标支持集中。尽管更尖锐的重加权可以减少下尾概率质量，但它也增加了对单一奖励模型在顶部进行的脆弱排名的依赖。我们提出了**TUP（截断坏项，提升好项策略）**，该方法从支持集中移除低排名补全，并仅对保留的上尾部分进行可调尖锐度的重加权。TUP允许与提示无关的闭式归一化，并可通过二元交叉熵完全离线训练，使用偏移截断胜率作为软标签，蒸馏到参考模型的对数似然比作为logits。理论上，在特定假设下，我们证明对于任意未知的神谕奖励，最佳单调排名重加权可以通过下尾截断规则匹配，这为移除下尾而非仅降低其权重提供了形式化支持。实证上，我们展示了TUP与强大的离线对齐基线方法具有竞争力。  

## 1 引言  
推理时选择方法（如Best-of-N）已成为提升语言模型生成质量的有效工具\[15, 4\]。BoN从基础策略中采样多个补全，使用奖励模型评分，并返回得分最高的结果\[17, 1, 36\]。代价在于计算量，因为每个提示需要生成多个补全并进行奖励评估\[5\]。这催生了一系列BoN风格蒸馏工作\[34, 26, 3, 14, 8, 4\]，其中不是在部署时进行选择，而是训练单一策略来内化推理时选择器的行为。  

近期BoN风格蒸馏工作的一个关键洞见是，对于补全的相关训练信号并非其孤立的原始奖励，而是其在采样池内的相对排名。该相对排名通常被称为补全的*胜率*。因此，现有蒸馏方法旨在最大化胜率，同时惩罚与基础策略的KL距离\[26, 3, 14\]。这导致策略采用平滑的全支持重加权，偏爱更高排名的补全，同时保持每个补全在目标支持集中。更尖锐的重加权会降低明显低排名补全的概率质量，但也会将更多质量集中在奖励模型排名的顶部。结果是，过度锐化可能使策略偏离参考模型更远（更小的KL惩罚），增加奖励破解的风险\[28, 38, 19, 31\]。  

由于奖励模型是未知神谕或人类偏好的不完美代理\[10, 21, 15, 20, 29, 45, 11, 31\]，近期研究表明，粗粒度偏好往往在不同评估者间保持一致，而排名靠前的补全之间的细粒度区分则一致性较低\[23, 43, 24\]。因此，我们主张BoN风格蒸馏策略的设计应区分两个选择：（i）下尾部分应从支持集中移除多少；（ii）保留的高排名补全应被提升多少权重。  

为此，我们提出**TUP**，一种用于BoN风格蒸馏的截断坏项、提升好项策略。TUP通过偏移截断胜率变换将这两个选择解耦。胜率低于阈值的补全被赋予零概率质量，而高于阈值的补全则根据其相对排名进行软重加权。阈值参数控制下尾截断，而上尾尖锐度参数控制保留集合内的提升强度。通过为下尾截断和上尾尖锐度引入独立参数，该设计保留了BoN偏爱更高排名补全的原则，同时（i）防止概率质量被分配给明显不理想的补全，（ii）避免因顶部可能的排名错误而导致的过度锐化。由此产生的目标策略从离线数据训练仍然简单。每个提示与从参考策略采样并由奖励模型评分的补全池配对。从这些分数中，我们计算池内经验排名，并将其转换为偏移截断胜率标签；通过减去截断参数，我们为下尾补全分配精确的零标签。结合logit变换，我们使用二元交叉熵（BCE）损失训练蒸馏语言模型。  

##### 贡献。本文做出三项关键贡献：  
1. **用于BoN风格蒸馏的截断坏项、提升好项策略。** 我们提出一种基于排名的目标策略，使用胜率阈值移除低排名补全，并对保留的上尾部分进行软重加权。  
2. **下尾截断和尾内提升的理论支持。** 在特定假设下，我们以神谕胜率为性能标准，给出TUP的两个理论证明。首先，我们证明最佳硬下尾截断规则可以匹配按代理奖励胜率重加权补全的最佳单调策略。其次，固定截断阈值后，我们证明当更高的代理奖励胜率在保留尾部内倾向于对应更高的神谕胜率时，平滑提升保留补全的权重可以进一步提高神谕胜率。  
3. **基准结果和消融实验。** 我们在QRPO基准测试中评估TUP，使用UltraFeedback和Magpie Air训练Llama-8B Tülu 3 SFT，使用Magpie Air训练Mistral-7B-Instruct-v0.2，并与四种领先的离线对齐基线方法进行比较。我们使用三种不同的奖励模型评估性能，包括在留出数据和AlpacaEval\[9\]上的评估。TUP和实验的完整实现在：github.com/yarinbar/truncate-bad-upweight-good。  

## 2 前提知识与相关工作  
参见图1：TUP示意图。我们将语言模型表示为策略\(\pi(y \mid x)\)，将提示\(x \in \mathcal{X}\)映射到补全\(y \in \mathcal{Y}\)上的分布。提示从\(\mathcal{X}\)上的分布\(\mathcal{D}\)抽取。令\(\pi_{\text{ref}}(y \mid x)\)表示参考策略，例如通过监督微调获得，并令\(r(x,y) \in \mathbb{R}\)为奖励函数。标准的后训练对齐目标是寻找一个策略，在保持与参考策略接近的同时最大化预期奖励：  

\[
\pi_{r,\beta}^{\star}(\cdot \| x) = \arg\max_{\pi} \mathbb{E}_{y \sim \pi(\cdot \| x)} \bigl[ r(x,y) \bigr] - \beta D_{\text{KL}}\bigl( \pi(\cdot \| x) \,\|\, \pi_{\text{ref}}(\cdot \| x) \bigr), \tag{1}
\]  

其中\(\beta > 0\)控制奖励与KL的权衡。此KL正则化奖励最大化目标是RLHF中语言模型对齐的标准形式\[6, 37, 46, 30, 2\]，通常通过在线强化学习方法（如PPO\[33\]和GRPO\[35\]）进行优化。  

另一种范式（本文采用）直接建立在公式(1)的闭式解上，即Gibbs策略：  

\[
\pi_{r,\beta}^{\star}(y \| x) = \frac{1}{Z(x)} \pi_{\text{ref}}(y \| x) \exp\left( \frac{r(x,y)}{\beta} \right), \quad Z(x) = \mathbb{E}_{Y \sim \pi_{\text{ref}}(\cdot \| x)} \left[ \exp\left( \frac{r(x,Y)}{\beta} \right) \right], \tag{2}
\]  

其中\(Z(x)\)是配分函数。整理公式(2)得：  

\[
r(x,y) = \beta \log Z(x) + \beta \log \frac{\pi_{r,\beta}^{\star}(y \| x)}{\pi_{\text{ref}}(y \| x)}.
\]  

该关系暗示了一种直接的拟合策略。将未知最优策略\(\pi_{r,\beta}^{\star}\)替换为参数化模型\(\pi_{\theta}\)，可以通过回归公式(3)右侧来预测奖励，从而训练策略。难点在于配分函数\(Z(x)\)依赖于提示且难以计算。  

偏好优化方法以不同但相关的方式避免了计算\(Z(x)\)的需要。DPO\[32\]、IPO\[13\]和SimPO\[27\]从成对或二元偏好信号中学习，而REBEL\[12\]回归补全之间的相对奖励差异。尽管标准后训练对齐通常将原始奖励\(r(x,y)\)视为优化量，但这对于BoN等推理时选择算法并非最自然的选择。近期的推理感知和基于排名的方法（包括InfAlign\[3\]和QRPO\[26\]）提倡用相对排名量代替原始奖励值。一个自然的选择是群体*胜率*：  

\[
w_r(x,y) := \mathbb{P}_{Y' \sim \pi_{\text{ref}}(\cdot | x)} \bigl( r(x,Y') \leq r(x,y) \bigr), \tag{4}
\]  

它衡量\(y\)击败独立参考样本的频率。实践中，可以从有限补全池中估计胜率。这种胜率参数化不仅有用，因为它匹配了BoN风格选择的相对性质，还因为它消除了公式(3)中依赖于提示的配分函数。当\(Y \sim \pi_{\text{ref}}(\cdot \mid x)\)且条件奖励分布连续时，概率积分变换意味着\(w_r(x,Y) \sim U[0,1]\)。因此，对于任何非负重加权函数\(g: [0,1] \to [0,\infty)\)且积分有限，基于排名的策略  

\[
\pi_g(y \mid x) = \frac{1}{Z_g} g\bigl(w_r(x,y)\bigr) \pi_{\mathrm{ref}}(y \mid x), \qquad Z_g = \int_{0}^{1} g(u) \, du, \tag{5}
\]  

的归一化因子\(Z_g\)与提示\(x\)无关。许多基于排名的策略遵循此构造。例如，经典的Best-of-N诱导\(g_{\mathrm{BoN},N}(w) \propto w^{N-1}\)\[14, 4\]，标准QRPO形式诱导\(g_{\mathrm{QRPO},\beta}(w) \propto e^{w/\beta}\)\[26\]。两者都在整个支持集上分配正权重，留下胜率应如何转换为目标策略重加权的问题。  

这种全支持性质阐明了我们的目标与相关对齐方法的区别。迭代BoN蒸馏方法（如BOND\[34\]和Faster-WIND\[41\]）也旨在近似或蒸馏Best-of-N类型行为，但它们继承了平滑重加权观点，而非定义硬下尾移除规则。其他目标（包括RAFT\[8\]、SLiC-HF\[44\]和RRHF\[42\]）使用奖励过滤或相对偏好排名构建训练信号，但它们未指定提示无关的排名变换密度，该密度首先截断下尾，然后软重加权保留的上尾。我们的工作使截断和重加权操作显式化，产生有限的提示无关归一化和一个目标，该目标可以作为一个软分类问题拟合，使用变换后的胜率作为标签。  

## 3 方法  
图2：奖励模型一致性（左）和\(\lambda\)-截断在UltraFeedback参考补全上的成本-收益权衡（右）。我们将用作代理奖励模型的ArmoRM与独立的Skywork-v2-Llama奖励模型进行比较。左图：ArmoRM放置在顶部或底部\(p\)区域的补全中，Skywork也放置在同一区域的比例，作为\(p\)的函数。右图：\(\lambda\)-截断丢弃Skywork前25%补全或保留Skywork后25%补全的概率。交叉点标记这些错误概率相等的位置。  

### 3.1 提出的截断坏项、提升好项目标策略  
通过代理胜率的平滑单调函数对\(\pi_{\mathrm{ref}}\)进行重加权的类Gibbs策略（公式5），如QRPO\[26\]，面临我们上面讨论的尖锐度权衡。图2（左）提供了经验说明，激励我们选择将下尾截断与上尾提升解耦。