Multi-Headは全部必要か？ — Attentionヘッドの冗長性・重要度・プルーニング

Multi-Head Attention の教科書的な説明はこうです——「複数のヘッドが異なる観点（文法・意味・位置など）で並行して注目することで、豊かな表現が得られる」。8ヘッドや16ヘッドという設計は、それだけの「観点」が必要だから、と思いたくなります。

ところが2019年、この直感を揺さぶる実験結果が相次いで報告されました。訓練済みTransformerのヘッドの大半は、推論時に削除しても性能がほとんど落ちないのです。Michel らの論文タイトルはそのものずばり “Are Sixteen Heads Really Better than One?”（16ヘッドは本当に1ヘッドより良いのか？）。機械翻訳のTransformerやBERTで、大多数のヘッドを1個ずつ外しても性能低下はほぼゼロ、層によっては1ヘッドだけ残して全部削っても性能を維持できるケースすらありました。

上の図がこの記事の主題です。訓練済みモデルのヘッドを調べると、実際に働いているのは少数の「重要ヘッド」（色付き）で、残りの大半（グレー）は削ってもほぼ無傷。この記事では、この冗長性を自分の手で再現しながら、ヘッドの重要度をどう測るか、どこまで削れるか、生き残るヘッドは何をしているのかを掘り下げます。

なぜこれを学ぶ価値があるのでしょうか。第一に、モデル圧縮・推論高速化の基礎だからです。ヘッドプルーニングは、性能をほぼ保ったまま計算量とメモリを削る実用技術で、MQA/GQAのようなヘッド共有アーキテクチャにも思想が受け継がれています。第二に、attentionの解釈に直結するからです。「どのヘッドが重要か」を測る技術は、「モデルの中で何が起きているか」を覗く窓になります。

本記事の内容

• ゲート変数 $\xi_h$ によるヘッド除去の定式化
• ヘッド重要度の2つの測り方 — マスク法（leave-one-out）と勾配法（テイラー近似）の導出
• 小型Transformerでの再現実験 — 冗長性・重要度の一致・累積プルーニングカーブを実測
• 生き残るヘッドの役割（Voitaらの分類）と、なぜ冗長になるのかの考察

前提知識

この記事を読む前に、以下の記事を読んでおくと理解が深まります。

• Multi-Head Attention を詳しく解説
• アテンションを直感で理解する

ヘッドを「切れる」ようにする — ゲート変数による定式化

ヘッドを削る実験をするには、まず「ヘッドを削る」操作を数式で定義する必要があります。Multi-Head Attention の出力は、各ヘッドの出力の和として書けます。$H$ 個のヘッドを持つMHAは

$$ \mathrm{MHA}(x) = \sum_{h=1}^{H} \mathrm{Attn}_h(x)\, W_O^{(h)} $$

です（$W_O^{(h)}$ は出力射影 $W_O$ のヘッド $h$ に対応するブロック）。「結合(concat)して射影」という通常の説明と等価ですが、和の形で書くとヘッドが独立な加算的部品であることが見やすくなります。

ここに、ヘッドごとのゲート変数 $\xi_h$ を挟みます。

$$ \mathrm{MHA}(x) = \sum_{h=1}^{H} \xi_h \cdot \mathrm{Attn}_h(x)\, W_O^{(h)}, \qquad \xi_h \in \{0, 1\} $$

$\xi_h = 1$ なら通常動作、$\xi_h = 0$ ならヘッド $h$ の出力が完全に消え、そのヘッドをプルーニング（刈り取り）した状態になります。再訓練は不要で、スイッチを切るだけです。

図はゲート付きMHAの模式図です。各ヘッドの出力にスイッチ $\xi_h$ が付いていて、$\xi_3 = 0$ にするとヘッド3だけが回路から切り離されます。この $\xi$ は単なる実験装置ではなく、次節で見るように重要度を微分で測るためのプローブとしても働きます。

ヘッド重要度の2つの測り方

ヘッドを削る前に、「どのヘッドが重要か」を測る方法が要ります。代表的な方法は2つあります。

方法1: マスク法（leave-one-out）。 最も素直な方法です。ヘッド $h$ を1個だけ切って（$\xi_h = 0$）、評価データで性能を測り直す。重要度は精度低下 $\Delta \mathrm{Acc}_h = \mathrm{Acc}(\text{full}) – \mathrm{Acc}(\xi_h = 0)$ で定義します。正確ですが、ヘッドの数だけ評価を回す必要があり、大きなモデルでは高コストです。

方法2: 勾配法（テイラー近似）。 Michel らが使った効率的な方法です。損失 $L$ を $\xi_h$ の関数と見て、$\xi_h = 1$（ヘッドあり）から $\xi_h = 0$（ヘッドなし）に変えたときの損失変化を1次のテイラー展開で近似します。

