我发现很多人学注意力机制是直接从Transformer源码看起对着softmax(QK^T)V那一行反复琢磨却始终觉得Q、K、V三个矩阵像是天上掉下来的。我自己也是直到有一次手写Nadaraya-Watson核回归把指数核函数往公式里一代输出来一条和注意力加权结果几乎一样的曲线我才意识到注意力机制的本质就是一次带学习的核回归。这篇博文想把这条少有人讲的线索讲清楚从核回归的数学形式出发一步步推到你每天在用的Transformer注意力架构最后给出可以直接跑通的PyTorch实验代码和一些训练阶段的工程细节。1. 为什么注意力机制要追溯到Nadaraya-Watson核回归1.1 核回归在统计里解决什么问题先说一个最朴素的统计问题。假设你有一堆散点横坐标是样本位置纵坐标是样本值现在要在任意一个新位置做预测传统回归会强行假设一个函数形式比如线性回归假设yaxb。可如果真实关系是一条复杂的波浪线你很难提前知道该用几次多项式。Nadaraya-Watson核回归的思路完全不同它不假设函数形式而是让每个样本根据“与当前查询点的距离”来分配权重距离越近贡献越大。这个“距离”由一个核函数定义最常见的是高斯核$$K_h(x - x_i) \exp\left(-\frac{(x - x_i)^2}{2h^2}\right)$$其中h是带宽控制影响力衰减的速度。预测公式是$$\hat{y}(x) \frac{\sum_{i1}^{n} K_h(x - x_i) y_i}{\sum_{i1}^{n} K_h(x - x_i)}$$说白了就是把所有样本的y_i做一次加权平均权重看谁和预测点更“像”。这个思路放到生活里特别容易理解。你去一个陌生城市租房想知道某个小区大概什么价位你不会傻到去查全国均价而是优先参考这个小区里已经成交的房源价格再看旁边商圈评估综合水平。Nadaraya-Watson核回归做的事情就是这个只不过把“相近的房源”量化成了数学上的核权重。1.2 把核回归公式改写成注意力公式现在把核回归公式的符号换一换你马上就会看到非常眼熟的东西。核回归里查询点x对应注意力机制里的query简称q样本点xi对应key简称k样本标签yi对应value简称v。权重计算过程其实就两步先算相似度再归一化。$$权重_i \frac{\exp(-\frac{|q - k_i|^2}{2h^2})}{\sum_j \exp(-\frac{|q - k_j|^2}{2h^2})}$$这一步和softmax没有任何区别只是把softmax里的logit换成了距离的负半平方。最终输出$$\hat{y} \sum_i 权重_i \cdot v_i$$这就是一个标准的注意力输出。Transformer里的注意力公式写成softmax(QK^T/√d)V本质上也是“算相似度、归一化、做加权平均”这三板斧唯一区别是它把原始输入先做了一次线性变换再用点积而非欧氏距离来算相似度。所以很多人第一次看到Transformer时会觉得QKV是凭空设计出来的组合但如果从核回归出发你会发现这套结构几乎是被数学推着走的你要做加权平均就必须有相似度函数要有相似度函数就必须有代表索引的key和代表内容的value。架构设计的自由度其实只体现在“相似度怎么算”“内容怎么映射”这两件事上。1.3 一张表说清楚核回归和注意力的对应关系为了后面展开方便我直接把两套名词对照起来Nadaraya-Watson核回归注意力机制作用查询点xquery q当前要预测/检索的对象样本坐标x_ikey k_i用来与查询做匹配的索引样本标签y_ivalue v_i真正被聚合的内容核函数K_hexp(q·k/√d)相似度度量函数带宽h√d或温度参数控制权重分布的尖锐程度分母归一化softmax保证所有权重之和为1这张表的最后两行特别重要。很多人会忽略带宽h但它在核回归里决定了模型的成败放到注意力机制里它就变成了你训练时经常要调的temperature和缩放因子。理解了这种对应后面调试模型时会少走很多弯路。2. 从核回归到注意力机制的两次升级可学习投影与缩放因子2.1 第一次升级把原始样本投影到查询、键、值空间核回归直接用原始坐标算距离但现实中“距离近”并不等于“应该关注”。举个例子对文本“苹果公司发布新品”这句话“苹果”和“公司”在向量空间里可能距离不远但它们之间的语义关系并不一定比“苹果”和“发布会”更值得关注。这需要模型自己学会什么是最重要的匹配依据。于是Transformer引入了三个可学习矩阵W_Q、W_K、W_V分别将输入映射到查询空间、键空间、值空间$$q_i x_i W_Q,\quad k_i x_i W_K,\quad v_i x_i W_V$$这一步等于把核回归里的“距离度量”从固定的欧氏距离换成了可学习的相似度函数。模型在训练过程中会调整W_Q和W_K让“需要互相注意”的token在点积空间里更靠近让“没关系”的token远离。同样的道理W_V决定真正应该从每个token中提取什么内容是只取表面信息还是取更深层的语义。你可以这样理解核回归像是一个拿固定尺子量距离的人而注意力机制给了你一把可以按任务打磨的尺子。同样是量距离不同任务下“近”的定义不同可学习投影就是为了让尺子适配任务。2.2 第二次升级为什么核函数收敛到“点积softmax”核回归里可以用各式各样的核但注意力机制最终统一到指数点积核背后有一个非常自然的逻辑当两个向量方向一致、长度别差太多时点积越大代表它们越相似。而且点积运算可以充分利用GPU的矩阵乘法效率远高于逐元素计算距离。但直接用点积会带来一个数值问题这也是缩放因子1/√d出现的原因。假设q和k的每个分量都是均值0、方差1的独立随机变量那么它们的点积展开后是$$q \cdot k \sum_{i1}^d q_i k_i$$点积的均值是0方差呢因为q_i和k_i相互独立E(q_i²)1、E(k_i²)1所以$$\text{Var}(q \cdot k) \sum_{i1}^d E(q_i^2)E(k_i^2) d$$当维度d增大到128甚至512时点积的方差会达到128甚至512这意味着logit的数值会分散得很开。做过softmax的都知道logit之间的差距一旦拉大softmax就会往one-hot方向跑梯度迅速趋近于0训练非常困难。除以√d之后方差重新回落到1这个量级softmax函数的输入分布更集中梯度也更稳定。这一步在代码里只是一个除法但背后的原因值得记住它直接关系到你手写注意力时要不要加缩放以及初期训练要不要调温度。2.3 核宽带宽与温度其实是同一个旋钮核回归里的带宽h控制的是“影响力半径”。h很小的时候几乎只有一个样本会被选中曲线会贴着噪声走过拟合h很大的时候所有样本权重都差不多曲线退化成一条均值线欠拟合。注意力机制里温度参数τ对softmax分布的影响一模一样τ越小分布越尖锐模型越来越像“只挑一个最强匹配的token”τ越大分布越平坦所有token雨露均沾输出趋向全局平均信息。在Transformer里这个旋钮默认被固定为√d_k但实际训练中可以将其暴露出来用温度缩放的方式控制注意力分布的熵。尤其在训练初期模型还没学到合理的语义空间一上来就用很低的温度很容易导致注意力坍缩成少数几个token后面怎么训练都很难恢复。所以很多预训练框架会配合warmup策略本质上就是让模型先用“较宽的核”探索全局信息再逐步收紧到有效匹配。3. 手写一次Nadaraya-Watson核回归亲眼看到权重从哪来3.1 造一个带噪声的波浪形数据集先动手做一个最简单的实验环境。我会用PyTorch造一条带噪声的波浪曲线模拟所有真实数据都有的“规律噪声”结构。import torch import torch.nn as nn import matplotlib.pyplot as plt torch.manual_seed(42) # 训练样本200个落在[-3, 3]上的点 x_train torch.linspace(-3, 3, 200) y_true torch.sin(x_train) 0.5 * torch.sin(2 * x_train) y_train y_true 0.3 * torch.randn(x_train.shape[0])预测点直接取一条密集的x_test线方便最后绘制拟合曲线。数据生成这块没什么讲究关键是要确保x_train和y_train一一对应后面的核回归本质上就是拿着这些已知的(x_i, y_i)对来预测新x处的结果。3.2 先写最朴素的版本理解每一步在算什么Nadaraya-Watson核回归最朴素的写法就是一层循环对每个测试点重新算一遍所有训练点的权重def nadaraya_watson_naive(x_query, x_train, y_train, h): 朴素实现逐个query计算核权重 x_query x_query.reshape(-1) outputs torch.zeros_like(x_query) for i in range(x_query.shape[0]): # 高斯核距离越近权重越大 weights torch.exp(-((x_query[i] - x_train) ** 2) / (2 * h * h)) weights weights / weights.sum() # 归一化等价于softmax outputs[i] (weights * y_train).sum() # 加权平均 return outputs这段代码里最关键的是中间那两行先算权重再归一化最后加权求和。你会发现它和注意力机制的执行过程完全一致只是把里面的exp参数从点积换成了带带宽的高斯距离。跑起来之后你会看到一条不错的平滑预测线噪声点对曲线的影响基本被平滑掉了。3.3 矩阵化写法15行代码复刻softmax注意力上面那种循环写法在数据量大的时候跑得极慢我们把它改写成矩阵形式顺便暴露出它和Transformer注意力的同构关系def nadaraya_watson_mat(x_query, x_train, y_train, h): 矩阵化实现一次性算出所有query的权重矩阵 # (N, M) 距离矩阵每行是一个query对所有key的距离 logits -((x_query[:, None] - x_train[None, :]) ** 2) / (2 * h * h) weights torch.softmax(logits, dim-1) # 按行做归一化 y_hat weights y_train # 加权求和 return y_hat这里x_query[:, None]把形状变成(N, 1)x_train[None, :]变成(1, M)广播之后得到(N, M)的距离矩阵。第N行第M列表示第N个查询点对第M个训练点的相似度。对每一行做softmax就得到一组和为1的注意力权重再右乘y_train完成聚合。这个逻辑对应到Transformer里就是把“高斯距离”替换成“点积”把“y_train”替换成“经W_V映射后的value矩阵V”仅此而已。如果你愿意把高斯距离展开成平方项再把常数项吸收到缩放里你甚至能把核回归改写成带QKV投影的形式。3.4 带宽h的敏感性过平滑和过拟合一眼可见参数h是核回归里唯一需要调的旋钮它的影响非常直观建议你实际把h设成0.05、0.5、10这三个值跑一遍h0.05时预测曲线几乎穿过每一个噪声点拟合了太多噪声细节这就是过拟合。放到注意力机制里相当于温度极低、注意力分布极度尖锐模型对单个token极其敏感一有波动就跟着跑。h0.5时曲线既保留了波峰波谷的趋势又没被噪声带偏是比较理想的平滑程度。h10时所有训练点的权重都趋于相同预测结果几乎就是全体y_train的平均值完全丢失了局部结构。这个实验强烈建议亲手跑一次因为你会发现模型复杂度根本不是越“紧”越好太紧反而会记住噪声太松又学不到规律。这个感受在调大模型注意力头、温度参数时迁移过来同样成立。3.5 把h换成可学习矩阵核回归立刻变成参数化注意力核回归的h是固定的但我们可以让它可学习。更进一步把“输入x本身”换成经过线性层映射之后的q、k、v就得到了参数化版本class ParamAttentionLayer(nn.Module): def __init__(self, in_features, d_k): super().__init__() self.w_q nn.Linear(in_features, d_k, biasFalse) self.w_k nn.Linear(in_features, d_k, biasFalse) self.w_v nn.Linear(in_features, d_k, biasFalse) self.scale d_k ** 0.5 def forward(self, x): q self.w_q(x) k self.w_k(x) v self.w_v(x) logits q k.transpose(-2, -1) / self.scale weights torch.softmax(logits, dim-1) return weights v这个类已经非常接近Transformer中单头注意力的最小实现了和核回归的差别只剩下“相似度函数”和“投影方式”。如果你用第3章的带噪声数据去训练这个层让它学习用sin曲线重建y式目标你会观察到它收敛后得到的权重模式和核回归自动学到的带宽模式高度相似。这就是为什么我说Transformer的注意力机制并非全新发明而是核回归朝着“可学习方向”自然演化的结果。4. 用核回归的视角重新拆解Transformer注意力架构4.1 QKV为什么必须分开索引和内容是两回事理解QKV最简单的方式是拿图书检索打比方你手里有一个检索词这个检索词是query图书馆里的索引导航是key真正想读的书页内容才是value。你用query去匹配key找到相关位置再从对应位置取value内容。很多人会问既然key和value都来自同一个输入序列为什么非得搞两套向量因为在实际语义场景里“决定该不该关注”和“关注的时候该提取什么”是两码事。一个token可能作为索引时很重要但作为内容时需要经过进一步变换才能提供有效信息。把q、k、v分开等于给模型两套参数去分别优化“匹配维度”和“内容维度”。如果不分开只让原始token既能当索引又能当内容匹配和提取就会互相牵制模型的表达能力会打个折扣。核回归里的y_train和x_train也天然是分开的判断位置远近用的是x_i最终加权平均用的却是y_i。同理注意力机制里匹配用的是k_i聚合用的却是v_i。这层设计在数学上是清晰的、漂亮的。4.2 多头注意力并行跑多组不同“带宽”的核回归多头注意力经常被解释成“让模型从不同表示子空间学习”这个说法宏观上没问题但真正动手写过的人应该更认同另一个视角每个头就是一组独立的核回归拥有自己的投影、自己的带宽、自己的权重分配模式。不同头在做不同角度的相似度判断。有的头可能更关注相邻token的局部结构相当于h比较小注意力窗口比较聚焦有的头可能更关注跨越长距离的修饰关系相当于h比较大全局信息也会纳入聚合。把这些头的结果拼接起来相当于同时跑多组不同带宽、不同投影方向的自适应核回归最后把结果合流。这也解释了为什么增大head数并不总是好事。如果你的模型只有8个头但训练数据不够充分很可能其中几个头重叠地学成了同一个模式浪费参数。我的做法是训练早期打印每个头的注意力熵如果发现多个头的熵分布几乎一致就会适当削减head数让每个头更“专注”。4.3 掩码的本质对核回归加局部性约束Transformer里有两种常见掩码一种是padding掩码一种是因果掩码。从核回归的视角看它们都是在限制“加权平均的参与范围”。因果掩码要解决的是自回归问题预测第t个位置时当前位置之后的token是未来信息不能提前看到。实现时直接把未来位置的logits设成一个很小的负数比如负无穷softmax之后这些位置的权重自然变成0。这就像核回归里把超出某个范围的样本点权重设为0只允许使用当前时刻之前的样本参与加权平均。padding掩码的处理方式同理被padding的位置没有有效语义信息不应该作为key参与匹配更不应该污染value聚合所以在logits层面将它屏蔽。理解和核回归的关联之后你会更容易记住掩码的正确位置它必须加在softmax之前、改的是logits而不是在softmax之后把权重手动刷成0。这个细节我在第5章会展开说因为它是一个极其常见且隐蔽的坑。5. 训练和推理阶段注意力机制最容易翻车的四个细节5.1 掩码加错位置softmax之后置零是不对的很多人写自回归注意力时会先算出完整权重矩阵再用布尔掩码把未来位置的权重手动置成0。这个做法在直观上好像没问题但实际会产生两个隐患。第一softmax保证每一行权重之和为1你把某些位置置0之后这一行的和就变成小于1聚合结果相当于把所有有效位置的value按一个缩小的比例混合数值整体变小等于在你完全不知情的情况下改变了输出尺度。第二如果你的实现没有在置0之后重新归一化梯度流经这些被置零的位置时会变得不连续反向传播的行为和原始概率分布不再匹配训练稳定性大打折扣。正确做法是在softmax之前通过masked_fill将需要屏蔽的位置填充为负无穷logits q k.transpose(-2, -1) / scale logits logits.masked_fill(~causal_mask, float(-inf)) attn torch.softmax(logits, dim-1) output attn v看完核回归的推导后这个逻辑一目了然权重本来就是归一化的你只能在归一化之前调控参与项而不是事后篡改归一化结果。5.2 缩放因子忘不得一个除号背后的稳定性手写注意力的时候最容易跳过的细节就是除以√d_k。刚跑小数据时可能没什么感觉因为数据小、维度低、logits的方差还不大。一旦你换成128维或更高的head维度点积结果的方差会成倍放大。你可以自己做一个验证实验d 128 q torch.randn(32, d) k torch.randn(32, d) logits q k.T print(logits.std().item()) # 大约11.3 logits_scaled logits / (d ** 0.5) print(logits_scaled.std().item()) # 大约1.0d128时标准差异常放大到11倍这意味着logits的分布拖尾很长softmax会有大量输出逼近0或1反向传播过程中这些位置的梯度极其微小。训练前期本身就脆弱再摊上这个缩放错误loss曲线会经常出现平台期模型看起来“学不进去”。如果你是在用现成框架这一步大概率已经帮你处理好如果你像我早期一样总喜欢自己造轮子千万别删掉这个除号。5.3 注意力坍缩所有头同时失效的现场训练Transformer时有一个非常隐蔽的失败模式loss在下降但是模型几乎不学习位置之间的互动而是把所有头的注意力权重逼向一个one-hot分布每个query都死盯某个固定token比如序列的第一个token或上一个token。从核回归的视角看这相当于把带宽h调到了极端小模型退化成只按“最近邻”做决策的非参数模型完全丧失了对全局上下文建模的能力。诊断这个问题的工具是注意力熵。每个位置的attention分布求熵如果熵长期趋近0或者远低于同层其他头就要警惕# 传入注意力权重矩阵 (B, H, N, N) entropy -(attn * torch.log(attn 1e-9)).sum(dim-1) # 如果entropy均值很低意味着注意力过于尖锐遇到这种坍缩我通常先调低注意力logits的温度放大初始熵值其次检查初始化方式让W_Q、W_K的初始值分布别太尖锐最后再看是否学习率过大导致模型早年走极端。这个排查顺序按照“由简到繁”多数情况都不用动网络结构就能把注意力拉回正轨。5.4 KV cache带来的训练-推理不对称最后说一个训练时感受不到、推理时才暴露的问题。训练阶段整个序列是一次性送进模型的所有位置的key、value都是一次算好的但推理阶段尤其是自回归生成时每生成一个新tokenquery都会变理论上要重新计算它对全序列所有key的匹配权重。如果每次都重新算一遍全部key和value复杂度会随序列长度平方上涨非常浪费。标准解法是KV cache把已经算过的key和value缓存在显存里每步只用新token的query去和缓存的K、V做注意力计算。这件事放在核回归里其实很自然核回归的训练集是一旦给定就不再变的预测新点时你只需要带着训练集里的(x_i, y_i)去算新权重就行不需要重新生成历史样本。对应到推理时的Transformer过去位置已经算好的k、v就是固定样本缓存复用即可。想深入优化推理时延的人很多功夫实际上都花在怎么更高效地组织这份“核回归样本集”上。我在实际使用中发现很多人对注意力机制和Transformer里面QKV的理解偏差根源在于从来没有从统计学的角度看过加权平均这件事。如果你能花一个下午把Nadaraya-Watson核回归的代码跑一遍调节几个不同的带宽h再回头去看Transformer源码你会觉得那些曾经显得跳脱的softmax(QK^T/√d)V突然就前后对上了。这个实验值得做而且做完之后线性注意力、随机特征核近似、KV cache压缩等进阶话题都会变得好理解得多。