🦜

Transformer に触れてみる (3) — GPT-2 もどき

に公開
Python
機械学習
transformer
GPT
tech

目的

OpenAI のとても有名な研究に Improving Language Understanding
by Generative Pre-Training と Language Models are Unsupervised Multitask Learners があると思う。いわゆる GPT(-1) と GPT-2 だ。GPT-2 のリファレンス実装は gpt-2 にある。

また、この GPT-2 を応用したアプリケーションにニューラルかな漢字変換エンジン「Zenzai」が存在し、macOS の「ライブ変換」のようなことが実現できている。（see ニューラルかな漢字変換エンジン「Zenzai」をazooKey on macOSに搭載します）

さて、今回も Transformer に触れてみる (1) と Transformer に触れてみる (2) — ViT もどき の続きであるが、「Zenzai」がローカルで動いていることから、お手頃なモデルであると判断し、大胆にも GPT-2 に手を出してみたいというものだ。

再び Vibe coding

懲りずに GPT-4.1 にお願いして学習素材を作ってもらうが、前々回の Transformer に触れてみる (1) の MiniFormer が再利用できるはずである。よって、前々回のソースコードをコンテキストとして与えて作成してもらった。

この判断の理由としては、Improving Language Understanding
by Generative Pre-Training を読むと、3.1 Unsupervised pre-training に

In our experiments, we use a multi-layer Transformer decoder [34] for the language model, which is a variant of the transformer [62]. This model applies a multi-headed self-attention operation over the input context tokens followed by position-wise feedforward layers to produce an output distribution over target tokens:

\begin{align*}h_0 &= U W_e + W_p \\ h_l &= \operatorname{transformer\_block}(h_{l−1}) \,\forall i \in [1, n] \\ P(u) &= \operatorname{softmax}(h_n W_e^T)\end{align*}\tag{2}

where U = (u_{−k}, \ldots , u_{−1}) is the context vector of tokens, n is the token embedding matrix, and W_p is the position embedding matrix.

とあるので、Transformer のデコーダを使えばいけるのだろうと思ったからである。さらに、GPT-2 論文 Language Models are Unsupervised Multitask Learners を読むと、2.3. Model に

We use a Transformer (Vaswani et al., 2017) based architecture for our LMs. The model largely follows the details of the OpenAI GPT model (Radford et al., 2018) with a few modifications. Layer normalization (Ba et al., 2016) was moved to the input of each sub-block, similar to a pre-activation residual network (He et al., 2016) and an additional layer normalization was added after the final selfattention block. A modified initialization which accounts for the accumulation on the residual path with model depth is used. We scale the weights of residual layers at initialization by a factor of 1/\sqrt{N} where N is the number of residual layers. The vocabulary is expanded to 50,257. We also increase the context size from 512 to 1024 tokens and a larger batchsize of 512 is used.

という感じなので、GPT-1 のレイヤ正規化の位置を変えたりすれば良いのだろうと感じた。

こういった内容を踏まえ、発注をし GPT-4.1 に作らせたものを Gemini 2.5 Flash に確認してもらったり、リファレンス実装 gpt-2 の比較をしてもらって教えてもらうといったことを試した。

作ったもの

  • タスクの内容: 先頭の数文字をプロンプトとして与えると、そこから動物、果物、野菜っぽい名前っぽい単語を捏造して返す。
  • ベースの Transformer の実装を Transformer に触れてみる (1) の MiniFormer とする。従って今回も、シンプルさを優先し、Attention はシングルヘッド、1 層のみとする。
  • Language Models are Unsupervised Multitask Learners のアーキテクチャを大幅に簡易化した MiniGPT2 を用いる。
  • データセットには Transformer に触れてみる (1) で用意した、動物、果物、野菜の名前からなる数百個の単語を利用。強い傾向があるわけでも、大量のデータがあるわけでもないので、モデルの学習における汎化性能は大きくは期待できない。

実装

基本的に MiniGPT2 以外は Transformer に触れてみる (1) と同じである。「# 5. MiniGPT2本体（Pre-LN）」で「★」を加えている辺りと generate メソッドが増えている点が変更点である。ほとんど MiniFormer である。「★」の変更についての幾つかの説明は後述する。

準備やデータセット

!pip install -qU bertviz

from __future__ import annotations

import json
import random
import string
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import numpy as np


SEED = 42
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(SEED)

