Part V · 第 17 章网络架构 · 17.3

Transformer 与注意力机制 (Attention)

注意力机制让每个 token 按内容去“查”其他所有 token：用查询 (query) 和每个键 (key) 做点积，经 softmax 得到权重，再对值 (value) 加权平均。它不在乎输入有几个、按什么顺序排，一步就能把任意两个 token 连起来，代价是 token 数的平方。本节从一次软查找出发，在浏览器里现场训练一个用注意力从桌上“取物”的小模型，再依次讲多头、为什么除以 √d、位置编码、Transformer 块（残差与 LayerNorm）、因果掩码与 KV 缓存，最后数一数机器人相机图像变成 token 之后要付出多少计算：这正是 ACT 和 VLA 模型的地基。

约 50 分钟前置：17.1 MLP 与 CNN · 17.2 RNN 与状态空间模型 · 16.3 反向传播 · 1.7 点积 7 个交互

1为什么机器人需要它

看一个机械臂策略在某一个控制周期里拿到的输入：两三台相机的图像、7 个关节角、夹爪开合、一句语言指令“把红色杯子放进盒子”，也许还有过去几帧的历史。前两节的网络各有难处。MLP 要求输入是一条长度固定、顺序固定的向量：桌上多一个杯子，或两个物体换了位置，它就得重新学。CNN 在图像上共享权重，但每层只看一个小邻域，左相机里的杯子和右相机里的盒子要很多层之后才“见面”。RNN 与状态空间模型把历史压进一个固定大小的状态，并且按时间一步步算。

策略真正需要的是按内容路由信息：“夹爪这个 token 应该去读取红色杯子所在的那块图像”，不管桌上有几个物体、它们在输入里排第几。注意力机制就做这件事：每个 token 发出一个查询，和所有 token 的键比较相似度，再按相似度对它们的值加权平均。输入可以是任意多个 token，任意两个 token 之间一步就连上，所有位置并行计算；代价是要算所有 token 对，计算量与 token 数 $n$ 的平方成正比。

把注意力和逐 token 的 MLP、残差连接、LayerNorm 叠成很多层，就是 2017 年提出的 Transformer。如今它是语言模型 (GPT)、视觉模型 (ViT) 和许多机器人策略的骨架：ACT 用它把多台相机的特征和关节角融合成一段动作（22.2 节），RT-1 用它处理图像 token 并输出动作 token，视觉-语言-动作 (VLA) 模型则把图像、文字和动作都当成 token 处理（Part VII）。本节从一次查找讲到一个完整的块，并现场训练一个注意力模型。

2直觉：注意力是一次“软”的字典查找

编程里的字典按键精确匹配：table["green"] 返回绿色物体的位置，键不完全一样就什么也查不到。注意力把它变“软”：拿查询 $\hl{q}{\mathbf q}$ 和每一个键 $\hl{k}{\mathbf k_i}$ 做点积（1.7 节：点积衡量两个向量有多同向），得到分数 $s_i=\mathbf q\cdot\mathbf k_i/\sqrt d$；分数经过 softmax（16.1 节）变成一组非负、加起来等于 1 的权重 $\hl{alpha}{\alpha_i}$；输出是值 $\hl{v}{\mathbf v_i}$ 的加权平均 $\hat{\mathbf y}=\sum_i\alpha_i\mathbf v_i$。

下图是一张桌子上的 5 个物体。每个物体的键描述“它看起来是什么样”（左图，一个二维特征，这里按颜色排成一圈），值是“它在哪里”（中图，桌面上的位置）。查询是“我在找什么”。拖动琥珀色的查询箭头：它指向哪个键，那个物体的权重就大；箭头越长，分数差距越大，权重越集中，接近一次“硬”查找；箭头很短时所有分数都接近 0，权重接近均匀，输出就是所有位置的平均。因为权重非负且和为 1，输出 $\hat{\mathbf y}$ 永远落在这些值围成的凸包里。

FIG 1一次软查找：查询找键，取回值的加权平均

左：键空间（$d=2$）。琥珀色箭头是查询 $\mathbf q$（拖动箭头尖），五个彩色圆点是五个物体的键 $\mathbf k_i$（也可拖动），旁边的数是分数 $s_i$。中：桌面，彩色方块是值 $\mathbf v_i$（物体位置，可拖动）；从输出 $\hat{\mathbf y}$（琥珀色空心圆）连向每个物体的线越粗，权重越大。右：五个权重 $\alpha_i$。“硬查找”只取分数最大的那个值：输出会跳变，而且它对 $\mathbf q$ 的梯度处处为零（读数里的 $\|\partial\hat{\mathbf y}/\partial\mathbf q\|$），所以无法用反向传播学习查询。

试一试

（softmax 模式）把查询拉长、指向某一个键，直到它的权重超过 0.9：这已接近硬查找。再把箭头缩回原点附近，直到每个权重都低于 0.3：输出变成所有位置的平均。
让查询指向两个相邻键的中间，并拉长到两个最大的权重都超过 0.4：输出落在两个物体之间的连线上，这是硬查找给不出的答案。

核心洞见

