Transformer Decoderは、系列を自己回帰的に生成するためのアーキテクチャです。GPTシリーズをはじめとする大規模言語モデル（LLM）の基盤となっています。
Decoderの特徴は、Masked Self-Attentionによって未来の情報を遮断しながら、過去の文脈のみから次のトークンを予測する点にあります。
本記事の内容
- Decoderの役割と自己回帰生成の仕組み
- Masked Self-Attentionの数式と必要性
- Cross-Attention（Encoder-Decoder Attention）の仕組み
- Decoder層の構成要素
- 自己回帰生成のプロセス
- PyTorchでの完全実装
前提知識
この記事を読む前に、以下の記事を読んでおくと理解が深まります。
Decoderの役割 — 自己回帰生成
自己回帰とは、過去に生成した出力を入力として、次の出力を順次生成していく方式です。
$$ P(\bm{y}) = \prod_{t=1}^{T} P(y_t \mid y_1, y_2, \dots, y_{t-1}) $$
これを連鎖律（chain rule）による分解と呼びます。
Masked Self-Attentionの仕組み
なぜマスクが必要か
Decoderでは、各位置 $t$ のトークンは位置 $t+1, t+2, \dots$ の情報を参照してはいけません。推論時には未来のトークンはまだ生成されていないからです。
因果マスクの数式
$$ \text{MaskedAttention}(\bm{Q}, \bm{K}, \bm{V}) = \text{softmax}\left(\frac{\bm{Q}\bm{K}^\top}{\sqrt{d_k}} + \bm{M}\right)\bm{V} $$
マスク行列 $\bm{M}$ は：
$$ M_{ij} = \begin{cases} 0 & \text{if } j \le i \\ -\infty & \text{if } j > i \end{cases} $$
$-\infty$ の位置はsoftmax後に確率0になります。
Cross-Attention
Encoder-Decoder構造において、DecoderはEncoderの出力を参照します。
| 成分 | 由来 | 役割 |
|---|---|---|
| Query ($\bm{Q}$) | Decoderの前層出力 | 「何を知りたいか」 |
| Key ($\bm{K}$) | Encoderの最終出力 | 「どこに情報があるか」 |
| Value ($\bm{V}$) | Encoderの最終出力 | 「具体的な情報内容」 |
Decoder層の構成要素
各Decoderブロックは3つのサブレイヤーで構成されます：
- Masked Self-Attention
- Cross-Attention（Encoder-Decoder Attention）
- Feed-Forward Network
PyTorchでの完全実装
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
def generate_causal_mask(size):
"""因果マスクを生成する"""
mask = torch.triu(torch.ones(size, size), diagonal=1)
mask = mask.masked_fill(mask == 1, float('-inf'))
return mask
class MaskedMultiHeadSelfAttention(nn.Module):
def __init__(self, d_model, num_heads, dropout=0.1):
super().__init__()
assert d_model % num_heads == 0
self.d_model = d_model
self.num_heads = num_heads
self.d_k = d_model // num_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.dropout = nn.Dropout(dropout)
def forward(self, x, mask=None):
batch_size, seq_len, _ = x.size()
Q = self.W_q(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
K = self.W_k(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
V = self.W_v(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
if mask is not None:
scores = scores + mask
attn_weights = F.softmax(scores, dim=-1)
attn_weights = self.dropout(attn_weights)
context = torch.matmul(attn_weights, V)
context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
output = self.W_o(context)
return output, attn_weights
class MultiHeadCrossAttention(nn.Module):
def __init__(self, d_model, num_heads, dropout=0.1):
super().__init__()
assert d_model % num_heads == 0
self.d_model = d_model
self.num_heads = num_heads
self.d_k = d_model // num_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.dropout = nn.Dropout(dropout)
def forward(self, dec_hidden, enc_output, enc_mask=None):
batch_size, tgt_len, _ = dec_hidden.size()
src_len = enc_output.size(1)
Q = self.W_q(dec_hidden).view(batch_size, tgt_len, self.num_heads, self.d_k).transpose(1, 2)
K = self.W_k(enc_output).view(batch_size, src_len, self.num_heads, self.d_k).transpose(1, 2)
V = self.W_v(enc_output).view(batch_size, src_len, self.num_heads, self.d_k).transpose(1, 2)
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
if enc_mask is not None:
scores = scores.masked_fill(enc_mask == 0, float('-inf'))
attn_weights = F.softmax(scores, dim=-1)
attn_weights = self.dropout(attn_weights)
context = torch.matmul(attn_weights, V)
context = context.transpose(1, 2).contiguous().view(batch_size, tgt_len, self.d_model)
output = self.W_o(context)
return output, attn_weights
class TransformerDecoderBlock(nn.Module):
def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
super().__init__()
self.masked_self_attn = MaskedMultiHeadSelfAttention(d_model, num_heads, dropout)
self.norm1 = nn.LayerNorm(d_model)
self.cross_attn = MultiHeadCrossAttention(d_model, num_heads, dropout)
self.norm2 = nn.LayerNorm(d_model)
self.ffn = nn.Sequential(
nn.Linear(d_model, d_ff), nn.ReLU(), nn.Dropout(dropout),
nn.Linear(d_ff, d_model), nn.Dropout(dropout)
)
self.norm3 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, x, enc_output, tgt_mask=None, enc_mask=None):
attn_out, _ = self.masked_self_attn(x, tgt_mask)
x = self.norm1(x + self.dropout(attn_out))
cross_out, _ = self.cross_attn(x, enc_output, enc_mask)
x = self.norm2(x + self.dropout(cross_out))
ffn_out = self.ffn(x)
x = self.norm3(x + ffn_out)
return x
class TransformerDecoder(nn.Module):
def __init__(self, vocab_size, d_model=512, num_heads=8, d_ff=2048, num_layers=6, dropout=0.1, max_len=5000):
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.layers = nn.ModuleList([
TransformerDecoderBlock(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)
])
self.output_projection = nn.Linear(d_model, vocab_size)
def forward(self, tgt, enc_output, tgt_mask=None, enc_mask=None):
x = self.embedding(tgt) * math.sqrt(self.d_model)
x = self.pos_encoding(x)
if tgt_mask is None:
tgt_mask = generate_causal_mask(tgt.size(1)).to(tgt.device)
for layer in self.layers:
x = layer(x, enc_output, tgt_mask, enc_mask)
logits = self.output_projection(x)
return logits
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)
まとめ
本記事では、Transformer Decoderの構造とMasked Self-Attentionの仕組みについて解説しました。
- 自己回帰生成: Decoderは連鎖律に基づき、過去のトークンを条件として次のトークンの確率分布を予測する
- Masked Self-Attention: 因果マスクにより未来の情報を遮断し、学習時でも並列計算を可能にする
- Cross-Attention: Encoder-Decoder構造において、DecoderがEncoderの出力を参照するための機構
- Decoder層の構成: Masked Self-Attention、Cross-Attention、FFNの3つのサブレイヤー
次のステップとして、以下の記事も参考にしてください。