Attention (Q, K, V) = softmax (\frac{Q K ^{⊤}}{d _{k}}) V

トランスフォーマーの注意機構

Scaled dot-product attention

論文の題は、ビートルズの歌から 5 秒で思いついたという。

文のどの言葉に目を向けるかを、クエリとキーの内積から softmax で重みにして決める式。√dₖ で割るのが要で、翻訳から大規模言語モデルまでの土台になった。

年代: 2017年
人物: アシシュ・ヴァスワニ、ノーム・シャジール、ニキ・パーマー、ヤコブ・ウスコライト、リオン・ジョーンズ、エイダン・ゴメス、ウカシュ・カイザー、イリア・ポロスキン、ドミトリー・バダナウ、チョ・ギョンヒョン、ヨシュア・ベンジオ

数学情報科学言語学物理学生物学

逸話

5 秒で思いついた題

2017 年の春、カリフォルニア州マウンテンビューの Google で、研究者たちが 1 本の論文を仕上げていた。出す先は、その年の 12 月に開かれる機械学習の国際会議 NIPS(いまの NeurIPS)で、締め切りは 5 月 19 日だった。記者スティーブン・レヴィは 2024 年、8 人全員に話を聞き、『WIRED』に記事を書いた。締め切り前の 2 週間を、インターンだったエイダン・ゴメスは「みんな寝ていなかった」と振り返る。

締め切りの 2 晩ほど前、ヤコブ・ウスコライトが、論文に題が要ると気づいた。チームは、それまで翻訳の定番だった LSTM(長短期記憶)などの再帰型ニューラルネットワークを手放し、注意(attention)という 1 つの仕組みに賭けていた。リオン・ジョーンズは、ビートルズの歌「愛こそはすべて」の原題「All You Need Is Love」を思い出し、「Attention Is All You Need」ではどうかと言った。ジョーンズはこう話す。「私はイギリス人だから。考えたのは文字どおり 5 秒ほど。使われるとは思わなかった」。

英語からフランス語への翻訳の数字が出たのは、提出の 5 分ほど前だったと、ニキ・パーマーは語っている。記事によれば、論文を送ったのは締め切りの 2 分ほど前だった。

8 人の名前には、すべて同じ印が付いている。脚注は「貢献は同等。並びは無作為」と書き、続けて一人ひとりの役目を記した。再帰型のネットワークを自己注意に置き換えることを提案したのはウスコライト。最初のモデルを設計して組んだのはアシシュ・ヴァスワニとイリア・ポロスキン。この記事の式、つまり内積を $d_{k}$ で割る注意と、それを何組も並べる多ヘッド注意を提案したのはノーム・シャジールである。

査読の評価は、パーマーによれば「好意的」「とても好意的」「まあ良い」で、発表は夜のポスターの枠になった。

起源

1 行の重み付け

注意の考えは、2014 年の機械翻訳の研究にさかのぼる。当時のニューラル機械翻訳は、元の文を再帰型ニューラルネットワーク(RNN)で 1 語ずつ読み、決まった長さの 1 本のベクトルに詰めてから、訳文を 1 語ずつ書き出していた。文が長いと、訳の質が急に落ちた。

2014 年 9 月、ドイツのジェイコブス大学ブレーメンのドミトリー・バダナウと、モントリオール大学のチョ・ギョンヒョン、ヨシュア・ベンジオは、1 本への詰め込みが壁だと考え、それをやめる論文を出した。訳語を 1 つ書くたびに、元の文の各単語に重み $α_{ij}$ を割り振り、重み付きの平均をとって使う。重みは

α_{ij} = \frac{exp ( e _{ij} )}{\sum _{k} exp ( e _{ik} )}

