アテンションメカニズム：腫瘍微小環境における文脈的な判断

このセクションを終えて

セクション#5でマークし、後回しにしたTransformerの中核であるアテンションについて深く掘り下げていきましょう。なぜ必要なのか？ 何をするのか？ そして、最新のLLMが処理する12.8万トークンの文脈は、どのように計算されるのでしょうか？

このセクションは、物語のクライマックスです。セクション#1～#5のコンポーネント（次の単語予測、3362万個のパラメータ、勾配降下法、バックプロパゲーション、埋め込み、残差接続、正規化）がアテンションと組み合わさり、Transformerが完成します。セクション#7以降で扱うプロンプトエンジニアリング、RAG、エージェントも、アテンションの特性に根ざしています。

セクション#1の「肝臓」問題に戻る

セクション#1の多義性の例を思い出してみましょう。

• 「私の肝臓が痛い」- 体の臓器（肝）
• 「味付けがちょうどいい」- 塩味の度合い
• 「3年が経った」- 時間の経過

セクション#5で見たように、「肝臓」という単語の埋め込みは、固定された単一のベクトルです。しかし、このベクトルは上記の3つの文で異なる意味で使用されています。訓練されたニューラルネットワークは、これらの3つのケースをどのように区別し、処理するのでしょうか？

その方法は次のとおりです。文脈内の他の単語が、この「肝臓」の表現に流れ込み、文脈的に適切な意味に変換されるのです。もし「痛い」という単語があれば、体の臓器の意味に傾き、「味付け」という単語があれば、塩味の意味に傾き、「3年」という単語があれば、時間の経過の意味に傾きます。この流れを担うのがアテンションです。

言い換えれば、アテンションはTransformer内で、単語の表現を文脈化するプロセスです。埋め込みベクトルが静的で事前に存在する表現であるとすれば、アテンション後のベクトルは、文中の生きた意味となります。

それでは、アテンションがどのようにこれを行うのかを学びましょう。

単純なアプローチの失敗：一様平均の問題

単語間の情報を伝える最も簡単な方法は、文中のすべてのトークンベクトルを単純に平均することです。

contextualized(肝臓) = (vec(肝臓) + vec(私) + vec(痛い) + vec(".")) / 4

問題点：すべての単語が等しい重みで扱われることです。「痛い」は「肝臓」の意味を決定する上で非常に重要ですが、「私」はそれほど重要ではありません。しかし、一様平均では、これらの違いを区別しません。

必要なのは：各単語が、状況に応じてどれだけ反映されるかを、差分的に重みづけする方法です。例えば、「痛い」80%、「3年」5%、「私」5%を「肝臓」の意味を決定するために使用します。

アテンションは、この「どれだけ反映するか」を計算します。各単語ペア（クエリ-キー）に対してアテンションの重みを計算し、それを使用して情報の重み付き平均を計算します。

病理医が病理スライド内のT細胞を判断する

セクション#2で見た病理スライドの例を用いて、生物学的なアナロジーを続けましょう。

あなたが病理スライド内の単一のT細胞を観察しているとします。このT細胞がどのような状態にあるのか（活性化しているか、疲弊しているか、または制御性T細胞であるか）を判断する必要があります。

この判断は、T細胞自体の形態だけを見て行うことはできません。周囲の細胞環境（腫瘍微小環境）を見る必要があります。

• このT細胞のすぐ隣に多くの腫瘍細胞がある場合？ → 腫瘍を攻撃しているか、または無効化されている状態である可能性が高い。
• 周囲に多くの他のT細胞と樹状細胞がある場合？ → ここは免疫応答が活性化されている場所です。
• 周囲が線維芽細胞とコラーゲンマトリックスによって厚くなっている場合？ → これは基質バリア内に閉じ込められていることを意味します。
• **制御性T細胞（Treg）**の近くにある場合？ → 抑制されている可能性が高い。

この判断において、周囲の各細胞の重要度は異なります。腫瘍細胞の密度は非常に重要ですが、基質の厚さはそれほど重要ではありません。病理医は無意識のうちに、各周囲の細胞に異なる重みを割り当て、判断を下します。

