小さな言語モデルをゼロから学習する — nanoGPT 級を MPS で回し、PPL・速度・メモリを自分で測る

Part 1〜4 では、すでにある重みを使って LLM がどう推論するかを覗いてきました。トークン化、attention、KV キャッシュ、サンプリング — どれも「学習済みのモデルをどう動かすか」の話でした。今回はその手前に戻ります。その重みは、そもそもどこから来るのでしょうか。

答えを確かめる一番早い方法は、自分で作ってみることです。nanoGPT 級（約1,080万パラメータ）の小さな Transformer を、TinyShakespeare というテキストでゼロから学習させます。train/val loss の推移・perplexity・所要時間・ピークメモリを、Apple Silicon の MPS で実測します。

この連載の現在地 — 後付けができない場所に来た

• ここまで分かっていること — Part 4 で greedy / temperature / top-p というサンプリングの選び方の違いと、self-consistency（N 回サンプリングして多数決）で正答率とコストがどう動くかを実測しました
• 今回新しく説明できること — 推論と学習の質的な違い、nanoGPT 級モデルの自作、train/val loss・perplexity・所要時間・ピークメモリの実測
• 今回はやらないこと — PHOTON のアーキテクチャの詳細や階層型との比較は Part 6 へ、fine-tuning のような既存モデルの改変は本連載の範囲外です
• PHOTON のどの主張につながるか — 「学習には後付けができない、独立したコストがかかる」という制約こそ、PHOTON が独自アーキテクチャとしてゼロから学習された理由の一端です。詳しい接続は Part 6 で行います

推論と学習の違い — forward だけか、forward + backward + optimizer step か

Part 1〜4 でずっと使ってきたのは 推論 でした。すでに決まった重みに対して、入力を forward に流すだけです。1 トークンごとに 1 回の計算で済みます。

学習 はこれと質的に違います。forward で予測を出したあと、正解との誤差（loss）を計算し、その誤差を逆向きに伝播させて各パラメータの勾配を求め（backward）、その勾配を使って重みそのものを少しずつ書き換えます（optimizer step）。これを何千回、何万回と繰り返します。パラメータが「固定された値」から「更新され続ける変数」に変わる、というのが本質的な違いです。

最小構成の GPT を組む — Part 2/3 の総復習

学習可能な Transformer を組むには、Part 2 で最小実装した self-attention と、position embedding・feed-forward・LayerNorm を組み合わせます。今回は decoder-only の Transformer ブロックを 6 層重ねます。

面白いのは、Part 3 で必須だった KV キャッシュが、学習では出てこない ことです。推論では「1 トークンずつ生成し、過去の計算を使い回す」ためにキャッシュが要りましたが、学習では毎回バッチ全体・系列全体を一度に forward します。キャッシュは推論時専用の最適化であり、学習ループには存在しないのです。

モデルの構成は Karpathy の nanoGPT が示す標準的な config を採用しました¹。

「約10.6M」という見立て（nanoGPT の公式記載）に対し、自分の実装で実際にカウントすると10,795,776でした。層構成やパラメータの数え方（bias の有無など）で数%のずれが出るのは自然なことなので、この記事では実測値をそのまま報告します。

データとトークナイザ — TinyShakespeare を1文字ずつ読む

Part 1 で「token ≠ 文字」ということを、strawberry の分割や日本語の byte fallback を通して確認しました。今回はあえてその裏側、char-level tokenizer（1文字=1トークン）を使います。学習ループそのものに集中したいので、サブワード分割のような複雑さを持ち込みたくないからです。

データは TinyShakespeare²。シェイクスピア作品からの抜粋、約40,000行・1.11MBの英語テキストです。出典は Karpathy の char-rnn リポジトリで、char-level モデルの教育用途として広く使われてきました。ただし正直に書いておくと、ライセンス表記は “More Information Needed” — Hugging Face のデータセットカードでも明確なライセンスは確認できませんでした²。出典と教育目的での利用実績を根拠に採用しています。

このテキストに含まれる文字は65種類（アルファベット大小・記号・改行など）。これがそのまま vocabulary になります。BPE のような複雑な語彙構築は不要で、「文字の集合を数える」だけで tokenizer が完成するのが char-level の強みです。

学習ループを実装する

学習ループの骨格はシンプルです。

1. TinyShakespeare をランダムな位置から block_size（256文字）ずつ切り出し、ミニバッチを作ります
2. モデルに forward させ、次の文字を予測させて cross entropy loss を計算します
3. loss.backward() で勾配を求め、AdamW optimizer で重みを更新します
4. 一定間隔ごとに、学習に使っていない検証データで val loss を測ります

MPS 特有の注意点として、bf16 autocast は不安定という報告があることを Phase 1 の調査で確認しました³。そのため今回は fp32 を既定にし、保険として PYTORCH_ENABLE_MPS_FALLBACK=1（未対応の演算があれば CPU に自動フォールバック）を設定しています。実際の学習実行では、このフォールバックは一度も発火しませんでした（run-meta の mps_fallback_triggered: false）。fp32 だけで最後まで安定して回りきった、ということです。

実測 — train loss は下がり続けるが、val loss はどこかで反転する

Apple M5 Pro 48GB、MPS、fp32で、5,000 iteration の学習を実際に走らせました。所要時間は約41分（2,445秒）です。この回はセッション内で並行して他の処理（コード修正・ファイル操作）も走らせていたため、iteration 3,500〜4,000 の区間だけ他区間の3〜8倍の時間がかかっています。その分を差し引いた「純粋な学習時間」は27分程度と見積もれますが、記事の数値は実測値（他プロセスの影響込み）をそのまま報告します — 実運用でも他の作業と並行して学習を回すことは普通にあるため、これも一つの現実的な計測条件です。

