Andrej Karpathy「Let's build GPT」解説シリーズ 第1動画・第3部
はじめに
第2部では、基本的なBigramモデルを実装しましたが、「直前の1文字しか見られない」という大きな制約がありました。第3部では、この問題を解決するSelf-Attentionメカニズムを実装し、真のTransformerアーキテクチャを構築していきます。
この記事は、Andrej Karpathy氏の「Let's build GPT: from scratch, in code, spelled out.」動画の解説シリーズ第3部(最終部)です。
Self-Attentionは、「系列内の全ての位置が全ての位置を参照できる」革新的な仕組みで、現代のGPTやBERTの核心技術です。実装を通して、このメカニズムを詳しく理解していきましょう。
Self-Attentionの数学的仕組み
課題:前のコンテキストをどう活用するか?
現在のBigramモデルでは直前の1文字しか見られませんが、実際の言語理解には「前のコンテキスト全体」が必要です。まずは、前のコンテキストの情報を単純に平均化する方法から始めてみましょう。
Version 1: 二重ループによる平均化
# 前のコンテクストから平均をとる version1
torch.manual_seed(1337)
B, T, C = 4, 8, 2 # バッチサイズ4、シーケンス長8、特徴次元2
x = torch.randn(B, T, C)
# bag of words
xbow = torch.zeros((B, T, C))
for b in range(B):
for t in range(T):
xprev = x[b, :t+1] # (t+1, C) - 0〜t番目まで
xbow[b, t] = torch.mean(xprev, 0) # 時系列方向の平均
なぜ「bag of words」なのか?
- 各時刻で「その時刻までの全ての情報の平均」が格納される
- 単語の順序情報は失われるが、これまでのコンテキストの概要は把握できる
- これがSelf-Attentionの基本的なアイデアの簡単な例
Version 2: 行列演算による効率化
# 前のコンテクストから平均をとる version2
wei = torch.tril(torch.ones(T, T)) # 下三角行列
wei = wei / wei.sum(1, keepdim=True) # 各行の合計を1に正規化
xbow2 = wei @ x # 行列積
weiの形状:
tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
[0.5000, 0.5000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
[0.3333, 0.3333, 0.3333, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
[0.2500, 0.2500, 0.2500, 0.2500, 0.0000, 0.0000, 0.0000, 0.0000],
...])
- 各行の合計が1になっており、ピラミッド型の重み行列
- これをxと掛け合わせることで、これまでのコンテキストの平均が得られる
Version 3: Softmaxを使った実装
# 前のコンテクストから平均をとる version3
tril = torch.tril(torch.ones(T, T)) # 下三角行列
wei = torch.zeros((T, T))
wei = wei.masked_fill(tril == 0, float('-inf')) # 未来の情報を-infでマスク
wei = F.softmax(wei, dim=-1) # softmaxで確率分布化
xbow3 = wei @ x # 行列積
なぜ「マスク付きsoftmax」なのか?
- softmaxの基本: 入力ベクトルを確率分布(合計が1)に変換
-
マスクの役割: 未来の情報(tより後の時刻)を
-inf
で埋める -
-infの効果:
exp(-inf) = 0
なので、マスクされた部分の重みは0になる - 結果: 過去と現在の情報だけで重み付き平均を計算
この方法は、実際のTransformerのself-attentionに最も近い実装です!
Self-Attentionの実装
Self-Attentionとは?
Self-Attentionは、「系列内の各トークンが他の全てのトークンとの関係性を学習し、重要な情報に注目する」仕組みです。
基本概念
- 系列内の関係性: 同じシーケンス内のトークン同士が相互作用
- 注目メカニズム: 重要な情報により多くの重みを割り当て
- 並列処理: 全ての位置を同時に処理可能
クエリ、キー、バリューの概念
# version4: self-attention
head_size = 16
key = nn.Linear(C, head_size, bias=False)
query = nn.Linear(C, head_size, bias=False)
value = nn.Linear(C, head_size, bias=False)
k = key(x) # (B, T, head_size)
q = query(x) # (B, T, head_size)
v = value(x) # (B, T, head_size)
3つの要素の役割
- クエリ(Query): 「今注目しているトークン」が「何を求めているか」を表す
- キー(Key): 各トークンが「どんな特徴を持っているか」を表す
- バリュー(Value): 実際に取り出したい「情報の内容」を表す
Attention重みの計算
wei = q @ k.transpose(-2, -1) * head_size**-0.5 # Scaled Dot-Product
wei = wei.masked_fill(tril[:T, :T] == 0, float('-inf')) # 因果性マスク
wei = F.softmax(wei, dim=-1) # 注目度の正規化
out = wei @ v # 重み付けされた情報の取得
なぜスケーリング(head_size**-0.5
)が必要?
- クエリとキーの内積は
head_size
に比例して大きくなる - 値が大きすぎるとsoftmaxが極端な分布になってしまう
-
√head_size
で割ることで、適切な範囲に正規化
Self-Attentionヘッドの実装
class Head(nn.Module):
"""one head of self-attention"""
def __init__(self, head_size):
super().__init__()
self.key = nn.Linear(n_embed, head_size, bias=False)
self.query = nn.Linear(n_embed, head_size, bias=False)
self.value = nn.Linear(n_embed, head_size, bias=False)
self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))
self.dropout = nn.Dropout(dropout)
def forward(self, x):
B, T, C = x.shape
k = self.key(x) # (B, T, head_size)
q = self.query(x) # (B, T, head_size)
# Scaled Dot-Product Attention
wei = q @ k.transpose(-2, -1) * C**-0.5 # (B, T, T)
wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
wei = F.softmax(wei, dim=-1)
wei = self.dropout(wei)
# 重み付きバリューの取得
v = self.value(x)
out = wei @ v
return out
register_bufferの役割
self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))
- モデルのパラメータではないが、モデルと一緒に管理したいテンソルを登録
- GPU/CPUの移動や保存時に自動的に処理される
-
tril
は学習不要な固定マスクなので、bufferとして登録
マルチヘッドアテンション
なぜ「マルチヘッド」なのか?
1つのSelf-Attentionヘッドでは、1種類の「注目の仕方」しか学習できません。しかし、言語には様々な関係性があります:
- 文法的関係: 主語と動詞、修飾語と被修飾語
- 意味的関係: 同義語、対義語、連想
- 位置的関係: 近接、遠距離依存
マルチヘッドアテンションでは、複数の異なる視点で同時に文脈を捉えます。
class MultiHeadAttention(nn.Module):
def __init__(self, num_heads, head_size):
super().__init__()
self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
self.proj = nn.Linear(n_embed, n_embed)
self.dropout = nn.Dropout(dropout)
def forward(self, x):
# 各ヘッドの出力を結合
out = torch.cat([h(x) for h in self.heads], dim=-1)
out = self.dropout(self.proj(out))
return out
重要なポイント:
- 各ヘッドが並列に異なる特徴を学習
- 最後に
torch.cat
で結合して元の次元数に戻す - プロジェクション層で情報を統合
FeedForward層とTransformerブロック
FeedForward層の役割
class FeedFoward(nn.Module):
def __init__(self, n_embed):
super().__init__()
self.net = nn.Sequential(
nn.Linear(n_embed, 4 * n_embed), # 拡張
nn.ReLU(),
nn.Linear(4 * n_embed, n_embed), # 縮小
nn.Dropout(dropout),
)
def forward(self, x):
return self.net(x)
FeedForwardの特徴
- 位置ごとの独立処理: 各トークンを独立して変換(系列の文脈は見ない)
- 非線形変換: ReLUにより複雑な特徴変換が可能
- 次元の拡張と縮小: 内部で4倍に拡張してから元に戻す
Transformerブロック
class Block(nn.Module):
def __init__(self, n_embed, n_head):
super().__init__()
head_size = n_embed // n_head
self.sa = MultiHeadAttention(n_head, head_size)
self.ffwd = FeedFoward(n_embed)
self.ln1 = nn.LayerNorm(n_embed)
self.ln2 = nn.LayerNorm(n_embed)
def forward(self, x):
x = x + self.sa(self.ln1(x)) # 残差接続 + Layer Norm
x = x + self.ffwd(self.ln2(x)) # 残差接続 + Layer Norm
return x
重要な要素の解説
1. 残差接続(Residual Connection)
x = x + self.sa(self.ln1(x))
- 入力
x
をそのまま出力に加算 - 効果: 勾配の流れを改善し、深いネットワークでも学習が安定
- アイデア: 「変化分だけを学習する」方が効率的
2. Layer Normalization
self.ln1 = nn.LayerNorm(n_embed)
- 各層の入力を正規化(平均0、分散1)
- 効果: 学習の安定化、勾配消失/爆発の防止
- 配置: 各サブ層の前に配置(Pre-LN)
3. Dropout
self.dropout = nn.Dropout(dropout)
- 訓練時にランダムにユニットを無効化
- 効果: 過学習の防止、汎化性能の向上
- Self-Attentionでの適用: 注目重みにもDropoutを適用
GPTモデル
最終的なアーキテクチャ
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)
self.blocks = nn.Sequential(
Block(n_embed, n_head=4),
Block(n_embed, n_head=4),
Block(n_embed, n_head=4),
nn.LayerNorm(n_embed),
)
self.sa_head = MultiheadAttention(4, n_embed//4)
self.ffwd = FeedFoward(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)
pos_emb = self.position_embedding_table(torch.arange(T, device=device))
x = tok_emb + pos_emb
# Transformerブロック
x = self.blocks(x)
logits = self.lm_head(x)
# 損失計算
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
アーキテクチャの特徴
- 多層構造: 複数のTransformerブロックを積み重ね
- 位置埋め込み: トークンの位置情報を追加
- 最終正規化: 出力前のLayer Normalization
- 言語モデルヘッド: 内部表現を語彙確率に変換
Self-AttentionとCross-Attentionの違い
Self-Attention(今回実装したもの)
- 定義: 同じ系列内のトークン同士が相互作用
- 用途: 文脈理解、長距離依存の捉える
- 例: 「I love you」で、各単語が同じ文の他の単語に注目
Cross-Attention(参考)
- 定義: 異なる系列間で注目し合う仕組み
- 用途: 機械翻訳、質問応答システム
- 例: 翻訳時にソース言語とターゲット言語間でアテンション
まとめ
第3部では、GPTの核心技術であるSelf-AttentionとTransformerアーキテクチャを実装しました:
- Self-Attention: 系列内の全位置間の関係性を捉える革新的メカニズム
- マルチヘッドアテンション: 複数の視点から並列に特徴を学習
- Transformerブロック: 残差接続、Layer Norm、FeedForwardの組み合わせ
- 完全なGPT: 現代的なアーキテクチャの実装
この実装で学んだこと
- 注目メカニズム: 重要な情報に動的に注目する仕組み
- 並列処理: RNNの逐次処理の制約を克服
- スケーラビリティ: 層を深くしても学習が安定する設計
3部にわたるシリーズを通して、GPTの基本的な仕組みから最新のTransformerアーキテクチャまでを実装できました。これで「GPTがどう動いているか」を深く理解できたはずです!
シリーズ継続予定
これにて第1回動画「Let’s build GPT: from scratch, in code, spelled out.」の解説は終了です!
次回は、今まで学習したGPTがどのようなプロセスを経て私たちが普段使うChatGPTのようになるのかについて第2動画で学んでいきたいと思います。
-
第2動画: State of GPT | BRK216HFS
お楽しみに!
(この記事は研究室インターンで取り組みました:https://kojima-r.github.io/kojima/)