これを定量化することで、アテンションの正確な構造を導き出すことができます。3つのコンポーネントが必要です。

• クエリ： 判断されているT細胞が「何を」知りたいのか？ 例えば、「私はCD8+ T細胞であり、現在腫瘍を攻撃しているのか、または疲弊しているのかを知りたい」
• キー： 周囲の各細胞が「どのような情報」を提供できるのか？ 例えば、腫瘍細胞は「PD-L1発現レベル」を提供でき、樹状細胞は「抗原提示状態」を提供でき、基質細胞は「物理的バリアの程度」を提供できます。
• バリュー： 実際に伝達される情報のコンテンツ。 例えば、腫瘍細胞のPD-L1発現レベル、樹状細胞の活性化状態など。

現在のT細胞のクエリは、周囲の各細胞のキーと比較され、「どの程度一致するか」がスコア化されます。一致度が高い細胞のバリューは、このT細胞の判断に強く反映され、一致しない細胞は無視されます。

これがアテンションの骨格です。

クエリ、キー、バリュー：言葉の三つの顔

Transformerのアテンションも、まさにこの三角形の構造を持っています。

文中の各単語ベクトルは、三つの派生ベクトルに変換されます。

• クエリ (Q): 「この単語は、文脈の中で何を知りたいのか？」
• キー (K): 「この単語は、他の単語に対してどのように自己を宣伝するのか？」
• バリュー (V): 「この単語は、実際にどのような情報を伝達するのか？」

それぞれは、単語の埋め込みベクトルに、学習された三つの重み行列 W_Q、W_K、W_V を掛けることで得られます。

Q_i = W_Q · x_i
K_i = W_K · x_i
V_i = W_V · x_i

x_i は、i番目の単語の埋め込みベクトルです。これらの三つのベクトルは、通常、元のベクトルよりも次元数が小さくなります（例：元の4096次元 → Q/K/V それぞれ128次元）。

直感: 単語ベクトルは、三つの「顔」を持っていると考えることができます。クエリの顔は他の単語に「あなたは私にとって役に立ちますか？」と質問します。キーの顔は他の単語からの質問に「私はこの情報を持っています」と答えます。バリューの顔は実際の情報を伝達します。

Q、K、V の三つの顔が異なることは重要です。単語が自己を宣伝する方法（キー）と、他の単語を調査する方法（クエリ）は異なり、実際に伝達する情報（バリュー）も、この二つとは異なる場合があります。この分離こそが、アテンションが非常に柔軟な情報フローを表現することを可能にしています。

スケール付きドット積アテンションの公式

それでは、実際の計算を見てみましょう。単語 i が単語 j に対して計算する アテンションスコア は次のとおりです。

score(i, j) = ⟨Q_i, K_j⟩ / sqrt(d_k)

これは、二つのベクトルのドット積です。二つのベクトルの方向が似ている場合、スコアは高くなり、異なる場合はスコアは低くなります。

sqrt(d_k) で割る理由。QとKは d_k 次元のベクトルですが、ベクトルがランダムな場合、d_k が大きくなるにつれて、ドット積の大きさは、おおよそ sqrt(d_k) に比例して増加します。この除算を行わないと、スコアはベクトル次元の増加とともに発散し、その後のソフトマックスが極端な値になります。このスケーリングは、Transformerのトレーニング安定性の隠れた立役者です。正確な理由は、付録A.2で導き出すことができます。

次に、単語 i の観点から、文中のすべての単語に対してスコアを計算しました。これらのスコアは、ソフトマックス を使用して確率分布に変換されます。

α(i, j) = exp(score(i, j)) / Σ_k exp(score(i, k))

Σ_j α(i, j) = 1。各単語は、自分のアテンションの100％を他の単語に割り当てます。

ソフトマックスは、セクション#1の付録A.2とセクション#3の付録A.4で既に見た関数です。ここでは、単語間の競争を表す確率分布を作成します。

最後に、バリューベクトルは、この重みを使用して重み付けされ、平均化されます。

