首页 > 教程攻略 > ai教程 >从零理解 Transformer

从零理解 Transformer

来源:互联网 时间:2026-06-30 07:22:06

从零理解 Transformer:注意力机制全解析

Transformer 架构彻底改写了自然语言处理的技术版图——从 BERT 到 GPT-4,从 T5 到 LLaMA,几乎所有现代大语言模型都长在 Transformer 的根上。但说实话,很多开发者的理解还停在“调 API”层面。本文从直觉出发,带你真正吃透注意力机制的核心原理,再通过代码实现把认知夯实。

从零理解 Transformer

为什么需要 Attention

传统的 RNN 和 LSTM 在处理序列时,有两个绕不开的硬伤。第一个是长期依赖——序列越走越长,早期信息在层层传递中逐渐被稀释,就像传话游戏一样,第一个人说的话传到第十个人耳朵里,已经不知道变成什么了。第二个是没法并行——每一步计算都依赖上一步的隐状态,训练速度慢得让人心焦,在 GPU 上只能串行计算,硬件性能完全发挥不出来。

Attention 机制用一个极其聪明的思路同时把这两个问题给解决了:让序列中的每个位置都能直接“看见”所有其他位置,计算它们之间的相关性权重。信息不用再一步步传递——全局视野,一步到位。这个设计理念简洁到近乎粗暴,但它的力量之强,整个 Transformer 的核心就这一句话。

Self-Attention 的直觉理解

在钻进公式之前,先建立直觉。想象你在读这么一句话:“昨天在图书馆看的那本书非常有趣,我推荐你也读一下。”当你的大脑处理“有趣”这个词时,你会很自然地联想起“书”——这就是注意力。你不需要按顺序把每个词都回忆一遍,而是直接跳到最相关的那个词。

Self-Attention 干的就是这件事:对于序列里的每个词,计算它跟其他所有词的相关性分数,然后用这些分数对所有词的信息做加权平均。这样一来,每个词的表示都融入了整个序列的上下文信息。

Self-Attention 的核心公式

Self-Attention 的数学表达简洁又优雅:

Attention(Q, K, V) = softmax(QK^T / √d_k) V

三个矩阵——Query(查询)、Key(键)、Value(值)——都来自同一个输入序列,通过不同的线性变换得到。直观理解:Q 代表“我在找什么”,K 代表“我有什么标签”,V 代表“我的实际内容是什么”。

计算过程分四步。第一步,算 Q 和 K 的点积,得到一个注意力分数矩阵——每一对词之间都有一个分数,表示它们应该互相关注多少。第二步,除以 √d_k——这是关键的缩放操作,d_k 是 Key 向量的维度,防止点积值过大导致 softmax 的梯度消失。第三步,对每一行做 softmax 归一化,把分数变成概率分布,每一行的和是 1。第四步,用这些概率权重对 V 做加权求和,得到每个位置的最终输出。

为什么除以 √d_k 这么要紧?假设 d_k = 64,Q 和 K 每个元素都是均值为 0、方差为 1 的随机变量。它们的点积是 64 个独立乘积的和,方差就变成了 64(也就是 d_k)。除以 √d_k 把方差拉回到 1,让 softmax 的输入待在合理范围内。不做这个缩放,点积值会太大,softmax 输出会偏向 one-hot——只有最大值的位置得 1,其他位置得 0,梯度几乎为零,模型根本学不动。

多头注意力的艺术

单头注意力只能捕获一种类型的关系——比如语法上的主谓搭配。但语言中的关系是多维的:语法结构、语义相似性、指代关系、长距离依赖、局部上下文,不一而足。多头注意力并行运行多个独立的注意力计算,每个“头”可以关注不同的模式。

具体实现上,假设想要 h 个头,每个头的维度是 d_k = d_model / h。输入 X 分别通过 h 组不同的 W_Q、W_K、W_V 矩阵投影到 h 个子空间,每个子空间独立计算 Attention,最后把所有头的输出沿特征维度拼接起来,再通过一个线性变换 W_O 融合。

