アテンションのヘッドが記憶を共有し、それを何倍も小さくする仕組み。 · Source

会話の記憶がふくらむ——だからヘッドはそれを共有する。

各アテンションのヘッドは、これまでの一語ごとに自分用のメモ——KeysとValues——を抱える。ヘッドも層も語数も積み重なれば、毎ステップ読み返すその山が本当のボトルネックになる。解決策はちょっと小生意気だ。ほとんどのヘッドは自分専用のコピーなど要らない。いくつかの共有のびんに何本ものストローを挿すように、各ヘッドはなお自分なりに吸い上げるが、その記憶は皆で共有している。

どのヘッドも過去のコピーを自分専用に抱え込む。

M = 2\,L\,n\,H\,d_{\text{head}}\,b \;=\; 2\,L\,n\,d_{\text{model}}\,b

キャッシュをモデルの会話の記憶だと思ってほしい。標準的なアテンションでは、どのヘッドも自分のKeysとValuesを——層ごと、語ごとに——抱える。長い宴会テーブルのどの席にも同じびんが置かれているように、同じ形が並びに沿って何度も複製される。平たく言えば、記憶はヘッド×層×語で増え、次の一語を足すためだけにその全部を読み返す。

もし全ヘッドが一つの記憶を共有したら？

\mathrm{head}_i = \mathrm{softmax}\!\left(\dfrac{Q_i K^\top}{\sqrt{d_{\text{head}}}}\right) V \qquad i = 1,\dots,H

ここで大胆な一手。各ヘッドのQuery——自分の問い——はそのままに、すべてのヘッドが一つの共有KeyとValueから読むようにする。一台のエスプレッソマシンがバリスタの長い列をさばくように、各人はなお自分のカップを淹れるが、源は一つ。平たく言えば、問いは各ヘッド固有のまま、共有されるのは記憶だけ——だからヘッド数が増えても記憶は増えない。

だが共有の記憶が一つだけでは大ざっぱすぎる。

一つの記憶をすべてのヘッドで共有するのはやり過ぎだ。各ヘッドは自分の問いに合わせたメモを持っていた。単一の組に押し込めば、鋭いものほどぼやけ、学習も揺らぎかねない。合唱団まるごとが一本のマイクに群がるように、歌声は聞こえても、それぞれの声が同じチャンネルに平板化する。記憶は節約できた——が、モデルの幅を静かにいくらか失った。

なら一つを共有せず、小さなグループで共有する。

\mathrm{head}_i = \mathrm{softmax}\!\left(\dfrac{Q_i\,K_{g(i)}^\top}{\sqrt{d_{\text{head}}}}\right) V_{g(i)},\qquad g(i)=\left\lceil iG/H\right\rceil,\quad 1\le G\le H

解決策は中をとる。ヘッドをいくつかのグループに分け、同じグループのヘッドは一つのKeyとValueを共有しつつ、グループ同士は別のまま。夜明けの相乗りバンのように、各自に一台でも、全員に一台のバスでもなく——数台のバンに、それぞれ小さな一団。一つのつまみが決める——H個のヘッドに対しG個の共有記憶として、G = Hはふつうのアテンション、G = 1は大ざっぱな極端、その間の数グループがちょうどいい塩梅。

一語ごとに、運ぶ量が八分の一。

M = 2\,L\,n\,n_{kv}\,d_{\text{head}}\,b,\qquad \text{shrink}=\dfrac{H}{n_{kv}}\quad (n_{kv}=H,\,G,\,1)

さて、その見返り。H個ではなくG個の共有記憶にすれば、キャッシュはH ∕ G倍——しばしば八分の一——に縮む。しかも二重に効く。長い文脈が収まり、さらに一語生成するたびにモデルはその記憶ぜんぶを遅い記憶域から読み直す。配送バンが同じ箱を八つではなく、共有の一箱だけ運ぶように、どの往復も軽くなる。記憶は小さく、往来は軽く、答えは同じ。

しかも、すでに訓練済みのモデルを後から作り替えられる。

K^{(g)} = \dfrac{G}{H}\sum_{i\,\in\,g} K_i \qquad \big(\,V^{(g)}\text{ likewise},\ |g|=H/G\,\big)

いちばんの利点は、ゼロから訓練し直さなくていいこと。各ヘッドが自分の記憶を持つ完成済みモデルを取り、各グループの記憶を平均して一つに落とす——あとは少しだけ微調整して馴染ませる。いくつもの細い流れが一つの淀みに合わさるように、ほぼ同じ無数の筋が一つになる。うまくいくのは、グループ内のそれらの記憶が初めからほぼ同じだったから——共有してもほとんど損なわれない理由がまさにそれだ。

百のヘッド。ひと握りの記憶。

グループ化クエリがうまくいくのは、何十ものヘッドが実はほぼ同じメモを抱えていたから——自分の問いを立てるには十分ちがい、答えを共有するには十分そっくり。では、本当はどれほど別物だったのか。ヘッドの個性は、覚えていることではなく、問うことだけにあったのかもしれない。🌱 ならば、視点はいったいどこに宿るのだろう。