Transformers Explained Visually

Transformerとは何か？

Transformerは、人工知能へのアプローチを根本的に変えたニューラルネットワークアーキテクチャです。Transformerは、2017年の画期的な論文「Attention is All You Need」で初めて導入されて以来、深層学習モデルのデファクトスタンダードとなり、OpenAIのGPT、MetaのLlama、GoogleのGeminiといったテキスト生成モデルを支えています。テキスト以外にも、音声生成、画像認識、タンパク質構造予測、さらにはゲームプレイなど、数多くの分野でその汎用性を示しています。

基本的に、テキスト生成Transformerモデルは次のトークン予測の原理に基づいて動作します。ユーザーからのテキストプロンプトが与えられたとき、この入力に続く最も確率の高い次のトークン（単語または単語の一部）は何でしょうか？Transformerの核心的な革新と強力さは、セルフアテンションメカニズムの使用にあり、これにより、以前のアーキテクチャよりも効果的にシーケンス全体を処理し、長距離の依存関係を捉えることができます。

GPT-2モデルファミリーは、テキスト生成Transformerの著名な例です。Transformer Explainerは、GPT-2（small）モデルによって強化されており、これは1億2400万個のパラメータを持っています。最新または最も強力なTransformerモデルではありませんが、現在の最先端モデルに見られる多くのアーキテクチャコンポーネントと原理を共有しており、基本を理解するための理想的な出発点となります。

Transformerアーキテクチャ

テキスト生成Transformerは、以下の3つの主要コンポーネントで構成されています。

1. Embedding: テキスト入力は、トークンと呼ばれるより小さな単位（単語またはサブワード）に分割されます。これらのトークンは、単語の意味的な意味を捉える埋め込み（embedding）と呼ばれる数値ベクトルに変換されます。
2. Transformer Block: モデルの基本的な構成要素であり、入力データを処理・変換します。各ブロックには以下が含まれます。
   • Attention Mechanism: Transformerブロックのコアコンポーネントです。トークン同士が他のトークンと通信することを可能にし、文脈情報と単語間の関係を捉えます。
   • MLP (Multilayer Perceptron) Layer: 各トークンに独立して作用するフィードフォワードネットワークです。アテンション層の目的がトークン間の情報ルーティングであるのに対し、MLPの目的は各トークンの表現を洗練することです。
3. Output Probabilities: 最終的な線形層とソフトマックス層が、処理された埋め込みを確率に変換し、シーケンス内の次のトークンを予測できるようにします。

Embedding

Transformerモデルを使用してテキストを生成したいとしましょう。次のようなプロンプトを入力します: “Data visualization empowers users to”。この入力は、モデルが理解し処理できる形式に変換される必要があります。そこで埋め込み（embedding）が登場します。これはテキストをモデルが処理できる数値表現に変換します。プロンプトを埋め込みに変換するには、1）入力をトークン化し、2）トークン埋め込みを取得し、3）位置情報を追加し、最後に4）トークンと位置エンコーディングを加算して最終的な埋め込みを得る必要があります。それぞれのステップがどのように行われるか見てみましょう。

［画像: Figure 1. Embedding層のビューを拡大し、入力プロンプトがベクトル表現にどのように変換されるかを示しています。プロセスには、(1) Tokenization、(2) Token Embedding、(3) Positional Encoding、(4) Final Embedding が含まれます。］

ステップ1: トークン化

トークン化とは、入力テキストをトークンと呼ばれる、より小さく管理しやすい断片に分割するプロセスです。これらのトークンは単語またはサブワードになります。"Data"と"visualization"は一意のトークンに対応しますが、"empowers"という単語は2つのトークンに分割されます。トークンの完全な語彙は、モデルのトレーニング前に決定されます。GPT-2の語彙には50,257個の一意のトークンがあります。入力テキストを個別のIDを持つトークンに分割したので、埋め込みからそれらのベクトル表現を取得できます。

ステップ2. トークン埋め込み