$$ L(\xi_h = 0) – L(\xi_h = 1) \approx -\frac{\partial L}{\partial \xi_h}\bigg|_{\xi_h=1} \cdot (0 – 1) = \frac{\partial L}{\partial \xi_h} $$

この変化の大きさの期待値をとったものが勾配ベース重要度です:

$$ I_h = \mathbb{E}_{x}\left| \xi_h \frac{\partial L(x)}{\partial \xi_h} \right| $$

順伝播と逆伝播を1回ずつ回すだけで全ヘッドの重要度が一括で手に入ります。$\xi$ を掛け算で挟んでおいたおかげで、「このヘッドを消したら損失がどれだけ動くか」が普通の勾配計算に化けるわけです。

図の左がマスク法（正確・高コスト）、右が勾配法（近似・低コスト）です。後の実験で、この2つの順位がよく一致することを確かめます。道具が揃ったので、実際に小型Transformerを訓練して冗長性を観察しましょう。

実験セットアップ — 2層×4ヘッドのTransformerを訓練する

再現実験の設定です。「トークン7が現れた後にトークン3が現れるか」を判定する合成系列分類タスク（語彙10、系列長16、ラベル50/50）を作り、ゲート付きMHAを組み込んだ小型Transformer（2層×4ヘッド、$d_{\text{model}}=64$、CLSトークンで分類）を訓練します。

import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(42)
np.random.seed(42)

VOCAB, CLS, SEQ_LEN = 10, 10, 16
A_TOK, B_TOK = 7, 3

def make_dataset(n, rng):
    xs, ys = [], []
    n_pos = n_neg = 0
    while len(xs) < n:
        seq = rng.integers(0, VOCAB, size=SEQ_LEN)
        pos_a = np.where(seq == A_TOK)[0]
        pos_b = np.where(seq == B_TOK)[0]
        label = int(len(pos_a) > 0 and len(pos_b) > 0
                    and pos_a.min() < pos_b.max())
        if label == 1 and n_pos < n // 2:
            xs.append(seq); ys.append(1); n_pos += 1
        elif label == 0 and n_neg < n - n // 2:
            xs.append(seq); ys.append(0); n_neg += 1
    return torch.tensor(np.array(xs)), torch.tensor(ys)

class GatedMHA(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.h, self.dh = n_heads, d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model)
        self.out = nn.Linear(d_model, d_model)
        self.gates = torch.ones(n_heads)   # ゲート ξ_h（差し替え可能）

    def forward(self, x):
        B, T, D = x.shape
        qkv = self.qkv(x).view(B, T, 3, self.h, self.dh).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]
        att = ((q @ k.transpose(-2, -1)) / math.sqrt(self.dh)).softmax(dim=-1)
        head_out = (att @ v) * self.gates.view(1, self.h, 1, 1)  # ← ξ_h を掛ける
        return self.out(head_out.transpose(1, 2).reshape(B, T, D))

GatedMHA の要点は1行だけ——各ヘッドの出力に self.gates（$\xi_h$）を掛けている部分です。このタスクは「7の位置と3の位置の前後関係」を見る必要があるので、attentionが本質的に働きます。12エポック訓練すると、テスト精度は 1.000（完全正解）に達しました。

訓練曲線です。損失は順調に下がり、検証精度は5エポック目あたりで1.0に到達します。フル性能のモデルが手に入ったので、いよいよヘッドを1個ずつ切ってみます。

実験1: ヘッドを1個ずつ外す — 大半は削ってもほぼ無傷

全8ヘッド（2層×4ヘッド）について、1個だけ $\xi_h = 0$ にしてテスト精度の低下を測ります。

def accuracy(model, X, y):
    model.eval()
    with torch.no_grad():
        return (model(X).argmax(-1) == y).float().mean().item()

acc_full = accuracy(model, X_test, y_test)   # 1.000
loo_drop = np.zeros((2, 4))
for l in range(2):
    for h in range(4):
        g = torch.ones(2, 4)
        g[l, h] = 0.0
        model.set_gates(g)                    # ヘッド(l,h)だけ切る
        loo_drop[l, h] = acc_full - accuracy(model, X_test, y_test)
model.reset_gates()

結果は論文の主張を鮮やかに再現しました。8ヘッド中6ヘッドは、外しても精度低下が 0.05未満（うち5ヘッドは0.02未満。第1層ヘッド1で0.005、第2層ヘッド4で0.004など、ほぼ誤差）。目立って痛いのは第2層ヘッド2（低下0.107）と第1層ヘッド3（低下0.086）の2つだけです。テスト精度1.000のモデルでも、実際に仕事をしているのは一部のヘッドで、残りは「いてもいなくても変わらない」ことが数字で見えました。

