アテンション機構を理解する クエリ、キー、バリューの正体

こんにちは。ゆうせいです。

前回の記事では、Transformerの入口と出口を追いました。単語がベクトルになり、最後にSoftmaxで確率に変わる。その流れを見ました。

残っているのが、真ん中です。アテンション機構と呼ばれる仕組みです。

2017年の論文「Attention is All You Need」で提案され、現在のAIブームを支えています。しかし、ここでつまずく人が非常に多い。私も最初は何度も読み返しました。

この記事では、アテンションが何をしているのかを、具体例から追います。理解に時間がかかっても、それが普通です。

結論

アテンション機構の要点は、三つです。

要点内容
文脈で意味を更新する単語のベクトルを、周囲の情報で書き換える
三つの行列が働くクエリが問い、キーが答え、バリューが中身を運ぶ
並列に動くことが本質多数の計算を同時に実行できる点が、規模の拡大を支えた

特に三番目を先に書いておきます。

アテンションが成功した理由は、特定の振る舞いを可能にしたことより、並列化しやすかったことにあります。GPUの性能を引き出せる構造だったから、規模を拡大できました。

以下、順に見ていきます。

1 何を解決したいのか

計算の詳細に入る前に、アテンションに期待する振る舞いを確認します。

同じ単語が違う意味を持つ

次の三つの表現を考えてください。

表現moleの意味
American true moleモグラ
one mole of carbon dioxideモル(物質量の単位)
take a biopsy of the moleほくろ

英語のmoleは、文脈によってまったく違う意味になります。

最初の段階では区別できない

ところが、Transformerの最初の処理では、この三つが同じベクトルになります。

埋め込み行列は、単語からベクトルへの対応表にすぎません。文脈を参照する仕組みがないためです。

次の段階で情報が流れ込む

周囲のベクトルから情報を受け取るのは、次の段階です。

そこで想定される状況は、次のようなものです。

埋め込み空間には、moleの複数の意味に対応する複数の方向があります。学習済みのアテンションブロックは、文脈に応じて、一般的な埋め込みに何を足せばよいかを計算します。

足された結果、より具体的な方向を指すようになります。

もう一つの例

towerという単語を考えます。

この埋め込みは、おそらく漠然とした方向でしょう。大きくて高いもの全般に関係する位置にあります。

直前にEiffelがあれば、どうでしょうか。ベクトルを更新して、エッフェル塔をより明確に表す方向へ動かしたくなります。パリやフランス、鉄製のものに関連するベクトルと相関する方向です。

さらに直前にminiatureがあれば、また更新が必要です。もはや「大きくて高いもの」とは相関しない方向へ動かします。

単語の意味だけではない

アテンションの役割は、単語の意味を精密にすることだけではありません。

ある埋め込みに符号化された情報を、別の埋め込みへ移す。それが本質です。

移す先は、かなり離れた位置かもしれません。移す情報も、単語一つ分よりずっと豊かなものかもしれません。

極端な例

推理小説のほぼ全体を入力したとします。終盤近くまで進み、こう書かれています。

「したがって、犯人は」

次の単語を正確に予測するには、どうすればよいでしょうか。

系列の最後のベクトルは、もともと「は」という単語を埋め込んだだけのものでした。そのベクトルが、アテンションブロックを通るうちに更新され、個々の単語をはるかに超える情報を持つようになる必要があります。

文脈全体のうち、次の単語の予測に関わるすべての情報が、そこに符号化されていなければなりません。

練習問題1

「銀行」という単語のベクトルが、次の二つの文脈でどう異なる方向へ更新されるべきか説明してください。

(1) 川の銀行を歩いた
(2) 銀行で口座を開いた

解答を示します。

(1) では、土手や水辺に関連する方向へ更新されるべきです。地形や自然に関わるベクトルと相関する位置へ動きます。

(2) では、金融機関に関連する方向へ更新されるべきです。預金、融資、経済といったベクトルと相関する位置へ動きます。

補足します。埋め込み行列から取り出された時点では、どちらも同一のベクトルです。

区別が生まれるのは、アテンションによって周囲の情報が流れ込んだ後です。

練習問題2

アテンションの役割を「単語の意味を精密にすること」と説明すると、不十分である理由を述べてください。

解答を示します。

アテンションの本質は、ある埋め込みに含まれる情報を、別の埋め込みへ移すことです。

移す情報は、単語一つ分の意味に限りません。文章全体の論理、登場人物の関係、文体といった、より豊かな情報も含まれます。

また、情報を移す距離も、隣接する単語に限りません。数千トークン離れた位置からでも、関連があれば情報が流れます。

補足します。推理小説の例が、この点をよく示しています。

