Andrej Karpathy「Let's build GPT」解説シリーズ 第4動画
はじめに
この回では、Andrej Karpathy氏の「Let's reproduce GPT-2 (124M)」を参考に、GPT2をついにゼロから作り出すことに挑戦します。今回はその第一歩として、GPTモデルの根幹をなすTransformerアーキテクチャをPyTorchでスクラッチ実装していきます。
Hugging Faceのライブラリを使えば数行で呼び出せるGPT-2モデルですが、その内部は一体どうなっているのでしょうか？本記事では、GPT、Block、CausalSelfAttention、MLPといったコンポーネントを一つずつ組み立てながら、その仕組みを深掘りします。
なぜデコーダーだけで良いのか？
GPT（Generative Pre-trained Transformer）は、その名の通りTransformerのアーキテクチャに基づいています。特にGPT-2のようなモデルは、元々のTransformerのエンコーダー・デコーダー構造のうち、デコーダー部分のみを採用しているのが特徴です（よく「Decoder-only」と呼ばれます）。
なぜデコーダーだけなのでしょうか？
- エンコーダー: 入力文全体の文脈を双方向に読み取り、情報を圧縮・表現する役割（例: 文章の感情分析）。
- デコーダー: それまでの情報（とエンコーダーからの情報）を元に、次に来る単語を自己回帰的に予測する役割。
GPTは「次に来る単語を予測する」という言語モデリングに特化しているため、デコーダーの仕組みだけで十分なのです。入力されたテキストの続きを、ひたすら予測していくことで文章を生成します。
まずはHugging FaceのGPT-2モデルの内部を覗いて、どのような構造なのか確認してみましょう。
from transformers import GPT2LMHeadModel
model_hf = GPT2LMHeadModel.from_pretrained("gpt2") # 124M
sd_hf = model_hf.state_dict()
for k, v in sd_hf.items():
print(k, v.shape)
# 出力例
# transformer.wte.weight torch.Size([50257, 768]) # Word Token Embedding
# transformer.wpe.weight torch.Size([1024, 768]) # Word Position Embedding
# transformer.h.0.ln_1.weight torch.Size([768]) # 最初のTransformerブロックのLayerNorm
# ...
# transformer.ln_f.weight torch.Size([768]) # 最終LayerNorm
# lm_head.weight torch.Size([50257, 768]) # 出力層
この構造を、これから一つずつ実装していきます。
実装の全体像
GPTモデルは、いくつかのクラスの組み合わせで構築されます。全体像は以下のようになります。
-
MLP(多層パーセプトロン): 各トークンの特徴量を非線形変換し、表現力を高めます。 -
CausalSelfAttention: 各トークンが文脈中の「過去の」トークンだけを参照して、次に来る単語を予測するための情報を集約します。 -
Block:CausalSelfAttentionとMLPを組み合わせた、Transformerの基本単位。これを何層も積み重ねます。 -
GPT:Blockを束ね、Embedding層や最終出力層を追加してモデル全体を構成します。
(GPT2では右側の、Decoder部分を実装します)
まずは、全体の骨格となるGPTクラスのスケルトンから作ってみましょう。
class GPT(nn.Module):
def __init__(self, config):
super().__init__()
self.config = config
self.transformer = nn.ModuleDict(dict(
wte = nn.Embedding(config.vocab_size, config.n_embd),
wpe = nn.Embedding(config.block_size, config.n_embd),
h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
ln_f = nn.LayerNorm(config.n_embd),
))
self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
具体的な実装
1. MLP (Multi-Layer Perceptron)
MLPは、Attention層で集約された情報をさらに加工・変換する役割を担います。実装はシンプルですが、モデルの表現力を高めるために重要です。
import torch.nn as nn
from torch.nn import functional as F
class MLP(nn.Module):
def __init__(self, config):
super().__init__()
# 埋め込み次元を4倍に拡張する全結合層
self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd)
# GELU活性化関数。tanhによる近似版を使い、計算を高速化
self.gelu = nn.GELU(approximate='tanh')
# 元の埋め込み次元に戻す全結合層
self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd)
def forward(self, x):
x = self.c_fc(x)
x = self.gelu(x)
x = self.c_proj(x)
return x
なぜ次元を4倍に拡張するのか？
これは、一度特徴量をより高次元の空間に射影することで、より複雑なパターンを学習できるようにするためです。様々な特徴を捉えた後、再び元の次元に情報を圧縮します。「一度大きく広げて、いろんな方向に情報を変換し、またギュッとまとめる」というイメージです。
2. Causal Self-Attention
ここがGPTの最も重要な部分で、「Causal（因果的）」、つまり「未来の情報を見ない」という制約のついた自己注意機構です。
Self-Attentionについては、第3部の「GPTをゼロから実装して理解してみる（第3部：Self-AttentionとTransformerアーキテクチャ編）」で解説しています。
class CausalSelfAttention(nn.Module):
def __init__(self, config):
super().__init__()
assert config.n_embd % config.n_head == 0
# Query, Key, Valueをまとめて生成する全結合層
self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
# Attentionの結果を元の次元に戻す全結合層
self.c_proj = nn.Linear(config.n_embd, config.n_embd)
self.n_head = config.n_head
self.n_embd = config.n_embd
# causal maskを定義
self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
.view(1, 1, config.block_size, config.block_size))
def forward(self, x):
B, T, C = x.size() # Batch, Time, Channels
# 1. Q, K, Vの生成
qkv = self.c_attn(x)
q, k, v = qkv.split(self.n_embd, dim=2)
# 2. Multi-Head化
# (B, T, C) -> (B, T, n_head, head_size) -> (B, n_head, T, head_size)
k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
# 3. Attentionスコアの計算
# (B, n_head, T, hs) @ (B, n_head, hs, T) -> (B, n_head, T, T)
att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
# 4. マスキング
att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
# 5. Softmax
att = F.softmax(att, dim=-1)
# 6. Valueとの積
y = att @ v # (B, n_head, T, T) @ (B, n_head, T, hs) -> (B, n_head, T, hs)
# 7. Headの結合
y = y.transpose(1, 2).contiguous().view(B, T, C)
# 8. 出力層
y = self.c_proj(y)
return y
ポイント解説:
-
c_attn: 1つの全結合層でQuery, Key, Valueを一度に計算することで、効率化を図っています。 -
Multi-Head:
n_embd次元のベクトルをn_head個の小さなhead_sizeのベクトルに分割します。これにより、異なる「観点」から情報の関連性を捉えることができます。 -
bias(Causal Mask):torch.trilで作成した下三角行列です。未来のトークンに対応する部分を-infで埋めることで、Softmax後にその部分の重みが0になり、「未来を見ない」という制約を実現します。
3. Block
Blockは、CausalSelfAttentionとMLPを、それぞれLayer Normalizationと残差接続（Residual Connection）でつないだものです。残差接続（x = x + ...の部分）は、勾配が消失しにくくし、深いネットワークの学習を安定させるために不可欠です。
class Block(nn.Module):
def __init__(self, config):
super().__init__()
self.ln_1 = nn.LayerNorm(config.n_embd)
self.attn = CausalSelfAttention(config)
self.ln_2 = nn.LayerNorm(config.n_embd)
self.mlp = MLP(config)
def forward(self, x):
# Attention -> 残差接続
x = x + self.attn(self.ln_1(x))
# MLP -> 残差接続
x = x + self.mlp(self.ln_2(x))
return x
4. GPTモデルの完成
最後に、GPTクラスにforwardメソッドを実装し、すべての部品を統合します。
# GPTクラスに追加
def forward(self, idx, targets=None):
B, T = idx.size()
assert T <= self.config.block_size, f"Cannot forward sequence of length {T}, block size is only {self.config.block_size}"
# 1. Token EmbeddingとPosition Embedding
pos = torch.arange(0, T, dtype=torch.long, device=idx.device)
pos_emb = self.transformer.wpe(pos) # (T, n_embd)
tok_emb = self.transformer.wte(idx) # (B, T, n_embd)
x = tok_emb + pos_emb
# 2. Transformer Blockの実行
for block in self.transformer.h:
x = block(x)
# 3. 最終LayerNormと出力層
x = self.transformer.ln_f(x)
logits = self.lm_head(x) # (B, T, vocab_size)
# 4. 損失計算 (学習時のみ)
loss = None
if targets is not None:
loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
return logits, loss
入力されたトークンID (idx) は、wteで単語のベクトルに、wpeで位置のベクトルに変換され、足し合わされます。これがTransformerブロック群を通り、最終的にlm_headで各単語の出現確率（の対数であるロジット）に変換されます。
実験と検証
GPT-2（124Mモデル）と同じ設定で、今作ったモデルを初期化してみましょう。
@dataclass
class GPTConfig:
block_size: int = 1024
vocab_size: int = 50257
n_layer: int = 12
n_head: int = 12
n_embd: int = 768
# モデルを初期化
model = GPT(GPTConfig())
この（まだ全く学習していない）モデルにテキストを生成させると、どうなるでしょうか？
import tiktoken
enc = tiktoken.get_encoding("gpt2")
tokens = enc.encode("Hello, I'm a language model")
tokens = torch.tensor(tokens, dtype=torch.long).unsqueeze(0)
# 生成ループ
x = tokens
while x.size(1) < 30: # 30トークンまで生成
with torch.no_grad():
logits, _ = model(x)
logits = logits[:, -1, :] # 最後のトークンのロジットだけ使う
probs = F.softmax(logits, dim=-1)
# 確率上位50個からランダムに選ぶ (Top-kサンプリング)
topk_probs, topk_indices = torch.topk(probs, 50, dim=-1)
ix = torch.multinomial(topk_probs, 1)
xcol = torch.gather(topk_indices, -1, ix)
x = torch.cat((x, xcol), dim=1)
decoded = enc.decode(x[0].tolist())
print(">", decoded)
# 出力例:
# > Hello, I'm a language model troopers broadcasterurned fraudulent foreclosure protest permanclick cheesBir McCarthy Returnhoff Kashmir announced SOCasket HotSword Chaoku metaphonder
当然ながら、ランダムで意味不明な文字列が出力されました。モデルの重みがランダムに初期化されているためです。
よくあるミス
-
Causal Maskの重要性: Self-Attentionで未来のトークンをマスクし忘れると、モデルは単に「次の単語をカンニングする」だけになってしまい、何も学習できません。
is_causal=Trueや手動でのマスキングが不可欠です。 - Position Embedding: トークンの「順序」を伝える位置エンベディングを足し忘れると、モデルは単語の袋（Bag of Words）としてしか入力を認識できず、文法や構造を学べません。
- 重みの初期化: 今回はデフォルトの初期化を使いましたが、大規模モデルの学習を安定させるには、分散をコントロールした適切な初期化が極めて重要です。これは次回の記事で詳しく扱います。
まとめ
今回は、GPT-2で重要なTransformerアーキテクチャを、CausalSelfAttentionやMLPといった部品から一つずつ組み立てました。まだ学習していないためランダムな出力しかできませんが、モデルの骨格はこれで完成です。
次回は、モデルの学習のプロセスに入ります。損失関数、オプティマイザ、そしてデータセットからデータを供給するDataLoaderを実装し、モデルが意味のあるテキストを生成できるように訓練していきます。
(この記事は研究室インターンで取り組みました：https://kojima-r.github.io/kojima/)