NVIDIA、再学習不要でスループットを2.42倍にする拡散LLM「Nemotron TwoTower」を公開

NVIDIAの研究チームは、従来の自己回帰型デコーディングと比べて2.42倍高速にテキストを生成できるディスクリート拡散言語モデル「Nemotron-Labs-TwoTower」のオープンウェイトと学習コードを公開した。このモデルは、ベースラインのベンチマーク品質を98.7%維持しつつ、大規模な事前再学習を必要としない点が大きな特徴である。既存の自己回帰型モデルの資産を活かしながら、並列生成による高速化を低コストで実現する新たなアプローチとして注目される。

■事前再学習のコストを劇的に削減する「2タワー」アプローチ

現在の大規模言語モデル（LLM）開発における最大のコストは、膨大なトークンと数カ月ものGPU時間を要する大規模な事前学習である。NVIDIAが公開した「Nemotron-Labs-TwoTower（以下、TwoTower）」は、既存の自己回帰型（AR）モデルのチェックポイントを持つ開発者が、元の学習データのわずか一部（約8%）を用いて2つ目のデノイジング（ノイズ除去）ネットワークを学習させるだけで、拡散ベースの並列生成機能を追加できることを実証した。

このアーキテクチャは、モデルがすでに学習した内容を破棄せず、その上に並列生成能力を重ね合わせる。これにより、拡散言語モデルへの移行コストを劇的に引き下げることが可能になる。ウェイトはHugging Faceの「NVIDIA Nemotron TwoTower」コレクションにて、商用利用が可能な「NVIDIA Nemotron Open Model License」の下で公開されている。

■ボトルネックを解消する「2つのタワー」の仕組み

現在広く普及しているLLMは、テキストを左から右へ1トークンずつ自己回帰的に生成する。この方式では、バッチサイズが小さいインタラクティブなチャットやエージェントのループ処理において、推論ワークロードが演算性能ではなくメモリ帯域幅の制限を受けやすくなる。つまり、GPUは計算処理ではなく、データの移動を待つ時間に大半を費やしている。

ディスクリート拡散言語モデルは、1トークンずつの生成ではなく、ブロック単位での反復的なデノイジングによってこの問題を解決する。しかし、従来の拡散言語モデルでは、すでに確定したクリーンなコンテキストの表現と、生成中のノイズを含むブロックのデノイジングという、相反する2つの役割を1つのネットワークに同時に担わせていたため、性能が制限されていた。

TwoTowerの最大の特徴は、これらの役割を同じバックボーン上に構築された2つの独立したネットワークに明確に分割した点にある。それは、凍結された「自己回帰コンテキストタワー（AR context tower）」と、新たに学習された「拡散デノイザータワー（diffusion denoiser tower）」である。

■詳細なアーキテクチャと動作原理

両タワーは、300億パラメータのハイブリッドモデル「Nemotron-3-Nano-30B-A3B」を共通のバックボーンとして使用している。このバックボーンは、23のMamba-2レイヤー、6つの自己アテンションレイヤー、23の混合専門家（MoE）レイヤーの計52レイヤーで構成されている。推論時には、1トークンあたり約30億の活性パラメータのみが動作する。

第1のタワーである「ARコンテキストタワー」は、標準的な自己回帰モデルとして動作し、プロンプトと確定済みの出力トークンを左から右へと処理する。このタワーはTwoTowerの学習プロセス全体を通じて完全に凍結（フリーズ）される。

第2のタワーである「拡散デノイザー」は、学習が行われる領域である。生成中のノイズを含むトークンブロックを受け取り、マスクされた拡散目的関数を通じてそれらを洗練させる。このデノイザータワーは、各レイヤーにおいて、凍結されたコンテキストタワーの対応するレイヤーからクロスアテンションを通じて条件付けを受け取り、豊かな文脈表現を引き継ぐ仕組みとなっている。

このアプローチにより、デコーディングの初期段階で1ステップあたり複数のトークンを確定させることが可能になり、自己回帰型のベースラインと比較して実時間（ウォールクロック時間）で2.42倍の高速化を達成している。

■実用性を担保する「学習コスト」の低さ

デノイザータワーの学習には、約2.1兆トークンが使用された。これはバックボーンモデルの事前学習に使用された25兆トークンの約8%にすぎない。つまり、すでに大規模な自己回帰モデルの事前学習に投資した組織は、拡散ベースの高速化メリットを得るためにゼロからやり直す必要がないことを意味している。論文では、TwoTowerの手法が任意の事前学習済み自己回帰言語モデルに適用可能な汎用的アプローチであると説明されており、公開された学習コードも他のバックボーンへの適応をサポートしている。