一个 token 在注意力里扮演三个角色：查询是“我在找什么”，键是“我能被什么找到”，值是“被选中时我交出什么”。在真正的网络里三者由同一个 token 向量经过三个不同的可学习矩阵 $W_Q,W_K,W_V$ 得到，所以一个物体的键可以只编码它的颜色，值却只编码它的位置。网络要学的，就是这三个投影。

3动手：现场训练一个“取物”的注意力模型

现在让网络自己学出 $W_Q,W_K,W_V$。任务：80 cm × 80 cm 的桌上放着 2 到 5 个颜色各不相同的方块（共 6 种颜色），指令是“取绿色方块”，要输出它的位置。每个物体是一个 token $\mathbf x_i=[p_x,p_y,\text{颜色 one-hot}]$（8 个数），指令是另一个 token $\mathbf c$（12 个数：目标颜色的 one-hot，再加一个第二颜色的 one-hot，下面的“向量”任务用）。模型只有一层注意力：

$$ \mathbf q=W_Q\mathbf c,\quad \mathbf k_i=W_K\mathbf x_i,\quad \mathbf v_i=W_V\mathbf x_i,\quad \alpha=\operatorname{softmax}_i\!\Big(\frac{\mathbf q\cdot\mathbf k_i}{\sqrt{d_k}}\Big),\quad \hat{\mathbf y}=W_O\sum_i\alpha_i\mathbf v_i+\mathbf b,\qquad L=\tfrac12\|\hat{\mathbf y}-\mathbf y\|^2 $$

这里 $d_k=8$，一个头只有 182 个参数。查询来自一个 token（指令），键和值来自另一组 token（物体）：这叫交叉注意力 (cross-attention)。每一步取 32 个随机新场景，用手写的反向传播（16.3 节；第 5 节的读数里和有限差分核对过）求梯度，再用 Adam（6.3 节，学习率 0.01）更新。作为对照，同样的数据同时训练一个 MLP（60 → 64 → 64 → 2，ReLU，8194 个参数）：它把物体 token 依次写进 6 个固定的“槽位”，空槽填 0。误差在两组固定的测试场景上算：200 个 2 到 5 个物体的场景（训练时见过的数量），和 200 个 6 个物体的场景（训练时从没见过）。

第二个任务“A 相对 B 的向量”要输出 $\mathbf p_A-\mathbf p_B$，比如“从蓝色方块出发，绿色方块在哪个方向、多远”。可以选 1、2 或 4 个注意力头：每个头有自己的 $W_Q,W_K,W_V$，各做一次查找，结果拼起来再由 $W_O$ 线性组合，这就是多头注意力。

FIG 2现场训练：交叉注意力从桌上取物，对照 MLP

左：当前展示的一个场景（不参与训练，物体可以拖动）。彩色方块是物体，左上角写着指令。“取物”任务里虚线圆是正确答案，粗加号 ✚ 是注意力模型的预测，细叉 ✕ 是 MLP 的预测；“向量”任务里虚线箭头是正确的 $\mathbf p_A-\mathbf p_B$（从 B 画出），粗箭头（尖端 ✚）是注意力的预测，细箭头（尖端 ✕）是 MLP 的。每个物体外的圆环是注意力权重：环越大越粗，权重越大；多头时第 1 个头是实线环，第 2 个头是虚线环，第 3、4 个头是点线和点划线，旁边的数字依次是各头的权重。右：测试误差（cm，对数坐标）随训练步数的变化，琥珀色为注意力，紫色为 MLP；实线是 2 到 5 个物体的测试集，虚线是 6 个物体的测试集。

保持“取物、1 个头”，按“训练”：约 200 多步后（默认种子是 225 步，换种子在 175 到 250 步之间）注意力在 2 到 5 个物体上的测试误差降到 2 cm 以下，而此时 MLP 还在 20 cm 以上（约 25 cm）。
训练好之后，把展示场景的物体数调到 6（训练时从没出现过的数量）：注意力在 6 个物体的测试集上误差仍低于 3 cm，同样一组权重直接适用。
切到“A 相对 B 的向量”，用 1 个头训练至少 1500 步：误差停在 20 cm 以上下不去。再换成 2 个头训练：误差降到 3 cm 以下（默认种子约 425 步，换种子在 350 到 550 步之间）。

值得注意的地方

(1) 注意力不需要槽位。同一组权重对任意个数、任意顺序的物体都适用：键只和颜色有关，所以无论绿色方块排在第几个，查询都能找到它。MLP 却要为 6 个槽位分别学会同一条“找颜色匹配”的规则，第 6 个槽位在训练里永远是 0，所以 6 个物体的场景它完全没准备。 (2) 学出来的权重就是一次查找：取物任务训练 600 步后，在测试场景上目标物体的平均权重约为 0.97。 (3) 一个头只能做凸组合。单头的输出是 $W_O\sum_i\alpha_i\mathbf v_i$，权重全是非负、和为 1，没法同时给 A 一个“+1”、给 B 一个“−1”；实际训练里单头只能盯住其中一个物体，误差停在 28 cm 左右（默认种子）。两个头训练 1500 步后，一个头找 A（平均权重约 0.98），另一个头找 B（约 0.97），$W_O$ 再把两次查找的结果相减。多头的意义就是并行做几次不同的查找，再线性组合。