output_i = Σ_j α(i, j) · V_j

これは、単語 i のアテンションに基づく表現です。

これを文全体に対して行列形式で記述すると、はるかに簡潔になります。

Attention(Q, K, V) = softmax(Q · K^T / sqrt(d_k)) · V

この一行が、Transformerの核心です。オリジナルの論文のタイトル「Attention Is All You Need」は、この方程式を指しています。

マルチヘッドアテンション：複数の角度から並行して判断

単一のアテンションメカニズムでも情報フローは可能ですが、実際には常に 複数のアテンションを並行して 使用します。これは マルチヘッドアテンション と呼ばれます。

各ヘッドは、独自の W_Q、W_K、W_V を持ち、独立してアテンションを計算します。結果は連結され、その後、単一の出力重み行列 W_O を使用して再度結合されます。

head_h = Attention(Q · W_Q^h, K · W_K^h, V · W_V^h)
MultiHead = Concat(head_1, ..., head_H) · W_O

H はヘッドの数です。GPT-2 medium では16、GPT-3 では96、LLaMA-70B では64です。

複数のヘッドが必要なのはなぜですか？ 各ヘッドは、特定の「視点」に特化しています。たとえば、あるヘッドは「前の単語」に焦点を当て、別のヘッドは「文中のより前の名詞」に焦点を当て、別のヘッドは「対応する括弧」に焦点を当てる場合があります。複数のヘッドを使用することで、さまざまな文法および意味関係を同時に処理できます。

生物学的アナロジー：複数のマーカーを並行して判断する。病理学者がT細胞を評価する場合、単一の特性だけを見るわけではありません。CD8/CD4マーカー、活性化マーカー（CD69、HLA-DR）、消耗マーカー（PD-1、TIM-3、LAG-3）、および制御マーカー（FoxP3）を同時に並行して調べ、結果を組み合わせて判断します。各マーカーの検査は、1つのアテンションヘッドに対応し、最終的な組み合わせは W_O の組み合わせに対応します。この並行視点のアナロジーが、マルチヘッドアテンションの鍵です。

解釈可能性: Transformerのトレーニングされたアテンションヘッドを調べた研究では、異なるヘッドが実際に異なる言語パターンを処理することが示されています。たとえば、GPT-2のいくつかのヘッドは「前の単語」にのみ焦点を当て、他のヘッドは「同じ名詞の前の出現」を追跡します。最近のメカニスティックな解釈可能性の研究では、これらのヘッドごとの関数がカタログ化されています。

因果的マスク処理：未来を見ることができないという問題

LLMは次の単語を予測する（セクション#1）ことで訓練されます。この際に重要な制約として、まだ予測されていない未来の単語を参照してはならないという点があります。

たとえば、「猫は魚が好きです」という文を見た場合、次の単語である「に」を参照してはなりません。訓練は常に、「これまでに見られたもの」に基づいて次の単語を予測する必要があります。

アテンションは、原理的には文中のすべての単語を参照できる構造です。したがって、未来を参照しないようにする必要がある場合、それを明示的にブロックする必要があります。これを因果的マスク処理と呼びます。

方法は以下の通りです。スコア行列の上三角部分（単語iに対してj > iとなる箇所）を-∞に設定します。ソフトマックス関数を通過すると、-∞は確率0になるため、未来の単語は無視されます。

score(i, j) = ⟨Q_i, K_j⟩ / sqrt(d_k)  if j ≤ i
score(i, j) = -∞                      if j > i

マスク処理を施したこのアテンションは、因果的アテンションまたはデコーダーアテンションと呼ばれます。GPTのようなモデルはこの手法を使用します。マスク処理を施さないアテンションは、双方向アテンションと呼ばれます。BERTのようなモデルはこの手法を使用します。ほとんどのLLMは因果的です。

生物学的なアナロジー: これは少し微妙ですが、考えられるアナロジーです。発生生物学では、細胞が分化の決定を行う際、まだ発生していない下流の細胞からの情報を使用することはできません。それは、時間的な順序で上流の細胞からのシグナルのみを反映します。因果的マスク処理は、この時間的な因果関係を強制します。

