AI Infra学习之旅-手算一次 Transformer 的 Forward

这是"Transformer 原理深度系列"的第 6 篇。
前面 5 篇我们已经把 Transformer 的主要结构讲清楚了：
为什么需要 Transformer，文本如何变成向量，位置编码、Q/K/V、Attention、Multi-Head、Mask、FFN、残差、LayerNorm、Encoder/Decoder，以及训练与推理的区别。
但如果一直停留在结构解释层面，很多人还是会觉得：
“我知道它在干什么，但我没有真正看过它怎么算。”
所以这一篇我们不再扩概念，而是做一件很扎实的事情：
用一个极小的、可以手算的例子，把一次 Transformer 的 forward 从头到尾走一遍。

一、为什么一定要做一次"手算版" Transformer

很多人在学 Transformer 时，会出现一种典型现象：

  • 看公式时觉得懂了
  • 看结构图时也觉得懂了
  • 但一旦真的问：
    • 输入矩阵长什么样？
    • Q/K/V 到底怎么乘出来的？
    • 为什么 \(QK^T\) 是那个形状？
    • softmax 前后数值发生了什么？
    • 最终输出向量到底是怎么来的？

就会突然变得不扎实。

原因很简单：

Transformer 很容易学成"概念图谱"，但不容易自动学成"可计算对象"。

而真正掌握它，必须有一个过程：

从语言解释，走到矩阵计算；
从结构理解，走到数值流动。

所以这一篇的目标不是再讲一遍抽象原理，而是：

  • 选一个非常小的例子
  • 把所有矩阵尺寸压最小
  • 明确写出每一步数值
  • 让你真正"看到"一次 Attention 是怎样算出来的

二、我们这次只做一个最小可解释版本

为了把事情讲清楚，这一篇我们故意做很简化：

我们保留的东西

  • token 向量输入
  • Q/K/V 的线性变换
  • 单头 Self-Attention
  • 缩放
  • softmax
  • 对 V 的加权求和

我们暂时不展开的东西

  • Multi-Head Attention
  • FFN
  • 残差连接
  • LayerNorm
  • RoPE / 复杂位置编码
  • Batch 维度
  • 多层堆叠
  • 词表输出与 sampling

也就是说，我们只聚焦于一件事：

一个单头 Self-Attention 到底怎么从输入变成输出

三、先设定我们的玩具例子

我们假设一个极小的序列长度：

n=3n = 3

也就是一共有 3 个 token。
为了方便，你可以把它们想象成一句极短的话，例如：

  • token 1：A
  • token 2：B
  • token 3：C

我们再假设模型维度很小：

dmodel=2d_{\text{model}} = 2

也就是说，每个 token 只用一个 2 维向量表示。
于是输入矩阵 (H) 写成：

H=[100111]H= \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \end{bmatrix}

这里：

  • 第 1 个 token 的输入向量是 ([1,0])
  • 第 2 个 token 的输入向量是 ([0,1])
  • 第 3 个 token 的输入向量是 ([1,1])

你可以把它看成"已经经过 embedding + 简化位置处理之后的输入"。

四、第一步：定义 Q、K、V 的线性变换矩阵

我们先定义三个投影矩阵。
因为我们这次做极简例子，所以设：

dk=dv=2d_k = d_v = 2

并设：

WQ=[1001]W^Q= \begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix}

WK=[1101]W^K= \begin{bmatrix} 1 & 1 \\ 0 & 1 \end{bmatrix}

WV=[1011]W^V= \begin{bmatrix} 1 & 0 \\ 1 & 1 \end{bmatrix}

这样设计的好处是：

  • 数值不会太复杂
  • 但又不是完全退化成"全是单位矩阵"
  • 足够让你看到 Q/K/V 确实是不同投影

五、第二步：计算 Q、K、V

根据定义：

Q=HWQ,K=HWK,V=HWVQ = HW^Q,\quad K = HW^K,\quad V = HW^V

1. 计算 Q

因为 (W^Q) 是单位矩阵，所以：

Q=HQ = H

因此：

Q=[100111]Q= \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \end{bmatrix}

2. 计算 K

K=[100111][1101]=[110112]K= \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \end{bmatrix} \begin{bmatrix} 1 & 1 \\ 0 & 1 \end{bmatrix} = \begin{bmatrix} 1 & 1 \\ 0 & 1 \\ 1 & 2 \end{bmatrix}

所以：

  • 第 1 个 key 是 ([1,1])
  • 第 2 个 key 是 ([0,1])
  • 第 3 个 key 是 ([1,2])

3. 计算 V

