From bigrams to GPT-2, one component at a time (in Jax)
JAXを使用してGPT-2 Smallサイズの言語モデルをゼロから段階的に構築・訓練する詳細ガイド。バイグラムモデルから始め、トークン埋め込み、位置エンコーディング、マルチヘッド自己注意機構、順伝播層、層正規化を順に追加し、完全なTransformerを実現する。
背景メモ
- 本記事は、ゼロから大規模言語モデル（LLM）を構築・訓練する連載企画の一部で、今回が最終回。著者のGiles Thomasは英AIスタートアップ・Python Anywhereの創業者兼CTO。
- 記事では、GPT-2 Small（2019年にOpenAIが公開したTransformerベースの言語モデル）を、ディープラーニングフレームワークJAXを用いてフルスクラッチで実装している。
- JAXはGoogleが開発したGPU/TPU向け自動微分ライブラリで、研究コミュニティではPyTorchよりも高速な実行が可能な選択肢として認知されている。
- 「Bigram（バイグラム）」とは単純な2単語の連接確率モデルのことで、連載はこの極めて単純なモデルから出発し、自己注意機構や位置エンコーディングなどTransformerの各構成要素を一つずつ追加しながらGPT-2を再現していく教育的アプローチをとっている。
- この連載が価値を持つ背景として、2020年代半ばのLLMブームにおいて、実際にモデルを内部から理解しているエンジニアが不足しているという問題意識がある。