💬 NLP

Data-Efficient Autoregressive-to-Diffusion Language Models via On-Policy Distillation

本論文は、オンポリシー蒸留を通じて自己回帰型言語モデルを拡散言語モデルへと変換し、学習時と推論時のミスマッチを効果的に排除しながら、事前知識を保持しつつ学習トークン数を最大7,000倍削減する、データ効率の高いフレームワークであるOPDLMを紹介するものである。

原著者： Xingyu Su, Jacob Helwig, Shubham Parashar, Atharv Chagi, Lakshmi Jotsna, Degui Zhi, James Caverlee, Dileep Kalathil, Shuiwang Ji

公開日 2026-06-08

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

CC BY 4.0

原著者： Xingyu Su, Jacob Helwig, Shubham Parashar, Atharv Chagi, Lakshmi Jotsna, Degui Zhi, James Caverlee, Dileep Kalathil, Shuiwang Ji

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

想像してみてください。あなたは、世界クラスの素晴らしいシェフ（自己回帰モデル、または ARLM）を抱えています。このシェフは、一つひとつの料理を、一歩ずつ、次に何を入れるかを直前の材料を見ながら決めるという方法で、長年かけてその技術を磨いてきました。このシェフは非常に速く、多くの知識を持っていますが、一直線に料理することしかできません。

さて、このシェフに、「拡散（ディフュージョン）」という新しい革命的な調理スタイル（拡散言語モデル、または DLM）を学ばせたいとしましょう。この新しいスタイルでは、一歩ずつ料理を作る代わりに、シェフはランダムで正体不明の材料が入ったボウル（「マスクされた」シーケンス）からスタートし、一度に、かつ段階的に、完璧な料理へと洗練させていきます。このスタイルは、一度に複数の材料を推測できるため、より速く料理できる可能性があります。

問題：「翻訳」のギャップ
論文では、このステップ・バイ・ステップのシェフを「一度にすべてを行う」シェフに変えようとするこれまでの試みには、2つの大きな問題があったと説明しています。

記憶の喪失： ステップ・バイ・ステップの調理を強制的にやめさせ、一度に料理全体を推測させようとすると、シェフは長年の訓練で学んだ何千ものレシピを忘れてしまう傾向があります。それはまるで、熟練のピアニストに、突然目をつぶり、別の調（キー）で演奏しろと言うようなものです。彼らはその技術を失ってしまうかもしれません。
練習と本番のミスマッチ： 以前の方法では、シェフはランダムに材料をボウルに投げ込み、それを修正しようとする練習をしていました（ランダム・マスキング）。しかし、現実の世界（推論）では、シェフは自分の推測に対する自信に基づいて、実際に料理を洗練させていきます。練習が本番と一致していなかったため、いざという時にシェフは少し手際が悪くなってしまうのです。

解決策：OPDLM（「オンポリシー」のメンター）
著者らは、この両方の問題を解決する新しい手法である OPDLM（オンポリシー拡散言語モデル）を導入しています。これは、**オンポリシー蒸留（On-Policy Distillation）**というテクニックを用いた、スマートなトレーニングキャンプのようなものです。

その仕組みは、以下のシンプルな比喩で説明できます：

生徒と教師： 「生徒」は新しい拡散シェフです。「教師」は、元の、凍結されたステップ・バイ・ステップのシェフです（教師は、変更されることなくガラスケースの中に保管されています）。
「オンポリシー」のひねり： ランダムでバラバラな材料のボウルを使って練習する（従来の方法）代わりに、生徒シェフは、自分自身の新しい「一度にすべてを行う」スタイルを使って、フルコースの料理を作ることが許されます。彼らは、バラバラのボウルから完成した料理に至るまでの、自分自身の経路を生成します。
蒸留： 生徒が料理を完成させると、凍結された教師シェフが、その全く同じ料理を見て、「もし私がこの特定の料理を作っていたとしたら、足りない材料に対して、まさにこのように答えていただろう」と教えます。
学習： 生徒は自分の推測と教師の答えを比較し、そこから学びます。

なぜこれがゲームチェンジャーなのか
この手法により、生徒は自分が実際に生成した料理を用いて練習し（実際のパフォーマンスと一致させる）、かつ元のエキスパートから修正を受けるため、古い知識を忘れることなく、より速く学ぶことができます。

論文によれば、この方法は驚異的な効率性を誇ります：