では、この重要度は勾配法でも同じように見えるのでしょうか。

実験2: 勾配重要度はマスク法と一致するか

勾配ベース重要度 $I_h = \mathbb{E}|\xi_h\, \partial L/\partial \xi_h|$ を計算し、マスク法の精度低下と突き合わせます。

def grad_importance(model, X, y, bs=512):
    model.eval()
    imp = torch.zeros(2, 4)
    n_batches = 0
    for i in range(0, len(X), bs):
        gates = [torch.ones(4, requires_grad=True) for _ in range(2)]
        for l, blk in enumerate(model.blocks):
            blk.mha.gates = gates[l]
        loss = F.cross_entropy(model(X[i:i+bs]), y[i:i+bs])
        loss.backward()
        for l in range(2):
            imp[l] += (gates[l] * gates[l].grad).abs().detach()
        n_batches += 1
    model.reset_gates()
    return imp / n_batches

横軸が勾配ベースの相対重要度、縦軸がマスク法の精度低下です。右上に浮かんでいるのが先ほどの重要ヘッド2つ（L2H2・L1H3）で、2つの方法が同じヘッドを指差しています。順位の一致度はスピアマン相関 0.833。1回の逆伝播で済む勾配法が、8回の再評価が必要なマスク法の良い代用になることが確認できました。

層×ヘッドのヒートマップでも見ておきましょう。

左が勾配重要度、右がマスク法。どちらの地図でも同じマスが濃く光っており、重要度の構造が方法によらず安定していることがわかります。1個ずつ外すだけなら無傷——ではまとめて外していったらどこで壊れるのでしょうか。

実験3: 累積プルーニング — 崖はどこにあるか

重要度の低い順にヘッドを1個ずつ累積的に削除し、テスト精度の変化を追います。比較対象として、ランダムな順で削除した場合（5回平均）も測ります。

