技术摘要：DISA - 用于分布匹配 LLM-RL 的离线重要性采样

1. 问题陈述

现代推理智能体越来越被评估为生成多个有效解决方案路径（例如，多样化的思维链、计划或工具使用轨迹）的能力，而不仅仅是生成单个正确答案。标准的强化学习（RL）目标，如近端策略优化（PPO）或组相对策略优化（GRPO），侧重于奖励最大化。虽然这些方法在单答案准确性方面行之有效，但它们往往遭受模式崩溃的困扰，即策略将概率质量集中在单一的高奖励轨迹上，从而消除了鲁棒评估（例如 pass@k 指标）所需的多样化有效解决方案。

分布匹配 RL提供了一种替代方案，通过训练策略以拟合目标分布，其中轨迹概率与指数化奖励成正比：
$$\tilde{\pi}(o | q) = \frac{1}{Z(q)} \pi_{\text{ref}}(o | q) \exp(\beta \tilde{r}(q, o))$$
其中 $Z(q)$ 是配分函数（一个依赖于提示词 $q$ 的归一化常数）。

现有分布匹配方法（例如 FlowRL、GFlowNets）的核心瓶颈在于策略（$\pi_\theta$）与配分函数估计器（$Z_\phi$）之间的在线耦合。这些方法使用联合损失（例如轨迹平衡）在共享的在线滚动数据上共同训练这两个组件。这种耦合造成了两个关键问题：

1. 误差传播：配分函数估计器 $Z_\phi$ 中的校准误差直接流入策略梯度，扭曲了更新过程。
2. 缺乏诊断：配分函数没有独立的质量信号；其性能与策略的损失纠缠在一起，使得无法诊断失败是源于策略还是配分函数估计。

2. 方法论：解耦重要性采样锚定（DISA）

作者提出了DISA，这是一个三阶段流水线，将配分函数的估计与策略优化循环解耦。关键见解在于，配分函数 $Z(q)$ 仅依赖于参考策略 $\pi_{\text{ref}}$、奖励函数和逆温度 $\beta$，而不依赖于可训练的策略 $\pi_\theta$。因此，它可以在 RL 训练开始之前进行离线估计。

阶段 1：离线重要性采样（IS）估计

对于一组离线提示词 $\mathcal{D}_Z$，该方法从一个强大的提议模型 $p_T(\cdot|q)$（例如，更大的教师模型）中抽取 $N$ 条提议轨迹。

• 重要性权重：对于每条轨迹 $o_i$，计算重要性权重：$w(q, o_i) = \frac{\pi_{\text{ref}}(o_i|q)}{p_T(o_i|q)} \exp(\beta \tilde{r}(q, o_i))$。
• 对数空间聚合：为了确保数值稳定性并减少偏差，提示词 $q$ 的配分函数估计值在对数空间中使用 logsumexp 聚合器计算：
  $$\log \hat{Z}_{\text{IS}}(q) = \text{logsumexp}_i \left( \log \pi_{\text{ref}}(o_i|q) - \log p_T(o_i|q) + \beta \tilde{r}(q, o_i) \right) - \log N$$
  这在线性尺度上提供了 $Z(q)$ 的无偏估计，并在对数尺度上具有受控的 $O(1/N)$ 偏差。

阶段 2：通过回归进行摊销

每个提示词的估计值 $\log \hat{Z}_{\text{IS}}(q)$ 作为标签，用于训练一个轻量级的提示词条件回归器 $g_\psi(q)$（例如，一个小型 MLP）。

• 回归器通过最小二乘回归进行训练，以最小化其预测值与离线 IS 标签之间的平方误差。
• 一旦训练完成，回归器参数 $\psi^*$ 将被冻结。这创建了一个可重用的、离线的配分函数估计，能够泛化到初始估计集之外的提示词。

阶段 3：锚定在线策略学习

在在线 RL 阶段，冻结的回归器 $g_{\psi^*}(q)$ 在轨迹平衡（TB）损失中取代了可学习的配分函数 $Z_\phi(q)$。

• 目标：优化策略 $\pi_\theta$ 以最小化残差：
  $$\mathcal{L}_{\text{DISA}}(\theta) = \mathbb{E} \left[ \sum_i \omega_i \left( g_{\psi^*}(q) + \log \pi_\theta(o_i|q) - \log \pi_{\text{ref}}(o_i|q) - \beta \tilde{r}(q, o_i) \right)^2 \right]$$
