Transformerアーキテクチャの全体像をわかりやすく解説

Transformerは、2017年にGoogleの研究チームがNeurIPS論文 “Attention Is All You Need” (Vaswani et al.) で発表したニューラルネットワークアーキテクチャです。RNNやLSTMに頼らず、Self-Attentionのみで系列間の依存関係を捉えるという革新的なアイデアにより、自然言語処理(NLP)の世界に大きなブレークスルーをもたらしました。

Transformerの登場以降、BERT、GPT、T5、Vision Transformer(ViT)など、ほぼすべての最先端モデルがこのアーキテクチャを基盤としています。Transformerを理解することは、現代の深層学習を学ぶ上で避けて通れないステップです。

本記事の内容

  • Transformerの全体アーキテクチャ(Encoder-Decoder構造)
  • 位置エンコーディングの数式と直感的な意味
  • Encoderブロックの構成(Self-Attention + FFN + 残差接続 + LayerNorm)
  • Decoderブロックの構成(Masked Self-Attention + Cross-Attention + FFN)
  • マスク付きAttentionの仕組み
  • 学習テクニック(ラベル平滑化、warmup付き学習率スケジューリング)
  • BERT vs GPT:Encoder-only vs Decoder-only
  • PyTorchによるTransformerの主要コンポーネントのスクラッチ実装

前提知識

この記事を読む前に、以下の記事を読んでおくと理解が深まります。

Transformerの全体アーキテクチャ

Transformerは大きくEncoderDecoderの2つの部分から構成されます。原論文では、Encoderが6層、Decoderが6層のブロックを積み重ねた構造を採用しています。

全体の処理の流れは以下のとおりです。

  1. 入力トークン列を埋め込みベクトルに変換し、位置エンコーディングを加算する
  2. Encoderが入力系列を処理し、文脈を反映した表現(隠れ状態の列)を出力する
  3. DecoderがEncoderの出力を参照しながら、出力系列を1トークンずつ自己回帰的に生成する
  4. 最終層の出力を線形変換 + Softmaxで語彙上の確率分布に変換する

これを擬似的に書くと、次のようになります。

[入力トークン列]
    ↓
[入力埋め込み + 位置エンコーディング]
    ↓
┌─────────────────────────┐
│ Encoder Block ×N        │
│  ├ Multi-Head Self-Attn │
│  ├ Add & LayerNorm      │
│  ├ Feed-Forward Network │
│  └ Add & LayerNorm      │
└─────────────────────────┘
    ↓ Encoderの出力
    ↓(Cross-Attentionで参照)
┌─────────────────────────────┐
│ Decoder Block ×N            │
│  ├ Masked Multi-Head Self-Attn │
│  ├ Add & LayerNorm             │
│  ├ Multi-Head Cross-Attn       │
│  ├ Add & LayerNorm             │
│  ├ Feed-Forward Network        │
│  └ Add & LayerNorm             │
└─────────────────────────────┘
    ↓
[線形変換 + Softmax]
    ↓
[出力確率分布]

Encoderは入力系列全体を双方向に処理できるのに対し、Decoderは未来のトークンを参照できないよう因果マスクがかかっている点が大きな違いです。

位置エンコーディング

なぜ位置エンコーディングが必要か

Self-Attentionは入力の全要素間の関連度を一度に計算するため、本質的に順序の情報を持ちません。つまり、入力トークンをどのように並べ替えても、Self-Attentionの出力は(重みが同じならば)変わりません。これは集合に対する演算としては美しいですが、自然言語のように語順が意味を決定する問題では致命的です。

例えば「犬が猫を追いかける」と「猫が犬を追いかける」は単語集合としては同じですが、意味はまったく異なります。この語順情報を補うために、入力埋め込みに位置の情報を加算するのが位置エンコーディング(Positional Encoding)です。

sin/cos関数による位置エンコーディング

