スクラッチからのLLM構築、パート34b――バイグラムからGPT-2へ、ひとつずつコンポーネントを追加（JAX）
本記事は長期連載の集大成。著者はSebastian Raschkaの書籍を元に、ノートのみを頼りにスクラッチからLLMを構築・訓練。PyTorch版のコードを一切参照せず、JAXを使用。当初は入力と同じ系列を出力するだけの「A-to-Aモデル」から出発し、LayerNormやTransformerブロックなどを段階的に追加。最終的にGPT-2 Small相当のモデルをRTX 3090で37時間15分かけて訓練し、損失3.418を達成。これは同等のPyTorchモデル（3.538）や元のGPT-2 small（3.499）を上回る結果となった。
背景メモ
このブログ記事は、GPT-2スモール相当の言語モデルを、PyTorchではなくJAX（Googleが開発した機械学習フレームワーク）でゼロから実装・訓練した記録の集大成。著者はSebastian Raschkaの書籍を参考にしながらも、あえて「本も過去のコードも見ずに」独自実装に挑んでいる。
- **JAX**: 数値計算と機械学習のためのPythonライブラリ。関数型プログラミングに近いスタイルで、自動微分やGPU/TPUでの高速実行が特徴。PyTorchと並ぶ主要フレームワーク。著者は「PyTorchコードを丸写ししていないことを確認するため」JAXを選んだ。
- **Flax / NNX**: JAX上でニューラルネットワークを構築するための高レベルライブラリ。記事内ではNNXのAPI（Linear, Embed, LayerNormなど）が使われている。
- **GPT-2 スモール**: OpenAIが2019年に公開した、パラメータ数約1.24億の比較的小型な言語モデル。現在の大規模LLM（数百億〜数千億パラメータ）の出発点としてよく参照される。
- **LayerNorm / マルチヘッド・アテンション**: Transformerモデルの中核構成要素。LayerNormは訓練の安定化、マルチヘッド・アテンションは文中の単語間の関係を学習する仕組み。
- **Perplexity / Loss**: 言語モデルの性能評価指標。loss（交差エントロピー損失）が低いほど、モデルが次のトークンを正確に予測できていることを示す。著者のモデルは最終的にloss 3.42を達成し、元のOpenAIのGPT-2スモール（3.50）を上回った。
この記事は、まず「入力と同じものを出力するだけのモデル（A-to-A）→ ビッグラム（直前1トークンだけから次のトークンを予測）」という極めてシンプルな出発点からスタートし、LayerNormなどの部品を1つずつ追加して性能を段階的に改善する過程を追っている。