为什么有效?因为每个头有自己的参数矩阵,不同的初始化会让它们学会关注不同的模式。一个头可能学会关注相邻词的关系,另一个头关注主谓搭配,还有一个头盯着所有代词和它们的指代对象。这种多角度同时处理的能力,正是 Transformer 超越 RNN 的关键所在。

位置编码的演进

Attention 机制本身对位置不敏感——“我爱你”和“你爱我”在 Self-Attention 眼里是完全一样的词向量组合,只是顺序不同。位置编码就是给每个位置注入独一无二的信号。

最初 Transformer 论文用的是正弦位置编码。对于位置 pos 和维度 i,正弦和余弦函数交替:

PE(pos, 2i) = sin(pos / 10000^(2i/d_model))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))

这个设计有好几个巧妙之处。每个位置都有唯一编码。不同位置的编码可以通过线性变换相互关联——这让模型学相对位置关系更容易。正弦函数的周期性保证编码值始终在 [-1, 1] 之间,训练稳定。而且它不用学习,理论上可以外推到训练时没见过的序列长度。

后来发展出了可学习的位置嵌入(Learned Positional Embedding)——直接把位置当做一个可训练的 Embedding 层。BERT 和 GPT 早期版本都这么用。优点是灵活,模型自己学最优的位置表示。缺点是无法外推到更长的序列。

RoPE(旋转位置编码)是近年来最重要的位置编码创新,已经被 LLaMA、Qwen、ChatGLM 等主流开源模型采用。核心思想是通过旋转变换把位置信息编进注意力计算里,让 Q 和 K 的点积自然地包含相对位置信息。RoPE 外推性好,数值稳定,已经成为大模型位置编码的事实标准。

残差连接和层归一化

Transformer 里的残差连接是另一个关键设计。每个子层(Self-Attention 和 Feed-Forward)的输出不是直接传给下一层,而是加上该子层的输入:LayerNorm(x + Sublayer(x))。

残差连接的作用是让梯度可以“抄近路”直接传播到前面的层。想象一个 100 层的网络,第 1 层的梯度得穿过 99 个激活函数和线性变换才能更新参数。经过这么多层,梯度要么消失(趋近于 0)要么爆炸(变得巨大)。残差连接提供了一条“高速公路”,让梯度可以无衰减地流动。

层归一化(Layer Normalization)则是稳定训练的另一个法宝。跟 Batch Normalization 沿 batch 维度归一化不同,Layer Normalization 沿特征维度归一化——对每个样本的所有特征计算均值和方差做标准化。这种设计更适合 NLP:序列长度可变,batch 里句子长度不一致时,Batch Norm 在 padding 位置的统计量没意义。而 Layer Norm 对每个位置独立计算,不受序列长度和 batch 大小的影响。

Feed-Forward Network:知识的存储器

每个 Transformer 层除了 Self-Attention,还有一个位置独立的前馈网络:

FFN(x) = ReLU(xW1 + b1)W2 + b2

在现代模型(比如 GPT)里更常用 GELU 或 SwiGLU 激活函数。

这个前馈网络看着简单,但作用大得很。有研究认为,Self-Attention 负责“检索”——从序列里找出相关信息并聚合起来;FFN 负责“存储”——把大量知识编码在巨大的参数矩阵里。FFN 的参数量通常占模型总参数的 2/3 以上——举个例子,一个 12 层的 Transformer,中间维度通常是 3072(d_model 是 768),W1 和 W2 每个都是 768×3072 的参数矩阵。

这也是为什么大模型能存下海量知识——这些知识并不存在于某个显式的数据库里,而是分布式地编码在每一层 FFN 的权重矩阵中。当你问 GPT “法国的首都是什么”,它不是在查表,而是这些知识早就写在了网络参数里。

从 Encoder 到 Decoder:不同的注意力掩码

Transformer 最初是为机器翻译设计的,包含一个 Encoder 和一个 Decoder。Encoder 用双向 Self-Attention——每个位置能看到所有其他位置(包括前面和后面的词)。这种设计适合理解任务,比如 BERT 只用 Encoder 来做分类、问答和序列标注。

Decoder 用单向(因果)Self-Attention——每个位置只能看它之前的位置(包括自己),后面的位置被 mask 掉。这是通过注意力掩码矩阵实现的:在 softmax 之前,把未来位置的分数设为负无穷(-∞),softmax 后这些位置的权重就成了 0。

