Transformer 架构
Transformer(Vaswani et al., “Attention Is All You Need”, 2017)是深度学习三大架构族中最年轻的一员(与 CNN、RNN 并列,见 深度学习概览),也是当今大语言模型(LLM)的架构基础。从 GPT 到 LLaMA、从 BERT 到 Vision Transformer,几乎所有 2018 年后的里程碑模型都建立在 Transformer 之上。本页系统讲解自注意力、多头注意力、位置编码、三大架构变体(Encoder/Decoder/Encoder-Decoder)、因果掩码,以及 2024-2025 年的最新架构创新(Mamba/SSM、Flash Attention 2/3、Ring Attention 等)。注意力机制的前置讲解见 注意力机制,RNN 的时序建模见 RNN 循环神经网络。
- Transformer = 全员圆桌会议:所有词同时到场,每个人可以同时”看”到所有人,直接计算”谁和我最相关”。没有 RNN 接力赛的距离衰减问题(见 RNN 循环神经网络),还能完全并行计算——这是 Transformer 比 RNN 快、效果好的根本原因。
- 自注意力 = 查字典:每个词生成三个角色——Query(“我在找什么”)、Key(“我能提供什么”)、Value(“我携带的信息”)。然后拿自己的 Query 去和所有人的 Key 匹配,匹配度越高就越关注对方的 Value。
- 多头注意力 = 多个角度同时看:一个头可能关注语法关系(主语-谓语),另一个头关注语义关系(同义词),第三个头关注位置关系(相邻词)。多角度并行后合并,理解更全面。
为什么 Transformer 能取代 RNN?
Section titled “为什么 Transformer 能取代 RNN?”RNN 处理序列时必须逐 token 递进——第 3 个词必须等第 2 个词处理完才能开始,无法并行。更致命的是梯度消失/爆炸:当序列很长时,第一个词的信号传到第 100 个词时已经极其微弱——模型”记不住远处”。
Transformer 用自注意力一步到位:任意两个位置之间距离恒为 1(只需一次矩阵乘法)。这既解决了并行问题(整个序列同时计算),又解决了长距离依赖问题——但代价是 的计算和内存复杂度。
自注意力的数学推导
Section titled “自注意力的数学推导”给定输入序列 ( 个 token,每个 维),自注意力通过三个可学习的权重矩阵 生成 Query、Key、Value:
缩放点积注意力(Scaled Dot-Product Attention)的计算公式:
逐步拆解:
- : 的矩阵,第 个元素是第 个位置的 Query 与第 个位置的 Key 的点积——衡量”位置 对位置 有多关注”。
- :缩放——防止点积值过大导致 softmax 进入梯度饱和区。
- softmax:沿 (Key)方向归一化,将分数变成概率权重(每行和为 1)。
- :用权重对 Value 加权求和,得到每个位置的最终表示。
为什么要除以 ? 假设 和 的元素是均值 0、方差 1 的独立随机变量,则 的方差为 。当 较大时(如 64 或 128),点积值会很大,softmax 的输入进入饱和区——输出接近 one-hot,梯度几乎为零,训练停滞。除以 将方差缩回 1 附近,保证梯度健康。这就是名字中 “Scaled(缩放)” 的由来。
单个注意力头只能学习一种关注模式。多头注意力(Multi-Head Attention)并行运行 组 Q/K/V,让模型从不同子空间关注不同关系:
其中每个 ,。实践中 ,,。8 个头并行后拼接回 维,再做一次线性映射。
为什么要多头? 比如处理”这只猫没有追那只狗,因为它太胆小了”——一个头可能关注”它→猫”的指代消解关系,另一个头关注”追→猫/狗”的主谓关系,还有一个头关注”胆小→没有”的否定修饰关系。不同头自动学习到不同维度的语言结构,这比单一注意力强大得多。
自注意力本身没有顺序概念(打乱输入序列的 token 顺序,注意力权重的排列也相应置换,但每个 token 的输出内容不变——因为 Q·K^T 只看内容不看位置)。必须注入位置信息告诉模型”谁在前谁在后”。
主流位置编码方案
Section titled “主流位置编码方案”| 方案 | 来源 | 核心思想 | 外推能力 |
|---|---|---|---|
| 正弦/余弦 | 原论文 2017 | 用不同频率的 sin/cos 生成固定位置向量 | 可外推但效果一般 |
| 可学习位置编码 | BERT/GPT | 位置向量作为可训练参数 | 不可外推(训练多长推理就多长) |
| ALiBi | Press et al. 2022 | 根据距离对注意力分数线性衰减 | 天然外推 |
| RoPE(旋转位置编码) | Su et al. 2021 | 在 Q/K 上施加旋转矩阵编码相对位置 | 可外推,现代大模型主流 |
RoPE(Rotary Position Embedding) 的核心思想:将位置 的 Query 向量 视为复数,乘以 (旋转 角度)。两个位置 和 的 Query-Key 点积只依赖于相对位置 :
为什么 RoPE 成为现代大模型的主流选择? 因为它编码的是相对位置而非绝对位置,且天然支持外推——推理时可以使用比训练时更长的序列。LLaMA、Qwen、DeepSeek、Mistral 等几乎所有 2023 年后的开源大模型都采用 RoPE。
三大架构变体
Section titled “三大架构变体”Transformer 原始论文是 Encoder-Decoder 架构(用于翻译),但后续发展出三大分野:
| 架构 | 代表模型 | 核心组件 | 注意力方向 | 擅长任务 |
|---|---|---|---|---|
| Encoder-Only | BERT、RoBERTa | 双向自注意力 | 每个 token 看全部 | 文本理解:分类、NER、阅读理解 |
| Decoder-Only | GPT、LLaMA、Qwen | 因果自注意力(掩码) | 每个 token 只看过去 | 文本生成:对话、续写、代码补全 |
| Encoder-Decoder | T5、BART、原始 Transformer | Encoder + Cross-attention | 编码器双向 + 解码器因果 | 序列到序列:翻译、摘要 |
当前大语言模型的主流是 Decoder-Only(GPT/LLaMA/Qwen/DeepSeek),因为它可以无限自回归生成,且扩大规模后理解能力同样出色。Decoder-Only 的核心优势是训练效率:不需要 Encoder-Decoder 的交叉注意力,架构更简洁;且通过”下一个 token 预测”任务可以利用海量无标注文本做自监督训练。
Cross-Attention(交叉注意力)
Section titled “Cross-Attention(交叉注意力)”在 Encoder-Decoder 架构中,解码器的注意力分为两种:
- 自注意力:解码器内部,当前位置看之前的已生成内容(带因果掩码)
- Cross-attention(交叉注意力):解码器的 Query 来自解码器自身,Key 和 Value 来自编码器的输出——即解码器在每一步”回看”编码器的表示,决定关注源序列的哪些部分
这正是 Bahdanau 注意力思想的直接延续——只不过编码器从 LSTM 换成了 Transformer Encoder。Cross-attention 也是多模态模型(如 BLIP-2、LLaVA)中连接视觉编码器和语言模型的桥梁,详见视觉问答与图像描述。
因果掩码 Causal Mask
Section titled “因果掩码 Causal Mask”Decoder-Only 模型训练时是一次性输入整条序列(而非逐 token 生成),但生成时第 个位置不能看到 及以后的”未来”信息。解决方案是因果掩码:
掩码后的注意力分数矩阵(对角线以上设为 ,softmax 后变 0):
pos1 pos2 pos3pos1 0.6 -∞ -∞pos2 0.3 0.5 -∞pos3 0.1 0.2 0.7因果掩码让解码器在训练时也能像推理时一样”只看过去”——这就是 GPT 系列能高效训练的关键。没有因果掩码,解码器在训练时会”偷看答案”(未来 token),导致推理时性能骤降。
KV-Cache:推理加速的关键
Section titled “KV-Cache:推理加速的关键”在自回归生成(逐 token 生成)时,每生成一个新 token,之前 token 的 Key 和 Value 不变——可以缓存起来避免重复计算。这就是 KV-Cache:
- 第一次生成:计算所有已生成 token 的 K/V,缓存
- 每生成一个新 token:只计算新 token 的 Q(1 个向量),用缓存的 K/V 做注意力
- 注意力计算从 降为
KV-Cache 的代价:缓存大小正比于序列长度 × 模型层数 × 隐藏维度。对于 70B 模型、32K 上下文,KV-Cache 可能占用数十 GB 显存——这是长上下文推理的主要瓶颈。PagedAttention(vLLM)和 GQA / MQA(Grouped/Multi-Query Attention)是 2023-2024 年降低 KV-Cache 的关键技术。
GQA 与 MQA:降低 KV-Cache
Section titled “GQA 与 MQA:降低 KV-Cache”- MQA(Multi-Query Attention):所有注意力头共享同一组 K 和 V(只保留 Q 的多头),KV-Cache 降低 倍,但精度有损。
- GQA(Grouped-Query Attention):折中方案——将 个 Q 头分成 组,每组共享一对 K/V。 时退化为 MQA, 时退化为标准 MHA。LLaMA-2/3、Mistral 等现代模型默认使用 GQA。
LayerNorm 的位置:Pre-LN vs Post-LN
Section titled “LayerNorm 的位置:Pre-LN vs Post-LN”原始 Transformer 使用 Post-LN(残差连接后再做 LayerNorm)。后续研究发现 Pre-LN(先做 LayerNorm 再做注意力/FFN)训练更稳定——不需要 learning rate warmup,可以训练更深的模型。现代实现几乎都用 Pre-LN。
FFN:前馈网络
Section titled “FFN:前馈网络”每个位置经过注意力层后,还要过一个前馈网络(Feed-Forward Network, FFN)——通常是两层线性变换 + 非线性激活:
FFN 的隐藏维度通常是模型维度的 4 倍(如 时 FFN 维度为 2048)。不同模型用不同的激活函数:原始 Transformer 用 ReLU,GPT-2 用 GELU,LLaMA 用 SwiGLU(一种 GLU 变体,效果优于 ReLU/GELU)。
FFN 是模型的”知识库”:注意力层负责信息路由(“谁关注谁”),FFN 负责信息变换(“被关注的信息怎么处理”)。研究表明,LLM 的事实知识大量存储在 FFN 的权重矩阵中。
自注意力 Q/K/V 计算
Section titled “自注意力 Q/K/V 计算”2024-2025 年架构前沿
Section titled “2024-2025 年架构前沿”Flash Attention 2/3
Section titled “Flash Attention 2/3”Flash Attention(Tri Dao, 2022-2024)是 IO 感知的精确注意力加速实现——数学上完全等价于标准注意力,但通过优化 GPU 显存读写(SRAM ↔ HBM 的 IO 模式),在不改变结果的前提下将速度提升 2-4 倍、显存降低 5-10 倍。Flash Attention 2(2023)优化了并行度,Flash Attention 3(2024)针对 H100 GPU 的异步特性进一步优化。2024 年几乎所有大模型训练都使用 Flash Attention。
Ring Attention 与长上下文
Section titled “Ring Attention 与长上下文”2024 年的长上下文竞赛中,Ring Attention(Liu et al., 2023)允许将超长序列的注意力计算分布到多个 GPU 上——每个 GPU 只负责一部分序列,通过环形通信传递 KV 块。这使训练 100 万+ token 的上下文成为可能。
Mamba 与状态空间模型(SSM)
Section titled “Mamba 与状态空间模型(SSM)”Mamba(Gu & Dao, 2023)是基于状态空间模型(State Space Model, SSM)的序列建模方法,被一些研究者视为 Transformer 的潜在挑战者:
- 核心思想:用选择性状态空间模型替代自注意力,计算复杂度 (线性于序列长度),而 Transformer 是 。
- 优势:推理速度快(无 KV-Cache 膨胀问题),长序列训练高效。
- 局限:2024 年研究表明 Mamba 在召回任务(如” needle in a haystack”)上不如 Attention,且扩大规模后在语言建模质量上尚未超越 Transformer。
- 混合架构:Jamba(AI21, 2024)将 Mamba 层和 Attention 层混合,取长补短。
2025 年现状:Transformer 仍然是绝对主流,Mamba/SSM 作为补充架构在特定场景(超长序列、低延迟推理)有潜力。大多数”打败 Transformer”的声明在更大规模上无法复现。
线性注意力与高效 Transformer
Section titled “线性注意力与高效 Transformer”为解决 复杂度,大量工作探索线性注意力(Linear Attention)——用核函数近似 softmax,将复杂度降为 。代表方法:Performer(2020)、Linformer(2020)、RWKV(2023-2024)。RWKV 是一种结合了 RNN 并行训练优势和 Transformer 表达能力的”RNN-CNN-Transformer 混合”架构,在 2024 年获得了社区关注。
ViT 视觉 Transformer
Section titled “ViT 视觉 Transformer”ViT(Dosovitskiy et al., 2020)的洞察:图像也可以当序列处理。把图像切成固定大小的 patch(如 16×16),每个 patch 拉平后经过线性投影变成一个 token,加一个可学习的 [CLS] token 做分类,然后送入标准 Transformer Encoder。详见视觉Transformer ViT。
完整的多头注意力实现
Section titled “完整的多头注意力实现”从零实现缩放点积注意力 + 多头机制(不依赖 nn.MultiheadAttention):
import torchimport torch.nn as nnimport torch.nn.functional as Fimport math
class MultiHeadAttention(nn.Module): def __init__(self, d_model=512, num_heads=8): super().__init__() self.h = num_heads self.d_k = d_model // num_heads # 每个头的维度 self.W_q = nn.Linear(d_model, d_model) # Query 投影 self.W_k = nn.Linear(d_model, d_model) # Key 投影 self.W_v = nn.Linear(d_model, d_model) # Value 投影 self.W_o = nn.Linear(d_model, d_model) # 输出投影
def forward(self, x, mask=None): B, N, D = x.shape Q = self.W_q(x).view(B, N, self.h, self.d_k).transpose(1, 2) # (B,h,N,d_k) K = self.W_k(x).view(B, N, self.h, self.d_k).transpose(1, 2) V = self.W_v(x).view(B, N, self.h, self.d_k).transpose(1, 2) scores = Q @ K.transpose(-2, -1) / math.sqrt(self.d_k) # 缩放点积 if mask is not None: scores = scores.masked_fill(mask == 0, float('-inf')) # 因果掩码 attn = F.softmax(scores, dim=-1) # 注意力权重 out = (attn @ V).transpose(1, 2).contiguous().view(B, N, D) return self.W_o(out) # 线性映射回原始维度
mha = MultiHeadAttention(d_model=512, num_heads=8)x = torch.randn(2, 10, 512) # batch=2, seq_len=10, d_model=512print(mha(x).shape) # 输出: torch.Size([2, 10, 512])带因果掩码的 Decoder
Section titled “带因果掩码的 Decoder”def create_causal_mask(seq_len): """创建因果掩码:上三角为 0(屏蔽未来),下三角和对角线为 1""" mask = torch.tril(torch.ones(seq_len, seq_len)) # 下三角矩阵 return mask.unsqueeze(0).unsqueeze(0) # (1, 1, N, N) 广播到 (B, h, N, N)
seq_len = 5mask = create_causal_mask(seq_len)print("因果掩码 (5×5):")print(mask[0, 0])# tensor([[1., 0., 0., 0., 0.],# [1., 1., 0., 0., 0.],# [1., 1., 1., 0., 0.],# [1., 1., 1., 1., 0.],# [1., 1., 1., 1., 1.]])
# 用在 Decoder-Only 模型中decoder = MultiHeadAttention(d_model=64, num_heads=4)x = torch.randn(1, seq_len, 64)out = decoder(x, mask=mask) # 第 t 个位置只看到 1..t使用 nn.TransformerEncoder 做分类
Section titled “使用 nn.TransformerEncoder 做分类”import torchimport torch.nn as nn
class TransformerClassifier(nn.Module): def __init__(self, vocab_size, d_model=64, nhead=4, num_classes=2): super().__init__() self.embedding = nn.Embedding(vocab_size, d_model) self.pos = nn.Parameter(torch.randn(1, 512, d_model) * 0.01) # 可学习位置编码 self.encoder = nn.TransformerEncoderLayer( d_model=d_model, nhead=nhead, dim_feedforward=128, batch_first=True) self.fc = nn.Linear(d_model, num_classes)
def forward(self, x): emb = self.embedding(x) + self.pos[:, :x.size(1)] # 词向量 + 位置编码 out = self.encoder(emb) # 自注意力编码整条序列 return self.fc(out.mean(dim=1)) # 平均池化 → 分类
model = TransformerClassifier(vocab_size=10000)print(model)- 位置编码不可省略:自注意力本身没有顺序感知,不加位置编码的 Transformer 对输入顺序完全不敏感。选择位置编码时,如需支持长序列外推,优先考虑 RoPE。
- Transformer 计算量是 :序列长度翻倍,注意力矩阵计算量变 4 倍,显存也变 4 倍。长文本(> 4096 token)需要特殊优化(Flash Attention、滑动窗口注意力、稀疏注意力等)。
- Decoder-Only 是当前主流:GPT/LLaMA/Qwen/DeepSeek 都是纯 Decoder 架构。Encoder-Only(BERT 系列)适合理解任务但不适合生成,Encoder-Decoder(T5)在翻译摘要上仍有优势。
- 因果掩码只在训练/生成时需要:Decoder 训练时用掩码保证”不看未来”,推理时因为逐 token 生成天然不会看到未来。Encoder 则不需要掩码——双向看全部上下文。
- ViT 需要大数据:ViT 在 ImageNet 级别数据上不如 CNN,但在 JFT-300M 等超大数据集上才能超过 CNN——它缺少 CNN 的归纳偏置(局部性、平移不变性),需要更多数据来学习这些。
- 使用 Flash Attention:2024 年训练任何 Transformer 模型都应使用 Flash Attention——它不改变数学结果,但大幅提升速度和降低显存。PyTorch 2.0+ 的
F.scaled_dot_product_attention内置了 Flash Attention。 - GQA 降低推理成本:如果做推理部署,使用 GQA(如 LLaMA-3 的 8 组 GQA)可以在几乎不损失质量的情况下大幅减少 KV-Cache 显存占用。
- 长上下文模型的”二次平方化”:2024 年开源模型的标准上下文从 4K 增长到 128K 甚至 1M(Qwen2、Llama-3.1)。关键技术:RoPE 基频调整(YaRN/NTK-aware)、Ring Attention、稀疏注意力。
- 大语言模型(LLM):GPT、LLaMA、Qwen、DeepSeek 全部基于 Transformer Decoder-Only 架构。详见 语言模型演进。
- 机器翻译:Google 翻译和 DeepL 用 Transformer Encoder-Decoder 实现上百种语言的高质量互译,全面替代了早期的统计翻译和 LSTM seq2seq。
- 文本理解:BERT 系列用于搜索排序(Google Search)、情感分析、命名实体识别等,是 NLP 理解任务的标准工具。
- 代码补全:GitHub Copilot 基于 GPT(Codex),用 Transformer 预测下一行代码,已成为开发者日常编程的标配辅助工具。
- 蛋白质结构预测:Google DeepMind 的 AlphaFold 2 用 Transformer 替代传统进化分析,将蛋白质三维结构预测精度提升到实验级。
- Whisper 语音识别:OpenAI 的 Whisper 以 Encoder-Decoder Transformer 实现多语言语音识别和翻译,开源后被广泛集成到字幕自动生成、会议转写等产品中。详见 Whisper 与 wav2vec2。
- 视觉理解:ViT 及其变体(Swin Transformer、DETR)应用于图像分类、目标检测、图像分割——Transformer 已成为视觉领域与 CNN 并列的主流架构。详见 视觉Transformer ViT。
典型类库与工具
Section titled “典型类库与工具”| 类库 | 语言 | 说明 |
|---|---|---|
| PyTorch | Python | nn.TransformerEncoderLayer / nn.MultiheadAttention / F.scaled_dot_product_attention 内置 |
| HuggingFace Transformers | Python | 提供 GPT/BERT/T5 等海量预训练 Transformer 模型与统一 API |
| TensorFlow / Keras | Python | tf.keras.layers.MultiHeadAttention 可快速构建 Transformer |
| timm | Python | 视觉 Transformer(ViT/Swin)的预训练模型集合 |
| Flash Attention | Python/CUDA | IO 感知的精确注意力加速实现,大模型训练标配 |
| xformers | Python | Meta 开源的高效 Transformer 组件库,支持记忆高效注意力 |
| vLLM | Python | 高效 LLM 推理引擎,PagedAttention 优化 KV-Cache 管理 |
| JAX / Flax | Python | Google 出品的高性能数值计算框架,其函数式设计非常适合大模型并行训练,是 TPU 生态的首选 |
| Tensor2Tensor | Python | Google 配合”Attention Is All You Need”论文开源的原始 Transformer 实现,具有历史参考价值 |
| 术语 | 英文 | 解释 |
|---|---|---|
| 自注意力 | Self-Attention | 序列中每个位置直接计算与所有其他位置相关性的机制,Transformer 核心 |
| 缩放点积注意力 | Scaled Dot-Product Attention | Q·K^T 除以根号 d_k 再 softmax 后加权 V 的注意力计算方式 |
| 多头注意力 | Multi-Head Attention | 并行运行多组注意力再拼接,让模型从不同子空间关注不同关系 |
| 位置编码 | Positional Encoding | 注入到输入中的序号信息,弥补自注意力本身没有顺序感知的缺陷 |
| 旋转位置编码 | RoPE (Rotary Position Embedding) | 在 Q/K 上施加旋转矩阵编码相对位置,支持外推,现代大模型主流 |
| 因果掩码 | Causal Mask | Decoder 中遮蔽未来位置的掩码,确保生成时不偷看后续 token |
| 交叉注意力 | Cross-Attention | Decoder 的 Q 来自自身、K/V 来自 Encoder 输出的注意力机制 |
| 编码器 | Encoder | Transformer 双向自注意力模块,适合理解任务(BERT) |
| 解码器 | Decoder | 带因果掩码的自注意力模块,适合生成任务(GPT) |
| 残差连接 | Residual Connection | 跳跃连接将输入直接加到输出上,缓解深层网络的梯度退化 |
| Q / K / V | Query / Key / Value | 注意力的三种角色:Q 是”我在找什么”,K 是”我有什么”,V 是”我携带的实际信息” |
| KV 缓存 | KV-Cache | 推理时缓存已生成 token 的 Key / Value,避免每生成一个新词就重算全部历史,大幅加速自回归生成 |
| Flash Attention | Flash Attention | 通过优化 GPU 显存读写顺序来加速注意力计算的技术,在不改变数学结果的前提下显著降低显存占用 |
| 分组查询注意力 | GQA (Grouped-Query Attention) | 多组 Query 共享一组 K/V,在精度和效率之间折中,降低 KV-Cache |
| 前馈网络 | Feed-Forward Network (FFN) | Transformer 中逐位置应用的两层 MLP,是模型”知识存储”的主要位置 |
| 层归一化 | Layer Normalization | 对每个样本的特征维度做归一化,稳定深层 Transformer 训练 |
| 状态空间模型 | State Space Model (SSM) | Mamba 等替代注意力的序列建模方法,复杂度 O(n) |
| Ring Attention | Ring Attention | 将超长序列的注意力计算分布到多 GPU 的并行策略 |
- Transformer 谱系:原始 Transformer(Vaswani et al., Google, NeurIPS 2017)源于 NLP 机器翻译 → BERT(2018,Encoder-Only,双向理解)→ GPT 系列(2018–至今,Decoder-Only,自回归生成)→ T5(2019,Encoder-Decoder,统一为 text-to-text)。视觉分支:ViT(2020,图像切 patch 做 token)、Swin Transformer(2021,层级化注意力)、DETR(2020,端到端检测)。
- 注意力机制的起源:自注意力源于 Bahdanau 注意力(2014),最初是 RNN seq2seq 的辅助组件,详见 RNN 循环神经网络 和 注意力机制。
- 位置编码前沿:RoPE(Su et al. 2021,旋转位置编码)、ALiBi(Press et al. 2022,线性偏置注意力外推)、YaRN(Peng et al. 2023,RoPE 外推)。
- 高效注意力:Flash Attention(Tri Dao, 2022-2024,IO 感知加速)、Ring Attention(Liu et al. 2023,超长上下文分布式训练)、GQA(Ainslie et al. 2023,分组查询注意力)。
- 替代架构:Mamba(Gu & Dao 2023,选择性状态空间模型)、RWKV(Peng 2023,线性注意力 RNN 混合)、Jamba(AI21 2024,Mamba+Attention 混合)。详见 Mamba与状态空间模型。
- 与大语言模型的关系:Transformer 是 LLM 的架构基础。GPT 本质上是”Transformer Decoder + 自回归生成”,详见 语言模型演进、自回归模型 和 提示工程。
- 现状与挑战:Transformer 是事实上的 AI 统一架构。主要局限 复杂度催生 Mamba 等状态空间模型(SSM)挑战者,但尚未撼动主导地位。