长序列中 Transformers 的高级注意力优化机制

注意力机制 (Attention Mechanism) 作为 Transformer 架构的核心组件，通过动态分配权重帮助模型捕捉长距离依赖关系，使每个词元能够聚焦输入序列中最相关的信息。

但是，随着序列长度的不断增加，每个词元需要与呈平方级增长的其他词元竞争注意力得分，导致注意力分数被过度稀释，这种稀释效应会削弱模型对关键特征的聚焦能力，使得生成的上下文表示逐渐趋于平均化，尤其损害远距离词元间的语义关联性。与此同时，长序列中大量非关键信息的累积形成噪声干扰，进一步分散注意力资源的有效分配，最终导致模型在信息筛选和语义整合上面临双重挑战。

本篇博客将简要介绍面向长序列建模的注意力机制的高级优化技巧，事实上，这些优化技巧已经被广泛用于现如今的 LLM 的训练和推理之中。本博客仅起到简单归纳和提纲挈领的作用。

序列长度的影响

为了理解较长的序列是如何稀释注意力得分和增加噪音的，我们需要深入研究 Transformers 等模型中使用的注意力机制的数学原理。Transformer 中的注意机制基于缩放点积注意，其定义为：

$$Attention(Q,K,V) = \textsf{softmax}(\frac{QK^T}{\sqrt{d_k}})V$$

$Q$(Query)，$K$(Key) 和 $V$(Value) 是由输入嵌入导出的矩阵。$d_k$ 是向量的维数，用于缩放点积以防止可能破坏 softmax 函数稳定的大值。考虑一个简单的例子，其中Q和K是相同的，每个元素都同样相关:

$$
QK^T = \begin{bmatrix}
q_1 \cdot k_1 & q_1 \cdot k_2 & \cdots & q_1 \cdot k_n \\
q_2 \cdot k_1 & q_2 \cdot k_2 & \cdots & q_2 \cdot k_n \\
\vdots & \vdots & \ddots & \vdots \\
q_n \cdot k_1 & q_n \cdot k_2 & \cdots & q_n \cdot k_n \\
\end{bmatrix}
$$

随序列长度 $n$ 的增加，矩阵 $QK^T$ (在应用 softmax 之前) 中每一行的总和增加，因为添加了更多的项，这可能会导致这样一种情况，任何单个 $k_j$ 对给定 $q_i$ 的影响都会减弱，因为它更接近于平均值：

$$
\textsf{softmax}(\frac{QK^T}{\sqrt{d_k}})_{ij} \approx \frac{\textsf{exp}(\frac{q_i \cdot k_j}{\sqrt{d_k}})}{\sum_{l=1}^{n}\textsf{exp}(\frac{q_i \cdot k_l}{\sqrt{d_k}})}
$$

• n 越大，分母越大，将注意力分散到更多的词元上。这种稀释降低了模型专注于最相关项的能力。
• 此外，较长的序列通常包含与正在处理的当前上下文不太相关的片段。这些不太相关或嘈杂的片段仍然会计算注意力机制中的点积。随着 $n$ 的增加，$q_i$ 与表示噪声或不太相关的信息的几个 $k_j$ 一致的概率也会增加。这种噪音影响了softmax 函数有效地优先考虑最相关的能力，从而降低了注意力驱动的上下文理解的整体质量。

减量优化方法

减量优化方式试图通过某种手段减轻计算量和任务量，使模型更容易注意到它应该注意到的东西。形象的比喻是：不给学生布置过多的课后作业，将他有限的注意力用在更有价值的地方。

局部敏感哈希 (Locality-Sensitive Hashing, LSH)

通过限制词元之间的交互数量来减少计算需求。将词元分散到桶中，仅计算桶内交互，从而简化注意力矩阵。每个词元被投影到一个由哈希函数定义的低维空间中，且注意力只在桶内计算：

$$h(q) = \lfloor\frac{q \cdot r}{w}\rfloor \qquad Attention(Q,K,V) = \underset{b \in \textsf{Buckets}}{\oplus} softmax(\frac{Q_bk_b^T}{\sqrt{d_k}})V_b$$

这种机制选择性地集中了模型的计算资源，将整体复杂度从 $O(n^2)$ 降低到 $O(n \log n)$。

低秩注意力 (Low-Rank Attention)

低秩注意力是一种优化注意力机制的方法，通过将注意力矩阵分解为低秩矩阵，有效地简化计算过程。低秩分解假设交互空间可以被更小的子空间有效捕获，减少了对完整 $n \times n$ 注意力计算的需要。

$$A \approx UV^T$$

