💻 computer science

Accelerating Visual On-Policy Distillation with Batched Speculative Jacobi Rollouts

本論文は、並列マルチトークン処理と適応型実行モードを可能にすることで、蒸留フレームワークを変更することなく生成品質を維持しながら学習時間を大幅に短縮する、視覚的オンポリシー蒸留を加速させるバッチ化された投機的Jacobiデコーディング（Speculative Jacobi Decoding）バックエンドであるHB-SJDを紹介するものである。

原著者： Bingqi Shan, Zhehao Yu, Kenhong Lin, Baoquan Zhang

公開日 2026-08-20

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Bingqi Shan, Zhehao Yu, Kenhong Lin, Baoquan Zhang

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

人工知能の世界では、驚くほど美しい画像を生成できるほど賢いモデルを作ることと、それらを高速に動作させるためにサイズを小さく保つこととの間で、絶え間ない綱引きが行われています。多くの場合「教師」と呼ばれる大規模なシステムは、一つの出発点から画像全体を構築するように、一度に極めて小さな断片を一つずつ予測していくことで、信じられないほど詳細な画像を作り出すことができます。しかし、このステップ・バイ・ステップのプロセスは遅く、コストがかかります。これらのシステムを高速化するために、研究者たちは「知識蒸留（knowledge distillation）」という手法を用います。これは、より小さく単純な「生徒」モデルが、教師の模倣を学習する手法です。このプロセスのより新しく効果的なバージョンである「オンポリシー蒸留（on-policy distillation）」では、生徒に対して、リアルタイムで自分自身の訓練用サンプルを生成しながら練習することを求めます。これにより、生徒は単に古い例を暗記するのではなく、実際に芸術作品を創造する際に自分がどのように振る舞うかを正確に学習できるのです。問題は、この練習セッションが非常に遅いことです。生徒は依然として、画像を一つ一つの小さな断片から構築しなければならず、これが学習プロセス全体を遅らせるボトルネックとなっています。

深センの研究チームは、生徒の学習方法の根本的なルールを変えることなく、このボトルネックを打破する方法を見出しました。彼らは「HB-SJD」と呼ばれる、生徒の練習セッションのためのより高速なエンジンとして機能する新しい手法を導入しました。この新システムは、生徒に画像のトークン（あるいは微細な視覚的単位）を一つずつ構築させる代わりに、複数の将来の断片を同時に提案し、それらの推測が正しいかどうかをチェックすることを可能にします。これは、学生がテストを受けている様子を想像してみてください。一つの答えを書き、先生が採点するのを待ち、それから次の答えを書くのではなく、学生が一ページ分の答えをまとめて書き、先生がどれが正解かを素早く判定するようなものです。生徒は正解を保持し、間違ったものだけを書き直すことで、より速く前へと進むことができます。

研究者たちは、単にこの高速な手法を画像のグループに対して同時に実行するだけでは不十分であることを発見しました。標準的なセットアップでは、もし一つの画像の生成が早く終わったとしても、システムが次のステップに進む前に、グループ内の最も遅い画像が終了するのを待たなければなりません。この待機時間は貴重な計算能力を無駄にしていました。これを解決するために、チームは各画像が独自のスピードで進行できるようなシステムを設計しました。もし画像が完了したら、その画像は脇に退き、システムは停止することなく残りの画像に対して作業を続けます。さらに、彼らは、稼働している画像の数に応じて、システムの最適な実行方法が変わることも発見しました。多くの画像が処理されているとき、システムはフルグループをアクティブに保つことで最もよく機能します。しかし、画像が完了してグループが小さくなるにつれて、システムは空のスロットを管理するオーバーヘッドを避けるために、残りのアクティブな画像のみを処理する、より軽量なモードへと切り替わります。

