从零手写Transformer：注意力机制与多头设计深度拆解

1. 从零理解Transformer：为什么它值得你花时间

如果你最近在关注AI领域的动态，大概率已经被“Transformer”这个词反复轰炸过了。不管是大语言模型、图像生成、语音识别，还是推荐系统，背后几乎都有它的影子。但很多刚接触的朋友会有一个困惑：这东西到底是怎么回事？为什么它能把之前的RNN、LSTM按在地上摩擦？我刚开始学的时候也是一头雾水，论文翻了好几遍才慢慢摸到门道。

这篇文章想做的事情很简单：把Transformer和注意力机制的核心逻辑掰开揉碎，用从业者之间聊天的口吻讲清楚。我不会堆砌公式吓人，也不会只停留在“它很厉害”这种空话上。我会从设计动机讲起，一步步拆解Self-Attention到底在算什么，Multi-Head Attention为什么要“多头”，位置编码为什么不能用简单的数字递增，以及在实际写代码时哪些坑最容易踩。适合有一定深度学习基础、想真正搞懂Transformer而不是只会调包的朋友。读完你至少能做到：自己手写一个单层Transformer的forward过程，知道每个矩阵乘法的维度怎么来的，遇到shape不匹配的报错能自己排查。

先说结论：Transformer的核心贡献不是某个单点技术，而是一整套“用注意力替代循环”的设计哲学。它把序列建模从“按时间步递归”变成了“全局并行加权”。这个转变带来的不只是速度提升，更重要的是它让模型真正具备了捕捉长距离依赖的能力。下面我们一层层剥开看。

2. 注意力机制的前世今生：从“看哪里”到“怎么加权”

2.1 注意力机制的直觉来源

注意力机制最早不是为NLP发明的。在图像领域，人类看一张图不会像素级均匀扫描，而是会聚焦在关键区域。比如你看一张猫的照片，目光会自然落在猫的眼睛、耳朵、胡须上，背景的墙壁和地板就被“忽略”了。这种“有选择地聚焦”就是注意力的朴素直觉。

2014年前后，Bahdanau等人在机器翻译任务中引入了注意力，解决了Encoder-Decoder架构中“固定长度上下文向量”的瓶颈。之前的做法是把整个源句子压缩成一个固定维度的向量，句子一长信息就丢了。加了注意力之后，解码器每生成一个词，都可以回头“看”一遍编码器的所有隐藏状态，动态决定当前最该关注源句子的哪个部分。这个思路在当时是革命性的。

但早期的注意力还是依附在RNN之上的。真正让注意力“独立门户”的，是2017年那篇著名的论文。它提出了一个大胆的假设：既然注意力这么好用，那还要RNN干什么？直接把注意力作为核心组件，堆叠起来不就行了？这就是Transformer的由来。

2.2 Scaled Dot-Product Attention的计算逻辑

Transformer里用的注意力叫“缩放点积注意力”（Scaled Dot-Product Attention）。名字听着唬人，拆开看就三样东西：Query（查询）、Key（键）、Value（值）。你可以把它想象成一个软性的字典查询。传统字典是你给一个key，返回对应的value，一一对应。但注意力是“软”的：你给一个query，它跟所有key算相似度，得到一组权重，然后对所有value做加权求和。

具体计算分四步：

1. 算相似度：Query矩阵和Key矩阵的转置做矩阵乘法，得到形状为(序列长度, 序列长度)的分数矩阵。每个元素表示第i个位置对第j个位置的“关注程度”。
2. 缩放：把分数除以sqrt(d_k)，其中d_k是Key的维度。这一步非常关键，后面会详细解释为什么。
3. Softmax归一化：对每一行做softmax，把分数变成概率分布，所有值在0到1之间且和为1。
4. 加权求和：用归一化后的权重对Value矩阵加权求和，得到输出。

用公式写就是：

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

这个公式值得你反复看几遍，因为Transformer里所有的注意力变体都是它的变形。我建议你在纸上手推一遍维度变化：假设Q是(L, d_k)，K是(L, d_k)，那么QK^T是(L, L)，softmax后还是(L, L)，再乘V（(L, d_v)），输出就是(L, d_v)。整个过程没有循环，全是矩阵乘法，这就是它能并行的根本原因。

2.3 为什么要除以sqrt(d_k)

这是面试高频题，也是很多人第一次看论文时的疑惑。假设Q和K的每个元素都是独立同分布的，均值为0，方差为1。那么它们的点积q·k = sum(q_i * k_i)的方差就是d_k。当d_k很大时（比如512），点积的结果会变得非常大，导致softmax的输入落入饱和区。softmax在饱和区梯度极小，训练会变得非常困难。

