NVIDIA、Nemotron-Labs-TwoTowerを公開：凍結された自己回帰バックボーン上に構築されたオープンウェイト拡散言語モデル

NVIDIAは、事前学習済みの自己回帰（AR）バックボーン上に構築された拡散言語モデル「Nemotron-Labs-TwoTower」をリリースしました。このモデルはオープンウェイトで提供され、NVIDIA Nemotronオープンモデルライセンスの下で利用可能です。今回のリリースは、テキスト生成におけるスループットのボトルネックを解決することを目的としています。

従来のARモデルはトークンを1つずつ復号するため、逐次処理が生成スループットを制限していました。一方、離散拡散言語モデルは並列にトークンを生成し、反復的に洗練します。しかし、ほとんどの拡散言語モデルは単一のネットワークで2つの役割（クリーントークンの表現と劣化トークンのデノイジング）を担っていました。TwoTowerはこれらの役割を2つのタワー（「コンテキストタワー」と「デノイザータワー」）に分離します。

TwoTowerは、Mamba-2、自己注意、混合エキスパート（MoE）層を組み合わせたオープンウェイトのハイブリッドバックボーン「Nemotron-3-Nano-30B-A3B」を基盤としています。各タワーは52層（23のMamba-2層、6の自己注意層、23のMoE層）で構成されます。公開されたチェックポイントは両方のタワーを含み、合計約600億パラメータ、各タワーのアクティブパラメータはトークンあたり約30億です。MoEは128のルーティング可能なエキスパートを持ち、そのうち6つがアクティブになり、2つの共有エキスパートが加わります。

両タワーは同じバックボーンチェックポイントのコピーから開始されますが、デノイザータワーのみがトレーニングされ、ARコンテキストタワーは凍結されたままです。デノイザーは約2.1Tトークンでトレーニングされましたが、バックボーンは25Tトークンで事前トレーニングされています。

推論では、ARコンテキストタワーはプロンプトとコミット済みトークンに対して因果的に動作し、レイヤーごとのKVキャッシュと最終的なMamba-2状態を生成します。拡散デノイザータワーはノイズの多いブロックを反復的に洗練します。ブロック内では双方向の注意を使用し、過去のクリーンブロックに対して因果的です。タワーはレイヤーごとに接続されており、デノイザータワーの各レイヤーはコンテキストタワーの対応するレイヤーにクロスアテンションを送ります。このレイヤー整列されたクロスアテンションにより、バックボーンの表現へのマルチスケールアクセスが可能になります。さらに、Mamba-2層はコンテキストタワーのMamba状態から初期状態をシードされ、拡散タイムステップはadaLN-single時間条件付けを介して各レイヤーを変調します。このadaLNモジュールは約150万パラメータしか追加しません。

生成はブロックごとに進行します。各ブロックはS個の[MASK]トークンで始まり、デノイザーがTステップかけて洗練した後、コミットされます。その後、コンテキストタワーがコミット済みトークンを処理してキャッシュを更新します。これにより、複数のデノイジングステップが1トークン復号よりも高速になる理由が説明されます。AR復号は1ステップで正確に1トークンをコミットしますが、TwoTowerは洗練の初期段階で1ステップあたり複数のトークンをコミットします。

2×H100 GPUでのBF16評価では、デフォルトの動作点（信頼度マスキング、閾値γ=0.8、ブロックサイズS=16）において、TwoTowerはMMLU、ARC-Challenge、WinoGrandeなどのタスクでARベースラインに近い性能を示し、総合品質は98.7%、生成スループットは2.42倍でした。一般知識タスクではベースラインとの差は約1ポイント以内、コードと数学ではやや低下し、常識と多言語では同等かわずかに改善されました。γを下げると1ステップあたりのコミットトークン数が増えてスループットが向上しますが、品質は低下します。

このチェックポイントは3つの推論モードを提供します：完全な2タワー拡散（2GPU必要、各GPU約59GB BF16メモリ）、模擬AR、およびAR（単一の80GB GPUで実行可能）。最も直接的なユースケースはバッチ生成の高速化です。例えば、合成テキスト生成では、わずかな品質低下と引き換えに大幅なスループット向上を得られます。ユーザーは信頼度閾値を調整して品質と速度のトレードオフを制御できます。また、コンテキストタワーはLMヘッドを保持しているため、投機的復号、検証、ARスコアリングに使用でき、チームは1つのチェックポイントから複数のモードを実行できます。

主な強み：オープンウェイトで商用利用可能；デフォルト動作点で品質98.7%を維持しつつスループット2.42倍；1つのチェックポイントで3つの復号モードをサポート；デノイザーのトレーニングコストが低い。弱点：完全な2タワー推論には2GPUが必要；コードと数学の性能低下が大きい；固定モデル重みのメモリフットプリントが大きい；リリース版はベースモデルであり、命令チューニングやアライメント未実施。