GPT 系列只用 Decoder,因为它的目标是自回归生成——根据前面的词预测下一个词。训练时,整个序列一次性输入,但通过因果掩码保证每个位置的预测只看它之前的内容,不会“作弊”看到答案。推理时,模型逐个生成 token,每次生成的新 token 被拼接到序列末尾,作为下一次生成的新输入。

Encoder-Decoder 注意力(也叫 Cross-Attention)是原始 Transformer 的第三个注意力类型。Decoder 的 Q 来自自己的输出,但 K 和 V 来自 Encoder 的最终输出。这样 Decoder 在生成每个词时,既能看自己已经生成的内容(通过因果 Self-Attention),又能参考源语言的全部信息(通过 Cross-Attention)。T5 和 BART 等 seq2seq 模型保留了这种结构。

代码实现:从零构建 Self-Attention

理解的最好方式是动手实现。下面用 PyTorch 实现一个单头 Self-Attention 模块:

import torch
import torch.nn as nn
import math

class SelfAttention(nn.Module):
def __init__(self, d_model=512):
super().__init__()
self.d_model = d_model
self.W_Q = nn.Linear(d_model, d_model, bias=False)
self.W_K = nn.Linear(d_model, d_model, bias=False)
self.W_V = nn.Linear(d_model, d_model, bias=False)

def forward(self, x, mask=None):
# x: (batch, seq_len, d_model)
Q = self.W_Q(x) # (batch, seq_len, d_model)
K = self.W_K(x)
V = self.W_V(x)

# 计算注意力分数
scores = torch.matmul(Q, K.transpose(-2, -1)) # (batch, seq, seq)
scores = scores / math.sqrt(self.d_model) # 缩放

# 因果掩码(可选)
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)

# softmax 归一化
attn_weights = torch.softmax(scores, dim=-1)

# 加权求和
output = torch.matmul(attn_weights, V)

return output, attn_weights

把这段代码复制到 Jupyter Notebook 里跑一下,传入一个随机张量,观察注意力权重的分布——你会发现不同的头确实学到了不同的关注模式。

Transformer 的变体和未来

Transformer 诞生八年来,发展出了众多变体。效率方面,Reformer 用局部敏感哈希(LSH)把注意力复杂度从 O(n²) 降到 O(n log n);Linformer 证明注意力矩阵是低秩的,可以用线性投影近似;FlashAttention 从硬件角度优化,通过分块计算和重排内存访问模式,在 GPU 上实现了数倍加速。

架构方面,Mamba 和状态空间模型(SSM)正在挑战 Transformer 的统治地位。它们用线性时间的状态空间模型替代了二次复杂度的 Self-Attention,在处理超长序列时优势明显。不过到目前为止,Transformer 仍然是效果最好、生态最完善的架构。

写在最后

理解 Transformer 不是一蹴而就的事——这可能是深度学习领域最精妙的设计之一。建议从三个方面入手:第一,把 Self-Attention 的计算过程用手算一遍(取一个 3×4 的矩阵,一步步推演),建立肌肉记忆;第二,用 PyTorch 或纯 NumPy 实现一个微型 Transformer,喂入几个随机句子,观察每一层的输出变化;第三,读原始论文《Attention Is All You Need》——虽然发表多年,但论文里的每个设计选择都值得仔细品味。

当你真正理解了 Transformer,你就掌握了现代 AI 最重要的一块拼图。从 ChatGPT 到你手机上的输入法联想,从代码补全到蛋白质结构预测,Transformer 正在重塑我们与信息交互的方式。

深入 Attention 的数学本质

上一节了解了 Attention 的基本公式,但还有很多细节值得深入挖一挖。

Q、K、V 的由来

很多人第一次接触 Self-Attention 时最大的困惑是:Q、K、V 到底是什么?为什么叫这些名字?其实这些术语来自信息检索领域。在搜索引擎里,用户输入查询(Query),系统将查询与数据库中每个文档的键(Key)进行匹配,计算相似度分数,然后根据分数返回对应的值(Value)——也就是文档内容。

