【收藏必看】大模型长序列处理优化：Transformer注意力机制的4大突破性技术

原创于 2025-12-29 11:24:51 发布 · 434 阅读

CC 4.0 BY-SA版权

文章标签：

#人工智能 #学习 #transformer #语言模型 #ai #程序员 #深度学习

在 Transformer 结构中，自注意力机制的计算和存储复杂度与序列长度成平方关系，这导致其在长序列建模时消耗大量的显存和计算资源。因此，如何在保证建模能力的前提下优化注意力机制的时空复杂度、提升计算效率，成为大语言模型的重要研究方向。当前的优化思路主要包括两类：一是通过近似注意力降低计算和内存需求，例如稀疏近似和低秩近似方法；二是结合硬件特性，探索更高效的实现方式。

1. 稀疏注意力机制

在对训练好的 Transformer 的注意力矩阵进行分析时发现，矩阵中存在大量稀疏结构，因此可以通过减少 Query-Key 的计算对数来降低复杂度。这类方法被称为稀疏注意力（Sparse Attention）。根据稀疏化的依据不同，稀疏注意力可分为基于位置和基于内容两大类。

其中，基于位置的稀疏注意力常见的形式包括：

全局注意力（Global Attention）：引入少量全局节点，用于建模长距离依赖。
带状注意力（Band Attention）：利用数据的局部性，仅让 Query 与相邻若干节点交互。
膨胀注意力（Dilated Attention）：类似于 CNN 中的空洞卷积，通过设置间隔扩大感受野。
随机注意力（Random Attention）：通过随机采样实现非局部信息交互。
局部块注意力（Block Local Attention）：将序列划分为多个不重叠块，信息交互限制在块内。

在实际应用中，稀疏注意力机制往往采用上述几种方式的组合。例如：

Star-Transformer：结合带状注意力和全局注意力，仅保留一个全局节点，并采用宽度为 3 的带状注意力。非相邻节点通过全局节点间接连接，相邻节点则直接相连。
Longformer：结合带状注意力和内部全局节点注意力，并在部分层中使用膨胀窗口注意力，以扩大感受野而不增加计算量。
ETC（Extended Transformer Construction）：结合带状注意力和外部全局节点注意力，并引入掩码机制以处理结构化输入，同时采用对比预测编码（CPC）进行预训练。
BigBird：采用带状注意力、全局注意力和随机注意力的混合方式来近似全连接注意力。研究表明，稀疏编码器和解码器的组合能够模拟任意图灵机，这也从理论上解释了稀疏注意力模型取得良好性能的原因。

2. FlashAttention

在 NVIDIA GPU 中，不同类型的显存具有差异化的速度、容量和访问权限，这取决于它们位于 GPU 芯片内部还是板卡上的 RAM 存储芯片。总体来看，GPU 的显存可以分为六类：全局内存、本地内存、共享存储、寄存器、常量内存和纹理内存（如下图所示）。其中，全局内存、本地内存、共享存储和寄存器均支持读写操作。

全局内存与本地内存：二者基于板卡上的高带宽显存（HBM），容量大。全局内存可被所有线程访问，本地内存仅限当前线程使用。例如，NVIDIA H100 的全局内存容量高达 80GB，理论带宽为 3.35TB/s，但当所有线程并发访问时，实际可用带宽远低于峰值。
共享存储与寄存器：位于 GPU 芯片上，容量较小，但速度远快于全局内存。共享存储支持线程块内并行访问，而寄存器则仅限单线程使用。在 H100 中，每个线程块可分配的共享存储空间为 228KB。

在自注意力计算中，显存带宽常成为性能瓶颈。传统的 Attention 计算流程为：S = QK, P = Softmax(S), O = P V

该过程中，Q 和 K 首先从全局内存中读取，计算得到的 S 写回全局内存，再次读取进行 Softmax 得到 P，随后写入全局内存，最后读取 P 和 V 计算输出 O。频繁的全局内存读写严重占用带宽，而 GPU 的算力远超内存带宽，导致整体效率受限。

FlashAttention 针对这一瓶颈进行了优化。其核心思想是：充分利用 GPU 芯片上的高速共享存储（SRAM），最大限度减少对 HBM 的读写。为此，FlashAttention 需要在不完整加载整个输入矩阵的情况下完成 Softmax 计算，并在反向传播中避免存储中间注意力矩阵。

在标准 Attention 中，Softmax 按行计算，必须先得到 S 的完整一行结果，再与 V 做矩阵乘。而在 FlashAttention 中，输入被划分为小块，算法在块内进行多次迭代传递，以增量的方式完成 Softmax。这样无需将 S 和 P 整体写入全局内存，而是仅存储前向传播的 Softmax 归一化因子，反向传播时在片上快速重算注意力值。

这种方式极大减少了全局内存访问量，显著提升了效率。虽然需要额外的计算开销（FLOPS 增加），但总体运行速度更快、显存占用更低。

3. 多查询注意力

多查询注意力（Multi-Query Attention, MQA） 是多头注意力的一种改进形式。与传统多头注意力不同，MQA 中的多个注意力头共享同一组键（Key）和值（Value）矩阵，每个头仅保留独立的查询（Query）参数。这样一来，键和值矩阵只需存储一份，大幅减少了显存开销，从而显著提升计算和存储效率。

需要注意的是，MQA 改变了注意力机制的结构，因此模型通常需要在训练初期就引入该机制。不过，已有研究表明，可以通过对已训练模型进行微调来增加 MQA 的支持，而且只需大约 5% 的原始训练数据即可取得较好效果。

目前，已有不少开源模型采用了多查询注意力，例如 Falcon、SantaCoder 和 StarCoder 等。

多头潜在注意力

多头潜在注意力（Multi-Head Latent Attention，MLA） 是 DeepSeek-V2 提出的注意力优化方法，核心目标是减少 KV 缓存 的大小，从而缓解推理过程中的显存和通信瓶颈。

在传统多头注意力（MHA）中，每个注意力头都需要单独存储完整的键（Key）和值（Value），这会占用大量显存。MLA 的做法是：先对键和值进行低秩压缩，得到一个更小的潜在表示，在推理时只需要缓存这个压缩后的表示，再通过上投影恢复即可。这种方式能在显著减少 KV 缓存的同时，保持模型的表达能力。

进一步优化时，查询（Query）也可以进行低秩压缩，以减少训练过程中的激活内存。这样，MLA 在推理和训练阶段都比传统注意力更高效。

研究表明，MLA 的表现力优于 组查询注意力（GQA）。在相同的 KV 缓存大小下，MLA 能展现出更大的多样性，而 GQA 因为组内的头是复制的，容易导致信息冗余。基于这一点，学者们提出了 TransMLA 方法，可以将现有的 GQA 模型（如 LLaMA、Qwen、Mixtral）转换为 MLA 模型，并通过少量训练进一步提升性能。