「犯人は」の次を予測するには、物語全体から集めた情報が必要です。単語の意味を精密にするという説明では、この働きが捉えられません。

2 具体例を設定する

計算を追うため、単純な例を用意します。

例文

「a fluffy blue creature roamed the verdant forest」

ふわふわした青い生き物が、緑豊かな森をさまよった、という文です。

注目する更新

ここでは、形容詞が対応する名詞の意味を調整する、という更新だけを考えます。

fluffyとblueが、creatureの意味を更新する。この動きを追います。

断っておくこと

この例は、説明のために作ったものです。

実際のアテンションヘッドが、こういう分業をしているとは限りません。深層学習の常で、真の振る舞いは解読が困難です。膨大なパラメータが、コスト関数を最小化するよう調整された結果だからです。

ただし、パラメータで満たされた行列を一つずつ見ていくとき、想像上の具体例があると理解が進みます。そのための設定です。

これから扱うもの

これから説明するのは、アテンションヘッド1個分の処理です。

実際のアテンションブロックは、多数のヘッドを並列に動かします。その話は後半で扱います。

埋め込みの記号

各単語の埋め込みを、E と表します。

正確には、この埋め込みは単語の意味だけでなく、位置の情報も符号化しています。位置の符号化には詳しい話がありますが、ここでは「何の単語か」と「どこにあるか」の両方が分かる、という点だけ押さえてください。

3 クエリ 問いを立てる

名詞が問いかける

処理の最初の段階を、擬人化して考えます。

creatureという名詞が、こう問いかけます。

「私の前に、形容詞はいますか」

問いをベクトルにする

この問いも、ベクトルとして符号化されます。

これをクエリと呼びます。

次元は小さい

クエリベクトルの次元は、埋め込みベクトルよりずっと小さくなります。

GPT-3では128次元です。埋め込みの12,288次元と比べると、100分の1程度です。

計算方法

クエリ行列 W_Q を、埋め込みに掛けます。

Q = W_Q E

この行列を、文脈内のすべての埋め込みに掛けます。トークンごとに1本のクエリベクトルができます。

行列の中身

W_Q の要素は、モデルのパラメータです。学習によってデータから決まります。

実際のアテンションヘッドで、この行列が何をしているかを読み解くのは困難です。

ここでは、想像上の役割を仮定します。名詞の埋め込みを、「前の位置にある形容詞を探す」という概念を表す方向へ写像する。そういう働きをすると考えてみます。

名詞以外の埋め込みに何をするのかは、分かりません。同時に別の目的を果たそうとしているのかもしれません。

練習問題3

クエリベクトルの次元が、埋め込みベクトルの次元より小さいことの利点を説明してください。

解答を示します。

パラメータ数と計算量を抑えられます。

クエリ行列の大きさは、クエリの次元と埋め込みの次元の積で決まります。クエリの次元を小さくすれば、その分だけパラメータが減ります。

また、後で計算する内積も、次元が小さいほど高速になります。

補足します。多数のヘッドを並列に動かすことを考えると、この効率は重要です。

GPT-3では96個のヘッドが並列に動きます。1個あたりの規模を抑えることで、全体が現実的な大きさに収まります。

4 キー 答えを用意する

二つ目の行列

クエリと同時に、キー行列 W_K が働きます。

これも、すべての埋め込みに掛けられます。

K = W_K E

得られるベクトルをキーと呼びます。

役割

キーは、クエリへの答えになり得るものだと考えてください。

キー行列も調整可能なパラメータで満たされており、クエリ行列と同じく、埋め込みを同じ小さい次元の空間へ写像します。

一致するとはどういうことか

キーとクエリが密接に揃っているとき、そのキーはクエリに合致していると考えます。

今回の例で言えば、キー行列がfluffyやblueといった形容詞を、creatureが生み出したクエリと近い方向のベクトルへ写像する、という状況です。

図解案

研修資料では、二つの空間を描くと伝わります。

左に埋め込み空間(12,288次元)、右にクエリキー空間(128次元)を置きます。そして、W_Q W_K という二本の矢印で結びます。

同じ小さい空間に、クエリとキーが降りてくる。この構図が要点です。

練習問題4

クエリ行列とキー行列は、どちらも埋め込みを同じ小さい空間へ写像します。なぜ同じ空間である必要があるのですか。

解答を示します。

クエリとキーの内積を計算するためです。

内積は、同じ次元のベクトルどうしでなければ計算できません。異なる空間に写像すると、比較ができなくなります。

補足します。同じ空間であっても、写像の仕方は異なります。

W_Q W_K は別の行列であり、学習によって別の値を取ります。問いを表す方向と、答えを表す方向が、それぞれ独立に学習されます。

