2026-09-13·★ 会員·論文·14分で読めます

論文解説: Online Draft Co-Training for Speculative Decoding in Large-Scale, Long-Context RL Post-Training — 122B・256Kトークンで下書きモデルを方針と一緒に学習させる

RL学習の時間を食うロールアウト生成を投機的デコーディングで速くし、下書きモデルを方針と一緒に学習し続ける。NVIDIAがコンテキスト並列とパイプライン並列の2つの壁をどう越えたかを、論文本文に沿って1から解説。

要点

本論文は、大規模・長文脈のRLポストトレーニングで下書きモデルを方針と一緒にオンライン学習する手法を扱う。
RLポストトレーニングの実時間はロールアウト生成に支配されることが多く、投機的デコーディングがこれを速くする。
投機的デコーディングは棄却サンプリングで検証するため、出力の分布は校閲者ひとりで書いた場合と同じである。
RLでは方針が毎ステップ変わるので、固定の下書きは整合を失い受理長が落ちる。
壁は2つあり、標準CPが分岐注意を扱えない点と、方針の中間特徴がPPの複数ステージにまたがる点である。

#speculative-decoding #reinforcement-learning #post-training #context-parallel #pipeline-parallel #ring-attention #nemo-rl #inference

対象textタスクinference

Online Draft Co-Training for Speculative Decoding in Large-Scale, Long-Context RL Post-Training

一次資料 — この記事の根拠

論文の発表 2026-09-07→この解説の公開 2026-09-13同月

Online Draft Co-Training for Speculative Decoding in Large-ScaleZili Wang, Zhaopeng Qiu, Yuekai Zhang ほか · 2026-09-07 · v1"arXiv:2609.07108論文ページ·PDF

https://arxiv.org/abs/2609.07108"Long-Context RL Post-Training

原文の要旨（Abstract）を読む

Speculative decoding accelerates rollout generation, which dominates the cost of reinforcement learning (RL) post-training. Online co-training can further increase the draft's accuracy, yielding greater speedups. However, scaling this approach to co-training on large models with long contexts poses two obstacles: (1) branch attention is unsupported by standard causal context-parallel (CP) implementations, and (2) target features span across pipeline-parallel (PP) stages. We address both with an end-to-end system for large-scale online draft co-training. For CP, we extend packed, load-balanced zigzag ring attention by merging rank-local branch attention with causal main-sequence attention. For PP, TapChannel transports intermediate target features across stages via a separate path, leaving the pipeline schedule unaffected. Experiments demonstrate that co-trained drafts closely track the policy baseline while delivering substantial rollout and end-to-end speedups across model scales up to 122B. Our CP design achieves strong scaling at 256K tokens with significant memory savings over prior work, and our PP transport incurs modest overhead. Code can be found at https://github.com/NVIDIA-NeMo/RL/issues/3698.

この論文は何を解決したのか

原題は "Online Draft Co-Training for Speculative Decoding in Large-Scale, Long-Context RL Post-Training"（arXiv:2609.07108、NVIDIA、Technical Report）。著者は Zili Wang、Zhaopeng Qiu、Yuekai Zhang、Shuang Yu、Junjie Lai で、2026年9月7日公開です。

要旨を日本語でまとめます。強化学習（RL）によるポストトレーニングでは、ロールアウト生成（方針モデルに答えを書かせる工程）がコストの大半を占める。投機的デコーディングはこれを速くし、さらに下書き（draft）モデルを方針と一緒にオンラインで学習し続ければ、下書きの的中率が上がってもっと速くなる。ところが大きなモデルと長いコンテキストでは2つの壁がある。(1) 下書き学習に必要な「枝分かれした注意（branch attention）」を標準の因果コンテキスト並列（CP）が扱えない。(2) 下書きが必要とする方針の中間特徴が、パイプライン並列（PP）の複数ステージにまたがる。論文はこれを、パック済み負荷分散ジグザグ・リング注意の拡張と、パイプラインのスケジュールを変えずに特徴を横流しする TapChannel で解決する。実験では、共学習した下書きが方針の学習曲線をほぼ保ちつつ122Bまでのモデルで学習全体を速くし、CPは256Kトークンでもスケールし、PP側の追加負担は控えめだった、というものです。

