arXiv (ML)AI
拡散言語モデルの軌跡バランス調整:報酬追求を超えた事後学習手法
Beyond Mode-Seeking RL: Trajectory-Balance Post-Training for Diffusion Language Models
この記事についてAIに質問する →
日本語要約青い用語にマウスを合わせると解説が表示されます
拡散言語モデル(Diffusion Language Models)はオートレグレッシブモデルに代わる有望な選択肢として注目されていますが、これらのモデルに対する事後学習手法の多くは報酬最大化を目指す目的関数を採用してきました。しかし、こうしたアプローチには重大な課題が存在します。研究者たちが「軌跡ロッキング」と呼ぶ失敗モードで、報酬駆動型の更新によってモデルが限定的な復号化経路に確率質量を過度に集中させ、繰り返しサンプリング時に正しい代替解決策を見落とす傾向が強まってしまいます。
これに対処するため、研究チームはTraFL(Trajectory Flow baLancing)という軌跡バランス型の目的関数を提案しました。このアプローチはポリシーを報酬加重された目標分布に誘導するもので、凍結された参照モデルをアンカーとして使用します。拡散言語モデルに実装するため、彼らは拡散互換の文字列レベル代理関数と学習可能なプロンプト依存の正規化を開発しました。これにより、計算効率と実用性を両立させることができます。
数学的推論とコード生成のベンチマークにおける評価結果は印象的です。TraFLは全てのベンチマーク・長さの設定でベースモデルを上回る改善を達成した唯一の事後学習手法となりました。さらに重要なことに、サンプリング予算が増加してもその利益が持続するという特性を示しており、これはモデルの安定性と信頼性を示唆しています。
また、テストセットに対する転移可能性も確認されました。Minerva Mathベンチマークではベースモデルを継続して上回る性能を維持し、LiveCodeBenchの全ての難易度レベルにおいて最高性能を達成しています。こうした結果は、TraFLが単なる学術的な改善にとどまらず、実用的な価値を持つ手法であることを強く示唆しています。