原論文では、学習パラメータを使わず、以下のsin/cos関数で位置エンコーディングを定義しています。

$$ \text{PE}(pos, 2i) = \sin\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right) $$

$$ \text{PE}(pos, 2i+1) = \cos\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right) $$

ここで、$pos$ はトークンの位置(0始まり)、$i$ は埋め込みベクトルの次元のインデックス($i = 0, 1, \dots, d_{\text{model}}/2 – 1$)、$d_{\text{model}}$ は埋め込みの次元数(原論文では512)です。

各次元 $2i$ にはsin関数を、次元 $2i+1$ にはcos関数を割り当てます。次元ごとに異なる周波数の波を使うことで、各位置に一意のパターンが割り当てられます。

数式の直感的な意味

この式を直感的に理解するために、指数部分を分解してみましょう。

$$ \frac{pos}{10000^{2i/d_{\text{model}}}} = pos \cdot 10000^{-2i/d_{\text{model}}} $$

$10000^{-2i/d_{\text{model}}}$ は次元 $i$ が大きくなるほど小さくなります。つまり、

  • 低次元($i$ が小さい):周波数が高く、近い位置間の差を捉える
  • 高次元($i$ が大きい):周波数が低く、遠い位置間の差を捉える

これは二進数の各桁が異なる周期で0と1を繰り返すのと似たアイデアです。最下位ビットは1ごとに反転し、次のビットは2ごとに反転し、…というように、各次元が異なるスケールの位置情報をエンコードします。

また、sin/cosを使う理由として、相対位置を線形変換で表現できるという性質があります。三角関数の加法定理により、

$$ \begin{align} \sin(pos + k) &= \sin(pos)\cos(k) + \cos(pos)\sin(k) \\ \cos(pos + k) &= \cos(pos)\cos(k) – \sin(pos)\sin(k) \end{align} $$

が成り立つため、位置 $pos$ のエンコーディングから位置 $pos + k$ のエンコーディングへの変換は、$k$ のみに依存する固定の線形変換(回転行列)で表現できます。これにより、モデルが相対位置の関係を学習しやすくなります。

位置エンコーディングは入力埋め込みに加算されます。

$$ \bm{z}_{\text{input}} = \text{Embedding}(\text{token}) + \text{PE}(pos) $$

Encoderブロック

Encoderは同一構造のブロックを $N$ 層(原論文では $N=6$)積み重ねたものです。各Encoderブロックは以下の2つのサブレイヤーで構成されます。

サブレイヤー1:Multi-Head Self-Attention

入力系列の各トークンが、同じ系列内の全トークンとの関連度を計算し、文脈に応じた表現を得ます。マルチヘッドAttentionの詳細はSelf-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) $$

Encoderでは $\bm{Q} = \bm{K} = \bm{V}$ (すべて同じ入力)です。

