2017年に Vaswani et al. が発表した論文「Attention Is All You Need」は、Transformer アーキテクチャを提案し、自然言語処理のみならず深層学習全体に革命を起こしました。RNN や CNN を一切使わず、Attention 機構のみで系列変換を行うという大胆な設計は、並列計算の効率性と長距離依存の捕捉能力において従来手法を大きく凌駕しました。
本記事では、この原論文の核心となるアイデアを数式レベルで徹底的に解読します。特に、Scaled Dot-Product Attention における $\sqrt{d_k}$ スケーリングの数学的根拠、位置エンコーディングの設計原理など、「なぜそうなっているのか」を厳密に導出します。
本記事の内容
- RNN / CNN の限界と Transformer の動機
- Scaled Dot-Product Attention の定式化と $\sqrt{d_k}$ スケーリングの数学的根拠
- マルチヘッド Attention の定式化
- Encoder / Decoder の構造（残差接続 + LayerNorm）
- 位置エンコーディングの設計原理と相対位置の線形変換可能性の証明
- 学習率スケジューリング（warmup）
- Python で Transformer Encoder のスクラッチ実装
前提知識
この記事を読む前に、以下の記事を読んでおくと理解が深まります。
RNN / CNN の限界
RNN の問題点
RNN（Recurrent Neural Network）は系列データを逐次的に処理します。時刻 $t$ の隠れ状態 $\bm{h}_t$ は
$$ \bm{h}_t = f(\bm{h}_{t-1}, \bm{x}_t) $$
と定義され、$\bm{h}_t$ の計算には $\bm{h}_{t-1}$ が必要です。この逐次依存性は以下の問題を引き起こします。
- 並列化の困難: $\bm{h}_1, \bm{h}_2, \dots, \bm{h}_T$ を並列に計算できない。計算量は $O(T)$ のステップが必要。
- 長距離依存の消失: 勾配が多数のステップを通じて伝搬するため、勾配消失・勾配爆発が発生しやすい。LSTM や GRU はこれを緩和するがメモリ上限がある。
- 計算効率: 系列長 $T$ に対して $O(T)$ の逐次ステップが必要で、GPU の並列性を活かせない。
CNN の問題点
CNN を系列処理に用いる場合（WaveNet, ConvS2S 等）、畳み込みカーネルの局所性により長距離依存の捕捉にはカーネルの積み重ね（$O(\log T)$ or $O(T/k)$ 層）が必要になります。
Transformer の解決策
Transformer はSelf-Attentionにより、系列中の任意の2点間を1層で直接接続します。計算の逐次ステップは $O(1)$（系列長に依存しない）であり、全位置の計算を完全に並列化できます。ただし、Attention の計算量は $O(T^2 \cdot d)$ であり、系列長に対して二次の計算量が必要です。
| 手法 | 層あたり計算量 | 逐次ステップ | 最大パス長 |
|---|---|---|---|
| RNN | $O(T \cdot d^2)$ | $O(T)$ | $O(T)$ |
| CNN | $O(T \cdot k \cdot d^2)$ | $O(1)$ | $O(\log_k T)$ |
| Self-Attention | $O(T^2 \cdot d)$ | $O(1)$ | $O(1)$ |
Scaled Dot-Product Attention
定義
Transformer の核心は Scaled Dot-Product Attention です。クエリ $\bm{Q} \in \mathbb{R}^{T_q \times d_k}$、キー $\bm{K} \in \mathbb{R}^{T_k \times d_k}$、バリュー $\bm{V} \in \mathbb{R}^{T_k \times d_v}$ に対して
$$ \text{Attention}(\bm{Q}, \bm{K}, \bm{V}) = \text{softmax}\left(\frac{\bm{Q}\bm{K}^T}{\sqrt{d_k}}\right)\bm{V} $$
と定義されます。$\bm{Q}\bm{K}^T \in \mathbb{R}^{T_q \times T_k}$ はクエリとキーの内積行列であり、softmax は各行（各クエリ）に対して独立に適用されます。
$\sqrt{d_k}$ スケーリングの数学的根拠
なぜ $\sqrt{d_k}$ で割るのでしょうか。これには明確な数学的根拠があります。
クエリベクトル $\bm{q} \in \mathbb{R}^{d_k}$ とキーベクトル $\bm{k} \in \mathbb{R}^{d_k}$ の各成分が独立に平均0、分散1の分布に従うと仮定します。すなわち
$$ q_i \sim (0, 1), \quad k_i \sim (0, 1), \quad i = 1, \dots, d_k $$
内積 $\bm{q} \cdot \bm{k} = \sum_{i=1}^{d_k} q_i k_i$ の期待値と分散を計算しましょう。
期待値:
$$ \begin{align} \mathbb{E}[\bm{q} \cdot \bm{k}] &= \mathbb{E}\left[\sum_{i=1}^{d_k} q_i k_i\right] \\ &= \sum_{i=1}^{d_k} \mathbb{E}[q_i] \cdot \mathbb{E}[k_i] \quad (\because q_i \text{ と } k_i \text{ は独立}) \\ &= \sum_{i=1}^{d_k} 0 \cdot 0 = 0 \end{align} $$
分散:
$$ \begin{align} \text{Var}[\bm{q} \cdot \bm{k}] &= \text{Var}\left[\sum_{i=1}^{d_k} q_i k_i\right] \\ &= \sum_{i=1}^{d_k} \text{Var}[q_i k_i] \quad (\because \text{各項は独立}) \end{align} $$
各項 $q_i k_i$ の分散を求めます。$q_i$ と $k_i$ は独立で $\mathbb{E}[q_i] = \mathbb{E}[k_i] = 0$、$\text{Var}[q_i] = \text{Var}[k_i] = 1$ なので
$$ \begin{align} \text{Var}[q_i k_i] &= \mathbb{E}[q_i^2 k_i^2] – (\mathbb{E}[q_i k_i])^2 \\ &= \mathbb{E}[q_i^2] \cdot \mathbb{E}[k_i^2] – 0^2 \\ &= (\text{Var}[q_i] + (\mathbb{E}[q_i])^2)(\text{Var}[k_i] + (\mathbb{E}[k_i])^2) \\ &= (1 + 0)(1 + 0) = 1 \end{align} $$
したがって
$$ \text{Var}[\bm{q} \cdot \bm{k}] = \sum_{i=1}^{d_k} 1 = d_k $$
つまり、内積の分散は $d_k$ に比例して増大します。$d_k$ が大きいとき、内積の値は非常に大きな正や負の値を取り得ます。
softmax 関数 $\text{softmax}(z_j) = \frac{e^{z_j}}{\sum_i e^{z_i}}$ の勾配を考えると、$z_j$ の値が大きい領域では $e^{z_j}$ が他の項に比べて圧倒的に大きくなり、出力がほぼワンホット（1つの要素だけが1、他が0）になります。この飽和領域では勾配が極めて小さくなり、学習が停滞します。
$\sqrt{d_k}$ でスケーリングすることで
$$ \text{Var}\left[\frac{\bm{q} \cdot \bm{k}}{\sqrt{d_k}}\right] = \frac{1}{d_k} \cdot d_k = 1 $$
となり、内積の分散が $d_k$ に依存せず常に1に保たれます。これにより、softmax が飽和しにくい領域で動作し、安定した勾配が得られます。
Attention の直感的理解
Attention の動作を直感的に理解しましょう。$\bm{A} = \text{softmax}(\bm{Q}\bm{K}^T / \sqrt{d_k})$ とすると
$$ \text{Attention}(\bm{Q}, \bm{K}, \bm{V}) = \bm{A}\bm{V} $$
$\bm{A}$ の $(i, j)$ 成分 $a_{ij}$ は、クエリ $\bm{q}_i$ とキー $\bm{k}_j$ の類似度に基づく重みです。出力の $i$ 行目は
$$ \text{output}_i = \sum_{j=1}^{T_k} a_{ij} \bm{v}_j $$
つまり、各クエリに対して、関連性の高いバリューを重み付き平均する操作です。これは「情報検索」のアナロジーで理解できます。クエリで検索し、キーとの一致度に応じてバリューを取得する仕組みです。
マルチヘッド Attention
定式化
単一の Attention では、$d_{\text{model}}$ 次元の空間で1つの「関連性のパターン」しか捉えられません。マルチヘッド Attention は、異なる表現部分空間で複数の Attention を並列に計算します。
$$ \text{MultiHead}(\bm{Q}, \bm{K}, \bm{V}) = \text{Concat}(\text{head}_1, \dots, \text{head}_h) \bm{W}^O $$
ここで各ヘッドは
$$ \text{head}_i = \text{Attention}(\bm{Q}\bm{W}_i^Q, \bm{K}\bm{W}_i^K, \bm{V}\bm{W}_i^V) $$
パラメータの次元は
$$ \bm{W}_i^Q \in \mathbb{R}^{d_{\text{model}} \times d_k}, \quad \bm{W}_i^K \in \mathbb{R}^{d_{\text{model}} \times d_k}, \quad \bm{W}_i^V \in \mathbb{R}^{d_{\text{model}} \times d_v}, \quad \bm{W}^O \in \mathbb{R}^{hd_v \times d_{\text{model}}} $$
原論文では $h = 8$, $d_k = d_v = d_{\text{model}} / h = 64$ としています。
計算量の保存
1ヘッドあたりの次元を $d_k = d_v = d_{\text{model}} / h$ とするため、$h$ ヘッドの計算量の合計は
$$ h \times O(T^2 \cdot d_k) = h \times O\left(T^2 \cdot \frac{d_{\text{model}}}{h}\right) = O(T^2 \cdot d_{\text{model}}) $$
これは単一ヘッドで $d_k = d_{\text{model}}$ とした場合の計算量と同じです。つまり、マルチヘッドにすることで計算量を増やさずに表現力を向上させています。
Encoder / Decoder の構造
Encoder
Encoder は $N = 6$ 層の同一構造のレイヤーを積み重ねます。各レイヤーは以下の2つのサブレイヤーで構成されます。
- マルチヘッド Self-Attention
- Position-wise Feed-Forward Network (FFN)
各サブレイヤーには残差接続（residual connection）と Layer Normalization が適用されます。
$$ \text{output} = \text{LayerNorm}(\bm{x} + \text{Sublayer}(\bm{x})) $$
残差接続は深いネットワークでの勾配の伝搬を助け、LayerNorm は各層の出力の分布を安定させます。
Position-wise FFN
FFN は各位置に独立に適用される2層の全結合ネットワークです。
$$ \text{FFN}(\bm{x}) = \text{ReLU}(\bm{x}\bm{W}_1 + \bm{b}_1)\bm{W}_2 + \bm{b}_2 $$
ここで $\bm{W}_1 \in \mathbb{R}^{d_{\text{model}} \times d_{ff}}$, $\bm{W}_2 \in \mathbb{R}^{d_{ff} \times d_{\text{model}}}$ で、原論文では $d_{ff} = 2048$, $d_{\text{model}} = 512$ です。
内部次元 $d_{ff}$ が $d_{\text{model}}$ の4倍に拡張されるため、FFN は各位置のトークン表現に対して非線形変換を行い、Attention で集約された情報を処理する役割を担います。
Layer Normalization
Layer Normalization は、各サンプルの各位置について、特徴次元方向で正規化を行います。入力 $\bm{x} \in \mathbb{R}^{d}$ に対して
$$ \text{LayerNorm}(\bm{x}) = \bm{\gamma} \odot \frac{\bm{x} – \mu}{\sqrt{\sigma^2 + \epsilon}} + \bm{\beta} $$
ここで
$$ \mu = \frac{1}{d} \sum_{i=1}^{d} x_i, \quad \sigma^2 = \frac{1}{d} \sum_{i=1}^{d} (x_i – \mu)^2 $$
$\bm{\gamma}, \bm{\beta} \in \mathbb{R}^d$ は学習可能なパラメータ、$\epsilon$ は数値安定性のための小さな定数です。
Decoder
Decoder も $N = 6$ 層ですが、各レイヤーは3つのサブレイヤーを持ちます。
- Masked マルチヘッド Self-Attention: 未来の位置を参照しないようマスクを適用
- マルチヘッド Cross-Attention: Encoder の出力をキー・バリューとして使用
- Position-wise FFN
Masked Self-Attention では、位置 $i$ のクエリが位置 $j > i$ のキーを参照できないように、Attention スコアの該当位置を $-\infty$ に設定します。
$$ \text{mask}_{ij} = \begin{cases} 0 & \text{if } j \leq i \\ -\infty & \text{if } j > i \end{cases} $$
$$ \text{MaskedAttention}(\bm{Q}, \bm{K}, \bm{V}) = \text{softmax}\left(\frac{\bm{Q}\bm{K}^T}{\sqrt{d_k}} + \text{mask}\right)\bm{V} $$
$-\infty$ を加えることで、softmax 後の該当位置の重みが0になります。
位置エンコーディング
なぜ位置情報が必要か
Self-Attention は入力の順序に不変（permutation invariant）です。これを厳密に示しましょう。
入力行列 $\bm{X} \in \mathbb{R}^{T \times d}$ の行を置換行列 $\bm{P}$ で並べ替えると $\bm{P}\bm{X}$ になります。Self-Attention の出力は
$$ \begin{align} \text{SA}(\bm{P}\bm{X}) &= \text{softmax}\left(\frac{(\bm{P}\bm{X}\bm{W}^Q)(\bm{P}\bm{X}\bm{W}^K)^T}{\sqrt{d_k}}\right)(\bm{P}\bm{X}\bm{W}^V) \\ &= \text{softmax}\left(\frac{\bm{P}\bm{X}\bm{W}^Q(\bm{W}^K)^T\bm{X}^T\bm{P}^T}{\sqrt{d_k}}\right)\bm{P}\bm{X}\bm{W}^V \\ &= \bm{P}\,\text{softmax}\left(\frac{\bm{X}\bm{W}^Q(\bm{W}^K)^T\bm{X}^T}{\sqrt{d_k}}\right)\bm{X}\bm{W}^V \\ &= \bm{P}\,\text{SA}(\bm{X}) \end{align} $$
3行目で、$\bm{P}\bm{A}\bm{P}^T$ に softmax を行方向に適用した結果が $\bm{P}\,\text{softmax}(\bm{A})\bm{P}^T$ であること、さらに $\bm{P}\,\text{softmax}(\bm{A})\bm{P}^T \cdot \bm{P}\bm{B} = \bm{P}\,\text{softmax}(\bm{A})\bm{B}$ を用いました。
つまり、入力の順序を変えても、出力は同じ置換を受けるだけで、Self-Attention 自体は位置情報を区別できません。自然言語では語順が意味を決定する（”犬が人を噛む” と “人が犬を噛む”は異なる）ため、位置情報の注入が不可欠です。
正弦波位置エンコーディング
原論文では、以下の正弦波関数による位置エンコーディングを提案しています。
$$ \begin{align} PE(pos, 2i) &= \sin\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right) \\ PE(pos, 2i+1) &= \cos\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right) \end{align} $$
ここで $pos$ はトークンの位置（$0, 1, 2, \dots$）、$i$ は次元のインデックス（$0, 1, \dots, d_{\text{model}}/2 – 1$）です。
各次元ペア $(2i, 2i+1)$ は周波数 $\omega_i = 1/10000^{2i/d_{\text{model}}}$ の正弦波と余弦波を形成します。$i = 0$ では波長が $2\pi$ と短く、$i$ が大きくなるにつれて波長は $2\pi \times 10000$ まで長くなります。
相対位置の線形変換可能性の証明
正弦波位置エンコーディングの重要な性質は、任意の固定オフセット $k$ に対して、$PE(pos + k)$ が $PE(pos)$ の線形変換で表せることです。これを証明しましょう。
次元ペア $(2i, 2i+1)$ に注目し、$\omega_i = 1/10000^{2i/d_{\text{model}}}$ と略記します。
$$ \begin{pmatrix} PE(pos+k, 2i) \\ PE(pos+k, 2i+1) \end{pmatrix} = \begin{pmatrix} \sin(\omega_i(pos+k)) \\ \cos(\omega_i(pos+k)) \end{pmatrix} $$
三角関数の加法定理を適用します。
$$ \begin{align} \sin(\omega_i(pos+k)) &= \sin(\omega_i \cdot pos) \cos(\omega_i \cdot k) + \cos(\omega_i \cdot pos) \sin(\omega_i \cdot k) \\ \cos(\omega_i(pos+k)) &= \cos(\omega_i \cdot pos) \cos(\omega_i \cdot k) – \sin(\omega_i \cdot pos) \sin(\omega_i \cdot k) \end{align} $$
これを行列形式で書くと
$$ \begin{pmatrix} PE(pos+k, 2i) \\ PE(pos+k, 2i+1) \end{pmatrix} = \begin{pmatrix} \cos(\omega_i k) & \sin(\omega_i k) \\ -\sin(\omega_i k) & \cos(\omega_i k) \end{pmatrix} \begin{pmatrix} PE(pos, 2i) \\ PE(pos, 2i+1) \end{pmatrix} $$
右辺の $2 \times 2$ 行列は回転行列 $\bm{R}_k^{(i)}$ であり、$pos$ に依存せず $k$ のみで決まります。
$$ \bm{R}_k^{(i)} = \begin{pmatrix} \cos(\omega_i k) & \sin(\omega_i k) \\ -\sin(\omega_i k) & \cos(\omega_i k) \end{pmatrix} $$
全次元をまとめると、$PE(pos + k)$ はブロック対角行列 $\bm{R}_k = \text{diag}(\bm{R}_k^{(0)}, \bm{R}_k^{(1)}, \dots)$ を用いて
$$ PE(pos + k) = \bm{R}_k \cdot PE(pos) $$
と書けます。これは相対位置 $k$ が線形変換として表現できることを意味し、モデルが相対位置情報を学習しやすくなります。
位置エンコーディングの内積
$PE(pos)$ と $PE(pos + k)$ の内積を計算してみましょう。
$$ \begin{align} PE(pos)^T PE(pos+k) &= \sum_{i=0}^{d_{\text{model}}/2 – 1} \left[\sin(\omega_i \cdot pos) \sin(\omega_i(pos+k)) + \cos(\omega_i \cdot pos)\cos(\omega_i(pos+k))\right] \\ &= \sum_{i=0}^{d_{\text{model}}/2 – 1} \cos(\omega_i \cdot k) \end{align} $$
最後の等号で積和の公式 $\cos(\alpha – \beta) = \cos\alpha\cos\beta + \sin\alpha\sin\beta$ を用いました。結果は相対位置 $k$ のみに依存し、絶対位置 $pos$ に依存しません。これは、位置エンコーディングの内積が位置の相対関係を反映するという望ましい性質です。
学習率スケジューリング（Warmup）
原論文では、以下の学習率スケジュールを使用します。
$$ lr = d_{\text{model}}^{-0.5} \cdot \min(\text{step}^{-0.5}, \text{step} \cdot \text{warmup\_steps}^{-1.5}) $$
この式は2つのフェーズから構成されます。
- Warmup フェーズ（$\text{step} \leq \text{warmup\_steps}$）: 学習率は $\text{step}$ に比例して線形に増加
- Decay フェーズ（$\text{step} > \text{warmup\_steps}$）: 学習率は $\text{step}^{-0.5}$ に比例して減衰
Warmup が必要な理由は、学習初期にはパラメータがランダムに近い状態であり、大きな学習率で更新すると不安定になりやすいためです。特に Attention の重みが適切に分散していない初期段階で大きな更新を行うと、特定の位置に過度に集中する（崩壊する）リスクがあります。
原論文では $\text{warmup\_steps} = 4000$ を使用しています。
Python で Transformer Encoder のスクラッチ実装
import numpy as np
import matplotlib.pyplot as plt
def softmax(x, axis=-1):
"""数値安定な softmax"""
e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
return e_x / np.sum(e_x, axis=axis, keepdims=True)
def layer_norm(x, gamma, beta, eps=1e-6):
"""Layer Normalization"""
mean = np.mean(x, axis=-1, keepdims=True)
var = np.var(x, axis=-1, keepdims=True)
x_norm = (x - mean) / np.sqrt(var + eps)
return gamma * x_norm + beta
def scaled_dot_product_attention(Q, K, V, mask=None):
"""Scaled Dot-Product Attention"""
d_k = Q.shape[-1]
# Attention スコアの計算
scores = Q @ K.transpose(0, 2, 1) / np.sqrt(d_k) # (batch, T_q, T_k)
if mask is not None:
scores = scores + mask # mask は 0 or -inf
# softmax で重みを計算
attn_weights = softmax(scores, axis=-1)
# バリューの重み付き和
output = attn_weights @ V # (batch, T_q, d_v)
return output, attn_weights
class MultiHeadAttention:
"""マルチヘッド Attention"""
def __init__(self, d_model, n_heads):
self.d_model = d_model
self.n_heads = n_heads
self.d_k = d_model // n_heads
self.d_v = d_model // n_heads
# 重みの初期化（Xavier）
scale = np.sqrt(2.0 / (d_model + self.d_k))
self.W_Q = np.random.randn(n_heads, d_model, self.d_k) * scale
self.W_K = np.random.randn(n_heads, d_model, self.d_k) * scale
self.W_V = np.random.randn(n_heads, d_model, self.d_v) * scale
self.W_O = np.random.randn(n_heads * self.d_v, d_model) * scale
def forward(self, Q, K, V, mask=None):
batch_size, T_q, _ = Q.shape
T_k = K.shape[1]
all_heads = []
all_weights = []
for h in range(self.n_heads):
# 各ヘッドの Q, K, V を射影
Q_h = Q @ self.W_Q[h] # (batch, T_q, d_k)
K_h = K @ self.W_K[h] # (batch, T_k, d_k)
V_h = V @ self.W_V[h] # (batch, T_k, d_v)
# Scaled Dot-Product Attention
head_output, attn_w = scaled_dot_product_attention(Q_h, K_h, V_h, mask)
all_heads.append(head_output)
all_weights.append(attn_w)
# ヘッドを結合
concat = np.concatenate(all_heads, axis=-1) # (batch, T_q, n_heads * d_v)
# 出力射影
output = concat @ self.W_O # (batch, T_q, d_model)
return output, all_weights
class PositionWiseFFN:
"""Position-wise Feed-Forward Network"""
def __init__(self, d_model, d_ff):
scale1 = np.sqrt(2.0 / (d_model + d_ff))
scale2 = np.sqrt(2.0 / (d_ff + d_model))
self.W1 = np.random.randn(d_model, d_ff) * scale1
self.b1 = np.zeros(d_ff)
self.W2 = np.random.randn(d_ff, d_model) * scale2
self.b2 = np.zeros(d_model)
def forward(self, x):
# ReLU 活性化
hidden = np.maximum(0, x @ self.W1 + self.b1)
output = hidden @ self.W2 + self.b2
return output
class TransformerEncoderLayer:
"""Transformer Encoder の1層"""
def __init__(self, d_model, n_heads, d_ff):
self.mha = MultiHeadAttention(d_model, n_heads)
self.ffn = PositionWiseFFN(d_model, d_ff)
# LayerNorm のパラメータ
self.gamma1 = np.ones(d_model)
self.beta1 = np.zeros(d_model)
self.gamma2 = np.ones(d_model)
self.beta2 = np.zeros(d_model)
def forward(self, x, mask=None):
# Self-Attention + 残差接続 + LayerNorm
attn_out, attn_weights = self.mha.forward(x, x, x, mask)
x = layer_norm(x + attn_out, self.gamma1, self.beta1)
# FFN + 残差接続 + LayerNorm
ffn_out = self.ffn.forward(x)
x = layer_norm(x + ffn_out, self.gamma2, self.beta2)
return x, attn_weights
def positional_encoding(max_len, d_model):
"""正弦波位置エンコーディングを生成"""
PE = np.zeros((max_len, d_model))
position = np.arange(max_len)[:, np.newaxis] # (max_len, 1)
div_term = np.exp(np.arange(0, d_model, 2) * (-np.log(10000.0) / d_model))
PE[:, 0::2] = np.sin(position * div_term)
PE[:, 1::2] = np.cos(position * div_term)
return PE
class TransformerEncoder:
"""Transformer Encoder（複数層）"""
def __init__(self, n_layers, d_model, n_heads, d_ff, max_len=512):
self.layers = [
TransformerEncoderLayer(d_model, n_heads, d_ff)
for _ in range(n_layers)
]
self.pe = positional_encoding(max_len, d_model)
def forward(self, x):
T = x.shape[1]
# 位置エンコーディングを加算
x = x + self.pe[:T, :]
all_attn_weights = []
for layer in self.layers:
x, attn_weights = layer.forward(x)
all_attn_weights.append(attn_weights)
return x, all_attn_weights
# デモ: Transformer Encoder の動作確認
np.random.seed(42)
# ハイパーパラメータ
d_model = 64
n_heads = 4
d_ff = 256
n_layers = 2
seq_len = 10
batch_size = 1
# ランダム入力（本来は埋め込み層の出力）
x = np.random.randn(batch_size, seq_len, d_model)
# Transformer Encoder
encoder = TransformerEncoder(n_layers, d_model, n_heads, d_ff)
output, attn_weights_all = encoder.forward(x)
print(f"入力形状: {x.shape}")
print(f"出力形状: {output.shape}")
print(f"Attention 層数: {len(attn_weights_all)}")
print(f"各層のヘッド数: {len(attn_weights_all[0])}")
位置エンコーディングの可視化
import numpy as np
import matplotlib.pyplot as plt
def positional_encoding(max_len, d_model):
"""正弦波位置エンコーディング"""
PE = np.zeros((max_len, d_model))
position = np.arange(max_len)[:, np.newaxis]
div_term = np.exp(np.arange(0, d_model, 2) * (-np.log(10000.0) / d_model))
PE[:, 0::2] = np.sin(position * div_term)
PE[:, 1::2] = np.cos(position * div_term)
return PE
# 位置エンコーディングの生成
max_len = 100
d_model = 64
PE = positional_encoding(max_len, d_model)
# (1) 位置エンコーディングのヒートマップ
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
im = axes[0].imshow(PE, cmap='RdBu', aspect='auto', interpolation='nearest')
axes[0].set_xlabel('Dimension', fontsize=12)
axes[0].set_ylabel('Position', fontsize=12)
axes[0].set_title('Positional Encoding', fontsize=14)
plt.colorbar(im, ax=axes[0])
# (2) 位置エンコーディングの内積（相対位置への依存性）
dot_products = PE @ PE.T
axes[1].imshow(dot_products, cmap='viridis', aspect='auto')
axes[1].set_xlabel('Position j', fontsize=12)
axes[1].set_ylabel('Position i', fontsize=12)
axes[1].set_title('PE(i) · PE(j) - Depends on |i-j|', fontsize=14)
plt.colorbar(im, ax=axes[1])
plt.tight_layout()
plt.show()
Warmup 学習率スケジュールの可視化
import numpy as np
import matplotlib.pyplot as plt
def transformer_lr_schedule(step, d_model=512, warmup_steps=4000):
"""Transformer の学習率スケジュール"""
return d_model ** (-0.5) * min(step ** (-0.5), step * warmup_steps ** (-1.5))
# 学習率の推移を計算
steps = np.arange(1, 50001)
lrs = [transformer_lr_schedule(s) for s in steps]
plt.figure(figsize=(10, 5))
plt.plot(steps, lrs, 'b-', linewidth=2)
plt.axvline(x=4000, color='r', linestyle='--', label='warmup_steps=4000')
plt.xlabel('Training Step', fontsize=12)
plt.ylabel('Learning Rate', fontsize=12)
plt.title('Transformer Learning Rate Schedule (Warmup + Decay)', fontsize=14)
plt.legend(fontsize=12)
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()
Attention 重みの可視化
import numpy as np
import matplotlib.pyplot as plt
np.random.seed(42)
# 簡易的な文を想定（各トークンをランダムベクトルで表現）
tokens = ["The", "cat", "sat", "on", "the", "mat", ".", "<pad>", "<pad>", "<pad>"]
seq_len = len(tokens)
d_model = 64
# ランダムな埋め込み + 位置エンコーディング
def positional_encoding(max_len, d_model):
PE = np.zeros((max_len, d_model))
position = np.arange(max_len)[:, np.newaxis]
div_term = np.exp(np.arange(0, d_model, 2) * (-np.log(10000.0) / d_model))
PE[:, 0::2] = np.sin(position * div_term)
PE[:, 1::2] = np.cos(position * div_term)
return PE
x = np.random.randn(seq_len, d_model) + positional_encoding(seq_len, d_model)
# Self-Attention の計算（1ヘッド）
d_k = 16
W_Q = np.random.randn(d_model, d_k) * 0.1
W_K = np.random.randn(d_model, d_k) * 0.1
Q = x @ W_Q
K = x @ W_K
scores = Q @ K.T / np.sqrt(d_k)
# softmax
def softmax(x, axis=-1):
e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
return e_x / np.sum(e_x, axis=axis, keepdims=True)
attn_weights = softmax(scores)
# 可視化
plt.figure(figsize=(8, 7))
plt.imshow(attn_weights, cmap='Blues', interpolation='nearest')
plt.xticks(range(seq_len), tokens, rotation=45, fontsize=10)
plt.yticks(range(seq_len), tokens, fontsize=10)
plt.xlabel('Key', fontsize=12)
plt.ylabel('Query', fontsize=12)
plt.title('Self-Attention Weights (Single Head)', fontsize=14)
plt.colorbar(label='Attention Weight')
plt.tight_layout()
plt.show()
まとめ
本記事では、Transformer 原論文「Attention Is All You Need」を徹底的に解読しました。
- Scaled Dot-Product Attention において $\sqrt{d_k}$ で除算する理由は、内積の分散が $d_k$ に比例して増大するのを抑え、softmax の飽和を防ぐためである
- マルチヘッド Attention は計算量を増やさずに複数の部分空間で異なるパターンを捉える
- 位置エンコーディングは Self-Attention の置換不変性を補い、正弦波関数は相対位置を線形変換で表現できるという優れた性質を持つ
- 残差接続 + LayerNorm が深いネットワークの学習を安定させる
- Warmup 学習率スケジュール が学習初期の不安定性を防ぐ
次のステップとして、以下の記事も参考にしてください。