最近不少朋友说,看完 Transformer 的可视化概念讲解,对 Attention 公式能点头了,但一回到完整模型,整个人又懵了——数据到底是怎么从头流到尾的?“查询”“键”“值”这些东西又是怎么在编码器和解码器里各司其职的?这篇是 Transformer 可视化解说的第二部分,不再停留在概念层面,直接把模型拆开,跟着数据流走一遍。每一步我都尽量说清楚张量长什么样、形状怎么变、参数在哪里、训练和推理有什么区别。
内容适合刚弄懂自注意力、想继续往下捋完整原理的读者,也适合那些看论文有点吃力、想靠直觉抓住关键设计的开发者。如果你已经把 Attention 公式背得滚瓜烂熟,同样建议读一读解码器部分和资源掩码那些容易漏掉的实现细节,在真实训练和推理时它们会非常关键。
1. 先看全景:Transformer 的数据河流
很多讲 Transformer 的文章一上来就贴公式,新手很容易被吓退。换一种方式,我们把整个模型当成一条流水线,原材料是文本,半成品是向量,成品是概率分布。
整条流水线分成两段:编码器(Encoder)负责把输入文本“读”成一组特征表示;解码器(Decoder)负责根据这些特征和已经生成的内容,一个个“写”出目标文本。
编码器由若干完全相同的层堆叠而成,每层里面有两件大事:多头自注意力(Multi-Head Self-Attention)和前馈网络(Feed-Forward Network)。两件大事外面都裹着残差连接和层归一化,防止过深网络出现训练困难。
解码器结构上更复杂一点,每一层里有三件事:带掩码的多头自注意力、交叉注意力(Cross-Attention)、前馈网络。同样带残差和归一化。这里多出来的掩码自注意力,作用是让解码器在预测第 n 个词时,看不到第 n 个词之后的任何内容,这样模型才不会“作弊”。
数据流是这样的:英文原句先分词、转成 token id,查 Embedding 表变成向量,加上位置编码后进入编码器。经过多层堆叠,得到一组和输入位置一一对应的向量序列 Z。这组向量携带了整句话的语义和句法信息。解码器这边,目标句同样做嵌入和位置编码,先经过掩码自注意力,然后把前一层的输出作为“查询”,去关注编码器输出的 Z,提取与当前生成步骤相关的信息。最后过线性层加 Softmax,推算出下一个词的概率。
整体看,编码器的输出是解码器的“外部记忆”,解码器每一步生成都要去这个记忆里查资料。理解了这条主干管道,后续的细节都只是往里面注入具体的计算逻辑。
2. 第一步:从 token 到向量,再到带上位置信息
2.1 嵌入层:词表映射与向量化
Transformer 不会读字符串,只懂向量。先把句子切成 token,常见的是用字节对编码(BPE)分法,token 有点像“介于字符和词之间”的碎片,能处理未登录词,又不会把词表搞得太大。每个 token 对应一个唯一 ID,范围是词表大小 V。
接下来查嵌入表。嵌入矩阵的形状是V × d_model,d_model 是隐藏层维度,基础版 Transformer 用 512,超大模型可以用到 4096 甚至更高。查到之后,一个形状为[batch_size, sequence_length]的 token ID 矩阵,就变成了[batch_size, sequence_length, d_model]的浮点数张量。
从实操角度看,训练初期 Embedding 是随机初始化的,模型在训练过程中不断调整这些向量,让语义相近的词在向量空间里靠得近。一个值得注意的细节是:Embedding 的权重和 decoder 最后一层线性层经常是权重共享的,这样可以减少参数量,很多开源实现都默认这么做。
2.2 位置编码:为什么非加不可
词嵌入有个天然缺陷:它不区分位置。“张三让李四帮忙”和“李四让张三帮忙”,编码出的向量会因为词序变化产生不同结果吗?这里要澄清——自注意力计算时,如果没有任何位置信息,模型看到的是完全相同词在不同位置的同一种表示,根本无法区分词序。两个句子里的“张三”“李四”是不同 token,嵌入本身不同,模型倒不至于完全混乱,但重要的问题是“待在位置3的动词是否应该被位置5的词影响”。自注意力是对序列做全局加权求和,位置信息一旦丢失,纯注意力的输出就会变成对称结构:第 1 个词和第 5 个词互换位置,计算出的表示完全不变(如果它们自身嵌入相同且注意力模式同样对称)。
所以需要位置编码。原版用了一组正弦函数:
PE(pos, 2i) = sin(pos / 10000^(2i/d_model)) PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))这里的 pos 是 token 在序列里的位置,i 是向量维度下标。为什么用正弦波组合?直觉很简单:不同频率的正弦波叠加以后,任意一个位置的编码都是唯一的,而且相邻位置的编码在余弦距离上相近,这让模型更容易学到“相邻位置关系紧密”的规律。另一个好处是,如果训练时见过的最长序列是 512,理论上这个公式能外推到更长的序列,因为它是按公式算出来的,不需要查表。
实操中,位置编码的加法发生在输入嵌入之后:input_embedding + positional_encoding,直接把两组向量逐位相加。这里有个经常被忽略的细节:Embedding 后面的数值尺度比较小,如果不做处理,位置编码会直接把嵌入信息淹没。原版把嵌入权重乘以 sqrt(d_model),相当于把嵌入向量放大到和位置编码同一个量级,这个细节透露着早期实验者的手感。
也可以使用可学习的位置编码(每个位置学一个向量),效果在某些任务上也不错,但无法外推超长序列。对于长度固定、任务固定的场景,直接用可学习位置编码也是一种省事的选择,很多轻量模型就是这么干的。
3. 自注意力的四步计算:从矩阵角度拆开看
自注意力是所有后续计算的地基。我们把一个 token 的向量表示记作 x,它长度为 d_model。想得到注意力输出,需要生成三个新向量:查询向量 q、键向量 k、值向量 v。计算公式很简单,每个都来自 x 乘以一个不同的权重矩阵:
q = x @ W_Q k = x @ W_K v = x @ W_VW_Q、W_K、W_V 的形状都是d_model × d_k,这里 d_k 是每个头的维度。基础版里 d_model=512,多头数 h=8,那么 d_k = 512/8 = 64。也就是说,每个头只在 64 维子空间里做注意力,而不是整个 512 维,这样既降低了单次计算量,也让不同的头可以关注不同的子空间特征。
生成 q、k、v 之后,注意力分数的计算分三步:
- 计算 q 和所有 k 的点积,得到相似度分数。
- 除以 sqrt(d_k),把分数控制在稳定范围内。
- 做 Softmax 归一化成权重,然后对 v 做加权求和。
矩阵形式就是一个公式:
Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V有人会问:为什么要除以 sqrt(d_k) 而不是直接除 d_k?这个问题当初吃过亏的人不少。如果 d_k 很大,Q 和 K 的点积数值范围也会变大。因为每个维度上如果期望是 0、方差是 1,那么 d_k 个维度的和,方差大约会线性增长到 d_k,于是点积的标准差约为 sqrt(d_k)。除以它以后,方差回到 1 左右,Softmax 在高值域就不会把梯度压到几乎为零。梯度消失是训练的大敌,这一步很小,但价值非常大。
现在用一个微缩例子手动跑一遍。假设有两个 token,x1、x2,d_k=2,所以 q、k、v 都是二维。q=[0.5, -0.8],k1=[0.4, 0.6],q 和 k1 的点积是0.5*0.4 + (-0.8)*0.6 = 0.2 - 0.48 = -0.28,除以 sqrt(2)≈1.41 后变成-0.198。假设 q 与 k2 的点积算出来是 0.8,除以 1.41 后为 0.567。这两个分数输入 Softmax,得到权重约 [0.32, 0.68]。最后 v 就是这两个权重对 v1、v2 的加权和。这个结果会作为当前 token 的新表示,传给下一层。
整个过程中,Q、K、V 都来源于同一个输入 X,所以叫“自”注意力。不同 token 之间通过 QK^T 发生交互,Softmax 的行维是序列长度方向,每一行表示当前位置对其他位置的关注度。注意一个关键点:这里没有任何显式的循环结构,所有 token 的计算在数学上是并行进行的,这也是 Transformer 比 RNN 训练快得多的重要原因。
4. 多头注意力:让模型拥有多双眼睛
4.1 为什么是一个“头”不够
如果只用一组 Q、K、V,模型在一次注意力操作里只能学到一种“词与词之间关系”的模式。可是自然语言里关系是复杂的,有时要看句法依赖(主语-谓语),有时看指代关系(这个“它”指谁),有时看前后修饰(形容词修饰哪个名词)。
多头注意力就是给模型配了多套独立的 Q、K、V,每套负责看一类关系。模型在训练中自动分配注意力头,各有专长,这有点像团队里有人看全局、有人抓细节。实践里经常观察到一个头关注位置邻近的词,另一个头关注强语义关联的词。
多头计算的实现方式也不难。把输入 x 拷贝 h 份,分别乘以不同的 W_Q^i、W_K^i、W_V^i,生成长度为 d_k 的查询、键、值。然后每个头独立做一次缩放点积注意力,输出[batch, seq_len, d_k]。把所有头的输出在最后一个维度拼起来,得到[batch, seq_len, d_model],再经过一个输出权重矩阵 W_O,映射回d_model维度。
有一点容易困惑:每个头的次级空间是独立的,但它们共享同一个输入 x。所以不同头可以并行计算,Batch 里所有头拼接后还能继续走残差和归一化。
4.2 残差连接与 LayerNorm:让深层网络跑得稳
多头注意力输出后,先加回原来的输入,这就是残差连接。它保证了梯度可以从深层直接“走捷径”回到浅层,缓解梯度消失。有人这样类比:网络越深,梯度走的路越长,容易越来越弱;残差连接等于给梯度修了一条高架桥,车辆可以直接从终点一站返回起点。
然后是 LayerNorm。Transformer 里的归一化既不是 BatchNorm,也不是全局归一化,它对每个 token 的特征向量单独做归一化:在一个维度方向(d_model)上计算均值和方差,再减去均值除以标准差,最后乘上两个可学习的缩放和平移参数。这个选择有现实原因:NLP 里 batch 里每个序列长度不一样,BatchNorm 需要统计 batch 维度的分布,训练时非常不稳定;LayerNorm 对每个样本计算,既稳定又不依赖 batch 大小。换成大白话说,LayerNorm 让进入每一层的数据保持在一个可控的尺度内,训练更稳,对学习率的容忍度也更高。
一个完整的 Transformer 编码器层的内部顺序是:
- 多头注意力
- 残差连接
- LayerNorm
- 前馈网络
- 残差连接
- LayerNorm
这里有个实现上的经典纠结:是先算注意力再归一化(Post-Norm),还是先归一化再算注意力(Pre-Norm)。原版是 Post-Norm,也就是上面写的顺序。而很多现代实现(包括大家常用的开源大模型)用 Pre-Norm,先把输入归一化再过注意力。Pre-Norm 的好处是训练初期更稳定,尤其适用于极深的网络,但有一点轻微“人设”代价——最终表示没有经过归一化,输出尺度可能稍微不受控。实操中两者差别不会太大,但在超深网络上 Pre-Norm 往往更容易收敛,为了省心,很多项目首选它。
5. 前馈网络:一个不引入 token 交互的非线性变换
注意力部分本质上是对各个 token 的向量做加权求和。哪怕叠了很多层,如果不加入逐 token 的非线性变换,模型的能力就被限制在“线性组合”内。前馈网络(FFN)就是那个注入非线性、对每个位置做特征变换的组件。
FFN 的典型结构是两层全连接加一个非线性激活函数(原版用 ReLU):
FFN(x) = max(0, xW1 + b1)W2 + b2第一层把 d_model 维向量扩大到 d_ff 维,d_ff 通常是 d_model 的 4 倍,基础版里就是 2048 维。第二层再降回 512 维。这里的“前馈”意味着每个 token 独立经过同一个 MLP,token 之间不交互。你可以把它理解为:注意力负责“交流信息”,FFN 负责“独立消化信息”。
为什么要 4 倍那么大?一个解释是:注意力计算是矩阵乘法,信息主要在 token 之间融合;FFN 是逐 token 的,每个 token 需要充足的参数空间来记忆和转换自己在融入了上下文之后的新特征。参数规模一大,模型容量才够。比如 d_model=512、d_ff=2048 时,两层 FFN 的参数是512*2048 + 2048 + 2048*512 + 512,大概是 210 万参数,这在那个量级已经是很大的开销。所以 FFN 往往占了 Transformer 总参数量的一半以上。
另一个常见设计是把 ReLU 换成 GELU,现在的预训练模型几乎默认 GELU,因为它更平滑,梯度效果更好。如果有机会加一个 FFN 层的实验,你会发现绝大部分收益其实来自这个看似“无聊”的 MLP,而不是注意力头。这个反直觉的结论在不少消融实验里都被验证过——注意力头只剩两三个时模型还能凑合,去掉 FFN 直接崩盘。
6. 解码器:三块拼图和两种掩码
6.1 掩码自注意力:训练时的“防偷看机制”
解码器的第一件大事是掩码多头自注意力。注意,这里的输入是目标句子(训练时),在做“给定前面所有词,预测下一个词”这个任务。假设目标句子是“I love you”,模型要依次预测“love”“you”和结束符。如果在预测“love”时就把“you”的信息提前放进来,那模型等于直接看到了答案,学习就没有意义了。
掩码的具体做法是:在 QK^T 得到的注意力分数矩阵里,把矩阵右上角(所有位置 j > i 的地方)替换为负无穷大,这样 Softmax 计算时这些位置的权重会变成零。实际实现中常用一个上三角矩阵,填充一个很大的负数,比如 -1e9。
这个掩码有一个容易踩坑的地方:它会破坏序列的并行性吗?恰恰相反,掩码矩阵依然是一次性计算,只是把未来位置强制置零。所以训练时也能并行计算,不需要逐 token 生成。解码器并行地算出每个位置的上下文表示,然后对每个位置单独算损失。
6.2 交叉注意力:查询来自解码器,键值来自编码器
解码器的第二件大事是交叉注意力。这一步是编码器和解码器交互的唯一通道。它的查询 Q 来自解码器前一层的输出,键 K 和值 V 则来自编码器最后输出的序列 Z。用一句大白话说:解码器在生成某个词时,会拿着自己当前的“问题”,去编码器输出的整段上下文里寻找证据。
具体实现上,K 和 V 通常直接是编码器对输入句子的最终编码,Q 是解码器当前隐藏状态。比如在机器翻译里,解码器生成英文“love”时,会去中文原句“我爱你”里关注与“love”语义关系最紧密的那个词。所以交叉注意力的注意力权重,可以直接体现“对齐关系”的概念。
交叉注意力没有掩码,因为解码器完全可以看到编码器输出的全部内容。这里要和掩码自注意力区分开:掩码只限制解码器自己看目标序列,不限制它去查源序列。
6.3 推理时的自回归:和训练完全不同的走法
训练时解码器靠 teacher forcing,把真实目标句整体喂进去,加掩码防止偷看。推理时模型只能“一步一步来”:初始输入是起始符号<sos>,输出一个概率分布,取 argmax 得到第一个词,把这个词拼到输入末尾,再送给模型,得到第二个词,如此反复,直到输出<eos>。
这里有一个坑:推理时没有真实目标句,模型把“自己预测的上一步结果”当输入,一旦上一步错了,错误会不断累积。为了缓解,实际部署常用 beam search 保留多条候选路径,而不是纯贪心。还有人用温度调节采样增加多样性,不过做翻译任务基本都是 beam search 为主。如果你在做一个聊天机器人,训练和推理的差异会特别明显:训练时模型每步都能看到正确的历史,推理时一旦自己说错了话,后续输出往往更难救回来。多去训练时加入一些随机替换历史的 trick 能缓解这种 exposure bias,但这是个没有完美解的工程问题。
7. 训练与推理中的常见问题排查
7.1 损失不下降,先怀疑这几件事
训练 Transformer 经常遇到 loss 掉不下去的情况。先说基础检查:学习率是否过高或过低。Transformer 训练常用 warmup 策略,前几百步学习率从很小逐步升到预设值,避免早期步长太大导致震荡。如果你的 loss 在一个水平上横盘不动,先看学习率和 warmup。
其次是归一化实现是否选对了。Post-Norm 和 Pre-Norm 的收敛曲线差异在深度大时很明显。如果深度的网络陷入 loss 不降,尝试切换到 Pre-Norm 或调整残差初始化的比例。
还有一个经典问题:序列长度 padding 导致注意力权重不均匀。padding 的位置被补成 0,但在注意力中这些位置仍会和有效 token 计算相似度。需要显式传入 padding mask,把 padding 位置在 Softmax 前也置为负无穷。很多新手训练时漏了这个 mask,模型往往还能跑,只是效果上不去,十分迷惑人。
Softmax 温度也是一种“眼见为实”的调节手段。推理时常见的困惑是:模型给的分布太平,总是随机乱说;或者太尖,永远重复一句话。前者降低温度,后者升高温度。但我更建议先看看生成写的 beam size 和长度惩罚参数,很多问题根本不用动温度就能解决。
7.2 注意力可视化里到底看什么
做可视化分析时,很多人盯着注意力权重图看半天,却不知道什么算“正常”。我的建议是先看三个东西:对角线是不是普遍偏亮、attention head 之间有没有分化、交叉注意力的对齐是否合理。
对角线偏亮是普遍现象,这意味着大多数 token 强烈关注自己,本身不是问题,是注意力权重和方差的自然结果。重要的是对比不同层之间:浅层往往关注局部词序和邻近词,深层更关注语义关联和长距离依赖。如果所有层都只盯着自己,且多头没有分化(所有头长得一样),那通常是初始化有问题,或者训练充分度不够。不要轻易下结论说“这个注意力图很奇怪”,先问自己“有没有代表性”和“是不是单条样本噪声”。多取几十条样本、看平均模式,才能得出可靠的结论。
7.3 调试 Transformer 的几条实战建议
单靠肉眼看训练曲线很难定位问题。我的实战工具表大概是这样:
| 症状 | 优先检查方向 | 手法 |
|---|---|---|
| loss 不降或 diverging | 学习率 / warmup / 预归一化顺序 | 尝试降到原学习率 1/10,加 warmup |
| 效果差但 loss 低 | 数据问题、标签错位、padding mask 漏设 | 可视化一批样本,检查 mask 是否生效 |
| 推理重复严重 | beam size / 长度惩罚 / 温度 | 增大 beam size,加长度惩罚 |
| 多头同质化 | 初始化方差 / 模型容量 | 检查初始化,适当增大 d_model |
| 训练速度慢 | 序列长度过度 padding | 按批次动态 padding,或分桶排序 |
刚开始做模型实验时,最想分享的一条建议是:不要一上来就上超大模型,先用一个小数据集跑通整个流程,确认 loss 能下降、过拟合能出现。一旦小模型能过拟合到训练集接近 100%,再逐步扩大规模。如果小规模都学不进去,多半是代码或数据链路出了问题,加机器加参数都没用。
很多人第一次写 Transformer 训练代码,都会在 mask 上栽跟头,尤其是因果掩码中的负无穷值设置。如果值太小(比如 -1e9),在混合精度训练下可能会破坏数值稳定性,因为 float16 的有效范围有限。现代实现里有人用布尔掩码配合masked_fill操作,数值上更稳定,针对混合精度也更友好。我的经验是:能用布尔掩码就不要手工填充负无穷,除非你确认数值范围没问题。
另外,调整 batch size 和序列长度时,一定要同步检查显存占用。自注意力的复杂度是 O(n^2) 的,序列长度从 512 提到 1024,显存消耗会变成原来的四倍,不只是两倍。所以处理长文本时,与其死怼长度,不如先用 Sliding Window Attention 这类局部注意力方案做降采样。要记住:Transformer 的序列长度和 GPU 显存之间不是线性关系,是平方级关系。
8. 结尾
写到这里,Transformer 的原理已经从输入 embeddings 一路走到了解码器推理。个人经验是,如果能把“数据从输入到输出每一步张量形状的变化”完全理清,很多困惑自然会消失。为此,我会习惯性地一边看模型代码一边写出每个张量的形状标注,遇到不确定的地方就用一行打印看 actual shape,比盯着论文推公式快得多。
最后再分享一个小技巧:在分析一个新建模型时,先用一段特别短的序列(比如两个 token,长度 4 以内)跑一次前向,打印出每个中间张量的形状和方差范围。很多“模型不学习”的问题,在这个阶段就能发现,不用等训练半天之后才一脸茫然。动手把每个环节都验证一遍,比看懂任何讲解都更扎实。