Transformer注意力机制演进：从缩放点积到高效变体

1. 从一张注意力热力图说起：Transformer到底在算什么

很多人第一次接触Transformer，是从一张花花绿绿的注意力热力图开始的。图上每个词对其它词的关注程度用深浅不一的色块表示，看起来挺直观，但真要问一句“这张图到底是怎么算出来的”，能说清楚的人就不多了。我自己当年也是先跑通了transformer代码，再回头补的transformer模型理论，中间踩了不少坑，所以这篇就把transformer架构里核心算法的演进脉络，按我自己的理解从头捋一遍。

先把结论摆前面：Transformer的核心从来不是“堆层数”，而是注意力机制这一套信息路由方案在不断被重新设计。从最初的缩放点积注意力，到多头自注意力，再到各种稀疏、线性、分组变体，本质上都在回答同一个问题——怎么在可控的算力预算下，让序列里任意两个位置高效地交换信息。搞懂这条主线，你再看swin transformer、pdformer这类具体模型，就不会觉得是一堆孤立的trick。

这篇文章适合三类人：一是刚学完多头自注意力机制原理、想搞清楚演进逻辑的入门者；二是手上有transformer分类任务或transformer预测正弦数据这类小项目、想优化结构的实践者；三是需要给团队做技术选型、要判断该用标准注意力还是某种高效变体的工程师。我会尽量用生活化的类比把原理讲透，同时给出可以直接抄的参数和代码片段，transformer matlab 完整代码和Python实现都会涉及。

2. 注意力机制的原始设计：为什么是“缩放点积”

2.1 从“查字典”理解Q、K、V三个角色

注意力机制最容易被讲复杂。我习惯用一个类比：你去图书馆查资料，手里有一个查询词（Query），书架每本书的书脊标签（Key），以及书里的实际内容（Value）。你拿查询词去和每个书脊标签比对相似度，相似度高的书就多翻几页，最后把翻到的内容按相似度加权汇总——这就是一次注意力计算。

对应到公式上，就是那个经典的：

Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V

QK^T算的是查询和每个键的相似度得分，softmax把得分归一化成权重，最后用权重对V加权求和。整个过程没有循环、没有卷积，全是矩阵乘法，这也是Transformer能大规模并行、进而催生分布式架构训练的根本原因。

2.2 那个sqrt(d_k)到底在防什么

很多人抄代码时会把/ sqrt(d_k)直接写上，但说不清为什么。我实测过一个例子：假设d_k = 64，Q和K的每个分量都是均值0、方差1的独立随机变量，那么点积结果的方差会接近d_k，也就是64。方差一大，softmax的输入就会有几个特别大的值，输出就退化成近似one-hot的分布——梯度几乎为零，训练直接卡死。

除以sqrt(d_k)就是把方差拉回到1附近，让softmax工作在梯度健康的区间。这个细节在transformer模型详解里经常一笔带过，但它是原始论文里少数几个“必须保留”的设计之一。我试过把它去掉跑transformer预测正弦数据，loss曲线前几百步几乎不动，加回来立刻正常收敛，差别非常明显。

注意：如果你自己实现注意力，d_k一定要用Key的维度，不是Value的维度，也不是模型总维度。多头场景下是每个头的维度。

2.3 单头注意力的三个致命短板

原始的单头注意力能work，但很快暴露出三个问题，这也是后续演进的直接动因：

• 表达能力单一：一次注意力只能学到一种“相似度视角”。比如一句话里，有的关系是语法依赖，有的是语义指代，单头只能取一个折中。
• 对位置不敏感：注意力本身是置换等变的，打乱输入顺序结果不变，必须额外注入位置信息。
• 计算和内存是O(n²)：序列长度翻倍，注意力矩阵大小翻四倍，长序列直接爆显存。

这三个短板，分别对应了后面三条演进路线：多头化、位置编码、高效注意力。下面逐个拆。

3. 多头自注意力：把“一个视角”拆成“多个视角”

3.1 多头的本质是子空间并行

多头注意力机制听起来玄乎，其实做法很朴素：把d_model维的输入切成h份，每份维度是d_k = d_model / h，每一份独立做一次注意力，最后把h个结果拼接再过一个线性层。用生活类比，就是原来一个人查资料，现在派h个人分头查不同书架，最后汇总。

关键点在于：每个头有自己独立的W_q、W_k、W_v投影矩阵。这意味着每个头会把输入映射到不同的子空间，学到的相似度模式天然不同。我做过一个可视化实验，在transformer分类任务里，有的头明显关注相邻词，有的头关注句首的[CLS]，还有的头几乎只关注标点——分工非常清晰。

3.2 头数不是越多越好：一个实测的取舍

新手常犯的错是把头数往大了设，觉得越多越强。我实测过在d_model = 512固定时，头数从4加到16的情况：

原因不难理解：头太多，每个头的维度太小，投影矩阵能表达的信息被压缩得厉害，注意力分布趋于噪声。业界经验是d_k保持在64左右比较稳，这也是原始论文d_model=512, h=8的由来。所以选头数时，先定d_k，再反推h，而不是反过来。