除以sqrt(d_k)就是把方差拉回到1附近，让softmax的输入分布更合理，梯度更健康。你可以做个简单实验：随机生成两个512维向量，算点积，重复一万次，看看结果的分布范围。再除以sqrt(512)≈22.6，对比一下。这个操作看似简单，但少了它训练很容易崩。

注意：这里的d_k是Key的维度，不是序列长度。有些实现里会写成d_model，如果d_k等于d_model那没问题，但Multi-Head里每个头的d_k是d_model / num_heads，别搞混了。

3. Transformer整体架构拆解：编码器与解码器

3.1 编码器：六层堆叠的“理解机器”

原始论文里的编码器由6个相同的层堆叠而成。每一层包含两个子层：一个是Multi-Head Self-Attention，另一个是前馈神经网络（FFN）。每个子层外面都包了残差连接和Layer Normalization。

残差连接的作用不用多说，解决深层网络梯度消失。Layer Norm和Batch Norm的区别在于，Layer Norm是在特征维度上做归一化，跟batch大小无关，这对变长序列特别友好。Batch Norm在NLP里用起来很别扭，因为不同样本的序列长度不一样，batch统计量不稳定。

FFN其实就是一个两层的全连接网络，中间夹一个ReLU激活。第一层把维度从d_model扩展到d_ff（通常是4倍），第二层再压回d_model。比如d_model=512，d_ff=2048。这个设计让每个位置独立做非线性变换，增强表达能力。你可以理解为：Attention负责“混合”不同位置的信息，FFN负责“加工”每个位置自己的信息。

3.2 解码器：带掩码的“生成机器”

解码器也是6层堆叠，但每层有三个子层：Masked Multi-Head Self-Attention、Cross-Attention（也叫Encoder-Decoder Attention）、FFN。

Masked Self-Attention的关键在于“掩码”。训练时解码器是并行喂入整个目标序列的，但生成第t个词时不能看到t之后的词，否则就是作弊。所以要用一个上三角为负无穷的掩码矩阵，softmax之后这些位置的概率就变成0了。这个掩码机制保证了训练和推理的一致性。

Cross-Attention是连接编码器和解码器的桥梁。它的Query来自解码器上一层的输出，Key和Value来自编码器的最终输出。这样解码器每生成一个词，都能“查询”源序列的所有位置，动态决定关注哪里。这也是Transformer在翻译任务上表现优异的关键。

3.3 位置编码：给并行化补上顺序信息

Self-Attention有个致命问题：它对输入顺序完全不敏感。你把句子里的词打乱，注意力算出来的结果是一样的（只是位置对应关系变了）。但语言是有顺序的，“猫追老鼠”和“老鼠追猫”意思完全相反。所以必须显式注入位置信息。

原始论文用的是正弦位置编码：

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

其中pos是位置，i是维度索引。这个设计很巧妙：不同频率的正弦波组合，让每个位置有唯一的编码，而且相对位置关系可以通过线性变换表示。比如PE(pos+k)可以写成PE(pos)的线性函数，这让模型容易学到“相隔k个位置”这种关系。

后来也有很多变体，比如可学习的位置嵌入、相对位置编码、旋转位置编码（RoPE）等。RoPE在近年的LLM里非常流行，它通过旋转矩阵把位置信息编码进Query和Key的点积里，外推能力更好。如果你要自己实现，正弦编码是最简单的起点，但生产环境建议用RoPE或ALiBi。

实操心得：位置编码的维度必须和d_model一致，而且通常直接加到词嵌入上，不是拼接。加完之后再做dropout，再送入编码器第一层。

4. Multi-Head Attention：为什么要“多头”

4.1 单头注意力的局限

如果只用一组Q、K、V，模型只能学到一种“关注模式”。比如在翻译“The animal didn't cross the street because it was too tired”时，模型需要知道“it”指代“animal”。单头注意力可能学到这个指代关系，但可能就忽略了其他有用的关系，比如“cross”和“street”的动宾关系。

换句话说，单头注意力把所有的关系都压缩到一个加权平均里，信息会互相干扰。就像你用一个模糊的镜头看世界，只能看到一个大概的轮廓。

4.2 多头的拆分与合并

Multi-Head Attention的做法是：把d_model维的Q、K、V分别投影到h个低维子空间，每个子空间独立做注意力，最后把h个输出拼接起来，再过一个线性层。