GPT-2（small）は、語彙の各トークンを768次元のベクトルで表現します。ベクトルの次元はモデルによって異なります。これらの埋め込みベクトルは、形状が(50,257, 768)の行列に格納されており、約3900万個のパラメータを含んでいます！この広範な行列により、モデルは各トークンに意味的な意味を割り当てることができます。つまり、言語での使用法や意味が似ているトークンはこの高次元空間で互いに近くに配置され、意味が異なるトークンはより離れた場所に配置されます。

ステップ3. 位置エンコーディング

埋め込み層は、入力プロンプトにおける各トークンの位置に関する情報もエンコードします。異なるモデルは、位置エンコーディングのためにさまざまな方法を使用します。GPT-2は、位置エンコーディング行列をスクラッチから学習し、トレーニングプロセスに直接統合します。

ステップ4. 最終埋め込み

最後に、トークンエンコーディングと位置エンコーディングを加算して、最終的な埋め込み表現を取得します。この結合された表現は、トークンの意味的な意味と入力シーケンス内の位置の両方を捉えます。

Transformer Block

Transformerの処理の中核はTransformerブロックにあり、これはマルチヘッドセルフアテンションと多層パーセプトロン（MLP）層で構成されています。ほとんどのモデルは、このような複数のブロックが順番に積み重ねられた構造を持っています。トークン表現はレイヤーを通過するにつれて進化し、最初のブロックから最後のブロックまで、モデルが各トークンの複雑な理解を構築できるようにします。このレイヤー化されたアプローチは、入力の高次表現につながります。検討しているGPT-2（small）モデルは、このような12個のブロックで構成されています。

マルチヘッドセルフアテンション

セルフアテンションメカニズムにより、モデルはシーケンス内のトークン間の関係を捉えることができ、各トークンの表現は他のトークンの影響を受けます。複数のアテンションヘッドにより、モデルはこれらの関係を異なる視点から考慮できます。たとえば、1つのヘッドは短い構文リンクを捉え、別のヘッドはより広範な意味的コンテキストを追跡する可能性があります。次のセクションでは、マルチヘッドセルフアテンションがステップバイステップでどのように計算されるかを説明します。

ステップ1: クエリ、キー、バリュー行列

［画像: Figure 2. 元の埋め込みからクエリ、キー、バリュー行列を計算しています。］

各トークンの埋め込みベクトルは、クエリ（Q）、キー（K）、バリュー（V）の3つのベクトルに変換されます。これらのベクトルは、入力埋め込み行列をQ、K、Vの学習済み重み行列で乗算することによって派生します。これらの行列の背後にある直感を構築するのに役立つWeb検索の例えを以下に示します。

• クエリ (Q): 検索エンジンバーに入力する検索テキストです。これは、「より多くの情報を検索したい」トークンです。
• キー (K): 検索結果ウィンドウ内の各Webページのタイトルです。これは、クエリがアテンドできる可能性のあるトークンを表します。
• バリュー (V): 表示されるWebページの実際のコンテンツです。適切な検索語（クエリ）と関連性の高い結果（キー）を一致させた後、最も関連性の高いページのコンテンツ（バリュー）を取得したいのです。

これらのQKV値を使用することで、モデルはアテンションスコアを計算できます。これは、各トークンが予測を生成する際にどれだけ焦点を当てるべきかを決定します。

ステップ2: マルチヘッド分割

クエリ、キー、バリューベクトルは複数のヘッドに分割されます。GPT-2（small）の場合は12個のヘッドに分割されます。各ヘッドは埋め込みのセグメントを独立して処理し、異なる構文的および意味的関係を捉えます。この設計は、多様な言語的特徴の並列学習を促進し、モデルの表現力を高めます。

ステップ3: マスク付きセルフアテンション

各ヘッドでは、マスク付きセルフアテンション計算を実行します。このメカニズムにより、モデルは将来のトークンへのアクセスを防ぎながら、入力の関連部分に焦点を当てることでシーケンスを生成できます。

［画像: Figure 3. クエリ、キー、バリュー行列を使用してマスク付きセルフアテンションを計算しています。］