4从一层到一个 Transformer：四个关键零件

零件 1：为什么分数要除以 √d

真实模型里每个头的维度 $d_k$ 常是 64 或 128。假设查询和键的每个分量都是均值 0、方差 1 的独立随机数（初始化时大致如此），那么点积 $\mathbf q\cdot\mathbf k=\sum_{j=1}^{d_k}q_jk_j$ 是 $d_k$ 个方差为 1 的项之和，方差为 $d_k$，标准差 $\sqrt{d_k}$。分数的差距随维度变大，softmax 就几乎把全部权重给最大的那个：它饱和了。饱和的 softmax 的雅可比 $\operatorname{diag}(\alpha)-\alpha\alpha^\top$ 接近零，梯度传不过去，训练卡住。除以 $\hl{sqrtd}{\sqrt{d_k}}$ 把分数的方差拉回到 1，与维度无关。

FIG 3不缩放时，维度越高，softmax 越饱和

每次随机抽一个查询和 16 个键，分量都是独立的标准正态数，共 200 次（浏览器里现算，固定随机种子）。左：分数的分布，红色是不缩放的 $\mathbf q\cdot\mathbf k$，琥珀色是除以 $\sqrt{d_k}$ 之后的；横轴固定在 ±40，超出的计入最边上的格子。中：其中一次抽样的 16 个 softmax 权重（按当前是否缩放）。右：随维度 $d_k=1$ 到 1024 的平均最大权重（实线）和 softmax 雅可比的平均 Frobenius 范数（虚线），红色不缩放，琥珀色缩放；竖线标出当前维度。

关掉缩放，把维度拉到 256 或以上：平均最大权重超过 0.9，雅可比范数只有缩放时的三分之一左右。再打开缩放（维度仍 ≥ 256）：最大权重回到约 0.25，和 $d_k=1$ 时的 0.21 相差无几。

零件 2：位置编码，让注意力知道顺序

FIG 2 里“不在乎顺序”是优点，但对一条轨迹、一句话或一段动作序列，顺序就是信息。自注意力 (self-attention) 里查询、键、值都来自同一组 token；把输入打乱顺序，每个 token 的输出只是跟着换位置，数值完全不变（第 5 节证明）。所以一个没有位置信息的 Transformer 分不清“伸向杯子”和“从杯子缩回”：两条轨迹的路点集合相同，只是顺序相反。

办法是给每个 token 加上一个只依赖位置的向量。原始 Transformer 用正弦位置编码 $\hl{pe}{\mathrm{PE}(p)}$：第 $2i$ 和 $2i+1$ 维是 $\sin(\omega_ip)$ 和 $\cos(\omega_ip)$，频率 $\omega_i=10000^{-2i/d}$ 从每步 1 弧度一路降到很慢。由 $\sin a\sin b+\cos a\cos b=\cos(a-b)$，两个位置编码的点积只依赖位置差 $p-p'$。现在很多模型改用旋转位置编码 (RoPE)：不往 token 上加东西，而是把查询和键的每一对分量按位置旋转一个角度 $\hl{rope}{m\omega}$，于是分数只依赖相对位置 $n-m$。

FIG 4位置编码：正弦编码、RoPE，以及没有它时会发生什么

