- Transformerのアテンションはトークン埋め込みを文脈に合わせて更新し、同じ単語でも周囲の単語に応じて異なる意味を持つベクトルへ移動させるメカニズムである
- 1つのアテンションヘッドは各トークンからquery/key/valueベクトルを作り、key-queryの内積とsoftmaxによって単語間の関連度の重みであるattention patternを計算する
- GPT型の自己回帰モデルは、後ろのトークンが前のトークンに影響できないようmaskingを適用し、attention patternの大きさが文脈長の2乗になるため、大きなcontext windowの拡張は負担が大きい
- GPT-3の例では、key/query行列はそれぞれ1,572,864個のパラメータを持ち、value mapは低ランク変換に分けることで、1ヘッドあたり約630万パラメータとなる
- 複数のアテンションヘッドとブロックを繰り返すTransformerは、さまざまな文脈更新の方法を学習し、その成功の大きな柱はGPUで大量の計算を高速に処理できる並列化のしやすさにある
Transformerにおけるアテンションの役割
- Transformerは入力テキストを次トークン予測に利用し、入力はまずトークン化されて単語または単語の断片に分割される
- 各トークンは高次元ベクトルである埋め込みに変換される
- この埋め込み空間の方向は意味に対応しうる
- たとえば特定方向への移動が、男性名詞の埋め込みを対応する女性名詞の埋め込みへ移すことがある
- アテンションの目的は、初期埋め込みを段階的に調整し、個々の単語情報だけでなく、より豊かな文脈的意味を持たせることである
同じ単語が文脈によって変わる理由
- “American shrew mole”, “One mole of carbon dioxide”, “Take a biopsy of the mole” におけるmoleは、それぞれ異なる意味を持つ
- 最初の埋め込み段階では、moleのベクトルは文脈を見ないlookup tableに近いため、3つのケースですべて同じである
- 次の段階であるattention blockでは、周囲の埋め込みがmoleの埋め込みへ情報を渡し、その値を更新できる
- 十分に学習されたモデルは、moleの複数の意味を埋め込み空間の異なる方向と結びつけ、文脈に応じて一般的な埋め込みに何を加えるかを計算する
- “Eiffel tower” と “miniature Eiffel tower” のように、ある単語の埋め込みは近くの単語だけでなく、離れたトークンから来る情報によっても更新されうる
- 次単語予測には最後のベクトルだけが使われるため、長い入力の最後の単語埋め込みは、予測に必要な文脈全体の情報をある程度含んでいる必要がある
単一アテンションヘッドの計算の流れ
- 基本説明はsingle head of attentionを基準に進める
- 例文 “A fluffy blue creature roamed the verdant forest.” では、形容詞が対応する名詞の初期埋め込みを更新する状況を想定する
- この例はアテンションヘッドが取りうる動作を示すためのものである
- 実際のヘッドの挙動は、多数のパラメータがコスト関数を下げるよう調整された結果であり、解釈は難しい
- 初期埋め込みは単語情報だけでなく位置情報も含み、
\vec{E}で表される
- 目標は、既存の埋め込みから文脈を反映した新しい埋め込み
\vec{E}' を作ることである
-
Query
- 最初の段階で各トークン埋め込みにquery matrix
W_Q を掛け、queryベクトル \vec{Q} を作る
- 名詞が「前に形容詞があるか？」という質問を投げるように考えられる
W_Q の要素は学習されるモデルパラメータであり、実際に特定のヘッドが何をしているかは解釈しづらい- 例としては、名詞埋め込みを「前方位置の形容詞を探す」方向へ写像すると考えられる
-
Key
- 同時に各埋め込みにkey matrix
W_k を掛け、keyベクトル \vec{K} を作る
- keyはqueryへの潜在的な答えとみなせ、queryと同じより小さな次元の空間に置かれる
- keyとqueryの整列度合いは内積で測定する
- 内積が大きいほど2つのベクトルはより強く整列している
- fluffyとblueのkeyがcreatureのqueryとよく合えば、大きな正の値を持つ
- すべてのkey-queryペアの内積を計算すると、ある単語が別の単語の意味更新にどれだけ関連しているかを示すスコア格子ができる
Attention patternとsoftmax
- 内積スコアは
-\infty から \infty までの値を取りうるため、各列にsoftmaxを適用して0から1の間の値に正規化する
- 正規化された格子はattention patternと呼ばれる
- 各列は、左側の単語が上側の単語を更新する際にどれだけ関連しているかの重みとみなせる
- 元のTransformer論文では、これをより簡潔に表記している
Q と K はqueryベクトルとkeyベクトルの全配列であるK^TQ は可能なすべてのkey-query内積の格子を表す- 論文の表記ではqueryとkeyが行に置かれ、
QK^T の形で書かれるため、ここで説明した図とは対角線方向に反転した形になる
- 数値安定性のため、key-query空間次元の平方根である
\sqrt{d_k} で割る項が入る
- softmaxは式全体を包むが、意味としては各列ごとに適用される
Maskingとcontextサイズの制約
- 学習中、モデルは与えられたテキストから単一の次トークンだけを予測するのではなく、各部分系列の後に来る可能な次トークンも同時に予測する
- 1つのテキスト例が複数の学習例のように働くため効率が高い
- GPTの例では、後ろのトークンが前のトークンに影響すると次トークンの正解を漏らしてしまうため、maskingを使う
- softmaxの前に該当位置の値を負の無限大に設定する
- softmax後、その位置は0になり、列は正規化された状態を保つ
- 常にmaskingが適用されるアテンションだけがあるわけではないが、GPTの例では後ろのトークンが前のトークンに影響しないよう常に使われる
- attention patternの大きさはcontext sizeの2乗に等しい
- このためcontext sizeは大規模言語モデルの重要な制約になりうる
- より大きなcontext windowのために、アテンションメカニズムをさらにスケーラブルにする変種も登場しているが、ここでは基本形のみを扱う
Valueで埋め込みを実際に更新する方法
- attention patternは、どの単語がどの単語を更新するかについての重みを与え、次の段階は実際の埋め込み変化量を作ることである
- 各埋め込みにvalue matrix
W_V を掛けてvalueベクトルを作る
- valueベクトルは埋め込みと同じ高次元空間に置かれる
- 関連する単語が別の単語の意味を調整する際、どのような具体的変化を加えるべきかを表す
- 各列でvalueベクトルにattention patternの対応する重みを掛けてすべて足し合わせると、変化量
\Delta \vec{E} になる
- この変化量を元の埋め込みに足すと、文脈を反映した新しい埋め込み
\vec{E}' が作られる
- 例ではcreatureがfluffyとblueの情報を吸収し、“fluffy blue creature” に近い意味を持つようになる
- 同じ過程をすべての列に適用すると、トークン列全体に対する精緻化された埋め込みがアテンションブロックの外へ出てくる
- 単一のアテンションヘッドは、key matrix、query matrix、value matrix という3種類の学習パラメータ行列でパラメータ化される
GPT-3基準のパラメータ計算
- GPT-3の例では、keyとquery行列はそれぞれ埋め込み次元に対応する12,288個の列と、key-query空間次元に対応する128個の行を持つ
- value matrixを12,288×12,288の正方行列にすると、150,994,944個のパラメータが追加され、key/queryよりはるかに大きくなる
- 実際には、value mapを2つの小さな行列に分解してパラメータ数をkey/queryと近い規模に合わせるほうが効率的である
- 1つ目の行列は大きな埋め込み空間を128次元のような小さな空間へ落とす
- 2つ目の行列は小さな空間から再び埋め込み空間へ戻す
- 線形代数の観点では、value map全体を低ランク変換に制限していることになる
- この説明では2つの行列を
Value_\downarrow、Value_\uparrow と呼ぶが、これは慣例的な名称ではない
- 4つの行列を合計すると、1つのアテンションヘッドは約630万パラメータを持つ
Self-attentionとcross-attention
- ここまでの構造は、より正確にはself-attention headに当たる
- cross-attention headは異なる2つのデータ集合を処理するモデルで登場する
- たとえば翻訳モデルでは、keyが一方の言語から来て、queryが別の言語から来ることがある
- attention patternは、一方の言語の単語が他方の言語の単語にどう対応するかを表せる
- cross-attentionでは、keyとquery mapが異なるデータセットに作用する点がself-attentionと異なる
- 翻訳のような設定では、後ろのトークンが前のトークンに影響するという概念がないため、通常はmaskingがない
Multi-headed attentionと繰り返されるブロック
- 実際のattention blockは、複数のヘッドを並列実行するmulti-headed attentionで構成される
- GPT-3は各ブロック内で96個のアテンションヘッドを使う
- 96個の異なるkey/query行列が96個の異なるattention patternを作る
- 各ヘッドは独自のvalue行列でvalueベクトル列を作る
- 各トークン位置ごとに、すべてのヘッドが提案した変化量
\Delta \vec{E} を合算して元の埋め込みに加える
- 複数ヘッドを並列に実行することで、モデルは文脈が意味を変えるさまざまな方法を学習できる容量を持つ
- GPT-3基準では、96個のヘッドを持つmulti-headed attention block 1つは約6億パラメータを持つ
- 論文と実装では、各ヘッドの
Value_\uparrow に相当する行列は1つの大きなoutput matrixにまとめられ、multi-headed attention block全体に接続される
- 通常、特定ヘッドのvalue matrixと言うと、ここで
Value_\downarrow と呼んだ最初の射影段階を指す
より深いTransformerで意味が蓄積される仕組み
- Transformer内部のデータは1つのattention blockだけを通過するのではなく、複数のattention blockとmulti-layer perceptronを通る
- ある単語の埋め込みが文脈を一部吸収した後も、より洗練された周囲の埋め込みの影響を受ける機会が続く
- ネットワークが深くなるほど、各埋め込みは他の埋め込みからより多くの意味を取り込み、感情、トーン、皮肉かどうかといった、より高次の抽象的特徴を符号化できる容量を持つようになる
- GPT-3は96層を含み、key/query/value関連のパラメータは合計で580億個未満と説明される
- これはネットワーク全体のパラメータのおよそ3分の1で、残る大半はアテンションの間にあるブロックから生じる
- アテンションメカニズム成功の大きな要因は特定の1つの挙動ではなく、GPUで大量の計算を短時間に実行できる高い並列化可能性にある
- ディープラーニングでは、スケール拡大がモデル性能に大きな質的改善をもたらしうるという教訓があるため、拡張を可能にする並列化しやすいアーキテクチャは大きな利点を持つ
1件のコメント
Hacker News のコメント
量子化学と一部の機械学習をやってきた立場から見ると、この動画を見ていて Transformer モデルと量子力学の間の類似性がかなり目についた
量子力学では、物理系全体の状態は非常に高次元の正規化ベクトル、つまり Hilbert 空間の半直線としてエンコードされ、時間に伴う変化はおおまかに unitary 行列 U = exp(-iHt) と見なせる時間発展演算子が担う
動画では、次トークン予測は最後の文脈認識埋め込みベクトルだけから次の文脈認識埋め込みベクトルを計算して決まるというので、高次元ベクトルに線形の状態関数を適用した結果のように見える
システム全体の Hamiltonian を学習データからオフラインで生成したうえで、特定のサブシステムであるコンテキストウィンドウをその Hamiltonian に合う基底で再パラメータ化し、1ステップの時間発展を適用してから元の基底に戻すのに似ているように感じる
ただ、特定分野を研究してきた人にはあらゆる問題がその分野のハンマーに合う釘に見えることもあるので、この類似性が他の人にも見えるのか、それともかなり無理があるのか気になっている
演算子を1回適用することで宇宙の状態が Planck 時間1つ分進むのだとしたら、そのような宇宙と時間が連続的な宇宙との違いを私たちが観測できるのかも気になる
当時の物理学で起きていたこととのつながりは、今になってようやく思い浮かんだ
CodeEmporium の YouTube 動画のほうが追いやすかった: https://www.youtube.com/watch?v=Nw_PJdmydZY
Transformer は比喩で説明しにくく、実際なぜ動くのかについてのよい説明もないので、単にメカニズムを見せて解釈は見る人に委ねるほうがよいかもしれない
また、ドット積はベクトルが互いに射影されるものとして説明するほうがより単純
“the cat sat on the ...” の次に “mat” が来る理由は、データセットで最も頻繁に出てきた単語だからで、ニューラルネットワークはそうした頻度のモデルである
“London in UK” は知っているが “London in France” は知らないように見える理由も、データセットでは “UK” のほうがはるかに頻繁に出てくるから
アルゴリズム自体は、計算をハードウェアに合わせて整列させること以外に、特に興味深いことはしていない。価値はデータ内の条件付き確率構造から来ており、その構造は人々が互いに情報を伝えようとして単語を有用に並べた結果である
不完全な説明なのは分かっているが、何もないよりはましだと思う
簡単なリクエストを処理するときに LLM がどう動作するかを示す説得力のある可視化がある: https://bbycroft.net/llm
3blue1brown の詳しい説明をうまく補完してくれる
すばらしい動画。Q*K 行列乗算がなぜボトルネックなのかをよく示している。シーケンス、つまりコンテキストウィンドウ長が S なら、すべての query とすべての key の結果である SxS サイズの行列をメモリに保存しなければならないため
このボトルネックを改善する新しめのアイデアとして Ring Attention があり、この記事がよく説明している: https://learnandburn.ai/p/how-to-build-a-10m-token-context
その記事は自分が編集した
Unsloth では Flash Attention のおかげでメモリ使用量が二乗ではなく線形に増え、ファインチューニングは2倍速くなり、VRAM 使用量は80%減り、推論も2倍速くなる。ただし計算量は依然として O(N^2) である
長いコンテキストでは、Unsloth の最新リリースが HF+FA2 より4倍長いコンテキストを +1.9% のオーバーヘッドで収められ、H100 で 228K コンテキストまで可能
以前の記事 “But what is a GPT?” も本当に良い: https://www.3blue1brown.com/lessons/gpt
この動画のおかげで、アテンション機構は特定の関数というより、一種のメタ関数に近いのだと気づいた。
正しく理解できているなら、Attention + 学習済み重みは、Transformerがある程度任意の関数を学習できるようにし、その関数の中には scaled dot-product のようなマッチング機構が含まれる。
だから線形アテンションは、標準アテンションの能力に決して近づきにくいと思う。すべての入出力ペアを探索する二乗項が本質的な特徴だからだ。
この動画が理解しやすかったのは、アニメーションのおかげが大きい。話すタイミングに合わせて拡大・縮小し、展開していく作りがとてもよくできていた。
密接に関連する分野で働いているが、この動画はすぐにうちのチームのオンボーディング文書に入った。
可視化コードのかなりの部分がGitHubに上がっている点も重要: https://github.com/3b1b/videos/tree/master/_2024/transformers
ついに理解できた。他の動画はなぜあんなに分かりにくく作ったのか分からない。
教育能力と研究能力は、概して互いに排他的なスキルなのではないかと思うようになった。
第一に、標準用語がよくない。“attention”もかろうじて直感的という程度で、“self-attention”はさらに悪く、“key”と“value”は言うまでもない。
第二に、主要論文である Attention is All You Need や BERT 論文などは、うまく書かれていなかった。成果を貶めたいわけではなく、巨大なブレークスルーを含む影響力のある論文であっても説明が下手なことはあり得るし、実際そうだったと思う。
第三に、これらの構造はおおむね、あれこれ試してうまくはまるものを見つける形で発見された。この構造がうまく機能するはずだという予測に至る洞察の過程が先にあり、それを実験で検証したのではなく、最初から最後まで経験的だった。
そのため、なぜこれほどうまく機能するのかを完全には理解しておらず、あらゆる説明は事後的な合理化に近い。最近では、十分に調整すれば別の構造でも同程度にうまく機能し得ることを示唆する研究もある。完全に理解していないものを説明するのは難しい。
現在のアーキテクチャがどのように進化したのかを説明する参考資料があるのか気になる。ごく単純な中核アイデアから、有名な「all you need」論文まで続く流れを見てみたい。
そうでないと、多くの仕組みが突然現れたように感じられ、計算は多いのに直感が少ない。
Jeremy HowardがTwitterで、このアイデアのさまざまなバージョンを何度も見たと言っていたが、それは自然に出てくるアイデアだったという意味に聞こえる。このアイデアが他の場所でどのように登場したのかの例を見れば、直感を築けそうだ。
このRNN/LSTMアプローチには、固定サイズ表現という弱点と、出力の特定部分を生成するときに入力シーケンスのどの部分を使うべきか決めにくいという弱点があった。Bahdanauらは、エンコーダ・デコーダRNNにアテンション機構を組み合わせた構造でこれを解決し、最終状態だけでなくRNNのすべての過去状態を参照させた。
RNNは学習が非効率だったため、Jakob Uszkoreitは大規模な並列ハードウェアをよりうまく活用する方法を探し、言語は逐次的であるだけでなく階層的でもある点に注目した。各層でサブシーケンスのトークンを並列処理しつつ、Bahdanau式アテンションを維持して、トークン同士が互いを参照する自己アテンションによって次の階層を予測する層状構造を提案した。
初期実装は動作したが、当時の畳み込みのような他のアプローチより優れていたわけではなく、その後Noam Shazeerがそのアイデアを発展させ、はるかにうまく機能する構造を作った。そして不要な構成要素を取り除く実験を経て、元のTransformerになったと理解している。最終構造におけるkeyベースのアテンション形式を誰が考案したのかはよく分からない。
Attention is All You Need 論文の元のTransformerは、以前のRNNベースのアプローチにならって別々のエンコーダとデコーダを持っており、GoogleのBERTのような初期モデルにも使われた。しかし言語モデルにはこれは必須ではなかったため、OpenAIのGPTはデコーダ部分だけを使い、今では誰もが概ねこの方式を使っている。デコーダ専用Transformerでは、入力文が最下層に入り、その後各層を通るごとに一段階ずつ変換されて上部から出てくる。入力シーケンスの末尾には終了トークンが付き、それが出力シーケンスの次のトークン、つまり最後のトークンへと変換される。