Transformer 架构详解：Self-Attention 机制与 PyTorch 实现

本文最后更新于 32 天前，如有失效请评论区留言。

说到 deep learning 里面最重要的 architecture，Transformer 绝对是排在前列的。2017 年 Google 发了那篇 “Attention is All You Need” 之后，基本上整个 NLP 领域就被 Transformer 统治了，后来又扩展到 CV、语音、甚至蛋白质结构预测。今天就来聊聊 Transformer 的核心机制，顺便用 PyTorch 写个简化版的实现。

为什么需要 Transformer？

在 Transformer 之前，处理序列数据主要靠 RNN（Recurrent Neural Network）和它的变体 LSTM、GRU。这些模型有一个根本性的问题：它们是按顺序处理的，t 时刻的输出依赖于 t-1 时刻的 hidden state。

这就带来两个麻烦：

无法并行：必须一步一步算，序列长了训练就特别慢
长距离依赖问题：虽然 LSTM 通过门机制缓解了一些，但序列超过一定长度，前面的信息还是会”忘掉”

简单说就是，RNN 像是读文章时一个字一个字往后看，看到后面的时候前面某些细节已经记不清了。而 Transformer 像是把整篇文章摊开，每个字都能直接”看到”其他所有字。

Self-Attention：Transformer 的灵魂

Transformer 最核心的概念就是 Self-Attention（自注意力机制）。它的想法其实很直觉：对于序列中的每个位置，计算它和所有其他位置的”相关性”，然后用这个相关性做加权求和，得到新的 representation。

具体来说，每个 token 会生成三个向量：

Query (Q)：我在找什么信息？
Key (K)：我有什么信息可以提供？
Value (V)：我实际携带的信息是什么？

Attention 的计算公式：

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

那个 √d_k 是 scaling factor，防止点积太大导致 softmax 输出变成 one-hot。这个小 trick 虽然简单但很重要，不加的话训练会很不稳定。

Multi-Head Attention

一个 Attention head 只能关注一种”模式”的信息。比如在处理 “The cat sat on the mat” 这句话时，一个 head 可能关注语法关系（cat-sat），另一个可能关注位置关系（sat-mat）。所以 Transformer 用了 Multi-Head Attention，把 Q、K、V 分成多个 head 各自计算，最后 concat 起来。

用 PyTorch 实现大概是这样：

注意：实际使用中建议用 nn.MultiheadAttention，下面的代码是为了理解原理手动实现的。

import torch import torch.nn as nn import math

class MultiHeadAttention(nn.Module): def __init__(self, d_model, n_heads): super().__init__() assert d_model % n_heads == 0 self.d_k = d_model // n_heads self.n_heads = n_heads 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_o = nn.Linear(d_model, d_model) def forward(self, Q, K, V, mask=None): batch_size = Q.size(0) # Linear projection and reshape to (batch, heads, seq_len, d_k) Q = self.W_q(Q).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) K = self.W_k(K).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) V = self.W_v(V).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) # Scaled dot-product attention scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: scores = scores.masked_fill(mask == 0, -1e9) attn = torch.softmax(scores, dim=-1) output = torch.matmul(attn, V) # Concat heads and project output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.n_heads * self.d_k) return self.W_o(output)

Position Encoding

Self-Attention 有个特点：它是 permutation invariant 的，也就是说打乱输入顺序，输出的 attention score 完全一样。这显然不对，”狗咬人”和”人咬狗”意思完全不同。

所以 Transformer 需要在输入上加上 Position Encoding。原始论文用的是 sinusoidal 函数：

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

说实话一开始我觉得这个设计挺奇怪的，为什么不直接学一个 position embedding？后来发现确实可以用 learned position embedding（BERT 就是这么做的），效果差不多。但 sinusoidal 的好处是可以 extrapolate 到训练时没见过的序列长度。

Transformer 的整体结构

一个完整的 Transformer 包含 Encoder 和 Decoder 两部分：

Encoder：N 层堆叠，每层包含 Multi-Head Self-Attention + Feed Forward Network，加上 residual connection 和 layer normalization
Decoder：结构类似，但多了一个 cross-attention 层来 attend encoder 的输出，同时用 causal mask 防止看到未来的 token

不过后来的发展把这两个拆开了：

Encoder-only：比如 BERT，适合理解类任务（分类、NER）
Decoder-only：比如 GPT 系列，适合生成类任务
Encoder-Decoder：原始 Transformer 和 T5，适合 seq2seq 任务（翻译、摘要）

现在最火的 LLM 基本都是 decoder-only 架构，虽然结构简单但 scaling 效果特别好。

Feed Forward Network 和 Layer Norm

Transformer 每一层除了 attention，还有一个简单的两层 FFN：

FFN(x) = max(0, xW₁ + b₁)W₂ + b₂

就是两个 linear 层中间夹一个 ReLU。第一层通常把维度扩大 4 倍（比如 512 → 2048），第二层再缩回去。这个结构虽然简单但参数量占了整个模型的一半左右。

Layer Normalization 在每个 sub-layer 的输出上做 normalize，配合 residual connection 确保梯度能正常流动。这里有个细节：原始论文用的是 Post-LN（layer norm 在 residual 之后），后来的研究发现 Pre-LN（layer norm 在 sub-layer 之前）训练更稳定，现在大多数实现都用 Pre-LN 了。

一些实际使用的心得

在实际项目里用 Transformer，有几个点我觉得比较重要：

Learning rate schedule：原始论文用了 warmup + inverse sqrt decay，这个对训练稳定性很重要。现在一般用 cosine schedule + warmup
Gradient clipping：Transformer 训练容易出现 loss spike，clip gradient norm 到 1.0 左右能缓解
参数量 vs 数据量：Transformer 的 scaling law 说明了模型越大、数据越多、效果越好，但这意味着小数据集上大模型容易 overfit
Mixed precision training：用 fp16/bf16 训练能省很多显存，训练速度也更快

我个人的经验是，如果你刚开始学 Transformer，先不要急着去训大模型。把 Multi-Head Attention 和 positional encoding 的实现搞清楚，比什么都重要。很多看似复杂的 trick（Flash Attention、MQA、RoPE 等）都是在这个基础上的优化。

总结

Transformer 的核心思想就是 Self-Attention：让序列中的每个位置都能直接和其他位置交互，解决了 RNN 的并行化和长距离依赖问题。Multi-Head Attention 让模型能同时关注不同类型的信息，Position Encoding 加入了位置信息，FFN 和 Layer Norm 则提供了非线性变换和训练稳定性。

理解了这些基础组件，后面再看 BERT、GPT、ViT 这些模型就会觉得顺理成章了。好了，这篇文章就到这里，希望能对你理解 Transformer 有帮助。