技術的サマリー：f-軌道バランス

問題定義

大規模言語モデル（LLM）や生成フローネットワーク（GFlowNets）を含む生成モデルのアライメントと微調整は、一般的に、モデルの分布と目標分布との間の不一致を最小化することに依存している。強化学習（RL）による微調整においては、これは通常、KL ダイバージェンス（Kullback-Leibler divergence）ペナルティ subject に期待報酬を最大化する形式で定式化される。

この分野における重大な課題は、**モード探索（mode-seeking）とモードカバリング（mode-covering）**の行動間のトレードオフである。反転 KL ダイバージェンス（通常、軌道バランスまたは二乗 KL 損失を通じて実装される）などの標準的な目的関数は、モード探索的傾向を示し、モデルを少数の高確率モードに収束させる傾向がある。これは、創薬や多様な候補生成など、探索を必要とするタスクには不適切であり、これらのタスクでは「モードカバリング」（目標分布の多様性全体を捉えること）が望まれる。

さらに、既存の手法はしばしばオフポリシー妥当性に苦慮している。REINFORCE や PPO などの推定量は一般的であるが、高い分散に悩まされる。最近の二乗 KL 損失（$KL_{sq}$）を用いたアプローチは、オンポリシーにおいて低分散の勾配を提供し、オフポリシーでも妥当性を有するが、これらは本質的に反転 KL ダイバージェンスの性質に縛られており、探索と利用のトレードオフを制御する能力を制限している。

手法

著者は、二乗 KL 損失をf-ダイバージェンス全体に一般化する代理損失関数のファミリーであるf-軌道バランスを提案する。

理論的基盤

核心的な理論的貢献は、対数確率における並進不変な損失関数とf-ダイバージェンスとの間の 1 対 1 対応の確立である。

1. ダイバージェンスから損失へ: 任意の凸関数 $f$ によって定義される f-ダイバージェンス $D_f(p_\theta \| p^*)$ に対して、著者は対数確率の差 $\Delta_\theta(y) = \log p_\theta(y) - \log p^*(y)$ 上で定義された損失関数 $L_f$ を導出した：
   $$L_f(\Delta_\theta(y)) = \int_0^{\Delta_\theta(y)} (f'(e^t) - f'(1)) dt$$
   この損失には 2 つの重要な性質がある：

   • オンポリシー勾配等価性: サンプルがモデル方策（$\mu = p_\theta$）から抽出されるとき、$L_f$ の期待自動微分勾配は、対応する f-ダイバージェンス $D_f$ の勾配と一致する。
   • オフポリシー妥当性: 行動分布 $\mu$ からサンプリングされたオフポリシーデータ上で評価されるとき、目標のサポートが行動方策のサポートに含まれている限り、損失は同じ大域的最小解（$p_\theta = p^*$）を持つ有効な目的関数のままとなる。

2. 損失からダイバージェンスへ: 逆に、対数確率上の任意の狭義凸かつ並進不変な損失 $\ell$ は、特定の f-ダイバージェンスの最小化に対応する。本論文は、これらの写像が互いに逆であることを証明している。

実用的実装

トレーニングにおける実用的な課題に対処するため、著者はいくつかの拡張を導入している：

• DevGrad 損失（バッチ正規化）: 非正規化された目標や GFlowNets などで一般的に問題となる分配関数の扱い不能性に対処するため、著者はVarGrad推定量を一般化する。バッチ分散の代わりに、対数確率の差のバッチごとの一般化偏差を使用する。これには、バッチ上で損失を最小化することによって正規化定数 $\log Z$ を推定し、スコア関数の係数を効果的に中心化して分散を低減することが含まれる。
• テンパード損失: 小さな温度パラメータ $\beta$（ここで目標分布は $p^* \propto \pi_{ref} \exp(r/\beta)$）を扱う際の数値的オーバーフローを防ぐため、著者はテンパード損失を提案する。これは損失を $1/\beta$ でスケーリングし、テンパード分布に適用することで、勾配が有界かつ安定であることを保証する。
• GFlowNets 向けの f-軌道バランス: この枠組みは、$L_f$ 損失を用いて軌道バランス制約を強制することで GFlowNets に適用される。これにより、標準的な軌道バランス（反転 KL に対応）が任意の f-ダイバージェンスに一般化される。