自己アテンションとクロスアテンション

これまで、Q、K、Vがすべて文中の同じトークンのセットから得られるアテンションメカニズムについて議論してきました。これは自己アテンションと呼ばれます。これは、トークンが文の中で情報を引き出す構造です。

クロスアテンションは、Qが1つのセットから、KとVが別のセットから得られる場合です。たとえば、翻訳モデルでは、ターゲット言語のトークンがソース言語のトークンを参照する場合などです。あるいは、画像キャプション生成では、テキストのトークンが画像のパッチを参照する場合などです。これは、マルチモーダルモデルの基本的な構成要素です。

• 自己アテンション: Q、K、Vがすべて同じシーケンスから得られます。ほとんどのLLMはこの手法を使用します。
• クロスアテンション: Qが1つのシーケンスから、KとVが別のシーケンスから得られます。マルチモーダルモデルと翻訳モデルで使用されます。

元のTransformer論文のエンコーダー・デコーダーアーキテクチャでは、自己アテンションとクロスアテンションの両方を使用します。純粋なデコーダーLLM（GPTファミリーなど）は、自己アテンションのみを使用します。最近では、クロスアテンションが再び注目を集め、ビジョン-言語モデル（GPT-4VやClaude Visionなど）で使用されています。

KVキャッシュ — 推論速度を100倍に向上させる

訓練から推論に移行すると、興味深い最適化手法が登場します。

LLMがテキストを生成する場合、以下の手順で進みます。

• プロンプトを入力：「猫は」。
• モデルは次のトークンを予測：「食べる」。
• プロンプトは「猫は食べる」に更新され、モデルに再入力されます。
• 次の予測：「魚」。
• プロンプトは「猫は魚を食べる」に更新され、モデルに再入力されます。
• ...

各ステップでプロンプト全体を再計算すると、ステップtでの計算コストはO(t^2)になります。100個のトークンを生成するには、10,000単位の計算が必要になります。

しかし、前のステップで計算されたKとVの値は、次のステップでも同じです。因果的マスク処理のおかげで、過去のトークンのQ、K、Vの値は、新しいトークンが追加されても変化しません。したがって、これらの値をキャッシュすることができます。

これがKVキャッシュです。各レイヤーのKとVベクトルはGPUメモリに保存され、各ステップで新しいトークンのQ値のみを計算し、キャッシュされたKとVの値とともにアテンションメカニズムで使用します。これにより、1ステップあたりの計算コストがO(t)に削減されます。100個のトークンを生成するには、1,000単位の計算で済み、100倍の速度向上が実現します。

トレードオフ。GPUメモリを大量に消費します。96層 × 96ヘッド × 12288次元 × 128Kコンテキスト × バッチサイズ... これだけで簡単に数百GBに達する可能性があります。このため、大規模モデルで長いコンテキストウィンドウを使用することが、GPUメモリによって制限されます。

このKVキャッシュ最適化は、ChatGPTやClaudeが今日のリアルタイムな会話を実現できる重要な理由の一つです。キャッシュがない場合、応答には数十秒かかってしまいます。

GQAとMQA — KVキャッシュサイズを削減するための進化

KVキャッシュが非常に大きいため、新しいアイデアが登場しています。その一つは、複数のクエリヘッドが単一のキー/バリューヘッドを共有することです。

• マルチヘッドアテンション（MHA）: Q、K、VはそれぞれH個のヘッドを持ちます。これは基本的な形式です。
• マルチクエリアテンション（MQA）: QはH個のヘッドを持ちますが、KとVはそれぞれ1個のヘッドを持ちます。これは極端なアプローチであり、KVキャッシュサイズをHの係数で削減します。
• グループ化クエリアテンション（GQA）: 中間のアプローチです。QはH個のヘッドを持ち、KとVはそれぞれG個のヘッドを持ちます（G < H）。たとえば、H=32、G=8の場合、4つのQヘッドが1つのK/Vヘッドを共有します。

