学习笔记：《从零构建大模型》Q&A • 杨鸿肇

本笔记整理了阅读《从零构建大模型》教材和代码时产生的思考，并以问答形式，按照从简单到复杂、从技术到理念的顺序展开。

主要用于复习和加深理解，推荐快速浏览问题，有困惑的看一下回答，不建议全文顺序阅读。

学习资料

资料来源（GitHub 仓库）：

LLMs-from-scratch-CN：Sebastian Raschka 的原教学代码仓库（LLMs-from-scratch）的中文翻译版，包含中文 Markdown 笔记和带有中文注释的 Jupyter 代码；
Build-A-Large-Language-Model-CN：《Build a Large Language Model (From Scratch)》教材的中文翻译。

教材按照数据集处理、注意力机制、模型架构、预训练和后训练的顺序，带领读者从零开始搭建一个可以在本地运行的 GPT-2 风格教学模型。

图片来源：《Build a Large Language Model (From Scratch)》第 4 章。

一、神经网络基础

Q1. 什么是前向传播和反向传播？

前向传播（forward pass） 使用当前参数计算预测结果和损失（loss）：

x \rightarrow \hat{y} \rightarrow L

反向传播（backward pass） 从损失开始，利用链式法则计算各参数的梯度：

L \rightarrow \frac{\partial L}{\partial \theta}

最后由优化器更新参数：

\theta \leftarrow \theta - \eta \nabla_{\theta}L

其中， $\eta$ 是学习率。完整训练步骤可以概括为：前向传播 → 计算损失 → 反向传播 → 更新参数。

Q2. 神经网络中的梯度是数值近似出来的吗？

通常不是。PyTorch 等现代深度学习框架主要使用**自动微分（automatic differentiation）**计算梯度。用户定义前向传播时，框架会记录参与计算的张量操作并构建计算图；执行 backward() 时，再按照各算子的局部求导规则和链式法则反向累积梯度。

以中心有限差分为例：

\frac{\partial L}{\partial w} \approx\frac{L(w+\varepsilon)-L(w-\varepsilon)}{2\varepsilon}

它通常需要针对每个待检查参数额外执行前向计算，成本随参数数量快速增长，因此不适合大规模训练，主要用于梯度检查（gradient checking），以验证反向传播的实现。

训练中更常见的“近似”是用小批量（mini-batch）梯度估计整个数据集上的梯度，而不是用有限差分近似每个导数。这也是随机梯度下降（stochastic gradient descent，SGD）中“随机”的主要来源。

Q3. 为什么一次 `backward()` 就能得到所有参数的梯度？

反向传播会从后向前遍历计算图，并复用已经计算出的中间梯度。

例如：

a=w_1x, \qquad y=w_2a

已知 $\frac{\partial L}{\partial y}$ 后，可以计算 $\frac{\partial L}{\partial w_2}$ 和继续向前传播的 $\frac{\partial L}{\partial a}$ ，再由 $\frac{\partial L}{\partial a}$ 计算 $\frac{\partial L}{\partial w_1}$ 。

其本质是：链式法则 + 中间结果复用。因此，无须为每个参数分别重新运行一次网络。

Q4. 反向传播会因为层数增加而不断积累“求导误差”吗？

通常不会，因为自动微分使用的是基本算子的局部解析导数。当然，浮点运算仍然存在舍入误差，但深层网络更突出的困难是梯度在链式法则中连续相乘后的尺度变化。

如果许多局部导数的绝对值小于 $1$ ，例如 $0.5^{50}\rightarrow 0$ ，就容易产生梯度消失；如果乘积不断放大，例如 $1.5^{50}\rightarrow\infty$ ，则容易产生梯度爆炸。

因此，关键问题是梯度尺度的稳定性，而不是逐层使用有限差分导致的近似误差。

Q5. 深层网络如何缓解梯度消失和梯度爆炸？

常见手段包括：

使用合理的参数初始化和激活函数，如 ReLU、GELU、SiLU；
使用残差连接（residual connection）；
使用 LayerNorm 或 RMSNorm 等归一化方法；
在需要时使用梯度裁剪（gradient clipping）抑制梯度爆炸。

其中，残差连接把子层写成：

y=x+F(x)

对输入 $x$ 求导：

\frac{\partial y}{\partial x}=I+J_F(x)

其中， $I$ 是恒等矩阵， $J_F(x)$ 是 $F$ 对 $x$ 的雅可比（Jacobian）矩阵。于是：

\frac{\partial L}{\partial x} =\frac{\partial L}{\partial y}\left[I+J_F(x)\right]

与没有残差连接时只经过 $J_F(x)$ 相比， $I$ 提供了一条不必穿过子层内部变换的直接梯度路径。

对于连续的残差块 $x_{l+1}=x_l+F_l(x_l)$ ，梯度传播可写成：

\frac{\partial x_L}{\partial x_l} =\prod_{k=l}^{L-1}\left[I+J_{F_k}(x_k)\right]

每一项都包含恒等映射，因此梯度不再只能依赖多个 Jacobian 的纯乘积。这能显著缓解梯度消失，使深层网络更容易优化，但并不意味着梯度绝不会消失或爆炸：不同路径仍可能缩放或相互抵消。

在教材采用的 Pre-LayerNorm Transformer 中，注意力子层和前馈子层分别使用一次残差连接：

\begin{aligned} u_l &= x_l + \operatorname{Attention}\!\left(\operatorname{LayerNorm}(x_l)\right), \\ x_{l+1} &= u_l + \operatorname{FFN}\!\left(\operatorname{LayerNorm}(u_l)\right). \end{aligned}

这也解释了为什么各子层的输出维度最终必须回到 emb_dim：只有输入和输出形状一致，才能与残差分支逐元素相加。

Q6. 残差连接为什么不限制模型的表达能力？

残差连接并不保证模型得到相同的训练结果；它会改变参数化方式、优化路径和模型的归纳偏置。它不限制表达能力的直观原因是：若原本希望学习目标映射 $H(x)$ ，残差块可以改写为

y=x+F(x), \qquad F(x)=H(x)-x

也就是说，在输入输出维度兼容且残差分支容量足够时，模型仍然可以表示原目标映射，只是改为学习相对于输入的变化量 $F(x)$ 。当目标映射接近恒等映射时，这种参数化通常更容易优化；但它不意味着最终参数、损失或预测结果一定与无残差网络相同。

Q7. 所有神经网络操作都有解析梯度吗？