order_imp = np.argsort(imp.numpy().ravel())   # 重要度の低い順
accs = [acc_full]
g = torch.ones(2, 4)
for flat in order_imp:
    g[flat // 4, flat % 4] = 0.0
    model.set_gates(g)
    accs.append(accuracy(model, X_test, y_test))

このカーブが本記事のハイライトです。重要度順の削除（オレンジ）では、3ヘッド削っても精度 0.995 とほぼ無傷。8ヘッドの4割近くを消しても性能は実質満点のままです。4個目から緩やかに下がり始め（0.891）、6個目で 0.703、そして残り1個になると 0.582——重要ヘッドに手を付けた途端に崖が現れます。一方ランダム順（グレー）は1個目から 0.932 と即座に下がり始めます。「どれを削るか」を重要度で選ぶことがプルーニングの成否を分けることが、2本のカーブの差として現れています。

削っても平気なヘッドと、削ると崩れるヘッド。中身はどう違うのでしょうか。attention パターンを覗いてみます。

働くヘッドと遊ぶヘッド — 注意パターンの中身

勾配重要度が最大のヘッドと最小のヘッドについて、正例系列（7の後に3がある系列）に対する注意重み行列を可視化します。

左の最重要ヘッド（第2層ヘッド2）は、タスクに関係するトークン（7や3の位置）の列に注意が集中する構造的なパターンを持っています。分類に使うCLS行の注意も特定の位置に鋭く立っています。右の最も暇なヘッド（第2層ヘッド4）は、注意がぼんやり広がっているだけで、タスク固有の構造が見えません。「削っても平気」の正体は、そもそもタスクに効く情報を運んでいないことなのです。

Voita ら（ACL 2019）は機械翻訳Transformerでこの分析を体系的に行い、生き残る重要ヘッドには解釈可能な役割があることを示しました。

役割は大きく3種類——①位置ヘッド（常に隣接トークンを見る）、②構文ヘッド（動詞→主語のような文法的依存関係を見る）、③レアトークンヘッド（文中の低頻度語に注目する）。Voitaらは機械翻訳Transformerのエンコーダで体系的なプルーニングを行い、48ヘッド中38ヘッドを削ってもBLEUの低下が0.15に留まること、そして最後まで生き残るのがこの3種の「役割持ち」ヘッドであることを報告しています。冗長なヘッドを刈ると、機能の分業が浮かび上がるわけです。

では、Voitaらはどうやって「体系的に」刈ったのでしょうか。これまでの2手法とは違う、学習で刈る第3のアプローチを見ておきます。

第3の方法 — L0正則化で「学習しながら」刈る

マスク法と勾配法は、どちらも訓練済みモデルを事後的に調べる方法でした。Voitaらのアプローチは発想が違います。ゲート $\xi_h$ を0/1の定数ではなく確率的な変数として扱い、「なるべく多くのゲートが0になれ」という圧力をかけながら追加学習するのです。

理想的には、0でないゲートの個数（$L_0$ ノルム）

$$ L_0(\xi) = \sum_{h=1}^{H} \mathbb{1}[\xi_h \neq 0] $$

をペナルティとして損失に足したいところです。しかし個数を数える関数は階段状で微分できません。そこでVoitaらは Louizos らの hard concrete 分布を使います。各ゲートを、学習可能なパラメータ $\phi_h$ を持つ連続な確率分布からサンプルされる確率変数 $\xi_h \sim q(\xi_h \mid \phi_h)$ とし、この分布が「ちょうど0を取る確率」を解析的に計算できるように設計するのです。すると目的関数は

$$ \mathcal{L} = \mathcal{L}_{\text{task}} + \lambda \sum_{h=1}^{H} \left(1 – q(\xi_h = 0 \mid \phi_h)\right) $$

という微分可能な形になります。第2項は「ゲートが0にならない確率」の合計、つまり期待値の意味で生きているヘッドの本数です。$\lambda$ を大きくするほど強くヘッドが刈られ、タスク性能とのトレードオフを $\lambda$ 一つで調整できます。訓練が終わるとゲートはほぼ0/1に分離し、0のヘッドを物理的に取り除けます。

3つの方法を整理すると、こうなります。

重要な違いは、$L_0$ 法では刈られる過程で残りのヘッドが仕事を引き継げることです。事後プルーニングは「今の分業のまま」切りますが、$L_0$ 法は再配置込みで最適化するため、より深く刈れます。48→10ヘッドという大胆な圧縮がBLEU −0.15で済んだのは、この適応のおかげです。

最後に残る疑問はこれです——そもそも、なぜこんなに冗長になるのでしょうか。

なぜ冗長になるのか — 訓練初期に主役が決まる

Michel らはもうひとつ重要な観察を残しています。ヘッドの重要度分布は訓練のかなり早い段階で決まり、その後は固定化する傾向があるのです。私たちの実験でも、エポックごとに勾配重要度を記録すると同じ現象が見えます。

オレンジの2本（最終的な重要ヘッド）は訓練の序盤から相対重要度を伸ばし、早々に主役の座を確立します。グレーの6本はその後もずっと脇役のままです。この描像は宝くじ仮説（大きなネットワークの中に、初期値の巡り合わせで「当たり」の部分ネットワークが含まれており、訓練はそれを掘り出す）とよく整合します。多数のヘッドは「当たりを引く確率を上げるための試行回数」であり、訓練が終わって当たりが確定した後は、外れ券（冗長ヘッド）を捨てても結果は変わらない——冗長性は無駄ではなく、探索のための保険だったという見方ができます。

この見方は実務にも示唆を与えます。訓練時には多めのヘッドが必要（探索のため）でも、推論時には少数precisionヘッドで足りる（当たりは確定済み）。訓練と推論でアーキテクチャ要件が非対称だからこそ、事後プルーニングやMQA/GQAのようなヘッド削減が成立するのです。

まとめ

Multi-Head Attention の「多数のヘッド」の実態を、定式化・重要度測定・プルーニング実験で確かめました。

• ゲート変数 $\xi_h$ を挟むと、ヘッド除去がスイッチ操作になり、重要度が勾配 $I_h = \mathbb{E}|\xi_h\, \partial L/\partial \xi_h|$ で測れる
• 冗長性の再現：テスト精度1.000の2層×4ヘッドモデルで、8ヘッド中6ヘッドは外しても精度低下0.02未満。痛いのは2ヘッドだけ（最大低下0.107）
• 2つの重要度は一致：勾配法とマスク法の順位はスピアマン相関0.833。1回の逆伝播で重要度マップが得られる
• 累積プルーニング：重要度順なら3ヘッド削っても精度0.995とほぼ無傷、重要ヘッドに達すると崖。ランダム順は最初から劣化
• 生き残るヘッドには役割がある：位置・構文・レアトークン（Voita）。重要度は訓練初期に決まり、冗長ヘッドは「探索の保険」と解釈できる

「ヘッドが多い=観点が多い」という素朴な理解から、「ヘッドは冗長に用意され、少数が仕事をする」という実像へ。この視点は、モデル圧縮だけでなくattentionの解釈研究の出発点にもなります。次のステップとして、以下の記事も参考にしてください。

• Multi-Head Attention を詳しく解説
• MQA と GQA — ヘッドを共有して推論を速くする
• softmaxの温度とattentionエントロピー
• アテンションを直感で理解する