# ReSpec: Towards Optimizing Speculative Decoding in Reinforcement Learning Systems > [!info] Talk metadata > - 会議: [[MLSys2026]] Day 3 (May 20 / Wed)、Research Track Oral "LLM Training 2" セッション(15:15–15:30 PDT、Grand Ballroom 2) > - 登壇者: Qiaoling Chen(Nanyang Technological University) > - 共著者: Zijun Liu, Peng Sun, Shenggui Li, Guoteng Wang, Ziming Liu, Yonggang Wen, Siyuan Feng, Tianwei Zhang > - 所属: Nanyang Technological University, Singapore / Shanghai Qiji Zhifeng Co., Ltd. / Tsinghua University / National University of Singapore / Shanghai Innovation Institute > - スライド: https://mlsys.org/media/mlsys-2026/Slides/3836.pdf > - 連絡先: [email protected] > [!abstract] 概要(論文アブストラクトの忠実な日本語訳) > 強化学習(RL)による大規模言語モデル(LLM)の適応は、学習時間の 75% 以上を占める生成段階がボトルネックとなる。投機的復号(SD)はサービングシステムにおいて自己回帰生成を高速化するが、RL 学習下での挙動はほぼ未探索である。我々は、SD を RL システムに素朴に統合する際に障壁となる **3 つの重大なギャップ**を特定した。大バッチサイズでの高速化の逓減、アクターモデルの継続的更新に伴うドラフトモデルの陳腐化、そしてドラフトモデルに起因する方策最適化の劣化である。これらのギャップに対処するため、我々は SD を RL に適応させるシステム **ReSpec** を提案する。ReSpec は 3 つの相補的メカニズムを通じて対処する。SD 構成の動的チューニング、知識蒸留によるドラフトモデルの進化、およびロールアウト報酬による更新の重み付けである。Qwen モデル(3B--14B)での実験により、ReSpec は報酬の収束と学習安定性を維持しつつ最大 4.5 倍の高速化を達成し、RL ベース LLM 適応のための実用的な解決策を提供する。 ## 背景: RL 学習における生成段階のボトルネック - LLM の RL 学習は **生成(ロールアウト)・推論(報酬評価)・学習(パラメータ更新)** の 3 段階で構成される - 7B モデル・最大応答長 8K トークンの場合、**数学タスクで 83--86%、コードタスクで 70.9--75.5%** の壁時計時間が生成段階に費やされる(Table 1) - 推論段階は 1.7--14%、学習段階は 12.3--15.1% に過ぎず、生成段階の最適化が全体の律速を決定する - GRPO や DAPO のようなグループベースサンプリング手法は、プロンプトあたりの復号トークン数を倍増させ、生成段階のコストをさらに増大させる ## 投機的復号(SD)の RL 適用における 3 つのギャップ ### GAP 1: 大バッチサイズでの高速化の逓減 - RL 学習では GPU 利用率を高めるため大バッチサイズで復号を行うが、バッチサイズが大きいと GPU は既にほぼ飽和しており、SD が提供する追加並列性の効果は限定的である - H100 上の Qwen2.5-7B-Instruct / MTBench で測定すると、バッチサイズ 2 では最良 SD 構成で **1.46 倍** の高速化を得るが、同じ構成をバッチサイズ 32 に適用すると **0.76 倍**(減速)に転じる(Figure 3, 7) - 最適な SD ハイパーパラメータ $(s, t, n)$(投機ラウンド数・分岐ファクター・ラウンドあたりドラフト長)は **実効バッチサイズに強く依存** し、静的設定では一貫した高速化を得られない ### GAP 2: 学習中のドラフトモデルの陳腐化 - RL 学習ではアクターモデルが継続的に更新されるため、固定されたドラフトモデルは急速にアクターと乖離する - Qwen2.5-7B-Instruct / 数学データセットで EAGLE-3 ドラフトモデルの受理長を 100 RL ステップにわたり計測すると、**受理長は 4 付近から 2 付近まで低下** する(Figure 4) - サービングにおけるオンライン SD 手法は入力分布の変動に対処するが、RL 学習では **ターゲットモデル自体** が進化するという本質的に異なる課題がある ### GAP 3: ドラフトモデルに起因する方策最適化の劣化 - ロスレス SD は理論的には棄却サンプリングにより期待値でターゲット分布を保存するが、**実務的にはオンポリシー RL 学習で報酬の測定可能な低下** を引き起こす - 原因は 3 つ: (R1) GPU カーネル実装の差異による **非決定的検証パス**、(R2) ドラフトモデルの陳腐化による **探索の制限**、(R3) 非定常方策更新下での **受理確率の分散の指数的増大** - Qwen2.5-7B で素朴に EAGLE-3 を適用すると、約 100 更新ステップ後に報酬が目に見えて低下する(Figure 5) ## ReSpec のシステム設計 ReSpec は 2 つの主要コンポーネントで構成される。 ### Adaptive Server(GAP 1 への対処) 実効バッチサイズに応じて SD 構成を動的に切り替えるランタイムコントローラである。 - **Solver**: オフラインプロファイリングにより、各バッチサイズにおける各 SD 構成 $(s, t, n)$ のスループット高速化を予測する性能モデルを構築する。RL 実行時には実効バッチサイズを観測して最適構成を選択する - **Scheduler**: 投機実行モードと非投機実行モードの間を **ゼロオーバーヘッド** で切り替える。バッチサイズが大きいときは非投機(ターゲット高速復号)を選択し、シーケンスが完了してバッチが縮小すると投機モードに遷移する。遷移時には前回の非投機状態を KV キャッシュの拡張に利用し、復号カーネルの変更を不要とする ### Online Learner(GAP 2・GAP 3 への対処) ドラフトモデルをオンラインで継続的に更新するコンポーネントであり、3 つのモジュールから成る。 #### 報酬重み付き知識蒸留(Reward-Weighted KD) - 従来の KD は全ロールアウトを等しく扱うが、RL では低品質・低報酬の軌道が多く含まれるため、ドラフトモデルを低品質な振る舞いへ引き寄せてしまう - ReSpec は各サンプルの報酬 $r$ に基づく重み $w(r)$ で KL ダイバージェンスを重み付けする: $\mathcal{L}_{\mathrm{KD}}(x, y) = w(r) \sum_{t=1}^{T} \mathrm{KL}\bigl(\tilde{p}(\cdot \mid x, y_{<t}) \,\|\, q_\theta(\cdot \mid x, y_{<t})\bigr)$ - 高報酬の軌道はアクターの更新後分布をより正確に反映するため、これを優先的に学習することで **正のフィードバックループ**(低品質ドラフト → 低品質ロールアウト → さらに低品質ドラフト)を断ち切る - Qwen2.5-7B での比較: 報酬重みなし KD は約 150 ステップでゼロ近傍に崩壊、EAGLE-3 固定は約 175 ステップで劣化するのに対し、報酬重み付き KD は 200 ステップを通じて安定的に報酬が上昇する(Figure 10) #### 非同期更新オーバーラップ(Async Update Overlap) - 同期的なドラフトモデル更新は生成パイプラインにバブルを生む - ReSpec はリプレイバッファに蒸留ターゲットを蓄積し、$I$ イテレーションごとにまとめて更新することで最適化コストを償却する - ドラフトモデルの学習を RL パイプラインのアイドルスロットに重畳し、ターゲットロールアウトと並列実行することで、追加の壁時計遅延をほぼゼロに抑える(Figure 11) - Qwen2.5-3B・7B では Async-1(毎イテレーション非同期更新)が最高の報酬(0.60、0.42)を達成し、間隔を広げると劣化する。14B では同期更新でも高報酬(0.60)を達成し、大モデルほど陳腐化への耐性が高い(Figure 15) ## 実験評価 ### 実験設定 - **モデル**: Qwen-2.5 の 3B・7B・14B パラメータ - **データ**: 実世界の数学データセット - **アルゴリズム**: RL には GRPO、SD には EAGLE-3 - **エンジン**: VeRL(学習)+ SGLang(推論)。ReSpec 実装は約 2K LOC(Adaptive Server 500 LOC、Online Learner 1500 LOC) - **インフラ**: 2 ノード x 8 基 NVIDIA H100(80GB)、ノード内 NVLink 900 GB/s、ノード間 8x 400 Gbps RoCE ### 学習安定性(RQ1) - ReSpec は 3 スケール全てで非加速ベースラインと同等のバリデーションスコアを維持する - Qwen2.5-3B: 400 ステップ後もベースラインと整合(EAGLE-3 はステップ 400 後に 0.15 まで低下) - Qwen2.5-7B: ステップ 160--180 で 0.4 を維持(EAGLE-3 は 0.06--0.2 に早期崩壊) - Qwen2.5-14B: ベースラインの傾向を一貫して追跡(EAGLE-3 は長期学習で発散) ### エンドツーエンド高速化(RQ2) | モデル | 最大高速化 | 平均高速化 | |---|---|---| | Qwen2.5-3B | **4.53 倍** | **1.84 倍** | | Qwen2.5-7B | **2.41 倍** | **1.69 倍** | | Qwen2.5-14B | **2.60 倍** | **1.50 倍** | - 最大高速化は生成初期のバッチサイズが小さい段階(トークンレベル並列性が最大限活用可能)で観測される - 大モデルほど安定した高速化を示すが、ピーク高速化は小モデルほど顕著である ### コンポーネント別寄与(RQ2・Qwen-14B での分解) - ベースライン(SD なし): 1.00 倍 - + Reward-Weighted KD(Online Learner): **1.48 倍**(SD の安定化とドラフト受理長の向上) - + Adaptive Server: **1.66 倍**(ワークロードに応じた最効率 SD 構成の動的選択で追加 12%) - + Async Update Overlap: **1.78 倍**(ドラフト更新遅延の隠蔽でさらに加速) ## 重要な知見 - 投機的復号は RL 学習の生成段階を高速化する有望な技術であるが、素朴な適用は **学習安定性の破壊** と **方策劣化** を招くため、RL 固有の動的特性に適応したシステム設計が不可欠である - ReSpec の 3 メカニズム(適応的 SD 構成・報酬重み付き KD・非同期オーバーラップ)はそれぞれ直交的に効果を発揮し、合計で最大 4.5 倍の高速化を報酬収束を維持したまま達成する