我发现很多人学注意力机制是直接从Transformer源码看起,对着softmax(QK^T)V那一行反复琢磨,却始终觉得Q、K、V三个矩阵像是天上掉下来的。我自己也是,直到有一次手写Nadaraya-Watson核回归,把指数核函数往公式里一代,输出来一条和注意力加权结果几乎一样的曲线,我才意识到:注意力机制的本质就是一次带学习的核回归。这篇博文想把这条少有人讲的线索讲清楚,从核回归的数学形式出发,一步步推到你每天在用的Transformer注意力架构,最后给出可以直接跑通的PyTorch实验代码和一些训练阶段的工程细节。
1. 为什么注意力机制要追溯到Nadaraya-Watson核回归
1.1 核回归在统计里解决什么问题
先说一个最朴素的统计问题。假设你有一堆散点,横坐标是样本位置,纵坐标是样本值,现在要在任意一个新位置做预测,传统回归会强行假设一个函数形式,比如线性回归假设y=ax+b。可如果真实关系是一条复杂的波浪线,你很难提前知道该用几次多项式。
Nadaraya-Watson核回归的思路完全不同:它不假设函数形式,而是让每个样本根据“与当前查询点的距离”来分配权重,距离越近贡献越大。这个“距离”由一个核函数定义,最常见的是高斯核:
$$K_h(x - x_i) = \exp\left(-\frac{(x - x_i)^2}{2h^2}\right)$$
其中h是带宽,控制影响力衰减的速度。预测公式是:
$$\hat{y}(x) = \frac{\sum_{i=1}^{n} K_h(x - x_i) y_i}{\sum_{i=1}^{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核回归 | 注意力机制 | 作用 |
|---|---|---|
| 查询点x | query q | 当前要预测/检索的对象 |
| 样本坐标x_i | key k_i | 用来与查询做匹配的索引 |
| 样本标签y_i | value v_i | 真正被聚合的内容 |
| 核函数K_h | exp(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_{i=1}^d q_i k_i$$
点积的均值是0,方差呢?因为q_i和k_i相互独立,E(q_i²)=1、E(k_i²)=1,所以:
$$\text{Var}(q \cdot k) = \sum_{i=1}^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这三个值跑一遍:
h=0.05时,预测曲线几乎穿过每一个噪声点,拟合了太多噪声细节,这就是过拟合。放到注意力机制里,相当于温度极低、注意力分布极度尖锐,模型对单个token极其敏感,一有波动就跟着跑。h=0.5时,曲线既保留了波峰波谷的趋势,又没被噪声带偏,是比较理想的平滑程度。h=10时,所有训练点的权重都趋于相同,预测结果几乎就是全体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, bias=False) self.w_k = nn.Linear(in_features, d_k, bias=False) self.w_v = nn.Linear(in_features, d_k, bias=False) 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.0d=128时标准差异常放大到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都是一次算好的;但推理阶段尤其是自回归生成时,每生成一个新token,query都会变,理论上要重新计算它对全序列所有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压缩等进阶话题,都会变得好理解得多。