ANIMALS = [
  "cat",
  "caracal",
  "capybara",
  "canary",
  "cavy",
  "caiman",
  "cacomistle",
  "caribou",
  "cassowary",
  "caterpillar",
  "dog",
  ...
]  # 長いので割愛

FRUITS_VEGGIES = [
  "apple",
  "apricot",
  "avocado",
  "artichoke",
  "banana",
  "bilberry",
  "blackberry",
  "blueberry",
  "boysenberry",
  "breadfruit",
  "cantaloupe",
  "casaba",
  ...
]  # 長いので割愛


# 1. データセット用：動物名+果物・野菜名で計1000種弱
NAMES = ANIMALS + FRUITS_VEGGIES

# 2. 文字のボキャブラリ作成
ALL_CHARS = sorted(set("".join(NAMES)))
SPECIAL_TOKENS = ["<PAD>", "<BOS>", "<EOS>"]
ALL_TOKENS = SPECIAL_TOKENS + ALL_CHARS
VOCAB_SIZE = len(ALL_TOKENS)
CHAR2IDX = {ch: i for i, ch in enumerate(ALL_TOKENS)}
IDX2CHAR = {i: ch for ch, i in CHAR2IDX.items()}
PAD_IDX = CHAR2IDX["<PAD>"]
BOS_IDX = CHAR2IDX["<BOS>"]
EOS_IDX = CHAR2IDX["<EOS>"]

def encode_word(word, max_len):
    tokens = [BOS_IDX] + [CHAR2IDX[c] for c in word] + [EOS_IDX]
    tokens += [PAD_IDX] * (max_len - len(tokens))
    return tokens

def decode_tokens(tokens):
    chars = []
    for idx in tokens:
        if idx == EOS_IDX:
            break
        if idx >= len(IDX2CHAR):
            continue
        ch = IDX2CHAR[idx]
        if ch not in SPECIAL_TOKENS:
            chars.append(ch)
    return "".join(chars)

# 3. PyTorch Dataset
class NameDataset(Dataset):
    def __init__(self, words, max_len):
        self.max_len = max_len
        self.data = []
        for w in words:
            tokens = encode_word(w, max_len)
            self.data.append(tokens)
    def __len__(self):
        return len(self.data)
    def __getitem__(self, idx):
        tokens = self.data[idx]
        x = torch.tensor(tokens[:-1], dtype=torch.long)
        y = torch.tensor(tokens[1:], dtype=torch.long)
        return x, y

# 4. シンプルな位置エンコーディング
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) * (-np.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        if d_model > 1:
            pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe)
    def forward(self, x):
        # x + self.pe[:x.size(1)].unsqueeze(0)
        return x + self.pe[:x.size(1)]

# 5. MiniGPT2本体（Pre-LN）
class MiniGPT2(nn.Module):
    def __init__(self, vocab_size, d_model=32, max_len=16):
        super().__init__()
        self.d_model = d_model
        self.embed = nn.Embedding(vocab_size, d_model)
        self.pos_enc = PositionalEncoding(d_model, max_len)
        # Self-Attention (single head, 1層)
        # ★ バイアス項はなし。詳細な理由は後述。
        self.q_linear = nn.Linear(d_model, d_model, bias=False)
        self.k_linear = nn.Linear(d_model, d_model, bias=False)
        self.v_linear = nn.Linear(d_model, d_model, bias=False)
        self.attn_out = nn.Linear(d_model, d_model, bias=False)
        # FFN
        # ★ バイアス項はなし。詳細な理由は後述。
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_model, bias=False),
            nn.ReLU(),
            nn.Linear(d_model, d_model, bias=False)
        )
        self.ln1 = nn.LayerNorm(d_model)
        self.ln2 = nn.LayerNorm(d_model)
        self.max_len = max_len
        self.attn_weights = None

    def forward(self, x, return_attn=False):
        emb = self.embed(x)
        emb = self.pos_enc(emb)
        # ★ Pre-LN構造：LN→Attention→Add
        attn_in = self.ln1(emb)
        Q = self.q_linear(attn_in)
        K = self.k_linear(attn_in)
        V = self.v_linear(attn_in)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / np.sqrt(self.d_model)
        # causal mask
        mask = torch.triu(torch.ones(scores.size(-2), scores.size(-1)), diagonal=1).bool().to(x.device)
        scores = scores.masked_fill(mask, float('-inf'))
        attn = torch.softmax(scores, dim=-1)
        attn_out = torch.matmul(attn, V)
        attn_out = self.attn_out(attn_out)
        # ★ self.ln1 は Pre-LN構造に移動
        x1 = emb + attn_out
        # ★ Pre-LN→FFN→Add
        ffn_in = self.ln2(x1)
        x2 = x1 + self.ffn(ffn_in)
        # ★ weight tying
        logits = torch.matmul(x2, self.embed.weight.t())
        if return_attn:
            self.attn_weights = attn.detach().cpu().numpy()
            return logits, attn
        return logits

    def generate(self, start_tokens, eos_idx, pad_idx, max_gen=None, temperature=1.0):
        self.eval()
        max_gen = max_gen or self.max_len
        tokens = start_tokens.tolist()
        for _ in range(max_gen - len(tokens)):
            inp = torch.tensor(tokens, dtype=torch.long).unsqueeze(0).to(next(self.parameters()).device)
            logits = self.forward(inp)
            next_token_logits = logits[0, len(tokens)-1] / temperature
            probs = torch.softmax(next_token_logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1).item()
            if next_token == eos_idx:
                break
            tokens.append(next_token)
        while len(tokens) < self.max_len:
            tokens.append(pad_idx)
        return tokens

