TransformerのSelf-Attentionを理解するためのシンプルな実装

Transformerモデルの根幹にある「Self-Attention（自己注意）」は、各単語が文中の他の単語にどれくらい注意を払うかを計算します。今回は、その仕組みをシンプルなPythonコードで可視化しながら説明します。

目次

1. Self-Attentionとは？
2. 基本的な実装
3. 実行結果と可視化
4. まとめ

---

1. Self-Attentionとは？

簡単に言うと、Self-Attentionは「文中の各単語が他の単語をどれくらい重要だと思うか」を計算するメカニズムです。これにより、文の中での単語間の関係を捉えることができます。たとえば、「私はリンゴが好き」という文で、"リンゴ"と"好き"は密接に関係していますよね。Self-Attentionはそのような関係を計算します。

---

2. 基本的な実装

まずは、Self-Attentionを構成する要素を簡単に見ていきましょう。

必要な要素

• Query（Q）: 各単語が他の単語を「どれくらい見るか？」という問い。
• Key（K）: 単語が他の単語にどれくらいの重要度を持つかを表す。
• Value（V）: 実際の情報（特徴量）。

Transformerでは、これらを使って注意スコアを計算し、文のコンテキストを捉えます。

以下のコードでは、3つの単語がそれぞれ4次元のベクトルとして表されており、それぞれに対してQuery, Key, Valueを計算します。

import numpy as np

# 入力：3単語、各単語は4次元のベクトル（埋め込み）
x = np.array([
    [1, 0, 1, 0],  # 単語1
    [0, 2, 0, 2],  # 単語2
    [1, 1, 1, 1],  # 単語3
])

# 重み行列（実際は学習されるパラメータ）
W_q = np.random.rand(4, 4)  # Query
W_k = np.random.rand(4, 4)  # Key
W_v = np.random.rand(4, 4)  # Value

# Query, Key, Value の計算
Q = x @ W_q  # (3,4)
K = x @ W_k  # (3,4)
V = x @ W_v  # (3,4)

# Attentionスコアの計算（スケールドドット積）
d_k = Q.shape[-1]
scores = Q @ K.T / np.sqrt(d_k)  # (3,3)

# ソフトマックスで正規化（行ごとに）
def softmax(x):
    e_x = np.exp(x - np.max(x, axis=-1, keepdims=True))
    return e_x / np.sum(e_x, axis=-1, keepdims=True)

attn_weights = softmax(scores)  # (3,3)

# 重み付き和（各単語が他の単語をどれだけ「見るか」を反映）
output = attn_weights @ V  # (3,4)

# 結果表示
output

---

3. 実行結果と可視化

実行結果を可視化することで、Self-Attentionがどのように機能しているのかを直感的に理解できます。

可視化内容

1. Query (Q)、Key (K)、Value (V)の行列
   それぞれの行列が、単語ごとにどのような特徴を持っているのかを示します。

2. Attention Weights（注意重み）
   各単語が他の単語をどれくらい「見ているか」の重みを表した行列です。これにより、文脈における単語同士の関係性が分かります。

実際に可視化した結果がこちらです：

図は次の情報を示しています：

• 左上：Query（Q）の行列
• 右上：Key（K）の行列
• 左下：Value（V）の行列
• 右下：Attention Weights（注意重み）

また、最終的に得られた出力ベクトルは、単語同士の注意の度合いを反映した加重平均です。

---

4. まとめ

Self-Attentionは、文中の単語がどれくらい他の単語を重視するかを計算する重要なメカニズムです。これにより、Transformerは長距離依存関係を学習し、文脈に基づいた情報を適切に反映させることができます。

今回のシンプルな実装と可視化を通して、その基本的な動作を理解できたかと思います。もちろん、実際のTransformerモデルではさらに多くの工夫がありますが、Self-Attentionの本質はこのようにシンプルに表現できます。