MilesがOPDを導入:蒸留をポストトレーニングのプリミティブへ

MilesがOPDを導入:蒸留をポストトレーニングのプリミティブへ

Milesチームは最近、On-Policy Distillation(OPD)をコア機能としてシステムに統合した。現在、OPDはMilesのrolloutおよびトレーニングフローに組み込まれており、ユーザーは教師モデルのシグナルのみで学生モデルを訓練することも、教師モデルの指導とGRPO/PPO-styleの強化学習目標を組み合わせることも可能だ。

今回のアップデートの主なポイントは、OPDをMilesの再利用可能なトレーニングプリミティブとして抽象化すること、疎なteacher-scoringフローによって不要な密なlogprob負荷を回避すること、そして単一の8×NVIDIA B200ノード上でQwen3.5-35B-A3Bの自己蒸留を検証することである。実験では、タスク固有のrewardが一切ない状態で、OPDが教師モデルの短い推論行動をベースとなる学生モデルへ転移させながら、学生モデルのタスク性能を維持できることが示された。

Miles and OPD workflow

OPDがMilesの第一級トレーニングプリミティブに

Milesは現在、OPDを追加のreverse-KLトレーニングシグナルとしてサポートしている。これにより、以下の2つの主要なトレーニングモードをカバーできる:

  • Pure distillation:学生モデルが完全に教師シグナルのみで訓練される。タスクrewardを0に設定し、モデルの更新はすべてトークンごとのstudent-teacher reverse-KLシグナルによって駆動される。
  • RL-augmented distillation:OPDシグナルをタスクrewardおよびGRPO/PPO-styleのadvantage estimatorと組み合わせ、学生モデルが環境フィードバックと教師の指導の両方から同時に学習する。

rollout側では、MilesはSGLangによってホストされた教師モデルを通じて、sampled-token OPDtop-k OPDの2種類のOPDをサポートしている。sampled-token OPDでは、教師モデルが各応答位置で学生モデルが実際にサンプリングしたトークンにスコアを付与する。サンプリングにはランダム性があるため、このトークンは学生モデルの確率が最も高いトークンとは限らない。

top-k OPDはさらに各位置での複数の候補トークンとそれに対応する学生モデルのlogprobsを保持する。教師モデルが同じ候補トークンにスコアを付与した後、Milesは整合された学生と教師のlogprobsを組み合わせてトークンごとのreverse-KLトレーニングシグナルを構築する。この設計により、OPDは一度限りの蒸留パスではなく、Miles内で繰り返し再利用可能なトレーニングコンポーネントとなる。

Top-k OPD:より豊かな師生分布のアライメント

単一のサンプリングトークンのみを見る推定方式と比較して、top-k OPDはより充実した学生—教師の比較を構築できる。学生のrollout過程において、Milesは各応答位置で設定可能な数の候補トークンIDと、対応する学生モデルのlogprobsを記録する。

位置tの候補トークンにスコアを付与する際、教師モデルは元のプロンプトと学生モデルがt以前に生成済みのトークンを因果的プレフィックスとして使用し、その位置での指定候補トークンのteacher logprobsを返す。

この実装は設定可能なtop-k蒸留戦略もサポートしている。ユーザーは候補セットの構築方法や、reverse-KLシグナル計算時の候補トークンへの重み付け方法をカスタマイズできる。これにより、異なる蒸留レシピが同一のrollout、teacher-scoring、トレーニングインフラを共有できる。

疎なStudent-to-Teacher Scoring:不要な負荷の回避

今回の実装における重要なシステム改善の一つは、疎なstudent-to-teacher scoringフローの導入だ。

OPDでは、学生のrolloutが位置ごとのtop-kテーブルを生成する。各生成位置について、MilesはKL計算に必要な候補トークンIDと学生logprobsを正確に把握している。そのため、教師モデルは実際には対応する因果的プレフィックスの下でこれらの候補トークンにスコアを付与するだけでよい。

初期のOPD実装では、scoringフローはまずすべての位置のtop-kトークンIDのグローバルな和集合を構築し、教師モデルに各応答位置でこの完全な和集合にスコアを付与させる必要があった。これは正確性を保証するが、密なR × |U|の中間負荷が生じる(Rは応答長、UはグローバルなトークンIDの和集合)。|U|R × Kに近づく可能性があるため、旧来のパスは最終的なOPD計算がO(RK)個の値のみを必要とする場合でも、O(R²K)スケールのJSONレスポンスを実体化して解析する可能性があった。

新しいフローは位置ごとの疎な候補トークンスコアリングに変更された。Milesは教師モデルに位置ごとのトークンIDテーブルを送信し、教師モデルは各位置の候補セットに必要なlogprobsのみを返す。これにより、教師のレスポンス負荷はOPDの実際の計算と完全に一致する:Milesが保持するのはOPDに必要な疎なR × Kの値のみとなる。

長い応答の推論タスクでは、この点が特に重要だ。teacher scoringは無意味な負荷の構築、転送、解析によって遅延させられるべきではない。疎なscoringはOPDパイプラインをより直接的にし、高スループットのrolloutとトレーニングに適したものにする。

検証実験:NVIDIA B200上でのQwen3.5自己蒸留