訓練ループ

# 6. 訓練ループ
def train(model, loader, optimizer, criterion, device):
    model.train()
    total_loss = 0
    for x, y in loader:
        x = x.to(device)
        y = y.to(device)
        optimizer.zero_grad()
        logits = model(x)
        loss = criterion(logits.view(-1, VOCAB_SIZE), y.view(-1))
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    return total_loss / len(loader)

def evaluate(model, loader, criterion, device):
    model.eval()
    total_loss = 0
    with torch.no_grad():
        for x, y in loader:
            x = x.to(device)
            y = y.to(device)
            logits = model(x)
            loss = criterion(logits.view(-1, VOCAB_SIZE), y.view(-1))
            total_loss += loss.item()
    return total_loss / len(loader)

学習

%%time

# 設定
max_word_len = max(len(w) for w in NAMES) + 2 # BOS, EOS
batch_size = 16
d_model = 32
n_epochs = 100
device = "cuda" if torch.cuda.is_available() else "cpu"

# データ分割
random.seed(42)
random.shuffle(NAMES)
split = int(len(NAMES) * 0.8)
train_words = NAMES[:split]
test_words = NAMES[split:]

train_ds = NameDataset(train_words, max_word_len)
test_ds = NameDataset(test_words, max_word_len)
train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_ds, batch_size=batch_size, shuffle=False)

# モデル
model = MiniGPT2(VOCAB_SIZE, d_model, max_word_len).to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss(ignore_index=PAD_IDX)

# 学習
for epoch in range(1, n_epochs+1):
    train_loss = train(model, train_loader, optimizer, criterion, device)
    test_loss = evaluate(model, test_loader, criterion, device)
    if epoch % 5 == 0:
        print(f"Epoch {epoch:2d}: train loss={train_loss:.4f}, test loss={test_loss:.4f}")

Epoch 5: train loss=3.2279, test loss=3.1276
...
Epoch 100: train loss=2.2263, test loss=2.4894
CPU times: user 18.7 s, sys: 710 ms, total: 19.4 s
Wall time: 21 s

生成例

# 7. 単語生成
def sample_generate(prompt, model, max_word_len, temperature=1.0):
    start_tokens = [BOS_IDX] + [CHAR2IDX[c] for c in prompt]
    start_tokens = torch.tensor(start_tokens, dtype=torch.long)
    out_tokens = model.generate(start_tokens, eos_idx=EOS_IDX, pad_idx=PAD_IDX, max_gen=max_word_len, temperature=temperature)
    return decode_tokens(out_tokens[1:])  # BOSを除く

# ↓ 例： "ca" で始まる動物名っぽいものを生成
print("生成例:", sample_generate("ca", model, max_word_len, temperature=0.8))

生成例: categow

print("生成例:", sample_generate("do", model, max_word_len, temperature=0.8))

生成例: dock

print("生成例:", sample_generate("app", model, max_word_len, temperature=0.8))

生成例: appot

print("生成例:", sample_generate("sh", model, max_word_len, temperature=0.8))

生成例: shanelon

