在当今主流大模型LLM架构中**注意力机制Attention Mechanism**是最核心的计算模块之一。无论是文本生成模型、视觉模型还是多模态模型几乎都建立在以注意力为基础的结构之上。自从 Ashish Vaswani 等人在 2017 年提出 Google 的论文《Attention Is All You Need》以来Transformer 架构已成为大模型的标准范式。本文将系统总结注意力机制的原理、结构演化与工程优化方向。一、为什么需要注意力机制在传统序列模型如 RNN、LSTM中模型需要将前文压缩到一个固定长度的隐状态中这种“信息瓶颈”会导致长距离依赖难以建模梯度消失或爆炸并行计算效率低注意力机制的核心思想是在处理当前 token 时动态地关注序列中的不同位置并为其分配不同权重。换句话说模型不再依赖单一隐向量而是对所有历史信息进行加权聚合。二、注意力的基本计算公式1. Query / Key / Value注意力机制的输入通常是三组向量Query (Q)当前查询向量Key (K)用于匹配的索引向量Value (V)实际信息载体2. Scaled Dot-Product Attention标准注意力公式为Attention(Q,K,V)softmax(QKTdk)V Attention(Q,K,V)softmax\left(\frac{QK^T}{\sqrt{d_k}}\right)VAttention(Q,K,V)softmax(dkQKT)V解释Q 与 K 做点积计算相似度除以dk\sqrt{d_k}dk做缩放防止梯度过大经 softmax 得到注意力权重对 V 加权求和本质上这是一个可学习的加权平均机制。PyTorch 实现基础版importtorchimporttorch.nn.functionalasFdefscaled_dot_product_attention(Q,K,V,maskNone): Q, K, V: (batch, heads, seq_len, dim) d_kQ.size(-1)# 1. 计算相似度scorestorch.matmul(Q,K.transpose(-2,-1))/torch.sqrt(torch.tensor(d_k,dtypetorch.float32))# 2. mask用于因果或paddingifmaskisnotNone:scoresscores.masked_fill(mask0,float(-inf))# 3. softmaxattn_weightsF.softmax(scores,dim-1)# 4. 加权求和outputtorch.matmul(attn_weights,V)returnoutput,attn_weights三、自注意力Self-Attention在大模型中最重要的是Self-Attention。特点Q、K、V 来自同一个序列每个 token 都可以“看到”其他 token能直接建模长距离依赖例如“The animal didn’t cross the street becauseitwas tired.”模型可以通过注意力将it指向animal。四、多头注意力Multi-Head Attention单头注意力可能只捕获某一类关系因此 Transformer 引入MultiHead(Q,K,V)Concat(head1,...,headh)WO MultiHead(Q,K,V)Concat(head_1,...,head_h)W^OMultiHead(Q,K,V)Concat(head1,...,headh)WO每个 head 在不同的线性空间中计算注意力。优势不同 head 关注不同语义模式增强表达能力类似“多视角观察”Multi-Head Attention 实现importtorch.nnasnnclassMultiHeadAttention(nn.Module):def__init__(self,d_model,num_heads):super().__init__()assertd_model%num_heads0self.num_headsnum_heads self.head_dimd_model//num_heads self.qkv_projnn.Linear(d_model,3*d_model)self.out_projnn.Linear(d_model,d_model)defforward(self,x,maskNone):B,T,Cx.shape qkvself.qkv_proj(x)qkvqkv.reshape(B,T,3,self.num_heads,self.head_dim)qkvqkv.permute(2,0,3,1,4)Q,K,Vqkv[0],qkv[1],qkv[2]out,attnscaled_dot_product_attention(Q,K,V,mask)outout.transpose(1,2).reshape(B,T,C)returnself.out_proj(out)五、注意力优化方向所有注意力优化本质都围绕三件事少算减少计算量少存减少显存占用让长度变长支持更长上下文方向一少算降低计算复杂度 目标把 Attention 的计算复杂度从O(n2)→O(nlogn) 或 O(n) O(n^2) \rightarrow O(n\log n) \text{ 或 } O(n)O(n2)→O(nlogn)或O(n)1️⃣ 稀疏注意力Sparse Attention代表模型LongformerBigBird核心原理不再计算完整 N×N 注意力矩阵而是只计算局部窗口块内 attention少量全局 token 原理图标准 Attention满矩阵█ █ █ █ █ █ █ █ █ █ █ █ █ █ █ █ █ █ █ █ █ █ █ █ █稀疏 Attention带状 全局█ █ . . . █ █ █ . . . █ █ █ . . . █ █ █ . . . █ █.表示不计算。✅ 优点理论复杂度下降显存占用减少❌ 缺点远距离依赖能力下降稀疏模式难设计GPU 不友好 落地情况文档任务、长文本建模❌ 主流 LLM 已基本不用纯稀疏结构2️⃣ 线性注意力Linear Attention核心原理把softmax(QKT)V softmax(QK^T)Vsoftmax(QKT)V改写为(Qϕ(K)T)(ϕ(V)) (Q\phi(K)^T)(\phi(V))(Qϕ(K)T)(ϕ(V))避免构造 N×N 矩阵。 原理图标准流程Q × Kᵀ → N×N矩阵 → softmax → ×V线性 Attentionφ(K)ᵀ × φ(V) → 中间聚合 Q × 上述结果 → 输出等价于先把序列维度“压掉”✅ 优点理论 O(n)超长序列潜力大❌ 缺点softmax 被近似精度下降明显训练不稳定 落地情况学术研究活跃❌ 商业 LLM 几乎不用方向二少存降低显存 IO这是当前工业界真正成功的方向。1️⃣ FlashAttention提出团队Stanford University核心原理不改变数学结构而是改变执行方式分块计算不存完整 attention matrix在 SRAM 内做 softmaxFlashAttention 的“图”和前面完全不一样因为它结构没变只是计算方式变了逻辑结构和标准 Attention 一样Q × Kᵀ → softmax → × V物理执行图关键HBM (显存) ├─ Load Q block ├─ Load K block ├─ Load V block └─ Compute softmax V inside SRAM ↓ Write OutputBlocked Attention 示意[Q1] x [K1] → partial softmax → partial output [Q1] x [K2] → partial softmax → accumulate [Q1] x [K3] → ...减少了什么❌ 没减少计算量✅ 大幅减少显存 IO✅ 不存整张 attention matrix为什么是王者数学完全等价精度 0 损失GPU 友好落地PyTorch 2.x 默认所有主流 LLM2️⃣ KV Cache推理期核心优化无 KV Cache每步都重算Step t: [Token1 ... Token t] → Q K V → Attention有 KV CacheCache: K1 K2 K3 ... K(t-1) V1 V2 V3 ... V(t-1) Step t: Qt × [K1...K(t-1)] → output示意图┌───────────────┐ New Q ──▶│ KV Cache │──▶ Attention │ K1 K2 ... Kt │ │ V1 V2 ... Vt │ └───────────────┘减少了什么❌ 历史 K/V 不再重复计算代价KV cache 显存线性增长落地所有 GPT 类模型推理必备3️⃣ MQA / GQA减少 KV 的“宽度”代表模型LLaMA标准 Multi-HeadQ1 K1 V1 Q2 K2 V2 Q3 K3 V3 Q4 K4 V4MQA极端Q1 ┐ Q2 ├── shared K, V Q3 ┤ Q4 ┘GQA折中(Q1 Q2) ── K1 V1 (Q3 Q4) ── K2 V2减少了什么KV Cache 从heads × seq_len × dim变成groups × seq_len × dim代价轻微表达能力下降落地LLaMA 2 / 3长上下文推理必用三、方向三让长度变长扩展上下文1️⃣ RoPE 外推原理使用旋转位置编码使位置关系具有外推能力。 示意标准位置编码pos1 pos2 pos3 ...RoPE向量按角度旋转 θ pos × 频率角度可扩展。 落地几乎所有开源 LLM2️⃣ 分布式 / Ring Attention原理把序列分布在多 GPU 上GPU0: tokens 0–8k GPU1: tokens 8k–16k GPU2: tokens 16k–24k通过 ring 传递 K/V。 示意图GPU0 → GPU1 → GPU2 → GPU3 ↑ ↓ └──────── ring ────────┘ 落地100k context企业级长文本系统大模型中的注意力优化本质上都围绕三个方向展开少算—— 降低计算复杂度少存—— 降低显存占用和内存访问变长—— 在资源可控的前提下支持更长上下文进一步抽象来看这些优化都是在平衡四种资源约束计算FLOPs、显存Memory、带宽IO、并行结构Parallelism。在当前硬件条件下带宽和显存往往比算力更稀缺因此工程上最成功的优化通常集中在“少存”和“少搬运”上而不是单纯减少理论计算量。结论注意力机制的演进已经从“算法问题”逐渐转向“系统工程问题”。 大模型基础123- 注意力机制，从数学原理到资源优化框架
在当今主流大模型LLM架构中**注意力机制Attention Mechanism**是最核心的计算模块之一。无论是文本生成模型、视觉模型还是多模态模型几乎都建立在以注意力为基础的结构之上。自从 Ashish Vaswani 等人在 2017 年提出 Google 的论文《Attention Is All You Need》以来Transformer 架构已成为大模型的标准范式。本文将系统总结注意力机制的原理、结构演化与工程优化方向。一、为什么需要注意力机制在传统序列模型如 RNN、LSTM中模型需要将前文压缩到一个固定长度的隐状态中这种“信息瓶颈”会导致长距离依赖难以建模梯度消失或爆炸并行计算效率低注意力机制的核心思想是在处理当前 token 时动态地关注序列中的不同位置并为其分配不同权重。换句话说模型不再依赖单一隐向量而是对所有历史信息进行加权聚合。二、注意力的基本计算公式1. Query / Key / Value注意力机制的输入通常是三组向量Query (Q)当前查询向量Key (K)用于匹配的索引向量Value (V)实际信息载体2. Scaled Dot-Product Attention标准注意力公式为Attention(Q,K,V)softmax(QKTdk)V Attention(Q,K,V)softmax\left(\frac{QK^T}{\sqrt{d_k}}\right)VAttention(Q,K,V)softmax(dkQKT)V解释Q 与 K 做点积计算相似度除以dk\sqrt{d_k}dk做缩放防止梯度过大经 softmax 得到注意力权重对 V 加权求和本质上这是一个可学习的加权平均机制。PyTorch 实现基础版importtorchimporttorch.nn.functionalasFdefscaled_dot_product_attention(Q,K,V,maskNone): Q, K, V: (batch, heads, seq_len, dim) d_kQ.size(-1)# 1. 计算相似度scorestorch.matmul(Q,K.transpose(-2,-1))/torch.sqrt(torch.tensor(d_k,dtypetorch.float32))# 2. mask用于因果或paddingifmaskisnotNone:scoresscores.masked_fill(mask0,float(-inf))# 3. softmaxattn_weightsF.softmax(scores,dim-1)# 4. 加权求和outputtorch.matmul(attn_weights,V)returnoutput,attn_weights三、自注意力Self-Attention在大模型中最重要的是Self-Attention。特点Q、K、V 来自同一个序列每个 token 都可以“看到”其他 token能直接建模长距离依赖例如“The animal didn’t cross the street becauseitwas tired.”模型可以通过注意力将it指向animal。四、多头注意力Multi-Head Attention单头注意力可能只捕获某一类关系因此 Transformer 引入MultiHead(Q,K,V)Concat(head1,...,headh)WO MultiHead(Q,K,V)Concat(head_1,...,head_h)W^OMultiHead(Q,K,V)Concat(head1,...,headh)WO每个 head 在不同的线性空间中计算注意力。优势不同 head 关注不同语义模式增强表达能力类似“多视角观察”Multi-Head Attention 实现importtorch.nnasnnclassMultiHeadAttention(nn.Module):def__init__(self,d_model,num_heads):super().__init__()assertd_model%num_heads0self.num_headsnum_heads self.head_dimd_model//num_heads self.qkv_projnn.Linear(d_model,3*d_model)self.out_projnn.Linear(d_model,d_model)defforward(self,x,maskNone):B,T,Cx.shape qkvself.qkv_proj(x)qkvqkv.reshape(B,T,3,self.num_heads,self.head_dim)qkvqkv.permute(2,0,3,1,4)Q,K,Vqkv[0],qkv[1],qkv[2]out,attnscaled_dot_product_attention(Q,K,V,mask)outout.transpose(1,2).reshape(B,T,C)returnself.out_proj(out)五、注意力优化方向所有注意力优化本质都围绕三件事少算减少计算量少存减少显存占用让长度变长支持更长上下文方向一少算降低计算复杂度 目标把 Attention 的计算复杂度从O(n2)→O(nlogn) 或 O(n) O(n^2) \rightarrow O(n\log n) \text{ 或 } O(n)O(n2)→O(nlogn)或O(n)1️⃣ 稀疏注意力Sparse Attention代表模型LongformerBigBird核心原理不再计算完整 N×N 注意力矩阵而是只计算局部窗口块内 attention少量全局 token 原理图标准 Attention满矩阵█ █ █ █ █ █ █ █ █ █ █ █ █ █ █ █ █ █ █ █ █ █ █ █ █稀疏 Attention带状 全局█ █ . . . █ █ █ . . . █ █ █ . . . █ █ █ . . . █ █.表示不计算。✅ 优点理论复杂度下降显存占用减少❌ 缺点远距离依赖能力下降稀疏模式难设计GPU 不友好 落地情况文档任务、长文本建模❌ 主流 LLM 已基本不用纯稀疏结构2️⃣ 线性注意力Linear Attention核心原理把softmax(QKT)V softmax(QK^T)Vsoftmax(QKT)V改写为(Qϕ(K)T)(ϕ(V)) (Q\phi(K)^T)(\phi(V))(Qϕ(K)T)(ϕ(V))避免构造 N×N 矩阵。 原理图标准流程Q × Kᵀ → N×N矩阵 → softmax → ×V线性 Attentionφ(K)ᵀ × φ(V) → 中间聚合 Q × 上述结果 → 输出等价于先把序列维度“压掉”✅ 优点理论 O(n)超长序列潜力大❌ 缺点softmax 被近似精度下降明显训练不稳定 落地情况学术研究活跃❌ 商业 LLM 几乎不用方向二少存降低显存 IO这是当前工业界真正成功的方向。1️⃣ FlashAttention提出团队Stanford University核心原理不改变数学结构而是改变执行方式分块计算不存完整 attention matrix在 SRAM 内做 softmaxFlashAttention 的“图”和前面完全不一样因为它结构没变只是计算方式变了逻辑结构和标准 Attention 一样Q × Kᵀ → softmax → × V物理执行图关键HBM (显存) ├─ Load Q block ├─ Load K block ├─ Load V block └─ Compute softmax V inside SRAM ↓ Write OutputBlocked Attention 示意[Q1] x [K1] → partial softmax → partial output [Q1] x [K2] → partial softmax → accumulate [Q1] x [K3] → ...减少了什么❌ 没减少计算量✅ 大幅减少显存 IO✅ 不存整张 attention matrix为什么是王者数学完全等价精度 0 损失GPU 友好落地PyTorch 2.x 默认所有主流 LLM2️⃣ KV Cache推理期核心优化无 KV Cache每步都重算Step t: [Token1 ... Token t] → Q K V → Attention有 KV CacheCache: K1 K2 K3 ... K(t-1) V1 V2 V3 ... V(t-1) Step t: Qt × [K1...K(t-1)] → output示意图┌───────────────┐ New Q ──▶│ KV Cache │──▶ Attention │ K1 K2 ... Kt │ │ V1 V2 ... Vt │ └───────────────┘减少了什么❌ 历史 K/V 不再重复计算代价KV cache 显存线性增长落地所有 GPT 类模型推理必备3️⃣ MQA / GQA减少 KV 的“宽度”代表模型LLaMA标准 Multi-HeadQ1 K1 V1 Q2 K2 V2 Q3 K3 V3 Q4 K4 V4MQA极端Q1 ┐ Q2 ├── shared K, V Q3 ┤ Q4 ┘GQA折中(Q1 Q2) ── K1 V1 (Q3 Q4) ── K2 V2减少了什么KV Cache 从heads × seq_len × dim变成groups × seq_len × dim代价轻微表达能力下降落地LLaMA 2 / 3长上下文推理必用三、方向三让长度变长扩展上下文1️⃣ RoPE 外推原理使用旋转位置编码使位置关系具有外推能力。 示意标准位置编码pos1 pos2 pos3 ...RoPE向量按角度旋转 θ pos × 频率角度可扩展。 落地几乎所有开源 LLM2️⃣ 分布式 / Ring Attention原理把序列分布在多 GPU 上GPU0: tokens 0–8k GPU1: tokens 8k–16k GPU2: tokens 16k–24k通过 ring 传递 K/V。 示意图GPU0 → GPU1 → GPU2 → GPU3 ↑ ↓ └──────── ring ────────┘ 落地100k context企业级长文本系统大模型中的注意力优化本质上都围绕三个方向展开少算—— 降低计算复杂度少存—— 降低显存占用和内存访问变长—— 在资源可控的前提下支持更长上下文进一步抽象来看这些优化都是在平衡四种资源约束计算FLOPs、显存Memory、带宽IO、并行结构Parallelism。在当前硬件条件下带宽和显存往往比算力更稀缺因此工程上最成功的优化通常集中在“少存”和“少搬运”上而不是单纯减少理论计算量。结论注意力机制的演进已经从“算法问题”逐渐转向“系统工程问题”。