AI Transformer架构详解:从自注意力到完整工作流程
Transformer 架构的核心是用自注意力机制让序列中每个位置直接获取全局信息,而不像循环神经网络那样逐步传递。
理解它的关键是跟踪一条数据从输入到输出经过了哪些变换、每一步张量形状如何变化。
下面按数据流向拆解,并给出可验证的代码片段。
先看整体:一个 Transformer 块里有什么
一个标准 Transformer 编码器块包含两个子层:多头自注意力层和前馈网络层。
每个子层外面都包着残差连接和层归一化。
解码器块多一个交叉注意力子层,用于让输出序列关注输入序列。
用列表概括单个编码器块的数据流:
- 输入嵌入加位置编码,得到形状
(batch, seq_len, d_model)的张量 - 多头自注意力计算,输出同形状张量
- 残差相加后做层归一化
- 前馈网络做两次线性变换加激活,输出同形状张量
- 再次残差相加和层归一化
关键结论:整个块不改变张量形状,只改变每个位置向量的内容。
形状不变是残差连接能成立的前提。
自注意力到底在算什么
自注意力把每个输入向量映射成三个向量:查询(Query)、键(Key)、值(Value)。
用查询和所有键做点积得到注意力分数,再经过 Softmax 归一化,最后用这些分数对值加权求和。
通俗理解:查询是“我在找什么”,键是“我有什么标签”,值是“我实际携带的信息”。
每个位置根据自己查询和所有位置键的匹配程度,决定从各个位置的值里取多少信息。
用 PyTorch 验证单头注意力的形状变化:
import torch
import torch.nn.functional as F
batch, seq_len, d_model = 2, 5, 16
x = torch.randn(batch, seq_len, d_model)
Wq = torch.nn.Linear(d_model, d_model)
Wk = torch.nn.Linear(d_model, d_model)
Wv = torch.nn.Linear(d_model, d_model)
Q, K, V = Wq(x), Wk(x), Wv(x)
scores = Q @ K.transpose(-2, -1) / (d_model ** 0.5)
attn = F.softmax(scores, dim=-1)
out = attn @ V
print(out.shape) # torch.Size([2, 5, 16])
运行结果应输出 torch.Size([2, 5, 16])。
如果形状不对,优先检查 transpose(-2, -1) 是否作用在键矩阵的最后两个维度上。
多头注意力只是把 d_model 拆成多个头并行计算,再把结果拼接回原维度。
比如 d_model=16、num_heads=4,每个头处理 4 维。注意:d_model 必须能被头数整除,否则拼接时会报维度错误。
位置编码和前馈网络各自解决什么问题
自注意力本身不区分词的顺序,交换两个词的位置,注意力分数可能只跟着交换而不体现先后。
位置编码把位置信息加到输入嵌入上,让模型能感知顺序。
常见做法是用正弦和余弦函数生成固定编码,也可以使用可学习的位置嵌入。
前馈网络是两层线性变换加一个非线性激活,对每个位置独立作用。
它的作用是增加模型容量和非线性表达能力,不负责位置之间的信息交换。
常见误解:以为前馈网络也在做词与词之间的交互。
实际上词间交互全部发生在注意力层,前馈层只处理单个位置向量。
残差连接和层归一化为什么不能省
残差连接把子层输入直接加到子层输出上,让梯度能更顺畅地传回浅层。
层归一化对每个样本的特征维度做归一化,稳定训练过程。
两者组合的顺序有 Pre-LN 和 Post-LN 两种常见变体,不同模型选择不同,排查训练不稳定时需要先确认实现用的是哪种顺序。
实操检查清单:
- 确认残差相加时两个张量形状完全一致
- 确认层归一化作用在特征维度而不是序列维度
- 如果训练出现梯度爆炸或损失不下降,先检查归一化位置和初始化方式
自己验证是否真的理解了架构
只看图容易产生“好像懂了”的错觉。
建议用上面的代码片段做三件事:
- 打印每个子层前后的张量形状,确认形状守恒
- 把
num_heads改成不能整除d_model的值,观察报错信息 - 交换输入序列中两个位置,对比注意力分数矩阵的变化
能解释清楚这三步的结果,说明对 AI Transformer 架构的数据流向已经有了可操作的认知,而不只是记住名词。
后续再去看 BERT、GPT 等具体模型时,差异点会更容易定位。