画像生成モデルを用いたテストにおいて、この新しいアプローチは非常に効果的であることが証明されました。研究者たちは、生徒が練習用画像を生成する時間をほぼ半分に短縮でき、従来のメソッドと比較して1.4倍から1.6倍の高速化を実現したことを発見しました。決定的なのは、このスピードアップが品質を犠牲にすることなく達成された点です。この新手法を用いて訓練された生徒モデルが生成する画像は、より遅い旧来の手法で訓練されたものと同様に、鮮明でリアルでした。このシステムは、画像が短くても長くても、処理される画像のグループが小さくても大きくても、また学習プロセスのどの段階においても、一貫して優れた性能を発揮しました。練習用データを生成するエンジンのみを置き換えることで、研究者たちは学習プロセス全体を大幅に高速化することに成功し、最終的な結果の品質を損なうことなく、人工知能の訓練を加速できることを証明しました。

技術要約：バッチ化された投機的ヤコビ・ロールアウトによる視覚的オンポリシー蒸留の加速化

1. 問題提起

視覚的自己回帰（AR）モデルは、画像生成における強力なスケーラビリティを示しています。知識蒸留は、大規模な教師モデルから小規模な生徒モデルへと能力を転移するための標準的な手法です。最近の進展であるオンポリシー蒸留（OPD）は、学習中に生徒自身に軌跡（trajectory）を生成させ、それを教師によって監督させることで、学習と推論の間のミスマッチを軽減し、学習を改善します。

しかし、OPDは重大な計算上のボトルネックを導入します。すなわち、生徒が教師による監督を受ける前に、オンラインでロールアウトのサフィックス（接尾辞）を生成しなければならないという点です。従来のARデコーディングはトークン単位で行われるため、 $L$ トークンのサフィックスに対して $L$ 回の逐次的なデコーディングステップを必要とします。このプロセスは学習のクリティカルパス上に直接存在するため、あらゆるオンポリシー更新において多大なコストを強います。投機的ヤコビデコーディング（SJD）は、補助的なドラフトモデルなしに複数のトークンを並列に処理する方法を提供しますが、オリジナルのSJDは単一シーケンスの推論用に設計されています。これを視覚的OPDのラージバッチ設定に直接適用するには、2つの課題があります。

並列性の喪失： 各画像を個別にSJDで実行すると、GPUのバッチ並列性が損なわれます。
同期オーバーヘッド： 全ての画像を同期的に処理（バッチ同期型）すると、高速な画像が共有ウィンドウの境界で低速な画像の完了を待たされることになり、効率性の利点が打ち消されます。

2. 手法：HB-SJD

著者らは、視覚的OPDのために特別に設計されたロールアウト・バックエンドである Hybrid Batched Speculative Jacobi Decoding (HB-SJD) を提案します。HB-SJDは、教師、蒸留目的関数、および最適化手順を変更することなく、標準的な生徒のロールアウトエンジンを置き換えます。

コアコンポーネント

画像ごとの独立した進捗： バッチ同期型SJDとは異なり、HB-SJDは各画像が自身のデコーディングの進捗に従って進むことを可能にします。画像は、検証ウィンドウの完了を待つためにバッチ内の最も遅いメンバーを待機する必要はありません。
バッチ化された検証： 独立した進捗にもかかわらず、HB-SJDはGPUの効率性を維持します。これは、アクティブな画像の最後の確定済みトークンとドラフトウィンドウを使用して、検証器の入力を構築します。画像ごとの位置インデックスとKVキャッシュの位置により、これらのヘテロジニアス（不均一）な入力を、単一のバッチ化された生徒フォワードパスでまとめて処理できます。
ハイブリッド実行戦略： 画像がロールアウトを終了するにつれて、アクティブな画像の数は減少します。HB-SJDは、ハードウェアで校正された閾値（ $\gamma$ $γ$ ）に基づいて、2つの実行モードを動的に切り替えます。
- フル実行（Full Execution）： 元の物理的なバッチ形状を維持します。終了した画像は論理的に非アクティブ（状態更新なし）として扱われますが、メモリの行を占有します。これは、アクティブなバッチが大きい場合に効率的です。
- コンパクト実行（Compact Execution）： 残っているアクティブな画像に対してのみ検証器を実行し、対応するKVキャッシュの行のみにアクセスします。これにより、バッチサイズが小さくなった際に、非アクティブな行を処理するオーバーヘッドを回避します。
初期化と検証：
- 履歴ベースの初期化： HB-SJDは、ランダムな初期化ではなく、固定のヒストリーオフセットを用いて以前に確定したトークンでドラフトウィンドウを埋めることで、空間的な局所性を活用します。
- 検証戦略： 本システムは、貪欲な検証（Greedy Verification）（生徒のargmax予測と一致するドラフトを受け入れる）と、確率的検証（Probabilistic Verification）（確率比に基づいてドラフトを受け入れ、非argmaxの提案を許可する）の両方をサポートしています。