假设d_model=512，h=8，那么每个头的维度就是512/8=64。每个头有自己的W_q、W_k、W_v投影矩阵，形状都是(512, 64)。8个头并行计算，每个头输出(L, 64)，拼接后变回(L, 512)，再乘W_o（(512, 512)）得到最终输出。

这样做的好处是：不同的头可以关注不同的模式。有的头可能关注语法依赖，有的头关注语义相似度，有的头关注位置邻近。论文里可视化过注意力权重，确实能看到不同头有不同的关注模式。

4.3 维度变化的完整推导

这部分我建议你亲手推一遍，因为代码里最容易出错的就是维度。假设：

• batch_size = 2
• seq_len = 10
• d_model = 512
• num_heads = 8
• d_k = d_v = 64

输入X形状(2, 10, 512)。经过W_q投影后还是(2, 10, 512)。然后reshape成(2, 10, 8, 64)，再transpose成(2, 8, 10, 64)。这样每个头就有独立的(10, 64)的Q矩阵。

QK^T得到(2, 8, 10, 10)，缩放、掩码、softmax后乘V得到(2, 8, 10, 64)。transpose回(2, 10, 8, 64)，reshape成(2, 10, 512)，最后过W_o。

注意：reshape和transpose的顺序不能错。如果先reshape成(2, 8, 10, 64)再transpose，和先transpose再reshape，结果完全不同。我见过太多人在这里debug半天。

5. 手把手实现一个单层Transformer

5.1 环境准备与依赖

我用PyTorch来演示，版本2.0以上。如果你用TensorFlow或JAX，逻辑一样，只是API不同。

pip install torch numpy

不需要装transformers库，我们要从零写，这样才能真正理解每一行在干什么。

5.2 实现Scaled Dot-Product Attention

import torch import torch.nn as nn import torch.nn.functional as F import math def scaled_dot_product_attention(q, k, v, mask=None): """ q, k, v: (batch, num_heads, seq_len, d_k) mask: (batch, 1, seq_len, seq_len) 或 (batch, num_heads, seq_len, seq_len) """ d_k = q.size(-1) scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k) if mask is not None: scores = scores.masked_fill(mask == 0, float('-inf')) attn_weights = F.softmax(scores, dim=-1) output = torch.matmul(attn_weights, v) return output, attn_weights

这段代码里，masked_fill用负无穷填充被掩码的位置，softmax之后这些位置就变成0。注意mask的语义：1表示保留，0表示屏蔽。如果你用mask == 1来填充，那就反了。

5.3 实现Multi-Head Attention模块

class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads): super().__init__() assert d_model % num_heads == 0, "d_model必须能被num_heads整除" self.d_model = d_model self.num_heads = num_heads self.d_k = d_model // num_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) # 投影并拆分多头 Q = self.W_q(q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) K = self.W_k(k).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) V = self.W_v(v).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) # 注意力计算 out, attn = scaled_dot_product_attention(Q, K, V, mask) # 合并多头 out = out.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) out = self.W_o(out) return out, attn

contiguous()这一步很关键。transpose之后张量在内存里不连续，直接view会报错。先contiguous再view，这是PyTorch的经典坑。

5.4 实现位置编码

class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len=5000): super().__init__() pe = torch.zeros(max_len, d_model) position = torch.arange(0, max_len).unsqueeze(1).float() div_term = torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)) pe[:, 0::2] = torch.sin(position * div_term) pe[:, 1::2] = torch.cos(position * div_term) pe = pe.unsqueeze(0) # (1, max_len, d_model) self.register_buffer('pe', pe) def forward(self, x): # x: (batch, seq_len, d_model) return x + self.pe[:, :x.size(1), :]

register_buffer让pe成为模型的一部分但不参与梯度更新，保存和加载模型时会自动带上。如果你用普通属性，保存模型时可能会丢。

5.5 组装编码器层

class EncoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout=0.1): super().__init__() self.self_attn = MultiHeadAttention(d_model, num_heads) self.ffn = nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Linear(d_ff, d_model) ) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.dropout = nn.Dropout(dropout) def forward(self, x, mask=None): # 子层1：自注意力 + 残差 + LayerNorm attn_out, _ = self.self_attn(x, x, x, mask) x = self.norm1(x + self.dropout(attn_out)) # 子层2：FFN + 残差 + LayerNorm ffn_out = self.ffn(x) x = self.norm2(x + self.dropout(ffn_out)) return x

