Transformer：逐步追踪一次前向与反向

一个 decoder-only Transformer 可以看作反复交替进行两件事：注意力让当前位置读取允许的历史，前馈网络加工每个位置的特征。残差保留原信息，归一化控制数值尺度，位置编码打破序列的排列对称性。这里逐项对照 labs/llm/model.py，避免把模型理解成一串只需照抄的层名。

形状是第一份规格书

设 $B = 2, T = 5, d = 8, H = 2$ ，则每头维度 $d_{h} = 4$ 。embedding 后 $X : [2, 5, 8]$ ，乘 $W_{Q}, W_{K}, W_{V} : [8, 8]$ ，得到三个 [2,5,8] 张量。先 reshape 成 [2,5,2,4]，再交换时间和头维，才得到 [2,2,5,4]。直接 reshape 成 [2,2,5,4] 会把原有元素分错组；形状相同并不代表语义相同。

S = \frac{Q K^{T}}{\sqrt{d_{h}}} + M, P = softmax (S), O = P V .

$S, P$ 的形状为 [B,H,T,T]，最后一维是被读取的 key 位置；softmax 应沿它归一化。 $O : [B, H, T, d_{h}]$ 经逆变换恢复 [B,T,d]，再乘输出投影 $W_{O}$ 。缩放因子来自内积的方差：若分量独立、方差约一， $d_{h}$ 项的内积方差约为 $d_{h}$ ；除以平方根避免维度越大 softmax 越早饱和。这是初始化时的量级解释，不是训练后分布的严格假设。

因果 mask 是模型定义的一部分

训练时同时计算全部位置，但位置 $i$ 只能读 $j \leq i$ 。因此 $M_{i j} = 0$ 对允许项，其他项为负无穷。允许读取自己，是因为输入 $x_{i}$ 对应目标 $x_{i + 1}$ ；如果把目标放在输入同一位置又允许对角线，就直接泄漏答案。

python

scores = q @ k.swapaxes(-1, -2) / np.sqrt(q.shape[-1])
mask = np.tril(np.ones((t, t), dtype=bool))
scores = np.where(mask, scores, -np.inf)
probabilities = softmax(scores)

本地测试把输入后三个位置改变，检查更早位置的 logits 完全不变。这比观察注意力图更直接。注意 padding mask 与 causal mask 是两个条件；若一整行都被屏蔽，-inf - (-inf) 会变成 NaN，应在数据组织或内核中明确处理。我们的固定长度无 padding 实验，每行至少允许对角线。

RMSNorm、RoPE 与 SwiGLU

RMSNorm 沿特征维度计算

RMSNorm (x) = g ⊙ x / \sqrt{\frac{1}{d} \sum_{i} x_{i}^{2} + ϵ} .

它没有 LayerNorm 的减均值步骤。增益 $g$ 可学习； $ϵ$ 防止接近零向量造成巨大比例。pre-norm block 写成 $R = X + Attention (RMSNorm (X))$ ， $Y = R + FFN (RMSNorm (R))$ 。反传既经过变换分支，也经过残差恒等分支，漏掉其中一条通常能在梯度检查中发现。

RoPE 对 query/key 的每两个分量进行旋转：

[\begin{matrix} q_{2 r}^{'} \\ q_{2 r + 1}^{'} \end{matrix}] = [\begin{matrix} \cos (m θ_{r}) & - \sin (m θ_{r}) \\ \sin (m θ_{r}) & \cos (m θ_{r}) \end{matrix}] [\begin{matrix} q_{2 r} \\ q_{2 r + 1} \end{matrix}], θ_{r} = 10000^{- 2 r / d_{h}} .

位置 $m, n$ 的 query/key 内积包含相对角度 $(n - m) θ_{r}$ 。它不是把一个位置向量直接加到 embedding；value 也不做相同旋转。旋转矩阵正交，反向传播对输入施加逆旋转，因此本地 rope(..., inverse=True) 可用于梯度。这一实现要求每头维度为偶数，配置校验会拒绝不满足条件的模型。