3.3 手写一个多头注意力的最小实现

光看原理不够，transformer手写一遍印象才深。下面是我常用的最小实现，去掉了工程优化，只保留核心逻辑：

import torch import torch.nn as nn import math class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads): super().__init__() assert d_model % num_heads == 0 self.d_model = d_model self.num_heads = num_heads self.d_k = d_model // num_heads # 一次性投影，再切分，比每个头单独建层更高效 self.W_q = nn.Linear(d_model, d_model) self.W_k = nn.Linear(d_model, d_model) self.W_v = nn.Linear(d_model, d_model) self.W_o = nn.Linear(d_model, d_model) def forward(self, x, mask=None): B, L, _ = x.shape # 投影后reshape成 (B, h, L, d_k) q = self.W_q(x).view(B, L, self.num_heads, self.d_k).transpose(1, 2) k = self.W_k(x).view(B, L, self.num_heads, self.d_k).transpose(1, 2) v = self.W_v(x).view(B, L, self.num_heads, self.d_k).transpose(1, 2) # 缩放点积 scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: scores = scores.masked_fill(mask == 0, -1e9) attn = torch.softmax(scores, dim=-1) out = torch.matmul(attn, v) # (B, h, L, d_k) out = out.transpose(1, 2).contiguous().view(B, L, self.d_model) return self.W_o(out)

这段代码有两个容易写错的地方：一是transpose之后一定要contiguous()再view，否则会报错；二是mask的填充值用-1e9而不是-inf，避免某些后端出现NaN。这两个坑我在transformer代码调试阶段都踩过。

3.4 位置编码：给置换等变的注意力补上顺序感

前面说过注意力对顺序不敏感，所以必须注入位置信息。原始方案是正弦位置编码，用不同频率的sin/cos函数生成一个和词向量同维的位置向量，直接相加。它的好处是不需要训练、可以外推到比训练更长的序列。

后来演进出了可学习位置编码、相对位置编码、旋转位置编码（RoPE）等。我个人的经验是：短序列任务用正弦或可学习都行，长序列或需要外推的场景优先考虑相对位置或RoPE。在transformer预测正弦数据这种任务里，位置编码的影响其实不大，因为正弦信号本身的周期性让模型容易学到模式；但在自然语言的长文档任务里，位置编码的选择能带来好几个点的差距。

4. 高效注意力演进：从O(n²)到可接受

4.1 长序列为什么是注意力的死穴

标准注意力的计算和显存都是O(n²)。n=512时注意力矩阵才26万个数，无所谓；n=8192时就是6700万，单头单层就要几百MB，多层多头直接爆。这就是为什么早期Transformer处理长文本要靠截断，也是分布式架构训练必须引入模型并行的原因之一。

围绕这个瓶颈，演进出了几条主要路线，我按“改动程度”从轻到重排一下：

• 稀疏注意力：不计算完整的n×n矩阵，只算部分位置对。比如局部窗口+全局token。
• 线性注意力：通过核函数近似，把softmax(QK^T)V的计算顺序换一下，复杂度降到O(n)。
• 低秩/原型近似：用少量“原型”代表整段序列，减少Key/Value数量。
• 分组查询注意力（GQA/MQA）：多个Query头共享同一组Key/Value头，主要省推理时的KV Cache显存。

4.2 稀疏注意力：Swin Transformer的窗口思路

swin transformer是稀疏注意力里最值得研究的一个。它的核心是把图像切成不重叠的窗口，注意力只在窗口内计算，复杂度从O(n²)降到O(n)。但纯窗口会切断窗口之间的联系，所以它又引入了移位窗口（shifted window）：下一层把窗口整体偏移半个窗口大小，让信息能跨窗口流动。

这个设计在remote sensing semantic segmentation这类任务里特别吃香，因为遥感图像分辨率高、token多，标准注意力根本跑不动。我复现过一个多层级多模态融合的Transformer用于遥感分割，窗口注意力几乎是标配。它的代价是实现复杂度上升，尤其是移位窗口的mask逻辑，写错一个索引结果就全乱。

4.3 线性注意力：换个乘法顺序

线性注意力的思路很巧妙。标准注意力是softmax(QK^T)V，先算n×n再乘V。如果能把softmax近似成一个核函数φ(·)，使得φ(Q)φ(K)^T可以结合，那么利用矩阵乘法结合律：

(φ(Q) φ(K)^T) V = φ(Q) (φ(K)^T V)

右边先算φ(K)^T V，得到一个d_k × d_v的矩阵，再和φ(Q)相乘，复杂度就变成O(n·d²)，和序列长度线性相关。代价是近似带来的精度损失，短序列上往往不如标准注意力，长序列才有优势。

提示：线性注意力不是万能药。我实测在n<512时，标准注意力又快又准，没必要换；只有n上千且显存吃紧时，线性变体才值得考虑。