训练路径中的大多数基本操作都有明确的局部导数，例如：

加、减；
矩阵乘法；
Softmax；
GELU；
LayerNorm；
Attention。

“能够参与自动微分”也不意味着函数必须处处可导。常见的特殊情况包括：

某些分段函数只在少数点不可导，框架会按照约定选取一个可用的次梯度；
argmax 和离散采样无法直接使用普通梯度；
某些随机采样可以借助重参数化等技巧训练。

Q8. 反向传播中的局部梯度计算可以用 GPU 加速吗？

可以。反向传播仍主要由大规模矩阵乘法、逐元素运算（element-wise operation）和归约（reduction）组成，非常适合使用 GPU 和 Tensor Cores 并行计算。

例如，在线性层 $Y=XW$ 中，令 $G=\frac{\partial L}{\partial Y}$ ，则：

\begin{aligned} \frac{\partial L}{\partial X} &= GW^{\mathsf T}, \\ \frac{\partial L}{\partial W} &= X^{\mathsf T}G. \end{aligned}

对于非线性算子，以 ReLU 和 Softmax 为例，也有对应的加速方法。ReLU 是逐元素运算：

\operatorname{ReLU}(x)=\max(0,x), \qquad \frac{\partial\operatorname{ReLU}(x)}{\partial x} =\begin{cases} 0, & x<0, \\ 1, & x>0. \end{cases}

在 $x=0$ 处，框架会采用约定的次梯度。反向传播时，只需将上游梯度逐元素乘以该局部梯度。

Softmax 会把一组分数归一化为概率：

s_i=\frac{e^{z_i}}{\sum_j e^{z_j}}, \qquad \frac{\partial s_i}{\partial z_j}=s_i\left(\delta_{ij}-s_j\right)

其中， $\delta_{ij}$ 是 Kronecker delta。Softmax 的各输出彼此关联，但其前向归约和反向梯度仍可由并行 GPU 内核高效完成。

二、RNN 与 Transformer

Q9. RNN 的核心原理是什么？

循环神经网络（recurrent neural network，RNN）使用不断更新的隐藏状态保存历史信息：

h_t=f\!\left(W_xx_t+W_hh_{t-1}+b\right)

因此， $h_t$ 可以看作对截至时刻 $t$ 的历史序列进行压缩后得到的状态。其核心特征是：历史信息不断压缩进一个固定维度的隐藏状态。

Q10. RNN 为什么难以处理长距离依赖？

RNN 中的信息和梯度都需要沿着时间步逐步传递：

h_1 \rightarrow h_2 \rightarrow \cdots \rightarrow h_T

反向传播涉及多个状态转移雅可比矩阵的连乘：

\frac{\partial L}{\partial h_1} =\frac{\partial L}{\partial h_T} \prod_t\frac{\partial h_t}{\partial h_{t-1}}

路径越长，越容易出现梯度消失或梯度爆炸。LSTM 和 GRU 的主要目的之一，就是改善长期信息和梯度的传播。

Q11. LSTM 和 GRU 如何改善长距离依赖？

长短期记忆网络（long short-term memory，LSTM）引入独立的细胞状态和门控机制。其核心状态更新可以简化为：

c_t=f_t\odot c_{t-1}+i_t\odot\tilde{c}_t

其中，遗忘门 $f_t$ 控制保留多少旧信息，输入门 $i_t$ 控制写入多少新信息。加法更新为信息和梯度提供了比普通 RNN 更直接的时间路径。

门控循环单元（gated recurrent unit，GRU）将状态结构进一步简化。若用 $z_t$ 表示保留旧状态的比例，其更新可写成：

h_t=z_t\odot h_{t-1}+(1-z_t)\odot\tilde{h}_t

更新门在旧状态与候选状态之间进行插值。LSTM 和 GRU 不能彻底消除长序列上的梯度问题，但能让模型更有选择地保留、遗忘和更新信息。

Q12. Transformer 与 RNN 最本质的区别是什么？

RNN：把过去的信息逐步压缩进隐藏状态；
Transformer：通过 Attention，让每个 token 直接聚合可见范围内其他 token 的信息。

RNN 中， $A$ 的信息到达 $E$ 通常要逐步经过中间状态：

1
A → B → C → D → E

在 Transformer 的一个注意力层中， $E$ 可以直接读取 $A$ 、 $B$ 、 $C$ 、 $D$ 的信息：

1
A ─────────────→ E
2
B ─────────────→ E
3
C ─────────────→ E
4
D ─────────────→ E

对于 GPT 使用的因果注意力（causal attention），每个位置只能访问自己及之前的 token，不能访问未来 token。

Q13. 为什么 Transformer 更适合 GPU？

RNN 存在时间步依赖：

h_t=f(h_{t-1},x_t)

必须先计算 $h_{t-1}$ ，才能计算 $h_t$ 。Transformer 训练时则可以把整个序列组成矩阵 $X$ ，并在同一层内同时计算：

\begin{aligned} Q &= XW_Q, & K &= XW_K, & V &= XW_V. \end{aligned}

大量工作由规则的矩阵乘法构成，因此更适合并行计算。不同 Transformer 层之间依然要顺序执行，并非整个网络可以一次并行完成。

2017 年提出的原始 Transformer 采用 Encoder-Decoder 架构。Encoder 先通过双向 Self-Attention 编码完整输入；Decoder 则先对已经生成的目标序列（训练时为右移后的目标序列）执行 Causal Self-Attention，再执行 Cross-Attention。在 Cross-Attention 中，Query 来自 Decoder 的中间表示，Key 和 Value 来自 Encoder 的输出。直观上，前者回答“此前生成了什么”，后者回答“应该从输入中获取什么”。

GPT 的 Decoder-only 架构没有独立的 Encoder，Q、K、V 都由同一序列的表示投影得到，因此不包含 Encoder-Decoder Cross-Attention，主要使用 Causal Self-Attention 和前馈网络逐层处理上下文。

Q15. GPT 也能泛化出翻译能力，Encoder-Decoder 架构还有什么意义？

能完成某项任务，并不意味着它一定采用最适合该任务结构的架构。Encoder-Decoder 明确区分“源输入”和“目标输出”：Encoder 可以双向读取完整输入并生成一次可复用的表示，Decoder 再通过 Cross-Attention 按需读取这些表示并自回归生成结果。

