Transformer 与注意力机制 (Attention)Transformers & Attention

看一个机械臂策略在某一个控制周期里拿到的输入：两三台相机的图像、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）。本节从一次查找讲到一个完整的块，并现场训练一个注意力模型。

Look at what a robot-arm policy receives in one control cycle: images from two or three cameras, 7 joint angles, the gripper opening, a language instruction "put the red mug in the box", perhaps a few frames of history. The networks of the previous two lessons each struggle with this. An MLP wants one vector of fixed length in a fixed order: put one more mug on the table or swap two objects in the input and it must learn again. A CNN shares weights across an image, but each layer looks only at a small neighbourhood, so the mug in the left camera and the box in the right camera only "meet" many layers later. RNNs and state-space models squeeze the history into a fixed-size state and compute step by step in time.

What a policy really needs is to route information by content: "the gripper token should read from the image patch where the red mug is", however many objects there are and wherever they sit in the input. Attention does exactly this: each token issues a query, compares it with the keys of all tokens, and takes the average of their values weighted by the similarity. The input can be any number of tokens, any two tokens are connected in a single step, and all positions are computed in parallel; the price is that all pairs of tokens must be evaluated, so the work grows with the square of the token count \(n\).

Stack attention with a per-token MLP, residual connections and LayerNorm into many layers and you have the transformer, introduced in 2017. Today it is the backbone of language models (GPT), vision models (ViT) and many robot policies: ACT uses it to fuse the features of several cameras with the joint angles into a chunk of actions (Lesson 22.2), RT-1 runs a transformer over image tokens and outputs action tokens, and vision-language-action (VLA) models treat images, words and actions all as tokens (Part VII). This lesson goes from one lookup to a complete block and trains an attention model live.

编程里的字典按键精确匹配：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}\) 永远落在这些值围成的凸包里。

A dictionary in a program matches its key exactly: table["green"] returns the position of the green object, and a key that is not identical finds nothing. Attention makes this "soft": dot the query \(\hl{q}{\mathbf q}\) with every key \(\hl{k}{\mathbf k_i}\) (Lesson 1.7: the dot product measures how aligned two vectors are) to get scores \(s_i=\mathbf q\cdot\mathbf k_i/\sqrt d\); a softmax (Lesson 16.1) turns the scores into non-negative weights \(\hl{alpha}{\alpha_i}\) that add up to 1; the output is the weighted average of the values \(\hl{v}{\mathbf v_i}\), \(\hat{\mathbf y}=\sum_i\alpha_i\mathbf v_i\).

The figure shows 5 objects on a table. Each object's key describes "what it looks like" (left, a 2D feature, here arranged around a circle by colour), its value is "where it is" (middle, a position on the table). The query is "what I am looking for". Drag the amber query arrow: the key it points at gets a large weight; the longer the arrow, the larger the gaps between scores and the more concentrated the weights, approaching a "hard" lookup; when the arrow is very short all scores are near 0, the weights are nearly uniform and the output is the average of all positions. Because the weights are non-negative and sum to 1, the output \(\hat{\mathbf y}\) always lies inside the convex hull of the values.

