画像や動画生成の中核技術である拡散モデルが、テキスト生成にも進出し始めている。NVIDIAの新たなTwo Tower方式は、既存の自己回帰モデルを複製し、一方を凍結した文脈処理用、もう一方をマスク除去型の生成用として再学習する構成である。ブロック単位の並列生成によって高速化を図りつつ、元モデルの性能を高水準で維持する一方、数学やコード生成での低下、ブロックサイズへの強い依存性、メモリ負荷などの課題も明らかになっている。
===
テキスト生成にも広がる拡散モデル
現在、拡散モデルは動画生成や画像生成の中核を担っています。しかし今、その状況が変わりつつあり、拡散モデルがテキストの世界にもやって来ています。
従来、言語は本質的に逐次的なものなので、言語生成には自己回帰モデルが使われてきました。しかし、この次トークン予測には大きな欠点があります。計算量とメモリの両方がボトルネックになるのです。その理由は、モデルがそれまでに生成したすべてのトークンを振り返る必要があるからです。
ところが今、拡散モデルがテキスト生成に進出することで、この状況が変わろうとしています。拡散モデルでは、逐次生成を並列生成に置き換えます。
もちろん、画像の場合のように単純にノイズを加えるのとまったく同じではありません。しかし最終的には、ブロックごとに並列生成する仕組みになります。
これまでにも、テキスト用の拡散モデルをゼロから学習させる試みは数多くありました。しかし最近、自己回帰モデルと拡散モデルを組み合わせる非常に興味深い新しいアプローチが登場しています。
それがNemotronのTwo Towerという新しい拡散言語モデルです。NVIDIAは2.4倍高速だと主張しています。
ただ、考えてみれば、速度はそもそも拡散モデルを使う大きな理由の一つです。そのため、私にとって本当に興味深いのはそこではありません。
興味深いのは、このネットワークの半分が一切学習されないことです。完全に凍結されていて、実はその凍結こそがNVIDIAの仕掛けの核心なのです。
NVIDIAは既存の自己回帰モデルを用意し、それを2つの異なるタワーとして複製しました。そして片方を凍結し、もう片方には空欄を埋める方法を学習させています。
この空欄を埋めるという部分が最も重要です。テキストにおける拡散モデルでは、ノイズの代わりにマスキングを使うからです。
この点については、動画の後半でもう一度戻ってきます。テキスト拡散について、多くの人が混乱している部分だからです。
では、このモデルをもう少し詳しく見ていきましょう。
Two Towerモデルの構造
現時点では、Hugging Face上で公開されているのはベースモデルだけです。Instruction調整版、RL版、チャット向けに調整されたバージョンは、まだ公開されていません。
現在、テキスト拡散は非常に注目されています。InceptionのモデルやGoogle Gemini関連のモデル、さらには拡散型のGemmaまで登場しています。そのため、NVIDIAがこの分野で最初というわけではありません。
重要なのは、このモデルを具体的にどう構築したのかです。そこが、これまで見てきた方式とは大きく異なっています。
NVIDIAはNemotron 3 Nanoモデルを用意し、そのまったく同じモデルを複製しました。
各タワーは52層で動作します。両方のアーキテクチャは完全に同じで、Mamba 2と自己注意機構を組み合わせています。またMoE部分では、各トークンにつき128個のエキスパートから6個を使用し、それに共有エキスパートが1個加わります。
そのため、各タワーで実際に活性化されるパラメータは約30億です。
モデル全体として見ると、総パラメータ数は約600億になります。つまり、完全精度で動かすにはH100が数台必要になる規模です。
ここで興味深いのは、この構造変更によって性能が失われるのかどうかです。そして実際、そこがこのモデルの非常に面白い部分です。
元の自己回帰モデルからどれだけ性能を維持できるのか
総合評価では、元の自己回帰モデルの品質を98.7％維持できています。ただし、私は総合値そのものはそれほど重要ではないと思っています。
本当に面白くなるのは、個別のベンチマークを見たときです。
一般知識のような分野では、性能を非常によく維持しています。多言語性能は実際には3ポイント上がっていて、これは少し興味深い結果です。
一方、コード性能は約2ポイント低下し、数学性能は3ポイント低下しています。
つまり、このアーキテクチャを使うことには確かな代償があります。そして、性能低下が起こる場所もある意味では予想どおりです。
数学やコードのような分野では、並列生成中にたった1つのトークンを間違えただけでも、それが全体の性能に影響を与える可能性があるからです。
1つのネットワークに2つの仕事をさせない
このアプローチは、拡散モデルが抱えている最大の問題の一つを解決しています。
従来の方式では、1つのネットワークが同時に2つの仕事をしていました。すでに存在するクリーンなテキストを表現する仕事と、生成途中の乱れたテキストをデノイズする仕事です。
論文では、この2つの仕事が同じ重みを異なる方向へ引っ張るため、その結果として両方の性能が悪化すると説明されています。
たとえるなら、一人の人間に基調講演をさせながら、同時に自分自身の同時通訳までやらせるようなものです。
そこでNVIDIAは、仕事を2人の専門家に分けました。それが先ほど見た2つのタワーです。
講演者はそのまま話し続けます。これが基本的にコンテキストタワーです。
ここが面白いところなのですが、このタワーはプロンプトを読み込み、自己回帰モデルと同じように、すべてを厳密に左から右へ処理します。そして、このタワーは一切学習されません。時間が止まったかのように、そのまま凍結されています。
一方、通訳ブースにいるのがデノイザーです。開始時点では同じチェックポイントを使っていますが、こちらは異なるテキスト生成方式のために完全に再学習されています。
2つのタワーはどう通信するのか
では、この2つのタワーは実際にどうやって互いに情報をやり取りするのでしょうか。
これは、この論文の中でも私が特に気に入っている部分です。
従来の設計では、多くの場合、コンテキストネットワーク最上部の最終隠れ状態、あるいはその要約情報だけを受け渡していました。
しかしTwo Towerでは、2つのタワーを層ごとに接続します。
つまり、デノイザーの第1層はコンテキストタワーの第1層にクロスアテンションし、第2層は第2層に接続するというように、対応する層同士が通信します。
同じ形をした2本のタワーがあり、屋上同士を1本のケーブルでつなぐのではなく、すべての階に空中連絡橋が架かっている様子を想像してください。
これが機能するのは、両方のモデルがまったく同じ状態から始まっているからです。
つまり、互いに通信する各層が、同じ言語を話しているわけです。
デノイザーには、小さなEDMタイマーも追加されています。これは画像拡散モデルから取り入れたアイデアで、現在のブロックがどの程度破損しているかをモデルに伝えます。
これによって追加されるパラメータは150万個です。600億パラメータのネットワーク全体と比べれば、ほとんど無視できる規模です。
16トークンを並列で生成する仕組み
では、実際の生成方法を見ていきましょう。これは画像の生成方法とはかなり異なります。
この場合、デノイザーは単に1つのトークンを予測するわけではありません。
16個のスロットで構成されたブロックが与えられ、最初はすべてがマスクトークンになっています。
そこからループ処理を行い、16個の位置をすべて同時に並列で予測します。そして、信頼度が一定のしきい値を超えた位置を確定します。
次の反復では、自信を持てなかった部分を再びマスクし、もう一度実行します。その時点では、すでに確定したトークンがアンカーとして残っています。
これを理解するうえで非常に良い例が、数独です。
数独では、左上から順番に1マスずつ埋めていくわけではありません。確実に分かるマスから埋めていきます。そして、1つ数字を確定するたびに、その周囲に入る可能性のある数字が絞られていきます。
では、この仕組みがTwo Tower方式ではどう動くのでしょうか。
1つのブロックが完全にクリーンになると、凍結されたタワーがそれを読み込み、キャッシュを更新します。そして、新しい16個のマスクからなるブロックが現れます。
自己回帰モデルは1回のパスで1トークンを生成しますが、拡散モデルはブロック単位で生成しているわけです。
実際の生成順序はほぼ左から右になる
公開されている生成トレースを見ると、この動作を実際に確認できます。
各ブロックでは、最初の数ステップのうちに、確定されるトークンが一気に増えていきます。
興味深いのは、モデルは原理上どの順番でもトークンを確定できるにもかかわらず、実際の確定順序はほぼ完全な左から右への三角形を形成することです。
これはアーキテクチャを考えると、かなり興味深い結果です。
52層のうち23層が基本的にMambaで、アテンション層はわずか6層しかありません。
そのため、この左から右へ進もうとする性質は、ある意味では納得できます。
新しい能力を獲得したわけではない
ここで考えるべきことがあります。
この2つのタワーがまったく同じモデルのコピーなら、新しい能力が生まれているのでしょうか。
実際のところ、新しい能力は特にありません。
デノイザーは2.1兆トークンで学習されています。これは、元のNemotron Nanoモデルの学習に使われた25兆トークンのごく一部です。
つまり、これは本質的にはデコード方式の変換です。モデルそのものが賢くなったわけではありません。
NVIDIAはアブレーション研究も行っています。
まず、通常の方法でバックボーンの学習をそのまま継続する方法を試しました。しかし、結果は悪化しました。
一般知識では約10％低下し、数学では7.8％低下しました。
さらに、2つのタワーの重みを共有し、1組の重みに両方の仕事をさせる方法も試しました。
これも、ほとんど大失敗でした。一般的なタスクでは26ポイント低下し、その他の分野でも全体的に悪化しました。
こうしたアブレーション研究の結果を見ると、ネットワークを凍結する現在の方式のほうが合理的だったことが分かります。
ただし、それにも代償があります。
高速化と引き換えに必要になるメモリ
この方式によって生成速度は向上します。
しかしその一方で、300億パラメータ級の重みを持つタワーが2つ、常にメモリ上に存在することになります。
ただし、シーケンス全体にわたってキャッシュを保持するのは凍結されたタワーだけです。
そのため、シーケンス長が長くなるにつれて増えるメモリ量は、通常の自己回帰モデルと同じような伸び方をします。
ブロックサイズを変えると生成が崩壊する
論文では、このモデルが特定の状況で非常に壊れやすいことを示す、かなり興味深い結果も紹介されています。
公開されたチェックポイントは、ブロックサイズ16で学習されています。
これを64に変更してサンプリングすると、たった1つの設定を変更しただけなのに、生成性能が完全に崩壊しました。
たとえばHumanEvalは76から約20％まで低下しました。
そしてGSM8Kは2.2％まで落ちています。これはかなり衝撃的です。22％ではありません。2.2％です。
ところが一方で、多肢選択式のMMLUは、78％付近からほとんど動きませんでした。
つまり、知識そのものはモデルの中に存在しています。
崩壊しているのは、基本的に生成プロセスそのものです。そして、この生成プロセスは使用するパラメータや設定に対して非常に敏感です。
一方で、ブロックサイズを16より小さくすると、性能自体は維持されることも分かりました。ただし、その場合はかなり遅くなります。
今後のInstruction版とRL版に注目
忘れてはいけないのは、これはベースモデルだということです。
NVIDIAは、ポストトレーニング済みのバージョンも公開する予定です。
そのため、モデルの実際の能力をきちんとテストするには、それらの公開を待つ必要があります。
とはいえ、これはかなり刺激的で新規性のあるアーキテクチャです。
Instruction版やRL版のモデルが登場したとき、どのような性能を見せるのか確認する必要があります。
NVIDIAは、モデル開発において、こうした新しくて少し異色なアーキテクチャを本当に積極的に推進しています。
たとえばNemotron自体も、MambaとTransformerを組み合わせたハイブリッドモデルです。
新しいアイデアをここまで積極的に試している企業は、ほかにはあまり見かけません。
強いて言えば、DeepSeekのような企業は良い例かもしれません。中国企業はさまざまな分野で多くのイノベーションを進めています。
ともかく、この内容が役に立ったならうれしいです。ご視聴ありがとうございました。それではいつものように、また次回お会いしましょう。
===
コメント