自注意力机制这个概念，我第一次真正搞明白它，不是在读论文的时候，而是在手写了一遍矩阵运算之后。当时对着一堆公式看了半天，总觉得"Q乘K再softmax再乘V"这个流程像是某种魔法，直到我用一个只有6个词的句子，把每一步的张量形状和数值都打印出来，才突然意识到：这东西本质上就是一套"带权重的信息检索系统"。每个词都在问其他词"你跟我有多相关"，然后根据相关程度把对方的信息按比例拿过来。就这么简单，但威力巨大。
这篇内容我想做的事情，是把自注意力机制从"能背出公式"推进到"能给别人讲明白、能自己动手实现、能判断哪里会出问题"。不管你是刚开始接触大模型、被各种教程绕晕的新手，还是已经用过不少框架但没深究过底层的老手，我都希望这篇东西能让你对自注意力有一个更扎实的理解。我会从最直觉的层面切入，逐步拆解到矩阵运算、多头机制、位置编码、复杂度分析，最后落到实际实现中那些文档里不会写的坑。
1. 从"每个词都想跟其他词对话"说起
1.1 为什么循环结构走不通了
在自注意力成为主流之前，处理序列数据主要靠循环神经网络（RNN）及其变体。RNN的思路很直白：从左到右一个一个读词，每读一个就更新一次内部状态，把这个状态当作"到目前为止我记住的所有信息"。这个设计在处理短句子时没问题，但一旦序列变长，麻烦就来了。
最核心的问题是信息传递路径太长。假设一个句子有50个词，第一个词的信息要传到第50个词的位置，中间要经过49次状态更新。每一次更新都是一次矩阵乘法加非线性变换，信息在传递过程中会不断被"稀释"或"扭曲"。这就像玩传话游戏，传的人越多，最后的话越离谱。梯度反向传播时也是同样的道理，梯度要穿过49层才能到达第一个词的位置，很容易出现梯度消失，导致前面的词学不到东西。
另一个问题是无法并行。RNN必须按顺序计算，第t步依赖第t-1步的结果，这意味着你没法把50个词同时扔给GPU去算。GPU最擅长的就是大规模并行矩阵运算，RNN这种串行结构等于把GPU的算力浪费了一大半。训练一个长序列模型，时间成本高得让人抓狂。
还有一个容易被忽略的点：RNN对"远距离依赖"的建模能力很弱。比如"我昨天在书店买的那本关于宇宙学的书，它非常精彩"这句话里，"它"指代的是"书"，中间隔了十几个词。RNN要把"书"的信息一路带到"它"的位置，中间经过那么多次变换，指代关系很容易丢失。而自注意力机制里，"它"可以直接和"书"建立连接，路径长度是1，不需要经过任何中间步骤。
1.2 自注意力的核心直觉：加权检索
自注意力的想法其实很像我们查资料的过程。假设你要写一篇关于"猫"的文章，你会去参考各种资料：一篇讲猫的习性的文章对你最有用，一篇讲狗的文章可能只有一点点参考价值，一篇讲量子力学的文章基本没用。你会根据"相关性"来决定从每份资料里吸收多少信息。
自注意力做的就是这件事。对于序列中的每一个词（称为查询词），它都会去"看"序列中的所有词（包括自己），计算自己跟每个词的相关性分数，然后根据这些分数对所有词的信息进行加权求和。相关性高的词，它的信息就被大量吸收；相关性低的词，它的信息几乎被忽略。
这里的关键在于：相关性不是预先定义好的，而是从数据里学出来的。每个词会生成三个向量：查询向量（Query）、键向量（Key）、值向量（Value）。查询向量代表"我在找什么"，键向量代表"我是什么"，值向量代表"我能提供什么信息"。查询词用自己的Q去和每个词的K做点积，得到的分数就是相关性。点积越大，说明这两个向量的方向越一致，相关性越高。
我常用一个类比来解释：想象一个大型图书馆，每个词是一本书。每本书都有一个"标签"（Key）描述自己的内容，也有一份"正文"（Value）提供实际信息。当你想找资料时，你带着自己的"问题"（Query）去匹配每本书的标签。标签跟你问题越匹配的书，你就越认真地读它的正文。最后你把所有书里读到的内容按匹配程度加权汇总，就得到了你的答案。
1.3 跟传统注意力机制的区别
注意力机制本身不是新东西，在自注意力之前，它已经被用在序列到序列（Seq2Seq）模型里了。但传统注意力有一个关键区别：它的Query来自解码器，Key和Value来自编码器，是跨序列的注意力。比如做机器翻译时，解码器在生成英文翻译的每个词时，会去"看"编码器输出的中文原文的每个位置，决定当前该关注原文的哪个部分。
自注意力的"自"字就体现在这里：Query、Key、Value全部来自同一个序列。序列自己跟自己做注意力，每个词都在问"这个序列里其他词跟我有什么关系"。这让模型能够捕捉序列内部的依赖关系，而不需要依赖外部的编码器-解码器结构。
这个区别看起来很小，但影响深远。跨序列注意力只能建模两个序列之间的关系，而自注意力可以建模序列内部任意两个位置之间的关系。堆叠多层自注意力之后，模型能够捕捉到非常复杂的内部结构，比如语法树、指代关系、语义角色等等。
2. Q、K、V到底是怎么算出来的
2.1 从词向量到三个投影矩阵
假设我们有一个长度为n的序列，每个词已经被转换成了一个d维的向量（这个转换通常通过嵌入层完成）。我们把整个序列表示成一个矩阵X，形状是n×d，每一行是一个词的向量。
接下来，我们需要三组可学习的权重矩阵：W_Q、W_K、W_V，它们的形状都是d×d_k（d_k是投影后的维度，通常等于d或者d除以头数）。用X分别乘以这三个矩阵，就得到了：
- Q = X · W_Q，形状n×d_k
- K = X · W_K，形状n×d_k
- V = X · W_V，形状n×d_v
这里的W_Q、W_K、W_V就是模型需要学习的参数。它们的作用是把原始的词向量投影到三个不同的"子空间"里。为什么需要三个不同的投影？因为"我在找什么"、"我是什么"、"我能提供什么"是三种不同的语义角色，用同一组表示来承担这三种角色会互相干扰。分开投影之后，每个角色都有自己独立的表示空间，模型能更灵活地学习。
我见过不少人问：为什么不直接用X本身当Q、K、V？理论上可以，但效果会差很多。因为如果Q和K是同一个向量，那么一个词跟自己的点积永远是最大的（自己跟自己最相关），这会引入很强的"自我偏好"偏差。而且不同词的向量在同一个空间里，点积的大小主要反映的是向量模长和方向的相似性，不一定对应我们想要的语义相关性。经过独立的线性投影之后，模型有了更大的自由度来学习"什么样的词应该跟什么样的词相关"。
2.2 点积注意力的计算流程
拿到Q、K、V之后，计算过程分四步：
第一步：算相关性分数。用Q乘以K的转置，得到一个n×n的矩阵，我们叫它注意力分数矩阵S。S[i][j]表示第i个词对第j个词的相关性分数。这里用的是点积，因为点积能很好地衡量两个向量的方向一致性。如果两个向量方向相同，点积就大；方向垂直，点积为零；方向相反，点积为负。
第二步：缩放。把S除以根号d_k。这一步非常关键，但很多人初学时容易忽略。为什么要缩放？因为当d_k比较大时，点积的结果会变得很大（每个维度乘积累加，维度越多，和的方差越大）。大的分数经过softmax之后，会变得非常"尖锐"——最大的那个分数接近1，其他的接近0。这会导致梯度几乎为零，训练不动。除以根号d_k可以把分数的方差拉回到合理范围，让softmax的输出更平滑，梯度更健康。
具体来说，如果Q和K的每个分量都是独立的、均值为0、方差为1的随机变量，那么它们的点积的方差就是d_k。除以根号d_k之后，方差变成1。这是一个很漂亮的数学性质，保证了不管d_k取多大，softmax的输入都在一个稳定的范围内。
第三步：softmax归一化。对S的每一行做softmax，得到注意力权重矩阵A。A的每一行加起来等于1，每个元素表示"当前词应该从对应位置的词那里吸收多少比例的信息"。softmax的作用是把任意实数分数转换成概率分布，同时放大高分、抑制低分，让注意力更集中。
第四步：加权求和。用A乘以V，得到输出矩阵，形状是n×d_v。输出的第i行就是第i个词的新表示，它是所有词的V向量的加权和，权重由A的第i行决定。
整个流程用公式写出来就是：
Attention(Q, K, V) = softmax(Q·K^T / √d_k) · V
这个公式看起来简洁，但每一步都有讲究。我在下面会逐一拆解那些容易踩坑的地方。
2.3 一个具体的手算例子
光看公式容易飘，我们拿一个极简的例子来走一遍。假设序列只有3个词，每个词的向量维度是2（实际中维度会大得多，但原理一样）。为了简化，假设W_Q、W_K、W_V都是单位矩阵，也就是说Q=K=V=X。
设X为：
- 词1：[1, 0]
- 词2：[0, 1]
- 词3：[1, 1]
算分数矩阵S = Q·K^T：
S[1][1] = 1×1 + 0×0 = 1 S[1][2] = 1×0 + 0×1 = 0 S[1][3] = 1×1 + 0×1 = 1 S[2][1] = 0×1 + 1×0 = 0 S[2][2] = 0×0 + 1×1 = 1 S[2][3] = 0×1 + 1×1 = 1 S[3][1] = 1×1 + 1×0 = 1 S[3][2] = 1×0 + 1×1 = 1 S[3][3] = 1×1 + 1×1 = 2
S = [[1,0,1],[0,1,1],[1,1,2]]
缩放：d_k=2，除以根号2约等于1.414，得到： S' = [[0.707, 0, 0.707], [0, 0.707, 0.707], [0.707, 0.707, 1.414]]
softmax（对每一行）：
第一行：exp(0.707)=2.028, exp(0)=1, exp(0.707)=2.028，和=5.056，权重=[0.401, 0.198, 0.401] 第二行：同理，权重=[0.198, 0.401, 0.401] 第三行：exp(0.707)=2.028, exp(0.707)=2.028, exp(1.414)=4.113，和=8.169，权重=[0.248, 0.248, 0.504]
加权求和 A·V：
输出词1 = 0.401×[1,0] + 0.198×[0,1] + 0.401×[1,1] = [0.802, 0.599] 输出词2 = 0.198×[1,0] + 0.401×[0,1] + 0.401×[1,1] = [0.599, 0.802] 输出词3 = 0.248×[1,0] + 0.248×[0,1] + 0.504×[1,1] = [0.752, 0.752]
可以看到，每个词的输出都是所有词信息的混合，混合比例由相关性决定。词3跟自己的相关性最高（分数2），所以它保留了自己最多的信息。词1和词2之间的相关性为0，所以它们几乎不互相吸收信息。
这个例子虽然简单，但把整个计算链路都走通了。你在实际实现时，把维度放大、加上多头、堆叠多层，本质逻辑完全一样。
3. 多头注意力：为什么要"分裂"成多个视角
3.1 单头注意力的表达能力瓶颈
如果只用一组Q、K、V，模型只能学到一种"相关性模式"。但语言中的关系是多种多样的：有语法上的主谓关系，有语义上的指代关系，有位置上的相邻关系，还有情感上的一致性关系等等。用一组投影矩阵去捕捉所有这些关系，就像用一把尺子量所有东西，力不从心。
具体来说，单头注意力在计算softmax之后，每个词对所有词的注意力权重被压缩成一组分布。这组分布只能表达"我最关注谁"，但没法同时表达"我在语法上关注A，在语义上关注B"。当一句话里存在多种需要同时建模的关系时，单头注意力会顾此失彼。
另一个问题是子空间利用不充分。d维的向量空间里，不同的语义特征可能分布在不同的子空间中。单头注意力在整个空间里做一次投影，可能会把不同性质的信号混在一起，降低信噪比。
3.2 多头机制的设计逻辑
多头注意力的思路很直接：既然一组不够，那就用多组。把d维的Q、K、V分别切成h份，每份维度是d/h，每一份独立做一次注意力计算，最后把h个输出拼接起来，再过一层线性变换。
举个例子，假设d=512，h=8，那么每个头的维度就是64。每个头有自己的W_Q、W_K、W_V（形状是512×64），独立计算注意力，得到8个n×64的输出。把这8个输出在最后一维拼接，得到n×512的矩阵，再乘以一个512×512的输出投影矩阵W_O，得到最终结果。
这里的关键设计是：每个头在不同的子空间里学习不同的注意力模式。有的头可能学会了关注相邻词，有的头可能学会了关注句法上的依赖，有的头可能学会了关注语义上的相似性。这种"分工合作"让模型能同时捕捉多种关系。
我见过一个很直观的实验：把训练好的多头注意力模型的每个头的注意力权重可视化出来，会发现不同的头确实呈现出不同的模式。有的头注意力矩阵接近对角线（关注自己或相邻位置），有的头呈现出明显的块状结构（关注同一短语内的词），还有的头会跨很远的位置建立连接（捕捉长距离依赖）。这说明多头机制确实在起作用，不是简单的冗余。
3.3 多头计算中的维度陷阱
多头注意力在实现时有一个很容易搞错的地方：维度的切分和拼接方式。
常见做法是：不真的把Q、K、V切成h份分别存储，而是直接用形状为d×d的权重矩阵一次性算出完整的Q、K、V（形状n×d），然后通过reshape操作把它变成n×h×(d/h)，再转置成h×n×(d/h)，这样每个头的数据就分开了。计算完注意力后，再把h个头的输出转置回来、reshape成n×d。
这个过程中，reshape和transpose的顺序非常关键。如果搞错了，虽然张量形状对得上，但每个头拿到的数据是错乱的，模型学出来的东西完全没有意义。而且这种错误不会报错，只会让效果变差，非常隐蔽。
我的经验是：在实现多头注意力时，一定要用一个小例子验证。比如设n=2、d=4、h=2，手动构造输入，打印出每个头拿到的Q、K、V，确认切分逻辑正确。这一步花五分钟，能省掉后面几小时的调试。
另一个坑是d_k的计算。在缩放那一步，除以的是每个头的维度d/h，而不是完整的d。如果用错了，缩放系数会差h倍，导致softmax的输出过于尖锐或过于平滑。这个错误同样不会报错，但会显著影响训练效果。
4. 位置编码：没有它，自注意力就是"词袋"
4.1 自注意力的排列不变性
自注意力有一个很重要的性质：它对输入的顺序不敏感。如果你把序列里的词打乱顺序，只要对应的Q、K、V也跟着打乱，注意力计算的结果只是跟着打乱，每个词的输出值不变。
这个性质来自注意力计算本身：它做的是加权求和，而加权求和是不区分顺序的。对于"猫追狗"和"狗追猫"这两个句子，如果只看自注意力的输出，模型无法区分它们的区别，因为两个句子里"猫"和"狗"的注意力权重关系是对称的。
这显然不行。语言是有顺序的，"我打你"和"你打我"意思完全相反。所以必须想办法把位置信息注入到模型里。
4.2 正弦位置编码的数学直觉
最经典的位置编码方案是正弦位置编码。对于位置pos和维度i，编码值定义为：
- 偶数维度：PE(pos, 2i) = sin(pos / 10000^(2i/d))
- 奇数维度：PE(pos, 2i+1) = cos(pos / 10000^(2i/d))
这个设计看起来很数学，但直觉其实不难理解。不同维度对应不同频率的正弦波，低维度变化快（高频），高维度变化慢（低频）。这就像用不同精度的尺子来标记位置：高频尺子能区分相邻位置，低频尺子能区分远距离位置。把所有这些尺子的读数拼在一起，就得到了一个位置的多尺度表示。
这个编码是直接加到词向量上的，而不是拼接。加法意味着位置信息和语义信息在同一个空间里混合，模型需要学会从中分离出两者。实践中证明这是可行的，而且比拼接更节省维度。
正弦编码还有一个好处：它可以外推到训练时没见过的长度。因为正弦函数是连续定义的，不管pos多大都有对应的值。虽然实际中外推效果有限（模型在训练长度之外的表现通常会下降），但至少不会像可学习位置编码那样直接报错。
4.3 可学习位置编码与旋转位置编码
除了正弦编码，还有两种常见方案：
可学习位置编码是把每个位置的编码当作可训练参数，初始化一个max_len×d的矩阵，让模型自己学。这种做法在BERT等模型里很常见。优点是灵活，模型可以根据任务自己决定怎么编码位置；缺点是没法外推到训练长度之外，而且参数量随最大长度线性增长。
**旋转位置编码（RoPE）**是近几年很流行的方案，被用在不少主流大模型里。它的思路不是"加"位置信息，而是通过旋转Q和K向量来"注入"位置信息。具体来说，对于位置pos的向量，把它按维度分成两两一组，每一组看作一个二维平面上的点，然后根据pos旋转一个角度。这样两个位置的点积就自然包含了它们的相对距离信息。
RoPE的优势在于：它编码的是相对位置而不是绝对位置，这更符合语言的特性（"相邻"这个概念比"第5个位置"更有意义）；而且它也有一定的外推能力。不过RoPE的实现细节比较多，尤其是和KV Cache结合时，需要仔细处理。
我在实际项目中的体会是：如果序列长度不超过训练长度，三种方案差别不大；如果需要处理超长序列，RoPE通常是更好的选择，但要注意它的外推需要配合一些技巧（比如位置插值）。
5. 复杂度分析与长序列的挑战
5.1 O(n²)的时间与空间复杂度
自注意力的计算复杂度是O(n²·d)，其中n是序列长度，d是维度。这个平方项来自注意力分数矩阵，它有n×n个元素，每个元素需要d次乘加运算。空间复杂度也是O(n²)，因为要存储这个分数矩阵（或者至少存储softmax之后的权重矩阵）。
当n比较小（比如512）时，n²=262144，还能接受。但当n变成8192时，n²=67 million，显存占用和计算量都会爆炸。这就是为什么早期的大模型上下文长度普遍限制在512或1024，不是不想做长，是算不起。
具体算一下：假设n=8192，d=512，用float32存储注意力矩阵，需要8192×8192×4字节=256MB。这还只是一个头、一层的数据。如果有32层、16个头，那就是256MB×32×16=131GB，显然不现实。实际实现中会通过分块计算、梯度检查点等技术来降低峰值显存，但计算量本身是省不掉的。
5.2 稀疏注意力与线性注意力的思路
为了缓解O(n²)的问题，研究者提出了很多改进方案，大致分两类：
稀疏注意力的思路是：不是所有词对之间都需要计算注意力，只计算一部分。比如只关注每个词周围的w个词（滑动窗口注意力），或者只关注一些全局token（全局注意力），或者用一些启发式规则选择性地计算。这样复杂度可以降到O(n·w)或O(n·√n)。缺点是可能漏掉一些重要的远距离依赖。
线性注意力的思路是：通过改变注意力的计算顺序或近似softmax，把复杂度降到O(n·d²)。比如先用K^T乘以V得到一个d×d的矩阵，再用Q去乘这个矩阵。这样复杂度就跟n成线性关系了。但这样做会损失一些表达能力，因为softmax的非线性被绕过了。
还有一类方法是分块计算，比如FlashAttention。它不改变注意力的数学定义，而是通过巧妙地安排计算顺序和利用GPU的SRAM，减少对HBM的读写次数，从而在保持精确计算的同时大幅提升速度。FlashAttention现在已经是很多框架的默认选项了，实测下来能提速2-4倍，显存占用也能降低好几倍。
5.3 KV Cache：推理加速的关键
在自回归生成（比如对话模型逐字输出）时，每生成一个新token，都需要用它去关注之前所有的token。如果每次都重新计算之前所有token的K和V，那就太浪费了。KV Cache的做法是：把之前算过的K和V缓存起来，新token只需要计算自己的Q、K、V，然后用Q去和缓存的K做注意力，把新的K、V追加到缓存里。
这样每步的计算复杂度从O(n²)降到了O(n)，生成速度大幅提升。但KV Cache也带来了显存问题：缓存的大小是2×n×d×层数×头数，当n很大时，缓存本身就会占用大量显存。这也是为什么长上下文推理对显存要求很高的原因之一。
我在实际部署时踩过一个坑：KV Cache的维度顺序。不同框架对KV Cache的存储格式要求不一样，有的要求[batch, head, seq, dim]，有的要求[batch, seq, head, dim]。如果搞错了，虽然不报错，但注意力会算错，生成的内容会变得莫名其妙。建议在接入新框架时，先用一个短序列对比一下有无Cache的输出是否一致。
6. 动手实现：从零写一个自注意力模块
6.1 用PyTorch实现单头注意力
下面是一个最基础的单头自注意力实现，我加了详细注释：
import torch import torch.nn as nn import math class SingleHeadAttention(nn.Module): def __init__(self, d_model): super().__init__() self.d_model = d_model # 三个投影矩阵，把输入投影到Q、K、V空间 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) def forward(self, x, mask=None): # x: [batch, seq_len, d_model] Q = self.W_q(x) # [batch, seq_len, d_model] K = self.W_k(x) V = self.W_v(x) # 算注意力分数 scores = torch.matmul(Q, K.transpose(-2, -1)) # [batch, seq_len, seq_len] # 缩放 scores = scores / math.sqrt(self.d_model) # 如果有mask，把不需要关注的位置设为负无穷 if mask is not None: scores = scores.masked_fill(mask == 0, float('-inf')) # softmax归一化 attn_weights = torch.softmax(scores, dim=-1) # 加权求和 output = torch.matmul(attn_weights, V) # [batch, seq_len, d_model] return output, attn_weights这段代码看起来简单，但有几个细节值得说：
mask的处理。在解码器里，为了防止模型"偷看"未来的token，需要用因果mask把上三角部分设为负无穷。这样softmax之后这些位置的权重就是0。注意要用float('-inf')而不是一个大负数，因为softmax对-inf的处理是精确的，对大负数可能会有数值误差。
缩放的位置。除以根号d_model要在softmax之前做，而且要在mask之前或之后都行，但建议在mask之前，这样逻辑更清晰。
transpose的维度。K.transpose(-2, -1)是把最后两个维度交换，从[batch, seq, d]变成[batch, d, seq]，这样跟Q做矩阵乘法才能得到[batch, seq, seq]。这个操作很常见，但初学者容易搞混。
6.2 扩展到多头注意力的完整代码
class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads): 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 # 一次性算出所有头的Q、K、V 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): batch_size, seq_len, _ = x.shape # 投影并分头 Q = self.W_q(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) K = self.W_k(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) V = self.W_v(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) # 现在形状都是 [batch, num_heads, seq_len, d_k] # 算注意力 scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: # mask需要扩展到 [batch, 1, seq_len, seq_len] 或类似形状 scores = scores.masked_fill(mask == 0, float('-inf')) attn_weights = torch.softmax(scores, dim=-1) output = torch.matmul(attn_weights, V) # [batch, num_heads, seq_len, d_k] # 合并多头 output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model) output = self.W_o(output) return output, attn_weights这段代码里，view和transpose的配合是关键。view(batch, seq, heads, d_k)把最后一维d_model拆成heads×d_k，然后transpose(1,2)把heads维度换到前面。合并时反过来操作。注意contiguous()的调用，因为transpose之后张量在内存里不连续，直接view会报错。
6.3 实现中容易出错的三个地方
第一个坑：mask的形状。mask的维度需要能广播到scores的形状[batch, heads, seq, seq]。常见的因果mask形状是[seq, seq]或[1, 1, seq, seq]，要确保它能正确广播。如果mask形状不对，PyTorch可能会静默地广播成错误的结果，而不是报错。
第二个坑：softmax的维度。必须对最后一维做softmax，也就是对每个query对所有key的分数做归一化。如果搞错了维度，归一化就完全错了。建议在代码里显式写dim=-1，不要依赖默认值。
第三个坑：数值稳定性。当scores里有-inf时，softmax可能会产生NaN。PyTorch的softmax实现已经处理了这种情况（会忽略-inf），但如果你自己手写softmax，要注意减去最大值再exp。另外，如果整行都是-inf（比如mask把所有位置都遮住了），softmax会输出NaN。这种情况在padding处理不当时会出现，需要特别注意。
7. 训练中的实际观察与调参经验
7.1 注意力权重的可视化与诊断
训练过程中，把注意力权重可视化出来是很有用的诊断手段。我通常会关注几个模式：
如果注意力权重几乎均匀分布（每个位置权重都接近1/n），说明模型还没学到有用的注意力模式，可能是学习率太小或者训练不够。如果注意力权重全部集中在一个位置（接近one-hot），可能是缩放系数不对或者初始化有问题。如果注意力权重呈现出有规律的结构（比如对角线附近高、或者某些列特别高），说明模型学到了有意义的东西。
还有一个常见的现象叫注意力坍缩：多个头学出了几乎相同的注意力模式，失去了多头应有的多样性。这通常发生在训练后期或者数据量不足时。缓解方法包括增加头数、加dropout、或者用一些正则化手段鼓励头之间的差异性。
7.2 学习率与初始化的影响
自注意力模块对初始化比较敏感。如果W_Q、W_K的初始化方差太大，点积的结果会很大，softmax会变得很尖锐，梯度消失。如果初始化方差太小，点积结果接近0，softmax接近均匀分布，模型学不到东西。
常用的做法是用Xavier初始化或Kaiming初始化，让每层的输出方差保持稳定。具体来说，对于线性层，PyTorch默认的初始化通常就够用，但如果你自己写矩阵乘法而不是用nn.Linear，就要手动初始化。
学习率方面，自注意力模块通常需要比CNN更小的学习率，因为它的参数更多、梯度更复杂。我一般会用warmup策略，前几千步线性增加学习率，之后再衰减。这样能避免训练初期的不稳定。
7.3 梯度检查与常见训练失败模式
训练自注意力模型时，有几种常见的失败模式：
loss不下降：可能是学习率太小、初始化太差、或者数据有问题。先检查梯度范数，如果梯度接近0，说明可能是初始化或缩放的问题；如果梯度爆炸，说明学习率太大或需要梯度裁剪。
loss下降但生成质量差：可能是过拟合，也可能是mask有问题（比如训练时用了因果mask但推理时没用，或者反过来）。检查训练和推理的mask逻辑是否一致。
训练后期loss突然飙升：可能是数值不稳定，比如某些batch的attention score太大导致softmax溢出。可以尝试降低学习率、加梯度裁剪、或者用更稳定的softmax实现。
我在实际项目里遇到过一次很隐蔽的bug：mask在训练时是正确的，但在推理时因为batch size变化导致广播出错，注意力算到了padding位置。结果就是模型在短序列上表现正常，长序列上生成的内容会突然变得混乱。排查了很久才发现是mask形状的问题。所以强烈建议在训练和推理时都加一个断言，检查mask的形状和值是否符合预期。
8. 从自注意力到完整Transformer的拼图
8.1 残差连接与层归一化的作用
自注意力模块本身只是一个"信息混合"的操作，它不包含非线性变换（softmax算一个，但主要是归一化作用）。要让模型有足够的表达能力，还需要加上前馈网络（FFN）、残差连接和层归一化。
残差连接的作用是让梯度能直接流过网络，缓解深层网络的梯度消失问题。具体做法是把自注意力的输出和输入相加：output = x + Attention(x)。这样即使注意力模块学得不好，至少还有一条恒等映射的路径，模型不会比不加注意力更差。
层归一化是对每个样本的特征维度做归一化，让激活值保持稳定的分布。它和BatchNorm的区别在于，LayerNorm不依赖batch维度，所以在序列长度变化、batch size变化时更稳定。Transformer里通常用Pre-LN（先归一化再进注意力）而不是Post-LN（先注意力再归一化），因为Pre-LN训练更稳定，不需要warmup也能收敛。
8.2 前馈网络的补充角色
FFN在Transformer里扮演的是"逐位置非线性变换"的角色。它对每个位置的向量独立做两次线性变换加一次激活，通常是先升维再降维（比如d→4d→d）。虽然它不涉及位置之间的交互，但它提供了非线性表达能力，让模型能学习更复杂的特征变换。
有一种理解是：注意力负责"收集信息"，FFN负责"处理信息"。注意力把相关的信息聚合到每个位置，FFN再对这些信息进行加工。两者交替堆叠，模型就能逐步构建出高层次的语义表示。
8.3 堆叠多层后的涌现行为
单层自注意力的能力有限，它只能建模一阶的关系。但堆叠多层之后，模型能建模高阶关系。比如第一层可能学到相邻词的关系，第二层能学到短语级别的关系，第三层能学到句子级别的关系，更深的层能学到跨句子的语义关系。
这种"逐层抽象"的机制是深度Transformer强大的关键。每一层都在前一层的表示基础上做进一步的注意力混合，逐步构建出越来越抽象的语义表示。这也是为什么大模型需要那么多层——层数不够，抽象能力就不够。
我在实验中发现一个有趣的现象：浅层注意力通常关注局部模式（相邻词、常见搭配），深层注意力关注全局模式（长距离依赖、语义角色）。如果你只训练几层，模型能处理简单的局部任务，但复杂的长距离推理任务表现会很差。这也解释了为什么大模型的"深度"和"宽度"都很重要。
9. 一些容易被忽略的工程细节
9.1 padding mask与causal mask的区别
这两种mask经常被混淆，但作用完全不同：
padding mask是用来遮住padding位置的。当batch里不同样本长度不一样时，短的样本会用padding补齐。这些padding位置没有实际意义，不应该参与注意力计算。padding mask的形状通常是[batch, seq]，需要扩展到[batch, 1, 1, seq]才能广播到注意力分数矩阵。
causal mask是用来防止解码器看到未来token的。它是一个上三角为0、下三角为1的矩阵，形状[seq, seq]。在训练时，它保证每个位置只能关注自己和之前的位置。
在解码器的自注意力里，这两种mask通常要同时使用。做法是把它们合并：先应用causal mask，再应用padding mask，或者用一个组合mask。注意合并时要用逻辑与（AND），而不是或（OR），否则会遮错位置。
9.2 dropout在注意力中的正确位置
Dropout在自注意力里有三个常见位置：
注意力权重上：在softmax之后、乘V之前，对注意力权重做dropout。这是最常用的做法，能防止模型过度依赖某些位置。
Q、K、V投影后：对投影结果做dropout，增加扰动。
输出投影后：对最终的输出做dropout。
我一般会在注意力权重和输出投影后各加一个dropout，概率0.1左右。注意不要在softmax之前做dropout，那样会破坏归一化，导致输出不稳定。
9.3 混合精度训练中的数值问题
用混合精度（fp16或bf16）训练时，自注意力模块容易出现数值问题。主要原因是softmax对精度敏感，fp16的动态范围有限，大的分数容易溢出。
常见的解决方案是：在softmax之前把分数转成fp32计算，算完再转回fp16。PyTorch的torch.cuda.amp会自动处理一部分，但如果你自己写kernel，就要手动处理。另外，缩放系数除以根号d_k也有助于缓解这个问题，因为它把分数的范围拉回到了fp16能表示的范围。
bf16的动态范围和fp32一样，所以数值问题比fp16少很多。如果硬件支持，我建议优先用bf16。
10. 我踩过的那些坑和最后的建议
说几个我在实现和使用自注意力时真实踩过的坑，希望能帮你省点时间。
第一个坑是缩放系数的位置。有一次我在实现一个变体时，把缩放放在了softmax之后，结果训练完全跑不动。后来才意识到，缩放必须在softmax之前，因为它的作用是控制softmax输入的方差。放在之后等于没缩放，softmax还是那么尖锐。
第二个坑是mask的dtype。mask如果是bool类型，masked_fill能正常工作；但如果是float类型，mask == 0可能因为浮点误差判断错误。建议统一用bool类型的mask，并且在创建时就确保值只有True和False。
第三个坑是KV Cache的更新顺序。在自回归生成时，新token的K、V要先追加到cache里，再用更新后的cache做注意力。如果顺序反了，新token就看不到自己，生成的内容会缺少当前token的信息。这个bug很隐蔽，因为输出看起来"差不多对"，只是质量差一点。
第四个坑是多头的维度顺序。不同框架对多头的维度排列不一样，有的是[batch, heads, seq, d_k]，有的是[batch, seq, heads, d_k]。在跨框架迁移代码时，一定要确认维度顺序，否则虽然不报错，但每个头拿到的数据是错乱的。
最后分享一个我个人的学习建议：如果你真的想搞懂自注意力，不要只看论文和博客，一定要自己动手写一遍。从单头开始，用一个小例子验证每一步的形状和数值，然后扩展到多头，再加上mask、dropout、残差、LayerNorm，一步步搭出一个完整的Transformer块。这个过程可能花你一个周末，但收获远比看十篇教程大。因为很多细节，只有在你亲手调试的时候才会真正理解为什么要那样设计。
另外，如果你在做实际项目，不要急着上最复杂的变体。标准的自注意力加上FlashAttention，在大多数场景下已经足够好了。先把基础版本跑通、调稳，再根据具体需求考虑优化。很多所谓的"改进"在特定场景下有效，但通用性未必比标准版本好。