主要な貢献

1. 一般化された損失ファミリー: オンポリシー勾配が対応する f-ダイバージェンスの勾配と一致し、かつオフポリシー妥当性を維持する、並進不変な代理損失（$L_f$）の一般ファミリーの導出。
2. 理論的等価性: 並進不変な対数確率損失と f-ダイバージェンスの逆関係の証明。これにより、モードカバリングなどの所望の性質を持つ新しい損失関数の体系的設計が可能となる。
3. 非正規化目標への拡張: DevGrad損失を介して VarGrad 手法を任意の f-ダイバージェンスに一般化し、RL 微調整における扱い不能な分配関数の問題に対処。
4. 探索の制御: f-ダイバージェンスの選択（$\alpha$-ダイバージェンスにおけるパラメータ $\alpha$ によってパラメータ化される）が、モード探索対モードカバリングの行動を直接制御することを示実証。低い $\alpha$ 値（例：フォワード KL、ヘルンガー）はモードカバリングを促進し、高い値（例：ピアソン $\chi^2$）はモード探索を促進する。

実験結果

著者は、4 つの異なるタスクにおいて損失ファミリーを評価した：

1. 合成グリッド（GFlowNets）: モードカバリングをテストするように設計された 2 次元グリッドタスクにおいて、低い $\alpha$（フォワード KL、ヘルンガー）を持つ損失は、標準的な軌道バランス（反転 KL）よりも有意に速く、かつ信頼性高く目標分布の 4 つのモードすべてを発見した。逆に、ピアソン $\chi^2$（$\alpha=2$）は強いモード探索行動を示し、最初に発見されたモードに留まる傾向があった。
2. 分子発見（SynFlowNet）: 合成的にアクセス可能な分子を生成するための SynFlowNet 枠組みに適用。$\alpha$ の調整により生成される分子の多様性を制御できることがわかった。アニーリングスケジュール（探索のために低い $\alpha$ から始め、利用のために $\alpha=1$ まで増加させる）が最善のトレードオフをもたらし、標準的な軌道バランスと比較して、より高い報酬を持つ分子をより高い多様性で生成した。
3. 拡散モデル: MNIST 数字の条件付きサンプリング（奇数対偶数）において、標準的な軌道バランスは目標モード全体で均等にサンプリングできず（特定の数字 0 や 6 の過剰サンプリング）、f-軌道バランス損失、特に低い $\alpha$ を持つものは、目標モード全体でより均一なサンプリングを達成した。
4. 非同期 LLM 微調整: これらの損失は、数学問題（GSM8k、MATH）に対する LLM（Qwen および OLMo ファミリー）の微調整に、50 ステップの非同期遅延を伴って適用された。f-ダイバージェンス損失は安定したトレーニングと一貫したエントロピー - 報酬のトレードオフを示したが、標準的な PPO 実装は同じ非同期条件下で不安定さを示した。

意義と主張

本論文は、二乗 KL 損失の有効性が孤立した現象ではなく、より広範な損失ファミリーの特定の事例であることを主張している。「オンポリシー勾配一致」と「オフポリシー妥当性」という性質を f-ダイバージェンス全体に拡張することにより、著者は生成モデルの微調整のための統一された枠組みを提供する。

主要な意義は、基盤となるトレーニングアルゴリズムを変更することなく、単に異なる代理損失関数を選択する（または $\alpha$ パラメータを調整する）ことで、探索と利用のトレードオフを明示的に制御できる点にある。これにより、実践者は報酬最大化のためのモード探索と、多様性および探索のためのモードカバリングの行動の間を選択でき、かつ現代の大規模非同期トレーニングに必要な安定性とオフポリシー互換性を維持できる。この枠組みは、高報酬候補の空間を網羅することが不可欠な創薬や探索的エージェントなどの応用において特に価値があると示唆されている。