大規模言語モデルの構造を3Dで視覚化してどんな計算が行われているのかを見やすく表示してくれるサイト「LLM Visualization」

ChatGPTなどのチャットAIは、内部で多数の計算を行って文章を生成しています。「LLM Visualization」は内部でどのようなパラメーターが保存されていてどういった計算が行われているのかを3D形式で見やすく視覚化してくれるサイトとなっています。

LLM Visualization
https://bbycroft.net/llm

サイトにアクセスするとこんな感じ。画面の左側に解説、右側に3Dモデルが設置されています。解説の「Continue」をクリックしてみます。

解説では約8万5000パラメーターのモデル「nano-gpt」で3種類の文字を並び替えるタスクを行い、大規模言語モデル内部の計算を確認していくとのこと。スペースキーを押すことで次のセクションへと進めます。

解説の段階に応じて3Dモデルの注目するべき場所がアップになるため、どの場所の話をしているのかが分かりやすくなっています。

3Dモデルのパーツにカーソルをのせると「どの構造に属しているパーツなのか」「行数」「列数」というパラメーターを確認できます。下図の「Token Embed」の場合は「Embedding(埋め込み)」を構成しており、行数が埋め込みのサイズであるチャンネル数を示す「C」、列数が語彙(ごい)数を示す「n_vocab」と表示されています。nano-gptは極めて小さいモデルであり、「A」「B」「C」という3つの文字しか扱えないためToken Embedの列数も3列ですむというわけです。

計算によって結果を求める部分にカーソルをあわせると計算式を見ることも可能。「Input Embed」は「Token Embed(トークン埋め込み)」と「Position Embed(位置埋め込み)」を足して求められます。なお、緑のブロックは入力を元に計算して値を求める部分で、青のブロックはトレーニング中に値を調整しておくパラメーターとなる部分とのこと。

構造が入り組んでいる場所では「どの構造に属しているのか」の表示が階層的に表示されます。

より詳しい内部構造を確認したい場合はズームインすればOK。

さまざまな計算を元に次の1文字の予測が「A」と出力されました。この予測を入力にフィードバックすることで次々と予測を行う事ができるというわけです。

サイトでは内部構造それぞれの詳細な解説も読めるので、気になる人は確認してみてください。

なお、解説では約8万5000パラメーターの「nano-gpt」を使用しましたが、nano-gptを1億2400万パラメーターのGPT-2(small)と比較すると下図の通りとなります。

そして約1750億パラメーターのGPT-3の大きさは下図の通り。3D化されているおかげでどれくらいモデルのサイズが異なるのかが一目瞭然です。