大语言模型中的注意力机制详解

一、什么是注意力机制？

注意力机制（Attention Mechanism）是 Transformer 架构的核心组成部分，它允许模型在处理信息时动态地关注输入序列中最相关的部分。这种机制极大地提升了模型对长距离依赖关系的理解能力，成为现代大语言模型（LLM）的关键技术之一。

核心公式回顾（Scaled Dot-Product Attention）：

Attention ( Q , K , V ) = softmax ( Q K T d k ) V \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V Attention(Q,K,V)=softmax(dk QKT)V

其中：

$Q$ ：Query 向量，表示当前 token 对其他 token 的关注度。
$K$ ：Key 向量，用于与其他 token 匹配。
$V$ ：Value 向量，代表 token 的语义信息。
$d_k$ ：Key 向量维度，用于防止点积过大导致 softmax 梯度消失。

二、常见的注意力机制详解

1. Self-Attention（自注意力）

原理：

Self-Attention 是最基础也是最常用的注意力机制。它通过 Query (Q)、Key (K) 和 Value (V) 向量之间的点积来衡量输入序列中不同位置之间的相关性。

公式：

Self-Attention ( X ) = softmax ( Q K T d k ) V \text{Self-Attention}(X) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V Self-Attention(X)=softmax(dk QKT)V

其中：

$Q = XW_Q$
$K = XW_K$
$V = XW_V$

特点：

全连接结构，建模全局依赖
计算复杂度为 $O(n\^2)$
可解释性强，支持可视化分析

应用场景：

BERT、GPT 系列、Transformer 等几乎所有 LLM 都使用 Self-Attention
适用于需要建模上下文语义的任务，如问答、翻译、摘要生成等

优点：

并行化能力强
支持长文本建模
可解释性强

缺点：

显存消耗高
对长序列效率低
容易过拟合局部模式

2. Cross-Attention（跨注意力）

原理：

Cross-Attention 发生在 Encoder-Decoder 架构中，用于解码器关注编码器输出的信息。

Query 来自 Decoder 输入，而 Key 和 Value 来自 Encoder 输出。

公式：

Cross-Attention ( Q d e c , K e n c , V e n c ) = softmax ( Q d e c K e n c T d k ) V e n c \text{Cross-Attention}(Q_{dec}, K_{enc}, V_{enc}) = \text{softmax}\left(\frac{Q_{dec} K_{enc}^T}{\sqrt{d_k}}\right)V_{enc} Cross-Attention(Qdec,Kenc,Venc)=softmax(dk QdecKencT)Venc

特点：

实现编码器与解码器之间的信息交互
在图像到文本生成、机器翻译中广泛应用

应用场景：

图像描述生成（如 CLIP、Flamingo）
机器翻译（Transformer）
多任务学习中的信息融合

优点：

有效结合两个模态的信息
提升模型对输入的理解能力

缺点：

仅适用于 Encoder-Decoder 结构
需要额外参数控制交互方式

3. Multi-Head Attention（多头注意力）

原理：

Multi-Head Attention 将输入映射到多个不同的子空间，并独立地进行多次 attention 操作，最后拼接结果并线性变换。

公式：

MultiHead ( Q , K , V ) = Concat ( head 1 , . . . , head h ) W O \text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, ..., \text{head}_h)W^O MultiHead(Q,K,V)=Concat(head1,...,headh)WO

其中：
head i = Attention ( Q W i Q , K W i K , V W i V ) \text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V) headi=Attention(QWiQ,KWiK,VWiV)

特点：

多个 head 学习不同角度的语义
增强表达能力和泛化能力

应用场景：

所有基于 Transformer 的模型（如 GPT、BERT、ChatGLM）
需要增强语义理解的场景

优点：

捕捉多样化的语义关系
提升模型表现
可扩展性强

缺点：

参数增加，训练成本上升
推理速度下降

4. Local Attention / Sliding Window Attention（局部/滑动窗口注意力）

原理：

Local Attention 只在固定大小的窗口内进行 attention 计算，而不是整个序列。

Sliding Window 则允许窗口在序列上滑动，从而覆盖更多上下文。

公式（简化）：