左上：$d=32$ 维正弦编码，横轴是位置 $p=0\ldots63$，纵轴是维度（0 在上）；越偏琥珀色值越接近 +1，越偏青色越接近 −1。方框标出当前位置。右上：当前位置与所有位置的相似度 $\mathrm{PE}(p)\cdot\mathrm{PE}(p')/(d/2)$：峰在 $p'=p$，形状随 $p$ 平移而不变。左下：RoPE 的一对分量（频率 $\omega=0.35$ 弧度/位置）：细虚线是原始的 $\mathbf q$（琥珀，可拖）和 $\mathbf k$（紫，可拖），实线箭头是旋转到位置 $m$ 的 $\mathbf q_m$ 和位置 $n$ 的 $\mathbf k_n$。右下：一条 6 个路点的伸手轨迹，经过一层随机初始化的自注意力（$d=8$，固定种子）再对所有 token 取平均，得到 8 维输出；琥珀色柱是按原顺序输入的输出，紫色柱是顺序反过来的输出。

位置编码选“无”，按“反转轨迹”：两组柱子完全重合（差 $10^{-16}$ 量级的舍入误差），网络分不出伸手和缩回。再选“正弦”并反转：两组柱子明显不同。
在 RoPE 图里记下分数 $\mathbf q_m\cdot\mathbf k_n$，然后用“m、n 同时 +1”把两个位置一起往后挪至少 3 步（不改变 $\mathbf q$、$\mathbf k$）：两个箭头都转了，分数却一点不变，因为 $n-m$ 没变。

零件 3：Transformer 块，残差与 LayerNorm

一个 Transformer 块有两个子层：先是多头自注意力，让 token 之间交换信息；再是对每个 token 独立作用的 MLP（通常隐藏层宽 $4d$），在每个 token 内部加工信息。每个子层外面包一条残差连接 $\hl{res}{\mathbf x+f(\mathbf x)}$（17.1 节的 ResNet 思想），以及一个 $\hl{ln}{\mathrm{LN}}$（LayerNorm：把每个 token 的向量减去自身均值、除以自身标准差，再乘以可学习的增益）。现在多数模型用 Pre-LN：先归一化再进子层；原始 Transformer 是 Post-LN：加完残差再归一化。

这些零件不是装饰。注意力的输出是值的加权平均，平均会让 token 彼此越来越像：只叠注意力、不加残差，几层之后所有 token 就变成同一个向量（“秩坍缩”，Dong 等 2021 年证明它随深度以双指数速度发生）。加上残差，每个 token 保留自己的“原件”；但没有归一化时，每层把残差流放大一个倍数，幅值随深度指数增长。下图叠了 24 个随机初始化的块（8 个 token，$d=16$，2 个头），让你逐个拆掉零件看后果。

FIG 5拆掉残差或 LayerNorm，深层网络会坍缩或爆炸

上：一个块的结构，随开关变化；关掉的零件变淡，关掉残差时跳线变成带 ✕ 的虚线。下左：各层的 token 多样性 $\|X-\mathbf 1\bar{\mathbf x}^\top\|_F/\|X\|_F$（0 表示所有 token 完全相同，对数坐标），蓝线。下中：各层残差流的均方根幅值（对数坐标），琥珀线。两图的竖线是滑块选的层。下右：该层 8 个 token 两两之间的余弦相似度，越偏琥珀色越接近 +1，越偏青色越接近 −1；全部一片琥珀色就是坍缩。权重都是随机初始化（未训练），“换一组权重”换随机种子。

关掉残差（深度 ≥ 6）：第 6 层时多样性已经降到 $10^{-10}$ 以下（实际是 $10^{-16}$ 量级的舍入误差），右下的相似度矩阵一片琥珀色：8 个 token 变成了同一个向量。
残差打开、LayerNorm 选“无”、MLP 打开，深度拉到 20 以上：幅值超过 $10^5$。然后改成 Pre-LN（深度仍 ≥ 20）：幅值回到 10 以下，多样性也保住了。

零件 4：因果掩码与 KV 缓存

注意力本身允许每个 token 看所有 token。要逐个生成序列（文字、或者一段动作 token），训练时就不能让第 $i$ 个位置偷看后面的答案：在 softmax 之前给 $j>i$ 的分数加上 $-\infty$，这就是因果掩码 $\hl{mask}{M}$，被屏蔽的权重恰好为 0。BERT 和 ViT 的编码器不加掩码（双向）；GPT 一类的解码器用因果掩码；一些 VLA 用“分块”掩码：观测前缀内部双向，动作 token 只看前缀和更早的动作。

因果掩码带来一个工程红利：既然早先的 token 永远看不到后来的 token，它们在每一层的键和值算完就不会再变，可以存起来，这就是 KV 缓存 $\hl{kv}{(K,V)}$。生成下一个 token 时只需要为它算一行（它的查询对所有已缓存的键），而不是把整张注意力矩阵从头再算一遍。双向注意力不能这样缓存：新 token 一来，旧 token 的表示也会变。

FIG 6掩码决定谁能看谁，KV 缓存决定每步要算多少

注意力矩阵：行是查询，列是键。序列是 6 个观测 token（4 个图像块 I、1 个本体感觉 P、1 个指令 T），后面逐个生成 8 个动作 token a₁…a₈。小 ✕ 是被掩码屏蔽的格子，空白是还不存在的 token。琥珀色格子是这一步要算的点积；青色格子早先算过，靠缓存的键和值，这一步不必重算；红色格子只在“双向 + 缓存”时出现：旧 token 现在本该看到新 token，缓存却没有更新它们，结果是错的。右侧计数的是累计点积次数（每层每头）。

用因果或分块掩码，分别在“缓存开”和“缓存关”下把 8 个动作 token 都生成完：因果掩码下缓存只需 105 次点积，不缓存要 525 次；序列越长，差距越大。

一个想法，几种看法

检索：注意力是可微的键值存储，每次读取是按相似度加权的“软”查找（FIG 1）。图：自注意力是在完全图上的消息传递，边权 $\alpha_{ij}$ 由内容现算；17.4 节的 GAT 就是只在给定的边上做注意力。 核平滑：$\hat{\mathbf y}=\sum_i\alpha_i\mathbf v_i$ 是以 $\exp(\mathbf q\cdot\mathbf k)$ 为核的 Nadaraya-Watson 回归。 对照 RNN：17.2 节的 RNN 把历史压进固定大小的状态，每步代价不变；带 KV 缓存的 Transformer 保留全部历史，每步代价随长度线性增长，换来任意两步之间的直接连接。

5再看数学

提示：把鼠标悬停在公式中带底色的符号上，图中对应的元素会被高亮（查询、键、值、权重、缩放、位置编码、RoPE、残差、LayerNorm、掩码、KV 缓存、平方项）。

缩放点积注意力

把 $n$ 个 token 排成矩阵 $X\in\R^{n\times d}$（每行一个 token）。三个投影给出 $Q=XW_Q$、$K=XW_K\in\R^{n\times d_k}$、$V=XW_V\in\R^{n\times d_v}$。注意力是

$$ \operatorname{Attn}(Q,K,V)=\underbrace{\operatorname{softmax}\!\Big(\frac{\hl{q}{Q}\,\hl{k}{K}^\top}{\hl{sqrtd}{\sqrt{d_k}}}+\hl{mask}{M}\Big)}_{\hl{alpha}{A}\ \in\ \R^{n\times n}}\hl{v}{V},\qquad A_{ij}=\frac{\exp(s_{ij})}{\sum_{j'}\exp(s_{ij'})},\quad s_{ij}=\frac{\mathbf q_i\cdot\mathbf k_j}{\sqrt{d_k}}+M_{ij} $$

