Andrej Karpathy「Let's build GPT」解説シリーズ 第1動画・第2部
はじめに
第1部では、データセットの準備とトークナイゼーションについて学びました。第2部では、いよいよ最初の言語モデルであるBigramモデルを実装していきます。
この記事は、Andrej Karpathy氏の「Let's build GPT: from scratch, in code, spelled out.」動画の解説シリーズ第2部です。
Bigramモデルは「前の1文字だけを見て次の文字を予測する」シンプルなモデルですが、言語モデルの基本概念を理解するには最適です。実装を通して、埋め込み層、損失計算、文章生成の仕組みを詳しく見ていきましょう。
Bigramモデルの実装
モデルの基本構造
import torch.nn as nn
from torch.nn import functional as F
class BigramLanguageModel(nn.Module):
def __init__(self, vocab_size):
super().__init__()
# 各トークンの埋め込みベクトルを学習するテーブル
self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)
def forward(self, idx, targets=None):
# idx: (B, T) バッチ化されたトークンID列
logits = self.token_embedding_table(idx) # (B, T, vocab_size)
if targets is None:
loss = None
else:
B, T, C = logits.shape
logits = logits.view(B*T, C)
targets = targets.view(B*T)
loss = F.cross_entropy(logits, targets)
return logits, loss
token_embedding_tableの役割
self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)
この埋め込み層は特殊な構造になっています：
-
入力次元:
vocab_size（65文字） -
出力次元:
vocab_size（65文字）
つまり、65×65の重み行列を持っており、各行が各トークンIDに対応するベクトルになっています。
初期状態の例（vocab_size=5の場合）:
行\列 0列目 1列目 2列目 3列目 4列目
0行目 0.12 -0.33 0.01 0.45 -0.22
1行目 -0.11 0.09 0.38 -0.27 0.14
2行目 0.07 0.21 -0.19 0.33 -0.05
3行目 0.29 -0.41 0.17 0.02 0.08
4行目 -0.15 0.27 0.11 -0.09 0.36
各行のベクトルが「そのトークンの特徴表現」として学習されます。
順伝播（forward）の詳細
パラメータの理解
- idx: バッチ化されたトークンID列（例：xb、形は(B, T)）
- targets: 正解ラベル（例：yb、形は(B, T)）
処理ステップ
-
埋め込み層通過:
logits = self.token_embedding_table(idx)→ (B, T, vocab_size) -
形状変換:
F.cross_entropyが(N, C)と(N,)の形式を要求するため - 損失計算: クロスエントロピー誤差で損失を計算
B, T, C = logits.shape
logits = logits.view(B*T, C) # (B*T, vocab_size)
targets = targets.view(B*T) # (B*T,)
loss = F.cross_entropy(logits, targets)
F.cross_entropyは、各位置での分類問題として扱うため、バッチと時系列をまとめて1次元のサンプルとして処理します。
文章生成（generate）メソッド
def generate(self, idx, max_new_tokens):
for _ in range(max_new_tokens):
# 現在の系列に対して予測を実行
logits, loss = self(idx) # (B, T, vocab_size)
logits = logits[:, -1, :] # 最新時刻のみ取得 (B, vocab_size)
probs = F.softmax(logits, dim=-1) # 確率分布に変換
idx_next = torch.multinomial(probs, num_samples=1) # サンプリング (B, 1)
idx = torch.cat((idx, idx_next), dim=1) # 系列に追加 (B, T+1)
return idx
重要な概念の解説
-
self(idx):self.forward(idx)の省略記法 -
dim=-1: 最後の次元（vocab_size方向）でsoftmax -
torch.multinomial: 確率分布に従ってサンプリング -
dim=1: シーケンス長方向（T方向）にconcatenate
処理フロー
- 予測値取得: モデルから次トークンの確率分布を取得
- 最新時刻の抽出: 系列の最後の予測のみを使用
- 確率分布変換: logitsをsoftmaxで正規化
- サンプリング: 確率に応じてランダムにトークンを選択
- 系列更新: 選択されたトークンを系列に追加
初期状態での動作確認
m = BigramLanguageModel(vocab_size)
logits, loss = m(xb, yb)
print(logits.shape) # torch.Size([32, 65])
print(loss) # 4.8786（初期状態では高い損失）
# 初期状態での文章生成
print(decode(m.generate(idx=torch.zeros((1, 1), dtype=torch.long), max_new_tokens=100)[0].tolist()))
出力例:
SKIcLT;AcELMoTbvZv C?nq-QE33:CJqkOKH-q;:la!oiywkHjgChzbQ?u!3bLIgwevmyFJGUGp
wnYWmnxKWWev-tDqXErVKLgJ
初期状態では出力がめちゃくちゃなのは、まだbackward()によるパラメータの更新を行っていないためです。重みがランダム初期化されているので、意味のある出力にはなりません。
トレーニングプロセス
最適化器の設定と訓練ループ
# オプティマイザーの設定
optimizer = torch.optim.AdamW(m.parameters(), lr=1e-3)
batch_size = 32
for steps in range(10000):
# バッチデータの取得
xb, yb = get_batch('train')
# 順伝播
logits, loss = m(xb, yb)
# 勾配リセット
optimizer.zero_grad(set_to_none=True)
# 逆伝播
loss.backward()
# パラメータ更新
optimizer.step()
print(loss.item()) # 最終的な損失値
各ステップの詳細解説
1. 勾配のリセット
optimizer.zero_grad(set_to_none=True)
なぜ毎回勾配をリセットする必要があるのか？
PyTorchでは、loss.backward()を呼ぶと各パラメータの.grad属性に勾配が加算されていきます。そのため、毎回の学習ステップの最初に勾配をリセットしないと、前回の勾配が残ってしまい、正しい学習ができません。
-
set_to_none=True: メモリ効率化のため.gradをNoneに設定
2. 訓練ループの基本構造
-
バッチ取得:
xb, yb = get_batch('train') -
順伝播:
logits, loss = m(xb, yb) -
勾配リセット:
optimizer.zero_grad() -
逆伝播:
loss.backward() -
パラメータ更新:
optimizer.step()
この5ステップが機械学習の基本的なトレーニングループです。
トレーニング後の結果
print(decode(m.generate(idx=torch.zeros((1, 1), dtype=torch.long), max_new_tokens=300)[0].tolist()))
出力例:
lso br. ave aviasurf my, yxMPZI ivee iuedrd whar ksth y h bora s be hese, woweee; the! KI 'de, ulseecherd d o blllando;LUCEO, oraingofof win!
RIfans picspeserer hee tha,
TOFonk? me ain ckntoty ded. bo'llll st ta d:
ELIS me hurf lal y, ma dus pe athouo
BEY:! Indy; by s afreanoo adicererupa anse tecor
結果の分析
トレーニング後の出力を見ると：
- 良い点: それっぽい英語の文章構造になっている
- 問題点: 単語や文全体の意味がまったく通じない
なぜこうなるのか？
現在のBigramモデルは最後の文字だけを見て次の文字を予測しているからです。人間の言語理解には、もっと長いコンテキスト（文脈）が必要です。
例えば：
- "The cat sat on the ___" → "mat"を予測するには、"cat"も考慮する必要がある
- でもBigramでは直前の"e"しか見えない
位置埋め込みの追加
なぜ位置情報が必要なのか？
言語モデルをより賢くするためには、「トークンの位置情報」も重要です。同じ単語でも文のどこに現れるかで意味が変わることがあります。
class BigramLanguageModel(nn.Module):
def __init__(self, vocab_size):
super().__init__()
self.token_embedding_table = nn.Embedding(vocab_size, n_embed)
self.position_embedding_table = nn.Embedding(block_size, n_embed)
self.lm_head = nn.Linear(n_embed, vocab_size)
def forward(self, idx, targets=None):
B, T = idx.shape
# トークン埋め込みと位置埋め込みを取得
tok_emb = self.token_embedding_table(idx) # (B, T, n_embed)
pos_emb = self.position_embedding_table(torch.arange(T, device=device)) # (T, n_embed)
# 加算で情報を組み合わせ
x = tok_emb + pos_emb # (B, T, n_embed)
# 最終的な予測値を計算
logits = self.lm_head(x) # (B, T, vocab_size)
# 損失計算は同じ
if targets is None:
loss = None
else:
B, T, C = logits.shape
logits = logits.view(B*T, C)
targets = targets.view(B*T)
loss = F.cross_entropy(logits, targets)
return logits, loss
新しい要素の解説
1. 埋め込み次元の分離
self.token_embedding_table = nn.Embedding(vocab_size, n_embed)
self.position_embedding_table = nn.Embedding(block_size, n_embed)
self.lm_head = nn.Linear(n_embed, vocab_size)
-
n_embed: 内部表現の次元数（例：64、128など） -
lm_head: 内部表現を語彙サイズのスコアに変換する全結合層
2. 位置埋め込みの仕組み
pos_emb = self.position_embedding_table(torch.arange(T, device=device))
-
torch.arange(T)で[0, 1, 2, ..., T-1]の位置インデックスを作成 - 各位置に対応する埋め込みベクトルを取得
3. 情報の組み合わせ
x = tok_emb + pos_emb
トークン情報と位置情報を加算で組み合わせます。
PyTorchのnn.Linearの動作
self.lm_head = nn.Linear(n_embed, vocab_size)
PyTorchのnn.Linearは、最後の次元（ここではn_embed）に対して線形変換を行います。つまり、(B, T, n_embed)の入力に対して(B, T, vocab_size)の出力を生成します。
まとめ
第2部では、基本的な言語モデルの実装について学びました：
- Bigramモデル: 前の1文字だけで次を予測する最もシンプルなモデル
- 埋め込み層: トークンを連続ベクトルに変換する仕組み
- トレーニングループ: 勾配降下法による学習プロセス
- 位置埋め込み: トークンの位置情報をモデルに与える方法
しかし、現在のモデルはまだ「直前の1文字」しか考慮できません。より賢い言語モデルを作るためには、もっと長いコンテキストを活用する必要があります。
第3部では、この問題を解決するSelf-Attentionメカニズムを実装し、真のTransformerアーキテクチャに近づけていきます！
(この記事は研究室インターンで取り組みました：https://kojima-r.github.io/kojima/)