5 内積で一致を測る

すべての組み合わせを計算する

各キーが各クエリにどれだけ合致するかを測ります。

方法は、内積です。可能なキーとクエリの組み合わせすべてについて、内積を計算します。

表として可視化する

格子状に点を並べた図を想像してください。

内積が大きいところほど、点が大きくなります。キーとクエリが揃っている場所です。

今回の例では、fluffyとblueが生み出したキーが、creatureのクエリと密接に揃っているはずです。その二か所の内積は、大きな正の値になります。

用語

機械学習の分野では、この状況を次のように言います。

fluffyとblueの埋め込みが、creatureの埋め込みに注意を向けている(attend to)。

関係のない組み合わせ

theという単語のキーと、creatureのクエリとの内積はどうでしょうか。

小さい値か、負の値になります。互いに無関係であることを反映しています。

この時点での値の範囲

この格子には、マイナス無限大からプラス無限大までの任意の実数が入り得ます。

各単語が、他の各単語の意味を更新するうえでどれだけ関連するかのスコアです。

練習問題5

クエリベクトル \mathbf{q} = (2, ; -1, ; 3) と、二つのキーベクトルの内積を計算し、どちらがより合致しているか判定してください。

\mathbf{k}_1 = (1, ; 0, ; 2), \quad \mathbf{k}_2 = (-1, ; 3, ; 0)

解答を示します。

\mathbf{k}_1 との内積を計算します。

2 \times 1 + (-1) \times 0 + 3 \times 2 = 2 + 0 + 6 = 8

\mathbf{k}_2 との内積を計算します。

2 \times (-1) + (-1) \times 3 + 3 \times 0 = -2 - 3 + 0 = -5

\mathbf{k}_1 のほうが合致しています。内積が大きな正の値であるためです。

補足します。\mathbf{k}_2 の内積は負でした。

これは、方向がむしろ反対を向いていることを意味します。Softmaxを通すと、この項の重みは非常に小さくなります。

6 Softmaxで正規化する

重み付き和のために

このスコアを、これから重み付き和に使います。

列に沿って、関連度で重み付けした和を取ります。

必要な条件

そのためには、値の範囲を整える必要があります。

条件内容
範囲各値が0以上1以下
合計各列の合計が1

確率分布のように振る舞ってほしいのです。

Softmaxを使う

前回の記事で扱った関数の出番です。

各列に対してSoftmaxを適用し、値を正規化します。

p_i = \frac{e^{z_i}}{\sum_j e^{z_j}}

アテンションパターン

正規化後の格子を、アテンションパターンと呼びます。

各列は、左側の単語が上側の単語にとってどれだけ関連するかを示す重みだと考えてください。

論文での表記

元のTransformer論文には、非常に簡潔な書き方があります。

\text{Attention}(Q, K, V) = \text{softmax}\left( \frac{QK^{T}}{\sqrt{d_k}} \right) V

記号を確認します。

記号意味
Q クエリベクトルの配列全体
K キーベクトルの配列全体
V バリューベクトルの配列(後述)
d_k クエリキー空間の次元

分子の意味

QK^{T} は、キーとクエリのすべての組み合わせの内積を、まとめて表したものです。

転置が現れる理由は、これまでの連載で扱った通りです。行と列の役割を入れ替えて、内積の形にしています。

なぜ平方根で割るのか

触れていなかった技術的な詳細があります。

数値の安定性のため、値を次元の平方根で割ります。

\sqrt{d_k}

GPT-3であれば、次の値です。

\sqrt{128} \approx 11.3

次元が大きいほど、内積の値も大きくなる傾向があります。そのままSoftmaxに入れると、分布が極端に尖ってしまいます。それを防ぐための調整です。

練習問題6

内積の値が (4, ; 1, ; 0) であるとき、Softmaxを適用して重みを求めてください。e^4 \approx 54.60 e^1 \approx 2.718 を使ってください。

解答を示します。

指数の値を並べます。

54.60, \quad 2.718, \quad 1

合計を求めます。

54.60 + 2.718 + 1 = 58.318

各値を割ります。

p_1 = \frac{54.60}{58.318} \approx 0.936

p_2 = \frac{2.718}{58.318} \approx 0.047

p_3 = \frac{1}{58.318} \approx 0.017

補足します。1番目が圧倒的な重みを持ちました。

この単語からの情報が、ほぼすべてを占めることになります。他の二つは、わずかに寄与するだけです。

練習問題7

内積を \sqrt{d_k} で割る理由を説明してください。d_k = 128 の場合の除数も求めてください。

解答を示します。

数値の安定性のためです。

