大規模・長文脈 RL 事後学習における Speculative Decoding のオンライン共同学習手法

論文概要

大規模モデルにおける強化学習（RL）事後学習のコストの大部分を占めるロールアウト生成に対し、Speculative Decoding とオンライン共同トレーニングを組み合わせることで高速化を図るエンドツーエンドのシステム手法が提案された。最大 122B パラメータまでのモデルスケールにおいて、ポリシーのベースライン性能を維持しながら、大幅なロールアウトおよびエンドツーエンドの高速化を実証している。

関連研究

Speculative Decoding の高速化にはドラフトモデルの精度向上が不可欠であり、そのためのオンライン共同トレーニングが検討されてきた。しかし、従来の手法を大規模モデルや長文脈環境へスケールさせる際には、標準的なコンテキスト並列化およびパイプライン並列化の実装上における制約という課題があった。

新規性と貢献

本研究の最大の貢献は、大規模・長文脈環境下でのオンライン共同トレーニングにおける2つのボトルネックをシステムレベルで解決した点にある。コンテキスト並列化（CP）における分岐アテンションの未サポート問題と、パイプライン並列化（PP）におけるステージ間の中間特徴量転送問題を同時に克服した。

提案手法の詳細

提案システムでは、CP に対してはランクローカルな分岐アテンションと因果メインシーケンスアテンションを統合する拡張型 zigzag ring attention を採用し、256K トークンでの強力なスケーリングと大幅なメモリ削減を達成している。また、PP に対しては、パイプラインスケジュールに影響を与えない独立したパス「TapChannel」を導入し、中間ターゲット特徴量をステージ間で効率的に転送する。

評価・考察

評価実験では、最大 122B パラメータのモデルスケールにおいて、共同トレーニングされたドラフトがポリシーベースラインに極めて近い追従性を示しつつ、ロールアウト生成およびエンドツーエンドで高い高速化を果たすことが確認された。CP 設計により 256K トークンにおける強スケーリングを実現し、従来手法と比較して顕著なメモリ削減効果が示されている。

応用例と今後の展望

本手法は、数百億パラメータ規模の大規模言語モデルにおける強化学習事後学習のコスト削減に直接寄与する。日本のAIインフラ・R&D部門において、数万トークン以上の長文脈処理や大規模分散学習を運用する際の基盤システム設計やスケーリング戦略において、大きな実装上の指針となる。

結論

本研究は、大規模・長文脈 RL 事後学習における Speculative Decoding のオンライン共同トレーニングのシステム的課題を、CP および PP の拡張により解決した。モデル性能の劣化を抑えつつ実用的な高速化とメモリ効率の向上を両立させている。

注釈

• Speculative Decoding: 小規模なドラフトモデルで高速にトークンを生成し、大規模モデルで一括検証することで生成を加速する手法。
• コンテキスト並列化 (CP): 長いコンテキスト（入力長）を複数のGPUに分散して処理する並列化手法。
• パイプライン並列化 (PP): モデルのレイヤーを複数のGPUステージに分割して逐次処理する並列化手法。