V=[100111][1011]=[101121]V= \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \end{bmatrix} \begin{bmatrix} 1 & 0 \\ 1 & 1 \end{bmatrix} = \begin{bmatrix} 1 & 0 \\ 1 & 1 \\ 2 & 1 \end{bmatrix}

所以：

  • 第 1 个 value 是 ([1,0])
  • 第 2 个 value 是 ([1,1])
  • 第 3 个 value 是 ([2,1])

六、到这里先停一下：我们已经得到了什么

现在我们手里有：

Q=[100111]K=[110112]V=[101121]Q= \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \end{bmatrix} \quad K= \begin{bmatrix} 1 & 1 \\ 0 & 1 \\ 1 & 2 \end{bmatrix} \quad V= \begin{bmatrix} 1 & 0 \\ 1 & 1 \\ 2 & 1 \end{bmatrix}

从直觉上，你可以理解为：

  • Q：每个位置"想找什么"
  • K：每个位置"适合被怎么匹配"
  • V：每个位置"真正携带什么内容"

接下来，Attention 的核心就要开始了。

七、第三步：计算原始相关性分数 (QK^T)

先写出 (K^T)：

KT=[101112]K^T= \begin{bmatrix} 1 & 0 & 1 \\ 1 & 1 & 2 \end{bmatrix}

然后：

QKT=[100111][101112]=[101112213]QK^T = \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \end{bmatrix} \begin{bmatrix} 1 & 0 & 1 \\ 1 & 1 & 2 \end{bmatrix} = \begin{bmatrix} 1 & 0 & 1 \\ 1 & 1 & 2 \\ 2 & 1 & 3 \end{bmatrix}

这就是 score matrix。

八、这个分数矩阵到底意味着什么

我们把它写出来：

S=[101112213]S= \begin{bmatrix} 1 & 0 & 1 \\ 1 & 1 & 2 \\ 2 & 1 & 3 \end{bmatrix}

它是一个 (3\times3) 矩阵。
第 (i,j) 个元素表示：

Sij=qi⋅kjS_{ij}=q_i\cdot k_j

也就是：

第 (i) 个位置，对第 (j) 个位置的原始关注分数。

例如：

第 1 行：([1,0,1])

表示第 1 个 token 对三个 token 的关注分数分别是：

  • 看 token 1：1
  • 看 token 2：0
  • 看 token 3：1

第 2 行：([1,1,2])

表示第 2 个 token 最关注 token 3。

第 3 行：([2,1,3])

表示第 3 个 token 对自身打分最高，其次是 token 1，再是 token 2。

所以到这里，attention 还没有开始"取内容"，它只是构建了一张：

谁该关注谁的原始打分表

九、第四步：做缩放 (\frac{1}{\sqrt{d_k}})

因为：

dk=2d_k = 2

所以：

dk=2≈1.414\sqrt{d_k} = \sqrt{2} \approx 1.414

于是缩放后的分数矩阵为：

S~=S2\tilde{S} = \frac{S}{\sqrt{2}}

近似得到：

S~≈[0.70700.7070.7070.7071.4141.4140.7072.121]\tilde{S} \approx \begin{bmatrix} 0.707 & 0 & 0.707 \\ 0.707 & 0.707 & 1.414 \\ 1.414 & 0.707 & 2.121 \end{bmatrix}

为什么要做这一步？
因为点积会随着维度变大而变大，softmax 对数值尺度很敏感。
所以这里的缩放，本质上是在：

控制 score 的数值范围，防止 softmax 过度尖锐。

十、第五步：对每一行做 softmax

现在我们对每一行分别做 softmax。

第 1 行 softmax

原始行为：

[0.707, 0, 0.707][0.707,\ 0,\ 0.707]

指数值近似为：

[e0.707, e0, e0.707]≈[2.028, 1, 2.028][e^{0.707},\ e^0,\ e^{0.707}] \approx [2.028,\ 1,\ 2.028]

总和约为：

5.0565.056

所以第 1 行权重大约是：

[0.401, 0.198, 0.401][0.401,\ 0.198,\ 0.401]

第 2 行 softmax

原始行为：

[0.707, 0.707, 1.414][0.707,\ 0.707,\ 1.414]

指数值近似为：

[2.028, 2.028, 4.113][2.028,\ 2.028,\ 4.113]

总和：

8.1698.169

所以第 2 行权重大约是：

[0.248, 0.248, 0.504][0.248,\ 0.248,\ 0.504]

第 3 行 softmax

原始行为：

[1.414, 0.707, 2.121][1.414,\ 0.707,\ 2.121]

指数值近似为：