MQAはキャッシュ削減を最大化しますが、品質が低下する可能性があります。GQAは、キャッシュ削減を大幅に行いながら、ほとんど同じ品質を維持する妥協案です。最新のモデル（LLaMA-2、LLaMA-3、Qwen、Mistralなど）はGQAを使用しています。これが現在、標準となっています。

Flash Attention — 実装によるメモリの壁の打破

アテンション計算におけるボトルネックは、実際の計算量ではなく、GPUメモリの帯域幅です。Q · K^Tはn × nの行列です（ここでnはシーケンス長です）。したがって、128Kのコンテキストでは、128000 × 128000の行列が必要になります。この大きな行列をHBM（GPUメモリ）とSRAM（オンチップキャッシュ）の間でやり取りするのにかかる時間は、実際の計算よりもはるかに長くなります。

Flash Attention (Dao et al., 2022) は、このデータのやり取りを最小限に抑える実装です。大きなアテンション行列をブロックに分割し、HBMに大きな行列を格納することなく、SRAM内で処理します。数学的な結果は同じですが、2倍から10倍高速化され、メモリ使用量がO(n²)からO(n)に削減されます。

Flash Attentionは、128K、200K、1Mトークンのコンテキストウィンドウが現実的になった主な理由です。現在、ほぼすべての大規模モデルは、トレーニングと推論の両方でFlash Attentionを使用しています。Flash Attention 2とFlash Attention 3によって、さらに改良が続けられています。

重要なポイント。アルゴリズム理論だけでなく、ハードウェアを意識した実装の最適化こそが、今日のLLMの拡大を推進しているのです。これは、深層学習システムのエンジニアリングの重要性を強調しています。

アテンションを超えて — 次のトークン予測へ

これまでに説明した要素を振り返りましょう。アテンションブロックはトークンベクトルにコンテキストを与え、次に#5で説明したFFN（フィードフォワードネットワーク）が、このコンテキスト化された表現をさらに処理します。いくつかのブロックを通過した後、言語モデリングヘッド（埋め込みに似たサイズのd × Vの行列）が、次のトークンの確率分布を生成します。

トレーニング中、この分布と実際の次のトークンの間のクロスエントロピーが損失として使用されます。#3と#4のアテンションパラメータへの逆伝播により、学習が可能です。その結果、アテンションは自然に「言語をコンテキスト化するための最も有用な方法」を学習します。

アテンション行列の解釈。トレーニングされたモデルのアテンションスコアα(i, j)を視覚化すると、特定の言語パターンが明らかになります。例えば：

• 前のトークンヘッド: スコアは対角線の下に集中しています。文法的な隣接性。
• 前の名詞参照ヘッド: スコアは代名詞から前の名詞を指しています。共参照解決。
• 文末ヘッド: スコアは文末トークンから文の先頭を指しています。文構造の要約。

この視覚化は、メカニスティックな解釈可能性研究の出発点となります。Anthropic、OpenAI、Googleは、アテンションパターン分析を使用して、LLMの内部動作を逆エンジニアリングしています。

生物学的応用シナリオ

シナリオ1 — ESMのアミノ酸接触予測

#4と#5で言及したESMプロテイン言語モデルのアテンションメカニズムは、注目すべき特性を示します。トレーニングされたESMのアテンション行列を調べると、アテンションが自然に3D構造で物理的に接触しているアミノ酸のペアに強く集中していることがわかります。

これは、モデルが「構造を予測する」ように明示的にトレーニングされていなくても、次のアミノ酸を予測するだけで、3D接触を自動的に学習することを意味します。このアテンションマップは、AlphaFoldの初期段階で接触予測ツールとして使用されました。

これは、アテンションの一般的な能力を示す事例です。「関連する情報に注意を払う」という原則は、言語だけでなく、生物学的シーケンスにも適用されます。

シナリオ2 — Enformerのエンハンサー・プロモーターの連携