对于自动语音识别（automatic speech recognition，ASR）这类输入与输出模态不同的任务，Encoder-Decoder 是很自然的选择，例如 Whisper 系列模型。对于翻译、摘要等具有明确输入和输出边界的文本任务，它也仍有很强的竞争力。Decoder-only 模型更通用、接口更统一，但 Encoder-Decoder 在特定任务的数据利用、计算组织和结构约束方面仍可能更合适。

四、Embedding

Q16. Embedding 的作用是什么？

Embedding 将离散的 token ID 映射为连续向量。设：

E\in\mathbb{R}^{N\times D}

其中：

$N$ ：词表大小；
$D$ ：隐藏维度（hidden dimension），也称嵌入维度（embedding dimension）。

对于 token $i$ ：

x_i=E[i], \qquad x_i\in\mathbb{R}^D

Embedding 为模型提供 token 的初始表示。

Q17. Embedding 如何学习 token 之间的语义关系？

Embedding 通常没有单独的“语义相似度损失”。它与整个模型一起通过下一个 token 预测损失训练：

L=-\log P\!\left(x_{t+1}\mid x_1,\ldots,x_t\right)

梯度会沿计算图传回 Embedding 矩阵 $E$ 。因此，Embedding 中的语义结构是模型长期优化语言建模目标后形成的结果。

Q18. Embedding 中的相关性和 Attention 中的相关性重复吗？

不重复。

Embedding：token 的静态初始表示；
Attention：当前上下文中的动态信息交互。

例如，bank 兼具“银行”和“河岸”的含义，其 token embedding 会同时承载与 river、ATM 等词相关的语义。在不同句子中，相同 token ID 对应的 token embedding 不变；经过 Attention 后，bank of the river 和 deposit money in the bank 中的 bank 才会形成不同的上下文表示。

Q19. Embedding 和 Attention 的参数是一起训练的吗？

通常是。一次反向传播会同时得到：

\frac{\partial L}{\partial E},\quad \frac{\partial L}{\partial W_Q},\quad \frac{\partial L}{\partial W_K},\quad \frac{\partial L}{\partial W_V},\quad\ldots

因此，Embedding、Attention 和 MLP 等模块通常通过同一个语言建模目标进行端到端优化。

五、Attention：Q、K、V

Q20. Attention 的核心公式是什么？

首先对输入 $X$ 进行三个线性投影：

\begin{aligned} Q &= XW_Q, \\ K &= XW_K, \\ V &= XW_V. \end{aligned}

缩放点积注意力（scaled dot-product attention）为：

A=\operatorname{softmax}\!\left(\frac{QK^{\mathsf T}}{\sqrt{d_k}}+M\right), \qquad O=AV

其中， $d_k$ 是 Key 的维度， $M$ 是可选的掩码。GPT 使用因果掩码把未来位置对应的分数设为负无穷，使 Softmax 后的权重为 $0$ ：

M = \begin{pmatrix} 0 & -\infty & -\infty & \cdots & -\infty \\ 0 & 0 & -\infty & \cdots & -\infty \\ 0 & 0 & 0 & \cdots & -\infty \\ \vdots & \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & 0 & \cdots & 0 \end{pmatrix}

Q21. Q、K、V 分别代表什么？

可以用下面的方式建立直觉：

Query（Q）：我想找什么；
Key（K）：我包含什么可供匹配的特征；
Value（V）：匹配后，我能提供什么信息。

因此，Attention 可以概括为：查询 → 匹配 → 信息聚合。

Q22. 为什么 Q 和 K 使用不同的投影矩阵？

因为信息检索关系通常具有方向性。分数：

s(a,b)=q_a k_b^{\mathsf T}

表示 token $a$ 是否应该从 token $b$ 获取信息，它没有理由等于 $s(b,a)$ 。

如果强制 $W_Q=W_K$ ，那么加入掩码和 Softmax 之前的原始分数矩阵会变为 $QQ^{\mathsf T}$ ，受到对称性约束。分别学习 $W_Q$ 和 $W_K$ ，可以让模型表达更灵活的有向匹配关系，例如当前 token 检索更早出现的实体。

因此，Attention 不只是固定空间中的相似度计算，而是可学习、可定向的信息匹配。

Q23. 为什么需要 V？为什么不能直接对原始 Embedding 加权？

理论上可以直接计算 $O=AX$ ，但这相当于找到某个 token 后，直接聚合它的当前表示。

加入 $V=XW_V$ 后，模型可以进一步学习“应该从这个 token 的表示中提取什么”。因此：

Q、K：决定从哪里获取信息；
V：决定获取和传递什么信息。

在 Multi-Head Attention 中，不同 Head 还可以从同一输入表示中提取不同类型的信息。

六、Multi-Head Attention

Q24. Multi-Head Attention 与单头 Attention 的核心区别是什么？

单头 Attention 只有一张注意力权重矩阵 $A$ ；Multi-Head Attention 则有多张独立计算的权重矩阵：

A^{(1)},A^{(2)},\ldots,A^{(H)}

因此，同一个 token 可以同时使用多套信息路由规则。可以粗略想象成：

1
Head 1 → 关注主谓关系
2
Head 2 → 关注实体
3
Head 3 → 关注时间
4
Head 4 → 关注远距离依赖

这只是直觉示例。不同 Head 确实可以学习不同模式，但并不保证形成如此明确、稳定或可解释的功能分工。

Q25. 增加 Head 数量会增加参数量吗？

如果总隐藏维度 $D$ 不变，标准 Multi-Head Attention 的参数量基本不会随 Head 数量增加。

因为单个 Head 的维度是：

d_h=\frac{D}{H}

所有 Head 合起来仍对应总宽度为 $D$ 的 Q、K、V 投影，再加一个输出投影。忽略 bias 后，其参数量约为：

3D^2+D^2=4D^2

所以，保持 $D$ 不变时，增加 Head 更像是重新划分固定的表示容量，而不是扩大总参数量。

Q26. 增加隐藏维度 D 和增加 Head 数 H 有什么区别？

二者控制的对象不同：增加 $D$ 会扩大总表示空间，并使许多权重矩阵的参数量按 $D^2$ 增长；增加 $H$ 且保持 $D$ 不变，则会把相同的总宽度划分成更多投影子空间，产生更多独立的注意力图，同时减小单个 Head 的维度。

可以记成：

$D$ ：总表示容量；
$H$ ：并行路由数量；
$d_h$ ：单个路由的子空间维度。

Q27. 为什么更宽的单头不能完全替代多头？