• ドット積: クエリとキー行列のドット積がアテンションスコアを決定し、すべての入力トークン間の関係を反映する正方行列を生成します。
• スケーリング・マスク: アテンションスコアはスケーリングされ、アテンション行列の上三角にマスクが適用され、モデルが将来のトークンにアクセスするのを防ぎ、これらの値を負の無限大に設定します。モデルは、未来を「覗き見」することなく次のトークンを予測する方法を学習する必要があります。
• ソフトマックス・ドロップアウト: マスキングとスケーリングの後、アテンションスコアはソフトマックス演算によって確率に変換され、その後、ドロップアウトでオプションで正則化されます。行列の各行は合計で1になり、その左側にある他のすべてのトークンとの関連性を示します。

ステップ4: 出力と連結

モデルはマスク付きセルフアテンションスコアを使用し、それをバリュー行列と乗算して、セルフアテンションメカニズムの最終出力を取得します。GPT-2には12個のセルフアテンションヘッドがあり、それぞれがトークン間の異なる関係を捉えます。これらのヘッドの出力は連結され、線形射影を通過します。

MLP: 多層パーセプトロン

［画像: Figure 4. MLP層を使用して、セルフアテンション表現をより高次元に射影し、モデルの表現能力を高めています。］

複数のセルフアテンションヘッドが入力トークン間の多様な関係を捉えた後、連結された出力は多層パーセプトロン（MLP）層を通過し、モデルの表現能力を高めます。MLPブロックは、間にGELU活性化関数を持つ2つの線形変換で構成されています。

最初の線形変換は、入力の次元を768から3072へと4倍に拡張します。この拡張ステップにより、モデルはトークン表現をより高次元空間に射影でき、元の次元では見えない、より豊かで複雑なパターンを捉えることができます。

次に、2番目の線形変換が次元を元のサイズである768に戻します。この圧縮ステップは、拡張ステップで導入された有用な非線形変換を保持しながら、表現を管理しやすいサイズに戻します。

トークン間で情報を統合するセルフアテンションメカニズムとは異なり、MLPはトークンを独立して処理し、各トークン表現をある空間から別の空間に単純にマッピングするだけで、モデル全体の容量を豊かにします。

出力確率

入力がすべてのTransformerブロックを通過して処理された後、出力は最終線形層を通過し、トークン予測の準備が整います。この層は、最終表現を50,257次元の空間に射影し、語彙の各トークンに対応するlogitと呼ばれる値が割り当てられます。どのトークンも次の単語になり得るため、このプロセスにより、これらのトークンを次の単語になる可能性によって単純にランク付けできます。次にソフトマックス関数を適用して、logitsを合計が1になる確率分布に変換します。これにより、その可能性に基づいて次のトークンをサンプリングできます。

［画像: Figure 5. モデルの出力logitsに基づいて、語彙の各トークンに確率が割り当てられます。これらの確率は、シーケンス内の次の単語となる各トークンの可能性を決定します。］

最後のステップは、この分布からサンプリングして次のトークンを生成することです。temperatureハイパーパラメータは、このプロセスで重要な役割を果たします。数学的には非常に単純な操作です。モデル出力のlogitsをtemperatureで割るだけです。

• temperature = 1: logitsを1で割っても、ソフトマックス出力には影響しません。
• temperature < 1: 低い温度は、確率分布をシャープにすることで、モデルをより自信があり決定論的にし、より予測可能な出力を生成します。
• temperature > 1: 高い温度は、よりソフトな確率分布を作成し、生成されるテキストに多くのランダム性を許容します。これは、モデルの「創造性」と呼ばれることもあります。

さらに、サンプリングプロセスは、top-kおよびtop-pパラメータを使用してさらに洗練させることができます。

• top-k sampling: 確率が最も高い上位k個のトークンに候補トークンを制限し、可能性の低いオプションを除外します。
• top-p sampling: 累積確率が閾値pを超える最小のトークンセットを考慮し、最も可能性の高いトークンのみが寄与するようにしながらも、多様性を許容します。

temperature、top-k、top-pを調整することで、決定論的出力と多様な出力のバランスを取り、モデルの動作を特定のニーズに合わせて調整できます。

補助的なアーキテクチャ機能