次元が大きいほど、内積の値は大きくなる傾向があります。値が大きいままSoftmaxに入れると、最大値だけが1に近づき、他がほぼ0になります。分布が極端に尖ると、学習時の勾配が小さくなり、学習が進みにくくなります。

除数を計算します。

\sqrt{128} = \sqrt{64 \times 2} = 8\sqrt{2} \approx 11.3

補足します。前回の記事で扱った温度の話と関係があります。

割るという操作は、温度を上げることと同じ効果を持ちます。分布を平坦にする調整です。

7 マスキング 未来を見せない

学習時の事情

省いていた技術的な詳細があります。

学習の際、ある文章でモデルを動かし、正解の次の単語にどれだけ高い確率を割り当てたかに応じて、重みが調整されます。

ここで、効率を上げる工夫があります。

文章の各位置について、同時に次のトークンを予測させるのです。

具体例

今回の例文であれば、次のような予測を同時に行います。

位置まで予測対象
a fluffy bluecreature
a fluffy blue creatureroamed
a fluffy blue creature roamedthe

一つの訓練例が、実質的に多数の訓練例として働きます。

問題が生じる

ここで困ったことが起きます。

後の単語が、前の単語に影響を与えてはいけません。答えを漏らしてしまうからです。

creatureの次を予測させるとき、creatureのベクトルがroamedの情報を含んでいたら、予測は自明になります。

解決策

アテンションパターンの、後のトークンが前のトークンに影響する位置を、すべて0にする必要があります。

単純な方法では失敗する

その位置を0に設定すればよい、と思うかもしれません。

しかし、そうすると列の合計が1でなくなります。正規化が崩れます。

正しい方法

Softmaxを適用する前に、該当する要素をマイナス無限大にします。

e^{-\infty} = 0

すると、Softmaxの後で0になります。しかも、列の合計は1のままです。

この処理をマスキングといいます。

適用の範囲

アテンションにはマスキングを適用しない種類もあります。

GPTの例では、常に適用します。学習時ほど重要ではないものの、チャットボットとして動かす際にも適用されます。

練習問題8

マスキングにおいて、該当する要素を0にするのではなく、マイナス無限大にする理由を説明してください。

解答を示します。

Softmaxを適用した後も、列の合計を1に保つためです。

Softmaxの後で0に設定すると、合計が1未満になり、確率分布として成立しなくなります。

Softmaxの前にマイナス無限大を設定すれば、e^{-\infty} = 0 となって該当項が消え、残りの項で正規化されます。合計は1のままです。

補足します。実装では、マイナス無限大の代わりに非常に大きな負の値を使います。

浮動小数点の計算で無限大を扱うと、不具合が生じやすいためです。

練習問題9

学習時に、文章の各位置で同時に予測を行うことの利点を説明してください。

解答を示します。

一つの訓練例から、多数の訓練信号が得られます。

長さ n の文章であれば、最大で n-1 個の予測ができます。文章を1回通すだけで、それだけの学習が進みます。

補足します。この効率が、大規模な学習を可能にしています。

ただし、この方式を成立させるにはマスキングが必須です。後の単語が前に影響すれば、答えを見ながら予測することになり、学習が成立しません。

8 コンテキストサイズという壁

パターンの大きさ

アテンションパターンについて、もう一つ考えるべき点があります。

その大きさは、コンテキストサイズの二乗に等しくなります。

数値で確かめる

コンテキストサイズが2,048の場合を計算します。

2048^2 = 4194304

約419万個の値になります。

4,096に倍増させるとどうでしょうか。

4096^2 = 16777216

約1,678万個です。4倍になりました。

なぜ壁になるのか

長さを2倍にすると、計算量は4倍になります。

これが、大規模言語モデルにとってコンテキストサイズが大きな制約となる理由です。単純に拡大することはできません。

研究の動向

より大きなコンテキストウィンドウを求めて、近年はアテンション機構の変種が提案されています。

計算量を線形に抑える手法、一部のトークンだけに注目する手法など、さまざまな方向があります。

ここでは、基本に集中します。

練習問題10

コンテキストサイズを1,024から8,192に拡大した場合、アテンションパターンの要素数は何倍になりますか。

解答を示します。

倍率を計算します。

\frac{8192}{1024} = 8

要素数は二乗に比例するため、次の通りです。

8^2 = 64

64倍になります。

補足します。実際の要素数も計算しておきます。

1024^2 = 1048576

8192^2 = 67108864

約105万個から約6,711万個へ増えます。メモリと計算時間の両面で、大きな負担です。

9 バリュー 情報を運ぶ

まだ半分

ここまでで、どの単語がどの単語に関連するかが分かりました。

しかし、まだ埋め込みを更新していません。関連する相手へ、実際に情報を渡す必要があります。