で、 $e_{ij}$ は小さなニューラルネットワークが出す、訳語 $i$ と元の単語 $j$ の合い具合である。論文は $e_{ij}$ を「エネルギー」と呼び、こう書いた。「直感的には、これは復号器に注意の仕組みを持たせることである。復号器は、元の文のどの部分に注意を払うかを決める」。英語の European Economic Area は、フランス語では語順が逆の zone économique européenne になる。モデルは zone を訳すとき、2 語を飛び越えて Area に重みを置き、そこから 1 語ずつ戻った。

2015 年には、スタンフォード大学のミン=タン・ルオンらが、2 つのベクトルの内積で合い具合を測る簡単な方法も試した。それでも注意は RNN に付け足す部品で、1 語ずつ進む計算は並列にできなかった。

重みを作る関数 $e^{x_{j}} / \sum_{i} e^{x_{i}}$ は、もっと古い。1902 年、アメリカの物理学者ウィラード・ギブズは『統計力学の基本原理』で、エネルギー $ε$ の状態の確率を $e^{(ψ - ε) /Θ}$ とする分布を「カノニカル(正準)」と名づけ、割る数 $Θ$ が温度に似た性質を持つことを示した。合い具合を、符号を変えたエネルギーと読めば、重みは正準分布そのものである。softmax の呼び名は、イギリスの王立信号レーダー研究所のジョン・ブライドルが 1989 年の NIPS に出した論文にある。ネットワークの出力をこの形で確率にそろえ、最大の値を 1 つ選ぶ操作を微分できるようにしたものなので soft max と呼びたい、と書いた。

2017 年 6 月 12 日に Google Brain と Google Research で研究した 8 人(1 人はトロント大学の学生)が arXiv に出した論文は、RNN を丸ごと取り除いた。文の各単語が、自分の「クエリ(query)」のベクトルで、すべての単語の「キー(key)」のベクトルと内積をとり、その重みで「値(value)」のベクトルを混ぜる。これを行列で書いたのが、この記事の式である。割らない内積の注意は、 $d_{k}$ が大きいと足し算型の注意に負けるという報告があった。論文は理由を「 $d_{k}$ が大きいと内積の絶対値が大きくなり、softmax が勾配のきわめて小さい領域に押しやられるのではないか」と推し量り、脚注で内積の分散が $d_{k}$ になることを示した(証明)。

論文のモデルは、符号器と復号器にそれぞれ 6 層を重ね、各層のベクトルは 512 次元である。注意は 8 つの「ヘッド」に分かれ、各ヘッドは 512 を 8 で割った $d_{k} = 64$ 次元で計算する。GPU 8 枚の計算機 1 台で、基本のモデルは 12 時間、大きいモデルは 3.5 日学習した。機械の訳を人の訳と比べる点数 BLEU で、英語からドイツ語への翻訳(WMT 2014)は 28.4 となり、複数のモデルを組み合わせたものを含めて、それまでの最高を 2 点以上上回った。英語からフランス語は、6 月の最初の版で 41.0、のちの改訂版の要旨と表で 41.8 である。会議の論文集のウェブページの要旨は、ドイツ語 27.5・フランス語 41.1 である。

2020 年の言語モデル GPT-3 は、96 層・96 ヘッド・各ヘッド $d_{k} = 128$ のトランスフォーマーで、パラメーターは 1750 億個あった。レヴィの記事は、いまの対話型の AI や画像を作る AI の多くがトランスフォーマーを土台にしている、と書いている。

証明

クエリのベクトル $q$ と、 $n$ 個のキー $k_{1}, \dots, k_{n}$ 、バリュー(値) $v_{1}, \dots, v_{n}$ を考える。 $Q$ の 1 行が 1 つのクエリで、式は各行の計算をまとめて書いたものである。 $d_{k}$ はクエリとキーの次元である。

