チャプター 04 · アテンション · 12 min

アテンションこそすべて

すべてを変えるメカニズム。各トークンがコンテキストを理解するために他のすべてのトークンを見る方法。

代名詞と医師

前に見たこの文章に戻ろう:

「医師は看護師を解雇した。なぜなら彼女は疲れていたから……」

人間にとって、「彼女」は自然に看護師を指す。LLMにとってはそう明確ではない——彼女というトークンを処理する時点では、看護師という単語は8つ前の位置にある。どうやって2つを結びつけるのか?

それがアテンションの役割だ。

アテンションが存在する理由

2017年以前、言語モデルはほとんどが再帰的RNNLSTM)だった——トークンを1つずつ読み、これまでに見たすべてを要約した「隠れ状態」を伝播させていた。

問題:この隠れ状態はボトルネックだ。すべてがそこを通過しなければならない。文章が長くなるにつれ、古い情報は薄まる。そして学習は逐次的——100番目の単語を処理するには、前の99個を処理していなければならず、並列化が難しい。

論文Attention Is All You Need(Vaswani et al.、2017)はブレークスルーを提案した:

再帰はもうやめよう。各トークンが直接すべての他のトークンを、並列に見る。

これが現代のモデルを可能にするメカニズムだ。

直感

モデルの各層で、各トークンは3つの操作を行う:

  1. 残りの文章に質問をする(Queryベクトル)。
  2. 他のすべてのトークンが、自分が何であるかを要約したラベルを表示する(Keyベクトル)。
  3. トークンは質問を各ラベルと比較する——一致するところで、何らかの内容を取得する(Valueベクトル)。

結果:各トークンの新しい表現が、Q-Kマッチから来る重みで他のトークンの重み付き合計になる。

Attention(Q, K, V) = softmax(QKT / √dk) · V

数式を覚える必要はない——アイデアを覚えよう。各トークンは他のすべてを見て、興味深いと思うものをブレンドする。

Q、K、Vは実際どこから来るのか?

無から来るわけではない。各トークンについて、その埋め込みベクトル x を取り、訓練中に学習される3つの行列で掛け合わせる:

Q = x · W_Q
K = x · W_K
V = x · W_V

これら3つの行列 W_QW_KW_Vアテンションのパラメータだ。同じ層内のすべてのトークンで共有される——これがモデルが勾配降下によって調整する部分であり、それによって正しい「質問」が正しい「ラベル」を見つけるようになる。

数式中の √d_k は、次元が大きいときに内積が爆発するのを防ぐ役割を果たす。これがないと、ソフトマックス前の値が巨大になり、ソフトマックスが飽和し、勾配が死ぬ。技術的なディテールだが、必要な工夫だ。

複数の質問を同時に

一組の質問では不十分だ。トークンは、統語的な主語照応関係主動詞を同時に見る必要があるかもしれない。

だからマルチヘッドアテンションがある:単一のQ-K-Vシステムの代わりに、複数が並列に動く(通常8、16、32個)。各ヘッドは特定のタイプの関係に特化することを学ぶ。訓練済みモデルを見ると、次のことに特化したヘッドが見つかる:

  • ローカルアテンション(各トークンが自身または隣接するトークンを見る)
  • 主語-動詞の結びつき
  • 照応関係(代名詞からその指示対象へ)
  • 区切り(句読点、文の始まり/終わり)
  • や詩的構造
  • 名前をつけられないもの

探索してみよう

以下の可視化では、2つの文章について異なるヘッドのアテンションがどのように見えるかを示している。パターンは様式化されているが(実際の重みは訓練済みモデルから来ている)、各ヘッドは現在のモデルで実際に観察された動作に対応している。

各行は、あるトークンが他のすべてのトークンをどう見ているかを示す。一部のヘッドは構文(主語 ↔ 動詞)を追い、別のヘッドは意味(参照関係や先行詞)を捉える。いずれも手で書かれたものではなく、学習から自然に立ち現れる。

試してみる3つのこと:

  • 猫が眠るを「主語↔動詞」ヘッドで見ると、眠るの行に注目。最も強い重みはに向かう。動詞が主語を「見つけた」のだ。
  • 照応関係を「照応」ヘッドで見ると、の行に注目。最も強い重みは赤ちゃんを指し示す。これがまさに代名詞の謎を解くメカニズムだ。
  • どのヘッドでも、右上の三角形を見ると:グレーになっている。これが因果マスク——トークンは先行するトークンしか見られない。これがモデルにコピーではなく予測を強制する。

因果的か、双方向か?

すべてのアテンションが等しいわけではない。2つの体制がある。

双方向。 各トークンがすべての他のトークンを、後ろも前も見る。これがBERT(Google、2018)やT5のエンコーダー側で使われている方式だ。これらのモデルは文章を理解することに長けている——分類、抽出型Q&A、意味検索——が、トークン単位でテキストを生成することはできない。

因果的。 各トークンは先行するものしか見られない。これが先ほど見た三角形のマスクだ。この制約こそが自己回帰生成を可能にする:次の単語を予測するには、モデルは過去だけから動作しなければならない。

GPT、Claude、Llama、Gemini、Mistral——一般向けのLLMはすべて因果的アテンションを使っている。このマスクこそが、モデルに記述するだけでなく予測する能力を与えるものだ。

アテンションは高コスト

このエレガントさにはコストがある。長さnの列に対して、アテンション行列の計算には**O(n²)**の操作が必要だ。コンテキストサイズを2倍にすると、コストは4倍になる。

だからGPT-2ではコンテキストウィンドウが2,048トークンに制限され、GPT-3.5では8,192トークンになり、今日の20万トークンに到達するためにはアルゴリズムの工夫(FlashAttention、スライディングウィンドウ、スパースアテンション)が必要だった。これらの技法には第 18 章で戻ってくる——KVキャッシュのメモリと結びつく形で。

アテンションは強力だが、現代のLLMの主なボトルネックでもある。

次へ

アテンション単体では言語モデルにならない。それを連続したブロックに積み重ね、フィードフォワード計算層、正規化、残差接続を追加する必要がある。それが完全なトランスフォーマーアーキテクチャ——次の章のテーマだ。

更新日

Attention:すべてを変える機構 · Step by Token