CS231n笔记10：Attention与Transformer

P1：Atention机制

引入

课上是使用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_Q}$ （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时注意力矩阵大小会增长到1.192TB，GPU无法存储这么多的数据。

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

序列处理方法的对比

前面忘了，后面忘了，反正 Attention is all you need。

P2：Transformer

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改进

Pre-Norm Transformer

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