クエリとキーの内積を $d_{k}$ で割って、点数にする。
$s_{j} = \frac{q \cdot k _{j}}{d _{k}}$
点数を softmax で重みに変える。 $e^{x}$ はいつも正なので、重みはすべて正で、足すと 1 になる。重みは $n$ 個のキーの上の確率分布である。
$a_{j} = \frac{e ^{s_{j}}}{\sum _{i = 1}^{n} e ^{s_{i}}}, a_{j} > 0, j = 1 \sum n a_{j} = 1$
出力は、値の重み付きの平均である。重みを確率とみた、値の期待値でもある。
$o = j = 1 \sum n a_{j} v_{j}$
2 つの重みの比は、点数の差だけで決まる。すべての点数に同じ数を足しても、重みは変わらない。
$\frac{a _{j}}{a _{i}} = e^{s_{j} - s_{i}}$
割る数を $τ > 0$ と書き、内積を $z_{j} = q \cdot k_{j}$ とする。分母と分子を $e^{z_{j} / τ}$ で割ると
$a_{j} = \frac{1}{\sum _{i = 1}^{n} e ^{(z_{i} - z_{j}) / τ}}$
$z_{j}$ がただ 1 つの最大なら、 $τ \to 0$ で $i \neq = j$ の項はすべて 0 に近づき、 $a_{j} \to 1$ となる。最大の 1 つだけを選ぶ操作に戻る。逆に $τ \to \infty$ では、どの項も 1 に近づく。
$τ \to 0 lim a_{j} = 1, τ \to \infty lim a_{j} = \frac{1}{n}$
重みを点数で微分する。分数の微分の公式から、 $i = j$ のとき $a_{j} - a_{j}^{2}$ 、 $i \neq = j$ のとき $- a_{j} a_{i}$ が出る。 $δ_{ij}$ は $i = j$ なら 1、そうでなければ 0 である。
$\frac{\partial a _{j}}{\partial s _{i}} = a_{j} (δ_{ij} - a_{i})$
1 つの重みが 1 に近く、ほかが 0 に近いと、7 のどの項も 0 に近い。点数を少し動かしても重みはほとんど動かず、学習の手がかりになる勾配が消える。論文の言う「勾配のきわめて小さい領域」はここである。
$a_{j} \approx 1 \Rightarrow a_{j} (1 - a_{j}) \approx 0$

次に、なぜ $d_{k}$ で割るのかを示す。論文の脚注にならい、1 つのクエリ $q = (q_{1}, \dots, q_{d_{k}})$ と 1 つのキー $k = (k_{1}, \dots, k_{d_{k}})$ の成分は、すべて互いに独立で、平均 0・分散 1 とする。

独立な量の積の平均は、平均の積である。内積の各項の平均は 0 になる。
$E [q_{i} k_{i}] = E [q_{i}] E [k_{i}] = 0$
分散 1・平均 0 なので $E [q_{i}^{2}] = E [k_{i}^{2}] = 1$ 。 $q_{i}^{2}$ と $k_{i}^{2}$ も独立なので、各項の分散は
$Var (q_{i} k_{i}) = E [q_{i}^{2} k_{i}^{2}] - (E [q_{i} k_{i}])^{2} = E [q_{i}^{2}] E [k_{i}^{2}] = 1$
番号の違う項どうしは独立なので、和の分散は分散の和になる。内積の平均は 0、分散は $d_{k}$ である。
$Var (q \cdot k) = i = 1 \sum d_{k} Var (q_{i} k_{i}) = d_{k}$
$d_{k}$ で割れば、次元によらず分散は 1 に戻る。
$Var (\frac{q \cdot k}{d _{k}}) = \frac{d _{k}}{( d _{k} ) ^{2}} = 1$

割らなければ、次元が大きいほど点数の差が開き、4 の比が指数で開いて、重みは 8 の状態に落ちる。 $d_{k}$ は 5 の $τ$ 、つまり正準分布の温度にあたる。論文はこの割り方を脚注の計算による見当として示したもので、最もよい割り方だと証明したわけではない。

応用