サブレイヤー2:Position-wise Feed-Forward Network(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_{\text{model}} = 512$、$d_{ff} = 2048$ と設定されています。つまり、内部で一度4倍に拡張してから元の次元に圧縮します。

「Position-wise」とは、系列の各位置に同じ重みの全結合層を独立に適用するという意味です。位置間の情報交換はSelf-Attentionが担当し、FFNは各位置の表現を非線形変換で豊かにする役割を担います。

残差接続とLayer Normalization

各サブレイヤーには残差接続(Residual Connection)とLayer Normalizationが適用されます。

$$ \text{output} = \text{LayerNorm}(\bm{x} + \text{SubLayer}(\bm{x})) $$

残差接続は入力 $\bm{x}$ をサブレイヤーの出力に直接加算する仕組みです。これにより勾配が直接的に下層に流れるため、深いネットワークでも安定して学習できます。ResNetで有名になった手法で、Transformerでも不可欠な要素です。

Layer Normalizationは、各サンプルの各位置ごとに、特徴量次元方向で正規化を行います。入力ベクトル $\bm{x} = (x_1, x_2, \dots, x_d)$ に対して、

$$ \mu = \frac{1}{d}\sum_{j=1}^{d} x_j, \quad \sigma^2 = \frac{1}{d}\sum_{j=1}^{d}(x_j – \mu)^2 $$

$$ \text{LayerNorm}(\bm{x}) = \frac{\bm{x} – \mu}{\sqrt{\sigma^2 + \epsilon}} \odot \bm{\gamma} + \bm{\beta} $$

ここで $\bm{\gamma}$, $\bm{\beta}$ は学習可能なパラメータ、$\epsilon$ はゼロ除算を防ぐ微小値、$\odot$ は要素ごとの積です。Batch Normalizationとは異なり、バッチ方向ではなく特徴量方向で正規化するため、系列長やバッチサイズに依存しない安定した正規化が実現できます。

したがって、Encoderブロック全体の処理は次のように書けます。

$$ \bm{h}_1 = \text{LayerNorm}(\bm{x} + \text{MultiHeadSelfAttn}(\bm{x})) $$

$$ \bm{h}_2 = \text{LayerNorm}(\bm{h}_1 + \text{FFN}(\bm{h}_1)) $$

Decoderブロック

Decoderも同一構造のブロックを $N$ 層積み重ねます。Encoderとの違いは、3つのサブレイヤーを持つことです。

サブレイヤー1:Masked Multi-Head Self-Attention

Decoderの自己注意では、因果マスク(Causal Mask)を適用します。これにより、位置 $t$ のトークンは位置 $1, 2, \dots, t$ の情報のみを参照でき、未来の位置 $t+1, t+2, \dots$ の情報は遮断されます。マスクの詳細は次節で解説します。

サブレイヤー2:Multi-Head Cross-Attention(Encoder-Decoder Attention)

Cross-Attentionでは、Decoderの各位置がEncoderの出力全体を参照します。

$$ \text{CrossAttn}(\bm{Q}, \bm{K}, \bm{V}) = \text{MultiHead}(\bm{Q}_{\text{dec}}, \bm{K}_{\text{enc}}, \bm{V}_{\text{enc}}) $$

ここで、QueryはDecoderの前のサブレイヤーの出力から、Key・ValueはEncoderの最終出力から取られます。これにより、Decoderは「出力を生成する際に入力系列のどこに注目すべきか」を学習できます。

例えば機械翻訳において、英語の “I love cats” を日本語に翻訳する場合、「猫」を出力するときに “cats” に強く注目し、「好き」を出力するときに “love” に強く注目する、といった対応関係をCross-Attentionが学習します。

サブレイヤー3:Position-wise FFN

Encoderと同じ構造のFFNです。

各サブレイヤーに残差接続とLayerNormが適用される点もEncoderと同じです。Decoderブロック全体をまとめると、

$$ \bm{d}_1 = \text{LayerNorm}(\bm{y} + \text{MaskedMultiHeadSelfAttn}(\bm{y})) $$

$$ \bm{d}_2 = \text{LayerNorm}(\bm{d}_1 + \text{MultiHeadCrossAttn}(\bm{d}_1, \bm{m}_{\text{enc}}, \bm{m}_{\text{enc}})) $$

$$ \bm{d}_3 = \text{LayerNorm}(\bm{d}_2 + \text{FFN}(\bm{d}_2)) $$

ここで $\bm{y}$ はDecoderへの入力(出力埋め込み + 位置エンコーディング)、$\bm{m}_{\text{enc}}$ はEncoderの最終出力です。

マスク付きAttention(因果マスク)

Decoderでは、自己回帰的な生成を行うため、各位置が未来の情報を参照してはいけません。これを実現するのが因果マスク(Causal Mask)です。

Scaled Dot-Product Attentionのスコア計算において、マスクを次のように適用します。

$$ \text{Attention}(\bm{Q}, \bm{K}, \bm{V}) = \text{softmax}\left(\frac{\bm{Q}\bm{K}^T}{\sqrt{d_k}} + \bm{M}\right)\bm{V} $$

ここでマスク行列 $\bm{M}$ は上三角部分が $-\infty$、それ以外が $0$ の行列です。

$$ M_{ij} = \begin{cases} 0 & \text{if } j \le i \\ -\infty & \text{if } j > i \end{cases} $$

具体的に、系列長4の場合のマスクは次のようになります。

$$ \bm{M} = \begin{pmatrix} 0 & -\infty & -\infty & -\infty \\ 0 & 0 & -\infty & -\infty \\ 0 & 0 & 0 & -\infty \\ 0 & 0 & 0 & 0 \end{pmatrix} $$

Softmaxの入力に $-\infty$ を加えると、その位置のSoftmax出力は $0$ になります。

$$ \text{softmax}(-\infty) = \frac{e^{-\infty}}{\sum_j e^{s_j}} = \frac{0}{\sum_j e^{s_j}} = 0 $$

これにより、位置 $i$ のトークンは位置 $j > i$(自分より未来の位置)の情報をまったく参照しなくなります。学習時には、正解系列全体を一度にDecoderに入力し、因果マスクによって未来の情報を遮断することで、各位置で独立に次トークン予測を行えます。これにより、RNNのように逐次的に処理する必要がなく、学習を完全に並列化できます。

学習テクニック

ラベル平滑化(Label Smoothing)

通常の分類では、正解クラスの確率を1、それ以外を0とするone-hotラベルを使います。ラベル平滑化では、正解クラスの確率を $1 – \epsilon$ に下げ、残りの確率 $\epsilon$ を他のクラスに均等に配分します。

$$ y_{\text{smooth}}(k) = \begin{cases} 1 – \epsilon + \frac{\epsilon}{K} & \text{if } k = k^* \quad (\text{正解クラス}) \\ \frac{\epsilon}{K} & \text{otherwise} \end{cases} $$

ここで $K$ は語彙サイズ、$\epsilon$ は平滑化パラメータ(原論文では $\epsilon = 0.1$)、$k^*$ は正解クラスのインデックスです。

ラベル平滑化は、モデルが正解に対して過度に自信を持つ(Softmax出力が極端に1に近づく)ことを防ぎ、汎化性能を向上させます。

学習率スケジューリング(Warmup)

原論文では、以下の学習率スケジュールを採用しています。

$$ lr = d_{\text{model}}^{-0.5} \cdot \min\left(\text{step}^{-0.5}, \quad \text{step} \cdot \text{warmup\_steps}^{-1.5}\right) $$

この式は2つの項の最小値を取ります。

  • $\text{step} \cdot \text{warmup\_steps}^{-1.5}$:学習初期に学習率を線形に増加させる
  • $\text{step}^{-0.5}$:ウォームアップ後に学習率を逆平方根で減衰させる

原論文では $\text{warmup\_steps} = 4000$ が使われています。学習初期はパラメータがランダムに初期化されているため、大きな学習率で更新すると不安定になりやすいです。ウォームアップにより、最初は小さな学習率で慎重にパラメータを調整し、モデルが安定してきたら学習率を上げて効率的に学習を進め、その後は徐々に学習率を下げて微調整を行います。

BERT vs GPT:Encoder-only vs Decoder-only

原論文のTransformerはEncoder-Decoder構造ですが、その後の発展でEncoderのみまたはDecoderのみを使うアーキテクチャが主流になっています。

BERT(Encoder-only)

BERT(Bidirectional Encoder Representations from Transformers)は、TransformerのEncoder部分のみを使います。

  • 双方向:因果マスクがないため、各トークンが系列全体(左右両方)を参照できる
  • 事前学習タスク:Masked Language Model(MLM)— 入力の一部をマスクし、そのトークンを予測する
  • 適用先:文分類、固有表現抽出、質問応答など、入力全体を理解する判別タスクに強い

BERTのMLMでは、入力トークン列の約15%をランダムにマスクし、そのマスクされたトークンを周囲の文脈から予測します。

$$ \mathcal{L}_{\text{MLM}} = -\sum_{t \in \mathcal{M}} \log P(x_t \mid \bm{x}_{\setminus\mathcal{M}}) $$

ここで $\mathcal{M}$ はマスクされた位置の集合、$\bm{x}_{\setminus\mathcal{M}}$ はマスクされていないトークン列です。

GPT(Decoder-only)

GPT(Generative Pre-trained Transformer)は、TransformerのDecoder部分のみを使います(ただしCross-Attention層は除去)。

  • 単方向(左から右):因果マスクにより、各トークンは自分より左のトークンのみを参照する
  • 事前学習タスク:次トークン予測(Causal Language Modeling)— 文脈から次のトークンを予測する
  • 適用先:文章生成、対話、コード生成など、逐次的に出力する生成タスクに強い

GPTの学習目標は、

$$ \mathcal{L}_{\text{CLM}} = -\sum_{t=1}^{T} \log P(x_t \mid x_1, x_2, \dots, x_{t-1}) $$

です。各位置 $t$ で、それまでの文脈 $x_1, \dots, x_{t-1}$ から $x_t$ を予測します。

比較まとめ

項目 BERT (Encoder-only) GPT (Decoder-only) 原論文Transformer
注意の方向 双方向 単方向(左→右) Enc:双方向, Dec:単方向
マスク なし 因果マスク Decのみ因果マスク
事前学習 MLM 次トークン予測 教師あり(翻訳等)
得意なタスク 判別(分類、抽出) 生成(文章、コード) 系列変換(翻訳)

PyTorchによる実装

Transformerの主要コンポーネントをPyTorchでスクラッチ実装してみましょう。

位置エンコーディング

import torch
import torch.nn as nn
import math

class PositionalEncoding(nn.Module):
    """sin/cos関数による位置エンコーディング"""
    def __init__(self, d_model, max_len=5000, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)

        # (max_len, d_model) の位置エンコーディング行列を事前計算
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)  # (max_len, 1)
        # 10000^{-2i/d_model} を計算
        div_term = torch.exp(
            torch.arange(0, d_model, 2, dtype=torch.float) * (-math.log(10000.0) / d_model)
        )  # (d_model/2,)

        pe[:, 0::2] = torch.sin(position * div_term)  # 偶数次元: sin
        pe[:, 1::2] = torch.cos(position * div_term)  # 奇数次元: cos
        pe = pe.unsqueeze(0)  # (1, max_len, d_model)

        # 学習対象ではないのでbufferとして登録
        self.register_buffer('pe', pe)

    def forward(self, x):
        """x: (batch_size, seq_len, d_model)"""
        x = x + self.pe[:, :x.size(1), :]
        return self.dropout(x)