Self-Attention 借用了这个框架,但做了一个关键的创新:Q、K、V 都来自同一个输入,通过不同的权重矩阵做线性投影。这意味着每个词同时扮演三个角色——它在查询别人(Q),也在被别人查询(K),还携带着实际要传递的信息(V)。这种设计让每个词能在统一的向量空间里跟其他词交互。

点积注意力的优势

为什么用点积来衡量相似度,而不是其他距离度量?点积有几个优势:计算高效(GPU 上的矩阵乘法极度优化),可微分(可以端到端训练),而且两个向量方向一致时点积值最大——正好符合“相似度”的直觉。

但点积也有局限。两个向量的点积值受向量维度的影响——维度越高,点积的方差越大。这就是为什么公式里要除以 √d_k。相比之下,加性注意力(Additive Attention)使用一个小型前馈网络来计算相似度分数,虽然计算量大一些,但不受维度影响。原论文比较了两种方式,发现当 d_k 较小时两者效果相当,当 d_k 较大时点积注意力的优势很明显。

Softmax 的温度参数

严格来说,公式里的 softmax 可以带一个温度参数 T:

softmax(score_i / T)

T 控制概率分布的“尖锐”程度。T 越小,分布越尖锐(最大值更突出);T 越大,分布越平滑(趋近均匀分布)。Transformer 里 T = √d_k 是一个精心设计的默认值,在“充分关注重要位置”和“保留一定不确定性”之间取得了平衡。

Transformer 的训练技巧

训练 Transformer 不是件轻松的事,尤其是模型变深之后。下面几点是关键的实践经验。

学习率预热(Warmup)

Transformer 通常不会一上来就用全速学习率。原论文提出一个学习率调度策略:在前 warmup_steps 步中,学习率线性增加到一个峰值,然后按步数的平方反比逐步衰减。这个策略被称为“Noam 调度器”,以第一作者的名字命名。

为什么需要预热?训练开始时,所有参数都是随机初始化的,注意力权重接近均匀分布。如果一开始就用大学习率,模型可能陷入次优的注意力模式——比如过度关注某些位置而忽略其他位置。预热给了模型一个“探索期”,让它在学习率较小时建立合理的注意力分布,然后再加速收敛。

标签平滑(Label Smoothing)

在分类任务中,通常用 one-hot 标签和交叉熵损失。但 Transformer 论文用了标签平滑——把真实标签的概率从 1.0 “平滑”到 0.9,把剩下的 0.1 均匀分配给所有其他类别。

这么做的目的是防止模型变得过于自信。如果模型认为正确答案是 1.0(绝对确定),它会对错误的预测分配极大的惩罚,导致梯度非常陡峭。标签平滑让模型保持一点“谦逊”,训练时更稳定,泛化性能也更好。不过近年的大模型训练实践中,标签平滑用得少了,因为现代正则化技术(比如 dropout、数据增强)已经足够强。

Dropout 的位置

Transformer 在三个位置用了 Dropout:Embedding 层之后、每个子层的输出上(在残差连接加和之后、LayerNorm 之前)、以及注意力权重上。注意力权重上的 Dropout 尤其独特——它在 softmax 之后随机丢弃一些注意力连接,迫使模型不要过度依赖某几个特定位置的交互。这相当于在说“你不能总盯着这一个词看,也要看看其他词”。

梯度裁剪

Transformer 训练中偶尔会出现梯度爆炸——特别是在训练初期或遇到异常输入时。梯度裁剪很简单:设定一个阈值(比如 1.0),如果梯度的 L2 范数超过这个阈值,就把它缩放到等于阈值。这保证了每一步参数更新的幅度不会太大,防止训练崩溃。

Transformer 的推理优化

训练完成后,Transformer 的推理阶段也有不少优化技巧。

KV 缓存(KV Cache)

在自回归生成中,每次生成新 token 都需要重新计算整个序列的 Self-Attention。但仔细观察会发现:对于之前已经生成的 token,它们的 K 和 V 在每一步都是相同的——只是多加了一个新 token 而已。KV Cache 就是把每一步计算出的 K 和 V 缓存起来,下一步只需要计算新 token 的 Q、K、V,然后和缓存的 K、V 拼在一起做 Attention。这样每一步的复杂度从 O(n²) 降到 O(n),推理速度提升明显。