现在让网络自己学出 \(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，下面的“向量”任务用）。模型只有一层注意力：

Now let the network learn \(W_Q,W_K,W_V\) itself. The task: an 80 cm × 80 cm table holds 2 to 5 blocks of different colours (6 colours in all), the instruction is "fetch the green block", and the model must output its position. Each object is a token \(\mathbf x_i=[p_x,p_y,\text{colour one-hot}]\) (8 numbers); the instruction is another token \(\mathbf c\) (12 numbers: a one-hot of the target colour, plus a one-hot of a second colour used by the "vector" task below). The model is a single attention layer:

这里 \(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\) 线性组合，这就是多头注意力。

Here \(d_k=8\), and one head has only 182 parameters. The query comes from one token (the instruction) while keys and values come from another set of tokens (the objects): this is cross-attention. Every step draws 32 fresh random scenes, computes the gradient by hand-written backpropagation (Lesson 16.3; checked against finite differences in the readout of Section 5) and updates with Adam (Lesson 6.3, learning rate 0.01). For comparison an MLP (60 → 64 → 64 → 2, ReLU, 8194 parameters) trains on the same data at the same time: it writes the object tokens one after another into 6 fixed "slots" and fills empty slots with 0. Errors are measured on two fixed test sets: 200 scenes with 2 to 5 objects (counts seen in training) and 200 scenes with 6 objects (a count never seen in training).

The second task, "vector from B to A", must output \(\mathbf p_A-\mathbf p_B\), for example "starting from the blue block, in which direction and how far is the green block". You can pick 1, 2 or 4 attention heads: each head has its own \(W_Q,W_K,W_V\) and does its own lookup, and the results are concatenated and combined linearly by \(W_O\): this is multi-head attention.

真实模型里每个头的维度 \(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，与维度无关。

In real models each head has dimension \(d_k\) of 64 or 128. Suppose every component of the query and key is an independent random number with mean 0 and variance 1 (roughly true at initialization). Then the dot product \(\mathbf q\cdot\mathbf k=\sum_{j=1}^{d_k}q_jk_j\) is a sum of \(d_k\) terms of variance 1, so its variance is \(d_k\) and its standard deviation \(\sqrt{d_k}\). The gaps between scores grow with dimension, and the softmax gives almost all the weight to the largest one: it saturates. The Jacobian of a saturated softmax, \(\operatorname{diag}(\alpha)-\alpha\alpha^\top\), is close to zero, gradients do not get through and training stalls. Dividing by \(\hl{sqrtd}{\sqrt{d_k}}\) brings the score variance back to 1, independent of the dimension.

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\)。

In FIG 2 "not caring about order" was a strength, but for a trajectory, a sentence or a sequence of actions the order is information. In self-attention the queries, keys and values all come from the same set of tokens; shuffle the input and each token's output just moves along with it, numerically unchanged (proved in Section 5). So a transformer without positional information cannot tell "reach towards the mug" from "retreat from the mug": the two trajectories have the same set of waypoints in opposite order.

The remedy is to add to each token a vector that depends only on its position. The original transformer uses a sinusoidal positional encoding \(\hl{pe}{\mathrm{PE}(p)}\): dimensions \(2i\) and \(2i+1\) are \(\sin(\omega_ip)\) and \(\cos(\omega_ip)\) with frequencies \(\omega_i=10000^{-2i/d}\), from 1 radian per step down to very slow. Since \(\sin a\sin b+\cos a\cos b=\cos(a-b)\), the dot product of two encodings depends only on the offset \(p-p'\). Many current models use the rotary positional encoding (RoPE) instead: nothing is added to the token; each pair of query and key components is rotated by an angle \(\hl{rope}{m\omega}\) proportional to the position, so the score depends only on the relative position \(n-m\).

一个 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 个头），让你逐个拆掉零件看后果。

A transformer block has two sublayers: first multi-head self-attention, which exchanges information between tokens; then an MLP applied to each token independently (hidden width usually \(4d\)), which processes information inside each token. Each sublayer is wrapped in a residual connection \(\hl{res}{\mathbf x+f(\mathbf x)}\) (the ResNet idea of Lesson 17.1) and an \(\hl{ln}{\mathrm{LN}}\) (LayerNorm: subtract each token vector's own mean, divide by its own standard deviation, multiply by a learnable gain). Most models today use Pre-LN, normalizing before the sublayer; the original transformer used Post-LN, normalizing after the residual addition.

These parts are not decoration. The output of attention is a weighted average of values, and averaging makes tokens more and more alike: stack attention alone without residuals and after a few layers every token is the same vector ("rank collapse"; Dong et al. proved in 2021 that it happens doubly exponentially fast with depth). Residuals let every token keep its own "original"; but without normalization each layer multiplies the residual stream by some factor, and its magnitude grows exponentially with depth. The figure stacks 24 randomly initialized blocks (8 tokens, \(d=16\), 2 heads) and lets you remove the parts one by one to see the consequences.

注意力本身允许每个 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 的表示也会变。

Attention by itself lets every token see every token. To generate a sequence one token at a time (words, or a run of action tokens), training must not let position \(i\) peek at the answers that come later: add \(-\infty\) to the scores with \(j>i\) before the softmax. This is the causal mask \(\hl{mask}{M}\); the masked weights become exactly 0. The encoders of BERT and ViT use no mask (bidirectional); GPT-style decoders use the causal mask; some VLAs use a "block" mask: bidirectional within the observation prefix, while action tokens see only the prefix and earlier actions.

The causal mask pays an engineering dividend: since earlier tokens never see later ones, their keys and values at every layer never change once computed and can be stored: the KV cache \(\hl{kv}{(K,V)}\). Generating the next token then needs just one new row (its query against all cached keys) instead of recomputing the whole attention matrix. Bidirectional attention cannot be cached this way: when a new token arrives, the old tokens' representations change too.

把 \(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}\)。注意力是

Stack \(n\) tokens into a matrix \(X\in\R^{n\times d}\) (one token per row). Three projections give \(Q=XW_Q\), \(K=XW_K\in\R^{n\times d_k}\) and \(V=XW_V\in\R^{n\times d_v}\). Attention is

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\)。

The softmax runs along rows (for each query, normalize over all keys), so every row of \(A\) is non-negative and sums to 1, and row \(i\) of the output is \(\sum_jA_{ij}\mathbf v_j\). The mask \(M_{ij}\) is 0 or \(-\infty\): the causal mask puts \(-\infty\) at \(j>i\), and padding tokens in variable-length batches are blocked the same way. In self-attention \(Q,K,V\) come from the same \(X\); in cross-attention \(Q=YW_Q\) comes from another set of tokens \(Y\) (the instruction of FIG 2, the action queries in ACT's decoder) while \(K,V\) come from \(X\). Why divide by \(\sqrt{d_k}\): if the \(q_j,k_j\) are independent with mean 0 and variance 1, then \(\operatorname{Var}(\mathbf q\cdot\mathbf k)=\sum_j\operatorname{Var}(q_jk_j)=d_k\).

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)\)，于是