Encoderブロック

class TransformerEncoderBlock(nn.Module):
    """Transformer Encoderブロック(1層分)"""
    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Linear(d_ff, d_model),
        )
        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, src_mask=None, src_key_padding_mask=None):
        """
        x: (batch_size, seq_len, d_model)
        """
        # サブレイヤー1: Multi-Head Self-Attention + 残差接続 + LayerNorm
        attn_out, _ = self.self_attn(x, x, x, attn_mask=src_mask,
                                      key_padding_mask=src_key_padding_mask)
        x = self.norm1(x + self.dropout1(attn_out))

        # サブレイヤー2: FFN + 残差接続 + LayerNorm
        ffn_out = self.ffn(x)
        x = self.norm2(x + self.dropout2(ffn_out))
        return x

Decoderブロック

class TransformerDecoderBlock(nn.Module):
    """Transformer Decoderブロック(1層分)"""
    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
        super().__init__()
        self.masked_self_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True)
        self.cross_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Linear(d_ff, d_model),
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)

    def forward(self, x, enc_output, tgt_mask=None, tgt_key_padding_mask=None,
                memory_key_padding_mask=None):
        """
        x: (batch_size, tgt_len, d_model) - Decoderの入力
        enc_output: (batch_size, src_len, d_model) - Encoderの出力
        tgt_mask: 因果マスク
        """
        # サブレイヤー1: Masked Multi-Head Self-Attention
        attn_out, _ = self.masked_self_attn(x, x, x, attn_mask=tgt_mask,
                                             key_padding_mask=tgt_key_padding_mask)
        x = self.norm1(x + self.dropout1(attn_out))

        # サブレイヤー2: Multi-Head Cross-Attention(Encoder出力を参照)
        cross_out, _ = self.cross_attn(x, enc_output, enc_output,
                                        key_padding_mask=memory_key_padding_mask)
        x = self.norm2(x + self.dropout2(cross_out))

        # サブレイヤー3: FFN
        ffn_out = self.ffn(x)
        x = self.norm3(x + self.dropout3(ffn_out))
        return x

