テクニカル・サマリー: Nemotron-Labs-Diffusion

問題提起

自己回帰（AR）言語モデル（LM）は、厳密に逐次的なトークン単位のデコーディングプロセスによって根本的な制限を受けており、これが推論の並列性を制限し、特に低バッチサイズの設定においてリソースの未利用を招いています。拡散（Diffusion）LMは、1回のフォワードパスで複数トークンをデコードするマルチトークン・デコーディングを可能にすることで、並列生成を実現する有望な代替案として浮上していますが、精度と学習効率の面では依然としてARモデルに遅れをとっています。既存の拡散LMは、同等の性能に達するために大幅に多くのデータを必要とする傾向があり、実用的な効率と精度のトレードオフにおいてマルチトークン予測（MTP）手法を上回ることに苦慮しています。さらに、ARと拡散のパラダイム間の関係は不明確です。これらが競合するアプローチなのか、あるいは両方の強みを活用するために調和させることができるのかは定まっていません。

メソドロジー

著者らは、AR、拡散、および自己投機（self-speculation）デコーディングを単一のアーキテクチャ内に統合した、三モード言語モデルファミリーであるNemotron-Labs-Diffusionを導入します。コアとなる手法は、結合学習フレームワークと柔軟な推論戦略です。

1. 結合AR-拡散学習

モデルは、AR次トークン損失とブロック単位の拡散デノイジング損失の加重結合を用いて学習されます：
$$\mathcal{L}(\theta) = \mathcal{L}_{AR}(\theta) + \alpha \mathcal{L}_{diff}(\theta)$$

• 二段階学習: 強固な左から右への言語的事前分布を確立するために、著者らは二段階戦略を採用しています。ステージ1ではAR目的関数のみで学習を行います（$\alpha=0$）。ステージ2では拡散の監督（$\alpha=0.3$）を導入し、ARの事前分布を上書きすることなく並列デコーディング機能を統合します。
• グローバル損失平均化: 拡散学習におけるノイズを含むトークン数の変動による最適化を安定させるため、著者らはシーケンス単位の平均化ではなく、バッチ全体でのグローバル損失平均化戦略を利用し、ノイズを含むトークンが少ないサンプルによる勾配の分散を防いでいます。
• アテンション・パターン: モデルは、構造化されたアテンションマスクを持つデュアルストリーム入力（ノイズありおよびクリーンなビュー）を利用します。決定的なのは、クリーンなストリームが厳格な因果的（causal）マスクを強制している点であり、これによりラベル漏洩なしにクリーンなトークン上でAR目的関数を計算できます。一方で、ノイズを含むストリームは、ブロック内での双方向アテンションを使用して並列デノイジングを行います。

2. 三モード推論

学習済みモデルは、推論時に選択可能な3つのデコーディングモードをサポートしています：

• ARモード: 因果的アテンションを用いた標準的な左から右への生成。このモードは高コンカレンシーなクラウド・サービングに最適化されています。
• 拡散モード: 並列的なブロック単位のデノイジング。モデルは信頼度閾値に基づいて、反復的に複数のトークンをデノイズします。軽量化された学習済みサンプラーを使用し、トップ1の予測が正しいかどうかを予測することで、1フォワードあたりのトークン数（TPF）とエラー率のトレードオフを最適化します。
• 自己投機（Self-Speculation）モード: 拡散パスウェイが複数の候補トークンを並列にドラフトし、ARパスウェイがそれらを検証するハイブリッドアプローチです。
  • 線形自己投機: 拡散が$k$個のトークンをドラフトし、ARが第2のフォワードパスでそれらを検証します。著者らは、拡散ドラフターをAR検証器により適合させるために、LoRAアダプターのチューニングを行うことでこれを強化しています。
  • 二次（Quadratic）自己投機: 単一のフォワードパス内で全ての受理可能ポジションを準備するために二次デコーディングを使用するバリアントですが、カーネル最適化の面で現在課題に直面しています。

3. 光速（Speed-of-Light: SOL）分析