何をしたいのか

fluffyの埋め込みが、creatureに変化を起こす。その変化によって、creatureが12,288次元空間の別の場所へ移動する。移動先は、ふわふわした生き物をより明確に表す位置です。

三つ目の行列

バリュー行列 W_V を使います。

これを、最初の単語の埋め込みに掛けます。fluffyの埋め込みです。

\mathbf{v} = W_V E_{\text{fluffy}}

得られるのがバリューベクトルです。

埋め込みに足す

このバリューベクトルを、二つ目の単語の埋め込みに足します。creatureの埋め込みです。

したがって、バリューベクトルは埋め込みと同じ高次元空間に存在します。

意味の解釈

バリュー行列を単語の埋め込みに掛けるとき、次のような問いを立てていると考えてください。

「この単語が他の何かの意味を調整するうえで関連するなら、それを反映させるために、その何かの埋め込みに何を足すべきか」

計算の流れ

アテンションパターンが計算できたら、クエリとキーは役目を終えます。

次に、バリュー行列をすべての埋め込みに掛け、バリューベクトルの列を作ります。

各バリューベクトルは、対応するキーと関連づいていると考えてください。

重み付き和を取る

図の各列について、バリューベクトルに対応する重みを掛けます。

creatureの埋め込みの下では、fluffyとblueのバリューベクトルが大きな割合で加わります。他のバリューベクトルは0になるか、ほぼ0になります。

更新する

列の中の、重み付けされたバリューをすべて足し合わせます。

これが、足したい変化量です。\Delta E と表します。

これを元の埋め込みに足します。

E_{\text{new}} = E_{\text{creature}} + \Delta E

結果として得られるのが、より洗練されたベクトルです。ふわふわした青い生き物という、文脈豊かな意味を符号化しています。

すべての列で行う

一つの埋め込みだけではありません。

図のすべての列について、同じ重み付き和を計算します。変化量の列ができます。

それらを対応する埋め込みに足すと、洗練された埋め込みの列がアテンションブロックから出てきます。

練習問題11

三つのバリューベクトルと、対応する重みが次の通りでした。変化量 \Delta E を求めてください。

単語バリューベクトル重み
fluffy(2, ; 0, ; 1) 0.5
blue(0, ; 4, ; 2) 0.4
the(1, ; 1, ; 1) 0.1

解答を示します。

各成分ごとに、重み付き和を計算します。

第1成分です。

0.5 \times 2 + 0.4 \times 0 + 0.1 \times 1 = 1.0 + 0 + 0.1 = 1.1

第2成分です。

0.5 \times 0 + 0.4 \times 4 + 0.1 \times 1 = 0 + 1.6 + 0.1 = 1.7

第3成分です。

0.5 \times 1 + 0.4 \times 2 + 0.1 \times 1 = 0.5 + 0.8 + 0.1 = 1.4

\Delta E = (1.1, ; 1.7, ; 1.4)

補足します。theの寄与は小さくなりました。重みが0.1しかないためです。

アテンションパターンが、どの情報をどれだけ取り込むかを制御しています。

練習問題12

クエリ、キー、バリューの三つの行列の役割を、それぞれ一文で説明してください。

解答を示します。

クエリ行列は、各単語が「何を探しているか」を表すベクトルを作ります。

キー行列は、各単語が「何を提供できるか」を表すベクトルを作ります。

バリュー行列は、実際に渡す情報の中身を作ります。

補足します。クエリとキーは、どの単語がどの単語に関連するかを決めるために使われます。

バリューは、関連が決まった後で、実際に何を渡すかを決めます。役割が分かれています。

10 バリュー行列の工夫

素朴に考えると

ここまでの説明では、バリュー行列は正方行列になります。

入力も出力も、12,288次元の埋め込み空間にあるためです。

12288 \times 12288 = 150994944

約1億5,000万個のパラメータです。

キーとクエリと比べる

キー行列とクエリ行列は、12,288列と128行です。

12288 \times 128 = 1572864

約157万個です。

バリューだけが、100倍近く大きくなってしまいます。

実際の設計

そうすることも可能です。しかし、実際にはもっと効率的な方法が取られます。

バリュー行列を、二つの小さな行列の積に分解します。

二段階の写像

段階内容大きさ
1埋め込み空間から小さい空間へ128 \times 12288
2小さい空間から埋め込み空間へ12288 \times 128

一段目を値下げ行列、二段目を値上げ行列と呼ぶことにします。これは一般的な呼称ではありませんが、説明のために使います。

線形代数の用語で言えば

全体のバリュー写像を、低ランク変換に制約していることになります。

情報を一度狭い空間に通すことで、パラメータ数を抑えています。

パラメータ数を数える