Transformer全体

class Transformer(nn.Module):
    """Transformerモデル全体"""
    def __init__(self, src_vocab_size, tgt_vocab_size, d_model=512,
                 n_heads=8, d_ff=2048, n_layers=6, dropout=0.1, max_len=5000):
        super().__init__()

        # 埋め込み層
        self.src_embedding = nn.Embedding(src_vocab_size, d_model)
        self.tgt_embedding = nn.Embedding(tgt_vocab_size, d_model)
        self.positional_encoding = PositionalEncoding(d_model, max_len, dropout)

        # 埋め込みのスケーリング(原論文に準拠)
        self.d_model = d_model

        # Encoder・Decoderブロック
        self.encoder_layers = nn.ModuleList([
            TransformerEncoderBlock(d_model, n_heads, d_ff, dropout)
            for _ in range(n_layers)
        ])
        self.decoder_layers = nn.ModuleList([
            TransformerDecoderBlock(d_model, n_heads, d_ff, dropout)
            for _ in range(n_layers)
        ])

        # 出力層
        self.output_projection = nn.Linear(d_model, tgt_vocab_size)

    def generate_causal_mask(self, size):
        """因果マスクを生成"""
        mask = torch.triu(torch.ones(size, size), diagonal=1).bool()
        return mask.float().masked_fill(mask, float('-inf'))

    def encode(self, src, src_mask=None, src_key_padding_mask=None):
        """Encoder処理"""
        x = self.src_embedding(src) * math.sqrt(self.d_model)
        x = self.positional_encoding(x)
        for layer in self.encoder_layers:
            x = layer(x, src_mask, src_key_padding_mask)
        return x

    def decode(self, tgt, enc_output, tgt_mask=None, tgt_key_padding_mask=None,
               memory_key_padding_mask=None):
        """Decoder処理"""
        x = self.tgt_embedding(tgt) * math.sqrt(self.d_model)
        x = self.positional_encoding(x)
        for layer in self.decoder_layers:
            x = layer(x, enc_output, tgt_mask, tgt_key_padding_mask,
                      memory_key_padding_mask)
        return x

    def forward(self, src, tgt, src_mask=None, tgt_mask=None,
                src_key_padding_mask=None, tgt_key_padding_mask=None,
                memory_key_padding_mask=None):
        """
        src: (batch_size, src_len) - 入力トークンID列
        tgt: (batch_size, tgt_len) - 出力トークンID列
        """
        # 因果マスクの自動生成
        if tgt_mask is None:
            tgt_mask = self.generate_causal_mask(tgt.size(1)).to(tgt.device)

        # Encoder
        enc_output = self.encode(src, src_mask, src_key_padding_mask)

        # Decoder
        dec_output = self.decode(tgt, enc_output, tgt_mask, tgt_key_padding_mask,
                                  memory_key_padding_mask)

        # 出力を語彙サイズの確率分布に変換
        logits = self.output_projection(dec_output)
        return logits

