Distributionally Robust Token Optimization in RLHF¶
会议: NeurIPS 2026
arXiv: 2604.08577
OpenReview: https://openreview.net/forum?id=9zSXz3TUrv
代码: https://osf.io/
领域: LLM对齐
关键词: RLHF, 强化学习, 分布鲁棒优化, 令牌级优化, 跨分布泛化
一句话总结¶
针对多步推理任务中大语言模型对浅层输入扰动脆弱的问题，本文将分布鲁棒优化（DRO）引入令牌级 RLHF 框架，提出在片段划分上构建 KL 与 \(\chi^2\) 散度不确定性集合的 DRTO 方法，在无需额外计算开销的前提下大幅提升多步数学与代码推理的分布外泛化能力。
研究背景与动机¶
大语言模型通过人类反馈强化学习（RLHF）在各类复杂生成与逻辑推理任务中取得了显著进展，其中基于近端策略优化（PPO）与直接偏好优化（DPO）的后训练管线已成为主流范式。然而，这类对齐模型在部署时对提示词表述、语法格式或符号表示的微小扰动表现出显著的脆弱性。特别是在多步数学和代码推理等长链条任务中，表面的提示词变化极易引发灾难性的逻辑断裂与错误累积。由于对齐偏好数据与实际测试环境之间必然存在分布偏移，传统的序列级对齐算法往往表现出严重的训练不稳定性和脆弱的跨分布泛化表现。
解决此问题的核心矛盾在于信用分配粒度与分布鲁棒性之间的权衡。传统的分布鲁棒优化尝试在离线 DPO 或轨迹级 PPO 上施加对抗重加权，但整个序列的标量奖励平摊到所有令牌上，无法准确定位长思维链中真正导致推理偏航的关键少数步骤；若直接在单一令牌上做对抗扰动，又会面临巨大的方差与离散噪声干扰。最近提出的令牌级对齐框架 RTO（Reinforced Token Optimization）虽然利用 DPO 隐式策略构造出密集的令牌级优势信号，但在复杂小批量样本中依然极易受到局部极端损失与难样本片段的冲击，导致策略优化轨迹震荡。
本文的切入角度是：将每个生成序列沿低置信度分界点自适应划分为连续的语义/逻辑片段（spans），在片段平均损失构成的名义分布周围建立散度不确定性集合，推导闭式对偶上界来实现低方差、高聚焦的鲁棒重加权。核心 idea：将分布鲁棒优化引入令牌级 RLHF，基于采样策略低置信度切分点构建响应片段并定义 KL 与 \(\chi^2\) 散度不确定性集合，通过对抗重加权自适应放大高损失关键推理片段的梯度贡献，实现零额外推断开销的高鲁棒性对齐。
方法详解¶
整体框架¶
DRTO 的整体优化流程基于在线策略采样（on-policy rollout）的小批量训练展开。在每个训练外循环中，旧策略模型 \(\pi_{\mathrm{old}}\) 针对当前批量提示词采样生成回答序列。系统首先利用离线训练得到的 DPO 隐式参考策略计算每个令牌位置的密集形状奖励，并经由广义优势估计（GAE）生成令牌级截断代理损失。接着，算法依据旧策略的生成置信度概率在关键转折点处将回答切割为变长片段，并在整个微批次的片段损失集合上求解 DRO 对抗极大化问题，最终以对抗加权后的代理损失执行演员网络（actor）的梯度更新。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入提示词微批次 & 采样回答"] --> B["低置信度概率切点划分片段"]
B --> C["令牌级密集奖励与片段代理损失聚合"]
C --> D["散度不确定性集合对抗加权 (KL / χ² 闭式求解)"]
D --> E["稳健演员网络策略更新 (PPO-Update)"]关键设计¶
1. 低置信度切点的自适应片段划分：定位长链推理的逻辑分支点 为了避免全序列轨迹平均掩盖局部错误，同时防止单令牌级的方差发散，DRTO 提出了基于模型自身生成置信度的自适应片段切割机制。对于回答长度为 \(H_i\) 的第 \(i\) 个样本，系统记录采样时的模型概率 \(p^{\mathrm{old}}_{i,h}\)。当某位置满足 \(p^{\mathrm{old}}_{i,h} < p_{\mathrm{cut}}\) 时，即被认定为推理逻辑转换点或语法分界线，在此处进行片段截断，同时设置最小片段长度约束（如 16）与单序列最大片段数上限（如 8）。对于得到的连续片段 \(\mathcal{G}_{i,s}\)，其片段损失由该区间内令牌级 PPO 代理损失取算术平均得到： $\(L_{i,s}(\theta) := \frac{1}{|\mathcal{G}_{i,s}|} \sum_{h \in \mathcal{G}_{i,s}} \ell_{i,h}(\theta)\)$ 其中 \(\ell_{i,h}(\theta) := -\min(\rho_{i,h}(\theta)\widehat{A}_{i,h}, \mathrm{clip}(\rho_{i,h}(\theta), 1-\epsilon, 1+\epsilon)\widehat{A}_{i,h})\)。此设计使得重加权作用于承载完整局部逻辑推演的语义块上，准确圈定错误高发区间。
2. 响应均衡名义分布与 KL 散度对抗重加权：平滑指数倾斜聚焦难样本 在构建的批次片段集合 \(\mathcal{S}_{\mathcal{B}} = \{(i,s): i \in [n], s \in [S_i]\}\) 上，若直接平铺片段会导致长回答对目标过度支配。DRTO 定义了响应均衡的名义分布 \(\mathbb{P}_{\mathcal{S}}(i,s) := \frac{1}{n S_i}\)，赋予每个回答样本相同的 \(1/n\) 总体权重，再在内部均分给各个片段。在 KL 散度球 \(\Omega^{\mathrm{KL}}_\rho(\mathbb{P}_{\mathcal{S}}) = \{\mathbb{Q} : D_{\mathrm{KL}}(\mathbb{Q} \parallel \mathbb{P}_{\mathcal{S}}) \le \rho\}\) 约束下，内层极大化问题可通过 Gibbs 变分原理转化为严格的一维标量凸优化： $\(R_{\mathrm{KL}}(\theta; \rho) = \inf_{\eta > 0} \left\{ \frac{1}{\eta} \log \mathbb{E}_{(i,s) \sim \mathbb{P}_{\mathcal{S}}} \left[ \exp(\eta L_{i,s}(\theta)) \right] + \frac{\rho}{\eta} \right\}\)$ 最佳对抗分布展现出自然的指数倾斜（exponential tilt）形式 \(\mathbb{Q}^\star(i,s) \propto \mathbb{P}_{\mathcal{S}}(i,s) \exp(\eta^\star L_{i,s}(\theta))\)。通过数值一维线搜索求得最优温度 \(\eta^\star\) 后，对抗权重以阻断梯度（detach）形式直接赋权微批次片段损失，使模型对高损失困难片段产生平滑但坚决的自适应关注。
3. Pearson χ² 散度解析松弛：方差敏感的自稳边界代理 作为 KL 形式之外更轻量的代数替代方案，DRTO 在 Pearson \(\chi^2\) 散度不确定性球 \(\Omega^{\chi^2}_\rho(\mathbb{P}_{\mathcal{S}})\) 下推导了解析上界。记片段损失的名义均值为 \(\bar{L}(\theta) = \mathbb{E}_{\mathbb{P}_{\mathcal{S}}}[L_{i,s}(\theta)]\)，标准差为 \(\sigma(\theta) = \sqrt{\mathrm{Var}_{\mathbb{P}_{\mathcal{S}}}(L_{i,s}(\theta))}\)，由柯西-施瓦茨不等式可严格证明最坏情况期望损失满足上界 \(R_{\chi^2}(\theta; \rho) \le \bar{L}(\theta) + \sigma(\theta) \sqrt{\rho}\)。为保证数值稳定性与全域可微性，\(\chi^2\)-DRTO 最终构建的稳健替代目标函数为： $\(L_{\chi^2}(\theta) := \bar{L}(\theta) + \sqrt{\rho} \sqrt{\mathrm{Var}_{\mathbb{P}_{\mathcal{S}}}(L_{i,s}(\theta)) + \delta}\)$ 其中 \(\delta > 0\) 为微小的数值平滑系数。当非负性条件满足时该上界严格取等。这一形式直观表明，\(\chi^2\)-DRTO 在最小化名义平均损失的同时显式惩罚片段间损失方差，以零超额计算直接抑制不同推理片段上的优化不平衡。
损失函数 / 训练策略¶
DRTO 完整继承 RTO 的三部分密集奖励机制，每个令牌位置的奖励计算为： $\(r^{\mathrm{RTO}}_h(x, y_{1:h}) = \beta_1 \log \frac{\pi_{\mathrm{dpo}}(y_h \mid x, y_{1:h-1})}{\pi_{\mathrm{ref}}(y_h \mid x, y_{1:h-1})} - \beta_2 \log \frac{\pi_\theta(y_h \mid x, y_{1:h-1})}{\pi_{\mathrm{ref}}(y_h \mid x, y_{1:h-1})} + \beta_3 \mathbf{1}_{\{h = H_i\}} r(x, y_{1:H_i})\)$ 在此基础上，微批次中汇总的微观片段损失通过公式 (8)-(11) 计算出 KL 加权损失 \(L_{\mathrm{KL}}(\theta)\)，或通过公式 (16) 计算出方差正则损失 \(L_{\chi^2}(\theta)\)，并结合辅助 PTX 语言建模预训练保留损失更新参数。实验中 KL 散度半径 \(\rho=0.01\)，\(\chi^2\) 散度半径 \(\rho=0.1\)，\(p_{\mathrm{cut}}=0.05\)。
实验关键数据¶
主实验¶
在数学推理（Llama3-8B 主干，OpenMathInstruct 微调）与代码推理（Qwen3-8B 主干，UltraInteract 微调）两大主流基准上的跨分布泛化评估结果如下。数学评测包含提示词思维链扰动 GSM(CoT)、西语跨语言漂移 GSM(ES)、高难竞赛题 MATH-500 以及跨学科难度跃升 GPQA；代码评测考察贪婪 pass@1。
| 模型 | GSM8K (近分布) | GSM (CoT) (格式漂移) | GSM (ES) (语言漂移) | MATH-500 (高难OOD) | GPQA (学科跨越) |
|---|---|---|---|---|---|
| SFT | 61.4 (-11.8) | 64.8 (-9.6) | 63.6 (-2.4) | 17.4 (-0.6) | 24.8 (-1.5) |
| PPO | 73.2 (+0.0) | 75.2 (+0.8) | 65.6 (-0.4) | 17.0 (-1.0) | 25.8 (-0.5) |
| DPO | 73.9 (+0.7) | 76.4 (+2.0) | 66.4 (+0.4) | 20.1 (+2.1) | 20.7 (-5.6) |
| FlowRL | 75.9 (+2.7) | 76.2 (+1.8) | 66.0 (+0.0) | 19.2 (+1.2) | 28.8 (+2.5) |
| DAPO | 78.0 (+4.8) | 76.8 (+2.4) | 64.0 (-2.0) | 19.4 (+1.4) | 25.3 (-1.0) |
| RTO (基线) | 73.2 (基准) | 74.4 (基准) | 66.0 (基准) | 18.0 (基准) | 26.3 (基准) |
| \(\chi^2\)-DRTO | 78.7 (+5.5) | 79.5 (+5.1) | 66.8 (+0.8) | 22.4 (+4.4) | 27.3 (+1.0) |
| KL-DRTO | 77.3 (+4.1) | 77.4 (+3.0) | 67.6 (+1.6) | 18.8 (+0.8) | 26.3 (+0.0) |
在代码生成任务中，\(\chi^2\)-DRTO 在 HumanEval 达到 88.41%（相比 RTO 的 87.80% 提升 +0.61%），MBPP 达到 69.00%（+8.40%），LiveCodeBench 达到 39.96%（+2.12%）；KL-DRTO 在 LiveCodeBench 上斩获 40.54%（+2.70%），MBPP 达到 70.20%（+9.60%），全面缓解了传统 DPO 存在的严重离线过度优化衰退问题（DPO 在 LiveCodeBench 仅为 24.13%）。
消融实验¶
为了严格解耦令牌级密集奖励（Token-wise Reward）与片段级分布鲁棒（Span-level DRO）各自的独立贡献，论文进行了细致的结构消融对比（均基于数学推理基准）：
| 配置 | GSM8K (Acc) | MATH-500 (Acc) | GPQA (Acc) | 说明 |
|---|---|---|---|---|
| \(\chi^2\)-DRTO (完整) | 78.7 | 22.4 | 27.3 | 片段级 \(\chi^2\)-DRO + 令牌级密集奖励 |
| \(\chi^2\)-DRTO (无片段，整轨DRO) | 77.7 | 18.6 | 25.8 | 退化为整条回复单一权值，高难数学暴跌 3.8% |
| \(\chi^2\)-DRPPO (无令牌奖励) | 77.4 | 14.2 | 26.8 | 移除 RTO 密集奖励，高难推理暴跌 8.2% |
| KL-DRTO (完整) | 77.3 | 18.8 | 26.3 | 片段级 KL-DRO + 令牌级密集奖励 |
| KL-DRTO (无片段，整轨DRO) | 76.8 | 18.6 | 26.3 | 缺少局部细粒度聚焦，常规计算集轻微下滑 |
| KL-DRPPO (无令牌奖励) | 77.0 | 15.8 | 25.8 | 稀疏轨迹级奖励下难样本梯度严重稀释 |
在端到端耗时对比中，RTO 耗时 10h00m（1.00×），\(\chi^2\)-DRTO 耗时 11h36m（1.16×），KL-DRTO 耗时 11h41m（1.17×），而 DAPO 与 FlowRL 分别高达 3.59× 与 4.86×。
关键发现¶
- 片段划分不可或缺：整轨轨迹级 DRO 由于将正确推演与错误分支混杂平均，无法精准向逻辑崩塌点施加梯度压力；在 MATH-500 这类高深度多步推理上，去掉片段划分导致准确率由 22.4% 骤降至 18.6%。
- 密集奖励与鲁棒加权相辅相成：将 DRO 简单移至原生 PPO（DRPPO）时收益微弱且在难题上剧烈退化（MATH-500 仅 14.2%），表明密集令牌奖励为策略探索提供了必要的微观坐标系，而片段 DRO 则确保了宏观更新方向免受离群噪声干扰。
- 训练动力学高度稳定：监控训练全程发现，DRTO 变体的 Critic 损失更低且极其平滑，达成峰值回报 80% 所需步数仅为 PPO/RTO 的三分之一左右，同时 PTX 语言建模损失未受扰动，有效抵御了对齐漂移。
亮点与洞察¶
- 基于模型自身先验的无监督语义切分：利用采样时的低生成概率自动定位逻辑转折点，避免了依赖外部语法解析器或预设符号规则带来的分词偏差与计算开销。
- DRO 理论向强化学习内部组件的成功下放：以往 DRO 大多仅在离线样本重采或全轨迹宏观加权中发挥作用，DRTO 首次将其优雅嵌入到小批量内部的片段级梯度聚合中，证明了局部不确定性控制的巨大价值。
- 可低成本复用的损失方差正则项：\(\chi^2\)-DRTO 推导出的“均值加标准差”代理形式实现极简，计算量几乎为零，非常易于无缝插入其他基于 PPO 或 GRPO 的现代推理后训练管线。
局限与展望¶
- 作者承认的局限：目前评估集中在 8B 规模的模型（Llama3-8B 和 Qwen3-8B）以及单轮对话推理场景，尚未在更大参数规模（如 70B+）或复杂多轮交互轨迹中进行全量验证；同时该方法对对抗性攻击或系统性越狱缺乏理论安全上界。
- 方法与假设边界：概率截断阈值 \(p_{\mathrm{cut}}\) 及最大/最小片段长度为固定超参，对于不同语言或不同温度采样下的置信度标定可能存在一定分布敏感性。
- 改进方向：可进一步探索结合动态置信度百分位切分机制，并将片段级 DRO 拓展至带有过程奖励模型（PRM）和多智能体交互的自博弈强化学习场景中。
相关工作与启发¶
- vs RTO (Zhong et al., 2025)：RTO 开创了融合 DPO 隐式策略与 PPO 的令牌级密集奖励范式，但其标量平均损失对批次内难片段敏感；本文在此基础上引入片段级 DRO 极大化重加权，弥补了其分布外泛化短板。
- vs DrDPO (Wu et al., 2025) & \(\chi^2\)-Preference Optimization (Huang et al., 2025)：上述工作均在离线成对偏好数据阶段应用 DRO 缓解标签噪声与过度优化；DRTO 则首次在在线策略（on-policy）的令牌强化学习阶段引入片段级鲁棒优化。
- vs DAPO (Yu et al., 2025) & FlowRL (Zhu et al., 2025)：DAPO 和 FlowRL 虽也关注推理强化学习，但依赖复杂的重采样、动态剪枝或流匹配，训练耗时达基线的 3.5-5 倍；DRTO 以仅 1.16 倍的极低开销实现了相当甚至更优的跨分布鲁棒性。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次将片段级 DRO 与令牌级 RLHF 深度融合，理论推导严密且对偶形式优雅。
- 实验充分度: ⭐⭐⭐⭐⭐ 覆盖数学与代码两大关键推理域，多维度验证自然分布偏移与多项深度消融。
- 写作质量: ⭐⭐⭐⭐⭐ 动机阐述自然流畅，数学推导清晰完备，实验图表详实自洽。
- 价值: ⭐⭐⭐⭐⭐ 为大模型长链逻辑推理的稳定对齐提供了通用、低成本且理论支撑扎实的关键优化组件。