#5で言及したゲノムTransformerであるEnformerのアテンションメカニズムは、100kbのコンテキストウィンドウ内で、数kb離れたプロモーター（遺伝子開始部位）とエンハンサー（調節要素）をリンクします。トレーニング後、一部のアテンションヘッドは、エンハンサー・プロモーターペアを自動的に識別することを学習します。

従来、エンハンサー・プロモーターのマッチングには、3Dゲノム接触実験（Hi-C）が必要でしたが、Enformerはシーケンスのみに基づいてこの関係を予測できます。これは、アテンションの「長距離相互作用を学習する能力」によるものです。

シナリオ3 — 病理画像の細胞コンテキスト判断

これは、#2と#6のシナリオの実装例です。Vision Transformer（ViT）ベースの病理モデルは、スライド内の各パッチ（小さな正方形領域）に対してアテンションを計算し、腫瘍の微小環境を決定します。各T細胞パッチは、周囲の腫瘍細胞パッチおよび間質細胞パッチとのアテンション関係を確立し、これらを組み合わせてT細胞の状態を予測します。

このファミリーのモデル（例：HIPT、CTransPath）は、精密な腫瘍診断と予後予測において、病理医レベルのパフォーマンスを発揮し始めています。

重要なポイント

• アテンションは、トークン表現にコンテキストを与えるプロセスです。これは、#1の「バンク」の多義性の処理が行われる場所です。
• 各トークンには、Q、K、Vの3つの側面があります。QとKのドット積は、トークン間の「アテンション」を計算し、次にsoftmaxを使用して正規化し、Vの重み付け平均として情報を渡します。
• 公式：Attention(Q, K, V) = softmax(Q · K^T / sqrt(d)) · V。
• マルチヘッドアテンションは、複数の視点から並行して情報を処理します。これは、病理医のマルチマーカーパネル分析に似ています。
• 因果的マスキングは、未来の情報を使用することを防ぎます。これは、LLMのトレーニングに不可欠です。
• KVキャッシュは、推論を100倍高速化します。GQAは、キャッシュサイズを削減します。
• Flash Attentionは、128K+のコンテキストウィンドウを実用化します。ハードウェアを意識した実装こそが、今日のLLMの拡大を推進しています。
• アテンションは、言語だけでなく、タンパク質、ゲノム、画像にも適用される一般的な原理です。「関連する情報に注意を払う」という考え方は、ドメイン間で再利用できます。

📐 付録 — 専門家向けの数学的公式

難易度: 非常に難しい 対象読者: 線形代数、確率、数値解析に関する大学院レベルの知識を持つ読者。

A.1 スケーリングされたドット積アテンション：完全な公式

シーケンス長 n、埋め込み次元 d、Q/K 次元 d_k、V 次元 d_v。

入力:

• Q ∈ ℝ^{n × d_k} (クエリ)
• K ∈ ℝ^{n × d_k} (キー)
• V ∈ ℝ^{n × d_v} (値)

出力:

Attention(Q, K, V) = softmax(Q K^T / sqrt(d_k)) V

行ごとに適用されるソフトマックス:

softmax(M)_{ij} = exp(M_{ij}) / Σ_k exp(M_{ik})

各行は確率分布（合計は1）です。

最終的な出力形状: ℝ^{n × d_v}。各行は、対応する単語の文脈に応じた表現です。

A.2 なぜ sqrt(d_k) で割るのか、その導出

Q と K の各要素が、平均0、分散1の確率変数であると仮定します。

Q_i と K_j のドット積：

⟨Q_i, K_j⟩ = Σ_{k=1}^{d_k} Q_{ik} K_{jk}

各項 Q_{ik} K_{jk} が独立である場合、平均0、分散1を持ちます。合計の分散は次のようになります。

Var(⟨Q_i, K_j⟩) = d_k
Std(⟨Q_i, K_j⟩) = sqrt(d_k)

したがって、d_k が大きくなるにつれて、ドット積の大きさも sqrt(d_k) の係数で増加します。d_k = 128 の場合、標準偏差は約11.3です。

これを直接ソフトマックスに入力すると、大きな値が支配的になり、ソフトマックスの出力はほぼ ワンホット分布 になります。ソフトマックスの勾配はほぼ0になり、学習を妨げます。