動作確認

# ハイパーパラメータ
src_vocab_size = 1000
tgt_vocab_size = 1000
d_model = 512
n_heads = 8
d_ff = 2048
n_layers = 6

# モデル作成
model = Transformer(src_vocab_size, tgt_vocab_size, d_model, n_heads, d_ff, n_layers)

# ダミー入力
batch_size = 2
src_len = 10
tgt_len = 8
src = torch.randint(0, src_vocab_size, (batch_size, src_len))
tgt = torch.randint(0, tgt_vocab_size, (batch_size, tgt_len))

# 順伝播
logits = model(src, tgt)
print(f"入力 (src): {src.shape}")   # (2, 10)
print(f"入力 (tgt): {tgt.shape}")   # (2, 8)
print(f"出力 (logits): {logits.shape}")  # (2, 8, 1000)

# パラメータ数を確認
total_params = sum(p.numel() for p in model.parameters())
print(f"総パラメータ数: {total_params:,}")

上記を実行すると、入力のソース系列 (2, 10) とターゲット系列 (2, 8) に対して、各位置で語彙サイズ1000の確率分布(logits)が出力される (2, 8, 1000) 形状のテンソルが得られます。

学習率スケジューラの実装

import matplotlib.pyplot as plt

class TransformerLRScheduler:
    """原論文のwarmup付き学習率スケジューラ"""
    def __init__(self, optimizer, d_model, warmup_steps=4000):
        self.optimizer = optimizer
        self.d_model = d_model
        self.warmup_steps = warmup_steps
        self.step_num = 0

    def step(self):
        self.step_num += 1
        lr = self.d_model ** (-0.5) * min(
            self.step_num ** (-0.5),
            self.step_num * self.warmup_steps ** (-1.5)
        )
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = lr
        return lr

