我第一次认真读Transformer论文时,注意力那张图盯了半小时,每个箭头都认识,连起来就是不知道在干什么。后来帮别人调模型,发现大多数人卡住的都是同一个点:自注意力(self-attention)和多头注意力(multi-head attention)到底是怎么从一堆矩阵乘法里长出“注意力”的。这篇文章是这个Transformer系列的第2篇,我会把这两个概念拆成零件从头讲起。如果你已经听说过Q、K、V,也看过公式,但心里总有个“尺寸怎么变”的疙瘩,那今天这篇就是为你准备的。
1. 从接力棒到全串联:自注意力凭什么取代串行建模
1.1 为什么RNN处理长序列那么“拧巴”
在Transformer出现之前,处理序列的主流工具是RNN、LSTM这类循环神经网络。它们的运行方式可以想象成接力赛:一个词的信息要经过隐藏状态逐个传给下一个词,每个时间步只能拿到前一个时间步的压缩信息。两个词在句子里隔得越远,信息传递的次数就越多,路径就越长。路径一长,梯度要么指数衰减下去(梯度消失),要么指数爆炸上来(梯度爆炸)。LSTM和GRU用门控机制缓解了衰减问题,但根本上还是串行的——每一步都得等前面算完才能动,GPU并行能力被浪费了一大截。
CNN其实也常被用来做序列建模,但它的感受野是局部的,哪怕把卷积窗口开得很大,想覆盖长距离依赖也得靠加深网络层数或者使用膨胀卷积。层数一多,信号路径还是拉长了。也就是说,在RNN和CNN的框架里,两个相隔很远的词建立联系需要“翻山越岭”,沿途的信息损耗会让模型难以判断它们之间到底是强关联还是弱关联。
自注意力的想法非常“偷懒”:不走接力赛了,直接让每个词和序列里的所有词同时计算相关性,然后按相关性加权融合彼此的信息。这样一来,任何两个位置之间的交互路径长度都是1,不存在“传很多棒”的问题。这个特性后来被称为长程依赖的直接建模,也是Transformer能在机器翻译、文本生成等任务上大幅度超越循环模型的原因之一。
1.2 自注意力把所有位置拉到同一张桌子上
注意力的雏形其实就是加权求和。假设有一个序列 (x_1, x_2, ..., x_n),位置 (i) 的表示更新成:
[ z_i = \sum_{j} \alpha_{ij} x_j ]
其中 (\alpha_{ij}) 表示位置 (i) 对位置 (j) 的关注权重,所有权重加起来为1。那这个权重从哪里来?原始的做法是用点积相似度:两个向量方向越一致,点积越大,说明语义上越相关;再通过softmax归一化,变成一个概率分布。
这套机制有个非常关键的“灵魂”差异——权重不是固定的。CNN的卷积核是“一套模板走天下”,不管输入是什么,用的都是同一套参数;而注意力权重是输入自己算出来的。同一句话,换个语境,同一对词的关联强度就会发生变化。举个例子,“它很可爱”里的“它”到底指猫还是狗,注意力权重会根据上下文动态地指过去。这个动态特性才是注意力机制真正强大的地方。
顺着这个思路自然就引出了一个问题:点积计算需要“两个东西”,一个代表“当前词想找什么”,一个代表“被找词能提供什么”。这正是下一章要拆解的Q、K、V。
2. Q、K、V的来龙去脉与缩放点积公式的数学理由
2.1 从“查资料”理解Q、K、V
自注意力的完整公式长这样:
[ Attention(Q,K,V)=softmax(\frac{QK^T}{\sqrt{d_k}})V ]
这个公式看起来突然,实际上用“查资料”的场景特别容易理解。假设你在搜索引擎里输入一句话找文档:Q(Query)是你输入的搜索词,代表“我想找什么”;K(Key)是每个文档的索引标签,代表“我有什么”;V(Value)是文档的正文内容,代表“我实际提供的信息”。搜索引擎先用搜索词和所有文档的标签做匹配,找到匹配度最高的几个文档,然后把对应正文提取出来混合给你。
落在Transformer里,Q、K、V都来自同一个输入序列 (X),只是分别乘上不同的权重矩阵:
[ Q = XW_Q,\quad K = XW_K,\quad V = XW_V ]
输入向量维度是 (d_{model}),注意力头的工作维度是 (d_k)(实践中一般 (d_k = d_v),都等于head_dim)。于是:
- (W_Q) 的形状是 (d_{model} \times d_k),把每个词映射成“查询向量”;
- (W_K) 的形状是 (d_{model} \times d_k),把每个词映射成“索引向量”;
- (W_V) 的形状是 (d_{model} \times d_v),把每个词映射成“内容向量”。
看完矩阵形状,很多疑惑其实就解开了一半。(QK^T) 是在做两两词之间的点积相似度,得到形状为 (seq_len \times seq_len) 的分数矩阵。分数矩阵的第 (i) 行,就是第 (i) 个词对序列里所有词的关注分数。对这个矩阵的每一行做softmax,就得到了归一化权重;再乘上 (V),相当于按权重把其他位置的“内容向量”加权求和,得到当前位置的新表示。
2.2 一个手工算例子:让符号落下来
很多人看公式最怕的是符号飘在空中。我拿一个极小的例子走一遍。假设序列里只有三个词:“猫”“追”“狗”,每个词的词向量维度先临时设成4,注意力头维度 (d_k = 2)。简化后,三个词的Q、K、V分别是:
[ Q = \begin{bmatrix} 0.8 & 0.2 \ 0.5 & -0.3 \ -0.2 & 0.9 \end{bmatrix},\quad K = \begin{bmatrix} 0.6 & 0.1 \ -0.4 & 0.7 \ 0.2 & -0.5 \end{bmatrix} ]
计算第一行(“猫”对所有词的注意力分数):
- “猫”对“猫”:(0.8 \times 0.6 + 0.2 \times 0.1 = 0.50)
- “猫”对“追”:(0.8 \times (-0.4) + 0.2 \times 0.7 = -0.18)
- “猫”对“狗”:(0.8 \times 0.2 + 0.2 \times (-0.5) = 0.06)
每个分数除以 (\sqrt{d_k} = \sqrt{2} \approx 1.414),再softmax,得到大约 ([0.43, 0.26, 0.31])。也就是说,“猫”这个词的新表示,会由“猫本身40%多”“追26%左右”“狗31%左右”的内容混合而成。实际模型学出来的权重会复杂很多,但底层逻辑就是这么朴素:算相似度,归一化,加权求和。
建议读者拿到一个真实模型后,把某一层的平均注意力权重打印出来,跟这个手工例子对照着看,很快就能建立“注意力分数到底在做什么”的直觉。
2.3 为什么缩放偏偏是除以根号dk,而不是dk
公式里那个 (\frac{1}{\sqrt{d_k}}) 是最容易被忽略、但极为重要的细节。它解决的是一个数学上的“尺度失控”问题。
假设 (q) 和 (k) 的每个分量都是均值为0、方差为1的独立随机变量。那么两个向量做点积,结果的期望是0,但方差是 (d_k)。也就是说,向量维度越大,点积的绝对值普遍越大,很容易累积到几十甚至上百。而softmax函数对过大的输入非常敏感:当输入分数差异很大时,softmax输出会迅速逼近one-hot分布,也就是模型几乎只关注分数最高的那个词,其他位置梯度很小,训练效率直线下降。
除以 (\sqrt{d_k}) 后,点积的方差被拉回1附近,softmax还能保持一段动态范围较大的区域,让模型有空间去“犹豫”到底该关注谁。这个操作本身不复杂,但去掉它之后,小模型训练时最容易观察到的现象就是loss震荡、注意力矩阵变成稀疏尖刺。我自己试过一次,去掉缩放因子之后,同样的数据和超参数,模型收敛速度明显变慢,最终效果也差了一截。
3. 多头注意力:用多副眼镜同时看序列
3.1 单头注意力的局限:一种相似度不够用
单一注意力只能计算一种“相似度”,这在实际语料里是不够用的。词和词之间的关系其实是多维的:有语法上的主谓关系,有指代上的“它指向前文某个名词”,有局部修饰关系(形容词和名词挨得很近),还有跨越长距离的因果推理。如果只有一个头,模型就不得不把所有这些关系揉进同一个打分函数里,彼此干扰,最后哪个都学不利索。
多头注意力的思路很直接:多开几个独立的自注意力通道,每个通道用不同的随机初始化权重,让它们在训练中自然分化,各学各的关系模式。这就好比看一个问题时,既要用语法视角扫一遍,又要用语义视角扫一遍,还要用指代关系视角扫一遍;每个视角输出一份侧重点不同的结果,最后合并在一起做综合判断。
3.2 维度流:从d_model切开再拼回去
多头注意力的维度变化是很多人第一次上手时最容易绕晕的地方。假设 (d_{model}=512),头数 (h=8),那么每个头的维度就是 (d_k=d_v=64)。流程如下:
- 输入 (X) 形状是 ([seq_len, 512]);
- 分别乘上 (W_Q、W_K、W_V) 得到三组形状为 ([seq_len, 512]) 的矩阵;
- 把512维切成8段,每段64维,相当于8个头各自拿到自己的Q、K、V;
- 每个头独立完成缩放点积注意力,输出 ([seq_len, 64]);
- 把8个头的输出按最后一维拼接,回到 ([seq_len, 512]);
- 乘上输出投影矩阵 (W_O),完成头与头之间的信息交流。
这里最关键的认知是:头的拆分不是复制多份完整模型,而是把一个大维度切碎分给多个“专家”。每个头只能看到64维的子空间,增加的其实不是参数量,而是“看问题角度的数量”。有些初学者以为多头应该把模型参数扩大8倍,实际并不是这样,参数量只多了一个输出投影矩阵和Q/K/V内部的矩阵,整体开销相当可控。
3.3 头之间真的各司其职吗
最初可视化注意力权重时,大家确实观察到一些头有明确的功能倾向:有的头喜欢关注下一个词,有的头喜欢关注句号逗号,有的头擅长寻找指代词的反向先行词。这些可视化给了“头=功能模块”的直观印象,但后来研究发现事情没那么简单。
后续的剪枝实验发现,在预训练模型中删除掉相当一部分注意力头,很多下游任务几乎不掉点。这说明头之间往往存在大量冗余和协作,可能几个头一起配合才完成一类功能,单拎一个头出来并不能独立负责某种语义关系。所以我给新人的建议是:不要在训练前就硬性规定“第几个头必须学什么”,而是训练完后再做可视化观察。实践中常见的是,某些头确实表现出明显偏好,但更多的头看起来“忠诚度不高”,这都很正常。
在选头数时有个经验可以参考:head_dim一般取64附近。大模型常用 (d_{model}=768) 配12头,(d_{model}=1024) 配16头,都差不多能对上64。head_dim太小时,每个头的表达能力受限;head_dim太大时,缩放点积又容易陷入softmax饱和区,刚好撞上第2章讲的那个坑。小规模自建模型时,把head_dim设在16到64之间我都试过,效果差异不算特别大,但按64来最稳。
4. 没有顺序意识的模型,位置编码为什么绕不开
4.1 排列不变性与“词序隐形”问题
自注意力公式里其实没有“第一个词”和“最后一个词”的概念。把输入序列 (X) 的行顺序交换一下,注意力输出会跟着按同样方式交换排列,但内容本身不会发生语义变化。换句话说,如果不额外注入顺序信息,Transformer从数学上分不清“我打篮球”和“篮球打我”这两种完全不同的语义。
这个性质叫排列不变性,是自注意力的“天生短板”。人的阅读非常依赖词序,谁在动作前谁在动作后,决定了谁是施动者谁是受动者。没有词序信息,模型只能靠词向量硬猜,等于把最重要的句法线索扔掉了。因此Transformer必须在输入端给每个词附带一个“座位号”,位置编码本质上就是在做这件事。
最开始有人也尝试过用词嵌入加序号特征之类的简单办法,但效果都不够好。后来原论文给出了固定频率的正弦编码方案,才让位置信息以一种稳定、可计算的方式注入模型。
4.2 三种主流位置编码与工程取舍
位置编码方案到目前为止有好几代演进,我这里按理解和工程取舍把主流方案分成三类。
第一种是正弦/余弦固定位置编码。它给不同维度分配不同的周期,形成一个固定编码表,位置 (p) 的向量在不同维度上是不同频率的三角函数值。优点是不需要训练、理论上可以外推到训练中没见过的更长序列;缺点是它对局部相对位置关系的建模比较间接,需要模型自己从周期性里悟出来。
第二种是可学习位置编码。初始化一个 ([max_len, d_{model}]) 的位置表,跟着网络一起训练。这种做法简单直接,中小模型里非常常见,训练稳定、效果不错。但它把最大长度焊死在了训练配置里,序列一旦超出设定长度就没法外推,只能靠插值或切换方案。
第三种是相对位置编码思路,比如旋转位置编码(RoPE)。它不告诉模型“词在绝对位置p”,而是告诉模型“词i到词j偏移了多少”,把相对偏移信息直接叠进Q/K的计算里。这类方案在长文本外推上表现更好,也是当前大模型普遍采用的方向。
| 方案 | 核心思路 | 优点 | 需要注意的点 |
|---|---|---|---|
| 正弦/余弦固定 | 绝对位置 + 周期性 | 无需训练,可外推 | 对局部相对位置建模弱 |
| 可学习位置编码 | 绝对位置参数化 | 简单稳定 | 长度写死,外推差 |
| 相对位置/RoPE | 相对偏移进入注意力 | 长文本友好 | 实现复杂度稍高 |
工程上的建议是:中小规模项目想快速跑稳定,用可学习位置编码够了;一旦要做长文本方向,从一开始就用相对位置思路更省心。不要等训完再换编码方案,迁移成本很高。
5. 从公式到训练:掩码、缓存与稳定性的真实经验
5.1 真实场景里必须处理的两类掩码
纸上公式看起来清爽,一到实际训练就会碰到两个绕不开的坑:padding掩码和因果掩码。
padding掩码解决的是“batch里句子长短不一”的问题。短句子为了凑成矩阵,末尾会补零。补出来的位置没有任何真实语义,注意力应该完全忽略它们。做法是把这些位置的注意力分数加上一个极大的负数,比如 (-1e9),softmax之后权重就会趋近于0。
因果掩码解决的是自回归生成问题。训练和推理时,当前位置只能看到它左边的词,不能偷看未来。做法是把分数矩阵右上角的上三角区域全部加上 (-1e9),让“未来位置”在softmax之前就被压成无效信息。这个掩码也叫causal mask,是GPT类模型训练的核心组件之一。
softmax本身也有数值稳定性问题。当分数矩阵里出现特别大的值时,直接算 (e^x) 可能溢出成无穷大。稳妥的习惯是:先对每行分数减去该行最大值,再做指数运算。这一步在框架里往往被封装好了,但自己动手写注意力实现时一定要记得,否则跑长序列很容易突然冒出NaN。
5.2 推理时的KV Cache与FlashAttention优化逻辑
理解自注意力还不能只停留在训练视角。训练时一个句子所有位置可以并行计算,推理生成时却是逐词来的。每来一个新词,严格说可以把历史所有词重新算一遍Q/K/V注意力,但这样浪费非常大——因为之前的K/V计算根本不会变,白算一遍。
KV Cache的思路就是把已经算好的历史K和V缓存下来,每次生成新词时只需要用新词的Q去跟缓存中的K做匹配,用权重去加权缓存中的V。所以实际推理时,Transformer的内存大头往往不是模型权重,而是随上下文不断增长的KV缓存。上下文越长,缓存占用越大,这也是长文本对话显存吃紧的重要原因。
FlashAttention这类优化则是从访存角度重新设计了自注意力计算方式。它不再把完整的 (seq_len \times seq_len) 注意力矩阵写回显存再读出来,而是分块计算,让一部分计算在片上完成,大幅减少对HBM的访问。说白了,注意力矩阵是“算得快但存不下”的典型场景,用分块计算去换访存量,性价比极高。理解了这个背景,再看现代推理引擎里的各种内存管理策略就顺了。
5.3 一个从零搭出来的自注意力计算示例
下面我用numpy写一个最小可运行的自注意力实现,不依赖任何深度学习框架。建议你亲手跑一遍,把每个矩阵的shape打印出来,比看十篇文章都管用。
import numpy as np def softmax(x, mask=None): # 先减最大值,避免指数溢出;mask为0或很大的负数 x = x - np.max(x, axis=-1, keepdims=True) if mask is not None: x = x + mask e = np.exp(x) return e / np.sum(e, axis=-1, keepdims=True) def scaled_dot_product_attention(Q, K, V, mask=None): # Q/K/V shape: [seq_len, d_k] d_k = Q.shape[-1] scores = Q @ K.T / np.sqrt(d_k) # [seq_len, seq_len] weights = softmax(scores, mask) # 按行归一化 out = weights @ V # [seq_len, d_v] return out, weights # 造一组小输入走一遍 rng = np.random.default_rng(42) seq_len, d_model, d_k = 4, 8, 4 X = rng.normal(size=(seq_len, d_model)) W_Q = rng.normal(size=(d_model, d_k)) * 0.05 W_K = rng.normal(size=(d_model, d_k)) * 0.05 W_V = rng.normal(size=(d_model, d_k)) * 0.05 Q, K, V = X @ W_Q, X @ W_K, X @ W_V out, attn = scaled_dot_product_attention(Q, K, V) print("out shape:", out.shape) print("attn shape:", attn.shape) # 因果掩码示例 causal_mask = np.triu(np.full((seq_len, seq_len), -1e9), k=1) out_masked, attn_masked = scaled_dot_product_attention(Q, K, V, causal_mask) print("masked attn:\n", attn_masked)跑完这段代码,你会看到每个头要做的其实就这么点事。多头无非是把输入切成多个低维子空间,每个子空间重复上述流程,最后用输出投影拼回去。建议你再做一个小实验:把缩放因子去掉,观察注意力权重变化和训练波动,你会对这个公式记忆非常深刻。
回头说初始化,上面示例里的权重是随手乘的0.05,真实训练时Q/K/V矩阵的初始化方差需要和维度匹配,否则残差网络里容易出现方差漂移。很多人跑Transformer一开始训练不稳定,未必是学习率的问题,先检查一下各线性层的初始化是不是合理,往往能省很多调试时间。
我自己在实际使用中最大的体会是:把注意力从零实现一遍和直接调用框架完全是两个体验。拿到一个现成注意力实现时,建议先不要急着改代码,先把每个矩阵乘法前后的shape标出来,把Q、K、V、scores、weights、output这六样东西连成一幅自己的流程图;然后再去改缩放因子、加不同mask,亲眼看一下中间结果怎么变。这两件事做完,Transformer的大门才算真正推开。下一篇我会继续拆前馈网络和残差连接、层归一化之间的配合,那部分同样有不少细节坑等着。