一个 Head 无论多宽，在每个 Query 位置上仍然只产生一套 Softmax 注意力分布：

A_{i1},A_{i2},\ldots,A_{it}

多个 Head 则分别产生多套分布 $A^{(1)},A^{(2)},\ldots$ ，可以同时使用不同方式聚合不同位置的信息。

因此，核心差异不只是向量维度，还包括独立注意力图（attention map）的数量。

Q28. 为什么不同 Head 会学到不同信息？

训练时并没有显式指定“Head 1 学语法，Head 2 学实体”。不同 Head 可能产生差异，主要因为：

每个 Head 对应不同的投影参数子空间；
随机初始化打破了初始对称性；
每个 Head 独立计算注意力图；
所有 Head 通过最终损失共同接受梯度更新。

如果某个 Head 学到了其他 Head 没有提供、但有助于降低损失的信息，这种差异就可能在训练中得到强化。因此，Multi-Head Attention 提供了形成分工的结构条件，但不保证 Head 之间严格分工。

Q29. 不同 Head 是直接负责 Embedding 的不同维度吗？

不是简单地把原始 Embedding 坐标固定切成几段，例如：

1
Head 1 → 原始表示的第 1～64 维
2
Head 2 → 原始表示的第 65～128 维

每个 Head 的投影都可以写成 $q_h=xW_Q^{(h)}$ 。也就是说，每个 Head 都能从完整的 $D$ 维输入中学习不同的线性组合，之后才得到维度为 $d_h$ 的 Query、Key 和 Value。

因此，不同 Head 学习的是不同的投影子空间，而不是固定承包原始 Embedding 的某一段坐标。

Q30. 实际中如何选择 Head 数量？

通常会同时考虑隐藏维度 $D$ 、Head 数量 $H$ 和单 Head 维度 $d_h=D/H$ ，而不是孤立地选择 $H$ 。许多常见架构把 $d_h$ 设为 $64$ 或 $128$ ，但这只是常见设计，不是固定规则。

例如，教材中的 GPT-2 124M 配置使用 $D=768$ 、 $H=12$ ，所以 $d_h=64$ 。

现代模型还可能使用分组查询注意力（grouped-query attention，GQA），让多个 Query Head 共享较少的 Key/Value Head，以降低 KV cache 的显存占用和内存带宽成本。

七、Transformer 参数结构

Q31. 词表大小为 N、隐藏维度为 D 时，主要参数矩阵有哪些？

忽略 bias 时，经典 GPT 风格 Transformer 的主要参数矩阵包括：

Token Embedding： $E\in\mathbb{R}^{N\times D}$ ；
Attention： $W_Q,W_K,W_V,W_O\in\mathbb{R}^{D\times D}$ ；
经典前馈网络： $W_1\in\mathbb{R}^{D\times 4D}$ 、 $W_2\in\mathbb{R}^{4D\times D}$ ；
LM Head：按数学上的右乘记法， $W_{\mathrm{out}}\in\mathbb{R}^{D\times N}$ 。

如果使用权重绑定（weight tying）：

W_{\mathrm{out}}=E^{\mathsf T}

则输入 Embedding 和输出投影可以共享参数。需要注意，PyTorch 的 nn.Linear(D, N) 在内存中把权重存为 $N\times D$ ，这与上述右乘记法只是矩阵方向约定不同。

Q32. 一个经典 Transformer Block 大约有多少参数？

忽略 bias 和归一化层参数，并假设前馈网络的中间维度为 $4D$ ：

Attention：约 $4D^2$ ；
前馈网络：约 $D\times 4D+4D\times D=8D^2$ 。

所以每层约有：

4D^2+8D^2=12D^2

如果模型有 $L$ 层，并且输入、输出 Embedding 共享参数，则总参数量可粗略估算为：

P\approx ND+12LD^2

该公式忽略了位置 Embedding、bias、归一化参数等较小项，也不适用于采用门控 MLP、不同扩展比例或其他注意力结构的模型。

根据参数量和每个参数的存储位数，可以进一步估算仅包含模型权重时的大小。设参数量为 $P$ ，每个参数占 $b$ 位（bit），则：

S_{\mathrm{weights}}\approx\frac{Pb}{8}\ \text{bytes}, \qquad S_{\mathrm{GiB}}\approx\frac{Pb}{8\times 2^{30}}

例如，一个 10 亿（1B）参数模型仅计算权重时，FP32 约占 $4$ GB，FP16 或 BF16 约占 $2$ GB，INT8 约占 $1$ GB，INT4 约占 $0.5$ GB。这里的 GB 使用十进制近似，实际文件还可能包含量化比例、分组元数据和文件头。

这只是权重大小，不能直接代表训练或推理所需的全部内存。训练还需要梯度、优化器状态和激活值；推理还需要激活值、KV cache，以及运行时工作区。

八、Dropout

Q33. Dropout 的作用是什么？

训练时，Dropout 随机将部分激活置零。PyTorch 等框架通常使用 inverted dropout：

h'=\frac{m\odot h}{1-p}, \qquad m_i\sim\operatorname{Bernoulli}(1-p)

其核心目的是防止模型在训练数据上过度依赖某些固定特征或路径，从而降低过拟合风险、改善泛化能力。

Q34. Dropout 在反向传播时会发生什么？

如果某个激活在本轮被 Dropout，即 $m_i=0$ ，那么经过这条分支传回的梯度也为 $0$ 。它在这一轮既不参与后续前向计算，也不通过该分支接收反向梯度。

下一次前向传播时会重新采样 mask。

Q35. 推理时还使用 Dropout 吗？

常规推理时不使用，即训练时开启，推理时关闭。

因为 inverted dropout 已在训练时除以 $1-p$ ，保持激活的期望值不变，所以推理时直接使用完整网络即可。在 PyTorch 中，通常通过 model.train() 和 model.eval() 切换这两种行为。

九、设计理念速记

Q36. Transformer 的 Attention 从本质上是什么？

Attention 不只是“计算 token 相似度”，更准确地说，它是一种可学习的动态信息路由机制。普通线性层的连接方式由固定参数决定，而 Attention 会根据当前输入计算注意力权重，因此同一层面对不同句子、位置和上下文时，可以采用不同的信息聚合路径。

Q37. Embedding、Attention、MLP 各自解决什么问题？

Embedding：token 最初如何表示；
Attention：当前应该从哪些位置获取信息；
MLP：对每个位置聚合到的信息进行怎样的非线性特征变换。

