7. PyTorchによる最小実装と動作確認

実装で理解するTransformer

数式や概念を最も深く理解する方法は、実際にプログラミング言語で動くコードを書いてみることです。本章では、深層学習フレームワークの標準である PyTorch を使用して、Transformer の核心である「Scaled Dot-Product Attention」「Multi-Head Attention」「Transformer ブロック」をゼロから実装します。

外部の高度なライブラリに頼らず、テンソル（多次元配列）の形状（Shape）の変化を 1 ステップずつ追いながら実装していきます。

テンソルの形状（Shape）の約束事

コードを読む前に、データがどのような形で流れていくのかを整理しておきます。

batch_size : 一度に処理する文章の数（例: 2）
seq_len    : 1つの文章に含まれる単語（トークン）の数（例: 4）
d_model    : 1つの単語を表現するベクトルの次元数（例: 64）
num_heads  : マルチヘッドアテンションのヘッドの数（例: 4）
d_k        : 1つのヘッドあたりの次元数 (d_model // num_heads, 例: 16)

入力テンソルは通常、(batch_size, seq_len, d_model) という 3 次元の形状を持ちます。

Scaled Dot-Product Attentionの実装

まずは、Attention の基本計算を行う関数を実装します。

import math
import torch
import torch.nn as nn
import torch.nn.functional as F


def scaled_dot_product_attention(q, k, v, mask=None):
    """Scaled Dot-Product Attention の計算

    q, k, v の形状: (batch_size, num_heads, seq_len, d_k)
    mask の形状:    (batch_size, 1, seq_len, seq_len) または None
    """
    d_k = q.size(-1)

    # 1. Q と K の内積計算: (batch, heads, seq, d_k) x (batch, heads, d_k, seq)
    #    -> (batch, heads, seq, seq)
    scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)

    # 2. マスクの適用（未来の単語やパディングを遮断）
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)

    # 3. Softmax による確率（アテンション重み）への変換
    attention_weights = F.softmax(scores, dim=-1)

    # 4. Value との行列積
    #    (batch, heads, seq, seq) x (batch, heads, seq, d_k) -> (batch, heads, seq, d_k)
    output = torch.matmul(attention_weights, v)

    return output, attention_weights

k.transpose(-2, -1) は、Key テンソルの最後の 2 つの軸を入れ替える操作（転置）です。内積を計算した後、scores.masked_fill(mask == 0, -1e9) によって、マスク対象の位置のスコアをほぼ $-\infty$ に置き換えています。

Multi-Head Attentionの実装

次に、複数のヘッドに分割して並列計算を行う Multi-Head Attention モジュールをクラスとして定義します。

class MultiHeadAttention(nn.Module):

    def __init__(self, d_model=64, num_heads=4):
        super().__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        # Q, K, V を生成するための線形変換層
        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)

    def forward(self, q, k, v, mask=None):
        batch_size = q.size(0)

        # 1. 各入力を線形射影し、ヘッドごとに分割して形状変更
        # (batch, seq, d_model) -> (batch, seq, num_heads, d_k) -> (batch, num_heads, seq, d_k)
        q = (
            self.w_q(q)
            .view(batch_size, -1, self.num_heads, self.d_k)
            .transpose(1, 2)
        )
        k = (
            self.w_k(k)
            .view(batch_size, -1, self.num_heads, self.d_k)
            .transpose(1, 2)
        )
        v = (
            self.w_v(v)
            .view(batch_size, -1, self.num_heads, self.d_k)
            .transpose(1, 2)
        )

        # 2. Scaled Dot-Product Attention の計算
        out, attn_weights = scaled_dot_product_attention(q, k, v, mask=mask)

        # 3. ヘッドを結合して元の形状へ戻す
        # (batch, num_heads, seq, d_k) -> (batch, seq, num_heads, d_k) -> (batch, seq, d_model)
        out = (
            out.transpose(1, 2)
            .contiguous()
            .view(batch_size, -1, self.d_model)
        )

        # 4. 仕上げの線形変換
        return self.w_o(out)

このクラスのポイントは、4 つの独立した小さな線形層を作るのではなく、1 つの大きな線形層（d_model $\rightarrow$ d_model）で一括計算してから .view() でヘッドごとに切り分けている点です。こうすることで、GPU の計算効率が飛躍的に向上します。

Transformerブロック（エンコーダ層）の組み立て

Multi-Head Attention の後段に、Feed-Forward Network（FFN）と残差接続、LayerNorm を組み合わせて、完全な 1 つの「Transformer ブロック（エンコーダ層）」を作ります。ここでは現代の標準である Pre-LN 方式を採用します。

class TransformerEncoderBlock(nn.Module):

    def __init__(self, d_model=64, num_heads=4, d_ff=256, dropout=0.1):
        super().__init__()
        self.attention = MultiHeadAttention(d_model, num_heads)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

        # Position-wise Feed-Forward Network
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
        )
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # 1. Pre-LN + Self-Attention + 残差接続
        normed_x = self.norm1(x)
        attn_out = self.attention(normed_x, normed_x, normed_x, mask=mask)
        x = x + self.dropout(attn_out)

        # 2. Pre-LN + FFN + 残差接続
        normed_x = self.norm2(x)
        ffn_out = self.ffn(normed_x)
        x = x + self.dropout(ffn_out)

        return x

動作確認スクリプト

作成したブロックにダミーの入力データ（2 文、各文 4 単語、64 次元）を流し、正常に入出力の形状が保たれているかを確認します。

if __name__ == "__main__":
    # パラメータ設定
    batch_size = 2
    seq_len = 4
    d_model = 64

    # ランダムな入力テンソルを作成
    dummy_input = torch.randn(batch_size, seq_len, d_model)
    print(f"入力の形状: {dummy_input.shape}")

    # Transformer ブロックの初期化と実行
    encoder_block = TransformerEncoderBlock(d_model=d_model, num_heads=4)
    output = encoder_block(dummy_input)

    print(f"出力の形状: {output.shape}")
    assert (
        dummy_input.shape == output.shape
    ), "入出力のテンソル形状が一致しません"
    print("動作確認成功: 入力と同じ形状のテンソルが出力されました。")

このコードを実行すると、入力とまったく同一の形状 (2, 4, 64) の文脈化されたテンソルが出力されることが確認できます。このブロックを何層も縦に積み重ねることで、強力な表現力を持つ深層ニューラルネットワークが完成します。

次章では、この Transformer を実サービスや超大規模言語モデルで実用化するために欠かせない「計算効率化と高速化技術」を学びます。