The gradients in FIG 2 are the reverse-mode chain rule of Lesson 16.3; only the softmax step needs a moment of thought. Let one query's output be \(\hat{\mathbf y}=\sum_i\alpha_i\mathbf v_i\) with upstream gradient \(\mathbf g=\partial L/\partial\hat{\mathbf y}\). The softmax Jacobian is \(\partial\alpha_i/\partial s_j=\alpha_i(\delta_{ij}-\alpha_j)\), so

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

The middle formula explains FIG 3: a key's score gradient is \(\alpha_i\) times "how much better its value is than the weighted average". When the weight is concentrated on one key (\(\alpha_i\approx1\)) the bracket tends to 0, and every other key has \(\alpha_i\approx0\): either way the gradient vanishes. The readout above is computed when the page loads: for the 2-head model of FIG 2 on the "vector" task, the analytic gradient is compared parameter by parameter with central finite differences.

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

Each head has dimension \(d/H\), so \(H\) heads cost the same parameters and computation as one head of width \(d\); the difference is that \(H\) heads have \(H\) independent attention matrices and can do \(H\) different lookups at once. Every output row of a single head lies in the convex hull of the values (which is why the vector task of FIG 2 cannot be learned with one); \(W_O\) combines the heads linearly, so differences, sums and weighted combinations become expressible.

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

Let \(P\) be a permutation matrix (a reordering of the tokens). Without a mask, \(\operatorname{softmax}(PQ(PK)^\top)=P\operatorname{softmax}(QK^\top)P^\top\) (a row-wise softmax does not care if rows and columns are reordered together), and \(P^\top P=I\), so

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

The per-token MLP, the residuals and LayerNorm also commute with permutations, so a whole transformer without positional encoding is permutation-equivariant, and after averaging over tokens it is fully permutation-invariant (FIG 4, bottom right). The right-hand identity is FIG 2: in cross-attention, shuffling the objects leaves the instruction query's output unchanged. For the sinusoidal encoding, the dot product is