重复多个 Transformer Block 后，表示会经历：初始 token 表示 → 上下文化表示 → 更适合预测下一个 token 的内部表示。

Q38. 为什么前馈网络先升维，经过激活函数后再降维？

首先要澄清一个容易混淆的概念：教材代码中，先从 emb_dim 升到 $4\times\texttt{emb\_dim}$ 、经过 GELU、再降回 emb_dim 的是前馈网络（feed-forward network，FFN，也常称 MLP）子层，并非 LayerNorm 所做的变换。

代码对应的变换是：

\operatorname{FFN}(x) =\operatorname{GELU}(xW_1+b_1)W_2+b_2

其中， $W_1$ 把每个 token 的表示从 $D$ 维映射到 $4D$ 维， $W_2$ 再映射回 $D$ 维。这样设计主要有三个原因：

升维提高中间表示容量：模型可以在更宽的特征空间中组合和分解信息；
激活函数引入非线性：如果去掉 GELU，两次线性变换可以合并为一次线性变换，深度和升维的意义都会大幅减弱；
降维恢复残差流的形状：输出回到 $D$ 维后，才能与输入相加，并继续堆叠后续 Transformer Block。

$4D$ 是原始 Transformer 和 GPT-2 采用的经典扩展比例，是架构设计选择，而不是由数学推导得到的唯一最优值。现代模型也会使用其他比例或门控前馈结构。

LayerNorm 则保持输入形状不变，只沿特征维计算均值和方差，并使用可学习的缩放与偏移参数：

\operatorname{LayerNorm}(x) =\gamma\odot\frac{x-\mu}{\sqrt{\sigma^2+\varepsilon}}+\beta

附录 A：从教学架构到现代 LLM

Q39. 这套较早的 GPT-2 风格知识在今天还有价值吗？

有，而且仍然具有很高的学习价值，但需要区分基础原理和具体实现：

作为理解现代 LLM 的基础骨架，它依然有效；
作为当前生产模型的完整架构说明，它已经不够全面。

原书选择 GPT-2 风格的小型模型，是为了用较少的代码完整展示数据处理、Embedding、Attention、Transformer Block、预训练和微调。现代模型虽然替换或扩展了许多组件，但大部分变化仍然围绕这些基本对象展开。

例如，不理解标准 Multi-Head Attention，就很难准确理解 GQA、MLA 和稀疏注意力究竟压缩或省略了什么；不理解稠密前馈网络，也很难理解 MoE 为什么只激活部分专家。

以下内容没有因为模型架构更新而失去价值：

自动微分、反向传播和优化器更新参数的基本过程；
梯度消失、梯度爆炸以及残差连接提供的直接梯度路径；
Tokenization、token embedding 和上下文化表示；
Q、K、V、缩放点积、Softmax 和因果掩码；
Decoder-only 模型的自回归生成过程；
Attention 负责跨位置聚合信息，FFN 负责逐位置非线性变换；
训练损失、验证损失、采样温度和 top-k 等训练与生成概念。

这些知识解释的是“模型为什么能够工作”。现代架构的改进，更多是在回答“怎样以更低成本、更长上下文和更大容量工作”。

哪些具体实现已经发生变化？

教学模型中的实现	现代模型中的常见演进	不变的核心问题
绝对位置 Embedding	RoPE 及其长上下文扩展方法	如何向模型注入 token 的顺序和相对位置信息
标准 Multi-Head Attention（MHA）	MQA、GQA、MLA、局部或稀疏注意力	如何根据上下文动态检索和聚合信息
所有 Head 分别保存 K/V	多个 Query Head 共享或压缩 K/V	如何降低 KV cache 的显存和带宽成本
LayerNorm	RMSNorm，以及不同的 Pre-Norm/Post-Norm 组合	如何稳定激活值和深层网络训练
GELU + $4D$ 前馈网络	SwiGLU、GeGLU、其他扩展比例	如何增加逐 token 非线性变换的容量
每层使用稠密 FFN	Mixture-of-Experts（MoE）	如何在控制单 token 计算量的同时扩大参数容量
全局稠密 Attention	局部—全局混合、稀疏或线性/递归式 Attention	如何降低长序列的计算量和缓存成本
Dropout 是主要正则化手段之一	是否启用及比例取决于训练规模和配方	如何抑制过拟合并维持训练稳定性
$P\approx ND+12LD^2$	需要单独计算 GQA、MLA、MoE 等组件	如何估算参数量、激活量和实际计算成本
单一的下一个 token 预测目标	多 token 预测、指令微调、偏好优化和强化学习	如何学习基础能力，并让输出符合任务与人类偏好

表中右侧的技术并不是对基础原理的否定，而是在保留相同计算目标的前提下，改变参数组织、信息路由和缓存方式。

附录 B：第 2～5 章示例代码

1
# Copyright (c) Sebastian Raschka under Apache License 2.0 (see LICENSE.txt).
2
# Source for "Build a Large Language Model From Scratch"
3
#   - https://www.manning.com/books/build-a-large-language-model-from-scratch
4
# Code: https://github.com/rasbt/LLMs-from-scratch
5
#
6
# This file collects all the relevant code that we covered thus far
7
# throughout Chapters 2-6.
8
# This file can be run as a standalone script.
9

10

11
import matplotlib.pyplot as plt
12
from matplotlib.ticker import MaxNLocator
13
import numpy as np
14
import tiktoken
15
import torch
16
import torch.nn as nn
17
from torch.utils.data import Dataset, DataLoader
18

19

20
#####################################
21
# Chapter 2
22
#####################################
23

24

25
class GPTDatasetV1(Dataset):
26
    def __init__(self, txt, tokenizer, max_length, stride):
27
        self.tokenizer = tokenizer
28
        self.input_ids = []
29
        self.target_ids = []
30

31
        # Tokenize the entire text
32
        token_ids = tokenizer.encode(txt, allowed_special={"<|endoftext|>"})
33

34
        # Use a sliding window to chunk the book into overlapping sequences of max_length
35
        for i in range(0, len(token_ids) - max_length, stride):
36
            input_chunk = token_ids[i : i + max_length]
37
            target_chunk = token_ids[i + 1 : i + max_length + 1]
38
            self.input_ids.append(torch.tensor(input_chunk))
39
            self.target_ids.append(torch.tensor(target_chunk))
40

41
    def __len__(self):
42
        return len(self.input_ids)
43

44
    def __getitem__(self, idx):