SwiGLU 用两次升维投影形成门控： $u = X W_{1}$ 、 $v = X W_{3}$ ，输出 $(SiLU (u) ⊙ v) W_{2}$ 。若中间宽度 $f = 64, d = 32$ ，三个矩阵共有 $3 d f = 6144$ 参数。 $SiLU (u) = u σ (u)$ ，导数是 $σ (u) + u σ (u) (1 - σ (u))$ 。门控比单一 ReLU 多一个输入依赖分支，也改变相同宽度下的参数和计算成本。

flowchart TB
  X[输入 X] --> N[RMSNorm]
  N --> Q[Q K V 投影与 RoPE]
  Q --> A[因果多头注意力]
  A --> R[输出投影 加 X]
  X --> R
  R --> N2[RMSNorm]
  N2 --> F[SwiGLU]
  F --> Y[加 R 得 Y]
  R --> Y

查看流程图文本

flowchart TB
  X[输入 X] --> N[RMSNorm]
  N --> Q[Q K V 投影与 RoPE]
  Q --> A[因果多头注意力]
  A --> R[输出投影 加 X]
  X --> R
  R --> N2[RMSNorm]
  N2 --> F[SwiGLU]
  F --> Y[加 R 得 Y]
  R --> Y

注意力反向传播不神秘

若输出梯度是 $G_{O}$ ，从矩阵乘法得到 $G_{V} = P^{T} G_{O}$ 、 $G_{P} = G_{O} V^{T}$ 。softmax 的逐行向量雅可比积可写为

G_{S} = P ⊙ (G_{P} - \sum_{j} P_{j} G_{P, j}) .

再得 $G_{Q} = G_{S} K / \sqrt{d_{h}}$ 、 $G_{K} = G_{S}^{T} Q / \sqrt{d_{h}}$ 。不用构造 [T,T,T] 雅可比；这正是框架自动微分实际利用的结构。被 mask 的位置 $P = 0$ ，对应梯度自然为零。attention_backward 实现了这些式子，完整模型测试覆盖两层及全部参数种类。

参数与运行成本

无偏置、输入输出权重不共享时，参数主要是 $2 V d + L (4 d^{2} + 3 d f)$ ，再加每块两个 RMSNorm 的 $2 L d$ 与末尾 $d$ 。本地默认 $V = 28, d = 32, f = 64, L = 1$ 时应使用实际词表大小重新计算；不要先假设固定字符数。注意力分数的存储为 $O (B H T^{2})$ ，投影与 FFN 的计算近似为 $O (B T L d^{2})$ 。增加上下文与增加宽度触及不同瓶颈，这会在系统章继续展开。

自测

阅读本地实现时，可以从 loss 进入 forward，在每个 block 记录 z,q,k,v,joined,n,hidden 的形状，再沿 backward 反向寻找它们的使用位置。缓存不是额外的模型参数，而是本次前向为求导保留的中间结果；batch 改变后就应重新计算。参数量统计只计 params，训练峰值内存还要算这些缓存、梯度和优化器状态。多层网络每层有独立权重；在代码中误把同一数组引用放到多个 block，会无意形成跨层权重共享，有限差分虽可能通过，架构语义却已变化。

没有位置编码，因果 mask 是否已经给出完整位置信息？

mask 确实区分了可读取的集合大小，但不能替代明确、可控制的相对或绝对位置表示。模型仍需从内容与层间结构间接推断顺序；RoPE 直接让 query-key 匹配随相对距离变化。不要由“有三角形 mask”推导出“对所有顺序关系都足够”。

为什么损失下降不能替代因果性测试？

未来信息泄漏往往让训练 loss 降得更快，甚至趋于零。只有构造“改变未来、不改变过去输出”的不变量，才能针对性检测这个错误；还需另测输入目标错位。

来源：官方 A1 固定讲义、Transformer 原论文、RoFormer、RMSNorm。下一章：训练。

Transformer：逐步追踪一次前向与反向 ​

形状是第一份规格书 ​

因果 mask 是模型定义的一部分 ​

RMSNorm、RoPE 与 SwiGLU ​

注意力反向传播不神秘 ​

参数与运行成本 ​

自测 ​