Transformer Encoderは、入力系列を双方向に処理し、文脈を反映した表現を生成するニューラルネットワークです。BERTをはじめとする多くの自然言語処理モデルの基盤となっています。
本記事では、Transformer Encoderの各構成要素を数式とともに丁寧に解説し、PyTorchでゼロから実装します。
本記事の内容
- Encoderの役割と全体構造
- Encoder層の構成要素（Self-Attention, FFN, Add&Norm）
- 残差接続（Residual Connection）の数学的意味
- 複数層のスタッキングによる表現力の向上
- 入力の処理（埋め込み + 位置エンコーディング）
- PyTorchでの完全実装
前提知識
この記事を読む前に、以下の記事を読んでおくと理解が深まります。
Encoderの役割
Transformer Encoderの役割は、入力系列を受け取り、各トークンの文脈を考慮した表現（contextualized representation）を出力することです。
Encoderの特徴： – 双方向処理: 各トークンが系列全体（左右両方）を参照できる – 並列計算: RNNと異なり、全トークンを同時に処理できる – スタック構造: 同一構造の層を複数積み重ねて表現力を高める
Encoder層の全体構造
各Encoder層は、以下の2つのサブレイヤーで構成されます。
- Multi-Head Self-Attention
- Position-wise Feed-Forward Network (FFN)
各サブレイヤーには残差接続とLayer Normalizationが適用されます。
入力 x
│
├───────────────┐
▼ │
Multi-Head │
Self-Attention │
│ │
├───────◀───────┘ (残差接続)
▼
Layer Norm
│
├───────────────┐
▼ │
Feed-Forward │
Network │
│ │
├───────◀───────┘ (残差接続)
▼
Layer Norm
│
▼
出力
数式で表すと：
$$ \bm{H}_1 = \text{LayerNorm}(\bm{X} + \text{MultiHeadSelfAttn}(\bm{X})) $$
$$ \bm{H}_2 = \text{LayerNorm}(\bm{H}_1 + \text{FFN}(\bm{H}_1)) $$
残差接続（Residual Connection）
残差接続がある場合、勾配の伝播経路を考えると：
$$ \frac{\partial \bm{y}}{\partial \bm{x}} = \bm{I} + \frac{\partial F(\bm{x})}{\partial \bm{x}} $$
単位行列 $\bm{I}$ の項があるため、勾配は少なくとも1の大きさを保ちます。これにより、深い層にも勾配が直接伝わり、学習が安定します。
複数層のスタッキング
Transformer Encoderは、同一構造のEncoder層を $N$ 層積み重ねます。原論文では $N = 6$ です。
$$ \bm{H}^{(l)} = \text{EncoderLayer}^{(l)}(\bm{H}^{(l-1)}) \quad \text{for } l = 1, \dots, N $$
PyTorchでの完全実装
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
class ScaledDotProductAttention(nn.Module):
def __init__(self, dropout=0.1):
super().__init__()
self.dropout = nn.Dropout(dropout)
def forward(self, Q, K, V, mask=None):
d_k = Q.size(-1)
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
if mask is not None:
scores = scores.masked_fill(mask == 0, float('-inf'))
attn_weights = F.softmax(scores, dim=-1)
attn_weights = self.dropout(attn_weights)
output = torch.matmul(attn_weights, V)
return output, attn_weights
class MultiHeadAttention(nn.Module):
def __init__(self, d_model, n_heads, dropout=0.1):
super().__init__()
assert d_model % n_heads == 0
self.d_model = d_model
self.n_heads = n_heads
self.d_k = d_model // n_heads
self.W_q = nn.Linear(d_model, d_model)
self.W_k = nn.Linear(d_model, d_model)
self.W_v = nn.Linear(d_model, d_model)
self.W_o = nn.Linear(d_model, d_model)
self.attention = ScaledDotProductAttention(dropout)
def forward(self, x, mask=None):
batch_size, seq_len, _ = x.size()
Q = self.W_q(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
K = self.W_k(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
V = self.W_v(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
context, attn_weights = self.attention(Q, K, V, mask)
context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
output = self.W_o(context)
return output, attn_weights
class PositionwiseFeedForward(nn.Module):
def __init__(self, d_model, d_ff, dropout=0.1):
super().__init__()
self.linear1 = nn.Linear(d_model, d_ff)
self.linear2 = nn.Linear(d_ff, d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, x):
x = F.relu(self.linear1(x))
x = self.dropout(x)
x = self.linear2(x)
return x
class PositionalEncoding(nn.Module):
def __init__(self, d_model, max_len=5000, dropout=0.1):
super().__init__()
self.dropout = nn.Dropout(dropout)
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
div_term = torch.exp(torch.arange(0, d_model, 2, dtype=torch.float) * (-math.log(10000.0) / d_model))
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
pe = pe.unsqueeze(0)
self.register_buffer('pe', pe)
def forward(self, x):
x = x + self.pe[:, :x.size(1), :]
return self.dropout(x)
class TransformerEncoderLayer(nn.Module):
def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
super().__init__()
self.self_attn = MultiHeadAttention(d_model, n_heads, dropout)
self.ffn = PositionwiseFeedForward(d_model, d_ff, dropout)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.dropout1 = nn.Dropout(dropout)
self.dropout2 = nn.Dropout(dropout)
def forward(self, x, mask=None):
attn_out, attn_weights = self.self_attn(x, mask)
x = self.norm1(x + self.dropout1(attn_out))
ffn_out = self.ffn(x)
x = self.norm2(x + self.dropout2(ffn_out))
return x, attn_weights
class TransformerEncoder(nn.Module):
def __init__(self, vocab_size, d_model=512, n_heads=8, d_ff=2048, n_layers=6, max_len=5000, dropout=0.1):
super().__init__()
self.d_model = d_model
self.embedding = nn.Embedding(vocab_size, d_model)
self.pos_encoding = PositionalEncoding(d_model, max_len, dropout)
self.encoder_layers = nn.ModuleList([
TransformerEncoderLayer(d_model, n_heads, d_ff, dropout) for _ in range(n_layers)
])
self.scale = math.sqrt(d_model)
def forward(self, src, mask=None):
x = self.embedding(src) * self.scale
x = self.pos_encoding(x)
all_attn_weights = []
for layer in self.encoder_layers:
x, attn_weights = layer(x, mask)
all_attn_weights.append(attn_weights)
return x, all_attn_weights
# 動作確認
vocab_size = 10000
encoder = TransformerEncoder(vocab_size=vocab_size)
src = torch.randint(0, vocab_size, (2, 20))
output, attn_weights_list = encoder(src)
print(f"入力形状: {src.shape}")
print(f"出力形状: {output.shape}")
まとめ
本記事では、Transformer Encoderの構造と実装について解説しました。
- Encoderの役割: 入力系列を双方向に処理し、各トークンの文脈を反映した表現を生成する
- Encoder層の構成: Multi-Head Self-Attention + FFN の2つのサブレイヤーで構成される
- 残差接続: 勾配消失を防ぎ、深いネットワークの学習を安定させる
- Layer Normalization: 特徴量次元方向で正規化を行い、学習を安定させる
- 複数層のスタッキング: 同一構造の層を積み重ねることで、階層的な特徴抽出と表現力の向上を実現する
次のステップとして、以下の記事も参考にしてください。