softmax 按行做（对每个查询，在所有键上归一化），所以 $A$ 的每一行非负且和为 1，输出的第 $i$ 行是 $\sum_jA_{ij}\mathbf v_j$。掩码 $M_{ij}$ 取 0 或 $-\infty$：因果掩码在 $j>i$ 处取 $-\infty$，变长批次里的填充 token 也用同样的方法屏蔽。自注意力的 $Q,K,V$ 来自同一个 $X$；交叉注意力的 $Q=YW_Q$ 来自另一组 token $Y$（FIG 2 的指令，ACT 解码器里的动作查询），$K,V$ 来自 $X$。为什么除以 $\sqrt{d_k}$：若 $q_j,k_j$ 独立、均值 0、方差 1，则 $\operatorname{Var}(\mathbf q\cdot\mathbf k)=\sum_j\operatorname{Var}(q_jk_j)=d_k$。

反向传播穿过 softmax

FIG 2 的梯度就是 16.3 节的反向模式链式法则，只有 softmax 这一步需要想一想。设一个查询的输出 $\hat{\mathbf y}=\sum_i\alpha_i\mathbf v_i$，上游梯度 $\mathbf g=\partial L/\partial\hat{\mathbf y}$。softmax 的雅可比是 $\partial\alpha_i/\partial s_j=\alpha_i(\delta_{ij}-\alpha_j)$，于是

$$ \frac{\partial L}{\partial\mathbf v_i}=\alpha_i\mathbf g,\qquad \frac{\partial L}{\partial s_i}=\alpha_i\Big(\mathbf g\cdot\mathbf v_i-\sum_j\alpha_j\,\mathbf g\cdot\mathbf v_j\Big),\qquad \frac{\partial L}{\partial\mathbf q}=\sum_i\frac{\partial L}{\partial s_i}\frac{\mathbf k_i}{\sqrt{d_k}},\qquad \frac{\partial L}{\partial\mathbf k_i}=\frac{\partial L}{\partial s_i}\frac{\mathbf q}{\sqrt{d_k}} $$

中间那个式子说明了 FIG 3 的现象：一个键的分数梯度正比于“它的值比加权平均好多少”乘以 $\alpha_i$。权重集中在一个键上时（$\alpha_i\approx1$），括号趋于 0；其他键的 $\alpha_i\approx0$。两种情况梯度都消失。上面一行读数是页面加载时现算的：FIG 2 的 2 头模型在“向量”任务上，解析梯度和中心有限差分逐个参数比较。

多头注意力

$$ \mathbf z_h=\operatorname{Attn}\big(XW_Q^{(h)},XW_K^{(h)},XW_V^{(h)}\big)\in\R^{n\times d/H},\qquad \operatorname{MHA}(X)=\big[\mathbf z_1\ \cdots\ \mathbf z_H\big]\,W_O,\quad W_O\in\R^{d\times d} $$

每个头的维度是 $d/H$，所以 $H$ 个头的总参数和计算量与一个宽度为 $d$ 的头相同；不同的是，$H$ 个头有 $H$ 张独立的注意力矩阵，可以同时做 $H$ 次不同的查找。单个头的每一行输出都在值的凸包里（FIG 2 的向量任务因此学不会）；$W_O$ 把几个头的结果线性组合，就能表达差、和、加权组合。

置换等变性与位置编码

设 $P$ 是置换矩阵（把 token 重新排序）。没有掩码时，$\operatorname{softmax}(PQ(PK)^\top)=P\operatorname{softmax}(QK^\top)P^\top$（按行的 softmax 不在乎行和列一起换序），又 $P^\top P=I$，所以

$$ \operatorname{Attn}(PX)=P\,\operatorname{Attn}(X)\quad\text{(self-attention is permutation-equivariant)},\qquad \operatorname{Attn}(Y\!\leftarrow\!PX)=\operatorname{Attn}(Y\!\leftarrow\!X)\quad\text{(cross-attention output is invariant)} $$

逐 token 的 MLP、残差和 LayerNorm 也都与置换交换，所以没有位置编码的整个 Transformer 都是置换等变的，对 token 取平均之后就完全置换不变（FIG 4 右下）。上式右边说的是 FIG 2：交叉注意力里打乱物体的顺序，指令查询的输出完全不变。正弦编码的点积