3. 主な貢献

視覚的OPDのためのSJD： 本論文は、視覚的オンポリシー蒸留のためのロールアウトエンジンとして、投機的ヤコビデコーディングを導入しています。現在の生徒のヤコビ反復からの予測を再利用することで、別途訓練された補助ドラフトモデルを必要とせずに、マルチトークン・ロールアウトを可能にします。
HB-SJDアーキテクチャ： 著者らは、SJDを単一シーケンス推論からラージバッチの視覚的OPDへと拡張するためにHB-SJDを提案しています。これは、独立した画像ごとの進捗と効率的なバッチ実行をユニークに組み合わせ、異なる速度で進む画像がGPUの並列性を維持しながら進行することを可能にします。
ドロップイン統合： HB-SJDは、生徒のロールアウト・バックエンドのドロップイン・リプレースメントとして設計されています。既存の手法（GKD、VarKDなど）の蒸留目的や学習手順を変更することはありません。
ハイブリッド実行： 本手法は、ロールアウト中の減少するアクティブバッチサイズに適応する、フル/コンパクトのハイブリッド実行戦略を導入しています。

4. 実験結果

LlamaGen（BおよびLバリアント）を用い、ImageNetデータセットにおいて、GKDおよびVarKDフレームワーク内での標準的なキャッシュ付き自己回帰（AR）ロールアウトと比較する実験を行いました。

効率性： HB-SJDは、ロールアウトおよびエンドツーエンドの学習時間を大幅に短縮します。
- 加速率： モデルサイズや蒸曲法（例：LlamaGen-Bに対するVarKD、貪欲な検証を用いた場合）に応じて、1.48倍から1.65倍のロールアウト加速を実現しました。
- 学習時間： 総学習時間を約1.18倍短縮しました。
- アブレーション研究： 独立した進捗が極めて重要であることが示されました。バッチ同期型SJDは、同期待ちが発生するため、標準的なキャッシュ付きARよりも低速でした（0.939倍の加速）。ハイブリッド実行戦略は、フル実行のみの場合と比較して、さらに1.182倍の加速を提供しました。
品質： 生成品質はARのベースラインと同等でした。
- 指標： FID、Inception Score (IS)、Precision、Recallにおいて、無視できる程度の差しか見られませんでした。例えば、LlamaGen-BとGKD+HB-SJDの組み合わせは、ARベースラインのFID 4.96に対し、4.95のFIDを達成しました。
- 定性的評価： 視覚的サンプルは、同等のセマンティックな内容、物体の構造、および局所的な詳細を示しました。
堅牢性： 加速は以下の項目において一貫していました。
- 異なるロールアウト長（ロールアウトが長いほど加速が増大）。
- 異なるバッチサイズ（8から64）。
- 異なる学習ステージ（2Kから22Kステップ）。
- 異なるヤコビウィンドウサイズ（ $W=16$ がバランスの取れた設定として特定されました）。

5. 意義と主張

本論文は、HB-SJDが視覚的OPDにおける主要な緊張関係（生徒が生成した軌跡の必要性と、それらを生成する高いコストとの対立）に対処することを主張しています。ロールアウト・バックエンドを学習プロセスから切り離すことにより、HB-SJDは、オンポリシー蒸留の理論的利点や生成の忠実度を損なうことなく、大幅な加速を可能にします。

著者らは、彼らのアプローチが（変化するポリシーへの適合を必要とする）言語モデルの強化学習における投機的デコーディングとは異なり、SJDは同じ生徒の予測を再利用するため、別途のドラフトモデルを必要としない点で独特であることを強調しています。さらに、本手法は、既存の視覚的自己回帰蒸留パイプラインに統合して、生成の忠実度を維持しながら学習コストを削減するための、実用的なエンジニアリング・ソリューションとして提示されています。