著者らは、拡散デコーディングの理論的上限を定量化するために「光速（SOL）」分析を提案しています。再帰的な動的圧縮法を用いて、シリアルなデノイジング対象に対する一致する位置の最大の安全なサブセットを特定することで、AR検証なしで達成可能な最大TPFを推定しています。この分析により、現在の信頼度ベースのサンプリングでは、多くの並列性が活用されていないことが明らかになりました。

主な貢献

1. 統合された三モード・アーキテクチャ: AR、拡散、および自己投機デコーディングをシームレスに統合した初のモデルであり、展開ニーズ（例：高コンカレンシー vs 低コンカレンシー）に応じた動的な切り替えを可能にします。
2. 目的関数の相補性: ARと拡散の目的関数が競合するのではなく、相補的であることを示しました。ARは強力な言語的事前分布と将来の計画を提供し、拡散は並列生成を解禁します。結合された目的関数は、ARの精度を維持またはわずかに向上させつつ、高速な拡散モードを可能にします。
3. MTPに対する優位性: 自己投機（拡散ドラフト + AR検証）が、既存のMTP手法（Eagle3など）よりも、特に低コンカレンシーにおいて、受理率と実デバイスでの効率の両面で優れていることを示しました。
4. スケーラビリティと汎用性: モデルファミリーは3Bから14Bパラメータまでスケールし、Vision-Language Models (VLM) にも拡張され、一貫して最先端のオープンソースARおよび拡散LMを上回る性能を発揮します。

結果

• 精度と効率: Nemotron-Labs-Diffusion-8Bモデルは、一般的なベンチマーク、コーディング、数学においてQwen3-8Bと同等または優れた精度を達成しながら、拡散モードにおいて1フォワードあたり6倍多くのトークンをデコードします。
• スループットの向上: SPEED-Benchにおいて、本モデルはNVIDIA GB200 GPU上のSGLangを使用した場合、Qwen3-8Bよりも4倍高いスループットを達成しました。
• 自己投機の性能: LoRAチューニングを施した線形自己投機は、平均TPFとして5.99×（8Bモデル）および6.38×（14Bモデル）を達成し、ベースラインを大幅に上回りました。
• SOLのポテンシャル: SOL分析は、最適なサンプラーを使用した場合、拡散デコーディングが現在の自己投機手法よりも76.5%多くのトークンを1フォワードあたりに予測できる可能性を示しており、将来的な余地が多分にあることを示唆しています。
• VLMの性能: VLMバリアントは、ARベースラインと同等の精度を維持しながら、拡散モードで2.46×–3.15× TPF、長い応答に対しては最大7.45× TPFの線形自己投機を提供します。

意義と主張

本論文は、Nemotron-Labs-Diffusionが、ARと拡散は排他的なものではなく、それらを統合することでより柔軟で効率的な推論エンジンを作成できるというパラダイムシフトを象徴していると主張しています。著者らは以下のように述べています：

• 調和は可能である: ARと拡散の目的関数を統合することは、純粋な拡散モデルに伴う典型的な精度低下を回避する有望な経路です。
• 展開の柔軟性: 三モード設計により、高コンカレンシーの設定ではARモデルのドロップイン・リプレースメントとして機能し、低コンカレンシーのパーソナルAIシナリオでは自己投機を通じて大幅なスピードアップを提供できます。
• 将来の可能性: 「光速（SOL）」分析は、拡散デコーディングの真のポテンシャルが現在のサンプラーではまだ実現されていないことを示唆しています。現在の性能とSOLの天井との間のギャップは、サンプラーの最適化に関する将来の研究が、さらなる並列性を解き放ち、プレフィックスベースの投機的デコーディングの限界を凌駕する可能性があることを示しています。

結論として、現在のインフラストラクチャは線形自己投機を支持していますが、拡散デコーディングの理論的上限は、高度に並列な非プレフィックス生成へと向かう長期的な軌道を示唆しており、LLM推論の効率を根本的に変える可能性があります。