从零手写多头自注意力：PyTorch实现与避坑指南

注意力机制这几年被讲得太玄了。打开任何一篇教程，上来就是Q、K、V三个矩阵，然后一顿矩阵乘法，最后Softmax收尾，公式写得漂漂亮亮，但很多人看完还是不知道这东西到底在干嘛、为什么能work、自己动手写的时候该从哪一行开始。我见过不少朋友，Transformer的论文翻了三遍，PyTorch环境也搭好了，真让他从零写一个多头自注意力，还是卡在维度变换上——不是张量形状对不上，就是mask加错了位置，要么就是训练半天loss不降。

这篇东西就是冲着这个问题来的。我把它叫做"极简书"，不是因为它浅，而是因为我想把注意力机制从数学符号里拽出来，用最少的公式、最直白的代码，把它的骨架讲清楚。读完你应该能做到三件事：第一，能手写一个可运行的多头自注意力模块；第二，能说清楚Transformer里每个注意力层到底在"注意"什么；第三，知道在PyTorch里怎么把它跑起来、怎么避坑。适合刚入门深度学习的同学，也适合那些用过Transformer但没自己实现过的工程师。

1. 先把注意力机制从公式里拽出来

1.1 注意力本质上是一次"带权重的信息检索"

很多人第一次接触注意力，看到的是这个公式：

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

然后就开始懵。其实你换个角度想，这就是一次数据库查询。你手里有一个查询条件（Query），数据库里有一堆键值对（Key-Value），你想根据查询条件和每个Key的匹配程度，把对应的Value加权取回来。匹配程度高的，权重就大，取回来的信息就多；匹配程度低的，权重趋近于零，基本忽略。

生活里这种场景太多了。比如你在一个群里问"谁会修电脑"，群里十个人，每个人擅长的领域不同（这就是Key），你问的这句话就是Query。谁擅长修电脑，谁的回答（Value）就应该被重点听。注意力机制干的就是这件事——它让模型学会"该听谁的"。

这个类比不是随便说说的，它直接对应到代码里的三个线性变换：Query、Key、Value都是从输入X经过不同的权重矩阵投影出来的。为什么要投影？因为原始输入里，一个词既可能是"被查询的对象"，也可能是"提供信息的来源"，用同一套表示没法区分这两种角色，所以要用三组独立的参数把它们分开。

1.2 为什么是点积，为什么除以根号d_k

匹配程度怎么算？最直接的方式是点积。两个向量的点积越大，说明它们方向越接近，相关性越强。这比欧氏距离计算起来快得多，而且能直接用矩阵乘法批量处理，GPU最喜欢这种操作。

那为什么要除以sqrt(d_k)？这里有个很实际的工程原因。假设Q和K的每个分量都是均值0、方差1的独立随机变量，那么它们的点积的方差就是d_k。当维度d_k很大的时候（比如64、128），点积的结果会变得非常大。Softmax有个特性：输入值越大，输出越接近one-hot，梯度就越小。梯度一小，训练就卡住了。除以sqrt(d_k)就是把方差拉回到1附近，让Softmax的输出保持在一个梯度健康的范围内。

我实测过，在一个d_k=512的注意力层里，如果不做这个缩放，训练前几百步的梯度范数会掉到1e-6量级，基本等于没在学。加上缩放之后，梯度稳定在1e-2到1e-1之间，收敛速度差了好几倍。这不是理论推导，是跑出来的。

1.3 从单头到多头：为什么要"分裂"注意力

单头注意力只能学到一种"关注模式"。但语言里的关系是多种多样的——有的位置需要关注语法主谓关系，有的需要关注指代关系，有的需要关注位置邻近关系。一个头不够用，那就多来几个。

多头的做法是把d_model维的输入切成h份，每份维度是d_k = d_model / h，每个头独立做一次注意力，最后把结果拼起来再过一个线性层。这里有个容易搞混的点：多头不是把计算量翻了h倍，因为每个头的维度缩小了，总的计算量和单头全维度注意力是差不多的。它的价值在于让模型在多个子空间里并行地捕捉不同的关系模式。

你可以这样理解：单头注意力像一个人用一双眼睛看世界，多头注意力像八个人从不同角度看同一个场景，最后把八个人的观察汇总起来。每个人看到的细节不同，合起来就比一个人看得全面。

2. 手写多头自注意力：从张量形状开始拆