• 解耦：关键在于，梯度被停止流向 $g_{\psi^*}$。在策略更新期间，配分函数充当一个固定的、每个提示词特定的常数（即“锚点”）。

3. 主要贡献

1. 识别耦合为核心障碍：本文指出，配分函数估计与策略学习的纠缠是分布匹配 RL 中校准误差和诊断盲视的主要来源。
2. DISA 流水线：一种新颖的三阶段方法，通过重要性采样将配分函数估计移至离线，通过回归进行摊销，并将其冻结以用于在线 RL。
3. 理论保证：
   • 精确保持：如果离线配分函数是精确的，则 DISA 目标的全局最小值与理论上的奖励倾斜目标分布相匹配。
   • 带偏差的平稳性：即使存在仅提示词偏差（即误差仅依赖于 $q$ 而不依赖于轨迹），目标分布仍然是在线策略梯度流的平稳点。
   • 有界扰动：残差估计误差仅通过有界扰动包（通过柯西 - 施瓦茨不等式推导）进入策略目标，从而确保稳定性。
4. 实证验证：在数学和代码基准测试中进行了全面评估，证明 DISA 在保留分布匹配优势的同时，避免了在线耦合的陷阱。

4. 实验结果

作者在两个开源权重骨干模型（Qwen2.5-7B 和 Qwen3-4B-Base）上评估了 DISA，涵盖了六个数学基准（包括 AIME 2024/2025、AMC 2023、MATH-500）和三个代码基准（LiveCodeBench、CodeForces、HumanEval+）。

• 数学推理：
  • DISA 匹配或略微超过了在线耦合基线 FlowRL（例如，在 Qwen3-4B-Base 上，Mean@8 为 65.5 对比 64.8）。
  • 它显著优于奖励最大化基线 GRPO 和 GSPO（例如，在 Qwen3-4B-Base 上比 GRPO 高出 16.1 分）。
• 代码生成：
  • 在 pass@16（一个对解决方案多样性敏感的指标）上，DISA 在所有方法中取得了最高分数（例如，在 Qwen3-4B-Base 上为 51.5），优于 FlowRL、GSPO 和 GRPO。
  • 在 pass@1 上，DISA 保持了竞争力，与奖励最大化方法相比没有显著的性能下降。
• 多样性分析：
  • 基于 LLM 作为裁判的解决方案策略多样性评估（使用 1-5 李克特量表）显示，DISA 保留了骨干模型原始多样性的最大比例（3.72/5.0），优于 FlowRL（3.40），并显著优于 GRPO/GSPO（约 3.0）。
• 消融研究：
  • 提议强度：使用较弱的提议模型（Qwen3-4B 对比 235B）增加了离线标签的方差，导致性能下降，但回归器仍然收敛，证实了理论预测的偏差 - 方差权衡。
  • 逆温度（$\beta$）：敏感性研究证实，欠倾斜（$\beta=10$）会导致轻微的信号损失，而过倾斜（$\beta=20$）则会导致“方差悬崖”和模式崩溃，验证了重要性采样权重的理论分析。
• 蒸馏与 RL 对比：与在 DISA 离线估计所使用的相同高奖励轨迹上进行微调（SFT）相比，DISA 提供了高达 13.8 个 Mean@8 分 的额外增益，表明拟合软性的、奖励倾斜的分布能够恢复硬 SFT 目标所丢弃的景观信息。

5. 意义与主张

本文声称，DISA 代表了推理后训练 LLM 架构的重大转变。通过解耦配分函数估计与策略循环，DISA：

• 消除校准误差：防止配分函数估计误差扭曲策略梯度。
• 启用独立诊断：配分函数的质量可以通过回归器的验证损失与策略训练分开监控。
• 保持多样性：它成功维持了解决方案空间的多模态结构，这对于存在多个有效路径的复杂推理任务至关重要。
• 理论严谨性：该方法提供了形式化保证，即只要估计误差是有界的，离线估计过程就不会改变分布匹配目标的理论不动点。

作者得出结论，DISA 提供了一种稳健的、有理论依据的替代方案，用于在线耦合的分布匹配，在准确性和多样性方面均实现了最先进的性能，同时避免了共同训练配分函数和策略所带来的不稳定性。