はじめに
前稿では、GPT（Generative Pre-trained Transformer）を支える基盤技術である Transformer のアーキテクチャ、とりわけ Self-Attention 機構の数学的背景や、最尤推定に基づく自己教師あり学習のプロセスについて理論的概説を行いました。
また 2WINS では隔週で社内勉強会を開催しており、本稿は第15回勉強会「TransformerとGPTの仕組み」で扱った内容をベースに、GPT の内部構造理解を目的として microgpt を題材に学習した内容を発展させたものです。
理論から実践へと一歩進み、PyTorch などの深層学習フレームワークに依存せず、自動微分からモデルアーキテクチャまでを約200行の Python コードで実装した「microgpt」を対象に解説していきたいと思います。
今回扱うコードは、学習用の microgpt_train.py と生成用の microgpt_generate.py です。
前者では日本語の名前データを使って文字単位の GPT を学習し、後者では保存した重みを読み込んで新しい名前を生成します。
この記事では、次の流れで見ていきます。
- どんなデータをどうトークン化しているか
- GPT 本体がどう次の文字を予測するか
- 学習ループでどう重みを更新するか
- 保存した重みを使って推論する流れ
- 学習と推論がどうつながっているか
この記事で扱うコード
学習スクリプトでは、argparse で層数、埋め込み次元、ヘッド数、学習率、学習ステップ数などを受け取れるようになっており、実験条件をコマンドラインから変えられるようになっています。
また、学習後には重みを pickle 形式で weights/ 配下に保存し、生成スクリプト側でそのファイルを読み込めるようになっています。
生成スクリプトでは、--start で最初の一文字を指定したり、--num で生成件数を変えたり、--temperature で生成の多様性を調整したりできます。
つまり今回の実装は、「学習して終わり」ではなく、「学習済み重みを保存し、あとで推論に再利用する」までがひとまとまりになっています。
実行例は次のようなイメージです。※次の、「運用・可視化編：Wandbで学習をコントロールする」にて、こちらの条件で実行してみます。
# 学習（デフォルト設定でモデルを学習）
uv run microgpt_train.py
# ハイパーパラメータを変えて学習
# 埋め込み次元とヘッド数を増やして、少し大きめのモデルを試す
uv run microgpt_train.py --n_embd 64 --n_head 8
# 学習ステップ数と学習率を変えて、収束の仕方の違いを観察する
uv run microgpt_train.py --num_steps 2000 --lr 0.005
# 生成（学習済みの最新の重みからランダムに名前を生成）
uv run microgpt_generate.py
# 「あ」から始まる名前だけに絞って 10 個生成する
uv run microgpt_generate.py --start あ -n 10
# temperature を上げて、より多様でランダム性の高い名前を生成する
uv run microgpt_generate.py --temperature 0.8
また、以下にmicrogpt_train.py、microgpt_generate.pyのコード全文を示します。ここからは、こちらのコードをピックアップしていくような形で、解説していこうと思います。なお、リポジトリ内にデータセット(japanese_names.txt)も入っているのでぜひご使用下さい。
全体像をつかもう
まずは、この実装全体を大きな流れで見ておきます。
学習から推論までの流れは、次のように整理できます。
- 名前データを読み込む
- 文字単位で語彙を作り、文字を整数 ID に変換できるようにする
- GPT の重みをランダム初期化する
- 「これまでの文字列から次の文字を当てる」学習を繰り返す
- 学習済み重みを
.pklに保存する - 生成時にその重みを読み込む
- 1文字ずつ自己回帰的に次の文字をサンプリングし、名前を完成させる
この流れを見ると、学習と推論の違いは「重みを更新するかどうか」です。
学習では loss を計算して backward し、推論では保存済みの重みを使って次の文字を選ぶだけです。
1. データからトークンへ
この実装は文字単位でトークナイズします。
uchars = sorted(set(''.join(docs))) でデータ中の全文字を語彙化し、各文字に整数 ID を割り当てます。
通常の文字に加えて、文頭・文末を兼ねる区切りトークン BOS を1つ追加しています。
たとえば "さくら" は次のように変換されます。
[BOS, さ, く, ら, BOS]
学習では、このトークン列を1つずらして「入力→正解」のペアを作ります。
token_id = tokens[pos_id]
target_id = tokens[pos_id + 1]
logits = gpt(token_id, pos_id, keys, values)
loss_t = -softmax(logits)[target_id].log() # cross-entropy
BOS → さ、さ → く、く → ら、ら → BOS という次文字予測を繰り返すのが、学習の本質です。
2. PyTorch なしで動く仕組み
自動微分：Value クラス
この実装の最大の特徴は、自動微分を Value クラスで自作していることです。
各スカラー値に data（値）・grad（勾配）・_children・_local_grads を持たせ、
演算のたびに計算グラフのノードを作ります。
def __mul__(self, other):
# z = x * y → ∂z/∂x = y, ∂z/∂y = x
return Value(self.data * other.data, (self, other), (other.data, self.data))
backward() ではトポロジカルソートで逆向きにたどり、連鎖律で勾配を伝播します。
for v in reversed(topo):
for child, local_grad in zip(v._children, v._local_grads):
child.grad += local_grad * v.grad
PyTorch の .backward() が内部でやっていることを、最小限のコードで確認できます。
重みの初期化
重みは matrix() でガウス分布から初期化し、state_dict に名前付きで格納します。
state_dict = {
'wte': matrix(vocab_size, n_embd), # Token Embedding
'wpe': matrix(block_size, n_embd), # Position Embedding
'lm_head': matrix(vocab_size, n_embd), # 出力層
}
# 各層に Attention (wq/wk/wv/wo) と MLP (mlp_fc1/mlp_fc2) を追加
全パラメータを params にフラット展開しておくことで、optimizer からの一括更新が容易になります。
3. 学習の一本の流れ
順伝播：gpt() 関数
gpt(token_id, pos_id, keys, values) が GPT 本体です。
-
wte[token_id] + wpe[pos_id]でトークン埋め込みと位置埋め込みを合成 - RMSNorm → Multi-Head Self-Attention → 残差接続
- RMSNorm → MLP（
n_embd → 4×n_embd → n_embd）→ 残差接続 -
lm_headでvocab_size次元の logits を出力
正規化は LayerNorm ではなく、実装が簡潔な RMSNorm を使っています。
def rmsnorm(x):
ms = sum(xi * xi for xi in x) / len(x)
scale = (ms + 1e-5) ** -0.5
return [xi * scale for xi in x]
逆伝播と Adam 更新
loss.backward() で各パラメータへ勾配が流れた後、Adam で重みを更新します。
学習率は lr_t = lr * (1 - step / num_steps) で線形減衰させています。
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)
p.grad = 0 # 次ステップのためにリセット
重みの保存
学習後、state_dict から .data（float）だけを取り出して .pkl に保存します。
'state_dict': {
key: [[p.data for p in row] for row in mat]
for key, mat in state_dict.items()
}
Value オブジェクトごと保存しないことで、推論側は軽量な float 行列として扱えます。
4. 重みを読んで名前を生成する
推論スクリプトは .pkl を読み込み、uchars とハイパーパラメータを復元します。
gpt() の構造は学習時と同じですが、こちらは勾配を持たない float のみで計算します。
名前生成は generate_name() で行います。
# BOS を入力して生成開始
logits = gpt(BOS, pos_id, keys_cache, values_cache)
while pos_id < block_size:
probs = softmax([l / temperature for l in logits])
token_id = random.choices(range(vocab_size), weights=probs)
if token_id == BOS:
break # BOS が出たら終了
sample.append(uchars[token_id])
logits = gpt(token_id, pos_id, keys_cache, values_cache)
pos_id += 1
temperature を下げると確率の高い候補に収束し、上げるとランダム性が増します。
KV キャッシュ（keys_cache / values_cache）に過去トークンの K・V を蓄積することで、
各ステップで過去全体を再計算せずに済みます。
まとめ：学習と推論はひとつながり
| 学習 | 推論 | |
|---|---|---|
gpt() の使い方 |
logits → loss → backward | logits → サンプリング |
| 重み | 更新する | 固定（.pkl から読む） |
Value クラス |
必要（勾配計算） | 不要（float のみ） |
学習で最適化していた「次文字予測」を、推論では実行モードに切り替えているだけです。
同じ gpt() 関数がそのままつながっている点に、この実装のシンプルさがあります。
次回は W&B を使った loss・perplexity の可視化と実験管理を見ていきます。
5. この実装の学びどころ
今回の実装のよさは、Transformer を「使う」だけでなく、「何が必要か」をかなり透明な形で見せてくれるところにあります。
特に次の点は、学習用の題材としてかなり分かりやすいです。
- 文字単位のトークナイズから始まっていること。
- 自動微分を
Valueクラスで自作していること。 - Attention、MLP、残差接続、正規化といった GPT の基本部品が小さく整理されていること。
- 学習済み重みを保存し、推論専用コードで再利用していること。
一方で、実用的な大規模モデルと比べると、この実装はかなりミニマルです。
たとえばバイアス項はなく、正規化は RMSNorm に簡略化され、活性化関数も GeLU ではなく ReLU になっています。
ですが、こうした簡略化があるからこそ、「本質的に GPT は何をしているのか」が見えやすくなっています。
いきなりフレームワークの大規模実装に入るよりも、まずこの規模感で学習と推論の一本の流れを掴むのはかなり有効だと感じます。
おわりに
今回は、microgpt_train.py と microgpt_generate.py を通して、文字単位 GPT の学習と推論の流れを見てきました。
学習では「次の文字を当てる」ために重みを更新し、推論ではその重みを使って 1 文字ずつ自己回帰的に名前を生成している、というのが全体の本質です。
次回は、今回のコードにも入っている W&B の処理に注目して、loss・perplexity・learning rate の記録や、実験条件の管理をどう行うかを整理したいと思います。
モデル本体の理解に加えて、学習の観測や比較の仕組みまで見ると、実験コードとしての見通しもかなり良くなります。