拡散言語モデル（Diffusion Language Model, DLM）は、自己回帰モデルに代わる有望な手法として注目されています。しかし、DLMの後訓練フェーズにおいて、既存の方法は大部分が報酬最大化目的を採用しており、これが「軌道ロッキング」（trajectory locking）と呼ばれる重大な障害を引き起こします。具体的には、報酬駆動の更新によって確率質量が狭い範囲のノイズ除去経路に過度に集中し、繰り返しサンプリングを行った際に代替の正解をカバーする能力が低下します。この問題は、DLMの実際の応用における性能を大きく制限しています。
この問題に対処するため、Saba Ahmadiらは新しい後訓練手法「TraFL」（Trajectory Flow baLancing）を提案しました。TraFLの核心は、軌道バランス目的を使用してポリシーを報酬で傾斜したターゲット分布に向けて訓練し、その分布を凍結された参照モデルにアンカーすることです。これにより、報酬最大化のみを追求する際に生じるモード崩壊を効果的に回避します。さらに、この手法を拡散言語モデルに実用的にするため、拡散互換のシーケンスレベルのサロゲートと、学習可能なプロンプト依存の正規化項を導入しています。
数学推論およびコード生成のベンチマークにおいて、TraFLは評価されたすべての長さ設定でベースモデルを上回った唯一の後訓練手法であり、その改善はサンプリング予算が増加するにつれて持続します。特に注目すべきは、その改善が留出評価にも転移する点です：Minerva Mathではベースモデルを常に上回り、LiveCodeBenchのすべての難易度区分で最強の手法となりました。
本研究は、単純な報酬最大化ではなく軌道分布のバランスを取ることで、拡散言語モデルの探索能力と最終性能を向上させる新しい後訓練の方向性を示しています。今後、この手法はさらに広範な応用シナリオに拡張され、言語モデルの訓練と最適化に新たな視点を提供することが期待されます。