$$ \hl{pe}{\mathrm{PE}(p)}\cdot\mathrm{PE}(p')=\sum_{i=0}^{d/2-1}\big[\sin\omega_ip\,\sin\omega_ip'+\cos\omega_ip\,\cos\omega_ip'\big]=\sum_{i}\cos\big(\omega_i(p-p')\big),\qquad \omega_i=10000^{-2i/d} $$

只依赖 $p-p'$，而且 $\mathrm{PE}(p+\Delta)$ 是 $\mathrm{PE}(p)$ 在每对维度上的一个旋转，网络可以用线性运算表示“往后挪 $\Delta$”。RoPE 把这个想法直接用在查询和键上：令 $R(\theta)$ 是二维旋转，把 $\mathbf q$ 的第 $i$ 对分量乘以 $R(m\omega_i)$、把 $\mathbf k$ 的第 $i$ 对乘以 $R(n\omega_i)$，因为 $R(a)^\top R(b)=R(b-a)$：

$$ \mathbf q_m\cdot\mathbf k_n=\sum_i\big(R(\hl{rope}{m\omega_i})\mathbf q^{(i)}\big)^\top\big(R(n\omega_i)\mathbf k^{(i)}\big)=\sum_i\mathbf q^{(i)\top}R\big((n-m)\omega_i\big)\mathbf k^{(i)} $$

Transformer 块

