11.5. マルチヘッドアテンション¶
実際には、同じクエリ、キー、値の集合が与えられたとき、同じアテンション機構の異なる振る舞いから得られる知識をモデルに統合させたい場合がある。 たとえば、系列内のさまざまな範囲の依存関係（例：短距離依存と長距離依存）を捉えることである。 したがって、アテンション機構が、クエリ、キー、値の異なる表現部分空間を共同で利用できるようにすると有益かもしれない。
この目的のために、単一のアテンションプーリングを行う代わりに、 クエリ、キー、値を \(h\) 個の独立に学習された線形射影で変換できる。 そして、これら \(h\) 個の射影されたクエリ、キー、値を 並列にアテンションプーリングへ入力する。 最後に、 \(h\) 個のアテンションプーリング出力を連結し、 別の学習された線形射影で変換して 最終出力を生成する。 この設計は マルチヘッドアテンション と呼ばれ、 \(h\) 個のアテンションプーリング出力のそれぞれを 1 つの ヘッド と呼ぶ (Vaswani et al., 2017)。 学習可能な線形変換を行うために全結合層を用いると、 図 11.5.1 はマルチヘッドアテンションを説明している。
図 11.5.1 マルチヘッドアテンション。複数のヘッドを連結してから線形変換する。¶
from d2l import torch as d2l
import math
import torch
from torch import nn
from d2l import mxnet as d2l
import math
from mxnet import autograd, np, npx
from mxnet.gluon import nn
npx.set_np()
from d2l import jax as d2l
from flax import linen as nn
from jax import numpy as jnp
import jax
No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
from d2l import tensorflow as d2l
import tensorflow as tf
11.5.1. モデル¶
マルチヘッドアテンションの実装を示す前に、 このモデルを数学的に定式化しよう。 クエリ \(\mathbf{q} \in \mathbb{R}^{d_q}\)、 キー \(\mathbf{k} \in \mathbb{R}^{d_k}\)、 値 \(\mathbf{v} \in \mathbb{R}^{d_v}\) が与えられたとき、 各アテンションヘッド \(\mathbf{h}_i\) (\(i = 1, \ldots, h\)) は次のように計算される。
(11.5.1)¶\[\mathbf{h}_i = f(\mathbf W_i^{(q)}\mathbf q, \mathbf W_i^{(k)}\mathbf k,\mathbf W_i^{(v)}\mathbf v) \in \mathbb R^{p_v},\]
ここで \(\mathbf W_i^{(q)}\in\mathbb R^{p_q\times d_q}\)、 \(\mathbf W_i^{(k)}\in\mathbb R^{p_k\times d_k}\)、 \(\mathbf W_i^{(v)}\in\mathbb R^{p_v\times d_v}\) は学習可能なパラメータであり、 \(f\) はアテンションプーリングである。 たとえば、 11.3 章 にある 加法アテンションやスケールド・ドット積アテンションが該当する。 マルチヘッドアテンションの出力は、 \(h\) 個のヘッドの連結に対して 学習可能なパラメータ \(\mathbf W_o\in\mathbb R^{p_o\times h p_v}\) による別の線形変換である。
(11.5.2)¶\[\begin{split}\mathbf W_o \begin{bmatrix}\mathbf h_1\\\vdots\\\mathbf h_h\end{bmatrix} \in \mathbb{R}^{p_o}.\end{split}\]
この設計に基づけば、各ヘッドは入力の異なる部分に注意を向けることができる。 単純な重み付き平均よりも高度な関数を表現できる。
11.5.2. 実装¶
実装では、 マルチヘッドアテンションの各ヘッドに
スケールド・ドット積アテンションを選択 する。
計算コストとパラメータ化コストの大幅な増加を避けるため、
\(p_q = p_k = p_v = p_o / h\) とする。
クエリ、キー、値に対する線形変換の出力数を
\(p_q h = p_k h = p_v h = p_o\) に設定すれば、 \(h\)
個のヘッドを並列に計算できることに注意されたい。 以下の実装では、
\(p_o\) は引数 num_hiddens によって指定される。
class MultiHeadAttention(d2l.Module): #@save
"""Multi-head attention."""
def __init__(self, num_hiddens, num_heads, dropout, bias=False, **kwargs):
super().__init__()
self.num_heads = num_heads
self.attention = d2l.DotProductAttention(dropout)
self.W_q = nn.LazyLinear(num_hiddens, bias=bias)
self.W_k = nn.LazyLinear(num_hiddens, bias=bias)
self.W_v = nn.LazyLinear(num_hiddens, bias=bias)
self.W_o = nn.LazyLinear(num_hiddens, bias=bias)
def forward(self, queries, keys, values, valid_lens):
# Shape of queries, keys, or values:
# (batch_size, no. of queries or key-value pairs, num_hiddens)
# Shape of valid_lens: (batch_size,) or (batch_size, no. of queries)
# After transposing, shape of output queries, keys, or values:
# (batch_size * num_heads, no. of queries or key-value pairs,
# num_hiddens / num_heads)
queries = self.transpose_qkv(self.W_q(queries))
keys = self.transpose_qkv(self.W_k(keys))
values = self.transpose_qkv(self.W_v(values))
if valid_lens is not None:
# On axis 0, copy the first item (scalar or vector) for num_heads
# times, then copy the next item, and so on
valid_lens = torch.repeat_interleave(
valid_lens, repeats=self.num_heads, dim=0)
# Shape of output: (batch_size * num_heads, no. of queries,
# num_hiddens / num_heads)
output = self.attention(queries, keys, values, valid_lens)
# Shape of output_concat: (batch_size, no. of queries, num_hiddens)
output_concat = self.transpose_output(output)
return self.W_o(output_concat)
class MultiHeadAttention(d2l.Module): #@save
"""Multi-head attention."""
def __init__(self, num_hiddens, num_heads, dropout, use_bias=False,
**kwargs):
super().__init__()
self.num_heads = num_heads
self.attention = d2l.DotProductAttention(dropout)
self.W_q = nn.Dense(num_hiddens, use_bias=use_bias, flatten=False)
self.W_k = nn.Dense(num_hiddens, use_bias=use_bias, flatten=False)
self.W_v = nn.Dense(num_hiddens, use_bias=use_bias, flatten=False)
self.W_o = nn.Dense(num_hiddens, use_bias=use_bias, flatten=False)
def forward(self, queries, keys, values, valid_lens):
# Shape of queries, keys, or values:
# (batch_size, no. of queries or key-value pairs, num_hiddens)
# Shape of valid_lens: (batch_size,) or (batch_size, no. of queries)
# After transposing, shape of output queries, keys, or values:
# (batch_size * num_heads, no. of queries or key-value pairs,
# num_hiddens / num_heads)
queries = self.transpose_qkv(self.W_q(queries))
keys = self.transpose_qkv(self.W_k(keys))
values = self.transpose_qkv(self.W_v(values))
if valid_lens is not None:
# On axis 0, copy the first item (scalar or vector) for num_heads
# times, then copy the next item, and so on
valid_lens = valid_lens.repeat(self.num_heads, axis=0)
# Shape of output: (batch_size * num_heads, no. of queries,
# num_hiddens / num_heads)
output = self.attention(queries, keys, values, valid_lens)
# Shape of output_concat: (batch_size, no. of queries, num_hiddens)
output_concat = self.transpose_output(output)
return self.W_o(output_concat)
class MultiHeadAttention(nn.Module): #@save
num_hiddens: int
num_heads: int
dropout: float
bias: bool = False
def setup(self):
self.attention = d2l.DotProductAttention(self.dropout)
self.W_q = nn.Dense(self.num_hiddens, use_bias=self.bias)
self.W_k = nn.Dense(self.num_hiddens, use_bias=self.bias)
self.W_v = nn.Dense(self.num_hiddens, use_bias=self.bias)
self.W_o = nn.Dense(self.num_hiddens, use_bias=self.bias)
@nn.compact
def __call__(self, queries, keys, values, valid_lens, training=False):
# Shape of queries, keys, or values:
# (batch_size, no. of queries or key-value pairs, num_hiddens)
# Shape of valid_lens: (batch_size,) or (batch_size, no. of queries)
# After transposing, shape of output queries, keys, or values:
# (batch_size * num_heads, no. of queries or key-value pairs,
# num_hiddens / num_heads)
queries = self.transpose_qkv(self.W_q(queries))
keys = self.transpose_qkv(self.W_k(keys))
values = self.transpose_qkv(self.W_v(values))
if valid_lens is not None:
# On axis 0, copy the first item (scalar or vector) for num_heads
# times, then copy the next item, and so on
valid_lens = jnp.repeat(valid_lens, self.num_heads, axis=0)
# Shape of output: (batch_size * num_heads, no. of queries,
# num_hiddens / num_heads)
output, attention_weights = self.attention(
queries, keys, values, valid_lens, training=training)
# Shape of output_concat: (batch_size, no. of queries, num_hiddens)
output_concat = self.transpose_output(output)
return self.W_o(output_concat), attention_weights
class MultiHeadAttention(d2l.Module): #@save
"""Multi-head attention."""
def __init__(self, key_size, query_size, value_size, num_hiddens,
num_heads, dropout, bias=False, **kwargs):
super().__init__()
self.num_heads = num_heads
self.attention = d2l.DotProductAttention(dropout)
self.W_q = tf.keras.layers.Dense(num_hiddens, use_bias=bias)
self.W_k = tf.keras.layers.Dense(num_hiddens, use_bias=bias)
self.W_v = tf.keras.layers.Dense(num_hiddens, use_bias=bias)
self.W_o = tf.keras.layers.Dense(num_hiddens, use_bias=bias)
def call(self, queries, keys, values, valid_lens, **kwargs):
# Shape of queries, keys, or values:
# (batch_size, no. of queries or key-value pairs, num_hiddens)
# Shape of valid_lens: (batch_size,) or (batch_size, no. of queries)
# After transposing, shape of output queries, keys, or values:
# (batch_size * num_heads, no. of queries or key-value pairs,
# num_hiddens / num_heads)
queries = self.transpose_qkv(self.W_q(queries))
keys = self.transpose_qkv(self.W_k(keys))
values = self.transpose_qkv(self.W_v(values))
if valid_lens is not None:
# On axis 0, copy the first item (scalar or vector) for num_heads
# times, then copy the next item, and so on
valid_lens = tf.repeat(valid_lens, repeats=self.num_heads, axis=0)
# Shape of output: (batch_size * num_heads, no. of queries,
# num_hiddens / num_heads)
output = self.attention(queries, keys, values, valid_lens, **kwargs)
# Shape of output_concat: (batch_size, no. of queries, num_hiddens)
output_concat = self.transpose_output(output)
return self.W_o(output_concat)
複数ヘッドの 並列計算 を可能にするために、 上の MultiHeadAttention
クラスでは以下に定義する 2 つの転置メソッドを使う。 具体的には、
transpose_output メソッドは transpose_qkv
メソッドの操作を逆にする。
@d2l.add_to_class(MultiHeadAttention) #@save
def transpose_qkv(self, X):
"""Transposition for parallel computation of multiple attention heads."""
# Shape of input X: (batch_size, no. of queries or key-value pairs,
# num_hiddens). Shape of output X: (batch_size, no. of queries or
# key-value pairs, num_heads, num_hiddens / num_heads)
X = X.reshape(X.shape[0], X.shape[1], self.num_heads, -1)
# Shape of output X: (batch_size, num_heads, no. of queries or key-value
# pairs, num_hiddens / num_heads)
X = X.permute(0, 2, 1, 3)
# Shape of output: (batch_size * num_heads, no. of queries or key-value
# pairs, num_hiddens / num_heads)
return X.reshape(-1, X.shape[2], X.shape[3])
@d2l.add_to_class(MultiHeadAttention) #@save
def transpose_output(self, X):
"""Reverse the operation of transpose_qkv."""
X = X.reshape(-1, self.num_heads, X.shape[1], X.shape[2])
X = X.permute(0, 2, 1, 3)
return X.reshape(X.shape[0], X.shape[1], -1)
@d2l.add_to_class(MultiHeadAttention) #@save
def transpose_qkv(self, X):
"""Transposition for parallel computation of multiple attention heads."""
# Shape of input X: (batch_size, no. of queries or key-value pairs,
# num_hiddens). Shape of output X: (batch_size, no. of queries or
# key-value pairs, num_heads, num_hiddens / num_heads)
X = X.reshape(X.shape[0], X.shape[1], self.num_heads, -1)
# Shape of output X: (batch_size, num_heads, no. of queries or key-value
# pairs, num_hiddens / num_heads)
X = X.transpose(0, 2, 1, 3)
# Shape of output: (batch_size * num_heads, no. of queries or key-value
# pairs, num_hiddens / num_heads)
return X.reshape(-1, X.shape[2], X.shape[3])
@d2l.add_to_class(MultiHeadAttention) #@save
def transpose_output(self, X):
"""Reverse the operation of transpose_qkv."""
X = X.reshape(-1, self.num_heads, X.shape[1], X.shape[2])
X = X.transpose(0, 2, 1, 3)
return X.reshape(X.shape[0], X.shape[1], -1)
@d2l.add_to_class(MultiHeadAttention) #@save
def transpose_qkv(self, X):
"""Transposition for parallel computation of multiple attention heads."""
# Shape of input X: (batch_size, no. of queries or key-value pairs,
# num_hiddens). Shape of output X: (batch_size, no. of queries or
# key-value pairs, num_heads, num_hiddens / num_heads)
X = X.reshape((X.shape[0], X.shape[1], self.num_heads, -1))
# Shape of output X: (batch_size, num_heads, no. of queries or key-value
# pairs, num_hiddens / num_heads)
X = jnp.transpose(X, (0, 2, 1, 3))
# Shape of output: (batch_size * num_heads, no. of queries or key-value
# pairs, num_hiddens / num_heads)
return X.reshape((-1, X.shape[2], X.shape[3]))
@d2l.add_to_class(MultiHeadAttention) #@save
def transpose_output(self, X):
"""Reverse the operation of transpose_qkv."""
X = X.reshape((-1, self.num_heads, X.shape[1], X.shape[2]))
X = jnp.transpose(X, (0, 2, 1, 3))
return X.reshape((X.shape[0], X.shape[1], -1))
@d2l.add_to_class(MultiHeadAttention) #@save
def transpose_qkv(self, X):
"""Transposition for parallel computation of multiple attention heads."""
# Shape of input X: (batch_size, no. of queries or key-value pairs,
# num_hiddens). Shape of output X: (batch_size, no. of queries or
# key-value pairs, num_heads, num_hiddens / num_heads)
X = tf.reshape(X, shape=(X.shape[0], X.shape[1], self.num_heads, -1))
# Shape of output X: (batch_size, num_heads, no. of queries or key-value
# pairs, num_hiddens / num_heads)
X = tf.transpose(X, perm=(0, 2, 1, 3))
# Shape of output: (batch_size * num_heads, no. of queries or key-value
# pairs, num_hiddens / num_heads)
return tf.reshape(X, shape=(-1, X.shape[2], X.shape[3]))
@d2l.add_to_class(MultiHeadAttention) #@save
def transpose_output(self, X):
"""Reverse the operation of transpose_qkv."""
X = tf.reshape(X, shape=(-1, self.num_heads, X.shape[1], X.shape[2]))
X = tf.transpose(X, perm=(0, 2, 1, 3))
return tf.reshape(X, shape=(X.shape[0], X.shape[1], -1))
実装した MultiHeadAttention クラスを、
キーと値が同じであるおもちゃの例を使って テストしてみよう。 その結果、
マルチヘッドアテンションの出力形状は (batch_size, num_queries,
num_hiddens) になる。
num_hiddens, num_heads = 100, 5
attention = MultiHeadAttention(num_hiddens, num_heads, 0.5)
batch_size, num_queries, num_kvpairs = 2, 4, 6
valid_lens = d2l.tensor([3, 2])
X = d2l.ones((batch_size, num_queries, num_hiddens))
Y = d2l.ones((batch_size, num_kvpairs, num_hiddens))
d2l.check_shape(attention(X, Y, Y, valid_lens),
(batch_size, num_queries, num_hiddens))
num_hiddens, num_heads = 100, 5
attention = MultiHeadAttention(num_hiddens, num_heads, 0.5)
attention.initialize()
num_hiddens, num_heads = 100, 5
attention = MultiHeadAttention(num_hiddens, num_heads, 0.5)
num_hiddens, num_heads = 100, 5
attention = MultiHeadAttention(num_hiddens, num_hiddens, num_hiddens,
num_hiddens, num_heads, 0.5)
11.5.3. 要約¶
マルチヘッドアテンションは、 クエリ、キー、値の異なる表現部分空間を通じて、 同じアテンションプーリングの知識を統合する。 マルチヘッドアテンションの複数ヘッドを並列に計算するには、 適切なテンソル操作が必要である。
11.5.4. 演習¶
この実験における複数ヘッドのアテンション重みを可視化せよ。
マルチヘッドアテンションに基づく学習済みモデルがあり、予測速度を上げるために重要度の低いアテンションヘッドを剪定したいとする。アテンションヘッドの重要度を測定する実験をどのように設計できるか？