A i j = { q i k j T d k , ∣ i − j ∣ < w − ∞ , otherwise A_{ij} = \begin{cases} \frac{q_i k_j^T}{\sqrt{d_k}}, & |i - j| < w \\ -\infty, & \text{otherwise} \end{cases} Aij={dk qikjT,−∞,∣i−j∣<wotherwise

其中 $w$ ：窗口大小

特点：

减少显存占用
更适合处理长文本
局部建模，降低冗余计算

应用场景：

Longformer、BigBird、StripedHyena
文本摘要、对话系统、长文档建模

优点：

显存友好
适合长序列建模
推理速度快

缺点：

无法建模远距离依赖
信息传递受限于窗口长度

5. Sparse Attention（稀疏注意力）

原理：

Sparse Attention 仅选择部分位置进行 attention 计算，而非全连接。

例如：随机采样 key/value 或基于规则选择重要位置。

公式（示意）：

SparseAttention ( Q , K , V ) = ∑ j ∈ S ( i ) α i j v j \text{SparseAttention}(Q, K, V) = \sum_{j \in S(i)} \alpha_{ij} v_j SparseAttention(Q,K,V)=j∈S(i)∑αijvj

其中 $S(i)$ ：第 i 个 token 的稀疏采样集合

特点：

降低计算复杂度
支持更长的上下文
可灵活设计稀疏策略

应用场景：

BigBird、Performer、Longformer
处理超长文档、日志分析等

优点：

显存占用低
支持任意长度输入
可控性高（可设计稀疏模式）

缺点：

实现复杂
稀疏策略影响性能
可能丢失关键信息

6. FlashAttention（快速注意力）

原理：

FlashAttention 是 NVIDIA 提出的一种高效 attention 实现方式，通过优化内存访问顺序和减少中间缓存，显著提升 attention 的计算效率。

核心思想：

分块计算 attention（Tile-based）
利用 GPU 内存层级（Shared Memory + Registers）
减少 HBM 读写开销

公式不变，但实现方式优化：

FlashAttention ( Q , K , V ) = Efficient softmax and matmul on GPU \text{FlashAttention}(Q, K, V) = \text{Efficient softmax and matmul on GPU} FlashAttention(Q,K,V)=Efficient softmax and matmul on GPU

特点：

不改变 attention 行为
显存节省 2~3x
速度提升 2~3x

应用场景：

PyTorch、HuggingFace Transformers 已支持
所有需要加速 attention 的模型（LLaMA、ChatGLM、Phi-3 等）

优点：

快速推理
显存优化
与标准 attention 兼容

缺点：

依赖硬件（CUDA 加速）
对非 GPU 用户帮助有限

7. Grouped Query Attention (GQA)

原理：

将 Query 分成若干组，每组共享一组 Key 和 Value 向量。

公式（简化）：

GQA ( Q , K , V ) = Grouped version of MultiHead \text{GQA}(Q, K, V) = \text{Grouped version of MultiHead} GQA(Q,K,V)=Grouped version of MultiHead

即：

Query 按 head 分组
每组共用相同的 Key/Value

特点：

减少 KV Cache 占用
推理速度加快
是 MQA 的泛化形式

应用场景：

Llama 3、DeepSeek、InternLM
推理部署优化

优点：

显存占用低
推理速度快
性能接近 MHA

缺点：

表达能力略弱于 MHA
需要调整 head 分组策略

8. Multi-Query Attention (MQA)

原理：

所有 attention head 共享相同的 Key 和 Value 向量。

公式：

K = X W K , V = X W V ( 所有 head 共享 ) K = X W_K,\quad V = X W_V \quad (\text{所有 head 共享}) K=XWK,V=XWV(所有 head 共享)

Q i = X W Q i Q_i = X W_{Q_i} Qi=XWQi

特点：

Key/Value 不变，Query 每个 head 独立
极大减少 KV Cache 内存

应用场景：

Google PaLM、Llama 3、DeepSeek
大规模部署和推理服务

优点：

显存极低
推理速度快
适合流式生成

缺点：

表达能力受限
多样性降低

9. RoPE（Rotary Position Embedding）

原理：

RoPE 将位置信息以旋转矩阵的方式注入 Query 和 Key 向量中，使 attention score 自动考虑相对位置关系。

公式：

对于第 $i$ 个 token，其向量表示为：

q i ′ = R i q i , k j ′ = R j k j q_i' = R_i q_i,\quad k_j' = R_j k_j qi′=Riqi,kj′=Rjkj

最终 attention score 为：

( q i ′ ) T k j ′ (q_i')^T k_j' (qi′)Tkj′

特点：

支持外推至比训练时更长的上下文
不需位置嵌入层
保持 attention 的平移不变性

应用场景：

LLaMA、ChatGLM、InternLM、Phi-3、DeepSeek
长文本生成、代码生成等

优点：

支持任意长度
相对位置建模
简洁高效

缺点：

需要特定实现（如 CUDA kernel）
对某些任务效果可能不如 learned PE

10. ALiBi（Attention with Linear Biases）

原理：

ALiBi 使用线性偏置引导 attention score，不使用显式的位置嵌入。

公式：

score i j = q i k j T d k + m ⋅ ( j − i ) \text{score}_{ij} = \frac{q_i k_j^T}{\sqrt{d_k}} + m \cdot (j - i) scoreij=dk qikjT+m⋅(j−i)

其中 $m$ 是预设的缩放系数（通常取 1, 2, 4...），用于控制偏置强度。

特点：

不需要位置嵌入层
支持无限长上下文
位置信息以偏置形式引入

应用场景：

BLOOM、Falcon、Pythia 等
医疗、金融等需要长上下文的领域

优点：

支持任意长度输入
显存更低
无位置嵌入限制

缺点：

偏置系数需要调参
对短文本建模稍弱

三、注意力机制对比表

注意力类型	是否支持长文本	是否支持 KV Cache 优化	是否适合推理加速	是否支持位置建模
Self-Attention	❌	❌	❌	✅
Multi-Head	❌	❌	❌	✅
Sparse	✅	✅	✅	✅
FlashAttention	✅	✅	✅	✅
GQA	✅	✅	✅	✅
MQA	✅	✅	✅	✅
RoPE	✅	✅	✅	✅
ALiBi	✅	✅	✅	✅

四、总结

注意力机制是大语言模型的核心组成部分，决定了模型如何理解和生成文本。随着模型规模的扩大和应用场景的丰富，越来越多的注意力机制被提出，以解决长上下文、推理效率、内存瓶颈等问题。

选择合适的注意力机制，不仅能提升模型表现，还能显著优化推理速度和部署成本。