これで、四つの行列がすべて同じ大きさになりました。

行列パラメータ数
クエリ約157万
キー約157万
値下げ約157万
値上げ約157万
合計約630万

アテンションヘッド1個で、約630万個です。

論文での書かれ方

補足しておきます。

多くの論文では、書き方が少し異なります。

各ヘッドの値上げ行列が、まとめて一つの巨大な行列として扱われます。これを出力行列と呼び、マルチヘッドアテンションブロック全体に対応づけます。

そして、あるヘッドのバリュー行列と言ったとき、通常は値下げ行列のほうだけを指します。

他の資料を読むときに混乱しないよう、この点を知っておいてください。

練習問題13

バリュー行列を二つの小さな行列に分解することで、パラメータ数がどれだけ削減されるか計算してください。埋め込み次元12,288、小さい空間の次元128とします。

解答を示します。

分解しない場合を計算します。

12288 \times 12288 = 150994944

分解した場合を計算します。

12288 \times 128 \times 2 = 3145728

削減率を求めます。

\frac{3145728}{150994944} \approx 0.0208

約2%になります。つまり、約98%削減されます。

補足します。削減の代償として、表現力に制約が加わります。

低ランク変換しか表せなくなるためです。しかし、多数のヘッドを並列に動かすことを考えると、この制約は許容されます。

11 自己注意と交差注意

ここまでの名称

ここまで説明してきたものは、正確には自己注意ヘッドと呼ばれます。

英語ではself-attentionです。

別の種類

区別されるのが、交差注意です。英語ではcross-attentionです。

GPTの例には関係しませんが、知っておくと役立ちます。

何が違うのか

交差注意は、二つの異なる種類のデータを処理するモデルで使われます。

データ1データ2
翻訳ある言語のテキスト別の言語のテキスト
音声認識音声入力生成中の文字起こし

構造はほぼ同じ

交差注意ヘッドは、自己注意とほとんど同じ形をしています。

違いは一点だけです。キー写像とクエリ写像が、異なるデータ集合に作用します。

翻訳を行うモデルであれば、キーが一方の言語から、クエリが他方の言語から来るかもしれません。アテンションパターンは、一方の言語のどの単語が他方のどの単語に対応するかを表すことになります。

マスキングの有無

この設定では、通常マスキングは行われません。

後のトークンが前のトークンに影響するという概念が、そもそも存在しないためです。

練習問題14

自己注意と交差注意の違いを説明してください。また、交差注意でマスキングが不要な理由を述べてください。

解答を示します。

自己注意では、キーとクエリが同じデータ集合から作られます。同じ文章の中で、単語どうしが情報をやり取りします。

交差注意では、キーとクエリが異なるデータ集合から作られます。たとえば翻訳では、原文と訳文という別々のデータを結びつけます。

マスキングが不要な理由は、二つのデータ集合の間に時間的な前後関係がないためです。原文全体を参照して訳文を生成することは、答えを漏らすことになりません。

補足します。翻訳モデルでは、訳文の側には自己注意とマスキングが適用されます。

訳文を1語ずつ生成する際、未来の単語を見てはいけないためです。

12 マルチヘッドアテンション

一つでは足りない

ここまで扱ったのは、アテンションヘッド1個分です。

しかし、文脈が単語の意味に影響する方法は、形容詞と名詞の関係だけではありません。

さまざまな影響

例を挙げます。

「they crashed the」が「car」の前にあれば、その車の形状や状態に意味を持ちます。

同じ文章に「wizard」があれば、「Harry」はハリー・ポッターを指すかもしれません。一方、「Queen」「Sussex」「William」があれば、王子のほうを指すでしょう。

ヘッドごとに違う役割

想定される文脈の更新方法ごとに、キー行列とクエリ行列のパラメータは異なります。異なるアテンションパターンを捉えるためです。

そしてバリュー写像のパラメータも、何を足すべきかに応じて異なります。

実際の振る舞いは読めない

繰り返しになりますが、これらの写像の真の振る舞いは解読が困難です。

次のトークンを予測するという目的を最もよく達成するために、モデルが必要とする通りに重みが設定されます。

並列に動かす

Transformerのアテンションブロックは、マルチヘッドアテンションで構成されます。

多数のヘッドを並列に動かし、それぞれが独自のキー、クエリ、バリュー写像を持ちます。

GPT-3では、各ブロックに96個のアテンションヘッドがあります。

具体的に書き出す

96個あるということは、次を意味します。

項目個数
キー行列とクエリ行列96組
アテンションパターン96個
バリューベクトルの列96列

それぞれが、対応するアテンションパターンを重みとして足し合わされます。

変化量を合計する