45
        return self.input_ids[idx], self.target_ids[idx]
46

47

48
def create_dataloader_v1(
49
    txt,
50
    batch_size=4,
51
    max_length=256,
52
    stride=128,
53
    shuffle=True,
54
    drop_last=True,
55
    num_workers=0,
56
):
57
    # Initialize the tokenizer
58
    tokenizer = tiktoken.get_encoding("gpt2")
59

60
    # Create dataset
61
    dataset = GPTDatasetV1(txt, tokenizer, max_length, stride)
62

63
    # Create dataloader
64
    dataloader = DataLoader(
65
        dataset,
66
        batch_size=batch_size,
67
        shuffle=shuffle,
68
        drop_last=drop_last,
69
        num_workers=num_workers,
70
    )
71

72
    return dataloader
73

74

75
#####################################
76
# Chapter 3
77
#####################################
78
class MultiHeadAttention(nn.Module):
79
    def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
80
        super().__init__()
81
        assert d_out % num_heads == 0, "d_out must be divisible by n_heads"
82

83
        self.d_out = d_out
84
        self.num_heads = num_heads
85
        self.head_dim = (
86
            d_out // num_heads
87
        )  # Reduce the projection dim to match desired output dim
88

89
        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
90
        self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
91
        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
92
        self.out_proj = nn.Linear(d_out, d_out)  # Linear layer to combine head outputs
93
        self.dropout = nn.Dropout(dropout)
94
        self.register_buffer(
95
            "mask", torch.triu(torch.ones(context_length, context_length), diagonal=1)
96
        )
97

98
    def forward(self, x):
99
        b, num_tokens, d_in = x.shape
100

101
        keys = self.W_key(x)  # Shape: (b, num_tokens, d_out)
102
        queries = self.W_query(x)
103
        values = self.W_value(x)
104

105
        # We implicitly split the matrix by adding a `num_heads` dimension
106
        # Unroll last dim: (b, num_tokens, d_out) -> (b, num_tokens, num_heads, head_dim)
107
        keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)
108
        values = values.view(b, num_tokens, self.num_heads, self.head_dim)
109
        queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)
110

111
        # Transpose: (b, num_tokens, num_heads, head_dim) -> (b, num_heads, num_tokens, head_dim)
112
        keys = keys.transpose(1, 2)
113
        queries = queries.transpose(1, 2)
114
        values = values.transpose(1, 2)
115

116
        # Compute scaled dot-product attention (aka self-attention) with a causal mask
117
        attn_scores = queries @ keys.transpose(2, 3)  # Dot product for each head
118

119
        # Original mask truncated to the number of tokens and converted to boolean
120
        mask_bool = self.mask.bool()[:num_tokens, :num_tokens]
121

122
        # Use the mask to fill attention scores
123
        attn_scores.masked_fill_(mask_bool, -torch.inf)
124

125
        attn_weights = torch.softmax(attn_scores / keys.shape[-1] ** 0.5, dim=-1)
126
        attn_weights = self.dropout(attn_weights)
127

128
        # Shape: (b, num_tokens, num_heads, head_dim)
129
        context_vec = (attn_weights @ values).transpose(1, 2)
130

131
        # Combine heads, where self.d_out = self.num_heads * self.head_dim
132
        context_vec = context_vec.reshape(b, num_tokens, self.d_out)
133
        context_vec = self.out_proj(context_vec)  # optional projection
134

135
        return context_vec
136

137

138
#####################################
139
# Chapter 4
140
#####################################
141
class LayerNorm(nn.Module):
142
    def __init__(self, emb_dim):
143
        super().__init__()
144
        self.eps = 1e-5
145
        self.scale = nn.Parameter(torch.ones(emb_dim))
146
        self.shift = nn.Parameter(torch.zeros(emb_dim))
147

148
    def forward(self, x):
149
        mean = x.mean(dim=-1, keepdim=True)
150
        var = x.var(dim=-1, keepdim=True, unbiased=False)
151
        norm_x = (x - mean) / torch.sqrt(var + self.eps)
152
        return self.scale * norm_x + self.shift
153

154

155
class GELU(nn.Module):
156
    def __init__(self):
157
        super().__init__()
158

159
    def forward(self, x):
160
        return (
161
            0.5
162
            * x
163
            * (
164
                1
165
                + torch.tanh(
166
                    torch.sqrt(torch.tensor(2.0 / torch.pi))
167
                    * (x + 0.044715 * torch.pow(x, 3))
168
                )
169
            )
170
        )
171

172

173
class FeedForward(nn.Module):
174
    def __init__(self, cfg):
175
        super().__init__()
176
        self.layers = nn.Sequential(
177
            nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
178
            GELU(),
179
            nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
180
        )
181

182
    def forward(self, x):
183
        return self.layers(x)
184

185

186
class TransformerBlock(nn.Module):
187
    def __init__(self, cfg):
188
        super().__init__()
189
        self.att = MultiHeadAttention(
190
            d_in=cfg["emb_dim"],
191
            d_out=cfg["emb_dim"],
192
            context_length=cfg["context_length"],
193
            num_heads=cfg["n_heads"],
194
            dropout=cfg["drop_rate"],
195
            qkv_bias=cfg["qkv_bias"],
196
        )
197
        self.ff = FeedForward(cfg)
198
        self.norm1 = LayerNorm(cfg["emb_dim"])
199
        self.norm2 = LayerNorm(cfg["emb_dim"])
200
        self.drop_resid = nn.Dropout(cfg["drop_rate"])
201

202
    def forward(self, x):
203
        # Shortcut connection for attention block
204
        shortcut = x
205
        x = self.norm1(x)
206
        x = self.att(x)  # Shape [batch_size, num_tokens, emb_size]
207
        x = self.drop_resid(x)
208
        x = x + shortcut  # Add the original input back
209

210
        # Shortcut connection for feed-forward block
211
        shortcut = x
212
        x = self.norm2(x)
213
        x = self.ff(x)
214
        x = self.drop_resid(x)
215
        x = x + shortcut  # Add the original input back
216

217
        return x
218

219

220
class GPTModel(nn.Module):
221
    def __init__(self, cfg):
222
        super().__init__()
223
        self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
224
        self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
225
        self.drop_emb = nn.Dropout(cfg["drop_rate"])
226

227
        self.trf_blocks = nn.Sequential(
228
            *[TransformerBlock(cfg) for _ in range(cfg["n_layers"])]
229
        )
230

