システム統合型Speculative DecodingによるRL後学習ロールアウトの高速化
Accelerating RL Post-Training Rollouts via System-Integrated Speculative Decoding
要約
大規模言語モデル（LLM）のRL後学習（RL post-training）において、自己回帰的なロールアウト生成がボトルネックとなっている。既存の効率化手法はオフポリシー実行やリプレイ、低精度生成などでスループット改善を図るが、出力分布を変えてしまう場合がある。本研究ではSpeculative Decoding（投機的デコーディング）をロスレスな加速プリミティブとして活用し、ターゲットモデルの出力分布を保持しながらRLロールアウトを高速化する手法を提案する。vLLMバックエンドを持つNeMo-RLに実装し、同期・非同期パイプラインの両方に対応。事前学習済みMTPヘッドや小規模ドラフトモデル、Eagle3などの投機機構を幅広くサポートする。8Bスケールの同期RL環境下で推論後学習ワークロードにおいてロールアウトスループットが1.8倍に向上し、高忠実度シミュレータによる試算では235Bスケールで非同期RLと組み合わせると最大2.5倍のエンドツーエンド学習高速化が見込めると報告している。
筆者コメント
Speculative DecodingはDeepMindやGoogleが推論高速化として確立した技術だが、RL学習ループ内に組み込む試みは比較的新しい。RLトレーニングでは分布の一致（on-policy性）が重要なため、出力分布を変えないロスレス加速という特性は理論的に非常に相性が良く、この点が本研究の核心的な価値と考えられる。一方、Eagle3のような手法はRL後に適用されることを前提に設計されているため、RL学習中に使用した場合のドラフトモデルの品質維持や、ポリシー更新に伴うアクセプタンス率の変動については実運用上の懸念が残る。235Bスケールでの2.5x高速化はシミュレータベースの推計であり、実機での検証結果ではない点に注意が必要だ。NeMo-RLという特定のフレームワーク前提であるため、他のトレーニングスタック（例：Verl、OpenRLHF）への移植性や日本語モデルの学習への適用可能性は不明瞭と見られる。計算資源に制約のある組織にとっては8Bスケールで1.8xという実測値の方が現実的な参照点となるだろう。RL後学習の効率化はモデル開発コスト削減に直結するため、実務上の重要度は高いと考えられる。
※ このコメントは本サイト独自のものです。論文・記事の公式見解ではありません。