计算机视觉-10：Attention 与 Transformer

Attention 机制

Attention 的核心问题是：当模型要生成或更新某个位置的表示时，应该从哪些位置取信息、各取多少。它不像 RNN 那样必须按时间顺序一步步传递隐藏状态，而是直接计算当前位置与其他位置之间的相关性，再按相关性加权汇总信息。因此 Attention 特别擅长建模长距离依赖，也更适合并行计算。

引入

课上使用 RNN 处理翻译工作的例子来引入 Attention：

在传统的 Encoder-Decoder 架构中，Encoder 会将输入句子编码成一个固定长度的向量，然后 Decoder 根据这个向量生成输出句子。然而，这种方法在处理长句子时效果不佳，因为固定长度的向量可能无法捕捉到所有重要的信息。这个问题可以通过查看整个输入序列来解决，而不是仅仅依赖于一个固定长度的向量。

对于 Decoder，我们给出一个初始的隐藏状态 $s_0$ ，将这个隐藏状态与 Encoder 的所有隐藏状态进行比较，计算出每个隐藏状态的重要性权重（也就是注意力权重），这个权重告诉 Decoder 在生成下一个词时应该关注输入序列的哪些部分。然后，对这些状态进行加权求和，得到一个上下文向量 $c_1$ ，这个向量包含了输入序列中最相关的信息。最后，Decoder 使用这个上下文向量 $c_1$ ，输入 $y_0$ 和当前的隐藏状态 $s_0$ 来生成下一个词，随后重复整个过程来完成翻译。

以这个英语到法语的翻译为例，可以看出在这个二维的注意力图中，每个输入的单词通过注意力机制与输出的单词进行关联，有的子句之间的注意力是一条正对角线——这说明输入和输出在位置上是一一对应的；有的子句则是反对角线或者更加不规则的结构——这说明 Attention 机制可以捕捉到输入和输出之间的复杂关系，而不仅仅是位置上的对应关系。反映到翻译任务上就是能够学习不同语言之间的语法差异。

注意力层

现在将这个计算注意力权重的部分单独分离出来。从外部看，这个层接受若干个输入的向量 $\{x_i\}$ 以及一个隐藏状态 $s$ ，输出一系列权重分布。现在，对于 Attention 层，我们将输入的 $\{x_i\}$ 整合成一个矩阵 $X[N_X \times D_X]$ （ $N$ 为数量， $D$ 为维度），在先前的 RNN 中，这个序列代表 Decoder 要回看的序列，在这里称之为 Data vectors。而隐藏状态 $s$ 负责在 Data vectors 中查看哪些部分是重要的，我们称之为 Query vector。在 Attention 层中，使用多个 Query vector，也就是矩阵 $Q[N_Q \times D_Q]$ 。Attention 层最终给出输出 Output vector。

接下来要做的是对 Data vectors 做键值分离，得到 Key vectors 和 Value vectors。Key vectors 用于计算 Query vector 与 Data vectors 的相关程度，而 Value vectors 则是最终输出的内容（如在翻译任务中，键 $K$ 中可能存储的是源语言的语法结构信息，而值 $V$ 中存储的是源语言的语义信息）。我们将 Data vectors 通过线性变换得到 Key vectors 和 Value vectors，最终得到如下的结构，也就是交叉注意力层（Cross-Attention Layer）：

KV 分离的好处在于可以让模型更灵活地学习不同的表示，Key vectors 可以专注于捕捉输入序列中的结构信息，高效地计算注意力权重；而 Value vectors 则可以携带更多的语义信息，这样在计算注意力权重后，模型可以更有效地利用这些信息来生成更准确的输出，避免计算注意力和提供信息两个任务的冲突。

值得注意的是，在计算注意力分数 $E$ 时，通常使用缩放因子 $\sqrt{d_k}$ ，其中 $d_k$ 是 Key/Query 投影后的维度。这样做是为了防止点积方差随维度增大而变大，导致 Softmax 过早饱和（如个别项注意力接近 1 而其他项接近 0）。

自注意力层

更加常见的用法是自注意力层（Self-Attention Layer），在这种情况下，Query vectors、Key vectors 和 Value vectors 都是从同一组 Data vectors 中生成的。自注意力层允许模型在处理输入序列时，能够关注序列中的不同位置，从而捕捉到输入序列内部的依赖关系。这对于理解句子结构和语义非常重要。

自注意力中的 $N$ 是一致的，在进行 QKV 分离时将 $D_{in}$ 变成 $D_{out}$ ，最终输出的维度也是 $D_{out}$ 。在自注意力层中，每个位置的输出都是通过对输入序列中所有位置进行加权求和得到的，这些权重由 Query vectors 和 Key vectors 之间的相关程度决定。这使得模型能够捕捉到输入序列中的长距离依赖关系，从而更好地理解句子结构和语义。

在自注意力机制中，对于每一个输入向量的处理之间实际上是独立的（对输入的各个维度分别做 QKV 变换，而 Softmax 计算不受输入顺序的影响），交换两个输入的位置，输出也会对应地交换。即具有排列等变性（Permutation Equivariance）。这说明 Self-Attention 并没有内在的顺序信息，打乱之后原来序列的顺序信息就丢失了。为了让模型能够捕捉到输入序列的顺序信息，可以在输入中加入位置编码（Positional Encoding），这个向量是位置的固定函数。