只依赖 \(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)\)：

It depends only on \(p-p'\), and \(\mathrm{PE}(p+\Delta)\) is a rotation of \(\mathrm{PE}(p)\) in each pair of dimensions, so the network can express "shift by \(\Delta\)" with a linear map. RoPE applies the idea directly to queries and keys: with \(R(\theta)\) a 2D rotation, multiply the \(i\)-th pair of \(\mathbf q\) by \(R(m\omega_i)\) and the \(i\)-th pair of \(\mathbf k\) by \(R(n\omega_i)\); since \(R(a)^\top R(b)=R(b-a)\),

这是 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 通常需要学习率预热。

This is Pre-LN; Post-LN reads \(X'=\mathrm{LN}(X+\operatorname{MHA}(X))\). \(\mu,\sigma\) are the mean and standard deviation over the \(d\) components of one token (BatchNorm instead takes them over the batch for one feature), so LayerNorm does not depend on the batch size and works with a single sample at inference. In Pre-LN every layer adds to the residual stream a quantity of controlled size, \(f(\mathrm{LN}(\mathbf x))\), so the magnitude grows only slowly (about 7 after 24 layers in FIG 5); without LN the added term is \(f(\mathbf x)\), proportional to \(\mathbf x\), so each layer multiplies by a factor above 1 and the growth is exponential. Pre-LN also trains more easily (Xiong et al. 2020); the original Post-LN usually needs a learning-rate warm-up.

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

Counting a multiply-add as two floating-point operations, one layer's forward pass costs \(8nd^2\) for the projections \(Q,K,V,O\), \(16nd^2\) for an MLP of width \(4d\), and \(\hl{n2}{4n^2d}\) for attention itself (\(QK^\top\) and \(AV\)). The share of attention is therefore

\(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\) 个数的内存。

While \(n\) is below \(6d\), most of the work is in the matrix products that are linear in \(n\); only once the tokens outnumber \(6d\) does the quadratic term dominate (FIG 7). In memory, a naive implementation stores an \(n\times n\) attention matrix per head; FlashAttention computes it in tiles without ever writing the whole matrix to GPU memory and gives exactly the same result (not an approximation), though the compute is still \(n^2\). In generation the KV cache computes only the new row at step \(t\): the total cost drops from order \(n^3\) (recomputing the full matrix every step is \(\sum_t t^2\sim n^3/3\); skipping the masked half still leaves \(n^3/6\)) to \(\sum_t t\sim n^2/2\) (FIG 6), paid for by a cache of \(2Lnd\) numbers in memory.

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

Most of a robot policy's tokens come from cameras: every frame of every camera is cut into a batch of patches, plus proprioception (joint angles, gripper), the instruction and the actions to be predicted. Double the cameras, the resolution or the history frames and the token count doubles, while the quadratic term of attention quadruples; and the control loop wants every step computed within tens of milliseconds. The figure does the accounting: on the left a synthetic 224 × 224 robot camera image is cut into patches of the chosen size (hover a patch to see how many numbers it becomes), on the right the compute of one layer is plotted against the token count.

所以机器人系统在 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\) 乘以隐变量个数。

That is why robot systems are frugal with tokens. ACT (Zhao et al., 2023) compresses each camera image into a feature map with a ResNet-18, makes every location of the feature map a token, and feeds them with a joint-angle token and a CVAE latent token into a transformer encoder; the decoder uses a fixed set of query embeddings (one per future step) that read out a whole chunk of future actions at once by cross-attention (action chunking, Lesson 22.2). RT-1 (Brohan et al., 2022) uses TokenLearner to compress each image's features into just a few tokens before a transformer outputs discretized actions (Lesson 17.5). Perceiver (Jaegle et al., 2021) lets a small set of latent tokens cross-attend to a huge input, turning the cost from \(n^2\) into \(n\) times the number of latents.