Transformer架构详解：从自注意力到多头机制与位置编码

我最早读《Attention Is All You Need》那篇论文的时候心里想的是这就完了一个注意力机制加几个全连接层就把当时的机器翻译SOTA给刷新了说实话第一次看完架构图我是懵的真正让我开窍的是后面自己动手把每个张量的维度都算了一遍。这篇内容是我整理的Transformer学习笔记的第四模块——架构详解。它适合两类人一类是刚学完词向量和注意力机制、想从上往下把Transformer彻底吃透的初学者另一类是已经把BERT、GPT各种开源模型当工具用了很久却始终没空把底座拆开看一眼的工程师。我会先讲清楚Transformer到底解决了老架构的哪些痛点再逐个拆解自注意力、多头机制、位置编码、编码器解码器的设计差异最后聊聊它怎么从NLP一路扩张到图像和医学影像领域并附一版可以直接跑起来的简化实现代码。1. 序列建模的老问题为什么RNN和CNN不够用了想理解Transformer的设计不能只盯着一堆公式看。你得先回到2017年之前看看大家在处理序列数据时用的都是什么玩意儿又卡在了哪些地方。只有把旧方案的问题看清楚了才能明白Transformer每一个看似“绕弯子”的设计都是在治疗哪种病。1.1 RNN的两堵墙串行计算与远距离遗忘RNN循环神经网络的核心逻辑是“逐步扫描”我读一个词把它和上一个时刻的隐状态融合再输出给下一个时刻。这种时间步上的串行依赖直接带来了两个麻烦。第一个麻烦是效率。假设一句话有30个token那第30个token必须等前29个都算完才能开始。GPU并行计算的优势完全使不上劲。当时也有不少人用CNN做序列建模正是想着用卷积把句子一次性扫掉就是为了避开这个串行瓶颈。第二个麻烦更致命——长期依赖。RNN在数学上等价于一个不断乘以同一个转移矩阵的迭代过程如果这个矩阵的特征值小于1早期的信息就会被指数级地衰减掉。你读了50个词可能早就把第1个词的信息忘光了。LSTM和GRU用门控机制稍微缓解了这个问题本质上是在说“哪些东西要多记一会”但并没有从根上改变“信息要沿着时间步一路传递”的模式。梯度要穿过这么长一条时间上的路径才能回到起点要么爆炸要么消失。1.2 CNN的感受野天花板CNN做序列建模的思路是把文本当成一维向量用卷积核去扫。好处是卷积天然支持并行因为每个位置的卷积计算互不依赖。但问题是卷积核的感受野是有限的。你要让第i个位置看到第i50个位置的信息就得把卷积层堆得很深每层往外扩一点点。像WaveNet之类的模型甚至要靠膨胀卷积dilated convolution来增大感受野。本质上CNN在处理长距离依赖时是在用“深度”换“视野”带来的是网络层数加深、优化变难、堆层带来的计算开销。这就很尴尬并行是有了但全局建模能力丢了。1.3 Transformer的破局思路全序列直接互看Transformer的解题思路非常“暴力”我不沿着时间步一步步走而是把一句话里的所有token同时拿出来让每个token直接和序列里的所有其他token做信息交互。交互的方式就是注意力打分——不是我偷偷记下来前面的信息而是每次计算时都光明正大地回头去看整个序列。这一下子解决了两个问题。第一并行度拉满所有位置的注意力可以同时算只需要一次巨大的矩阵乘法。第二信息距离化为乌有任意两个位置之间都只有一次注意力计算的距离不管隔了多远我都能直接“看到你”。当然代价也随之而来。注意力机制是二次复杂度序列长度是n那两两配对就是n²。但至少在2017年那个时间点对于几万个token以内的文本序列来说这个代价是划算的、可以接受的。后面无数变体都在围绕“怎么把n²降下来”做文章但那是后话先把最基础的结构吃透再说。2. 自注意力机制拆解Q、K、V到底在算什么如果说Transformer是一栋楼自注意力就是它的承重墙。很多人第一次接触注意力时脑子里只有一张密密麻麻的热力图知道“模型在关注某几个词”但对Q、K、V这三兄弟到底在干什么完全没有手感。这一节我用最笨的办法把它彻底算一遍。2.1 用查资料打比方理解自注意力你可以把自注意力理解成一次“带着问题的查资料过程”。假设你正在写一篇行业分析报告桌上摊了一堆资料这些资料就是Token列表。写到某个段落时你脑子里会带着一个“问题”——Q。然后你去翻每份资料先看它的标题索引——K你觉得这份资料和我的问题匹配吗匹配程度就是一个权重。匹配之后你再把这份资料里的具体内容——V按权重取过来融合进你正在写的这段话里。在Transformer里每个token同时扮演三个角色它要向别人提问Q要回答别人的匹配度K还要提供自己的内容V。这就是“自”的含义——查询、索引、内容全来自同一批输入数据不需要外部知识。2.2 缩放点积注意力的完整计算流程把上面的直觉翻译成矩阵运算就四步。第一步把输入X分别乘上三个权重矩阵W_Q、W_K、W_V得到Q、K、V。X的形状是(batch, seq_len, d_model)乘完以后Q、K、V的形状不变仍是(batch, seq_len, d_model)。W_Q、W_K、W_V就是模型要学的参数它们负责决定“我应该从哪个角度提问”和“我应该从哪个角度展示内容”。第二步计算注意力分数矩阵Q乘以K的转置。具体来说[ \text{scores} Q \cdot K^T ]Q是(batch, seq_len, d_model)K^T是(batch, d_model, seq_len)乘完得到(batch, seq_len, seq_len)的矩阵。第i行第j列的元素就是“第i个token对第j个token的匹配程度”。第三步缩放加softmax。将scores除以(\sqrt{d_k})然后对每一行做softmax让每一行的权重之和变成1。理解起来就是把匹配分数变成概率分布。这一步也有个很直观的作用——让某些权重变得非常大、其他压到接近0这样模型在融合下游信息时可以有明确的选择偏好。第四步加权求和得到输出。用得到的注意力权重去加权V里的每一行[ \text{Attention}(Q,K,V) \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V ]输出形状还是(batch, seq_len, d_model)每个位置的向量都包含了来自全序列的信息权重则由相关性决定。2.3 为什么一定要除以根号d_k这个缩放是Transformer论文里特别不起眼、但实际特别关键的一个细节。不除以(\sqrt{d_k})会怎样答案是softmax会过早饱和。想象一下你有一堆200维的向量。如果随机初始化两个独立向量点积的结果方差大约是d_k的数量级——也就是说向量维度越高点积数值可能越大。当分数一大比如到几十甚至上百softmax里对应的(\exp(x))就会爆炸式增长导致注意力权重退化成接近one-hot几乎所有的概率都压给最大值那个位置其他位置的梯度变成0。梯度为0意味着模型学不动注意力彻底“死”掉。除以(\sqrt{d_k})之后分数的方差被拉回到1左右softmax的输入区间保持在合理范围梯度能顺畅传递。你可以理解为打分函数的分母就是一把尺子先把量纲统一再丢进softmax。这一行公式是训练稳定性的隐形功臣。3. 多头注意力与位置编码模型如何兼顾“多视角”和“顺序感”单层的自注意力能做的事情其实还很有限不然Transformer不需要堆那么多层。多头注意力是其中一个关键的设计。而Transformer和RNN还有个本质区别RNN天生是按顺序处理输入数据的所以“顺序”这件事是刻在骨子里的Transformer是同时把所有token摊开算的本质上它根本不知道谁先谁后。这两个问题——多视角和顺序感——就是这一节要解决的核心。3.1 多头注意力在做什么8个头各看各的模式所谓的“多头”就是把注意力计算拆成好几组并行运行。假设d_model是512你设置8个头那每个头分到的维度是64。每个头都各自拥有一套W_Q、W_K、W_V在64维的子空间里做一遍完整的注意力计算再把8个头的输出拼起来经过一个线性层压缩回512维。为什么这能起作用因为每个头可以在不同的子空间里捕捉不同类型的关系。有的头可能专门学习“语法依赖”动词和它的主语总是有高权重有的头可能专门学习“共指关系”代词和它所指代的名词还有的头专门捕捉局部邻近的信息或者跨句子的全局关联。你训练完一个真实的Transformer把8个头的注意力矩阵打印出来看会发现每个头的热力图长得都不一样这就是“多视角”的直观体现。多头还有个隐藏价值可以理解为一种鲁棒性。如果只有一套Q/K/V,那么注意力模式被限制在一个线性子空间里能表达的关系比较单一。8个头相当于8个平行的“专家”每个专家专攻一种关系模式最后汇总输出。没有多头Transformer的表达能力会显著缩水堆再深也堆不出那种效果。3.2 位置编码给模型补上“顺序感”位置编码解决的问题是自注意力本身是“对称”的——你把句子里任意两个词互换位置注意力计算的结果不会变词序信息完全丢失。这显然不行因为“我喜欢你不”和“你不喜欢我”完全是两个意思。Transformer的原始方案是加一个正弦余弦位置编码。具体形式是[ PE_{(pos, 2i)} \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right) ] [ PE_{(pos, 2i1)} \cos\left(\frac{pos}{10000^{2i/d_{model}}}\right) ]每个位置生成一个固定长度的向量直接用加法融合进token的embedding里。这样输入到后续注意力层的向量里就带上了“我在第几个位置”的信息。用正余弦的好处有二一是函数式生成可以外推到训练时没见过的更长序列二是不同频率的组合天然编码了相对位置关系——两个向量的内积可以反应位置差异。后来的实践证明这种固定编码的效果不错但存在上限。于是又出现了很多进阶方案比如可学习的位置编码直接把位置向量当作可训练参数BERT就是这么干的、相对位置编码、以及现在大模型里非常流行的旋转位置编码RoPE。3.3 从绝对位置到相对位置RoPE等进阶方案RoPE的思路很有意思它不显式给embedding加一串位置向量而是在计算Q和K的点积时把位置信息通过旋转矩阵“乘”进去。也就是说两个token的注意力分数天然由它们的相对距离决定距离越远衰减越明显距离近的分数更高。这更符合语言学的直觉离得近的词关系更密切。Swin Transformer和很多视觉模型则偏爱相对位置偏置给注意力分数加一个可学习的相对位置偏移矩阵。不同任务对位置信息的敏感度不一样文字讲究绝对的词序图像更多是相对的空间关系所以出现了这么多分化——这也是Transformer架构“换汤不换药”的典型例子核心还是一套注意力机制位置编码则像调味料一样随任务而变。4. 编码器里剩下的组件残差、归一化与FFN原始Transformer的编码器是一个可重复堆叠的Block。每个Block除了多头注意力还包括三个必要的配套组件残差连接、层归一化、以及一个前馈网络。外行看架构图可能觉得这是“凑数的边角料”但内行都知道没有这三样东西Transformer根本堆不起来。4.1 残差连接给梯度开了一条“高速公路”残差连接就是那条Add操作把多头注意力的输入和输出直接相加。也就是[ \text{Output} \text{LayerNorm}(x \text{Attention}(x)) ]为什么要做这个加法因为归一化也好、注意力加权也好都是对信息的“改造”。如果网络每一层都在原来的信息上做加法改造那随着层数加深原始信息会被稀释得所剩无几。残差连接的作用是保证每一层的输出永远包含“上一层的信息”加“这一层学到的增量”。在反向传播时这条直接相加的路径还有更大的价值——它相当于给梯度修了一条高速通道。不管网络有多深梯度都可以通过这条“短路”直接从最后一层传回第一层不用担心中间每一层的矩阵乘法把梯度连乘到消失。这也是为什么Transformer可以轻松堆到12层、24层甚至上百层而许多早期深度模型三层就开始训练不稳。4.2 层归一化为何选LN而不是BN归一化是为了解决“内部协变量偏移”的问题——每一层的输入分布会在训练中不断变化导致上层很难稳定学习。但Transformer选了层归一化LayerNorm而不是当时更流行的批归一化BatchNorm这件事有它的必然性。BN是在batch维度上做归一化也就是对同一个特征维度跨样本求均值和方差。但Transformer的输入是变长序列不同样本可能长度不一样。更关键的是在一个句子里词与词之间的均值方差本身就含有语义信息BN粗暴地在样本间拉齐会把这些信息打掉。另一个现实问题是训练时BN需要依赖一个足够大的batch统计量推理时又需要维护一份运行时统计这在序列任务里相当麻烦。LayerNorm则是单独对一个样本内部的所有特征维度算均值和方差。无论句子里有多少个token每个token的特征向量都会被拉到自己合理的分布上。这种归一化方式在变长序列和在线推理场景下都更稳定而且实现起来最简单。GPT、BERT、Llama这些模型无一例外地沿用了它只不过后来很多实现把归一化的位置从“先归一化再注意力”换成了“先注意力再归一化”Post-LN vs Pre-LN原因是在深度加深时Pre-LN的训练更稳。4.3 位置级前馈网络与完整Block数据流FFN是每个Block里的最后一环它其实就是一个只带一层隐藏层的MLP[ \text{FFN}(x) \text{ReLU}(xW_1 b_1)W_2 b_2 ]通常隐藏层维度是4倍d_model比如d_model是512那隐藏层就有2048个神经元。它独立地作用于序列中的每个token相当于对每个位置做一次非线性特征变换。自注意力负责“位置之间”的信息交换FFN负责“单个位置内部”的信息加工。两者交替出现组成了Transformer的核心循环。一个完整Block的数据流是输入x → 多头注意力 → 残差相加 → 层归一化 → FFN → 残差相加 → 层归一化 → 输出。原版架构里注意力之后先Add再Norm后来很多实现调换顺序但本质上都是这套“交换-加工-稳定”的循环。把几个这样的Block叠起来就是一个完整的编码器。你可以把它当成一个“越来越精炼的信息加工厂”每经过一层序列里的每个token都快变成一个带着全句上下文的浓缩语义向量。5. 解码器特殊在哪儿掩码、跨注意力与训练/推理差异原始Transformer的架构图里左侧是编码器右侧是解码器两者长得有点像但解码器里多了两个关键设计——掩码自注意力和编码器-解码器注意力。这两个设计决定了训练和推理时的行为完全不同。5.1 掩码自注意力挡住未来的token解码器的任务是一个词一个词地预测下一个词。假设你正在生成“今天天气真好”这句话已经写了“今天天气”接下来要预测“真”字。如果模型在预测的时候偷看了后面的内容那训练就没意义了——它只需要记住答案而不需要学会推理。掩码自注意力做的事情就是在计算注意力分数矩阵之后把矩阵右上角的所有位置都替换成负无穷。这样经过softmax之后那些“看向未来位置”的权重会变成0。第i个token只能看到位置1到i的信息看不到i之后的内容。这就是为什么它叫“因果掩码”也经常被叫做“因果注意力”。这个设计在机器翻译里尤其关键解码器生成目标语言的某个词时能看到此前已经生成的所有词但绝对不能提前看到还没生成的词。没有掩码训练阶段模型就能作弊训练再好一到推理就原形毕露。5.2 编码器-解码器注意力输入和输出的桥梁解码器里第二层注意力不叫自注意力了它叫Cross-Attention源语言的信息在这里第一次进入输出决策。它的Q来自解码器的上一层输出但K和V来自编码器的最终输出。也就是说在生成目标语言的每个词时模型都会回到源语言句子那里“查资料”问问源句子的哪些部分和当前这个词最相关。原始的机器翻译直觉就是这么实现的源句子被编码成一组语义向量解码器每走一步都在这里面检索最相关的信息再融合进自己的预测。5.3 训练时并行、推理时逐token生成这个地方如果搞不懂很多人会被训练和推理的速度差吓到。训练阶段解码器看到的是完整的目标句子只是掩码遮住了未来位置所有位置的预测都可以同时算出来——并行吃数据一次前向传播更新所有参数。推理阶段就完全不同了。模型手里只有“已经生成的词”输出当前预测的词然后把新词拼到输入序列后面再跑一次前向传播生成下一个词。这个过程严格串行每生成一个token都要重跑一遍整条解码链路。这也是为什么生成式大模型推理那么慢、那么吃显存的原因——不是模型参数大而是每个token都要重新算一遍且KV Cache只能缓存注意力里的K和V无法省掉所有的计算。6. Transformer的家族扩张从BERT/GPT到ViT、Swin与医学影像2017年 Transformer刚出来时大家还在用它做机器翻译。但没过多久研究者们发现编码器和解码器完全可以“各自独立发展”并且这套架构不挑食——文本能处理图像也能处理甚至时序信号也能处理。这一节梳理一下几大分支的关键改动。6.1 编码器与解码器的分道扬镳BERT和GPT是同年横空出世的但路径完全不同。BERT只用了Transformer的编码器训练方式是“完形填空”随机把一些token遮住让模型根据上下文猜被遮住的是什么。这种双向上下文建模特别适合理解类任务比如分类、抽取、相似度判断。GPT则只留解码器训练方式是“接龙游戏”给你上一堆词预测下一个词。这种单向建模天然适配生成类任务。从这里开始大家意识到Transformer不是一栋固定的房子而是可以拿不同楼层随意拼搭的框架。你是做理解就搭编码器你是做生成就搭解码器你要做翻译就两个都要。后来很多多模态模型也是在这个基础上给不同模态各配一个编码器再统一送进主干网络。6.2 ViT把图像切成16×16的PatchVision TransformerViT的关键改动只有一个前提图像怎么变成token序列。答案是“切块”。一张224×224的图片按16×16的patch大小去切可以得到196个块。每个块展平成向量经过一个线性投影变成embedding再和位置编码相加剩下就完全是一个标准的Transformer编码器了。这个改动在视觉领域是革命性的因为它彻底放弃了卷积。ViT让图像分类这种任务第一次可以和NLP共享同一套骨干架构。但它也有代价patch数量多的话序列长度就大自注意力的二次复杂度在图像上非常让人头痛——一张高分辨率图片动辄几千上万个patch原始注意力根本扛不住。6.3 Swin、轻量Transformer与MissFormer等医学分割变体Swin Transformer给出的解决方案是“移动窗口注意力”只在局部窗口内做自注意力然后通过窗口移动让不同窗口之间交换信息。这样一来复杂度从平方级降到了和token数量线性相关而且还能自然地构建出类似CNN那样的层次化特征金字塔——先低分辨率大窗口再高分辨率小窗口。这种方法在目标检测和分割任务上效果显著很多视觉通用模型后续采用类似思路。至于轻量Transformer核心思想是在保持架构骨架的情况下削减参数或计算量缩小模型宽度、减少层数、做蒸馏或者用线性注意力近似替代softmax注意力。轻量化的目的很现实——让Transformer能在手机端、嵌入式设备上跑起来医学影像分析就是典型的高价值场景。比如MissFormer这个面向2D医学图像分割的模型它保留了编码器-解码器架构但针对性改了特征融合方式用多个解码器分支去捕获不同尺度的特征在参数量不大的情况下拿到不错的分割精度。这类变体数不胜数但万变不离其宗架构的基础组件还是Q/K/V、多头注意力、残差、归一化改的只是“在哪些token上做注意力”“怎么组织层次”而已。7. 手写一个简化Transformer核心逻辑与代码演示看了这么多理论不动手写一遍等于白看。这一节我用PyTorch实现一个最小可跑的Transformer编码器。这段代码去掉所有工程化包装只保留最核心的逻辑适合逐行阅读。7.1 数据准备从词表到嵌入先准备一个最小的词表比如一句“hello world”拆成两个token加上特殊符号。实际项目中用的是BPE之类的分词器这里简化处理目的是把完整的数据流转起来。7.2 自注意力模块的PyTorch代码下面这个类是本文核心的自注意力实现。我特意把缩放、mask、多头的维度变换分开写方便对照公式。import torch import torch.nn as nn import torch.nn.functional as F class SelfAttention(nn.Module): def __init__(self, d_model, num_heads, maskFalse): super().__init__() assert d_model % num_heads 0 self.d_model d_model self.num_heads num_heads self.d_k d_model // num_heads self.mask mask self.w_q nn.Linear(d_model, d_model) self.w_k nn.Linear(d_model, d_model) self.w_v nn.Linear(d_model, d_model) self.w_out nn.Linear(d_model, d_model) def forward(self, x): batch, seq_len, _ x.shape # 1. 生成 Q/K/V Q self.w_q(x) # (batch, seq_len, d_model) K self.w_k(x) V self.w_v(x) # 2. 拆成多头形状变成 (batch, num_heads, seq_len, d_k) Q Q.view(batch, seq_len, self.num_heads, self.d_k).transpose(1, 2) K K.view(batch, seq_len, self.num_heads, self.d_k).transpose(1, 2) V V.view(batch, seq_len, self.num_heads, self.d_k).transpose(1, 2) # 3. 打分 缩放 scores Q K.transpose(-1, -2) / (self.d_k ** 0.5) # 4. 因果掩码挡住未来位置 if self.mask: mask torch.triu(torch.ones(seq_len, seq_len), diagonal1).bool() scores scores.masked_fill(mask.to(scores.device), float(-inf)) attn_weights F.softmax(scores, dim-1) # 5. 加权求和 out attn_weights V # (batch, num_heads, seq_len, d_k) # 6. 拼接多头输出 out out.transpose(1, 2).contiguous().view(batch, seq_len, self.d_model) return self.w_out(out)如果只想跑通“数据进数据出”的最小流程一个带注意力、FFN和残差归一化的完整编码器Block大概长这样class TransformerBlock(nn.Module): def __init__(self, d_model, num_heads, hidden_dim): super().__init__() self.attention SelfAttention(d_model, num_heads) self.norm1 nn.LayerNorm(d_model) self.ffn nn.Sequential( nn.Linear(d_model, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, d_model), ) self.norm2 nn.LayerNorm(d_model) def forward(self, x): x self.norm1(x self.attention(x)) x self.norm2(x self.ffn(x)) return x7.3 跑通后的几个心得我第一次照猫画虎写完这段代码踩了三个比较明显的坑都是那种“看理论根本发现不了、一跑就崩”的问题。第一个坑是维度转置。多头拆分后Q/K/V的形状是(batch, num_heads, seq_len, d_k)而不是(batch, num_heads, d_k, seq_len)。顺序搞反矩阵乘法的维度永远对不上。我当时的排查方法就是打印每一层的shape从第一个矩阵乘开始逐行对。第二个坑是因果掩码的方向。torch.triu生成的上三角矩阵默认把对角线右侧下标更大的位置置1。如果忘了diagonal1这个参数模型就连当前位置也看不见了。我当时跑出来的结果是每个序列位置输出都是一样的——因为没有任何信息能从当前位置进来。第三个坑是未分词的整句输入。如果你拿一个没有分词的句子直接喂给模型词表里会全是没见过的新词模型输出永远是一团随机数。这个问题不解决后面什么性能分析都是白搭。实际项目里一定要先跑通一个特别小的、词表完整的玩具数据确认前向传播和反向传播都没问题再换真实数据。结尾一点实在话一路把Transformer从“为什么需要它”拆到“怎么手写出来”我现在最大的感触是这架构的门槛真的不高但想完全吃透也不简单。难的不是某一步细节——自注意力的公式十分钟能看懂难的是把Q/K/V、多头、位置编码、残差、归一化、掩码这些零件组装成一个能稳定训练的系统每一步都是环环相扣的。你要是真想学好Transformer我建议别急着追着各种模型变体跑先自己写一个玩具版本跑通再去看ViT、Swin、GPT那些代码你会发现自己突然能看懂了——因为骨架是一样的变的只是肉。