文脈の各位置、つまり各トークンについて、96個のヘッドがそれぞれ変化量を提案します。

それらをすべて足し合わせ、元の埋め込みに加えます。

E_{\text{new}} = E + \sum_{h=1}^{96} \Delta E_h

これが、マルチヘッドアテンションブロックから出てくる、洗練された埋め込みの一つです。

何のためか

多数の異なるヘッドを並列に動かすことで、文脈が意味を変える多様な方法を学習する能力をモデルに与えています。

パラメータを数える

96ヘッドで、それぞれ四つの行列を持ちます。

6300000 \times 96 = 604800000

マルチヘッドアテンションブロック1個で、約6億個です。

練習問題15

GPT-3のアテンションヘッド1個が約630万個のパラメータを持つとき、96ヘッド分の合計を求めてください。

解答を示します。

掛け算を実行します。

6300000 \times 96 = 604800000

約6億480万個です。

補足します。これは1ブロック分です。

GPT-3には96層あるため、さらに96倍されます。次節で計算します。

13 全体の集計

層の数

GPT-3には、96個の異なる層があります。

アテンション関連の合計

キー、クエリ、バリューのパラメータ総数を計算します。

604800000 \times 96 = 58060800000

約580億個です。

全体に占める割合

GPT-3全体は1,750億個です。

\frac{5.8 \times 10^{10}}{1.75 \times 10^{11}} \approx 0.33

約3分の1です。

残りはどこにあるのか

注目を集めているのはアテンションですが、パラメータの多数派は別の場所にあります。

アテンションブロックの間に挟まれる、多層パーセプトロンブロックです。

次の記事で扱います。

練習問題16

アテンション関連のパラメータが約580億個、全体が1,750億個のとき、多層パーセプトロンなど他の部分のパラメータ数を求めてください。

解答を示します。

引き算を実行します。

175000000000 - 58060800000 = 116939200000

約1,169億個です。

割合を計算します。

\frac{1.169 \times 10^{11}}{1.75 \times 10^{11}} \approx 0.668

約67%です。

補足します。前回の記事で、埋め込みと逆埋め込みが約12億個でした。

1169 - 12 = 1157

約1,157億個が、多層パーセプトロンに使われていることになります。

14 なぜアテンションが成功したのか

能力だけではない

アテンション機構の成功を語るとき、どんな振る舞いを可能にしたかに注目しがちです。

しかし、それ以上に大きな理由があります。

並列化できること

アテンションは、極めて並列化しやすい構造です。

GPUを使えば、膨大な計算を短時間で実行できます。

規模が質を変える

過去10年から20年の深層学習における大きな教訓の一つは、規模そのものが性能に大きな質的改善をもたらすように見える、という点です。

したがって、並列化を可能にするアーキテクチャには、大きな利点があります。

前回までの記事とのつながり

この連載の第5回で、Transformer以前のモデルが単語を順番に処理していたことに触れました。

順番待ちが発生すると、GPUの利点を活かせません。アテンションは、この制約を外しました。

研修で伝えるべきこと

技術の成功要因は、能力だけではありません。

計算資源との相性、実装のしやすさ、既存の仕組みとの接続。こうした要素が、実際の普及を左右します。

この視点は、技術選定の場面でも役立ちます。

練習問題17

アテンション機構の成功要因として、並列化のしやすさが重要である理由を説明してください。

解答を示します。

深層学習では、モデルの規模を拡大することで性能が向上する傾向が知られています。

規模を拡大するには、膨大な計算が必要です。GPUは多数の計算を並列に実行することで高速化しますが、順番待ちが発生する処理では、その性能を引き出せません。

アテンションは並列化しやすい構造であるため、GPUの性能を最大限に活用でき、大規模な学習が現実的な時間で可能になりました。

補足します。ただし、コンテキストサイズの二乗に比例する計算量という制約は残ります。

長い文脈を扱う研究が、現在も続いている理由です。

15 研修での教え方

このテーマを扱う場合の設計案です。

90分の構成

時間内容形式
5分解決したい問題。moleの例講義
10分クエリとキー。問いと答え講義
10分内積とアテンションパターン。練習問題5講義と個人
10分Softmax。練習問題6講義と個人
10分マスキング。練習問題8講義と個人
15分バリュー。埋め込みの更新。練習問題11講義と個人
10分マルチヘッド。96個の並列講義
10分パラメータの集計。練習問題15と16個人
10分成功要因と限界。まとめ講義

教える順序の考え方

三つ、意識してください。

一つ目は、問題から入ることです。

moleやtowerの例を先に示してください。「同じ単語が文脈で意味を変える」という問題意識がないまま行列の話を始めると、何のための仕組みか分かりません。