Transformerモデルのパフォーマンスを向上させるいくつかの補助的なアーキテクチャ機能があります。これらはモデル全体のパフォーマンスにとって重要ですが、アーキテクチャのコアコンセプトを理解する上ではそれほど重要ではありません。Layer Normalization、Dropout、Residual Connectionsは、Transformerモデル、特にトレーニングフェーズにおいて重要なコンポーネントです。Layer Normalizationはトレーニングを安定させ、モデルの収束を速めます。Dropoutは、ニューロンをランダムに無効にすることで過学習を防ぎます。Residual Connectionsは、勾配がネットワークを直接流れることを可能にし、勾配消失問題を防ぐのに役立ちます。

Layer Normalization

Layer Normalizationは、トレーニングプロセスを安定させ、収束を改善するのに役立ちます。これは、特徴全体で入力を正規化することによって機能し、活性化の平均と分散が一貫していることを保証します。この正規化は、内部共変量シフトに関連する問題を軽減し、モデルがより効果的に学習できるようにし、初期重みへの感度を低下させます。Layer Normalizationは、セルフアテンションメカニズムの前とMLP層の前という、各Transformerブロックで2回適用されます。

Dropout

Dropoutは、トレーニング中にモデルの重みの一部をランダムにゼロに設定することにより、ニューラルネットワークの過学習を防ぐために使用される正則化技術です。これにより、モデルはより堅牢な特徴を学習し、特定のニューロンへの依存を減らし、ネットワークが新しい、未知のデータにうまく汎化するのを助けます。モデル推論中、Dropoutは無効になります。これは本質的に、トレーニングされたサブネットワークのアンサンブルを使用していることを意味し、モデルパフォーマンスの向上につながります。

Residual Connections

Residual connectionsは、2015年にResNetモデルで初めて導入されました。このアーキテクチャ上の革新は、非常に深いニューラルネットワークのトレーニングを可能にすることで、深層学習に革命をもたらしました。本質的に、Residual connectionsは1つ以上のレイヤーをバイパスするショートカットであり、レイヤーの入力をその出力に追加します。これは勾配消失問題を軽減するのに役立ち、複数のTransformerブロックが積み重ねられた深いネットワークのトレーニングを容易にします。GPT-2では、Residual connectionsは各Transformerブロック内で2回使用されます。MLPの前と後で、勾配がより簡単に流れ、バックプロパゲーション中に初期レイヤーが十分な更新を受け取ることを保証します。

インタラクティブ機能

Transformer Explainerはインタラクティブに構築されており、Transformerの内部動作を探求できます。以下は、操作できるインタラクティブ機能の一部です。

• 独自のテキストシーケンスを入力して、モデルがそれをどのように処理し、次の単語を予測するかを確認します。アテンションウェイト、中間計算を探索し、最終出力確率がどのように計算されるかを確認します。
• 温度スライダーを使用して、モデルの予測のランダム性を制御します。温度値を変更することで、モデルの出力をより決定論的またはより創造的にする方法を探索します。
• top-kおよびtop-pサンプリングメソッドを選択して、推論中のサンプリング動作を調整します。さまざまな値を試して、確率分布がどのように変化し、モデルの予測に影響するかを確認します。
• アテンションマップと対話して、モデルが入力シーケンスの異なるトークンにどのように焦点を当てるかを確認します。トークンにカーソルを合わせるとアテンションウェイトがハイライトされ、モデルがコンテキストと単語間の関係をどのように捉えるかを探索します。

ビデオチュートリアル

Transformer Explainerはどのように実装されていますか？

Transformer Explainerは、ブラウザで直接動作するライブGPT-2（small）モデルを搭載しています。このモデルは、Andrej KarpathyのnanoGPTプロジェクトのPyTorch実装から派生し、ブラウザでのシームレスな実行のためにONNX Runtimeに変換されています。インターフェースはJavaScriptを使用して構築されており、フロントエンドフレームワークとしてSvelte、動的な視覚化の作成にはD3.jsが使用されています。ユーザー入力に続いて、数値はリアルタイムで更新されます。

Transformer Explainerは誰が開発しましたか？

Transformer Explainerは、ジョージア工科大学のAeree Cho、Grace C. Kim、Alexander Karpekov、Alec Helbling、Jay Wang、Seongmin Lee、Benjamin Hoover、およびPolo Chauによって作成されました。