■ベンチマーク結果とトレードオフ

品質とスループットのトレードオフは、1ステップあたりに確定させるトークン数を決定する確信度パラメータ（γ）によって制御される。デフォルトの動作条件（確信度しきい値 γ = 0.8、ブロックサイズ 16、NVIDIA H100 GPU 2基を使用、BF16精度）において、知識、推論、言語タスクにわたる総合的なベンチマーク品質は、自己回帰ベースラインの1.3パーセンテージポイント以内に収まっている。

ただし、タスクの性質によって結果は異なる。双方向の文脈やグローバルな一貫性が重視されるタスク（常識推論や多言語評価など）では、自己回帰ベースラインと同等以上の品質に回復した。一方で、厳密な左から右への逐次的な推論が必要なタスクは影響を受けやすく、Pythonコード生成ベンチマーク「HumanEval」では79.27から75.58へとスコアが低下した。数学ベンチマークでも同様のわずかな低下が見られており、これは並列ブロック生成における既知の制限とされている。

■動作に必要なハードウェア要件

公開されたチェックポイントには両方のタワーが含まれており、合計で約600億パラメータ、ストレージ容量は約126GBに達する。フル機能の2タワー拡散デコーディングを実行するには、BF16精度でそれぞれ約59GBのメモリを持つGPUが2基必要となる。これは通常、NVIDIA H100やA100などの80GBデータセンター向けGPU 2基の構成に相当する。

GPUを1基しか持たない環境向けには、互換性テストやデバッグ用の「模擬ARモード」と、シングルGPU（80GB）で動作し元の自己回帰バックボーンと同等のパフォーマンスを発揮する「標準AR専用デコーディングモード」の2つのフォールバックモードが用意されている。これらは同一のウェイトからロードでき、設定の変更だけで切り替えが可能である。

■今後の課題と位置づけ

なお、今回公開されたチェックポイントは「ベースモデル」であり、指示調整（インストラクションチューニング）や、チャット、コーディングアシスタント、安全な展開のためのアライメント調整は施されていない。実用的なチャットモデルなどを構築するには、追加の微調整（ファインチューニング）や安全対策のアライメントが必要となる。

本研究は、Fitsum Reda氏、John Kamalu氏、Roger Waleffe氏、Mostofa Patwary氏、Mohammad Shoeybi氏、Bryan Catanzaro氏らによって進められ、他の研究チームがTwoTowerのアプローチを異なるバックボーンモデルに適応できるよう、ウェイトとともに学習コードも公開されている。

■注目ポイントQ&A

●拡散言語モデルとは何ですか？従来のLLMとどう違いますか？

従来のLLMはテキストを左から右へ1トークンずつ順番に生成（自己回帰）しますが、拡散言語モデルはノイズやマスクされた状態のブロックから開始し、複数ステップかけて並列にノイズを除去（デノイジング）しながらテキストを洗練させていきます。これにより、1回の処理で複数のトークンを同時に確定させることができるため、推論速度（スループット）を大幅に向上させることができます。

●なぜ拡散モデルはコードや数学のタスクで性能が低下するのですか？

コードや数学の生成では、「あるトークンの正しい値が、その直前や直後に来るトークンと厳密に連動している」という強い逐次的依存関係があります。拡散モデルがブロック内を並列かつ双方向のアテンションで処理する場合、この厳密な左から右への順序関係を完全に代替することが難しいため、わずかな精度低下が発生します。

●TwoTowerのアプローチはNemotron以外の言語モデルにも適用できますか？

はい、適用可能です。論文の著者らによると、TwoTowerは任意の事前学習済み自己回帰言語モデルに適用できる汎用的なアプローチです。既存モデルのウェイトを2つのタワー（凍結するコンテキスト用と、学習するデノイザー用）に複製できることが要件となります。NVIDIAは他のモデルへの適応を支援する学習コードも同時に公開しています。

●Nemotron TwoTowerを実行するにはどのようなハードウェアが必要ですか？

フル機能の2タワー拡散デコーディングを実行するには、BF16精度でそれぞれ約59GBのメモリを搭載したGPUが2基（NVIDIA H100やA100の80GBモデルなど計2基）必要です。ただし、GPUが1基のみの環境でも、通常の自己回帰（AR）モードであれば、シングル80GB GPUで動作させることができます。

元記事: NVIDIA Diffusion LLM Hits 2.42x Throughput Without Retraining: Nemotron TwoTower Released