技術的概要：分布一致 LLM-RL 向けのオフライン重要度サンプリング DISA

1. 問題定義

現代の推論エージェントは、単一の正解ではなく、複数の有効な解決経路（多様な思考連鎖、計画、またはツール使用の痕跡など）を生成する能力で評価されるようになりつつある。プロキシマル方策最適化（PPO）やグループ相対方策最適化（GRPO）などの標準的な強化学習（RL）の目的関数は、報酬最大化に焦点を当てている。単一の正解の精度に対しては効果的であるが、これらの手法はモード崩壊に陥りやすいため、堅牢な評価（例：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$）のオンライン結合である。これらの手法は、結合損失（例：経路バランス）を用いて共有されたオンラインロールアウト上で両コンポーネントを共同訓練する。この結合は以下の 2 つの重大な問題を生じさせる：

1. 誤差伝播: 分配関数推定量 $Z_\phi$ の較正誤差が直接方策勾配に流れ込み、更新を歪める。
2. 診断機能の欠如: 分配関数に対する独立した品質シグナルが存在しない。その性能は方策の損失と絡み合っているため、失敗の原因が方策にあるのか、それとも分配関数の推定にあるのかを診断することが不可能である。

2. 手法：分離型重要度サンプリング・アンカリング（DISA）

著者は、分配関数の推定を方策最適化ループから分離する 3 段階のパイプラインであるDISAを提案する。鍵となる洞察は、分配関数 $Z(q)$ が学習可能な方策 $\pi_\theta$ には依存せず、参照方策 $\pi_{\text{ref}}$、報酬関数、および逆温度 $\beta$ のみに依存することである。したがって、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$$
  これにより、線形スケールでは不偏推定量が得られ、対数スケールでは制御された $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 段階アプローチ。
3. 理論的保証:
   • 完全な保存: オフライン分配関数が正確であれば、DISA 目的関数の大域的最適解は、理論的な報酬傾斜ターゲット分布と一致する。
   • バイアス付き定常性: プロンプトのみのバイアス（誤差が経路ではなく $q$ のみに依存する場合）が存在しても、ターゲット分布はオンポリシー勾配フローの定常点のままである。
   • 有界摂動: 残差推定誤差は、コーシー・シュワルツの不等式から導出される有界摂動エンベロープを通じてのみ方策目的関数に入り込み、安定性を保証する。
4. 実証的検証: 数学およびコードベンチマークにおける包括的な評価により、DISA がオンライン結合の落とし穴を回避しつつ、分布一致の恩恵を保持することを示している。

4. 実験結果

著者は、6 つの数学ベンチマーク（AIME 2024/2025、AMC 2023、MATH-500 などを含む）および 3 つのコードベンチマーク（LiveCodeBench、CodeForces、HumanEval+）において、2 つのオープンウェイトバックボーン（Qwen2.5-7BおよびQwen3-4B-Base）に対して DISA を評価した。

• 数学的推論:
  • 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-as-judge 評価（1-5 リッカート尺度を使用）により、DISA はバックボーンの元の多様性の最大割合（3.72/5.0）を保持しており、FlowRL（3.40）および GRPO/GSPO（約 3.0）と比較して著しく優れていることが示された。
• アブレーション研究:
  • 提案の強さ: 弱い提案モデル（235B 対 Qwen3-4B）を使用すると、オフラインラベルの分散が増加し、性能が低下したが、回帰モデルは依然として収束した。これは理論が予測するバイアス - 分散のトレードオフを確認するものである。
  • 逆温度（$\beta$）: 感度分析により、アンダーティルティング（$\beta=10$）は軽微なシグナル損失をもたらすのに対し、オーバーティルティング（$\beta=20$）は「分散の崖」とモード崩壊を引き起こすことが確認され、重要度サンプリング重みの理論的分析が妥当であることが検証された。
• 蒸留対 RL: DISA のオフライン推定に使用された同じ高報酬経路に対するファインチューニング（SFT）と比較した場合、DISA は最大13.8 Mean@8 ポイントの追加的な gains を提供し、ハードな SFT ターゲットによって破棄されたランドスケープ情報を、ソフトな報酬傾斜分布に適合させることで回復できることを示した。

5. 意義と主張

本論文は、DISA が推論向けポストトレーニング LLM のアーキテクチャにおいて重要な転換点であると主張する。分配関数推定を方策ループから分離することにより、DISA は以下の点を実現する：

• 較正誤差の排除: 分配関数推定誤差が方策勾配を歪めることを防止する。
• 独立した診断の実現: 回帰モデルの検証損失を通じて、方策訓練とは別に分配関数の品質を監視できる。
• 多様性の保持: 複数の有効な経路が存在する複雑な推論タスクにおいて不可欠な、解決空間のマルチモーダル構造を成功裏に維持する。
• 理論的妥当性: 推定誤差が有界であれば、オフライン推定プロセスが分布一致目的関数の理論的固定点を変更しないことを形式的に保証する。

著者は、DISA がオンライン結合型分布一致に対する堅牢で理論的に裏付けられた代替案を提供し、分配関数と方策の共同訓練に伴う不安定性を伴わずに、精度と多様性の両方で最先端の性能を達成すると結論付けている。