这里用的是Post-LN结构，也就是先残差再加LayerNorm。原始论文就是这么做的。后来很多工作发现Pre-LN（先LayerNorm再子层）训练更稳定，不需要warmup。如果你自己训练，建议从Pre-LN开始试。

6. 训练与调试中的常见问题排查

6.1 损失不下降或变成NaN

这是最常见的问题。排查顺序如下：

现象	可能原因	解决方法
损失一开始就NaN	学习率太大	降到1e-4或更低，加warmup
损失震荡不降	梯度爆炸	加梯度裁剪，max_norm=1.0
损失缓慢下降后停滞	模型太小或数据太少	增大d_model或层数，检查数据质量
训练损失降但验证不降	过拟合	加dropout、权重衰减、早停

还有一个隐蔽的坑：位置编码的max_len设小了，序列一长就索引越界。或者mask的维度不对，广播之后形状错了但没报错，结果静默错误。

6.2 注意力权重全是均匀分布

如果你可视化注意力权重，发现每一行都差不多，说明模型没学到东西。可能的原因：

• 学习率太小，还没开始学
• 初始化有问题，比如全零初始化
• 数据本身没有可学的模式
• 层数太深，梯度传不到底层

我遇到过一次，排查半天发现是数据预处理时把标签和输入搞混了。所以调试时先在小数据集上过拟合，确保模型有能力学到东西，再上大规模数据。

6.3 显存不够用

Transformer的显存占用主要来自注意力矩阵，复杂度是O(L^2 * h)。序列长度翻倍，显存翻四倍。几个实用的优化手段：

• 用混合精度训练（AMP），显存减半，速度提升
• 梯度累积，小batch模拟大batch
• 梯度检查点，用时间换显存
• 如果序列特别长，考虑稀疏注意力或线性注意力变体

实操心得：调试时先把batch_size设为1，seq_len设为10，跑通整个流程。确认无误后再逐步放大。我见过太多人一上来就上全量数据，报错都不知道从哪查。

6.4 推理时生成重复文本

这是生成任务的经典问题。原因通常是解码策略太贪心。解决方法：

• 用beam search代替greedy decoding
• 加重复惩罚（repetition penalty）
• 用top-k或top-p采样增加多样性
• 检查训练数据里是否有大量重复模式

另外，如果训练时用了teacher forcing，推理时是自回归生成，两者分布不一致会导致曝光偏差。可以在训练后期逐渐减少teacher forcing的比例，缓解这个问题。

7. 从Transformer到现代大模型：哪些设计被保留了

7.1 被继承的核心设计

现在的主流大模型，不管是GPT系列、LLaMA系列还是其他，骨架都是Transformer的解码器部分。核心设计基本没变：

• Self-Attention作为核心混合机制
• 残差连接 + LayerNorm
• FFN作为位置独立的非线性变换
• 位置编码注入顺序信息

变化的主要是规模：层数从6层涨到几十上百层，d_model从512涨到几千甚至上万，注意力头数从8涨到几十。但底层逻辑还是那套。

7.2 被改进或替换的部分

一些细节做了优化：

• LayerNorm位置：从Post-LN改成Pre-LN，训练更稳定
• 激活函数：从ReLU换成GeLU、SwiGLU，效果更好
• 位置编码：从正弦编码换成RoPE、ALiBi，外推能力更强
• 注意力优化：FlashAttention、PagedAttention等，提升速度和显存效率
• 归一化：RMSNorm替代LayerNorm，计算更简单

这些改进都是在原有框架上的微调，没有颠覆性变化。所以你把原始Transformer搞懂了，再看这些变体会很快上手。

7.3 学习路径建议

如果你刚入门，我建议这个顺序：

1. 先手写一遍单层Transformer，确保每个矩阵维度都清楚
2. 在小数据集（比如字符级语言模型）上训练，观察损失和注意力权重
3. 读原始论文，对照代码理解每个设计选择
4. 再看现代大模型的技术报告，看哪些地方做了改进
5. 最后尝试复现一个小型GPT，从头训练一个能生成文本的模型

这个过程大概需要两到三周，每天投入两三个小时。不要跳过手写代码这一步，光看论文和博客是不够的。很多细节只有自己踩过坑才能真正记住。

我自己在第一次实现时，光是Multi-Head的维度变换就调了一整天。后来发现是transpose和view的顺序搞反了。这种错误看别人的代码永远学不会，必须自己写一遍。等你把单层Transformer跑通，再回头看那些大模型架构图，会发现它们只是把同样的积木堆得更高而已。