另外，在某些任务中，我们希望模型在生成输出时只能关注输入序列中某些特定的位置（例如在语言模型中，生成下一个词时只能关注前面的词）。为此，可以使用掩码 Mask 来限制注意力机制的计算范围，使得模型只能关注特定位置的信息。比较简单的方法是将每一个输入对后续输入的注意力分数设置为负无穷，这样在 Softmax 计算时这些位置的权重就会变为 0，从而实现对未来信息的屏蔽。

多头自注意力层

多头的意思是我们使用多个独立的自注意力机制来处理输入序列，每个注意力机制被称为一个头（Head）。每个头都有自己的一套 QKV 变换参数，这样可以让模型从不同的子空间中学习输入序列的不同方面的信息。最终，将所有头的输出进行拼接和线性变换，得到最终的输出。多头机制被认为是自注意力机制良好工作的关键。人们推测，这种机制能够让自注意力网络对不良的初始值设定更加鲁棒。

看这个图对大小的变化感觉不大直观，找到另一张图：

把下面这一坨计算放到图中看就很直观了。

如果输入的数量是 $N$ ，那么随着 $N$ 增加，显式存储注意力矩阵的空间复杂度是 $O(N^2)$ 。例如 $N=100k$ 、注意力头数 $H=64$ ，若为每个 head 存储一个 $N\times N$ 的注意力矩阵，即使用 fp16 也会达到 TB 量级，普通 GPU 无法承受。

为了解决这个问题，可以使用 Flash Attention 方法，通过分块计算和内存层级优化来加速 Transformer 自注意力机制的算法，核心思想是将大型注意力矩阵拆分成小块，在 GPU 高速缓存中逐块计算并迭代融合结果，避免存储完整 $QK^T$ 中间矩阵。这种方法将显存占用从 $O(N^2)$ 降至 $O(N)$ ，显著提升长序列处理效率。另外还有 KV Cache 和 Linear Attention 等 Attention 的优化方案，可以参考：ADS13：KV Cache

序列处理方法的对比

RNN 的优势是按时间顺序自然处理序列，内存占用与序列长度关系较温和，但缺点是难以并行，长距离依赖需要穿过很多时间步。CNN 可以并行处理不同位置，并通过堆叠卷积扩大感受野，但远距离信息需要多层传播。Self-Attention 则让任意两个位置在一层内直接交互，长距离建模能力强、并行性好；代价是注意力矩阵随序列长度呈 $O(N^2)$ 增长。

Transformer

Transformer 可以理解为“以 Self-Attention 为核心的序列建模架构”。它把 token 先映射成向量，再反复经过 Attention 和 MLP 两类子层：Attention 负责让不同位置交换信息，MLP 负责对每个位置的特征做非线性变换。残差连接和归一化层则负责稳定深层训练。

Transformer Block

在 Transformer Block 中，输入先进入 Self-Attention 层计算输入向量之间的关系（输入向量在 Transformer 块中只通过 Attention 层进行交互），其输出经过残差连接后进行层标准化。接下来，每个输出向量分别独立地进入各个 MLP 中，再经过一次残差连接和层标准化得到最终的输出。

Transformer 具有很强的可扩展性和并行性。另外，其核心计算部分只有六个矩阵乘法,其中 4 个来自 Attention 层（QKV 变换和输出变换），2 个来自 MLP 层（升维和降维）。因此，Transformer 的计算效率较高，适合处理大规模数据。

Transformer 就是一系列 Transformer Block 的堆叠。当用在 LLM 中，加上了输入的 Embedding 层和输出的线性变换层。

其中词汇表大小（Vocab size）指的是 NLP 中模型能识别和处理的唯一单词/子词数量，也就是所有可能的输入和输出 token 的总数。嵌入操作相当于在词典里进行查询，将每个 token 映射到一个高维空间中，得到一个固定维度的向量表示（ $[L] \to [L \times D_{in}]$ ）。输出层的线性变换则是将 Transformer Block 的输出向量映射到词汇表大小的维度（ $[L \times D_{in}] \to [L \times V]$ ），在整个词汇表的范围中进行预测。

ViT

在图像分类任务中，通过添加位置编码解决 Transformer 不能感知位置的问题，这样保留位置信息的同时，每个图像块都能看到其他所有的图像块上的信息。对于输入图像，首先将其划分为固定大小的图像块（patches），每个图像块被展平并通过线性变换映射到一个高维空间中，得到一个固定维度的向量表示。然后，将这些向量作为 Transformer 的输入，经过多个 Transformer Block 的处理后，得到一个全局的图像表示，最后通过一个线性变换层将这个表示映射到类别空间中进行分类。

Transformer 改进

Transformer 的基础结构很规整，因此很多改进都围绕三个目标展开：让深层训练更稳定，让同等计算量下表达能力更强，以及让长序列或大模型训练更省显存和计算。下面这些改动都没有改变“Attention + MLP + 残差 + 归一化”的主干，只是在归一化位置、前馈层形式或参数容量上做工程化增强。

Pre-Norm Transformer

在之前的 Transformer 架构中，标准化层被放置在残差连接之后，即 Post-LN。由于 LN 放在残差外，即使残差为 0，标准化层仍然会对输入进行缩放和平移，导致模型并不能通过学习恒等映射真正“跳过”这一层的处理而引发梯度问题。为了解决这个问题，将标准化层移到 MLP 和 Attention 之前，这样当残差为 0 时，模型可以通过学习恒等映射来跳过这一层的处理，让训练更加稳定。