2.1 输入输出的形状约定必须先定死

写注意力代码，90%的bug都出在张量形状上。所以在动手之前，必须把形状约定定死。我习惯用这套：

输入x的形状是(B, L, d_model)。经过Q、K、V三个线性层之后，形状不变，还是(B, L, d_model)。然后要把它拆成多头，变成(B, h, L, d_k)。注意这里的变换顺序：先view成(B, L, h, d_k)，再transpose成(B, h, L, d_k)。为什么要先view再transpose？因为view要求内存连续，直接transpose会打乱内存布局，后续计算可能报错。

计算完注意力之后，输出形状是(B, h, L, d_k)，需要transpose回(B, L, h, d_k)，再reshape成(B, L, d_model)。这一套变换我建议你拿张纸画一遍，比看十遍代码都管用。

2.2 一个能直接跑的多头注意力实现

下面这份代码是我自己反复用过的版本，去掉了所有花哨的东西，只保留核心逻辑：

import torch import torch.nn as nn import math class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads, dropout=0.1): super().__init__() assert d_model % num_heads == 0, "d_model必须能被num_heads整除" 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) self.dropout = nn.Dropout(dropout) def forward(self, x, mask=None): B, L, _ = x.shape 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, float('-inf')) attn = torch.softmax(scores, dim=-1) attn = self.dropout(attn) out = torch.matmul(attn, V) out = out.transpose(1, 2).contiguous().view(B, L, self.d_model) return self.W_o(out)

这段代码有几个细节值得说。contiguous()那一步很多人会漏掉，transpose之后张量在内存里不连续，直接view会报错。mask用float('-inf')填充而不是0，是因为Softmax之后exp(-inf)=0，这样被mask的位置权重严格为零；如果填0，Softmax之后还会有非零权重，信息就泄漏了。

2.3 自注意力和交叉注意力的区别就在Q从哪来

上面写的是自注意力，Q、K、V都来自同一个输入x。交叉注意力（cross-attention）的区别只有一个：Q来自解码器，K和V来自编码器的输出。代码上就是把W_q(x)换成W_q(decoder_input)，W_k和W_v作用在encoder_output上。其他逻辑完全一样。

这个区别在机器翻译里特别直观：解码器在生成第t个词的时候，它的Query是"我现在想生成什么"，而Key和Value是编码器对源句子的理解。模型通过注意力决定"生成这个词的时候，应该重点看源句子的哪几个词"。这就是所谓的对齐（alignment），只不过它是学出来的，不是人工标注的。

3. 位置编码：注意力为什么需要"补课"

3.1 自注意力本身是"位置盲"的

这是注意力机制一个反直觉的特性：如果你把输入序列的顺序打乱，自注意力的输出（在忽略位置编码的情况下）是不变的。因为它的计算本质是加权求和，而加权求和是对所有位置一视同仁的，没有先后概念。

但语言是有顺序的。"猫追老鼠"和"老鼠追猫"意思完全相反，但词袋是一样的。所以必须把位置信息注入进去。Transformer的做法是在输入embedding上加上一个位置编码向量，让每个位置的表示带上"我在第几个位置"的信息。

3.2 正弦位置编码的构造逻辑

原始Transformer用的是正弦位置编码：

def positional_encoding(max_len, d_model): pe = torch.zeros(max_len, d_model) position = torch.arange(0, max_len).unsqueeze(1).float() div_term = torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)) pe[:, 0::2] = torch.sin(position * div_term) pe[:, 1::2] = torch.cos(position * div_term) return pe

为什么用不同频率的正弦余弦？因为这样构造出来的编码有一个很好的性质：位置pos+k的编码可以表示成位置pos编码的线性变换。这意味着模型可以通过线性操作学到相对位置关系，而不需要为每个绝对位置单独记忆。这个设计在长序列上比可学习的位置embedding更稳，因为可学习的embedding在训练时没见过的长度上会失效。

不过现在很多新模型（比如一些视觉Transformer）改用可学习的位置embedding或者相对位置编码，效果也不差。选哪个取决于你的任务：如果序列长度固定且不长，可学习的更灵活；如果序列长度变化大，正弦编码更省心。

3.3 位置编码加在哪里，加几次

标准做法是在输入embedding上加一次，然后进入编码器堆叠。但也有一些变体，比如在每一层都加位置编码，或者在注意力分数上加相对位置偏置。我试过在每层都加，效果提升不明显，反而增加了计算量。对于大多数任务，入口加一次就够了。

