LLM強化学習のコストを80%削減するJackpot技法【論文】

Jackpot: 小さなモデルで大きなモデルを訓練する3つの核心

LLM強化学習において、ロールアウト生成が全体の費用の80%を占める^{[Jackpot論文]}。小さなモデルでロールアウトを代わりに作成すれば費用は減るが、2つのモデル間の分布の差(actor-policy mismatch)が訓練を不安定にする。

JackpotはOBRS(Optimal Budgeted Rejection Sampling)で解決した^{[Jackpot論文]}。小さなモデルが生成したトークンのうち、大きなモデルの分布に近いものだけを選んで訓練に使う。完璧な分布一致の代わりに、受容予算内の最適戦略を見つける。

Qwen3-1.7Bでロールアウトを生成し、Qwen3-8Bを訓練した結果、GSM8K 93.57%、MATH-500 82.65%を記録した^{[Jackpot論文]}。on-policy基準(93.29%、79.50%)と同等かそれ以上である。

既存のTISはMATH-500で76.45%にとどまり、後半部の不安定性も見られた。Jackpotは300ステップまで安定的な学習を維持した。

受容確率 a(x) = min(1, p_target / (lambda * p_inf))でトークンをフィルタリングする。top-k近似で演算量を減らし、既存の軌跡で動作するため追加のオーバーヘッドが少ない^[PPO論文]。

Q: Jackpotはどのような状況で有用か？

A: LLM強化学習でロールアウト費用を減らしたい場合に効果的である。訓練対象が大きく、小さなモデルをロールアウトに活用できる環境で有利である。モデルサイズの差が大きいほど、既存の方法に比べて安定性の利点が大きくなる。

Q: Actor-policy mismatchがなぜ問題なのか？

A: ロールアウトモデルと訓練モデルの分布が異なると、尤度比が希少トークンで急激に跳ね上がる。グラディエントが不安定になり、訓練が発散する可能性がある。非同期訓練よりもKLダイバージェンスが一桁以上大きい。

Q: 既存のimportance samplingと何が違うのか？

A: TISは尤度比を切り捨てて分散を減らすが、分布自体を矯正しない。OBRSはサンプルを選択的に受容または拒否し、ロールアウト分布自体を目標に近づける。この差が訓練安定性の格差として現れた。

この記事が役に立ったなら、AI Digesterを購読してください。