のような結果になった。何度かモデルを微調整しているが、その過程で得た他の例としては catelow, docket, appine, shoon や categothar, dotes, appinch, shabeer や categot, docketep, appincho, she であった。微妙なやつも混じっているが、「何かの名前っぽい雰囲気は感じられる。」という雰囲気はあるような気がする。

GPT-4.1 と Gemini Flash 2.5 による解説

適当に抜粋する。リファレンス実装 gpt-2 特に model.py と比較させた。自分では絶対に気が付かないようなところを教えてくれている。

主な整合ポイント

  1. 基本構造（デコーダ型Transformer、自己回帰）
  • 公式もあなたの実装も「Embedding → Positional Encoding → 複数ブロック（Attention＋FFN） → 出力層」という流れです。
  1. LayerNormの位置（Pre-LN）
  • 公式実装のblock()関数では、AttentionやMLP(FFN)の前にLayerNorm（norm(x, ...)）を適用（Pre-LN）。
  • あなたのMiniGPT2でもPre-LNに修正済み。
  1. 残差接続
  • 公式：x = x + a（Attention出力）、x = x + m（MLP出力）。
  • MiniGPT2も同じ（x1 = emb + attn_out、x2 = x1 + ffn(...)）。
  1. causal mask
  • 公式もPyTorchもtriuで上三角マスク実装、自己回帰を保証。
  1. 出力層

  • 公式実装（openai/gpt-2）

    • 最後のLayerNorm後、入力埋め込み行列wte（self.embed.weight相当）を転置して線形変換を行う（weight tying、重み共有）。
    • これによりパラメータ数削減と学習安定性が得られる。

  • MiniGPT2（この実装）

    • 最後のLayerNorm・FFN後、torch.matmul(x2, self.embed.weight.t()) でlogitsを計算し、埋め込み重みと出力重みを共有（weight tying）している。
    • 公式実装と同様にweight tyingを正しく導入しており、より“GPT-2らしい”挙動・構造となっている。
  1. Q, K, V, attn_out, FFN のバイアス無効

公式実装（openai/gpt-2/src/model.py より）

  • Q, K, V, attn_out, FFN（MLP）の全線形層は、バイアス項を持たない設計です。

    • 公式実装では、これらの線形変換は conv1d(..., ...) で実装されており、バイアス項は b = tf.get_variable('b', [nf], initializer=tf.constant_initializer(0)) により定義されていますが、weight tyingやLayerNormとの兼ね合い、および最新版PyTorch/TensorFlow実装では「バイアス項なし」がデファクトです。
    • OpenAIのPyTorch実装（transformersライブラリ等）でも bias=False で統一されています。

  • 具体的にバイアスが無効（bias=False）となるのは以下の部分：

    • Q, K, V: c_attn の分割
    • Attention出力: c_proj
    • FFNの各層: c_fc, c_proj

我々のMiniGPT2実装（PyTorch）

  • Q, K, V, attn_out, FFN（MLP）各線形層で bias=False を明示している。
  • これにより、GPT-2公式実装の設計思想（Q, K, V, attn_out, FFNすべてバイアス無し）と完全に整合しています。

補足

  • LayerNormやEmbeddingにはバイアスが存在しますが、これは公式実装でも同様です。
  • バイアスを無効にする主な理由は、「LayerNormとの組み合わせで不要」「パラメータ数削減」「学習挙動の安定化」などです。

今後の展望・応用例

  • 他の応用としては、GPT-4.1 によると「わざと 1 文字だけ間違えた単語をプロンプトにして、正しい単語を生成できるか試す」ことで typo 訂正タスクに転用できるかもとのことである。要は一種の Autoencoder による異常検知のような内容になるだろうか。正解とモデル出力の一致率や、Levenshtein距離（編集距離）などで評価すれば良いらしい。
  • モデルが単純なのでちょっときついかもしれないが、データセットをキャラ名からなるものに差し替えて、キャラクターの自動命名ツールとして使うというのもあるようだ。

まとめ

今回、漸く GPT-2 とはどういうものなのか少しだけ雰囲気が分かったような気がする。

元を辿れば一連の記事は ゼロから作るDeep Learning ❷ にもしも Extended Edition があったなら、こういう話が読みたかったというものである。「ゼロから作るDeep Learning ❷」を読んでいた頃はまだそこまで生成 AI が発達していなかったので、Transformer を含め、まともには手を出せない感覚だったが、現状までくると AI の力を借りて指導をしてもらえるようになった。

次は何をしようかな。

参考文献

GitHubで編集を提案

Discussion