Transformerベースのモデルは、機械翻訳、文書分類、質問応答など多くのタスクで驚異的な性能を発揮しています。しかし「なぜそのような出力を返したのか?」と問われると、モデル自身は何も答えてくれません。たとえば感情分析モデルが「この映画は素晴らしい演技と脚本が光る作品だ」を「ポジティブ」と判定したとき、モデルは「素晴らしい」に注目したのか、「演技」と「脚本」の組み合わせを捉えたのか — その内部のメカニズムを知りたいというのは自然な欲求です。
ここで手がかりとなるのが、Transformerの中核にあるAttention重みです。Attention重みは「各トークンが他のどのトークンにどれだけ注目しているか」を数値として持っています。この重みを可視化することで、モデルの判断過程を部分的に覗き見ることができます。
Attention重みの可視化と解釈を学ぶことで、以下のようなことが可能になります。
- モデルのデバッグ: 分類誤りの原因を特定する。モデルが無関係なトークンに注目していれば、学習データやモデル構造の改善に繋がります
- モデルの改善: 各ヘッドや各層の役割を分析し、不要なヘッドの枝刈り(pruning)や効率的なアーキテクチャ設計に活用できます
- 安全性と信頼性の向上: 医療や法律などのハイステークスな領域では、モデルの判断根拠を示すことが求められます。Attention可視化はその第一歩です
- 言語構造の発見: Attentionパターンを分析することで、モデルが自然に学習した構文的・意味的関係を発見できます
本記事の内容
- Attention重みの定義と意味の復習
- ヒートマップ、BERTVizなど代表的な可視化手法
- Hugging Face BERTからattention重みを抽出する方法
- matplotlibによる可視化の実装(単一ヘッド、全ヘッド比較、層別比較)
- よく観察されるAttentionパターンの解説
- 注意重みの解釈性の限界に関する議論
- Attention Rolloutの理論と実装
前提知識
この記事を読む前に、以下の記事を読んでおくと理解が深まります。
Attention重みとは何か(復習)
直感的な理解
自然言語の文を読むとき、人間は全ての単語に均等に注意を向けるわけではありません。「彼女はリンゴを食べた」という文で「何を食べたか?」と聞かれたら、「リンゴ」に注意が向きます。一方「誰が食べたか?」なら「彼女」に注目するでしょう。同じ文であっても、問いの種類(つまり文脈)によって注目すべき場所が変わるのです。
Transformerにおけるattention重みは、まさにこの「どこにどれだけ注目するか」を数値化したものです。あるトークン $i$ が別のトークン $j$ にどれだけ注目しているかを表す値 $\alpha_{ij}$ が、自動的に計算されます。
数学的な定義
前提記事で学んだScaled Dot-Product Attentionの式を改めて確認しましょう。入力系列からQuery $\bm{q}_i$、Key $\bm{k}_j$、Value $\bm{v}_j$ が計算されたとき、トークン $i$ からトークン $j$ へのattention重みは次のように定義されます。
まず、QueryとKeyの内積にスケーリングを施してスコアを計算します。
$$ e_{ij} = \frac{\bm{q}_i^\top \bm{k}_j}{\sqrt{d_k}} $$
ここで $d_k$ はKeyベクトルの次元数です。$\sqrt{d_k}$ で割るのは、内積の値が次元数に比例して大きくなるのを防ぎ、softmaxの勾配消失を避けるためです。
このスコアにsoftmaxを適用して、確率分布に変換します。
$$ \alpha_{ij} = \text{softmax}_j(e_{ij}) = \frac{\exp(e_{ij})}{\sum_{k=1}^{n} \exp(e_{ik})} $$
softmaxによって $\sum_j \alpha_{ij} = 1$ かつ $\alpha_{ij} \geq 0$ が保証されます。つまり、各トークン $i$ について、全トークンへの注目度が確率分布を成します。
重み行列の意味
系列長が $n$ のとき、全トークン対のattention重みをまとめると、行列 $\bm{A} \in \mathbb{R}^{n \times n}$ が得られます。
$$ \bm{A} = \text{softmax}\left(\frac{\bm{Q}\bm{K}^\top}{\sqrt{d_k}}\right) $$
この行列の $(i, j)$ 成分 $A_{ij}$ は「トークン $i$ がトークン $j$ にどれだけ注目しているか」を表します。行方向($i$ を固定)に見ると、トークン $i$ の注目先の分布が得られます。各行の和は1になるので、ヒートマップで可視化すれば、各トークンがどこに注意を分散させているかが一目でわかります。
Multi-Head Attentionにおける多様なパターン
Multi-Head Attentionでは、$h$ 個のヘッドがそれぞれ独立にattention重みを計算します。ヘッド $l$ のattention行列を $\bm{A}^{(l)} \in \mathbb{R}^{n \times n}$ とすると、モデル全体では $h$ 枚のattention行列を持ちます。
各ヘッドは異なる射影行列 $\bm{W}_Q^{(l)}, \bm{W}_K^{(l)}, \bm{W}_V^{(l)}$ を学習するため、それぞれが異なるパターン — たとえば構文的な関係(主語→動詞)、位置的な関係(隣接トークン)、意味的な関係(同義語・反義語)— に特化することができます。BERT-baseモデルの場合、12層 $\times$ 12ヘッド = 144枚のattention行列が存在し、それぞれが異なる視点から入力を分析しています。
ここまでattention重みの数学的な定義と構造を復習しました。では、これらの重み行列をどのように可視化すれば、人間にとって解釈しやすい形になるのでしょうか。次のセクションでは、代表的な可視化手法を紹介します。
Attention重みの可視化手法
Attention重みを可視化する方法はいくつか存在しますが、それぞれ得意な表現が異なります。目的に応じて使い分けることが重要です。
ヒートマップ(最も基本的)
Attention行列 $\bm{A} \in \mathbb{R}^{n \times n}$ は本質的に2次元配列なので、ヒートマップとして可視化するのが最も自然です。横軸にKeyトークン(注目先)、縦軸にQueryトークン(注目元)を取り、セルの色の濃さで重みの大きさを表します。
ヒートマップの利点は、行列全体の構造を俯瞰できることです。対角線上にattentionが集中していれば「各トークンは主に自分自身に注目している」とわかりますし、特定の列に値が集中していれば「あるトークンが全体から強く注目されている」と読み取れます。
ただし、ヒートマップは系列長が長くなるとセルが小さくなり、詳細が読みにくくなるという欠点があります。系列長が50以上の場合は、関心のあるトークンに絞った可視化が実用的です。
ヘッド別の可視化
Multi-Head Attentionでは複数のヘッドがそれぞれ異なるパターンを学習しています。全ヘッドをグリッド状に並べて比較することで、ヘッド間の役割の違いが明確になります。
たとえば、あるヘッドは対角成分が強く(自己参照パターン)、別のヘッドは特定のトークン([CLS]や[SEP])に集中し、さらに別のヘッドは隣接トークンへの注目が強い、といった違いが観察されます。これはモデルが多角的に情報を処理していることの証拠です。
層別の比較
Transformerの深い層と浅い層では、attentionのパターンが大きく異なることが知られています。同じヘッド番号を複数の層にわたって比較することで、層の深さに応じたパターンの変化を追跡できます。
一般的な傾向として、浅い層(第1〜3層あたり)では局所的なパターン(隣接トークンへの注目)が支配的であり、深い層になるにつれてより広域の意味的パターン(離れたトークン間の関係)が現れます。この層ごとの変化を可視化することは、モデルが情報をどのように段階的に処理しているかを理解する手がかりになります。
BERTVizの紹介
BERTViz は、Transformerのattention重みを対話的に可視化するためのオープンソースツールです。Jesse Vig氏によって開発され、Hugging Faceのモデルと直接連携できます。BERTVizには3つの主要なビューがあります。
Attention-head view(ヘッドビュー): 入力トークンを左右に並べ、attention重みを線で結んで表示します。線の太さがattention重みの大きさに対応し、特定のトークンがどこに注目しているかが直感的にわかります。層やヘッドの切り替えもインタラクティブに行えます。
Model view(モデルビュー): 全層・全ヘッドのattentionパターンをサムネイルで一覧表示します。モデル全体のattention構造を俯瞰でき、「この層のこのヘッドは特徴的なパターンを持っている」という発見に役立ちます。
Neuron view(ニューロンビュー): Query・Key・Valueの個々のニューロンレベルの値を可視化します。attention重みがどのようなQuery-Key相互作用から生じているかを詳細に分析できます。これは最も細粒度の可視化であり、ヘッドの内部メカニズムを理解するのに有用です。
BERTVizはJupyter Notebook上で動作し、数行のコードで上記のビューを生成できるため、研究やデバッグの場面で広く使われています。本記事では、BERTVizに頼らずmatplotlibで基本的な可視化を実装する方法に焦点を当てますが、対話的な探索が必要な場合にはBERTVizの利用を推奨します。
ここまで可視化手法の全体像を把握しました。では実際に、PyTorchとHugging Faceを使ってBERTモデルからattention重みを抽出する方法を見ていきましょう。
PyTorchでAttention重みを抽出する方法
output_attentions=Trueの使い方
Hugging FaceのTransformersライブラリでは、モデルの推論時に output_attentions=True を指定するだけでattention重みを取得できます。この引数をモデルのforwardメソッドに渡すと、通常の出力に加えて各層のattention行列がタプルとして返されます。
具体的には、返されるattentionの構造は以下のとおりです。
- 最外タプルの長さ: 層の数(BERT-baseなら12)
- 各要素のテンソル形状:
(batch_size, num_heads, seq_len, seq_len)
つまり、BERT-baseモデルで系列長16のバッチサイズ1の入力を処理した場合、各層のattentionテンソルの形状は (1, 12, 16, 16) になります。attentions[l][0, h, i, j] が「第 $l$ 層・第 $h$ ヘッドにおける、トークン $i$ からトークン $j$ へのattention重み」に対応します。
Hugging Face BERTからの抽出
実際にBERTモデルからattention重みを抽出するコードを見てみましょう。ここでは bert-base-uncased モデルを使います。
import torch
from transformers import BertTokenizer, BertModel
# トークナイザとモデルの準備
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
model = BertModel.from_pretrained("bert-base-uncased")
model.eval()
# 入力テキストのトークナイズ
text = "The cat sat on the mat because it was tired."
inputs = tokenizer(text, return_tensors="pt")
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
# attention重みを含む推論
with torch.no_grad():
outputs = model(**inputs, output_attentions=True)
# attention重みの取得
attentions = outputs.attentions # タプル(層数 = 12)
# 形状の確認
print(f"層の数: {len(attentions)}")
print(f"各層のテンソル形状: {attentions[0].shape}")
print(f"トークン一覧: {tokens}")
このコードを実行すると、以下のような出力が得られます。
層の数: 12
各層のテンソル形状: torch.Size([1, 12, 12, 12])
トークン一覧: ['[CLS]', 'the', 'cat', 'sat', 'on', 'the', 'mat', 'because', 'it', 'was', 'tired', '.', '[SEP]']
ここでテンソル形状の意味を確認しておきましょう。(1, 12, 12, 12) は、バッチサイズ1、12ヘッド、系列長12、系列長12です。ただし実際の系列長はトークナイザの出力に依存するため(特殊トークン [CLS] と [SEP] を含む)、テキストによって変わります。
抽出した重みのスタック
後の可視化で扱いやすいように、全層のattention重みを1つの4次元テンソルにスタックしておくと便利です。
import torch
# 全層のattentionを (num_layers, num_heads, seq_len, seq_len) にスタック
attention_tensor = torch.stack(attentions).squeeze(1)
print(f"スタック後の形状: {attention_tensor.shape}")
# → torch.Size([12, 12, seq_len, seq_len])
squeeze(1) でバッチ次元(サイズ1)を除去しています。結果として (12, 12, seq_len, seq_len) の形状、すなわち(層数, ヘッド数, 系列長, 系列長)のテンソルが得られます。この形状であれば、attention_tensor[layer, head] で特定の層・ヘッドのattention行列($n \times n$)に直接アクセスできます。
ここまでで、BERTからattention重みを抽出し、扱いやすい形に整形するところまで完了しました。次は、いよいよこれらの重みをmatplotlibで可視化していきます。
可視化の実装(matplotlib)
単一ヘッドのヒートマップ
まずは最も基本的な可視化として、特定の層・特定のヘッドのattention行列を1枚のヒートマップとして描画します。これにより、そのヘッドが入力系列のどのペアに注目しているかを一望できます。
import torch
import numpy as np
import matplotlib.pyplot as plt
from transformers import BertTokenizer, BertModel
# モデルとトークナイザの準備
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
model = BertModel.from_pretrained("bert-base-uncased")
model.eval()
# 入力テキスト
text = "The cat sat on the mat because it was tired."
inputs = tokenizer(text, return_tensors="pt")
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
# attention重みの取得
with torch.no_grad():
outputs = model(**inputs, output_attentions=True)
attentions = outputs.attentions
# 特定の層・ヘッドを選択(第6層・第8ヘッド)
layer_idx = 5 # 0-indexed
head_idx = 7 # 0-indexed
attn_matrix = attentions[layer_idx][0, head_idx].numpy()
# ヒートマップの描画
fig, ax = plt.subplots(figsize=(10, 8))
im = ax.imshow(attn_matrix, cmap="Blues", vmin=0, vmax=1)
ax.set_xticks(range(len(tokens)))
ax.set_yticks(range(len(tokens)))
ax.set_xticklabels(tokens, rotation=45, ha="right", fontsize=9)
ax.set_yticklabels(tokens, fontsize=9)
ax.set_xlabel("Key (attended to)")
ax.set_ylabel("Query (attending from)")
ax.set_title(f"Attention Heatmap — Layer {layer_idx+1}, Head {head_idx+1}")
# 各セルに数値を表示
for i in range(len(tokens)):
for j in range(len(tokens)):
val = attn_matrix[i, j]
color = "white" if val > 0.5 else "black"
ax.text(j, i, f"{val:.2f}", ha="center", va="center",
fontsize=7, color=color)
plt.colorbar(im, ax=ax, label="Attention weight")
plt.tight_layout()
plt.savefig("attention_heatmap_single.png", dpi=150, bbox_inches="tight")
plt.show()
このヒートマップから、いくつかの特徴を読み取ることができます。まず、対角線上のセルが比較的高い値を持っている場合、そのヘッドは「各トークンが自分自身に注目する」パターンを学習していることがわかります。また、特定の列(たとえば [CLS] や [SEP] の列)に値が集中している場合、そのトークンが「情報の集約点」として機能していることを示唆します。さらに、「it」の行で「cat」の列の値が高ければ、このヘッドが代名詞の照応関係を捉えている可能性があります。
全ヘッドのグリッド表示
1つのヘッドだけでは全体像がつかめません。次に、ある層の全12ヘッドを一度にグリッド表示して比較してみましょう。
import torch
import numpy as np
import matplotlib.pyplot as plt
from transformers import BertTokenizer, BertModel
# モデルとトークナイザの準備
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
model = BertModel.from_pretrained("bert-base-uncased")
model.eval()
# 入力テキスト
text = "The cat sat on the mat because it was tired."
inputs = tokenizer(text, return_tensors="pt")
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
# attention重みの取得
with torch.no_grad():
outputs = model(**inputs, output_attentions=True)
attentions = outputs.attentions
# 指定した層の全ヘッドをグリッド表示
layer_idx = 5 # 第6層
fig, axes = plt.subplots(3, 4, figsize=(20, 15))
for head_idx in range(12):
row, col = head_idx // 4, head_idx % 4
ax = axes[row, col]
attn_matrix = attentions[layer_idx][0, head_idx].numpy()
im = ax.imshow(attn_matrix, cmap="Blues", vmin=0, vmax=1)
ax.set_title(f"Head {head_idx + 1}", fontsize=11)
ax.set_xticks(range(len(tokens)))
ax.set_yticks(range(len(tokens)))
ax.set_xticklabels(tokens, rotation=45, ha="right", fontsize=6)
ax.set_yticklabels(tokens, fontsize=6)
fig.suptitle(f"All Attention Heads — Layer {layer_idx + 1}", fontsize=14, y=1.02)
plt.tight_layout()
plt.savefig("attention_heatmap_all_heads.png", dpi=150, bbox_inches="tight")
plt.show()
12枚のヒートマップを並べると、ヘッド間の多様性が鮮やかに浮かび上がります。あるヘッドは強い対角パターン(自己参照)を示し、別のヘッドは特定の列に値が集中([SEP]トークンへの集中)し、また別のヘッドは対角線の1つ上や1つ下にattentionが集中(隣接トークンへの注目)するなど、それぞれが異なる「視点」で入力を分析していることが見て取れます。この多様性こそが、Multi-Head Attentionが単一ヘッドよりも豊かな表現力を持つ理由です。
層ごとのAttentionパターンの変化
次に、ヘッドを固定して層の深さによるパターンの変化を追跡します。これにより、浅い層から深い層にかけて情報処理がどう変化するかを観察できます。
import torch
import numpy as np
import matplotlib.pyplot as plt
from transformers import BertTokenizer, BertModel
# モデルとトークナイザの準備
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
model = BertModel.from_pretrained("bert-base-uncased")
model.eval()
# 入力テキスト
text = "The cat sat on the mat because it was tired."
inputs = tokenizer(text, return_tensors="pt")
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
# attention重みの取得
with torch.no_grad():
outputs = model(**inputs, output_attentions=True)
attentions = outputs.attentions
# ヘッドを固定し、層ごとの変化を表示
head_idx = 0 # 第1ヘッド
selected_layers = [0, 3, 7, 11] # 第1, 4, 8, 12層
fig, axes = plt.subplots(1, 4, figsize=(24, 6))
for idx, layer_idx in enumerate(selected_layers):
ax = axes[idx]
attn_matrix = attentions[layer_idx][0, head_idx].numpy()
im = ax.imshow(attn_matrix, cmap="Blues", vmin=0, vmax=1)
ax.set_title(f"Layer {layer_idx + 1}", fontsize=12)
ax.set_xticks(range(len(tokens)))
ax.set_yticks(range(len(tokens)))
ax.set_xticklabels(tokens, rotation=45, ha="right", fontsize=7)
ax.set_yticklabels(tokens, fontsize=7)
fig.suptitle(f"Attention Pattern Across Layers — Head {head_idx + 1}",
fontsize=14, y=1.02)
plt.tight_layout()
plt.savefig("attention_layer_comparison.png", dpi=150, bbox_inches="tight")
plt.show()
この可視化から、層ごとのattentionパターンの進化を確認できます。一般的に、浅い層(Layer 1付近)では対角線や隣接トークンへの注目が支配的で、局所的な情報処理が行われています。中間層(Layer 4〜8)では、より広い範囲のトークン間関係が現れ始めます。深い層(Layer 12付近)では、特定のトークンへの集中やタスクに関連した意味的パターンが顕著になる傾向があります。ただし、この傾向はヘッドによって大きく異なるため、複数のヘッドを比較することが重要です。
特定トークンに対するAttention分布の棒グラフ
ヒートマップは行列全体を俯瞰するのに適していますが、特定のトークンに注目したい場合は棒グラフの方が読みやすくなります。あるQueryトークンが各Keyトークンにどれだけ注目しているかを棒グラフで表示してみましょう。
import torch
import numpy as np
import matplotlib.pyplot as plt
from transformers import BertTokenizer, BertModel
# モデルとトークナイザの準備
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
model = BertModel.from_pretrained("bert-base-uncased")
model.eval()
# 入力テキスト
text = "The cat sat on the mat because it was tired."
inputs = tokenizer(text, return_tensors="pt")
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
# attention重みの取得
with torch.no_grad():
outputs = model(**inputs, output_attentions=True)
attentions = outputs.attentions
# 「it」トークンのattention分布を棒グラフで表示
query_token = "it"
query_idx = tokens.index(query_token)
layer_idx = 5 # 第6層
head_idx = 7 # 第8ヘッド
attn_weights = attentions[layer_idx][0, head_idx, query_idx].numpy()
fig, ax = plt.subplots(figsize=(12, 5))
colors = ["#2196F3" if t != query_token else "#FF5722" for t in tokens]
bars = ax.bar(range(len(tokens)), attn_weights, color=colors, edgecolor="gray")
ax.set_xticks(range(len(tokens)))
ax.set_xticklabels(tokens, rotation=45, ha="right", fontsize=10)
ax.set_ylabel("Attention weight", fontsize=11)
ax.set_title(
f'Attention distribution from "{query_token}" — '
f"Layer {layer_idx+1}, Head {head_idx+1}",
fontsize=12,
)
ax.set_ylim(0, 1)
ax.axhline(y=1/len(tokens), color="red", linestyle="--", alpha=0.5,
label="Uniform distribution")
ax.legend()
plt.tight_layout()
plt.savefig("attention_bar_chart.png", dpi=150, bbox_inches="tight")
plt.show()
棒グラフの結果から、代名詞「it」がどのトークンに注目しているかを直接読み取ることができます。赤い破線は一様分布(全トークンに均等に注目する場合)の値を示しており、これより高い棒は「一様分布よりも強く注目されている」ことを意味します。「cat」への注目が高ければ、このヘッドが照応解析(coreference resolution)的な役割を果たしていることが示唆されます。一方、[SEP] への注目が高い場合は、そのヘッドは特定の意味的パターンよりも「集約トークン」への情報の流れを担っている可能性があります。
ここまでmatplotlibによる基本的な可視化を実装しました。次に、これらの可視化で実際に観察されることの多いattentionパターンの類型を整理します。
よく観察されるAttentionパターン
BERTなどのTransformerモデルのattention重みを可視化すると、いくつかの特徴的なパターンが繰り返し観察されます。Clark et al. (2019) の大規模な分析をはじめ、多くの研究がこれらのパターンを報告しています。ここでは代表的な6つのパターンを紹介します。
対角パターン(自分自身に注目)
ヒートマップの対角成分が強くなるパターンです。各トークンが主に自分自身に注目しており、文脈情報よりも自身のトークン表現を保持する役割を果たしていると考えられます。
このパターンは特に浅い層で頻繁に見られます。直感的には、最初の数層ではまだ文脈的な情報処理が本格化しておらず、各トークンが自身の入力表現をそのまま次の層に伝えている段階だと解釈できます。residual connection(残差接続)の存在によって、自己注目の重みが高くても情報の劣化は起きません。
[CLS]トークンと[SEP]トークンへの集中
BERTの多くのヘッドで、attention重みが [CLS] や [SEP] といった特殊トークンに集中するパターンが観察されます。特に [SEP] への集中は非常に一般的で、Clark et al. (2019) はBERT-baseの多くのヘッドでこの傾向を報告しています。
[CLS] トークンは文全体の要約表現を担うため、文全体から情報を集約するattentionパターンを持つのは合理的です。一方、[SEP] への集中については議論があります。[SEP] は文の境界を示す特殊トークンであり、意味的に重要な情報を持つわけではありません。この現象は、attentionの「ゴミ箱」仮説で説明されることがあります — 特定のトークンに注目する必要がないとき、モデルは情報量の少ない [SEP] に重みを「捨てる」ことで、他のトークンへの注目を抑制しているという解釈です。
位置ベースのパターン(前後のトークン)
対角線の1つ上や1つ下に沿ったattentionのパターンです。つまり、各トークンが直前のトークンまたは直後のトークンに注目しています。
このパターンは、言語のn-gramモデルに通じるものがあります。自然言語では隣接する単語間の関係(コロケーション、複合語、活用形など)が非常に重要であり、Transformerがこのような局所的パターンを学習するのは理にかなっています。位置ベースのパターンは浅い層で特に顕著で、局所的な構文情報の処理に関与していると考えられます。
構文的パターン(主語→動詞)
「主語のトークンが動詞に強いattentionを示す」「修飾語が被修飾語に注目する」といった、言語の構文構造を反映したパターンです。
たとえば「The large cat sat quietly」という文で、「cat」のattentionが「sat」に集中し、「large」のattentionが「cat」に集中するようなパターンが、特定のヘッドで観察されることがあります。Htut et al. (2019) やClark et al. (2019) は、BERTの特定のヘッドが依存構文木(dependency tree)のエッジに対応するattentionパターンを持つことを示しました。ただし、全ての構文関係がattentionに直接的に反映されるわけではなく、一部の関係(特に直接目的語や前置詞の補語)がより顕著に現れる傾向があります。
各層での役割の違い
Transformerの各層は、情報処理において異なる役割を果たしていることが示されています。
浅い層(第1〜3層): 主に局所的なパターンが支配的です。隣接トークンへの注目、自己参照(対角パターン)、位置ベースのパターンが多く見られます。これは、語のレベルでの構文的な情報(品詞、活用、コロケーション)を処理している段階と考えられます。
中間層(第4〜8層): 構文的なパターンが最も顕著に現れる層です。主語-動詞の関係、修飾関係、照応関係などが特定のヘッドに現れます。BERTの構文的な知識は、主にこの中間層に集中していることが複数の研究で示されています(Jawahar et al., 2019)。
深い層(第9〜12層): タスク固有の情報処理が行われます。ファインチューニング後のモデルでは、タスクに関連するトークンへのattentionが強くなります。また、[CLS] トークンへの注目が増加し、文全体の情報を集約する動きが強まります。未ファインチューニングのBERTでは、深い層でattentionパターンがやや拡散する傾向も見られます。
エントロピーによるパターンの定量化
上記のパターンを定量的に評価するために、attention分布のエントロピーを計算するのが有効です。エントロピーが低いほどattentionが特定のトークンに集中しており、エントロピーが高いほど一様分布に近い(=特定のパターンが弱い)ことを意味します。
import torch
import numpy as np
import matplotlib.pyplot as plt
from transformers import BertTokenizer, BertModel
# モデルとトークナイザの準備
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
model = BertModel.from_pretrained("bert-base-uncased")
model.eval()
# 入力テキスト
text = "The cat sat on the mat because it was tired."
inputs = tokenizer(text, return_tensors="pt")
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
# attention重みの取得
with torch.no_grad():
outputs = model(**inputs, output_attentions=True)
attentions = outputs.attentions
# 各層・各ヘッドのattentionエントロピーを計算
num_layers = len(attentions)
num_heads = attentions[0].shape[1]
entropy_matrix = np.zeros((num_layers, num_heads))
for layer_idx in range(num_layers):
for head_idx in range(num_heads):
# attention行列 (seq_len, seq_len)
attn = attentions[layer_idx][0, head_idx].numpy()
# 各行(各Queryトークン)のエントロピーを計算し平均
row_entropy = -np.sum(
attn * np.log(attn + 1e-12), axis=1
)
entropy_matrix[layer_idx, head_idx] = np.mean(row_entropy)
# 最大エントロピー(一様分布の場合)
max_entropy = np.log(len(tokens))
# ヒートマップで表示
fig, ax = plt.subplots(figsize=(12, 6))
im = ax.imshow(entropy_matrix, cmap="YlOrRd_r", vmin=0, vmax=max_entropy,
aspect="auto")
ax.set_xlabel("Head", fontsize=11)
ax.set_ylabel("Layer", fontsize=11)
ax.set_xticks(range(num_heads))
ax.set_yticks(range(num_layers))
ax.set_xticklabels([f"H{i+1}" for i in range(num_heads)])
ax.set_yticklabels([f"L{i+1}" for i in range(num_layers)])
ax.set_title("Attention Entropy per Layer and Head\n(lower = more focused)",
fontsize=13)
# セルに数値を表示
for i in range(num_layers):
for j in range(num_heads):
val = entropy_matrix[i, j]
color = "white" if val < max_entropy * 0.4 else "black"
ax.text(j, i, f"{val:.2f}", ha="center", va="center",
fontsize=6, color=color)
plt.colorbar(im, ax=ax, label="Mean entropy (nats)")
plt.tight_layout()
plt.savefig("attention_entropy.png", dpi=150, bbox_inches="tight")
plt.show()
エントロピーのヒートマップから、モデル全体のattentionの集中度合いを一望できます。エントロピーが低い(色が濃い)セルは、そのヘッドが特定のトークンに強く集中していることを意味します。一般的に、浅い層のヘッドの一部は低エントロピー(対角パターンや位置パターン)、[SEP] に集中するヘッドも低エントロピーになります。逆にエントロピーが高い(一様に近い)ヘッドは、明確なパターンを持たず、情報をまんべんなく収集する役割を担っている可能性があります。このエントロピー分析は、ヘッドの枝刈り(pruning)の際に「重要度の低いヘッド」を特定する指標としても活用されます。
可視化で観察されるパターンを整理できたところで、ここからはより慎重な議論に入ります。attention重みを「モデルの説明」として使うことにはどのような限界があるのでしょうか。
Attentionの解釈性の限界
直感的な問題提起
attention重みの可視化は視覚的に魅力的であり、「モデルが何に注目しているか」を示しているように見えます。しかし、ここで根本的な問いが生じます — attentionの重みが高いことは、そのトークンが出力に重要だということを本当に意味するのでしょうか?
たとえば、ある感情分析モデルが「This movie is not great」を「ネガティブ」と正しく判定し、attention可視化で「not」に高い重みが見えたとします。「not」に注目したから否定を捉えた、と解釈したくなりますが、本当にそうでしょうか。もし「not」へのattention重みをゼロに強制的に置き換えても、出力が変わらなかったとしたら、その解釈は誤りだったことになります。
“Attention is not Explanation”(Jain & Wallace, 2019)
Jain & Wallace (2019) は、attention重みを「説明」として扱うことに重要な疑問を投げかけた論文です。彼らの主張の核は以下の2点です。
主張1: Attentionとgradient-based attributionの不一致
入力の各トークンが出力にどれだけ寄与したかを測る手法として、勾配ベースの帰属法(gradient-based attribution)があります。これは出力を入力の各要素で偏微分したもので、出力に対する「真の」影響度の指標として使われます。Jain & Wallaceは、attention重みとgradient-based attributionの間に一貫した相関がないことを実験的に示しました。つまり、attentionが高いトークンが必ずしも出力に大きな影響を与えているわけではありません。
主張2: 代替attention分布の存在
attention重みを元の分布とは大きく異なる分布に変更しても、モデルの出力がほとんど変化しない場合があることを示しました。具体的には、元のattention分布と相関の低い「代替分布」を探索し、それを使っても同程度の出力が得られるケースが多数存在することを実験で確認しました。
これは深刻な問題です。もし同じ出力を生むattention分布が複数存在するなら、観測されたattention分布はモデルの判断根拠の「一つの可能な表現」に過ぎず、唯一の説明ではないことになります。
数式で見る問題の本質
この問題を数式で表現してみましょう。attention重みを $\bm{\alpha}$ 、Value行列を $\bm{V}$ とすると、attentionの出力は次のようになります。
$$ \bm{o} = \sum_{j} \alpha_j \bm{v}_j $$
ここで、$\alpha_j$ が大きくても $\bm{v}_j$ の方向や大きさによっては出力 $\bm{o}$ への寄与が小さい場合があります。逆に $\alpha_j$ が小さくても、$\bm{v}_j$ が大きければ最終出力への影響は無視できません。
さらに、Transformerでは残差接続(residual connection)により、attention層の出力は入力にそのまま加算されます。
$$ \bm{h}^{(l)} = \bm{h}^{(l-1)} + \text{Attention}(\bm{h}^{(l-1)}) $$
つまり、attention層がゼロを出力しても(attentionが何も寄与しなくても)、情報は残差接続を通じて次の層に伝わります。この事実は、attention重みだけを見てもモデル全体の情報フローを理解するには不十分であることを示しています。
代替手法: Gradient-based Attribution
attention重みの限界を補う手法として、勾配ベースの帰属法があります。
Vanilla Gradient: 出力 $y$ の入力 $\bm{x}_j$ に関する勾配 $\frac{\partial y}{\partial \bm{x}_j}$ を計算します。勾配の大きさが、そのトークンの出力への影響度を直接表します。
$$ \text{Attribution}(j) = \left\| \frac{\partial y}{\partial \bm{x}_j} \right\| $$
Integrated Gradients(Sundararajan et al., 2017): Vanilla Gradientの問題(飽和領域での勾配消失など)を解決するために、ベースライン $\bm{x}’$ から入力 $\bm{x}$ までの経路に沿って勾配を積分します。
$$ \text{IG}(j) = (\bm{x}_j – \bm{x}’_j) \times \int_0^1 \frac{\partial F(\bm{x}’ + t(\bm{x} – \bm{x}’))}{\partial \bm{x}_j} dt $$
$t = 0$ のときベースライン、$t = 1$ のとき実際の入力に対応し、経路全体の勾配の積分を取ることで公理的に望ましい性質(完全性、感度など)を満たす帰属法が得られます。Integrated Gradientsは計算コストが高い(複数のforward passが必要)ですが、より信頼性の高い帰属スコアを提供します。
“Attention is not not Explanation”(Wiegreffe & Pinter, 2019)
Jain & Wallaceの議論に対して、Wiegreffe & Pinter (2019) は反論を提示しました。彼らの主な主張は以下のとおりです。
反論1: 代替分布の存在は問題にならない
Jain & Wallaceは「同じ出力を生むattention分布が複数存在する」ことを問題視しましたが、Wiegreffe & Pinterは「複数の説明が存在すること自体は、元の説明が無意味であることを示さない」と反論しました。人間の意思決定でも、同じ結論に至る複数の推論パスが存在し得ますが、だからといって1つの推論パスが「説明にならない」とは言えません。
反論2: Attentionは弱い説明としては有効
彼らは「attention重みが完璧な説明ではない」としても、「全く説明力がない」わけではないことを実験的に示しました。具体的には、attention分布を固定したモデルの性能が、ランダムなattention分布を使った場合よりも有意に高いことを確認しました。これは、学習されたattention分布にはタスクに関連する情報が含まれていることを意味します。
反論3: adversarially trained attentionの限界
Jain & Wallaceの手法で見つかる「代替attention分布」は、元のモデルの他のパラメータ(Value行列、後続のFFN層など)を固定した状態で探索したものです。しかしモデルの全パラメータは相互作用的に学習されているため、attention分布だけを変更する操作は、モデルの本来の動作を反映していない可能性があります。
結論: 注意重みは参考にはなるが唯一の説明ではない
以上の議論を総合すると、attention重みの解釈性に関しては以下のような中庸な立場が妥当です。
- Attention重みは一定の説明力を持つ — 完全にランダムなわけではなく、言語構造やタスク関連情報を反映するパターンが確かに存在します
- ただし唯一の説明ではない — 同じ出力を生む代替分布が存在する以上、observedなattention分布だけを「モデルの判断根拠」とするのは危険です
- 他の手法と併用すべき — gradient-based attribution、Integrated Gradients、probing classifierなど、複数の解釈手法を併用して結論を導くのが望ましいです
- 層やヘッドによって解釈可能性は異なる — 全てのヘッドが解釈可能なパターンを持つわけではなく、特定のヘッドだけが人間にとって理解しやすいパターンを示します
attention可視化は「モデルの内部を覗く窓」としては有用ですが、「窓から見える景色が全て」ではないことを念頭に置く必要があります。
ここまでattention重みの解釈性の限界を理解しました。では、単一の層のattention行列だけを見るのではなく、複数の層を通じたattentionの伝播を追跡するにはどうすればよいでしょうか。次のセクションでは、Attention Rollout という手法を紹介します。
Attention Rollout
問題意識: 単一層のattentionでは不十分
これまでの可視化は、特定の1つの層のattention行列を観察するものでした。しかしTransformerは多層構造であり、入力トークンの情報は層を経るごとに混合されていきます。第1層で「cat」に注目した結果が第2層の表現に反映され、その表現がさらに第3層で別のトークンへの注目に影響する — というように、情報は層を跨いで伝播していきます。
つまり、最終層のattention行列だけを見ても、入力トークンから最終出力までの情報フローの全体像は把握できません。これを解決するために、Abnar & Zuidema (2020) が提案したAttention Rolloutという手法があります。
Attention Rolloutの計算方法
Attention Rolloutのアイデアは、各層のattention行列を順番に掛け合わせることで、入力から出力までの「累積的な注目度」を計算するというものです。
ただし、Transformerには各層に残差接続があります。残差接続は、attention層の入力をそのまま出力に加算するため、各トークンは自分自身の情報を必ず保持します。この効果をattention行列に組み込むために、単位行列を混合します。
ステップ1: 残差接続を考慮したattention行列の構築
各層 $l$ のattention行列 $\bm{A}_l$ に対して、残差接続の効果を反映した修正行列 $\hat{\bm{A}}_l$ を計算します。
$$ \hat{\bm{A}}_l = 0.5 \cdot \bm{A}_l + 0.5 \cdot \bm{I} $$
ここで $\bm{I}$ は単位行列です。0.5という係数は、残差接続が attention の出力と入力を均等に混合するという近似です。単位行列を加えることで「各トークンが自分自身にも50%の注目を残す」ことを表現しています。
ステップ2: 行の正規化
$\hat{\bm{A}}_l$ の各行が確率分布を維持するように正規化します。ステップ1で単位行列を加えた結果、各行の和が1より大きくなっているため、再正規化が必要です。
$$ \hat{\bm{A}}_l \leftarrow \frac{\hat{\bm{A}}_l}{\sum_j \hat{A}_{l,ij}} $$
ただし、$0.5 \bm{A}_l + 0.5 \bm{I}$ の場合は $\bm{A}_l$ の各行の和が1であり、$\bm{I}$ の各行の和も1なので、混合後の各行の和は自動的に1になります。したがって、この場合は明示的な再正規化は不要です。
ステップ3: 層にわたる行列の逐次積
全層にわたって修正attention行列の積を計算します。
$$ \bm{R} = \hat{\bm{A}}_L \cdot \hat{\bm{A}}_{L-1} \cdots \hat{\bm{A}}_2 \cdot \hat{\bm{A}}_1 = \prod_{l=1}^{L} \hat{\bm{A}}_l $$
$\bm{R} \in \mathbb{R}^{n \times n}$ のRollout行列で、$R_{ij}$ は「入力トークン $j$ から最終表現のトークン $i$ に至るまでに伝播した注目度の総量」を表します。
Multi-Head Attentionの扱い
Multi-Head Attentionでは、各層に複数のヘッドがあります。Attention Rolloutでは、各層のattention行列としてヘッドの平均を使います。
$$ \bm{A}_l = \frac{1}{h} \sum_{k=1}^{h} \bm{A}_l^{(k)} $$
これは単純化ですが、実用的にはヘッドの平均が全体的なattention傾向を捉えるのに十分であることが多いです。
実装
Attention Rolloutをpythonで実装し、結果を可視化してみましょう。
import torch
import numpy as np
import matplotlib.pyplot as plt
from transformers import BertTokenizer, BertModel
# モデルとトークナイザの準備
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
model = BertModel.from_pretrained("bert-base-uncased")
model.eval()
# 入力テキスト
text = "The cat sat on the mat because it was tired."
inputs = tokenizer(text, return_tensors="pt")
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
# attention重みの取得
with torch.no_grad():
outputs = model(**inputs, output_attentions=True)
attentions = outputs.attentions
# Attention Rolloutの計算
def compute_attention_rollout(attentions):
"""
全層のattention重みからAttention Rolloutを計算する。
Parameters
----------
attentions : tuple of torch.Tensor
各層のattention (batch, heads, seq_len, seq_len)
Returns
-------
rollout : np.ndarray (seq_len, seq_len)
Rollout行列
"""
num_layers = len(attentions)
seq_len = attentions[0].shape[-1]
# 単位行列
identity = np.eye(seq_len)
# 累積行列(初期値は単位行列)
rollout = identity.copy()
for layer_idx in range(num_layers):
# 全ヘッドの平均をとる
attn = attentions[layer_idx][0].numpy() # (heads, seq_len, seq_len)
attn_mean = attn.mean(axis=0) # (seq_len, seq_len)
# 残差接続を考慮して単位行列を加える
attn_hat = 0.5 * attn_mean + 0.5 * identity
# 逐次的に行列積を計算
rollout = attn_hat @ rollout
return rollout
rollout = compute_attention_rollout(attentions)
# Rollout行列のヒートマップ
fig, axes = plt.subplots(1, 2, figsize=(20, 8))
# 左: Rollout行列全体
ax = axes[0]
im = ax.imshow(rollout, cmap="Blues")
ax.set_xticks(range(len(tokens)))
ax.set_yticks(range(len(tokens)))
ax.set_xticklabels(tokens, rotation=45, ha="right", fontsize=9)
ax.set_yticklabels(tokens, fontsize=9)
ax.set_xlabel("Input token (source)")
ax.set_ylabel("Output token (target)")
ax.set_title("Attention Rollout Matrix", fontsize=12)
plt.colorbar(im, ax=ax, label="Rollout score")
# 右: [CLS]トークンの入力トークンへの注目度
ax = axes[1]
cls_rollout = rollout[0] # [CLS]は0番目
# [CLS]自身を除外して正規化
cls_rollout_normalized = cls_rollout / cls_rollout.sum()
colors = ["#FF5722" if t in ["[CLS]", "[SEP]"] else "#2196F3"
for t in tokens]
ax.bar(range(len(tokens)), cls_rollout_normalized, color=colors,
edgecolor="gray")
ax.set_xticks(range(len(tokens)))
ax.set_xticklabels(tokens, rotation=45, ha="right", fontsize=10)
ax.set_ylabel("Rollout score (normalized)", fontsize=11)
ax.set_title("[CLS] Token Attention Rollout", fontsize=12)
plt.tight_layout()
plt.savefig("attention_rollout.png", dpi=150, bbox_inches="tight")
plt.show()
Attention Rolloutの結果は、単一層のattention行列とは異なる全体像を提供します。左のヒートマップでは、対角成分が依然として強いですが、これは残差接続により各トークンが自身の情報を保持し続けることの反映です。注目すべきは対角線以外の成分であり、これが層を跨いだ情報の伝播を示しています。
右の棒グラフは、[CLS] トークンが最終的にどの入力トークンから情報を受け取っているかを示しています。[CLS] は文全体の要約表現を担うトークンなので、このRolloutスコアは「モデルが文のどの部分を重視しているか」の近似として使えます。特殊トークン(赤色のバー)への集中が見られる場合は、実際のコンテンツトークン(青色のバー)に着目して解釈するのが有用です。
単一層のattentionとRolloutの比較
Attention Rolloutの効果をより明確にするために、最終層のattentionとRolloutを並べて比較してみましょう。
import torch
import numpy as np
import matplotlib.pyplot as plt
from transformers import BertTokenizer, BertModel
# モデルとトークナイザの準備
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
model = BertModel.from_pretrained("bert-base-uncased")
model.eval()
# 入力テキスト
text = "The cat sat on the mat because it was tired."
inputs = tokenizer(text, return_tensors="pt")
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
# attention重みの取得
with torch.no_grad():
outputs = model(**inputs, output_attentions=True)
attentions = outputs.attentions
def compute_attention_rollout(attentions):
"""Attention Rolloutを計算する"""
num_layers = len(attentions)
seq_len = attentions[0].shape[-1]
identity = np.eye(seq_len)
rollout = identity.copy()
for layer_idx in range(num_layers):
attn = attentions[layer_idx][0].numpy()
attn_mean = attn.mean(axis=0)
attn_hat = 0.5 * attn_mean + 0.5 * identity
rollout = attn_hat @ rollout
return rollout
rollout = compute_attention_rollout(attentions)
# 最終層のattention(ヘッド平均)とRolloutを比較
last_layer_attn = attentions[-1][0].numpy().mean(axis=0) # 最終層のヘッド平均
# [CLS]トークンの注目先を比較
fig, axes = plt.subplots(1, 2, figsize=(18, 5))
# 最終層のattention
ax = axes[0]
cls_last = last_layer_attn[0]
cls_last_norm = cls_last / cls_last.sum()
ax.bar(range(len(tokens)), cls_last_norm, color="#42A5F5", edgecolor="gray")
ax.set_xticks(range(len(tokens)))
ax.set_xticklabels(tokens, rotation=45, ha="right", fontsize=9)
ax.set_ylabel("Attention weight (normalized)")
ax.set_title("[CLS] Attention — Last Layer Only", fontsize=12)
ax.set_ylim(0, max(cls_last_norm) * 1.2)
# Rollout
ax = axes[1]
cls_rollout = rollout[0]
cls_rollout_norm = cls_rollout / cls_rollout.sum()
ax.bar(range(len(tokens)), cls_rollout_norm, color="#66BB6A", edgecolor="gray")
ax.set_xticks(range(len(tokens)))
ax.set_xticklabels(tokens, rotation=45, ha="right", fontsize=9)
ax.set_ylabel("Rollout score (normalized)")
ax.set_title("[CLS] Attention — Rollout (All Layers)", fontsize=12)
ax.set_ylim(0, max(cls_rollout_norm) * 1.2)
plt.tight_layout()
plt.savefig("attention_last_vs_rollout.png", dpi=150, bbox_inches="tight")
plt.show()
左右の棒グラフを比較すると、最終層のattentionだけを見た場合とRolloutで全層を通じた伝播を考慮した場合で、注目度の分布が異なることがわかります。最終層のattentionは特定のトークンに集中しがちですが、Rolloutはより分散した分布を示す傾向があります。これは、浅い層での情報混合が累積的に反映されるためです。Rolloutは単一層のattentionよりも入力全体の寄与をバランスよく捉えるため、特に [CLS] トークンを使った分類タスクにおいて、モデルがどの入力トークンを重視しているかの分析に有用です。
ここまでで、単一層のattention可視化から、複数層にわたる情報伝播の追跡まで一通りのツールを揃えることができました。最後に、本記事のポイントをまとめましょう。
まとめ
本記事では、Transformerのattention重みの可視化と解釈について、理論・実装・議論の3つの側面から解説しました。
- Attention重みの構造: 各層・各ヘッドのattention行列 $\bm{A} \in \mathbb{R}^{n \times n}$ は、トークン間の注目度を表す確率行列です。Multi-Head Attentionにより複数の視点から入力が分析されます
- 可視化手法: ヒートマップ、ヘッド別グリッド表示、層別比較、棒グラフ、エントロピー分析など、目的に応じた可視化手法を実装しました
- 代表的なパターン: 対角パターン、[CLS]/[SEP]集中、位置ベースパターン、構文的パターンなどが観察されます。浅い層は構文的、深い層は意味的なパターンを示す傾向があります
- 解釈性の限界: Jain & Wallace (2019) とWiegreffe & Pinter (2019) の議論に見るように、attention重みは「参考になるが唯一の説明ではない」ことを理解する必要があります。gradient-based attribution等の併用が推奨されます
- Attention Rollout: 残差接続を考慮し、全層にわたるattentionの累積的伝播を計算することで、単一層の分析では得られない全体的な情報フローを把握できます
Attention可視化は、Transformerの「ブラックボックス」を覗くための最もアクセスしやすい手法ですが、それだけでモデルの動作を完全に理解できるわけではありません。より深い理解のためには、probing classifierによる各層の表現分析や、Integrated Gradientsなどのより精緻な帰属法と組み合わせることが重要です。
次のステップとして、以下の記事も参考にしてください。