ここに、想定していなかった発見がありました。train loss は最後まで単調に下がり続けます（3.66→0.80）。一方 val loss は iteration 2,250 付近で 1.4791 という最小値をつけたあと、そこから緩やかに反転して上昇に転じます。最終的な 5,000 iteration 時点では 1.5965 まで悪化していました。もし「5,000 iteration 学習した最終モデル」を素朴に採用していたら、実は途中の 2,250 iteration 時点のモデルより明確に悪いものを持ち帰ることになっていました。

これは典型的な過学習（overfitting）の兆候です。モデルが学習データの細部を丸暗記し始め、初めて見るデータへの汎化性能はむしろ落ちていきます。train loss だけを見ていては気づけない現象です。

nanoGPT の公式ベンチマークでは、同じ config を A100 GPU で学習させ、val loss 1.4697 という数字を報告しています¹。自分の MPS 実測での最小値は 1.4791 — 差はわずか0.6%程度で、かなり近い値に着地しました。ただし条件は完全に同一ではありません。公式実装は学習率のスケジューリング(cosine decayなど)を使っていますが、今回の実装は学習率を最後まで固定値のまま回しています。学習率を下げていく仕組みがない分、後半でモデルが訓練データに寄りすぎて反転が起きやすかった可能性があります。GPU と MPS というハードウェアの違いも当然あるので、単純な優劣比較はできません。

perplexity（exp(val_loss)）で見ると、最小値の時点で4.39です。「次の文字を予測するとき、実質的に4.39通りの選択肢から選んでいるのと同じ不確実性」という解釈です。bits-per-char では2.13 bit/文字でした。現代の言語モデルは英語で概ね0.9〜1.1 bit/文字程度まで到達すると言われます。Shannon の推定では、英語の理論的下限は1.0〜1.5 bit/文字程度です。この数字に対し、小さな自作モデルの2.13 bit/文字はまだ理論的下限には届いていません。ただし1,080万パラメータ・数十分の学習という規模を考えれば、妥当な位置に着地したと評価できます。

学習は推論よりさらに高くつく — 所要時間とメモリの実測

Part 1〜4 で扱ってきた推論は、数秒から長くても数分のオーダーでした。今回の学習は 約41分（2,445秒、他プロセスの影響込み。純粋な学習時間の見積もりは27分程度）。同じ「LLM を動かす」という括りでも、桁がまるで違います。

ピークメモリも実測しました。torch.mps には CUDA の max_memory_allocated() に相当する仕組みがないため、current_allocated_memory() を毎 iteration（backward と optimizer.step 直後）でサンプリングし、その最大値を近似値として使っています。今回の実測では約1.27GiB（1,360,550,400バイト）でした。1,080万パラメータのモデル自体は数十MB程度のはずなので、この大部分は勾配・optimizer の状態（AdamW は各パラメータにつき2つの追加状態を持つ）・バッチデータ・中間活性化が占めています。推論時は不要だったこれらの領域が、学習では常につきまとうコストです。

「学習は既存モデルに後付けできない」というのは、単なる仕組み上の制約ではなく、時間とメモリという具体的なコストの制約でもあります。数十分という所要時間は、他のどんな推論の最適化を積み重ねても迂回できない、学習という営みそのものの重さです。

PHOTON への接続 — ゼロから学習するという制約

今回、小さなモデルを1本学習させるだけで数十分・約1.27GiBのメモリを使いました。PHOTON は The Pile 134.2Bトークンという桁違いのデータで、DGX H200というハードウェアを使ってゼロから学習されています⁴。既存の Transformer に「あとから」PHOTON のアーキテクチャを移植することはできません — 学習からやり直す以外に道がないのです。この「学習必須・後付け不可」という制約こそ、PHOTON がなぜ独自のアーキテクチャとして一から作られたのかという問いへの、ささやかな伏線になります。詳しい接続は最終回で行います。

手元に残るもの と 次回予告

• 動くコード — tiny_gpt.py（decoder-only Transformer の自作 + 学習ループ + train/val loss・perplexity・bits-per-char・所要時間・ピークメモリの計測ハーネス）
• 測った数字 — 1,080万パラメータのモデルをTinyShakespeareで学習し、val loss は iteration 2,250で1.4791（PPL 4.39, BPC 2.13）まで下がったのち反転して悪化。所要時間は約41分（他プロセスの影響込み）、ピークメモリは約1.27GiB
• メンタルモデル — train loss は嘘をつかないが、モデルの本当の実力を測るのは val loss。学習には反転する瞬間があり、「長く学習させるほど良い」は必ずしも成り立たない

次回 Part 6（最終回）では、今回作ったモデルと計測ハーネスをそのまま引き継ぎ、階層型のチャンク構造に改造して vanilla との同条件比較を行います。ここまでの6部で積み上げてきた知識を総動員します — 逐次生成のコスト、O(T²) の壁、memory-bound な KV キャッシュ、test-time compute の経済、そして今回の学習コスト。PHOTON の「なぜ GPU 当たり最大475倍が成立し、なぜ品質が落ちるのか」に、自分の手元の数字で決着をつけます。

コードはサンプルコード（companion repo）の part-05 タグで再現できます。clone と各 Part のタグは、companion repo かこの記事末尾のシリーズナビからたどれます。

参考文献