1. 注意力评价函数到底在算什么,为什么它是注意力的命门
先抛一个我经常在组会被问到的问题:你写的注意力模块里,那行score = torch.matmul(q, k.transpose(-2, -1))到底在干嘛?很多同学能背出softmax(QK^T / sqrt(d)) V,但一旦追问“为什么是点积,换成别的行不行”“缩放因子为什么非要除以根号d”,就立刻卡壳。
这个被绝大多数教程一笔带过的环节,就是注意力中的评价函数(scoring function),也叫注意力打分函数、相似度函数。它的唯一职责是:给定一个查询向量query和一个被查询的向量key,输出一个实数值score,表示“这两个向量有多匹配”。这个分数经过softmax归一化之后,就变成了施加在value上的权重。
可以这样理解:假如你是组长,要分配一项任务给组里的三个人。你脑子里的“需求清单”是 query,每个人的“能力标签”是 key,你说的“你们仨各自拿多少工作量”就取决于你对每组(需求,能力)的匹配程度打分。这个打分方式就是评价函数。评分给多少分、怎么算匹配,直接决定了任务怎么分。在注意力机制里,同样一行代码,换一个评价函数,模型的表现可以天差地别。
所以别看它只是整套流程里“先算一个数”的小步骤,它才是注意力的内核。整篇文章我就围绕这“一个数怎么算”展开,把主流的加性模型、点积模型、缩放点积模型、双线性模型全部拆开讲清楚,配合推导、代码、踩坑记录和选型建议。无论你是在搭 Transformer、做多模态对齐,还是只想搞清楚为什么QK^T前面要挂个1/sqrt(d),这篇都能给你一个直接可以抄走的答案。
2. 五类主流评价函数:原理拆解与数学细节
2.1 加性模型:最早被广泛使用的注意力打分方式
先看最早被大规模验证可行的结构——加性注意力,它把query和key拼接(concatenate)之后送入一个单隐层网络,让网络自己去学“什么才算匹配”。
假设query的维度是d_q,key的维度是d_k,加性模型的打分公式长这样:
score(q, k) = v^T · tanh(W_q · q + W_k · k)其中W_q是d_h × d_q的矩阵,W_k是d_h × d_k的矩阵,v是d_h维的向量,d_h是隐层维度。这里的v^T相当于对tanh输出的每个维度做加权求和,把所有信息压成一个标量分数。
为什么要用tanh而不是ReLU?早期做机器翻译时实验发现,tanh的输出范围是(-1,1),对称且有界,训练更稳定;ReLU会把负值直接截断成 0,在这个场景下容易让很多(query, key)对的分数完全一样,学不出差异性。虽然现在也有用ReLU的变体,但经典结构还是tanh更稳。
从表达能力的角度看,加性模型是理论上最强的评价函数之一:因为它内部有一个可学习的隐层,原则上可以拟合任意复杂的“匹配关系”,不只是简单的向量内积。你甚至可以把它当成一个小的分类器,输入 (q, k) 的拼接特征,输出匹配分数。
但这东西的缺点也很明显:计算量比较大。假设序列长度是L,你需要对每一对 (q, k) 都做一次矩阵乘法和激活,复杂度是O(L^2 · d),而且这里的d不仅仅是向量维度,还多了一层隐层矩阵的运算。更麻烦的是,由于W_q·q和W_k·k可以提前算好,加性模型在解码时还能做点缓存优化,但在编码阶段的自注意力里优势不大。
2.2 点积模型:直觉上最自然,空域上最直接
既然要衡量“匹配”,很多人脑子里冒出来的第一个想法就是:向量内积。两个向量越“像”,它们在同一方向上的分量越大,内积自然越大。
点积评价函数是最简单的:
score(q, k) = q^T · k这个公式没有额外参数,计算极快,尤其在 GPU 上可以用矩阵乘法一次算完所有配对分数。Q的形状是(batch, seq_len_q, d_q),K的形状是(batch, seq_len_k, d_k),直接Q @ K^T得到(batch, seq_len_q, seq_len_k)的分数矩阵,效率极高。
但点积模型有个隐藏前提:它默认query和key的各个维度之间是独立且等权重的。实际情况往往不是这样——query的第 3 个维度可能比第 1 个维度重要得多,或者第 2 和第 5 个维度需要组合起来才有意义。点积模型无法学习这种特征重要性,它只能按原始坐标方向做内积。
还有一个让点积模型翻车的细节:如果向量维度d很大,内积的数值会自然变得很大。因为每个维度上都是“均值 0、方差 1”的随机变量(经过合理初始化后),d个独立随机变量的乘积之和,方差会随着d线性增长。方差大了,内积的绝对值就容易大到离谱,直接送进softmax之后,梯度会小到几乎消失。关于这个问题,下一节讲缩放点积时会详细算一笔账。
2.3 缩放点积模型:为什么偏偏除以根号d
缩放点积是 Transformer 的默认方案,公式是:
score(q, k) = q^T · k / sqrt(d_k)和普通点积唯一的区别就是除以sqrt(d_k)。这个sqrt(d_k)到底起什么作用?我把推导过程完整讲一遍。
先假设q和k都是d_k维向量,每个维度的随机变量独立且均值为 0、方差为 1(这是很多初始化方法的目标状态)。那么q_i · k_i的期望是 0,方差呢?
如果q_i和k_i独立,那么Var(q_i · k_i) = E[(q_i · k_i)^2] - (E[q_i · k_i])^2 = E[q_i^2] · E[k_i^2] - 0 = 1 · 1 = 1。
由于各个维度独立,整个内积q^T · k = Σ(q_i · k_i)的方差等于各维度方差之和,也就是d_k。
所以内积的方差是d_k,标准差是sqrt(d_k)。除以sqrt(d_k)之后,分数的方差被拉回1,这直接缓解了 softmax 饱和问题。softmax 在输入特别大或特别小时(比如 ±20 以上),梯度会逼近 0,模型几乎学不动。除以缩放因子之后,分数的尺度被控制在合理范围内,梯度能顺畅回传。
这也是为什么缩放因子一定要除以sqrt(d_k)而不是d_k:因为标准差是sqrt(d_k),不是d_k。如果除以d_k,方差就变成1/d_k,分数被压得太小,softmax 的输出会几乎均匀分布,注意力就分不出主次了。
2.4 双线性模型:给点积加一个可学习的“旋转”
如果觉得普通点积太死板,又不想上加性模型那么贵的结构,可以选双线性模型:
score(q, k) = q^T · W · k这里的W是一个d_q × d_k的矩阵,完全可学习。直观地说,相当于在q和k之间插入了一个线性变换,模型可以学习“先对 q 做某种变换,再和 k 做内积”。
这个东西的数学解释是:它可以改变q和k所在空间的度量方式。普通点积假设“坐标轴正交且等尺度”,而双线性模型允许模型学习一个更合适的坐标系——可以把某些维度放大、某些维度缩小,甚至让某些维度之间发生耦合。这在跨模态任务里非常有用,比如query是文本向量、key是图像区域向量,两者来自不同的嵌入空间,直接点积没有意义,但通过一个可学习的W可以把它们映射到一个公共语义空间再比较。
双线性模型和加性模型比,参数量少很多,计算也更快;和普通点积比,表达能力更强。它的问题在于W的初始化比较敏感。初始化太大会导致分数范围不可控,初始化太小又会让模型一开始几乎“看不到”任何匹配关系,训练前期损失不下降。我通常建议把W初始化为单位矩阵附近,这样起步时等价于普通点积,再让模型自己去学偏离。
2.5 其他值得留意的变体:拼接映射、余弦相似度、偏置注入
除了上面四种,工程里还能见到几类变体。
拼接映射:有人把query和key直接拼接,然后通过一个线性层score = w^T · [q; k]。没有激活函数,本质上是双线性模型的一种退化形式,表达能力弱一些,但实现最简洁。有些轻量级 attention 会用它。
余弦相似度:score = cos(q, k),等价于把两个向量先做 L2 归一化再做点积。好处是分数天然落在[-1, 1],数值尺度永远可控,不需要缩放因子;坏处是强制丢掉了向量的模长信息,有些场景下模长本身是重要的语义信息,丢掉会掉点。它适合用在文本检索、向量数据库召回这类本来就要算余弦相似度的任务上。
偏置注入:在打分函数之后加一个可学习的偏置项score = score + bias,或者在某些位置强制加上 mask 偏置(比如-inf)。这在图注意力网络(GAT)里很常见——把邻接矩阵信息以偏置形式注入到分数里,让模型只在有边的节点之间分配注意力。严格来说这不算新的评价函数,但对打分结果的影响非常大,后面聊掩码时会再展开。
3. 评价函数选型:复杂度、表达能力与场景的三角权衡
3.1 计算复杂度与内存开销对比
先给一张直接可以拿去写 PPT 的对比表,基于查询序列长度L_q、键序列长度L_k、特征维度d来算:
| 评价函数 | 公式 | 时间复杂度 | 额外参数 | 适用场景 |
|---|---|---|---|---|
| 点积 | q^T k | O(L_q·L_k·d) | 无 | 快速原型、序列长度较短、各维度独立性强 |
| 缩放点积 | q^T k / sqrt(d) | O(L_q·L_k·d) | 无 | Transformer 及绝大多数通用场景 |
| 双线性 | q^T W k | O(L_q·L_k·d) | d×d | 跨模态对齐、需要特征重定向 |
| 加性 | v^T tanh(W_q q + W_k k) | O(L_q·L_k·d²)级别 | 约2d² + d | 小规模序列、理论表达能力优先 |
| 余弦 | `q·k / ( | q |
注意加性模型的时间复杂度,我特意写了“O(L_q·L_k·d²)级别”。表面上看它也是每个(q,k)对做一次计算,但内部涉及隐层矩阵乘法和非线性激活,实际开销比点积系高一个档次。在序列长度 512、特征维度 512 的场景下,加性模型比缩放点积慢 3 到 5 倍是很正常的。
内存方面,所有需要显式构建(L_q, L_k)分数矩阵的方法,显存占用都是O(L_q·L_k)。这条在长序列场景下才是真正的瓶颈,评价函数本身反而不是关键。
3.2 表达能力与可解释性的差异
表达能力上,我的判断标准很简单:能不能表达非线性匹配关系。
- 点积和余弦都是固定的线性度量,不能学习“什么维度重要”;
- 双线性可以学维度重要性,但整体仍然是线性变换加内积,表达的是二次型关系;
- 加性模型带有
tanh非线性,理论上可以逼近任意匹配函数,表达能力天花板最高。
但表达能力越强,不代表效果越好。实际项目中见过太多加性模型在中小规模数据上反而跑不过缩放点积的情况。原因不复杂:可学习参数多了,优化难度跟着上来,数据量不够时模型容易过拟合到训练集上的噪声关系,泛化反而更差。
可解释性也有差异。点积和余弦的计算过程完全透明,你能直接把 attention 分数可视化出来,看看模型关注了哪些 token,这在 NLP 里是常见的调试手段。加性模型经过隐层变换和激活之后,分数来源就变得模糊了,可视化出来的热力图经常看不出明显规律。
3.3 按任务场景的实操选型建议
基于我自己做项目的经验,不同任务可以按下面的思路选:
通用 NLP 任务(文本分类、机器翻译、文本生成):默认缩放点积,没有之一。Transformer 系列已经把这个组合验证得非常充分了,硬件优化和算子融合也做得最到位,没必要为了炫技换别的。
视觉任务(图像分类、目标检测中的注意力):同样推荐缩放点积。视觉特征维度通常比较高(d可以到 1024),缩放因子对稳定训练帮助很大。很多视觉模型还喜欢在评价函数后面直接加一个可学习的温度系数(temperature),本质上就是score = q^T k / (sqrt(d) · t),让模型自己调节 softmax 的锐利程度。
跨模态对齐(图文匹配、语音文本对齐):优先试双线性模型。文本向量和图像/音频特征不在同一个空间,直接点积等于让两个坐标系强行对齐,结果通常不好。双线性模型的W相当于一个可学习的“桥梁”,效果显著优于普通点积,而且只增加d²参数,算力可以接受。
轻量级检索场景(向量数据库召回、候选生成):用余弦相似度。它天然和 FAISS 这类索引的度量方式一致,而且模型训练完可以直接把权重导出为 embedding,不需要额外适配。
图结构数据(GAT 等图注意力网络):加性模型反而是主流。因为图注意力里的节点特征维度和邻居数量通常不大,加性模型那点计算开销可以接受,而它强大的表达能力对建模节点间复杂的相互作用有实打实的好处。
4. 手写实现与调试实录:从理论到可运行的注意力模块
4.1 用 PyTorch 实现五种评价函数
理论讲再多,不如看代码。我贴一段可以直接运行的最小实现,包含五种评价函数,输入输出格式对齐,方便你在自己的模型里替换测试。
import torch import torch.nn as nn import torch.nn.functional as F import math class ScoreFunctions(nn.Module): def __init__(self, d_model, score_type='scaled_dot', bias=True): super().__init__() self.d_model = d_model self.score_type = score_type if score_type == 'additive': self.d_h = d_model self.W_q = nn.Linear(d_model, self.d_h, bias=bias) self.W_k = nn.Linear(d_model, self.d_h, bias=bias) self.v = nn.Linear(self.d_h, 1, bias=False) elif score_type == 'bilinear': self.W = nn.Parameter(torch.eye(d_model) + 0.01 * torch.randn(d_model, d_model)) elif score_type == 'concat': self.proj = nn.Linear(2 * d_model, 1, bias=False) def forward(self, q, k): # q: (batch, L_q, d), k: (batch, L_k, d) if self.score_type == 'dot': return torch.matmul(q, k.transpose(-2, -1)) elif self.score_type == 'scaled_dot': return torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_model) elif self.score_type == 'cosine': q_norm = F.normalize(q, dim=-1) k_norm = F.normalize(k, dim=-1) return torch.matmul(q_norm, k_norm.transpose(-2, -1)) elif self.score_type == 'bilinear': return torch.matmul(torch.matmul(q, self.W), k.transpose(-2, -1)) elif self.score_type == 'concat': q_expanded = q.unsqueeze(2).expand(-1, -1, k.size(1), -1) # (batch, L_q, L_k, d) k_expanded = k.unsqueeze(1).expand(-1, q.size(1), -1, -1) # (batch, L_q, L_k, d) concat = torch.cat([q_expanded, k_expanded], dim=-1) return self.proj(concat).squeeze(-1) elif self.score_type == 'additive': q_proj = self.W_q(q).unsqueeze(2) # (batch, L_q, 1, d_h) k_proj = self.W_k(k).unsqueeze(1) # (batch, 1, L_k, d_h) return self.v(torch.tanh(q_proj + k_proj)).squeeze(-1) # (batch, L_q, L_k) else: raise ValueError(f"Unknown score type: {self.score_type}")这里有一个很容易踩的坑,我特意在代码注释里标出来了:加性模型的前向传播不要用 for 循环逐对计算。初学者最容易写出来的版本就是双循环遍历所有 (i, j) 组合,数据量一上来直接被慢死。正确做法是借助广播机制,先把q扩展成(batch, L_q, 1, d_h),把k扩展成(batch, 1, L_k, d_h),两者相加时自动广播成(batch, L_q, L_k, d_h),一次完成所有位置的计算。
4.2 维度对齐与广播机制:最常见报错来源
我第一次写加性模型时,报错信息乱七八糟,后来把维度对齐规则彻底理清楚了,现在给你梳理一遍。
query和key的最后一维必须相同?不一定。双线性模型中W的维度是d_q × d_k,允许 query 和 key 的特征维度不同。其他模型里,余弦相似度和点积要求维度严格一致,加性模型因为先投影到一个公共隐层,理论上也可以接受不同维度。
完整的分数矩阵维度规则:
- 点积和缩放点积要求
d_q == d_k,输出(batch, L_q, L_k); - 加性模型要求
W_q的输入维等于d_q,W_k的输入维等于d_k,输出(batch, L_q, L_k); - 双线性模型要求
W的维度是d_q × d_k,输出(batch, L_q, L_k)。
我自己调试时最常用的排查手法:在 forward 里临时打印每个中间结果的 shape,从前往后逐个对。90% 的 bug 都是某个矩阵转置忘了,或者 unsqueeze/expand 的维度次序写反了。
4.3 掩码机制如何与评价函数配合
注意力机制里几乎不可能不碰 mask。padding mask 要求query不能 attend 到那些没有实际意义的 padding token;causal mask 要求 decoder 在生成时只能看到当前位置之前的信息。
实现方式是在评价函数算出的分数矩阵上做加法掩码:
def apply_mask(scores, mask): # mask 中为 1 的位置表示需要屏蔽 scores = scores.masked_fill(mask.unsqueeze(1).unsqueeze(1) == 1, float('-inf')) return scores把被屏蔽位置的分数设成-inf,经过 softmax 之后权重精确地变成 0,等价于该位置不参与注意力。这个设计比“把权重设成 0”更优雅,因为 softmax 的归一化会重新分配权重给未屏蔽的位置。
这里有个工程细节值得记住:mask 一定要在评价函数算完之后、softmax 之前应用。如果在评价函数之前就把 key 的位置换成零向量,等于往评价函数里塞了一堆无效信息,结果依然会受影响。正确顺序是:算分数 → 套 mask → softmax → 加权求和。
5. 常见问题与排查技巧实录
5.1 注意力分布过于尖锐或过于平滑怎么办
最常见的一个表现:训练初期 attention 权重几乎均匀分布,所有位置都是 0.1、0.1、0.1,模型看不出该关注哪。这通常不是评价函数写错了,而是分数尺度太小。检查一下缩放因子有没有生效,或者初始化权重是不是太接近 0。
反过来,如果 attention 权重出现“一枝独秀”的极端分布(比如某个位置权重 0.99,其他几乎为 0),容易导致梯度过度集中在一个 token 上,训练不稳定。可以加一个温度系数调节:
score = q^T k / (sqrt(d) · tau)tau大于 1 时分布变平滑,小于 1 时变尖锐。在很多可视化工具里它叫 temperature,调试时先从tau = 1.2开始试,看看 attention 热力图是否更合理。
5.2 梯度消失与数值不稳定的排查思路
如果用非缩放点积在高维空间训练 Transformer,很容易遇到 loss 不降、甚至直接 NaN 的情况。原因前面已经算过了:d大时内积方差也大,softmax 进入饱和区,反向传播梯度趋近于 0。
还有一类隐蔽的数值问题出现在混合精度训练里。fp16的表示范围比 fp32 窄很多,-inf掩码在 fp16 下偶尔会变成很大的负值或者 NaN。如果你在 AMP 模式下发现 attention 分数异常,先把核心打分部分强制切成fp32试试。PyTorch 里可以对特定模块用torch.cuda.amp.autocast(enabled=False)关掉自动混合精度。
5.3 加性模型收敛慢的排查
加性模型如果收敛明显偏慢,优先检查W_q和W_k的初始化方差是不是太大了。隐层输出一旦过大,tanh直接饱和,反向传到前面的梯度几乎为零。给这两个线性层设置更小的初始化方差,比如nn.init.normal_(w, std=0.01),能让训练前期更顺畅。
另外,使用加性模型时建议给v这层也做一点正则。因为v直接决定了“哪些隐层维度对分数有贡献”,它学崩了,整个打分函数就废了。我在项目里会在v上加 L2 正则,系数设在1e-5到1e-4之间,效果比不加稳定不少。
5.4 怎么判断评价函数学得好不好
一个很实用的经验指标:直接看 attention 热力图和任务表现的相关性。
比如在文本分类任务里,你预期模型关注关键主题词;在机器翻译里,翻译某个词时应该对齐原文的对应词。如果热力图完全看不出语义关联,说明评价函数可能没有学到有效特征——但这个判断要小心,有些任务(比如深层的 Transformer 层)attention 模式本来就不是人眼可解释的,不能一概而论。
另一个更客观的指标是看分数矩阵的熵。注意力分布的熵过高(接近均匀)说明模型“注意力涣散”;熵过低说明“过度专注”。在不同层之间对比熵的变化趋势,能帮你判断是评价函数的温度设置出了问题,还是整体网络结构需要调整。
5.5 我总结的几条避坑原则
我做了不少注意力相关的项目,把踩过的坑整理成几条原则:
- 别轻易把缩放因子改成别的东西。想看效果差异时,先固定缩放点积跑到一个 baseline,再换其他评价函数,否则你根本分不清是评价函数的差异还是训练不充分导致的差异。
- 注意力分数不能直接当置信度用。softmax 分数只表示相对重要性,不代表模型的“确信程度”。有人拿 attention 权重当可解释性证据,这需要非常谨慎。
- 不同层的 attention 模式本来就不同。浅层更关注局部邻近 token,深层更关注全局语义关系。如果你只在某一层换了评价函数,效果可能被其他层的模式掩盖,要做就做全套。
- 先跑小规模消融实验再大规模铺开。换评价函数时,先在一个小数据集上做 10 个 epoch 内的对比,看 loss 下降趋势和最终指标。这样可以节省大量试错时间。
6. 从评价函数出发的扩展思考
6.1 线性注意力对评价函数的改造
经典的QK^T打分方式复杂度是O(L²),序列一长就吃不消。所以出现了很多“线性注意力”方案,它们的核心思路都是换掉或近似掉评价函数。
最简单的一种做法:把softmax(QK^T)近似成φ(Q)·φ(K)^T,其中φ是一个非线性核映射。这样可以用矩阵结合律先算φ(K)^T V,把复杂度降到O(L)。评价函数在这里被替换成一个核函数,本质上是在另一个特征空间里算匹配度。效果比原始 softmax 差一些,但长序列推理时收益极大。
我在做长文档处理时试过这种方案,L = 4096的序列,普通缩放点积直接爆显存,线性注意力则轻松跑起来,速度提升非常明显。代价是效果下降了一些,尤其是在需要精细对齐的任务上。后来我用了折中方案:局部窗口用缩放点积,跨窗口用线性注意力,效果和性能都兼顾了。
6.2 稀疏注意力与 Top-K 选择
另一种想绕开O(L²)的思路是稀疏注意力。核心思路是:不要对所有 (query, key) 对打分,只对一部分候选打分。
具体到评价函数上,可以先算一个粗略的分数,然后只取 Top-K 的 key 参与完整的注意力计算。这个“粗略分数”和“最终分数”可以用不同的评价函数,甚至用同一个函数的两阶段近似。更进一步,固定窗口注意力、全局 token 注意力、随机注意力的混合,本质上都是在限制评价函数的“作用范围”,让模型在有限的候选集里做精确的匹配。
这部分工程实现比较复杂,但思路值得借鉴:评价函数不一定要一次性算完整矩阵,它可以和“候选筛选”结合起来。
6.3 我在实际项目中的几条心得
最后分享几条我做注意力相关项目攒下来的实际经验。
关于加性模型的隐藏高度:如果你确实需要加性模型的表达能力,d_h不需要设得跟d_model一样大。很多实验表明d_h = d_model // 2甚至更小就够用了,再往上收益递减,计算开销却线性增长。
关于双线性模型的正则:W矩阵如果完全自由训练,很容易出现秩不足的问题,也就是学出来的线性变换是“扁”的,丢掉了部分信息。可以考虑加一个正交性正则,让W^T W尽量接近单位矩阵,这样既保住了表达能力,又提升了训练稳定性。
关于做实验的习惯:换评价函数之前,先跑一个相同的 baseline,确认你当前的训练管线本身没问题。我以前有一次折腾了一个星期调双线性模型,最后发现是学习率调度器的 warmup 步数写错了,跟评价函数一点关系都没有。先把环境变量、数据迭代、mask 生成这些基础环节验证好,再动模型结构,真的是血泪教训。
关于打通理论和实现:纸上推导和实际工程之间的差距,往往比想象中大得多。比如数学上两个评价函数“等价”,但换个硬件、换种实现方式,数值误差就对结果产生了不可忽略的影响。所以做实验时,建议同时记录 baseline 和实验组的注意力热力图,肉眼对比一下,很多问题一眼就能看出来。
说回最本质的一点:评价函数是注意力机制的“口味偏好”。你让它偏向什么,它就更关注什么。理解清楚每种评价函数擅长和不擅长什么,才能在面对一个新任务时,少走很多弯路。这个方向看着小,但值得你花时间深入。