データの節約： 従来の手法よりも15倍から7,000倍少ない訓練例で済みます。他の手法が、学習のために10億冊の本がある図書館を読み込む必要があったとしても、この手法なら数十冊だけで済むかもしれません。
「事前学習」コストの不在： ゼロから新しい拡散シェフを訓練する必要はありません（それは高価で時間がかかります）。既存のエキスパートを取り上げ、この特定の訓練を与えるだけで、彼らは変貌を遂げます。
魔法の維持： 新しいモデルは、たとえそれが明示的に教えられていなくても、元のシェフが持っていた「思考力」や「多言語能力」を維持します。それは、訓練によって元の脳が保持されたため、シェフが自然にフランス語を話したり、複雑な謎解きをしたりすることを覚えているようなものです。

結果
論文は、この新しい「OPDLM」シェフが、数学、コーディング、一般的な知識のタスクにおいて、既存の最高峰のシェフと同等のパフォーマンスを発揮することを示しています。しかも、それらを達成するために費やした労力とデータは、ごくわずかなものです。この手法は、モデルの「脳」を作り変えるという困難なプロセスを、シンプルで効率的なポストトレーニングのアップグレードへと変えてしまうのです。

テクニカル・サマリー：オンポリシー蒸留によるデータ効率的な自己回帰から拡散への言語モデルへの変換

1. 問題提起

本論文は、学習済みの自己回帰型言語モデル（ARLM）を拡散言語モデル（DLM）へと変換する際に内在する課題に対処している。近年の研究では、因果的アテンションを双方向アテンションに置き換え、拡散目的関数で再学習することでARLMをDLMに変換する試みがなされているが、著者らは、学習効率と性能を制限する2つの決定的な「分布シフト」を特定している。

知識保持のミスマッチ（Knowledge-Retention Mismatch）: 次トークン予測目的（ARLM）から拡散目的（DLM）への移行は、元のARLMの事前学習中に獲得された膨大な知識を弱体化させ、あるいは消失させる可能性がある。
訓練・推論状態のミスマッチ（Train-Inference State Mismatch）: 標準的なDLMは、一様ランダムなマスキング（順方向プロセス）によって破損したシーケンスを用いて訓練される。しかし、推論時、モデルは信頼度ベースのデコーディング・ヒューリスティックによって導かれた「逆マスキング・トラジェトリ（reverse unmasking trajectories）」に従ってシーケンスを生成する。推論中に遭遇する状態は、訓練中に見たランダムな状態とは大きく異なるため、学習が最適化されない。

既存の変換手法は、これを事前学習の継続として扱うことが多く、これらのミスマッチを明示的に解決したり、ゼロからDLMを事前学習するために必要な天文学的な量のデータ（数兆トークン）を必要としたりする。

2. 手法：オンポリシー拡散言語モデル (OPDLM)

著者らは、オンポリシー蒸留 (On-Policy Distillation: OPD) を用いてARLMをDLMへと変換するフレームワークである OPDLM を提案している。このアプローチは、変換を事前学習の問題ではなく、ポストトレーニング（事後学習）の問題として再定義するものである。

コアメカニズム

OPDLMは、学生モデルが自身の訓練データを生成し、元のARLMが凍結された教師として機能する「セルフ-OPD（self-OPD）」戦略を採用している。

学生モデル (Student Model): 元のARLMの重みから初期化され、ブロック単位の因果的（双方向）アテンションを使用するように修正されたARLMチェックポイント。
教師モデル (Teacher Model): オリジナルの、凍結されたARLM。
訓練トラジェトリ (On-Policy Training Trajectory): ランダムなマスクをクリーンなデータセットからサンプリングする代わりに、学生DLMは自身の逆デコーディング・トラジェトリ ( $\tau \sim \Gamma^{rev}_{\theta}$ ) をサンプリングする。これらのトラジェトリは、モデルが推論中に実際に遭遇する状態を表している。
蒸留目的関数 (Distillation Objective):
- サンプリングされた逆トラジェトリから、部分的にデノイズされた状態 $x_t$ が選択される。
- ブロック内の各マスクされたトークン位置 $i$ について、学生の予測分布 $p_{\theta}$ は教師の分布に一致するように訓練される。
- 決定的なのは、教師の分布が、ターミナルシーケンス $x_0$ （そのトラジェトリの最終生成出力）のアンマスクされた因果的プレフィックス $x_{<b}^0$ に条件付けられている点である。具体的には、ターゲットは $p_{ARLM}(\cdot | x_{<b}^0, x_{<i}^0)$ となる。
- 損失関数は、マスクされた位置における教師のトークンレベルの分布と学生の分布の間のカルバック・ライブラー（KL）ダイバージェンスである。

