Attentionヘッドを増やしても計算は増えない　ヘッドは層の幅を分けたもの

ヘッドを増やせばモデルはそれだけ多くを見られる、と考えるのは、同じ予算のまま担当チームを8つから16に分ければ成果が倍になると考えるのと同じ計算違いです。Attentionのヘッドは層に追加される部品ではなく、層の幅を分割したものだからです。ヘッドを8から16に増やしても、各ヘッドの次元が半分になるだけで、計算量とパラメータ数はほぼ変わりません。この記事では、Attention層とAttentionヘッドの関係を、この分割の見方から整理します。

層は処理のブロック、ヘッドはその中の分割単位です

Transformerは、Vaswaniら（2017年）が提案した構造です。エンコーダーは複数の層を積み重ねたもので、各層はマルチヘッド・セルフアテンションとフィードフォワード・ニューラルネットワークでできています。Attention層は、各層の中でトークンどうしの関係を捉える部分です。

Attention層の計算は三段です。入力を重み行列 W_Q、W_K、W_V で変換してクエリ、キー、バリューを作る。クエリとキーの内積を √d_k で割り、ソフトマックスで重みにする。その重みでバリューを足し合わせる。Attentionヘッドは、この計算を独立に行うサブユニットで、自分専用の重みを持ち、ほかのヘッドと共有しません。最後に各ヘッドの出力を連結し、重み行列 W_O で変換して、層としての出力にまとめます。

ヘッドを増やすと、一つあたりの次元が減ります

ヘッドの次元は、入力の次元 d_model をヘッド数 h で割って決めます（d_k ＝ d_model ÷ h）。d_model が512の場合を並べます。

ヘッド数	各ヘッドの次元	W_Q 全体の大きさ
1	512	512 × 512
8	64	512 × 512（64列を8つ）
16	32	512 × 512（32列を16個）

W_Q、W_K、W_V の全体の大きさは、ヘッド数によらず同じです。スコアの計算も、ヘッドごとの量は次元に比例して小さくなるので、全ヘッドを合わせた量はほぼ変わりません。ヘッドは計算を足すのではなく、同じ計算を細かく分けています。

増えるのは計算量ではなく、注意の配り方の数です

ヘッドが1つなら、注意の分布は一つで、各トークンは一通りの重みでしか情報を集められません。ヘッドが8つなら、一つの層の中で8通りの重みで同時に集められます。文法的な関係を拾う重みと、意味の関連を拾う重みを、別々に持てるということです。

代わりに失うものもあります。一つのヘッドが扱える次元は小さくなり、スコア行列はヘッドごとに作るのでその分のメモリは増えます。独立しているので並列に計算でき、互いの配り方に引きずられませんが、学習で必ず違う役割に分かれる保証はありません。

GPT-3では、96層の各層に96のヘッドがある構成が示されており、総数は9,216です。似た働きのヘッドも出てくるため、ヘッドを刈り込む手法や、ヘッドどうしで一部を共有する手法が研究されています。ヘッドの数は、多いほどよいものではなく、幅の分け方の設計です。

担当を増やしても、総量が同じなら深さが減ります

組織でも同じ引き換えが起きます。新しい観点が必要になるたびに担当を一つ増やし、総工数はそのまま、という運用は、ヘッドを増やして一つあたりの次元を削るのと同じ形をしています。観点の数は増えますが、一つの観点に割ける深さは確実に減ります。拙著『GoogleのAI NotebookLM 意思決定の強化書』で、前提の要素として制約を独立に書き出すのは、この総量の制約が見えないまま観点だけが増えると、決めた後で破綻するからです。観点を足すときは、総量が増えるのか、既存の観点から削るのかを同時に書きます。

この比較が当てはまらない場面もあります。総量そのものを増やせるなら、ヘッドを増やしても一つあたりの深さは保てます。実際、大きなモデルは d_model とヘッド数を一緒に大きくしています。引き換えが生じるのは、総量が固定されているときです。

AIに頼むときも、一度の指示に十の観点を詰め込めば、それぞれの扱いは浅くなりがちです。どの観点を並べて読ませ、どれを深く掘らせるのかを分けて伝え、最後にどんな形でまとめるかを指定します。

あなたが最近AIに頼んだ指示には、観点がいくつ入っていたでしょうか。そのうち、本当に深く掘ってほしかったのはどれでしたか。

この記事の背景にある考え方は、拙著『GoogleのAI NotebookLM 意思決定の強化書』の4章4節「意思決定の設計図（前提リスト）を完成させる」に書きました。対象、目的、制約、評価基準を一枚の前提リストにまとめ、項目ごとに結論への影響度と根拠になる資料を書き添える節で、まとめた中身は、制約と評価基準が矛盾していないかを含めて人が目で確かめます。そこには「ここはAIに丸投げしません。」とあります。この節を読むと、観点を一つ足すとき、それが総量という制約とぶつかっていないかを、足す前に自分で確かめられるようになります。