学ぶマルチヘッドアテンションの実装方法 | トランスフォーマーコンポーネントの構築

メニューを表示するにはスワイプしてください


              12345678910111213141516171819202122232425262728293031323334353637383940
            
import numpy as np

# Toy input: 3 tokens, embedding size 6
x = np.array([
    [0.1, 0.2, 0.3, 0.4, 0.5, 0.6],
    [0.6, 0.5, 0.4, 0.3, 0.2, 0.1],
    [0.2, 0.1, 0.4, 0.3, 0.6, 0.5],
])

# Multi-head attention parameters
num_heads = 2
embed_size = x.shape[1]
head_dim = embed_size // num_heads

# Random weights for queries, keys, and values for each head
np.random.seed(42)
W_q = np.random.randn(num_heads, embed_size, head_dim)
W_k = np.random.randn(num_heads, embed_size, head_dim)
W_v = np.random.randn(num_heads, embed_size, head_dim)

attention_outputs = []

for head in range(num_heads):
    # Linear projections
    Q = x @ W_q[head]
    K = x @ W_k[head]
    V = x @ W_v[head]
    
    # Scaled dot-product attention
    scores = Q @ K.T / np.sqrt(head_dim)
    attention_weights = np.exp(scores) / np.sum(np.exp(scores), axis=-1, keepdims=True)
    head_output = attention_weights @ V
    attention_outputs.append(head_output)

# Concatenate outputs from all heads
multi_head_output = np.concatenate(attention_outputs, axis=-1)

print("Output from head 0:\n", attention_outputs[0])
print("Output from head 1:\n", attention_outputs[1])
print("Concatenated multi-head output:\n", multi_head_output)

マルチヘッドアテンションは、同じ自己注意メカニズムを複数のヘッドで独立して並列に適用することで機能します。各ヘッドはクエリ、キー、バリュー（W_q、W_k、W_v）用の独自の重み行列を使用し、入力を低次元の部分空間に射影します。この例の入力 x は3つのトークンを含み、それぞれ埋め込みサイズは6で、2つのヘッドに分割されます。

各ヘッドごとに、コードは標準的な自己注意の手順を繰り返します。クエリ（Q）、キー（K）、バリュー（V）行列を計算し、Q と K.T のドット積（適切にスケーリング）でアテンションスコアを算出し、ソフトマックス関数でアテンション重みを得ます。これらの重みを使ってバリューベクトルを合成し、各ヘッドの出力を生成します。

ここで新しいメカニズムはなく、マルチヘッドアテンションは単に各ヘッドごとに並列で自己注意処理を行うものです。すべてのヘッドの処理が終わった後、それぞれの出力は最後の次元で連結されます。この手法により、各ヘッドが入力の異なる側面に注目できるため、モデルはデータ内の多様な関係性を捉えることが可能になります。出力結果から、同じ自己注意メカニズムを複数回実行し、その結果を組み合わせることで、各トークンに対してより豊かな表現が得られることが示されています。

すべて明確でしたか？

フィードバックありがとうございます！

セクション 2. 章 2

AIに質問する

何でも質問するか、提案された質問の1つを試してチャットを始めてください

セクション 2. 章 2