231
        self.final_norm = LayerNorm(cfg["emb_dim"])
232
        self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)
233

234
    def forward(self, in_idx):
235
        batch_size, seq_len = in_idx.shape
236
        tok_embeds = self.tok_emb(in_idx)
237
        pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
238
        x = tok_embeds + pos_embeds  # Shape [batch_size, num_tokens, emb_size]
239
        x = self.drop_emb(x)
240
        x = self.trf_blocks(x)
241
        x = self.final_norm(x)
242
        logits = self.out_head(x)
243
        return logits
244

245

246
def generate_text_simple(model, idx, max_new_tokens, context_size):
247
    # idx is (B, T) array of indices in the current context
248
    for _ in range(max_new_tokens):
249
        # Crop current context if it exceeds the supported context size
250
        # E.g., if LLM supports only 5 tokens, and the context size is 10
251
        # then only the last 5 tokens are used as context
252
        idx_cond = idx[:, -context_size:]
253

254
        # Get the predictions
255
        with torch.no_grad():
256
            logits = model(idx_cond)
257

258
        # Focus only on the last time step
259
        # (batch, n_token, vocab_size) becomes (batch, vocab_size)
260
        logits = logits[:, -1, :]
261

262
        # Get the idx of the vocab entry with the highest logits value
263
        idx_next = torch.argmax(logits, dim=-1, keepdim=True)  # (batch, 1)
264

265
        # Append sampled index to the running sequence
266
        idx = torch.cat((idx, idx_next), dim=1)  # (batch, n_tokens+1)
267

268
    return idx
269

270

271
#####################################
272
# Chapter 5
273
#####################################
274
def generate(
275
    model, idx, max_new_tokens, context_size, temperature=0.0, top_k=None, eos_id=None
276
):
277

278
    # For-loop is the same as before: Get logits, and only focus on last time step
279
    for _ in range(max_new_tokens):
280
        idx_cond = idx[:, -context_size:]
281
        with torch.no_grad():
282
            logits = model(idx_cond)
283
        logits = logits[:, -1, :]
284

285
        # New: Filter logits with top_k sampling
286
        if top_k is not None:
287
            # Keep only top_k values
288
            top_logits, _ = torch.topk(logits, top_k)
289
            min_val = top_logits[:, -1]
290
            logits = torch.where(
291
                logits < min_val, torch.tensor(float("-inf")).to(logits.device), logits
292
            )
293

294
        # New: Apply temperature scaling
295
        if temperature > 0.0:
296
            logits = logits / temperature
297

298
            # Apply softmax to get probabilities
299
            probs = torch.softmax(logits, dim=-1)  # (batch_size, context_len)
300

301
            # Sample from the distribution
302
            idx_next = torch.multinomial(probs, num_samples=1)  # (batch_size, 1)
303

304
        # Otherwise same as before: get idx of the vocab entry with the highest logits value
305
        else:
306
            idx_next = torch.argmax(logits, dim=-1, keepdim=True)  # (batch_size, 1)
307

308
        if (
309
            idx_next == eos_id
310
        ):  # Stop generating early if end-of-sequence token is encountered and eos_id is specified
311
            break
312

313
        # Same as before: append sampled index to the running sequence
314
        idx = torch.cat((idx, idx_next), dim=1)  # (batch_size, num_tokens+1)
315

316
    return idx
317

318

319
def train_model_simple(
320
    model,
321
    train_loader,
322
    val_loader,
323
    optimizer,
324
    device,
325
    num_epochs,
326
    eval_freq,
327
    eval_iter,
328
    start_context,
329
    tokenizer,
330
):
331
    # Initialize lists to track losses and tokens seen
332
    train_losses, val_losses, track_tokens_seen = [], [], []
333
    tokens_seen, global_step = 0, -1
334

335
    # Main training loop
336
    for epoch in range(num_epochs):
337
        model.train()  # Set model to training mode
338

339
        for input_batch, target_batch in train_loader:
340
            optimizer.zero_grad()  # Reset loss gradients from previous batch iteration
341
            loss = calc_loss_batch(input_batch, target_batch, model, device)
342
            loss.backward()  # Calculate loss gradients
343
            optimizer.step()  # Update model weights using loss gradients
344
            tokens_seen += input_batch.numel()
345
            global_step += 1
346

347
            # Optional evaluation step
348
            if global_step % eval_freq == 0:
349
                train_loss, val_loss = evaluate_model(
350
                    model, train_loader, val_loader, device, eval_iter
351
                )
352
                train_losses.append(train_loss)
353
                val_losses.append(val_loss)
354
                track_tokens_seen.append(tokens_seen)
355
                print(
356
                    f"Ep {epoch + 1} (Step {global_step:06d}): "
357
                    f"Train loss {train_loss:.3f}, Val loss {val_loss:.3f}"
358
                )
359

360
        # Print a sample text after each epoch
361
        generate_and_print_sample(model, tokenizer, device, start_context)
362

363
    return train_losses, val_losses, track_tokens_seen
364

365

366
def evaluate_model(model, train_loader, val_loader, device, eval_iter):
367
    model.eval()
368
    with torch.no_grad():
369
        train_loss = calc_loss_loader(
370
            train_loader, model, device, num_batches=eval_iter
371
        )
372
        val_loss = calc_loss_loader(val_loader, model, device, num_batches=eval_iter)
373
    model.train()
374
    return train_loss, val_loss
375

376

377
def generate_and_print_sample(model, tokenizer, device, start_context):
378
    model.eval()
379
    context_size = model.pos_emb.weight.shape[0]
380
    encoded = text_to_token_ids(start_context, tokenizer).to(device)
381
    with torch.no_grad():
382
        token_ids = generate_text_simple(
383
            model=model, idx=encoded, max_new_tokens=50, context_size=context_size
384
        )
385
        decoded_text = token_ids_to_text(token_ids, tokenizer)
386
        print(decoded_text.replace("\n", " "))  # Compact print format
387
    model.train()
388

389

390
def assign(left, right):
391
    if left.shape != right.shape:
392
        raise ValueError(f"Shape mismatch. Left: {left.shape}, Right: {right.shape}")
393
    return torch.nn.Parameter(torch.tensor(right))
394

395

396
def load_weights_into_gpt(gpt, params):
397
    gpt.pos_emb.weight = assign(gpt.pos_emb.weight, params["wpe"])
398
    gpt.tok_emb.weight = assign(gpt.tok_emb.weight, params["wte"])
399

