从零开始写LLM，第34b部分——逐个组件从二元模型到GPT-2（JAX实现）
作者以JAX框架从零构建并训练了一个GPT-2 Small模型（参数量与原始论文一致），全程仅凭笔记，不参考任何现有代码。训练耗时37小时15分钟，在测试集上的损失（3.418784）优于等效PyTorch模型（3.538161）和原版GPT-2 Small（3.499677）。文章详细记录了从最简单的"映射到自身"模型开始，逐项添加位置编码、LayerNorm、多头注意力、Transformer块等组件，并验证每一步确实降低了损失的过程。
背景速读
- 这是一篇技术博客的长篇系列（从 2024 年 12 月始）的收官之作，作者 Giles Thomas 记录了自己从头构建并训练一个 GPT-2 Small 级别语言模型的完整过程。
- 作者此前用 PyTorch 按 Sebastian Raschka 的教材实现过一遍模型；本篇改用 **JAX**（Google 推出的高性能数值计算框架，尤其适合自动微分和 GPU/TPU 加速）以及 Flax 生态的 NNX 子库重写，目的是避免无意识模仿 PyTorch 代码，真正“从零构建”。
- 模型的起点是一个“A-to-A”结构（输入什么就输出什么），只包含词嵌入层和输出层，没有任何 Transformer 模块。作者先测试了简单改为“下一个 token 预测”的效果（类似传统 NLP 的 bigram 模型），损失从 10.82 降至约 5.9，作为后续改进的基准。
- 本文逐一代入 LayerNorm（层归一化）、位置嵌入、多头注意力、前馈网络等 GPT-2 组件，每加一块都观察损失是否下降。最终在 RTX 3090 上训练了约 37 小时，测试集损失 3.4188，优于作者此前 PyTorch 版（3.538）甚至优于 OpenAI 的原始 GPT-2 Small（3.4997）。
- 核心看点：展示了一条让每个新增组件都确实提升模型性能的组装路线，并用 JAX 实现了与 PyTorch 相当甚至略好的训练效率（未使用混合精度也更快）。