$$ X'=\hl{res}{X+}\operatorname{MHA}\big(\hl{ln}{\mathrm{LN}}(X)\big),\qquad X''=X'+\mathrm{MLP}\big(\mathrm{LN}(X')\big),\qquad \mathrm{LN}(\mathbf x)=\boldsymbol\gamma\odot\frac{\mathbf x-\mu(\mathbf x)\mathbf 1}{\sigma(\mathbf x)}+\boldsymbol\beta $$

这是 Pre-LN；Post-LN 写成 $X'=\mathrm{LN}(X+\operatorname{MHA}(X))$。$\mu,\sigma$ 是同一个 token 的 $d$ 个分量的均值和标准差（BatchNorm 则是在批次里对同一个特征求），所以 LayerNorm 不依赖批次大小，推理时一个样本也能用。Pre-LN 的残差流上每层只加一个大小受控的量 $f(\mathrm{LN}(\mathbf x))$，幅值增长得很慢（FIG 5 里 24 层后约为 7）；没有 LN 时加的是 $f(\mathbf x)$，它和 $\mathbf x$ 成比例，于是每层乘一个大于 1 的因子，指数增长。Pre-LN 也更好训练（Xiong 等 2020），原始 Post-LN 通常需要学习率预热。

计算量与内存

按乘加算两次浮点运算，一层的前向计算是投影 $Q,K,V,O$：$8nd^2$；宽 $4d$ 的 MLP：$16nd^2$；注意力本身（$QK^\top$ 和 $AV$）：$\hl{n2}{4n^2d}$。所以注意力所占比例是

$$ \frac{4n^2d}{24nd^2+4n^2d}=\frac{n}{n+6d},\qquad \text{KV cache}=2\cdot L\cdot n\cdot d\cdot b\ \text{bytes} \quad(L\ \text{layers},\ b\ \text{bytes per number}) $$

$n$ 小于 $6d$ 时，大头其实是那些和 $n$ 成线性的矩阵乘法；只有 token 多到超过 $6d$，平方项才占主导（FIG 7）。内存上，朴素实现要为每个头存一张 $n\times n$ 的注意力矩阵；FlashAttention 分块计算、不把它整张写进显存，结果完全相同（不是近似），但计算量仍是 $n^2$。生成时用 KV 缓存，第 $t$ 步只算新的一行：总代价从 $n^3$ 量级（每步重算整张矩阵是 $\sum_t t^2\sim n^3/3$，跳过被掩码的一半也还有 $n^3/6$）降到 $\sum_t t\sim n^2/2$（FIG 6），代价是缓存要占 $2Lnd$ 个数的内存。

ViT：图像也是一串 token

Vision Transformer 把 224 × 224 的图像切成 16 × 16 像素的块，得到 14 × 14 = 196 块；每块的 16·16·3 = 768 个像素值经一个线性层变成一个 $d$ 维 token，再加上可学习的位置编码，前面通常还放一个 [CLS] token 作为整张图的汇总。之后就是普通的 Transformer 编码器。线性切块等价于一个核大小和步长都是 16 的卷积（17.1 节），但之后每一层都是全局的：第一层起任意两块之间就能直接交换信息。

6机器人中的它：相机 token 有多贵？

机器人策略的 token 大多来自相机：每台相机每一帧切成一批图像块，再加上本体感觉（关节角、夹爪）、指令和要预测的动作。相机数、分辨率、历史帧数每翻一倍，token 数就翻一倍，而注意力的平方项翻四倍；控制回路又要求每一步在几十毫秒内算完。下图把这笔账算出来：左边是一张合成的 224 × 224 机器人相机图像按所选块大小切块（把鼠标移到某块上看它变成多少个数），右边是一层的计算量随 token 数的变化。

FIG 7从相机到 token：ViT 切块与 n² 之墙

左：一台相机的一帧，网格是 ViT 的切块。token 总数 = 相机数 × (224/块大小)² × 帧数 + 帧数（每帧一个本体感觉 token）+ 动作 token 数。右：一层前向计算量（FLOPs，对数-对数坐标），蓝线是与 $n$ 成线性的投影和 MLP（$24nd^2$），琥珀线是注意力（$4n^2d$），两者交于 $n=6d$（青色竖虚线）；圆点是当前设置。“小模型”取 $L=6$ 层、$d=512$；“7B 级”取 $L=32$、$d=4096$（与 Llama 2 7B 的层数和宽度相同）。KV 缓存按 fp16（每个数 2 字节）计。

在“小模型”下调相机数、块大小和帧数，让注意力占一层计算量的一半以上（token 数超过 $6d=3072$，例如 4 台相机、8 像素的块）。再切到“7B 级”：同样的 token 数下注意力只占一小部分，要到 24576 个 token 才过半。

所以机器人系统在 token 上很节省：ACT（Zhao 等，2023）用 ResNet-18 把每台相机的图像压成一张特征图，特征图上的每个位置是一个 token，再加上关节角 token 和 CVAE 隐变量 token 进 Transformer 编码器；解码器用一组固定的查询向量（每个未来时刻一个），通过交叉注意力一次读出整段未来动作（动作分块，22.2 节）。RT-1（Brohan 等，2022）用 TokenLearner 把每张图像的特征压缩到很少几个 token，再交给 Transformer 输出离散化的动作（17.5 节）。Perceiver（Jaegle 等，2021）让一小组隐变量 token 对海量输入做交叉注意力，代价从 $n^2$ 变成 $n$ 乘以隐变量个数。

真实系统里

模仿学习策略：ACT 是“编码器读观测、解码器用动作查询做交叉注意力”的典型结构，和 DETR 的目标查询是同一个思路；Diffusion Policy（Chi 等，2023）除了 CNN 版本，还有一个 Transformer 版本，用交叉注意力把观测条件送进去噪网络（22.2 节）。 序列决策：Decision Transformer（Chen 等，2021）把回报、状态、动作排成一条序列，用因果 Transformer 预测下一个动作。 VLA：RT-2（2023）把机器人动作写成 token，让一个视觉-语言大模型直接输出；OpenVLA（Kim 等，2024）基于 Llama 2 7B；π0（Physical Intelligence，2024）在一个视觉-语言模型旁边接一个用 flow matching 生成连续动作的“动作专家”，并使用分块的注意力掩码。这些模型的推理延迟主要来自上千个视觉 token 和 KV 缓存，第 VII 部分细讲。工程：PyTorch 的 scaled_dot_product_attention 会自动选用 FlashAttention 一类的融合实现；多查询注意力 (MQA) 和分组查询注意力 (GQA) 让多个头共享键和值，把 KV 缓存缩小数倍。

7常见误区

误区 1：忘了 √d，或把温度弄错

FIG 3：不缩放时 $d_k=256$ 的 softmax 平均把 0.93 的权重给一个键，雅可比的范数只有缩放时的三成，训练一开始梯度就很弱。同类错误还有：自己实现时把 $\sqrt{d_k}$ 写成 $d_k$（分数太平，注意力学不会集中），或者 $\mathbf q,\mathbf k$ 的初始化尺度太大。训练中注意力权重的熵突然塌到接近 0，常常是数值不稳定的信号，需要检查学习率、归一化，或对 $\mathbf q,\mathbf k$ 再做一次归一化。

误区 2：以为 Transformer 天生知道顺序

没有位置编码，自注意力把输入当集合（FIG 4）：一条轨迹和它的倒序给出同样的汇总输出。对物体集合这正合适（FIG 2），对时间序列、动作序列和图像块就是灾难。另一个坑是长度外推：用学习的位置编码或正弦编码在 50 帧上训练，推理时喂 200 帧，模型遇到的是从没见过的位置；RoPE 外推稍好但也不是免费的。训练和部署的序列长度、采样频率要一致，或者专门做外推的训练。

误区 3：掩码写错

最常见的三种：(a) 训练自回归模型时忘了因果掩码，模型直接“抄”后面的答案，训练损失漂亮得可疑，一到逐步生成就崩；(b) 变长批次里忘了屏蔽填充 token，填充位也分到了权重，结果随批次里最长的样本而变；(c) softmax 归一化的轴弄错（应在键的维度上，每个查询一行和为 1）。还有：用 KV 缓存时掩码必须真的是因果的（FIG 6 双向模式下缓存是错的），否则缓存的旧值与重算结果不一致。

误区 4：把注意力图当成解释

FIG 2 的单层模型里，权重确实就是“看了哪个物体”。但在深层多头的 Transformer 里，信息经过残差流、多个头和 $W_O$ 混合，某一层某一个头的注意力图很可能和最终决策没有直接关系；Jain 与 Wallace（2019）发现，换一组差别很大的注意力权重，预测有时几乎不变。把注意力热图叠在相机图像上做展示没问题，但要说“机器人因为看了这里才这样做”，需要做干预实验（遮挡、替换输入）。

误区 5：只盯着 n²，或完全忽视它

FIG 7：$n<6d$ 时一层计算量的大头是线性的矩阵乘法，这时砍 token 数的收益是线性的，换更小的模型宽度可能更有效。反过来，多相机、高分辨率、长历史一叠加，token 数很快过万，平方项和 KV 缓存就主宰延迟和显存。对一个必须以固定频率出动作的控制回路，要先算清楚预算：多少 token、多少层、多宽、能不能缓存、能不能一次输出一整段动作（ACT 的动作分块正是为此）。

8小练习

1. 查询和键都是 $d_k=64$ 维，分量独立、均值 0、方差 1。点积 $\mathbf q\cdot\mathbf k$ 的标准差约为多少？除以 $\sqrt{d_k}$ 之后呢？

每项 $q_jk_j$ 的方差是 $\operatorname{E}[q_j^2]\operatorname{E}[k_j^2]=1$，64 项独立相加方差为 64，标准差 8；除以 $\sqrt{64}=8$ 后标准差为 1。FIG 3 在 $d_k=64$ 时测得未缩放分数的标准差约 7.8。均值为 0 不代表每次都是 0（选项 D 混淆了均值和取值）。

2. 一个没有位置编码的 Transformer 编码器读入一条轨迹的 10 个路点，再对所有输出取平均得到一个向量。把路点顺序倒过来，这个向量会怎样？

自注意力、逐 token 的 MLP、残差和 LayerNorm 都与置换交换，所以每个路点的输出只是跟着换了位置；取平均与顺序无关，结果完全相同（FIG 4 右下，差 $10^{-16}$ 量级）。要区分伸手和缩回，必须加位置编码。

3. 为什么 KV 缓存可以用在 GPT 式的解码器上，却不能直接用在 BERT 式的双向编码器上？

缓存的前提是“算过的东西不会再变”。因果掩码保证第 $i$ 个 token 在每一层的表示只依赖 $1\ldots i$，所以它的 $\mathbf k,\mathbf v$ 可以永久保存，新 token 只需算自己的一行（FIG 6：105 对 525 次点积）。双向注意力里第 2 层起旧 token 的表示就依赖新 token，缓存的值是过时的。

4. FIG 2 的“向量”任务要输出 $\mathbf p_A-\mathbf p_B$。为什么单头注意力学不会，而两个头可以？

无论 $d_k$ 多宽、训练多久，单头的 $\sum_i\alpha_i\mathbf v_i$ 都在值的凸包里，加上共享的线性映射 $W_O$ 也无法让两个物体的位置以相反符号出现（实测误差停在约 28 cm）。两个头时一个头的权重集中在 A、一个在 B，$W_O$ 实现减法，误差降到 1 cm 以下。物体位置本来就在 token 里，所以选项 C 无关。

5. 一个 ViT 读 224 × 224 的图像。把块大小从 16 像素改成 8 像素，token 数和每层每头注意力矩阵的元素数各变成原来的多少倍？

块边长减半，每个方向的块数翻倍：$14\times14=196$ 变成 $28\times28=784$，是 4 倍；注意力矩阵是 $n\times n$，$196^2\approx3.8$ 万变成 $784^2\approx61$ 万，是 16 倍。与 $n$ 成线性的投影和 MLP 只变成 4 倍（FIG 7）。

9延伸阅读

D. Bahdanau, K. Cho & Y. Bengio, Neural Machine Translation by Jointly Learning to Align and Translate, ICLR 2015 — 注意力机制的起点：解码器对编码器状态做软对齐
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser & I. Polosukhin, Attention Is All You Need, NeurIPS 2017 — Transformer：缩放点积、多头、正弦位置编码
J. L. Ba, J. R. Kiros & G. E. Hinton, Layer Normalization, arXiv 2016; R. Xiong et al., On Layer Normalization in the Transformer Architecture, ICML 2020 — LayerNorm，以及 Pre-LN 为什么更好训练
Y. Dong, J.-B. Cordonnier & A. Loukas, Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth, ICML 2021 — FIG 5 的秩坍缩，以及残差和 MLP 的作用
J. Su, Y. Lu, S. Pan, A. Murtadha, B. Wen & Y. Liu, RoFormer: Enhanced Transformer with Rotary Position Embedding, arXiv 2021 (Neurocomputing 2024) — RoPE
A. Dosovitskiy et al., An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale, ICLR 2021 — ViT：图像块即 token
T. Dao, D. Y. Fu, S. Ermon, A. Rudra & C. Ré, FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, NeurIPS 2022 — 不存整张注意力矩阵的精确实现
A. Jaegle et al., Perceiver: General Perception with Iterative Attention, ICML 2021 — 用少量隐变量 token 做交叉注意力，绕开 n²
T. Z. Zhao, V. Kumar, S. Levine & C. Finn, Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware, RSS 2023 — ALOHA 与 ACT：多相机 token、动作查询、动作分块
L. Chen et al., Decision Transformer: Reinforcement Learning via Sequence Modeling, NeurIPS 2021 — 把决策写成因果序列建模
M. Phuong & M. Hutter, Formal Algorithms for Transformers, arXiv 2022 — 各种 Transformer 变体的精确伪代码