sqrt(d_k) で割ることで、標準偏差を1に近い状態に保ち、ソフトマックスがより滑らかな分布を出力できるようにします。これにより、勾配が流れ、学習が継続できます。

A.3 ソフトマックスの偏微分

p = softmax(z) の場合：

∂p_i/∂z_j = p_i · (δ_{ij} - p_j)

δ_{ij} はクロネッカーのデルタ（i=jの場合1、それ以外の場合は0）です。

クロスエントロピー損失と組み合わせる:

真の値 y (ワンホット) が与えられた場合、損失は L = -Σ y_i log p_i です。

連鎖律を使用します。

∂L/∂z_j = p_j - y_j

これは、セクションA.4でわかるように、非常に美しい結果です。アテンションの学習における勾配も、この原理を繰り返し利用しています。

A.4 マルチヘッドアテンション：公式

H 個のヘッド。各ヘッド h について：

Q_h = X · W_Q^h,  K_h = X · W_K^h,  V_h = X · W_V^h
head_h = Attention(Q_h, K_h, V_h)

W_Q^h, W_K^h ∈ ℝ^{d × d_k}, W_V^h ∈ ℝ^{d × d_v}。通常、d_k = d_v = d/H です。

連結:

MultiHead(X) = Concat(head_1, ..., head_H) · W_O

W_O ∈ ℝ^{H · d_v × d}。

パラメータ数:

• 各Q、K、V射影は d × d です。したがって、3 d^2 です。
• 出力射影は d^2 です。
• 合計: ブロックあたり 4 d^2 です。

これは、セクションA.10の計算と一致します。

A.5 因果マスクの実装

マスク行列：

M_{ij} = 0        if j ≤ i
M_{ij} = -∞       if j > i

アテンション計算：

Attention_masked(Q, K, V) = softmax((Q K^T / sqrt(d_k)) + M) V

-∞ はソフトマックスで exp(-∞) = 0 になり、未来の位置に確率0を割り当てます。

実際には、数値的な安定性を確保するために、-∞ ではなく大きな負の数（例：-1e9）を使用します。

A.6 KVキャッシュアルゴリズム

推論時:

• ステップ1：Q_1, K_1, V_1 を計算します。[K_1], [V_1] をキャッシュに保存します。
• ステップ2：新しい単語に対して Q_2, K_2, V_2 を計算します。キャッシュに追加します：[K_1, K_2], [V_1, V_2]。
• ステップ t：新しい単語の Q_t のみを計算します。キャッシュされたKとVを使用してアテンションを実行します。

メモリ:

• レイヤーごとのヘッド数：H、ヘッド次元：d_h、コンテキスト長：n、バッチサイズ：B
• 1つのレイヤーのKVキャッシュ：2 · B · n · H · d_h · bytes_per_value

例：LLaMA-70B、H=64, d_h=128, n=128K, B=1, fp16：

2 · 1 · 128000 · 64 · 128 · 2 = 4.19 GB/レイヤー

80レイヤーの場合、これは 335 GB です。これが、長いコンテキストを持つ大規模モデルがメモリによって制限される理由です。

A.7 GQA（グループ化クエリアテンション）

• Qヘッド数：H_Q
• K、Vヘッド数：H_KV
• グループサイズ：g = H_Q / H_KV

同じグループ内のQヘッドは、単一のKヘッドとVヘッドを共有します。