这里的 $U$ 和 $V$ 是秩较低的矩阵，大大降低了复杂度，增强了跨长序列的注意力的可管理性。这样注意力的计算就变为：

$$Attention(Q,K,V) = \textsf{softmax}(\frac{QU(KV^T)^T}{\sqrt{d_k}})V$$

现在市面上的许多公司都利用 RoLA 来训练大模型，因为它有效地降低了计算量，将计算负载从 $O(n^2)$ 降低到了 $O(nk)$。

分段注意力 (Segmented Attention)

通过将输入序列分割成较小的片段，并在这些片段上独立地计算注意力，从而减少计算的复杂度和内存需求。(不知道它与后来出现的 Paged Attention 有没有关联？)

$$Attention(Q,K,V) = \textsf{softmax}(\frac{Q_sK_s^T}{\sqrt{d_k}})V_s$$

在每个独立的片段上执行标准的注意力机制。这意味着，每个片段内的元素只与同一片段内的其他元素进行交互，而不是与整个序列的元素进行交互。它通常会与 KV-Cache 搭配使用，每次缓存一个片段的 KV 对以提升推理时速度。

在某些实现中，可能会在分段注意力之后添加一个步骤，以整合不同片段间的信息，确保全局上下文不会丢失。这可以通过另一层跨段注意力或简单的序列级操作（如汇聚或连接）来实现。

层次化注意力 (Hierarchical Attention)

这种注意力模型通过在不同的层次上逐级应用注意力机制，能够更有效地捕捉数据中的结构和关联。数据被组织成多个层次，例如，在文本处理中，可以将数据结构化为字、词、句子和段落等层次。模型首先在较低层次上计算注意力，然后将计算结果传递到更高层次。

$$Attention(Q,K,V) = \textsf{softmax}(\frac{QG(K^T)}{\sqrt{d_k}})G(V)$$

部分人可能将层次化注意力和 GQA 搞混。两者的关键区别在于：层次化注意力的 $G(\cdot)$ 不是简单的分组，而是聚合了跨段或层的输出函数，它可能包含额外的转换，以细化跨层的注意力过程。

效能优化方法

效能优化方式试图引入某些新的机制，来增强注意力机制的长序列容忍度和抗噪声能力。形象的比喻是：给学生配备高效的学习工具，来让他拥有更专注的注意力，能够处理更多、更难的任务。

递归记忆 (Reccurssion Memory)

在 Transformers 中加入记忆可以让他们记住过去的计算，增强他们在较长序列中保持上下文的能力。

$$h_t = Attention(h_{t-1},[h_{t-1};h_t],[h_{t-1};h_t]$$

这种递归调用参考了 RNN 的思想，集成了历史信息，提供了一个连续的上下文线索，对于理解超出当前处理窗口的序列至关重要。

路由注意力机制 (Routed Attention Mechanism)

带有路由的注意力机制是一种高级的神经网络架构，通常用于处理具有复杂内部结构或需要精细调整信息流动的应用中。这种方法结合了注意力机制的灵活性和动态路由协议的决策过程，从而实现一种更为有效的信息处理方式：

$$
\begin{aligned}
P_i &= \textsf{softmax}(W_rh_i) \\
Attention_i &= \underset{j}{\sum} P_{ij} Attention(Q_i,K_j,V_j)
\end{aligned}
$$

这里的 $W_r$ 是一个路由矩阵，它决定了不同注意路径上的概率分布，允许模型根据输入 $h_i$ 的性质动态地调整其注意焦点。与发生在 FFN 部分的 MoE 不同，它是发生在注意力计算的部分的。

相对位置编码 (Position Embedding)

位置编码的历史相当悠久，最早可以追述到 Transformer 初代。在当时，Google 使用了基于三角函数周期变换的绝对位置编码来引入位置信息。而现在，人们更多地使用相对位置编码的方式。

$$Attention(Q,K,V) = \textsf{softmax}(QK^T + S_{rel})V$$

$S_{rel}$ 表示相对位置偏差，允许模型根据标记的相对距离和排列调整其注意力，增强其处理不同序列长度和结构的能力。现阶段，被许多人认可的相对位置编码方式，是 Jianlin Su 提出的 RoPE (旋转位置编码)，它以绝对位置编码的方式引入相对位置编码，简单而高效，被用于各种大模型之中。

总结的部分

本文系统梳理了多种高效的注意力优化机制，这些方法显著提升了 Transformer 架构的性能表现：在降低计算复杂度的同时，有效增强了模型处理长序列时的上下文建模与生成能力。通过优化注意力计算模式，模型的处理效率和计算性能都得到了大大的提升。