実装効果を検証するため、チームはQwen3.5-35B-A3Bを用いた制御された自己蒸留実験を実施した。

実験では、教師モデルをまずReinforcement Learning with Verifiable Rewards(RLVR)で訓練し、DAPOの数学問題を解く際に明らかに短い回答を生成するようにした。学生モデルは対応するpre-RL base checkpointとした。このセットアップを選択した理由は、DAPO能力を直接蒸留しても顕著な変化が生じないためだ:ベースとなるQwen学生モデルは元のタスクでもすでに高い性能を持っている。そのため、チームは意図的に教師モデルに測定可能な行動の差異——より短いが依然として有効な推論——を持たせ、OPDがその特性をまだ持っていない学生モデルへこの行動を転移できるかをテストした。

OPDの効果を分離するため、本実験ではpure distillationを採用した:

  • Task reward:0
  • Training signal:top-1 student-teacher reverse-KL
  • Teacher:RLVRで訓練された短推論checkpoint
  • Student:pre-RL base checkpoint
  • Hardware:単一の8×NVIDIA B200ノード
  • Execution mode:student rollout、teacher scoring、student trainingを同一GPUプール上で時分割実行

実験の参照ベースラインは以下の通り:

Model / stageHeld-out DAPOMean response length
RLVR teacher0.88706,248
Initial base student0.845718,675

教師モデルはより短い応答で同じDAPOの問題を解くのに対し、ベースとなる学生モデルは明らかに長い解答を生成することがわかる。この実験が答えようとする核心的な問いは:純粋なOPDは教師モデルの効率的な推論行動を学生モデルへ転移させながら、held-out DAPOの性能を維持できるか?である。

結果:推論が大幅に短縮され、性能は低下せず

実験結果は、OPDが教師モデルの効率的な推論行動を転移できることを明確に示している。

held-out DAPOのスコアは0.8457から0.8945へと向上した。同時に、sampled rolloutの長さは最初の更新後に急速に低下し、約18.6kトークンから、その後の多くのrolloutでは5.5k~6.7kトークンへと減少した(期間中、一度だけ短時間のランダムなスパイクが見られた)。トークンごとのOPD reverse-KLも約0.045から0.010へと低下し、学生自身のrolloutにおいて学生の分布が教師の分布に近づいていることを示している。

3本の学習曲線が異なる角度から同じ結論を示している。まず、held-out DAPO性能が向上しており、OPD後に学生モデルのタスク性能が低下していないことがわかる。

Held-out DAPO performance over evaluation steps

次に、sampled rollout長が急速に低下しており、学生モデルが教師モデルの短い回答行動を学習したことを示している。

Sampled rollout length over rollout steps

最後に、トークンごとのOPD reverse-KLが継続的に低下しており、学生の分布が徐々に教師の分布に収束していることが確認される。

Per-token OPD reverse-KL over rollout steps

本実験では元のタスクrewardが0であるため、行動の変化はすべてOPD教師シグナルによって駆動されている。核心的な結論は:純粋なOPDは教師モデルの効率的な推論行動をベースとなる学生モデルへ転移でき、held-out DAPOの性能は低下するどころか向上したということだ。

Milesにとっての意味

OPDの追加により、Milesは検証可能なタスクrewardだけでは表現しにくいより多くのポストトレーニングシナリオをカバーできるようになった。

実際の多くのワークフローでは、タスクの成功率だけを追求するのではなく、推論長、回答スタイル、ドメイン固有の慣習、参照モデルとの分布アライメントなど、モデルの行動を最適化する必要がある。OPDにより、Milesはこれらの目標を教師主導のトレーニングシグナルとして表現しながら、既存のrolloutとRLトレーニングパイプラインをそのまま維持できる。

OPDはまた、multi-teacher consolidationへの道も開く。異なる専門家教師が異なるドメインや能力で指導を提供し、単一の学生モデルが相補的な行動を吸収してより汎用的なモデルへと成長できる。MilesがOPDを再利用可能なプリミティブとして扱っているため、同一の学生トレーニングパイプラインが異なる教師や蒸留レシピをサポートでき、各ドメインごとに独立したトレーニングシステムを構築する必要がない。複数教師によるOPDの検証は引き続き今後の課題だ。

さらに、OPDの実装は良好な組み合わせ可能性を維持している。ユーザーはpure distillationを実行することも、OPDを補助シグナルとしてGRPO/PPO-styleの目標と組み合わせることもできる。これは、教師の指導とタスクrewardが協調して機能する必要のあるポストトレーニングフローにとって非常に重要だ。

同様に重要なのは、このシステムが長い応答のワークロードに最適化されている点だ。疎な位置ごとのscoringとtop-k OPDにより、不要な密なteacher-scoring負荷を回避し、パイプライン全体を蒸留が実際に必要とする計算により近いものにしている。

今後の展開

現在の実装により、OPDはMilesの第一級機能として確立され、制御されたQwen3.5自己蒸留設定でpure OPDの効果が検証された。次のステップとして、チームはより複雑な蒸留レシピ、RL目標との組み合わせ、そして複数教師のシナリオにおけるOPDの汎化能力の検証など、検証範囲の拡大を継続的に計画している。