カリキュラム学習

変換直後の不安定性（因果的アテンションからブロック単位のアテンションへの構造的変化により、学生が低品質なトラジェトリを生成する可能性があること）に対処するため、著者らはロールアウト長カリキュラム (Rollout-Length Curriculum) を導入している。生成の長さは、ウォームアップステップにわたって最小値 ( $L_{min}$ ) から最大値 ( $L_{max}$ ) へと段階的に増加させ、学生が長いシーケンスに取り組む前に、一貫したコンテキスト上で教師に適合することを保証する。

3. 主な貢献

ARLMからDLMへの変換をポストトレーニングとして定式化: 本論文は、計算集約的な事前学習の段階を経ることなくDLMが得られることを示し、変換を効率的なポストトレーニング手順として位置づけている。
分布シフトの解決: 学生自身の逆トラジェトリ（オンポリシー）で訓練することにより、OPDLMは訓練と推論の状態のミスマッチを排除する。また、元のARLMを教師として使用することで、知識の保持を確実にする。
極めて高いデータ効率: 本手法は、確立されたDLMのベースライン（LLaDA、Dream、SDARなど）と比較して、15倍から7,000倍少ない訓練トークンで競争力のある性能を達成している。
汎用的およびタスク特化的な適用可能性: 本フレームワークは、汎用的なDLMとタスク特化型のモデル（数学用など）の両方を生み出すことができ、別途DLMの事前学習を必要としない。

4. 実験結果

著者らは、一般的な知識、数学、コーディングのベンチマークにおいて、OPDLM（4Bおよび8Bスケール）をLLaDA、Dream、SDAR、Fast-dLLM-v2などのベースラインと比較評価した。

性能 vs 効率: OPDLM-8Bは、わずか 0.0660億個の訓練トークン および 4.2 × 10¹⁸ FLOPs の計算量でありながら、AIME-24 (精度 14.7%) や GPQA-Diamond (精度 36.1%) といったベンチマークで強力な結果を残した。対照的に、ベースラインは数百億のトークンを必要とした（例：LLaDA-8Bは1,500Bトークンを使用）。
アブレーション研究:
- オンポリシー vs オフポリシー: オンポリシーのトラジェトリサンプリングを標準的なランダムマスキング（オフポリシー）に置き換えると、一貫した性能低下が見られた。これは、モデル自身の推論状態での訓練が成功の主要な要因であることを裏付けている。
- 知識保持: OPDLMは、変換中に明示的にターゲットとしていない、拡張思考（<thought> タグ内での推論）や多言語生成を含むゼロショット能力を保持した。これは、英語の精選コーパスのみで訓練されたにもかかわらず実現された。
推論効率: モデルはマルチトークンデコーディングをサポートしている。デコーディングの信頼度閾値とブロックサイズを調整することで、精度を適度に犠牲にしつつ、スループットが1ステップあたり2〜3トークンを超えるトレードオフが可能であることを示した。
タスク特化型の適応: 数学データ（MATHデータセット）のみで厳密に訓練した場合、OPDLM-MATHは標準的なDLMベースラインを上回り、検証可能な報酬を用いた強化学習（RL）を必要とせずに、専門化されたRL訓練モデル（TraDo）に匹敵する性能を示した。

5. 意義と主張

本論文は、OPDLMがDLM訓練の効率における新たなパレート境界を確立すると主張している。オンポリシー蒸留を活用することで、ARLMからDLMへの変換を最小限の計算リソースで達成できることを示し、ゼロからの事前学習という「法外なコスト」なしに、拡散ベースのテキスト生成を身近なものにできると述べている。

著者らは、このアプローチによって、コミュニティが将来のARLMを効率的に拡散モデルへと適応させることが可能になり、元のモデルの豊かな知識を保持しながら、拡散による並列デコーディングの利点を得られることを強調している。本研究は、自己回帰モデルの成熟度と、拡散モデルの潜在的な利点の間の溝を埋めるための実用的な経路を提示している。