SNSでAndrej Karpathyが公開したmicrogpt.pyというコードを見かけました。
GPTの訓練と推論を、外部ライブラリなしのPython単一ファイルで実装したものです。
約150行。
やることは「人名を学習して、実在しない新しい人名を生成する」というシンプルなタスクですが、ChatGPTの基盤技術であるGPTの動作原理がそのまま入っています。
生成AIの仕組みを「なんとなくのイメージ」ではなく、コードとして具体的に追える。
それがこのファイルの面白さです。
1. まず動かしてみる
Python環境さえあれば、依存パッケージなしでそのまま実行できます。
$ python microgpt.py実行すると、まず訓練データの読み込みとモデルの初期化が表示されます。
num docs: 32033
vocab size: 27
num params: 765932,033件の人名を読み込み、語彙サイズは英小文字26種と特殊トークン1つで27。パラメータ数は7,659個です。ChatGPTのパラメータが数千億個であることを考えると、極めて小さなモデルだとわかります。
続いて訓練が始まり、損失の値が徐々に下がっていきます。損失とは「モデルの予測と正解のずれ」を表す数値で、小さいほど予測が正確です。
step 1 / 1000 | loss 3.4340
step 2 / 1000 | loss 3.4528
...
step 999 / 1000 | loss 1.7869
step 1000 / 1000 | loss 2.2convey1000ステップの訓練が終わると、学習済みモデルが新しい名前を生成します。
--- inference (new, hallucinated names) ---
sample 1: Katede
sample 2: Jede
sample 3: Arianna
sample 4: Khi
sample 5: Maede
...Code language: JavaScript (javascript)“Katede”や”Jede”のような、実在しないがどこかそれらしい名前が出てきます。モデルが「英語の名前っぽい文字の並び方」を学んだ結果です。
1.1. 生成AIは何をしているのか
GPTがやっていることは、突き詰めると「直前までの文字列から、次の1文字を予測する」の繰り返しです。“Ell”まで読んだら次は”a”が来そう、“Ella”の次は文末が来そう、という確率を出力します。これを1文字ずつ繰り返せば、名前全体が生成されます。
ChatGPTが長い文章を生成するのも原理は同じで、予測単位が文字からトークンに変わり、パラメータ数が桁違いに大きくなるだけです。
2. コードの全体設計
microgpt.pyは上から順に5つの段階で構成されています。データ準備、自動微分エンジン、モデル定義、訓練ループ、推論。ここからはコードを引用しながら、各段階の要所を見ていきます。
2.1. データの準備
docs = [l.strip() for l in open('input.txt').read().strip().split('\n') if l.strip()]Code language: JavaScript (javascript)人名リストを読み込み、1行1名前のリストにします。
uchars = sorted(set(''.join(docs)))
BOS = len(uchars)
vocab_size = len(uchars) + 1Code language: JavaScript (javascript)全名前に登場するユニークな文字を集め、それぞれにIDを振ります。aなら0、bなら1、という対応表です。BOSは「文の始まり・終わり」を示す特殊トークンで、語彙サイズの末尾に追加されます。このトークンがあることで、モデルは「ここから名前が始まる」「ここで名前が終わる」を学習できます。
2.2. 自動微分エンジン
ニューラルネットの学習には「各パラメータを少し動かしたら損失がどう変わるか」を知る必要があります。この変化の割合を勾配と呼び、勾配を自動で計算する仕組みが自動微分です。
class Value:
def __init__(self, data, children=(), local_grads=()):
self.data = data
self.grad = 0
self._children = children
self._local_grads = local_gradsValueは1つのスカラー値を包むオブジェクトです。dataが値そのもの、gradにはあとから逆伝播で勾配が書き込まれます。_childrenは「この値を作るのに使った値」、_local_gradsは「その値に対する局所的な変化の割合」を記録しており、この2つが計算の履歴にあたります。
def __add__(self, other):
other = other if isinstance(other, Value) else Value(other)
return Value(self.data + other.data, (self, other), (1, 1))
def __mul__(self, other):
other = other if isinstance(other, Value) else Value(other)
return Value(self.data * other.data, (self, other), (other.data, self.data))Code language: PHP (php)足し算では両方の子に対する局所勾配が1。掛け算ではたすき掛けになり、a * bのaに対する勾配はbの値、bに対する勾配はaの値です。exp、log、reluなどもすべて同じパターンで定義されています。
def backward(self):
# トポロジカルソートで計算グラフを整列
...
self.grad = 1
for v in reversed(topo):
for child, local_grad in zip(v._children, v._local_grads):
child.grad += local_grad * v.gradCode language: PHP (php)逆伝播の処理です。逆伝播とは、損失から出発して計算の履歴を逆順にたどり、各パラメータの勾配を求める手続きのこと。各ノードで「局所勾配 × 親の勾配」を足し込んでいくだけで、全パラメータの勾配が求まります。PyTorchのautogradと同じことを、素朴なPythonで再実装した形です。
3. モデルの構造
n_embd = 16
n_head = 4
n_layer = 1
block_size = 16埋め込み次元は、各文字を表すベクトルの長さです。ここでは16次元。アテンションヘッドは後述するアテンション処理を並列に走らせる数で、4つ。レイヤーはモデルの処理を何段重ねるかで、1段。最大系列長は一度に処理できる文字数で、16文字。名前生成に特化した極小設定です。
state_dict = {
'wte': matrix(vocab_size, n_embd),
'wpe': matrix(block_size, n_embd),
'lm_head': matrix(vocab_size, n_embd)
}Code language: JavaScript (javascript)wteはトークン埋め込み行列で、各トークンIDを16次元のベクトルに変換します。wpeは位置埋め込みで、「系列中の何番目か」という情報をベクトルとして与えるもの。lm_headは最後に16次元ベクトルを語彙サイズの次元に射影し、次トークンの確率分布を出力します。
3.1. 入力の組み立て
tok_emb = state_dict['wte'][token_id]
pos_emb = state_dict['wpe'][pos_id]
x = [t + p for t, p in zip(tok_emb, pos_emb)]Code language: JavaScript (javascript)トークン埋め込みと位置埋め込みを足し合わせたものが、各トークンの初期表現になります。「どの文字か」と「何番目か」の情報が1つのベクトルに合成されるわけです。
3.2. アテンション
このコードで個人的に一番面白かった箇所です。アテンションとは、「今注目している文字が、過去のどの文字からどれだけ情報を受け取るか」を動的に決める仕組みです。
attn_logits = [
sum(q_h[j] * k_h[t][j] for j in range(head_dim)) / head_dim**0.5
for t in range(len(k_h))
]
attn_weights = softmax(attn_logits)
head_out = [
sum(attn_weights[t] * v_h[t][j] for t in range(len(v_h)))
for j in range(head_dim)
]各文字はクエリ、キー、バリューという3種類のベクトルに変換されます。クエリは「何を探しているか」、キーは「自分が何を持っているか」、バリューは「実際に渡す情報」に相当します。クエリとキーの内積でスコアを計算し、softmaxという関数で合計が1になるよう正規化してからバリューの加重平均を取ります。head_dim**0.5で割るのは、次元が大きいときにスコアが極端な値になるのを防ぐためです。フレームワークの行列演算に隠れがちなこの処理が、forループと四則演算だけで全部見えます。
3.3. MLP
x = linear(x, state_dict[f'layer{li}.mlp_fc1'])
x = [xi.relu() for xi in x]
x = linear(x, state_dict[f'layer{li}.mlp_fc2'])Code language: JavaScript (javascript)MLPは多層パーセプトロンの略で、単純な全結合ニューラルネットワークです。16次元を64次元に拡大してReLUという活性化関数で負の値をゼロに切り捨て、また16次元に戻します。アテンションが「どの情報を集めるか」を決めた後、MLPが「集めた情報をどう変換するか」を担います。両ブロックとも残差接続という手法で入力をそのまま足し戻しており、これが学習を安定させています。
4. 訓練ループ
tokens = [BOS] + [uchars.index(ch) for ch in doc] + [BOS]名前を1つ取り出し、前後にBOSを付けてトークン列にします。たとえば”emma”なら[BOS, 4, 12, 12, 0, BOS]のような数列になります。
for pos_id in range(n):
logits = gpt(token_id, pos_id, keys, values)
probs = softmax(logits)
loss_t = -probs[target_id].log()1トークンずつ順にモデルへ入力し、次トークンの予測確率を得ます。logitsはモデルが出力する生のスコアで、softmaxで確率に変換されます。正解トークンの確率の負の対数が損失。正解に高い確率を割り当てるほど損失は小さくなります。
loss.backward()
lr_t = learning_rate * (1 - step / num_steps)
for i, p in enumerate(params):
m[i] = beta1 * m[i] + (1 - beta1) * p.grad
v[i] = beta2 * v[i] + (1 - beta2) * p.grad ** 2
m_hat = m[i] / (1 - beta1 ** (step + 1))
v_hat = v[i] / (1 - beta2 ** (step + 1))
p.data -= lr_t * m_hat / (v_hat ** 0.5 + eps_adam)
p.grad = 0逆伝播で全パラメータの勾配を求めた後、Adamというオプティマイザで更新します。オプティマイザとは「勾配をもとにパラメータをどう動かすか」を決めるアルゴリズムのことです。Adamは勾配の移動平均と勾配の二乗の移動平均を追跡し、パラメータごとに更新幅を適応的に調整します。学習率は線形に減衰させ、訓練終盤では慎重に更新させています。
4.1. 推論
logits = gpt(token_id, pos_id, keys, values)
probs = softmax([l / temperature for l in logits])
token_id = random.choices(range(vocab_size), weights=[p.data for p in probs])[0]BOSから始めて、モデルが出力する確率分布から1文字ずつサンプリングし、再びBOSが出たら生成終了です。temperatureはlogitsを割る値で、小さいほど確率分布が尖り、最も確率の高いトークンが選ばれやすくなります。0.5は比較的保守的な設定で、訓練データに近い名前が生成されやすくなります。
5. 「あとは効率の問題」
ファイル冒頭のコメントにこうあります。
This file is the complete algorithm. Everything else is just efficiency.
PyTorchやGPUは、ここに書かれている計算を高速に実行するための道具にすぎません。アルゴリズム自体は150行のPythonに収まる。この事実が、生成AIの仕組みをぐっと身近にしてくれます。
興味がある方はぜひ手元で動かしてみてください。Python環境さえあれば、依存パッケージのインストールなしにそのまま実行できます。