二つ目は、擬人化を活用することです。

名詞が「前に形容詞はいますか」と問いかけ、形容詞が「私は形容詞です」と答える。この説明は、直感的に伝わります。

ただし、実際にそうなっているとは限らないと、必ず添えてください。

三つ目は、三つの行列の役割を表で示すことです。

クエリは問い、キーは答え、バリューは中身。この対応を繰り返し確認してください。混同が最も多い箇所です。

つまずきやすい点

四つあります。

一つ目は、クエリとキーの区別です。

どちらも同じ小さい空間へ写像されるため、違いが見えにくくなります。問いを表すか、答えを表すか。この役割の違いを強調してください。

二つ目は、なぜ内積なのかという疑問です。

内積が方向の一致を測る道具であること、そして重み付き和の形をしていて深層学習と相性が良いこと。この二点を説明してください。

三つ目は、マスキングの必要性です。

学習時に各位置で同時に予測していることを先に説明しないと、なぜ未来を隠す必要があるのか分かりません。

四つ目は、ヘッドの数です。

96個が並列に動くという事実が、イメージしにくいようです。それぞれが異なる種類の関係を捉えている、という説明を加えてください。

想定される質問

質問1 ヘッドごとの役割は決まっているのですか。

回答します。決まっていません。学習の結果として、そうなることが期待されるだけです。実際に調べると、解釈しやすいヘッドもあれば、何をしているか分からないヘッドもあります。

質問2 なぜ96個なのですか。

回答します。設計上の選択です。多いほど多様な関係を捉えられますが、計算量も増えます。試行錯誤で決められた値だと考えてください。

質問3 クエリとキーを分ける必要はあるのですか。

回答します。あります。問いと答えは非対称だからです。AがBに注目することと、BがAに注目することは別です。同じ行列を使うと、この非対称性が表現できません。

質問4 アテンションは人間の注意と同じものですか。

回答します。名前は同じですが、直接の対応はありません。どの情報に重みを置くかを決める仕組みという点で類似していますが、人間の注意の研究から導かれたものではありません。

まとめと次の学習ステップ

アテンション機構は、文脈に応じて単語のベクトルを更新する仕組みです。

三つの行列が働きます。クエリが「何を探しているか」を表し、キーが「何を提供できるか」を表し、バリューが実際に渡す情報を作ります。

クエリとキーの内積が、どの単語がどの単語に関連するかを決めます。その値をSoftmaxで正規化したものが、アテンションパターンです。

そして、バリューベクトルをそのパターンで重み付けして足し合わせ、元の埋め込みに加えます。これで、文脈を吸収した新しい埋め込みができます。

マスキングは、後の単語が前の単語に影響することを防ぎます。Softmaxの前にマイナス無限大を設定することで、正規化を保ったまま遮断できます。

GPT-3では、この処理が96個のヘッドで並列に行われ、それが96層繰り返されます。アテンション関連のパラメータは約580億個、全体の約3分の1です。

そして、この仕組みが成功した最大の理由は、並列化しやすかったことです。規模の拡大が性能を押し上げる分野において、この性質は決定的でした。

次の学習ステップとして、三つを提案します。

一つ目は、多層パーセプトロンブロックに進むことです。パラメータの3分の2が、そこにあります。アテンションほど注目されませんが、量としては主役です。

二つ目は、アテンションパターンを可視化することです。学習済みモデルのパターンを描画するツールが公開されています。どの単語がどの単語を見ているかを観察すると、この記事の内容が目の前で起きます。

三つ目は、位置の符号化を調べることです。この記事では触れませんでしたが、単語の位置をどう表現するかは重要な論点です。絶対位置、相対位置、回転位置埋め込みなど、複数の方式があります。

これまでの連載で、構造、学習、逆伝播、言語モデルの入口と出口、そしてアテンションを追ってきました。残るは多層パーセプトロンと学習の詳細です。土台は、すでにできています。

セイ・コンサルティング・グループでは新人エンジニア研修のアシスタント講師を募集しています。

投稿者プロフィール

山崎講師
山崎講師代表取締役
セイ・コンサルティング・グループ株式会社代表取締役。
岐阜県出身。
海外放浪の末、2000年創業、2004年会社設立。
IT企業向け人材育成研修歴業界歴20年以上。
すべての無駄を省いた費用対効果の高い「筋肉質」な研修を提供します!
この記事に間違い等ありましたらぜひお知らせください。

学生時代は趣味と実益を兼ねてリゾートバイトにいそしむ。長野県白馬村に始まり、志賀高原でのスキーインストラクター、沖縄石垣島、北海道トマム。高じてオーストラリアのゴールドコーストでツアーガイドなど。現在は野菜作りにはまっている。