# 学習率スケジュールの可視化
steps = list(range(1, 50001))
lrs = []
d_model = 512
warmup_steps = 4000
for s in steps:
    lr = d_model ** (-0.5) * min(s ** (-0.5), s * warmup_steps ** (-1.5))
    lrs.append(lr)

plt.figure(figsize=(10, 5))
plt.plot(steps, lrs, linewidth=2)
plt.axvline(x=warmup_steps, color='r', linestyle='--', label=f'warmup_steps={warmup_steps}')
plt.xlabel('Training Step')
plt.ylabel('Learning Rate')
plt.title('Transformer Learning Rate Schedule (Warmup)')
plt.legend()
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

ウォームアップ期間(4000ステップ)までは学習率が線形に増加し、その後は逆平方根に従って徐々に減衰していく曲線が確認できます。

まとめ

本記事では、Transformerアーキテクチャの全体像を解説しました。

  • Encoder-Decoder構造:Encoderが入力系列を双方向に処理し、Decoderが因果マスク付きで出力系列を自己回帰的に生成する
  • 位置エンコーディング:sin/cos関数により、Self-Attentionに欠けている順序情報を補う。各次元が異なる周波数を持ち、相対位置を線形変換で表現できる
  • 残差接続とLayer Normalization:深いネットワークの安定した学習を実現する
  • 因果マスク:上三角部分を $-\infty$ にすることで、Decoderが未来の情報を参照することを防ぐ
  • 学習テクニック:ラベル平滑化で過信を防ぎ、warmup付きスケジューリングで学習を安定させる
  • BERT vs GPT:Encoder-onlyは双方向の判別タスク、Decoder-onlyは単方向の生成タスクに適している

Transformerは現代の深層学習において最も重要なアーキテクチャの1つです。この全体像を理解した上で、各コンポーネントの詳細やBERT・GPTなどの発展モデルを学んでいくとよいでしょう。

次のステップとして、以下の記事も参考にしてください。