「それ」は何を指すか。 「ねこが魚を食べたそれはおいしかった」の 8 語で、自己注意を手で計算してみる。ベクトルは説明のために手で決めた 4 次元のもので、実在のモデルの値ではない。成分は順に、生き物・食べ物・述語・助詞らしさとする。キーはねこ $(2, 0, 0, 0)$ 、魚 $(1, 2, 0, 0)$ 、食べた $(0, 0, 2, 0)$ 、それ $(0, 0.5, 0, 1)$ 、おいしかった $(0, 1, 1.5, 0)$ で、3 つの助詞はどれも $(0, 0, 0, 2)$ である。「それ」のクエリは、食べ物を探す $q = (0, 3, 0, 0)$ とする。

内積は、魚が 6、おいしかったが 3、それが 1.5、ほかの 5 語が 0 になる。 $d_{k} = 4$ なので 2 で割ると、3・1.5・0.75・0 である。 $e^{3} \approx 20.09$ 、 $e^{1.5} \approx 4.48$ 、 $e^{0.75} \approx 2.12$ 、 $e^{0} = 1$ が 5 つで、和は 31.68。重みは魚が 0.634、おいしかったが 0.141、それが 0.067、ほかの 5 語が 0.032 ずつになる。値のベクトルをキーと同じとすると、「それ」の出力は $(0.70, 1.44, 0.28, 0.26)$ で、もとの $(0, 0.5, 0, 1)$ より食べ物らしさが増えている。魚から流れ込んだ分である。

割らなければ、魚の重みは 0.932 になり、ほかはほとんど 0 になる。8 で割ると 0.217 まで下がり、8 語の重みは平らに近づく(図示)。

次元が 64 なら。 論文の各ヘッドと同じ $d_{k} = 64$ では、成分がランダムでも、内積の標準偏差は 8 ある。キーが 2 つで、内積が標準偏差 1 つ分違うと、割らなければ重みの比は $e^{8} \approx 2981$ 倍で、重みは 0.9997 と 0.0003 に分かれる。証明の 8 の $a (1 - a)$ は 0.0003 しかない。 $64 = 8$ で割れば、比は $e \approx 2.72$ 倍(ネイピア数 e)、重みは 0.731 と 0.269、 $a (1 - a)$ は 0.197 で、学習の手がかりは約 590 倍大きい。

先を読まない仕掛け。 訳文を書く復号器は、まだ書いていない後ろの単語を見てはいけない。論文はその点数を、softmax に入れる前に $- \infty$ に置き換えた。 $e^{- \infty} = 0$ なので重みはちょうど 0 になり、残りの和はやはり 1 である。

長い文の計算。 論文の表 1 によれば、自己注意の 1 層の計算量は $n^{2} d$ 、RNN は $n d^{2}$ に比例する。50 語の文を 512 次元で扱うなら、自己注意が 128 万、RNN が 1311 万で、自己注意のほうが約 10 分の 1 で済む。ただし語の数 $n$ の 2 乗で増えるので、 $n$ が $d$ を超えると逆転する。一度に 2048 個のトークンまで読む GPT-3 では、注意を近くに限らない層の 1 つのヘッドだけで、点数の表は $204 8^{2} \approx 419$ 万マスになる。

図示

「ねこが魚を食べたそれはおいしかった」で、「それ」をクエリにしたときの重み。手作りの 4 次元のベクトルで、内積を √4 = 2 で割った。線の太さと棒の長さが重みを表す。

同じ点数を、1・2・8 で割ったときの 8 語の重み。塗りつぶした棒が魚、点線は 8 語が平らなときの 1/8。割る数が小さいほど 1 語に集まり、大きいほど平らになる。

成分が互いに独立な標準正規分布に従うときの、内積 q·k の分布。左は次元 4・16・64 のままで、次元とともに横に広がる。右は √dₖ で割ったもので、3 つとも分散 1 にそろい、次元が大きいほど正規分布の形に近づく(中心極限定理)。

$d_{k}$ で割ると、成分がランダムなときの点数の散らばりが、次元によらずそろう。形は次元が大きいほど正規分布に近づく。