有个坑要注意：位置编码的维度必须和d_model一致，而且要和embedding相加，不是拼接。相加意味着位置信息和语义信息在同一个空间里混合，模型需要自己学会区分它们。拼接的话维度会翻倍，后续所有层的参数都要调整，不划算。

4. 训练注意力模型时最容易踩的五个坑

4.1 mask用错方向，模型直接学废

mask是注意力里最容易出错的地方。有两种mask：padding mask和causal mask。padding mask用来忽略填充的无效位置，causal mask用来防止解码器看到未来的词。

padding mask的形状通常是(B, 1, 1, L)，在scores上广播。causal mask是一个下三角矩阵，形状(1, 1, L, L)。如果你把两者搞混，或者mask的True/False方向反了，模型要么看到不该看的信息，要么把有效信息也屏蔽了。

我踩过一次坑：mask里用1表示"保留"，0表示"屏蔽"，结果masked_fill(mask == 0, -inf)写成了masked_fill(mask == 1, -inf)，训练loss一直不降，查了半天才发现是mask反了。建议你写完mask之后，打印出来看一眼，确认对角线以下和以上的值符合预期。

4.2 学习率没做warmup，前期梯度爆炸

Transformer的训练对学习率很敏感。原始论文用了warmup策略：前几千步学习率线性增加，之后按步数的平方根倒数衰减。为什么？因为注意力层在初始化时，Q和K的点积方差很大，如果一上来就用大学习率，梯度会爆炸。

def lr_lambda(step): d_model = 512 warmup_steps = 4000 return d_model ** (-0.5) * min(step ** (-0.5), step * warmup_steps ** (-1.5))

这个策略我几乎在所有Transformer训练里都用，warmup步数根据任务规模调整，小任务1000步就够，大任务可以到10000步。不用warmup的话，前几百步loss经常直接飙到nan。

4.3 dropout加错位置，正则化效果打折

注意力里的dropout有两个位置：一个是Softmax之后对注意力权重做dropout，另一个是每个子层输出后做dropout。前者是原始论文的做法，目的是防止模型过度依赖某几个位置。后者是标准的残差dropout。

我见过有人在Q、K、V的线性层后面加dropout，这其实不太对。因为Q和K是用来算相似度的，对它们做dropout会引入噪声，让相似度计算不稳定。正确的做法是在注意力权重和最终输出上加。

4.4 忘记缩放，大维度下训练不动

前面说过要除以sqrt(d_k)，但实际写代码的时候很容易忘。尤其是当你把d_k设得比较小（比如32）的时候，不缩放好像也能跑，但一旦把模型放大到d_k=128或更大，问题就暴露了。我的建议是：不管维度大小，缩放一律加上，养成习惯。

4.5 显存不够时，先检查是不是注意力矩阵太大

注意力矩阵的形状是(B, h, L, L)，显存占用是B * h * L * L * 4字节。当L=512、h=8、B=32时，这个矩阵就是32*8*512*512*4 ≈ 268MB，还只是一层。如果堆叠12层，光注意力矩阵就3GB多。

显存不够的时候，优先考虑减小batch size或者序列长度，而不是减头数。因为减头数会降低模型容量，而减batch size只是影响训练稳定性。另外，现在有一些高效的注意力实现（比如flash-attn），通过分块计算避免显式存储完整的注意力矩阵，能省不少显存。不过那是另一个话题了，初学阶段先把标准实现搞明白。

5. 从注意力到Transformer：还差哪几块拼图

5.1 残差连接和层归一化不是可选项

注意力层本身只是一个信息混合器，它不改变信息的"量"，只改变信息的"分配"。要让深层网络能训练，必须加残差连接和层归一化。标准做法是LayerNorm(x + Sublayer(x))，也就是Post-LN。后来有人提出Pre-LN，把LayerNorm放在子层之前，训练更稳定，现在很多实现都用Pre-LN。

我自己的经验是：小模型（6层以下）用Post-LN和Pre-LN差别不大，但层数一多，Pre-LN明显更稳。如果你训练的时候发现loss震荡厉害，先试试把LayerNorm的位置换一下。

5.2 前馈网络承担了大部分参数