Beam Search 的困境

早期的文本生成广泛使用 Beam Search——每一步保留概率最高的 k 个候选序列。但研究发现,对于开放式生成任务(如对话、故事创作),Beam Search 往往产生更差的结果——文本变得重复、无聊、缺乏多样性。这是因为 Beam Search 追求概率最大化,但人类语言并不是总选概率最高的词。现代大模型通常用 Top-p(核采样)或 Top-k 采样,在概率较高的一批词里随机选择,平衡了质量和多样性。

量化与推理加速

大模型的参数量动辄数十亿甚至数千亿,完整的 FP32 推理需要巨大的显存和计算资源。量化技术把模型参数从 FP32 压缩到 INT8 甚至 INT4,大幅降低显存占用和计算延迟。常见的方案包括 GPTQ、AWQ 和 GGUF——后者被 Ollama 和 llama.cpp 广泛使用,让大模型可以在消费级硬件上运行。

从理论到工程:大模型时代的 Transformer

今天我们看到的 GPT-4、Claude、Gemini 等大模型,虽然核心架构仍然是 Transformer,但引入了大量工程创新。

混合精度训练

大模型训练使用 FP16 或 BF16 混合精度——大部分计算用半精度加速,关键操作用全精度保证数值稳定性。BF16 比 FP16 有更大的动态范围(与 FP32 相同的指数位),训练大模型时更稳定,已经成为主流选择。一个 175B 参数的模型用 FP32 需要 700GB 显存,用 FP16 只需要 350GB,再加上梯度累积和优化器状态,实际需要约 2800GB——所以大模型训练通常需要数百甚至数千张 GPU。

3D 并行策略

当模型大到单张 GPU 放不下时,需要并行策略。数据并行(Data Parallelism)是最基础的——每张 GPU 持有一份完整模型副本,处理不同的数据批次,然后在反向传播时同步梯度。ZeRO(Zero Redundancy Optimizer)优化了这个过程,把优化器状态、梯度、参数分片到不同 GPU 上,大幅减少冗余。

模型并行(Model Parallelism)把模型的不同层放在不同 GPU 上——GPU1 负责第 1-12 层,GPU2 负责第 13-24 层,数据按流水线方式传递。管道并行(Pipeline Parallelism)进一步把每个 batch 切分成多个 micro-batch,让不同 GPU 可以同时处理不同的 micro-batch,提高利用率。

张量并行(Tensor Parallelism)则是在层内拆分——把一个大的权重矩阵切分到多个 GPU 上,每个 GPU 计算一部分,然后合并结果。三种并行策略往往组合使用,Megatron-LM 和 DeepSpeed 是这方面的代表性框架。

RLHF 与对齐

基础的语言模型只学会了“预测下一个词”,但不能理解“什么样的回答是好的”。RLHF(基于人类反馈的强化学习)通过在预训练模型之上增加一个奖励模型(从人类偏好数据中学习),然后用 PPO 算法微调,使模型的输出更符合人类期望——更有用、更诚实、更无害。这个技术是 ChatGPT 成功的关键之一,也是大模型从“会说话”到“会聊天”的转折点。

总结与展望

Transformer 从 2017 年的一篇论文出发,在短短数年内重塑了整个 AI 领域。它的核心思想——用注意力机制替代循环,让序列中的每个元素都能直接访问所有其他元素——简单而深刻。

展望未来,几个方向值得关注。效率优化方面,FlashAttention 已经证明了硬件感知算法设计的巨大潜力,未来可能有更多针对特定硬件的 Transformer 变体。架构创新方面,状态空间模型(如 Mamba)正在缩小与 Transformer 的差距,Hybrid 架构(混合 Attention 和 SSM)可能是下一个突破口。多模态融合方面,Transformer 已经统一了文本、图像、音频、视频等多种模态,未来的基础模型将更加通用。

无论技术如何演进,理解 Transformer 中的每一个设计选择——为什么要缩放、为什么要多头、为什么要残差——都会让你在面对新技术时更有底气。因为所有的创新,都是站在这些基础设计之上的改进。