シミュレーション

左の列の単語をたたくと、それがクエリになり、右の列の各単語への重みが線の太さと棒の長さで出る。「割る数 τ」を小さくすると重みは 1 語に集まり、大きくすると 1/8 ずつの平らな分布に近づく。「ベクトル」をランダムに替えると、成分が平均 0・分散 1 の 16〜256 次元になり、「÷1」と「÷√dₖ」のボタンで、割らないと重みが 1 つに偏ることを確かめられる。「割る数を往復」では $τ$ がひとりでに行き来する。

広がる学問

情報科学 自己注意は、文のどの 2 語の間も 1 回の計算で結ぶ。論文の表 1 では、離れた 2 語をつなぐ道のりの長さが、RNN の $O (n)$ に対して $O (1)$ になる。2020 年の Vision Transformer は、画像を 16×16 画素の小片の列として、同じ式で扱った。
言語学 論文の付録には、6 層のうち 5 層目で、多くのヘッドが動詞 making から離れた more difficult に注意を向ける図や、代名詞 its の指すものを探すらしいヘッドの図がある。2019 年、スタンフォード大学のケヴィン・クラークらは、言語モデル BERT のヘッドに、動詞の目的語や、同じものを指す語句を高い精度で選ぶものがあることを示した。文法を教えずに現れた性質である。
物理学 softmax は統計力学の正準分布と同じ形で、割る数は温度にあたる(ボルツマンのエントロピー)。2020 年、フーベルト・ラムザウアーやゼップ・ホッホライターらは、連続の値をとるホップフィールド・ネットワーク(覚えた模様をエネルギーの谷に蓄える回路)の更新の式が、注意と同じ形になることを示した。ジョン・ホップフィールドは 2024 年、ジェフリー・ヒントンとともにノーベル物理学賞を受けた。
生物学 たんぱく質の立体構造を予測する AlphaFold2(2020 年)は、似た配列を並べた表と、アミノ酸どうしの組の表に、注意を何重にもかける。開発したデミス・ハサビスとジョン・ジャンパーは、2024 年のノーベル化学賞を受けた。
数学重みの和が 1 の平均という形は、1964 年にナダラヤとワトソンが別々に示したカーネル回帰(近いデータほど重くして平均する推定)と同じである。重みは確率分布なので、ばらつきはシャノンの情報エントロピーで測れる。分母の対数 $lo g \sum_{i} e^{s_{i}}$ は $max_{i} s_{i}$ と $max_{i} s_{i} + lo g n$ の間に入り、最大値をなめらかにした量として最適化に使われる。

参考文献

A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, I. Polosukhin, "Attention Is All You Need", Advances in Neural Information Processing Systems 30 (NIPS 2017). arXiv:1706.03762
D. Bahdanau, K. Cho, Y. Bengio, "Neural Machine Translation by Jointly Learning to Align and Translate", ICLR 2015. arXiv:1409.0473
M.-T. Luong, H. Pham, C. D. Manning, "Effective Approaches to Attention-based Neural Machine Translation", EMNLP 2015. arXiv:1508.04025
J. S. Bridle, "Training Stochastic Model Recognition Algorithms as Networks can Lead to Maximum Mutual Information Estimation of Parameters", Advances in Neural Information Processing Systems 2 (1990), 211-217.
J. W. Gibbs, Elementary Principles in Statistical Mechanics, Charles Scribner's Sons, 1902.
S. Levy, "8 Google Employees Invented Modern AI. Here's the Inside Story", WIRED, 2024年3月20日.
K. Clark, U. Khandelwal, O. Levy, C. D. Manning, "What Does BERT Look At? An Analysis of BERT's Attention", BlackboxNLP 2019 (arXiv:1906.04341); H. Ramsauer et al., "Hopfield Networks is All You Need", arXiv:2008.02217 (2020); T. B. Brown et al., "Language Models are Few-Shot Learners", arXiv:2005.14165 (2020).