1. 从一张注意力热力图说起Transformer到底在算什么第一次把Transformer的注意力权重可视化出来的时候我盯着那张热力图看了很久。输入是一句普通的中文模型在编码它这个代词时权重明显压在了前文的某个名词上。那一刻我意识到这个架构真正做的事情不是理解语言而是在高维空间里动态地决定当前这个位置应该看哪里、看多少。这就是Transformer的核心。它抛弃了循环结构抛弃了卷积核的局部假设用一套纯粹的注意力计算把序列中任意两个位置之间的关系都变成了可学习的权重。2017年那篇《Attention Is All You Need》出来的时候很多人第一反应是这不就是把Seq2Seq的注意力单独拎出来了吗但真正动手复现过的人才知道去掉RNN之后带来的并行性和可扩展性才是它统治接下来近十年深度学习的原因。这篇文章我想聊的不是Transformer是什么这种入门科普网上已经太多了。我想聊的是核心算法的演进脉络——从最原始的点积注意力到多头机制到位置编码的迭代到稀疏化、线性化、分组查询再到MoE这种把FFN层也拆开的路子。每一步演进背后都有明确的动机要么是算力瓶颈要么是长序列建模的困难要么是训练稳定性问题。我会把每一步的为什么讲清楚配上可以手算的小例子和关键参数的计算过程最后给一份实操中踩过的坑和排查清单。适合谁看如果你已经能默写Transformer的结构图但说不清楚为什么除以根号dk为什么多头要降维为什么后来的模型都用RoPE那这篇正好。如果你刚开始学也能看懂因为我会尽量用生活化的类比把数学直觉讲明白。2. 注意力机制的原始形态与它的三个致命问题2.1 缩放点积注意力一个除法背后的数值稳定性原始注意力公式写出来很简单Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) VQ是查询K是键V是值。用生活化的说法你手里有一个问题Q面前有一堆资料卡片K每张卡片上写着内容V。你先拿问题去和每张卡片的标题比对相似度QK^T相似度高的卡片你就多看一眼最后把所有卡片的内容按关注程度加权求和。那个除以根号dk的操作很多人初学时会忽略。我试过把它去掉在dk64的时候训练到几百步loss就开始震荡。原因很直接Q和K的每个分量如果都是均值0方差1的独立随机变量那么它们的点积的方差就是dk。dk64时点积的数值范围会落在正负二三十这个量级。softmax对输入的尺度极其敏感输入一大输出就趋近于one-hot梯度几乎消失。除以根号dk之后方差被拉回1softmax工作在梯度健康的区间。注意这个缩放因子是sqrt(d_k)而不是d_k因为方差是线性累加的标准差才是开根号。这个细节在面试里被问到的频率极高。2.2 三个致命问题串行、长程衰减、单视角原始注意力本身没问题但把它放进完整架构里配合当时的工程条件暴露出三个问题。第一个是位置信息的缺失。注意力是置换等变的你把输入序列打乱输出只是跟着打乱模型根本不知道谁在前谁在后。猫追老鼠和老鼠追猫在纯注意力眼里是一样的。所以必须额外注入位置信息这就引出了后面位置编码的整条演进线。第二个是单头注意力的表达瓶颈。一个softmax分布只能表达一种关注模式。但语言里的关系是多元的有的位置需要关注语法主谓关系有的需要关注指代有的需要关注语义相似。一个头只能选一种就像一个人同时只能盯一个方向。多头机制就是来解决这个的。第三个是计算复杂度随序列长度平方增长。QK^T是一个n×n的矩阵n是序列长度。n512时还好n8192时这个矩阵就是六千多万个元素显存直接爆炸。这条线催生了后面所有的稀疏注意力、线性注意力、FlashAttention等工作。2.3 一个可以手算的小例子假设序列长度n3dk2我们手动算一遍。Q [[1,0],[0,1],[1,1]]K [[1,0],[0,1],[1,1]]V [[1,2],[3,4],[5,6]]。先算QK^T位置0对位置01×10×01位置0对位置11×00×10位置0对位置21×10×11第一行就是[1, 0, 1]。除以sqrt(2)≈1.414得到[0.707, 0, 0.707]。softmax之后大约是[0.39, 0.22, 0.39]。也就是说位置0把注意力平均分给了位置0和位置2忽略了位置1。这个手算过程建议你自己拿纸推一遍比看十遍公式都管用。理解了单头的计算多头只是把它并行做几次再拼接。3. 多头注意力为什么要拆、怎么拆、拆完怎么合3.1 多头的本质是子空间投影多头注意力的做法是把dk维的Q、K、V分别用不同的线性层投影到更低的维度每个头在自己的低维子空间里独立做注意力最后把结果拼接再投影回来。为什么要降维假设模型维度d_model512头数h8。如果不降维每个头都在512维上做注意力8个头就是8倍计算量。实际做法是每个头降到512/864维8个头加起来的总计算量和单头在512维上做是一样的。所以多头不是更贵而是同样的预算下换8个视角看问题。这个设计背后的直觉不同的头会自发学到不同的关注模式。我实际训练过的小模型里确实能观察到有的头专门盯相邻位置类似局部卷积有的头盯句首的[CLS]或句尾有的头盯标点符号。这种分工不是人为设计的是训练自己涌现出来的。3.2 头数怎么选一个被低估的超参数头数h和每头维度dk_head的关系是dk_head d_model / h。这里有个经验性的取舍。h太小比如h1就退化成单头表达力不够。h太大比如d_model512时h64每头只有8维每个头的表达能力太弱注意力分布容易退化。实践中d_model512配h8d_model768配h12d_model1024配h16是比较常见的配置每头维度保持在64左右。我踩过的坑有一次为了省显存把h设成16但d_model只有256每头16维训练loss能降但验证集效果很差注意力图看起来非常杂乱。后来调回h8就好了。每头维度低于32时要警惕表达力不足的问题。3.3 拼接与输出投影的必要性多头算完之后是拼接concat然后过一个输出投影矩阵Wo。这个Wo经常被忽略但它不是可有可无的。拼接只是把8个头的输出在通道维度上摞起来各头之间还没有交互。Wo的作用是让不同头的信息在这个线性层里混合相当于一个汇总操作。去掉Wo各头就是完全独立的模型容量会下降。代码上大概是这个结构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, q, k, v, maskNone): batch_size q.size(0) # 投影并拆头: (B, n, d_model) - (B, h, n, d_k) Q self.W_q(q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) K self.W_k(k).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) V self.W_v(v).view(batch_size, -1, 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 F.softmax(scores, dim-1) out torch.matmul(attn, V) # (B, h, n, d_k) # 合并头 out out.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) return self.W_o(out)提示view之前一定要contiguous()否则transpose之后内存不连续view会报错。这个错误我见过太多次了。4. 位置编码的演进从正弦波到旋转位置编码4.1 正弦位置编码优雅但外推能力有限原始Transformer用的是正弦位置编码用不同频率的sin和cos函数生成位置向量直接加到词嵌入上。公式是PE(pos, 2i) sin(pos / 10000^(2i/d_model)) PE(pos, 2i1) cos(pos / 10000^(2i/d_model))这个设计的好处是它允许模型通过线性变换表达相对位置关系因为PE(posk)可以写成PE(pos)的线性函数。但实际用下来它的外推能力一般。训练时序列长度512推理时给1024效果会明显下降因为模型没见过那么大的位置值。4.2 可学习位置编码简单但无法外推BERT用的是可学习的位置嵌入就是一个nn.Embedding(max_len, d_model)。简单直接效果也不错。但缺点更明显max_len是硬编码的超过这个长度直接报错完全没有外推能力。而且位置嵌入和词嵌入一样是查表没有利用位置的连续性先验。4.3 相对位置编码与RoPE当前的主流选择真正解决外推问题的是相对位置编码的思路。核心想法是模型关心的不是这个token在第5位而是这两个token相距多远。T5用的相对位置偏置、Transformer-XL的相对编码都属于这一类。目前最主流的是旋转位置编码RoPE。它的做法不是加位置向量而是把Q和K向量按位置进行旋转。具体来说把d维向量两两配对每一对看作复平面上的一个点根据位置pos旋转一个角度。这样两个位置的内积只依赖于它们的相对距离天然具备相对位置的性质。RoPE的数学形式对第i对分量旋转角度θ_i pos / 10000^(2i/d)。用复数表示就是乘以e^(i·pos·θ_i)。实现上等价于def apply_rope(x, sin, cos): # x: (B, h, n, d_k) x1, x2 x[..., ::2], x[..., 1::2] # 旋转 rotated torch.stack([-x2, x1], dim-1).flatten(-2) return x * cos rotated * sinRoPE的好处是相对位置性质天然具备外推能力比正弦编码好很多而且不增加额外参数。现在主流的开源大模型基本都用RoPE或者它的变体比如NTK-aware插值、YaRN用来支持超长上下文。注意RoPE的频率基数是10000这个值在长上下文场景下会被调整。把基数调大比如改成500000可以让模型在更长序列上保持区分度这就是所谓的NTK-aware插值思路。4.4 ALiBi另一种极简方案ALiBiAttention with Linear Biases走的是另一条路不加位置编码而是在注意力分数上直接加一个和距离成正比的负偏置。距离越远惩罚越大。这个方案极其简单没有额外参数外推能力也不错。但它在需要精确位置信息的任务上比如某些结构化预测表现不如RoPE。5. 效率演进从O(n²)到线性与稀疏5.1 平方复杂度的瓶颈到底卡在哪标准注意力的计算量是O(n²·d)显存也是O(n²)。n2048时注意力矩阵是4M个元素还能接受。n32768时就是10亿个元素单是存这个矩阵就要几个GB。这就是长上下文的核心障碍。5.2 稀疏注意力只算该算的部分稀疏注意力的思路是不是所有位置对都需要计算。比如Longformer用了滑动窗口加全局token的模式每个token只关注附近w个token和少数几个全局token复杂度降到O(n·w)。BigBird在此基础上加了随机注意力理论上能逼近全注意力的表达能力。这类方法的坑在于稀疏模式是人为设计的可能和任务真正需要的关系不匹配。我试过在需要长距离精确检索的任务上用滑动窗口效果明显不如全注意力。5.3 线性注意力用核函数替换softmax线性注意力的核心技巧是把softmax(QK^T)V重新结合。利用矩阵乘法的结合律(QK^T)V Q(K^T V)。先算K^T V得到一个d×d的矩阵再和Q相乘复杂度就变成O(n·d²)和序列长度线性相关。但这里有个问题softmax是非线性的不能直接交换顺序。线性注意力用核函数φ(·)替换softmax把注意力写成φ(Q)(φ(K)^T V)的形式。代价是表达力下降因为核函数是对softmax的近似。5.4 FlashAttention不改变数学只改变计算顺序FlashAttention是这几年的一个重要工作它的思路和上面几种完全不同数学上完全等价于标准注意力但通过分块计算和重计算把显存占用从O(n²)降到O(n)同时利用GPU的SRAM减少HBM读写速度还更快。它的核心是tiling把Q、K、V分块每次只加载一块到高速缓存里算用在线softmax的技巧累积结果避免把完整的n×n矩阵写回显存。这个工作让我意识到很多时候瓶颈不在算法数学而在内存访问模式。实操建议如果你用的是PyTorch 2.0以上F.scaled_dot_product_attention会自动选择FlashAttention后端直接用它就行别自己手写。6. 架构层面的演进从标准块到MoE与分组查询6.1 分组查询注意力GQA推理加速的关键标准多头注意力在推理时有个问题KV cache。自回归生成时每生成一个token都要缓存之前所有token的K和V。头数h32、d_k128、序列长度4096时单层的KV cache就是32×4096×128×2×2字节≈64MB几十层下来就是几个GB。GQAGrouped Query Attention的做法是Q保持多头但K和V分成更少的组每组共享。比如32个Q头8个KV头每4个Q头共享一组KV。这样KV cache直接降到四分之一推理速度明显提升效果损失很小。现在Llama系列、Mistral等基本都用GQA。MQAMulti-Query Attention是GQA的极端情况只有1个KV头压缩最狠但效果损失也最大。6.2 MoE把FFN也拆开Transformer块里注意力之后是FFN层通常是两层线性加激活参数量占整个块的2/3左右。MoE混合专家的思路是把FFN拆成多个专家每个token只激活其中少数几个。这样做的好处是参数量可以做得极大但每个token的实际计算量不变。比如8个专家每次激活2个总参数量是8倍但FLOPs只有2倍。这就是为什么现在很多大模型能到千亿参数但推理成本可控。MoE的难点在负载均衡。如果所有token都涌向同一个专家其他专家就浪费了。所以需要加辅助损失来鼓励均匀分配。我调过的一个小MoE辅助损失系数设0.01时负载还算均衡设0.1时反而影响了主任务效果这个系数需要仔细调。架构变体核心改动主要收益主要代价MHA标准多头表达力强KV cache大MQA单KV头推理快效果略降GQA分组KV平衡需调组数MoEFFN分专家参数效率高负载均衡难线性注意力核函数替换长序列快表达力降6.3 归一化位置的迁移Pre-LN与Post-LN原始Transformer用的是Post-LN也就是残差之后做LayerNorm。但Post-LN在深层网络里训练不稳定需要warmup。后来GPT系列改成Pre-LN残差之前做归一化训练稳定多了warmup需求也降低。再后来RMSNorm出现了去掉了均值中心化只做缩放计算更快效果基本持平。现在主流大模型基本都用Pre-RMSNorm的组合。7. 实操中的常见问题与排查清单7.1 训练不收敛或loss震荡这是最常见的问题。排查顺序我一般是这样的先看学习率和warmup。Transformer对学习率很敏感没有warmup直接上大学习率前几百步很容易炸。warmup步数一般是总步数的1%到5%。再看初始化。Q、K、V的投影层初始化方差要合适太大太小都会影响。默认的xavier或kaiming一般够用但如果自己改过要检查。然后看缩放因子。确认除以的是sqrt(d_k)而不是d_k这个错误很隐蔽loss能降但收敛慢。最后看mask。padding mask和causal mask有没有正确应用mask的位置是不是-1e9而不是0。用0的话softmax之后padding位置还有权重会污染结果。7.2 注意力图看起来糊成一片如果可视化出来所有位置的注意力权重都差不多说明softmax没有拉开差距。可能的原因缩放因子没加导致数值太小、温度参数不对、或者模型还没训练充分。也有可能是d_k太大导致点积方差过大反而让softmax饱和。7.3 长序列显存爆炸优先用FlashAttention。如果硬件不支持考虑GQA减少KV cache或者用梯度检查点gradient checkpointing用时间换显存。再不行就上稀疏或线性注意力但要评估效果损失。7.4 位置编码外推失效如果推理长度超过训练长度效果骤降检查位置编码类型。正弦编码和可学习编码外推都差换成RoPE并配合NTK-aware插值。RoPE的基数调整是个经验活一般从10000往上调调到500000能支持到训练长度的几十倍。问题现象可能原因排查动作loss震荡学习率过大/无warmup加warmup降lr收敛慢缩放因子错误检查sqrt(d_k)注意力糊温度/缩放问题检查softmax输入尺度显存爆炸序列过长FlashAttention/GQA外推失效位置编码类型换RoPE插值推理慢KV cache大上GQA/MQA7.5 几个容易被忽略的细节第一个是dropout的位置。注意力权重上的dropout和残差上的dropout作用不同前者防止过度关注某些位置后者防止过拟合。原始论文里两个都有但现在的趋势是减少dropout因为训练数据量大了。第二个是权重共享。嵌入层和输出投影层共享权重能减少参数量在机器翻译这类任务上有效但在大模型上不一定。要看具体任务。第三个是初始化时输出投影的缩放。有些实现会把输出投影的初始权重除以sqrt(2×层数)用来稳定深层网络的训练。这个技巧在GPT-2的代码里能看到。8. 我个人的一些体会Transformer的演进史本质上是一部在表达力和计算效率之间反复权衡的历史。每一次改动都不是拍脑袋而是被某个具体的瓶颈逼出来的。理解了瓶颈就理解了为什么会有这个改动。我建议学的时候不要只背结构图而是拿一个小任务比如序列复制、简单翻译从头实现一遍把每个模块去掉试试看效果怎么变。我当年把位置编码去掉跑了一次模型完全学不会顺序相关的任务那个印象比看任何论文都深。另外现在的很多新架构其实是在Transformer基础上的微调核心的注意力机制没有变。与其追每一个新名词不如把注意力、位置编码、归一化这三块吃透剩下的都是组合和工程优化。真正难的不是理解公式而是知道在什么场景下该选哪个变体以及为什么。这个判断力只能靠实际跑实验积累没有捷径。