4.4 分组查询注意力：推理阶段的显存救星

多头自注意力机制原理在训练时没问题，但推理时有个隐藏成本：KV Cache。自回归生成时，每生成一个token都要缓存之前所有token的Key和Value，序列一长，这部分显存就非常可观。MQA（Multi-Query Attention）让所有Query头共享一组KV，GQA（Grouped-Query Attention）折中，几个Query头共享一组KV。

现在主流大模型基本都用GQA，因为它在显存和效果之间取得了很好的平衡。这个演进方向很务实——不是为了刷指标，而是为了推理时能省下真金白银的显存。

5. 架构层面的演进：不只是注意力

5.1 从Post-LN到Pre-LN：一个训练稳定性的关键改动

原始Transformer用的是Post-LN，也就是LayerNorm(x + Sublayer(x))。这个结构在层数少时没问题，但层数一深就难训练，需要warmup很久。后来大家发现改成Pre-LN，即x + Sublayer(LayerNorm(x))，训练稳定得多，几乎不需要warmup。

这个改动看起来小，但影响巨大。我做过对比，同样12层，Post-LN在没warmup时loss直接发散，Pre-LN则平稳下降。现在绝大多数实现默认都是Pre-LN。如果你在复现老论文的transformer代码时发现训不动，先检查一下LN的位置。

5.2 前馈网络：被低估的另一半

很多人聊Transformer只聊注意力，其实前馈网络（FFN）占了参数量的三分之二。标准FFN是两层线性加一个激活，中间维度通常是4 × d_model。演进方向主要有两个：一是把激活从ReLU换成GELU、Swish等更平滑的函数；二是引入门控机制，比如GLU变体，用一条支路做门控。

我实测下来，把FFN的中间维度从4倍降到2倍，参数量减半，效果掉得并不多，在算力紧张时是个不错的取舍。这个经验在transformer模型详解里很少提，但实际调参时很有用。

5.3 编码器-解码器 vs 纯解码器

原始Transformer是编码器-解码器结构，用于机器翻译。后来演进出了三条分支：纯编码器（如BERT，适合理解类任务）、纯解码器（如GPT系列，适合生成类任务）、以及各种混合结构。

选哪种取决于任务。做transformer分类任务，纯编码器最直接；做文本生成，纯解码器配因果mask；做翻译或摘要，编码器-解码器仍然合适。我见过有人拿纯解码器硬做分类，效果也能凑合，但纯属绕远路。

6. 实操中的常见问题与排查

6.1 训练不收敛的排查顺序

这是被问得最多的问题。我总结了一个排查顺序，基本能覆盖八成情况：

1. 检查LN位置：是不是Post-LN且没warmup？改成Pre-LN试试。
2. 检查缩放因子：sqrt(d_k)有没有漏？漏了必炸。
3. 检查mask：padding mask和causal mask有没有搞混？填充值是不是-1e9？
4. 检查学习率：Transformer对学习率敏感，太大直接发散，建议从1e-4起步配warmup。
5. 检查初始化：投影矩阵初始化方差是否合理，太大太小都会影响早期训练。

6.2 显存不够的几种救法

按性价比排序，我一般这么处理：

• 先降batch size，配合梯度累积保持等效batch。
• 开启混合精度训练，显存能省近一半。
• 用梯度检查点，用时间换显存。
• 序列太长就上稀疏或线性注意力。
• 实在不行再上模型并行，但工程复杂度陡增。

6.3 常见问题速查表

6.4 几个我踩过的坑

第一个坑是mask的维度。批量训练时mask要能广播到(B, h, L, L)，我一开始只写了(B, L)，结果广播错了维度，模型学了个寂寞。第二个坑是view和reshape的区别，在transpose之后必须用contiguous().view()或直接reshape()，否则内存不连续会报错。第三个坑是学习率warmup的步数，设太短没效果，设太长浪费训练时间，一般占总步数的5%到10%比较合适。

7. 演进背后的那条主线

把上面这些串起来看，Transformer的算法演进其实一直围绕三个约束在走：表达能力、计算复杂度、训练稳定性。多头和FFN是在提表达能力，稀疏和线性注意力是在压复杂度，Pre-LN和各种归一化是在保稳定性。任何一个新变体，基本都能归到这三条线里的某一条。

我自己在选型时的习惯是：先看任务序列长度，短序列直接用标准多头注意力，别折腾；长序列再考虑稀疏或线性；推理部署时优先看KV Cache，该上GQA就上。至于transformer预测正弦数据这种玩具任务，标准结构足够，重点是把位置编码和缩放因子这些基础打牢。

最后分享一个我常用的调试技巧：把注意力权重可视化出来看。如果某些头始终是均匀分布，说明它没学到东西，可以考虑减头；如果某些头高度集中在少数位置，说明它在做有效路由。这个观察比看loss曲线直观得多，尤其在transformer分类任务里，能帮你快速判断模型是不是真的在“看”该看的地方。