[4.113, 2.028, 8.339][4.113,\ 2.028,\ 8.339]

总和：

14.48014.480

所以第 3 行权重大约是：

[0.284, 0.140, 0.576][0.284,\ 0.140,\ 0.576]

十一、于是我们得到了注意力权重矩阵

A≈[0.4010.1980.4010.2480.2480.5040.2840.1400.576]A \approx \begin{bmatrix} 0.401 & 0.198 & 0.401 \\ 0.248 & 0.248 & 0.504 \\ 0.284 & 0.140 & 0.576 \end{bmatrix}

这是最关键的中间结果之一。

它的含义是：

  • 每一行加起来都等于 1
  • 每一行表示一个 token 如何把自己的注意力分配给全序列所有 token

例如：

第 1 行

第 1 个 token 会：

  • 40.1% 关注 token 1
  • 19.8% 关注 token 2
  • 40.1% 关注 token 3

第 2 行

第 2 个 token 最关注 token 3。

第 3 行

第 3 个 token 对自身关注最多，也会看不少 token 1。

所以 softmax 的本质不是神秘变换，而是：

把原始打分，转成可解释的注意力分配比例。

十二、第六步：用注意力权重去加权聚合 V

现在做最后一步：

O=AVO = AV

回忆一下：

V=[101121]V= \begin{bmatrix} 1 & 0 \\ 1 & 1 \\ 2 & 1 \end{bmatrix}

所以：

O≈[0.4010.1980.4010.2480.2480.5040.2840.1400.576][101121]O \approx \begin{bmatrix} 0.401 & 0.198 & 0.401 \\ 0.248 & 0.248 & 0.504 \\ 0.284 & 0.140 & 0.576 \end{bmatrix} \begin{bmatrix} 1 & 0 \\ 1 & 1 \\ 2 & 1 \end{bmatrix}

我们逐行来算。

第 1 个输出向量

o1=0.401[1,0]+0.198[1,1]+0.401[2,1]o_1 = 0.401[1,0] + 0.198[1,1] + 0.401[2,1]

第一维：

0.401×1+0.198×1+0.401×2=0.401+0.198+0.802=1.4010.401\times1 + 0.198\times1 + 0.401\times2 =0.401+0.198+0.802 =1.401

第二维：

0.401×0+0.198×1+0.401×1=0+0.198+0.401=0.5990.401\times0 + 0.198\times1 + 0.401\times1 =0+0.198+0.401 =0.599

所以：

o1≈[1.401, 0.599]o_1 \approx [1.401,\ 0.599]

第 2 个输出向量

o2=0.248[1,0]+0.248[1,1]+0.504[2,1]o_2 = 0.248[1,0] + 0.248[1,1] + 0.504[2,1]

第一维：

0.248+0.248+1.008=1.5040.248+0.248+1.008=1.504

第二维：

0+0.248+0.504=0.7520+0.248+0.504=0.752

所以：

o2≈[1.504, 0.752]o_2 \approx [1.504,\ 0.752]

第 3 个输出向量

o3=0.284[1,0]+0.140[1,1]+0.576[2,1]o_3 = 0.284[1,0] + 0.140[1,1] + 0.576[2,1]

第一维：

0.284+0.140+1.152=1.5760.284+0.140+1.152=1.576

第二维：

0+0.140+0.576=0.7160+0.140+0.576=0.716

所以：

o3≈[1.576, 0.716]o_3 \approx [1.576,\ 0.716]

十三、于是最终输出矩阵是

O≈[1.4010.5991.5040.7521.5760.716]O \approx \begin{bmatrix} 1.401 & 0.599 \\ 1.504 & 0.752 \\ 1.576 & 0.716 \end{bmatrix}

这就是一次单头 Self-Attention 的最终输出。

现在你已经真正看到了从：

  • 输入 (H)
  • 到 (Q,K,V)
  • 到 (QK^T)
  • 到缩放
  • 到 softmax
  • 到 (AV)

的完整数值流动。

十四、现在从意义上看这次输出到底发生了什么

我们再回头看一下最初输入：

H=[100111]H= \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \end{bmatrix}

经过 attention 后输出：

O≈[1.4010.5991.5040.7521.5760.716]O \approx \begin{bmatrix} 1.401 & 0.599 \\ 1.504 & 0.752 \\ 1.576 & 0.716 \end{bmatrix}

你会发现：

  • 每个 token 的表示都变了
  • 而且都不再只是自己的原始向量
  • 它们都混入了其他 token 的信息

这就是上下文化的本质。

例如第 1 个 token，原来只是：

[1,0][1,0]

现在变成了：