比喩: 速記者と校閲者

投機的デコーディング（SD）を絵にします。大きな言語モデルが1トークンずつ書くのは、重い校閲者が一語ずつ辞書を引くようなものです。1語ごとにモデル全体を動かすので、GPUの計算力を使い切れません。

そこで軽い速記者（下書きモデル）を横に置きます。速記者が次の数語をさっと書き、校閲者はそれを一度にまとめて検証し、最初に外れた語から自分で書き直す。検証は棄却サンプリングで行うので、出力の分布は校閲者ひとりで書いた場合と同じです（§1）。品質を落とさず、速記者が当たる限り速くなる。基礎は投機的デコーディングを1から解説で扱っています。

RLに持ち込む理由は、論文いわく「RLポストトレーニングの実時間はロールアウト生成に支配されることが多い」（§1）からです。GRPOのような手法は方針に答えを何本も生成させ、その報酬で方針を更新します。この生成が一番遅い。

なぜ「一緒に学習」するのか

厄介なのは、RLでは校閲者の書き方が毎ステップ変わることです。固定の速記者はだんだん癖から外れ、当たる語数（受理長、accepted length）が落ちます。論文は「固定の下書きは進化する方針との整合を失う」と書いています（§2）。

そこで速記者も一緒に学習し続ける、これが online co-training です。目的関数は §3.1 の式(1)。

\mathcal{L}(\theta,\phi) = \mathcal{L}_{\text{RL}}(\theta) + \lambda\,\mathcal{L}_{\text{draft}}\bigl(\phi;\,x,\,\operatorname{sg}(H_{\theta}(x))\bigr)

(1)

読み方は「方針 $\theta$ のRL損失に、下書き $\phi$ の損失を $\lambda$ 倍して足す」。 $x$ はロールアウトのトークン列、 $H_\theta(x)$ は方針の中間層から取った特徴、 $\operatorname{sg}$ は stop-gradient で「この特徴には勾配を流さない」印です。下書きは方針の中身を覗いて学ぶが、方針は下書きの都合で引きずられない。下書きは方針の最後のパイプラインステージに子モジュールとして置かれます（§3.1）。

2つの壁

ここまでは先行研究（FastGRPO、ReSpec、MTP-RL など、§2）にもありました。この論文の主題は「大規模・長文脈でやると何が壊れるか」です。

壁1: 分岐注意。EAGLE-3・DFlash・DSpark といった高性能な下書きは、学習時に枝分かれした注意を使います。EAGLE-3 は自分の予測を入力に戻して数ステップ先まで練習する Training-Time Test（TTT）で各位置から枝が伸び、DFlash は1回の順伝播でブロック単位に生成し、DSpark はそこに軽い Markov head を足します（§3.2）。共通の要求は、各枝が「本流の因果的な前置き」と「自分の枝の中のKV」の両方を見ること。長いコンテキストをGPU間で分割する標準CPは、素直な因果注意しか想定していません。

壁2: 特徴がステージをまたぐ。下書きは方針の中間層の特徴を入力にしますが、PPでは層が複数ステージに分かれ、下書きは最後のステージにしかいません。標準のパイプライン通信は隣同士しかつなぎません（§3.3）。

並列化の用語が曖昧なら分散学習を1からを先に。論文の立場は「先行研究は下書きのために並列レイアウトを変えるが、我々は方針のトポロジーを変えず、周りを合わせる」（§2）です。

論文解説: Online Draft Co-Training for Speculative Decoding in Large-Scale, Long-Context RL Post-Training — 122B・256Kトークンで下書きモデルを方針と一緒に学習させる

要点

この論文は何を解決したのか

比喩: 速記者と校閲者

なぜ「一緒に学習」するのか

2つの壁

ここから先は会員限定です

参考文献

このあとに出た解説

コメント