head_i = Attention(Q_i, K_{i // g}, V_{i // g})

KVキャッシュのサイズは、g の係数で削減されます。g = H_Q（すべてのQが1つのKVを共有する場合）の場合、MQAになります。

LLaMA-2-70B：H_Q = 64, H_KV = 8, g = 8。キャッシュは8倍節約されます。

A.8 Flashアテンションブロックアルゴリズム

重要なアイデア: n × n の大きなアテンション行列を、SRAMに全体を格納することなく、タイルを使用してブロックごとに計算します。

オンラインソフトマックス: ソフトマックス計算をブロックごとにストリーミングするためのアルゴリズム。各ブロックは、部分的な最大値と部分的な指数和を維持し、それらを結合します。

ブロックサイズ B_r（行）× B_c（列）：

• Qを B_r チャンクでロードします。
• KとVを B_c チャンクで反復処理します。
• 各（Q_block、K_block）について、部分的なスコアと部分的な指数和を計算します。
• 前のブロックからの結果と結合します（オンラインソフトマックスの公式）。

結果:

• 時間：理論的には、FLOPsは同じですが、実際の処理時間は2〜4倍高速化されます（メモリボトルネックが解決されます）。
• メモリ：O(n²) → O(n)
• 精度：同一（無視できる程度の数値誤差）。

A.9 スパースアテンションシリーズ

長いコンテキストを処理するための近似手法。

• スライディングウィンドウアテンション: 各トークンは、w トークンの近傍のみを参照します。計算量は O(nw) です。MistralとLongformerで採用されています。
• Longformer: スライディングウィンドウ + 少数グローバルなトークン。
• BigBird: スライディング + グローバル + ランダムアテンション。
• スパースアテンション（GPT-3）: 特定のパターン（ストライド、ファクタライズ）を使用してアテンションを実行します。

トレードオフ: 近似であるため、品質がいくらか低下します。Flashアテンションの成功により、現在はそれほど必要ではありません。ただし、非常に長いコンテキスト（1M+）の場合には依然として役立ちます。

A.10 アテンションの計算およびメモリの複雑性

シーケンス長 n、次元 d：

• 完全アテンション: 時間 O(n²d)、メモリ O(n² + nd)
• Flashアテンション: 時間 O(n²d)、メモリ O(n · d)（ブロックサイズのみがSRAMに存在）
• スライディングウィンドウ: 時間 O(nwd)、メモリ O(nd)

n = 128K の場合、n² は160億になります。Flashアテンションにより、この規模の行列の計算が可能になりました。最新の大規模言語モデルの長いコンテキスト機能は、このアルゴリズムとシステムコラボレーションの成果です。

参考文献

このセクションのすべてのコンテンツ、シナリオ、アナロジー、および図は、BioPlaygroundによって社内で作成されました。以下は、これらの概念を学習するのに役立つ可能性のある外部の参考文献です。

• オリジナルのTransformer論文（アテンションを定義）: Vaswani et al., "Attention Is All You Need" (NeurIPS 2017)
• Flash Attention: Dao et al., "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness" (NeurIPS 2022)
• Flash Attention 2: Dao, "FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning" (2023)
• Multi-Query Attention: Shazeer, "Fast Transformer Decoding: One Write-Head is All You Need" (2019)
• Grouped-Query Attention: Ainslie et al., "GQA: Training Generalized Multi-Query Transformer Models" (EMNLP 2023)
• アテンションヘッドの解釈: Elhage et al., "A Mathematical Framework for Transformer Circuits" (Anthropic 2021)
• アテンションヘッドカタログ: Olsson et al., "In-context Learning and Induction Heads" (Anthropic 2022)
• ESMコンタクト予測: Rao et al., "Transformer protein language models are unsupervised structure learners" (ICLR 2021)
• Vision Transformer: Dosovitskiy et al., "An Image is Worth 16x16 Words" (ICLR 2021)
• 深層学習の可視化チュートリアル: 3Blue1Brown "Deep Learning" Ch 6 & 7 (YouTube) — 教育的な参考資料として。

これで、原則セクションであるフェーズ1は終了しました（セクション6）。セクション7以降では、このトレーニングされたTransformerを実際にどのように活用するか、つまりプロンプトエンジニアリング、RAG、エージェントについて説明します。

次の概念

• Ep. #7 prompt-engineering — トレーニングされたTransformerのアテンションを効果的に誘導するように設計されたプロンプト。
• Ep. #8 rag-and-context — アテンションが処理できる情報を拡張するRAG。
• Ep. #11 hallucination-and-alignment — アテンションの失敗が、どのようにして幻覚につながるか。