400
    for b in range(len(params["blocks"])):
401
        q_w, k_w, v_w = np.split(
402
            (params["blocks"][b]["attn"]["c_attn"])["w"], 3, axis=-1
403
        )
404
        gpt.trf_blocks[b].att.W_query.weight = assign(
405
            gpt.trf_blocks[b].att.W_query.weight, q_w.T
406
        )
407
        gpt.trf_blocks[b].att.W_key.weight = assign(
408
            gpt.trf_blocks[b].att.W_key.weight, k_w.T
409
        )
410
        gpt.trf_blocks[b].att.W_value.weight = assign(
411
            gpt.trf_blocks[b].att.W_value.weight, v_w.T
412
        )
413

414
        q_b, k_b, v_b = np.split(
415
            (params["blocks"][b]["attn"]["c_attn"])["b"], 3, axis=-1
416
        )
417
        gpt.trf_blocks[b].att.W_query.bias = assign(
418
            gpt.trf_blocks[b].att.W_query.bias, q_b
419
        )
420
        gpt.trf_blocks[b].att.W_key.bias = assign(gpt.trf_blocks[b].att.W_key.bias, k_b)
421
        gpt.trf_blocks[b].att.W_value.bias = assign(
422
            gpt.trf_blocks[b].att.W_value.bias, v_b
423
        )
424

425
        gpt.trf_blocks[b].att.out_proj.weight = assign(
426
            gpt.trf_blocks[b].att.out_proj.weight,
427
            params["blocks"][b]["attn"]["c_proj"]["w"].T,
428
        )
429
        gpt.trf_blocks[b].att.out_proj.bias = assign(
430
            gpt.trf_blocks[b].att.out_proj.bias,
431
            params["blocks"][b]["attn"]["c_proj"]["b"],
432
        )
433

434
        gpt.trf_blocks[b].ff.layers[0].weight = assign(
435
            gpt.trf_blocks[b].ff.layers[0].weight,
436
            params["blocks"][b]["mlp"]["c_fc"]["w"].T,
437
        )
438
        gpt.trf_blocks[b].ff.layers[0].bias = assign(
439
            gpt.trf_blocks[b].ff.layers[0].bias, params["blocks"][b]["mlp"]["c_fc"]["b"]
440
        )
441
        gpt.trf_blocks[b].ff.layers[2].weight = assign(
442
            gpt.trf_blocks[b].ff.layers[2].weight,
443
            params["blocks"][b]["mlp"]["c_proj"]["w"].T,
444
        )
445
        gpt.trf_blocks[b].ff.layers[2].bias = assign(
446
            gpt.trf_blocks[b].ff.layers[2].bias,
447
            params["blocks"][b]["mlp"]["c_proj"]["b"],
448
        )
449

450
        gpt.trf_blocks[b].norm1.scale = assign(
451
            gpt.trf_blocks[b].norm1.scale, params["blocks"][b]["ln_1"]["g"]
452
        )
453
        gpt.trf_blocks[b].norm1.shift = assign(
454
            gpt.trf_blocks[b].norm1.shift, params["blocks"][b]["ln_1"]["b"]
455
        )
456
        gpt.trf_blocks[b].norm2.scale = assign(
457
            gpt.trf_blocks[b].norm2.scale, params["blocks"][b]["ln_2"]["g"]
458
        )
459
        gpt.trf_blocks[b].norm2.shift = assign(
460
            gpt.trf_blocks[b].norm2.shift, params["blocks"][b]["ln_2"]["b"]
461
        )
462

463
    gpt.final_norm.scale = assign(gpt.final_norm.scale, params["g"])
464
    gpt.final_norm.shift = assign(gpt.final_norm.shift, params["b"])
465
    gpt.out_head.weight = assign(gpt.out_head.weight, params["wte"])
466

467

468
def text_to_token_ids(text, tokenizer):
469
    encoded = tokenizer.encode(text, allowed_special={"<|endoftext|>"})
470
    encoded_tensor = torch.tensor(encoded).unsqueeze(0)  # add batch dimension
471
    return encoded_tensor
472

473

474
def token_ids_to_text(token_ids, tokenizer):
475
    flat = token_ids.squeeze(0)  # remove batch dimension
476
    return tokenizer.decode(flat.tolist())
477

478

479
def calc_loss_batch(input_batch, target_batch, model, device):
480
    input_batch, target_batch = input_batch.to(device), target_batch.to(device)
481
    logits = model(input_batch)
482
    loss = torch.nn.functional.cross_entropy(
483
        logits.flatten(0, 1), target_batch.flatten()
484
    )
485
    return loss
486

487

488
def calc_loss_loader(data_loader, model, device, num_batches=None):
489
    total_loss = 0.0
490
    if len(data_loader) == 0:
491
        return float("nan")
492
    elif num_batches is None:
493
        num_batches = len(data_loader)
494
    else:
495
        # Reduce the number of batches to match the total number of batches in the data loader
496
        # if num_batches exceeds the number of batches in the data loader
497
        num_batches = min(num_batches, len(data_loader))
498
    for i, (input_batch, target_batch) in enumerate(data_loader):
499
        if i < num_batches:
500
            loss = calc_loss_batch(input_batch, target_batch, model, device)
501
            total_loss += loss.item()
502
        else:
503
            break
504
    return total_loss / num_batches
505

506

507
def plot_losses(epochs_seen, tokens_seen, train_losses, val_losses, label="loss"):
508
    fig, ax1 = plt.subplots(figsize=(5, 3))
509

510
    # Plot training and validation loss against epochs
511
    ax1.plot(epochs_seen, train_losses, label=f"Training {label}")
512
    ax1.plot(epochs_seen, val_losses, linestyle="-.", label=f"Validation {label}")
513
    ax1.set_xlabel("Epochs")
514
    ax1.set_ylabel(label.capitalize())
515
    ax1.legend()
516
    ax1.xaxis.set_major_locator(
517
        MaxNLocator(integer=True)
518
    )  # only show integer labels on x-axis
519

520
    # Create a second x-axis for tokens seen
521
    ax2 = ax1.twiny()  # Create a second x-axis that shares the same y-axis
522
    ax2.plot(tokens_seen, train_losses, alpha=0)  # Invisible plot for aligning ticks
523
    ax2.set_xlabel("Tokens seen")
524

525
    fig.tight_layout()  # Adjust layout to make room
526
    plt.savefig(f"{label}-plot.pdf")
527
    plt.show()