Transformer每层里，注意力子层之后还有一个前馈网络（FFN），通常是两个线性层加一个激活函数，中间维度是d_model的4倍。这个FFN占了每层大约三分之二的参数量。它的作用是对每个位置独立地做非线性变换，可以理解为"在注意力混合完信息之后，每个位置自己再消化一下"。

有意思的是，最近有些研究发现FFN层可能充当了键值记忆的角色，存储了大量事实性知识。这也解释了为什么增大FFN维度比增大注意力头数对模型容量的提升更明显。

5.3 编码器和解码器的注意力有什么不同

编码器里只有自注意力，每个位置可以看到所有位置。解码器里有两种注意力：一个是带causal mask的自注意力，只能看当前位置及之前；另一个是交叉注意力，Q来自解码器，K和V来自编码器。

这个设计的原因很直接：生成任务必须保证自回归性质，不能偷看未来。而理解任务（比如分类）不需要这个限制，所以编码器可以用全可见的自注意力。

6. 注意力机制在视觉和时序任务里的变体

6.1 视觉Transformer把图像切成patch

Vision Transformer（ViT）的做法很暴力：把图像切成16x16的patch，每个patch展平成一个向量，加上位置编码，直接送进标准Transformer编码器。它不做卷积，完全靠注意力捕捉空间关系。

这个做法在小数据集上效果不如CNN，因为注意力没有卷积那样的局部性归纳偏置，需要大量数据才能学好。但在大数据集上，ViT能超过最好的CNN。后来Swin Transformer引入了窗口注意力和层次化结构，把局部性和全局性结合起来，在视觉任务上表现更好。

6.2 时序任务里的注意力需要处理变长和周期性

时序数据（比如传感器读数、股票价格）用注意力处理时，有几个特殊问题。一是序列可能非常长，标准注意力的O(L^2)复杂度吃不消；二是时序数据往往有周期性，需要位置编码能表达这种周期。

一些做法是用稀疏注意力，只计算局部窗口内的注意力，或者用固定的稀疏模式（比如每隔k个位置关注一次）。还有的做法是在位置编码里加入时间戳信息，让模型知道每个点对应一天中的哪个时刻、一周中的哪一天。

6.3 通道注意力和空间注意力是另一条路

在CNN时代，SE模块（Squeeze-and-Excitation）提出了一种通道注意力：对每个通道的特征图做全局池化，学出一组通道权重，然后加权回去。这跟Transformer的注意力不是一回事，但思想相通——都是让模型学会"该关注什么"。

后来CBAM把通道注意力和空间注意力结合起来，先算通道权重，再算空间权重。这类模块参数量小，插入到现有CNN里就能涨点，在工业界用得很多。如果你做的是图像任务，不一定非得上Transformer，加个注意力模块可能性价比更高。

7. 一些实操中的零碎经验

关于PyTorch环境，有个问题被问得特别多：conda activate报"无法将conda项识别"。这通常是因为conda没有加到系统PATH里，或者你用的是PowerShell但conda初始化的是cmd。解决办法是在Anaconda Prompt里运行conda init powershell，然后重启终端。如果还不行，直接用Anaconda Prompt操作，别折腾PATH了。

关于GPU，注意力计算是显存密集型而不是计算密集型。同样的FLOPs，注意力层比卷积层占更多显存。所以如果你的显卡显存不大（比如8G），模型层数和序列长度都要控制。我一般先用小配置跑通，再逐步放大，而不是一上来就拉满。

关于调试，注意力权重是最好的诊断工具。把attn矩阵可视化出来，看看模型在关注哪些位置。如果发现注意力全部集中在一个位置（比如第一个token），说明可能训练不充分或者初始化有问题。如果注意力均匀分布，说明模型还没学到有用的模式。正常情况下，注意力应该呈现出一定的稀疏性和结构性。

关于从零实现和用现成模块的取舍，我的建议是：学习阶段一定自己写一遍，哪怕写得丑。因为只有自己写过，才知道transpose和view的顺序为什么重要、mask为什么要用-inf、缩放为什么不能省。等工作里真正用的时候，直接用PyTorch的nn.MultiheadAttention或者现成的Transformer库，没必要重复造轮子。但前提是你知道那个轮子是怎么转的。

最后说一个我自己的体会：注意力机制之所以强大，不是因为它复杂，恰恰是因为它简单。它只做了一件事——根据相关性加权求和。但就是这么简单的一个操作，配合足够的深度和数据，就能涌现出惊人的能力。理解这一点，比记住公式重要得多。