はじめに
個人的な学習として LLM の基本機構を理解したいと思い、書いています。
本稿では、Multi-Head Attention について、原論文を参考に紐解いてみます。
Self-Attention
まず、Self-Attention は、入力シーケンス内の各トークン同士の関連性を計算し、それぞれのトークンが他のどのトークンに注目すべきかを決定する仕組みです。
これにより、文脈を考慮した特徴表現を生成できます。
自然言語処理では、文章は入力シーケンスとして扱われ、各単語はトークンとして表現されます。
Self-Attention において、各トークンは埋め込み（Embedding）によってベクトルへ変換され、この埋め込みに対して、Query(Q), Key(K), Value(V) の3種類のベクトルを生成します。
その後、各トークンのクエリベクトルと他のすべてのトークンのキーベクトルとの類似度を計算します。この計算には Scaled Dot-Product Attention という計算が用いられ、得られたスコアを Softmax 関数で正規化することで、各トークンへの注意の重み（Attention Weight）が求められます。
Attention(Q,K,V) = softmax(\frac{QK^T}{\sqrt{d_k}})
最後に、その重みを Value ベクトルへ適用することで、入力文の中で関連性の高い情報をより強く反映した新しい特徴表現が得られます。
これにより、例えば代名詞がどの名詞を指しているのかや、離れた単語同士の関係など、文脈を考慮した情報を効果的に学習できます。
Multi-Head Attention
Self-Attention では1種類の関連性しか学習できないため、Transformer では複数の Self-Attention を並列に実行する Multi-Head Attention が採用されており、構文的な関係や意味的な関係など、異なる種類の特徴を同時に学習できるようになっています。
Multi-Head Attention では、1回だけ Attention を計算する代わりに、Query（Q）、Key（K）、Value（V）をそれぞれ異なる学習済みの線形変換によって複数（h個）の低次元ベクトルに変換します。各 Head で並列に Attention を計算し、結果を結合することで、異なる特徴空間や位置関係を同時に捉えられるようになります。1
Multi-Head Attention consists of several attention layers running in parallel 1
Multi-Head Attention の実装例
Multi-Head Attention は、numpy を使って実装することができます。(ChatGPT を使用)
Multi-Head Attention の実装
import numpy as np
def softmax(x, axis=-1):
x = x - np.max(x, axis=axis, keepdims=True)
exp_x = np.exp(x)
return exp_x / np.sum(exp_x, axis=axis, keepdims=True)
class MultiHeadAttention:
def __init__(self, d_model, num_heads):
assert d_model % num_heads == 0
self.d_model = d_model
self.num_heads = num_heads
self.d_head = d_model // num_heads
self.Wq = np.random.randn(d_model, d_model) * 0.01
self.Wk = np.random.randn(d_model, d_model) * 0.01
self.Wv = np.random.randn(d_model, d_model) * 0.01
self.Wo = np.random.randn(d_model, d_model) * 0.01
def split_heads(self, x):
# x: (B, T, d_model)
B, T, _ = x.shape
x = x.reshape(B, T, self.num_heads, self.d_head)
return x.transpose(0, 2, 1, 3)
# (B, H, T, d_head)
def combine_heads(self, x):
# x: (B, H, T, d_head)
B, H, T, D = x.shape
x = x.transpose(0, 2, 1, 3)
return x.reshape(B, T, H * D)
# (B, T, d_model)
def forward(self, x):
# x: (B, T, d_model)
Q = x @ self.Wq
K = x @ self.Wk
V = x @ self.Wv
Q = self.split_heads(Q)
K = self.split_heads(K)
V = self.split_heads(V)
scores = Q @ K.transpose(0, 1, 3, 2)
scores = scores / np.sqrt(self.d_head)
# scores: (B, H, T, T)
attn_weights = softmax(scores, axis=-1)
# attn_weights: (B, H, T, T)
context = attn_weights @ V
# context: (B, H, T, d_head)
context = self.combine_heads(context)
# context: (B, T, d_model)
output = context @ self.Wo
# output: (B, T, d_model)
return output, attn_weights
B = 2
T = 4
d_model = 8
num_heads = 2
x = np.random.randn(B, T, d_model)
mha = MultiHeadAttention(d_model, num_heads)
output, attn_weights = mha.forward(x)
print(output.shape) # (2, 4, 8)
print(attn_weights.shape) # (2, 2, 4, 4)
以下、forward の処理について見てみます。
def forward(self, x):
# x: (B, T, d_model)
Q = x @ self.Wq
K = x @ self.Wk
V = x @ self.Wv
Q = self.split_heads(Q)
K = self.split_heads(K)
V = self.split_heads(V)
scores = Q @ K.transpose(0, 1, 3, 2)
scores = scores / np.sqrt(self.d_head)
attn_weights = softmax(scores, axis=-1)
context = attn_weights @ V
context = self.combine_heads(context)
output = context @ self.Wo
return output, attn_weights
まず、split_heads について、
def split_heads(self, x):
# x: (B, T, d_model)
B, T, _ = x.shape
x = x.reshape(B, T, self.num_heads, self.d_head)
return x.transpose(0, 2, 1, 3)
# (B, H, T, d_head)
1つの大きな特徴ベクトルを複数の Head に分けて、Head ごとに Attention を計算しやすい形に並べ替えています。例えば、d_model = 8, num_heads = 2 であれば、入力 (B, T, 8) を (B, T, 2, 4) にしています。
これは、Head0, Head1 に対して 4次元ずつに分けているという意味です。
次に、x.transpose(0, 2, 1, 3) により、(B, T, H, D) を (B, H, T, D) に並び替えます。これにより、Head ごとにトークン列をまとめて扱うことができるようになります。
Q = x @ self.Wq
K = x @ self.Wk
V = x @ self.Wv
Q = self.split_heads(Q)
K = self.split_heads(K)
V = self.split_heads(V)
ここで、Q, K, V: (B, T, d_model) → Q, K, V: (B, H, T, d_head) となります。
続いて、
scores = Q @ K.transpose(0, 1, 3, 2)
で、次元は以下のようになります。
K: (B, H, T, D)
K.transpose(0,1,3,2): (B, H, D, T)
scores: (B, H, T, T)
また、
attn_weights = softmax(scores, axis=-1)
context = attn_weights @ V
によって、
attn_weights: (B, H, T, T)
V: (B, H, T, D)
context: (B, H, T, D)
となります。
さらに、
def combine_heads(self, x):
# x: (B, H, T, d_head)
B, H, T, D = x.shape
x = x.transpose(0, 2, 1, 3)
return x.reshape(B, T, H * D)
# (B, T, d_model)
で、(B, H, T, D) → (B, T, H, D) → (B, T, d_model) に戻しています。
最後に、
output = context @ self.Wo
で Head の情報を再び線形変換し、Multi-Head Attention の出力にしています。
結局、Multi-Head Attention では、複数の Head が異なる重みをもち、同時に入力を読み学習していきます。これが、論文での different representation subspaces の意味です。各 Head は異なる特徴空間を見ることになります。
そのため、例えば、
d_model = 512
num_heads = 8
であれば、512次元を64次元ずつの8つのグループに分けて、8つの Head で異なる特徴空間から Attention を行う。ということになります。
まとめ
Multi-Head Attention の「Head」とは、それぞれが独立した Self-Attention を表します。各 Head は異なる学習済み重みを用いて Query・Key・Value を生成し、それぞれ独自の Attention を計算します。そのため、ある Head は文法的な関係、別の Head は意味的な関係、さらに別の Head は長距離の依存関係など、異なる種類の特徴を同時に学習できます。最後に、それらの結果を結合することで、単一の Self-Attention よりも豊かな特徴表現を得られます。