[1.401,0.599][1.401, 0.599]

这说明它已经吸收了：

  • token 2 的一部分信息
  • token 3 的一部分信息

所以 attention 的真正作用不是"把一个 token 替换掉"，而是：

让每个 token 的表示变成"它自己 + 它和全局上下文的关系"的混合结果。

十五、如果加入 causal mask，会发生什么

我们刚才算的是一个不带 mask 的单头 Self-Attention。
这更像 Encoder 风格，或者纯粹为了理解计算流程。

如果换成 Decoder-only 模型的 Causal Self-Attention，那么 score matrix 在 softmax 前要加 mask。

例如长度 3 的序列，mask 会让：

  • 第 1 行只能看第 1 列
  • 第 2 行只能看前 2 列
  • 第 3 行可以看前 3 列

也就是非法位置：变成 (-\infty)，再 softmax 归一化。

这样就会强制：

  • token 1 不能看未来
  • token 2 不能看 token 3

所以 causal mask 本质上是在改 attention 的信息流方向。

你也能从这里真正看到：

Decoder-only Transformer 的自回归性质，不是口头说的，而是通过 mask 在 score matrix 上硬编码进去的。

十六、为什么这个手算例子很重要

因为它让你把很多抽象说法真正对应到了数值对象上。

例如：

1. 你真正看到了 Q/K/V 不是神秘变量

而是：

Q=HWQ, K=HWK, V=HWVQ=HW^Q,\ K=HW^K,\ V=HW^V

2. 你真正看到了 (QK^T) 的形状

它就是一个 (n\times n) 的相关性矩阵。

3. 你真正看到了 softmax 的角色

它不是神饰，而是在做每个 query 的注意力分配归一化。

4. 你真正看到了输出为什么是加权和

因为：

oi=∑jaijvjo_i = \sum_j a_{ij} v_j

这是 attention 最本质的形式。

换句话说，这个例子让你从"会解释 Attention"走向了"能亲手算一次 Attention"。

十七、如果继续往后一步，真正的 Transformer block 还会加什么

我们这次只算了最核心的 attention 部分。
但真实 Transformer block 还会继续做：

  1. attention 输出后加残差
  2. 再过 LayerNorm（或先 LN 再过 attention，取决于 Pre-LN/Post-LN）
  3. 再过 FFN
  4. 再加残差
  5. 再进入下一层

也就是说，这篇算的是：

一层 Transformer Block 里最核心的那个——跨位置交互心脏

而不是完整 block 的所有后续处理。

如果你已经能把这一篇的例子算明白，再去看下一篇关于完整 block 的手算，就会扎实很多。

十八、从"线性代数角度"再把这一篇压缩一遍

如果你更偏爱数学，那这篇其实可以被压缩成下面这条链：

输入表示

H∈Rn×dmodelH \in \mathbb{R}^{n\times d_{\text{model}}}

三个线性投影

Q=HWQ, K=HWK, V=HWVQ=HW^Q,\ K=HW^K,\ V=HW^V

构造关系矩阵

S=QKTS = QK^T

缩放

S~=Sdk\tilde{S}=\frac{S}{\sqrt{d_k}}

归一化成注意力权重

A=softmax(S~)A=\text{softmax}(\tilde{S})

对 Value 做全局加权聚合

O=AVO = AV

所以 attention 的数学本体非常清晰：

先构造一张由 Query-Key 相似度决定的图，再沿着这张图对 Value 做加权传播。

十九、本篇真正要记住的四条主线

第一条：Q/K/V 都来自输入的线性投影

它们不是外部给的，而是同一输入在三个不同功能空间中的表示。

第二条：(QK^T) 构造的是"谁该关注谁"的关系矩阵

这一步本质上是在做两两匹配。

第三条：softmax 把打分变成注意力分配

每一行都变成一个对全序列的权重分布。

第四条：最终输出是对 V 的加权和

所以 attention 本质上是一个基于相似度的内容聚合器。

二十、用一句话压缩本篇

一次单头 Self-Attention 的 forward，本质上就是：把输入向量投影成 Q/K/V，用 Q 和 K 构造相关性矩阵，再把这个相关性归一化成注意力权重，并用这些权重对 V 做加权求和，从而得到新的上下文化表示。

二十一、下一篇预告

下一篇主题是：

手算一个完整 Transformer Block：加入残差、LayerNorm 与 FFN

会重点讲清楚：

  • 如何在 attention 输出后接入残差连接
  • LayerNorm 的手算步骤
  • FFN 的升维非线性压缩完整过程
  • 一层完整 Transformer Block 从头到尾的数值流动