Transformer 架构深度解析：从原理推导到 PyTorch 实现

重点提示：虽然这 6 层结构长得一模一样，但它们内部的 参数（权重矩阵）是完全独立训练的！千万别以为是同一个层循环了 6 次。就像公司里的 6 级审批流，每一级领导关注的重点都不一样，科长看格式，处长看预算，局长看方向。

三、拆解编码器：机器是如何'理解'语言的？

好了，我们要把镜头推近，拿着显微镜看看单个 Encoder 内部到底长什么样。每一个 Encoder 层，虽然只有两步操作，却蕴含了深度学习的最高智慧。

Encoder 的两大护法

每一个 Encoder Layer 包含两个核心子层：

1. Multi-Head Self-Attention（多头自注意力机制）：这是 Transformer 的灵魂！它让模型在处理某个词时，能同时关注句子里的其他词。比如处理'苹果'时，能关注到前面的'吃'和后面的'甜'。
2. Position-wise Feed-Forward Network（前馈神经网络）：这就相当于人的大脑皮层，负责对提取出的特征进行进一步的消化、整合和非线性变换。

数据流动的细节

假设我们输入一个形状为 (Batch_Size, Seq_Len, D_Model) 的张量（Tensor）：

1. 这个张量首先流经 Self-Attention 层，这时候每个词都和其他词'聊了聊天'，交换了信息，输出一个新的张量。
2. 这个新张量经过 Add & Norm（残差连接和归一化，后面会细讲）。
3. 然后流经 Feed-Forward 层，被全连接层映射一番。
4. 再经过一次 Add & Norm。
5. 最后输出给下一层 Encoder。

这种设计极其精妙，既保证了信息的充分交互，又通过残差连接防止了梯度消失，让模型可以堆得非常深。

四、灵魂核心：Self-Attention 自注意力机制

现在，我们进入整篇文章最烧脑、也最精彩的部分——自注意力机制。如果没有它，Transformer 就是一堆废铁。

为什么要'自注意力'？

举个经典的例子：

'The animal didn't cross the street because it was too tired.'

作为人类，你一眼就知道这个 'it' 指的是 'animal'，而不是 'street'。为什么？因为你的大脑在处理 'it' 时，下意识地联系了上下文，发现 'animal' 能 'tired'（累），而 'street' 不能。

对于计算机来说，'it' 只是一个 ASCII 码或者一个 ID。怎么让计算机学会这种'联系'？Self-Attention 就是干这个的。它让 'it' 这个词向量在计算时，能够吸收 'animal' 这个词向量的信息，从而让自己变成一个'指代了动物的它'。

更是'相亲'大会：Q、K、V 的爱恨情仇

为了实现这个机制，Transformer 创造性地引入了三个概念：Query (Q)、Key (K)、Value (V)。

很多同学看论文死活不理解这仨是啥。来，我们用相亲大会做个比喻：

• Query (Q) —— 择偶标准：这是你（当前的词）发出的信号。比如你在处理'帅哥'这个词，你的 Q 向量里可能包含'寻找高颜值'、'寻找男性'这样的特征请求。
• Key (K) —— 个人标签：这是场上其他嘉宾（句子里的其他词）举着的牌子。比如'我'这个词的 K 向量可能是'第一人称'、'主语'；'苹果'的 K 向量可能是'水果'、'红色'、'圆的'。
• Value (V) —— 灵魂内涵：这是嘉宾真正的内涵信息。如果配对成功，你就要把这个 V 向量领回家，融入到你自己的脑子里。

Self-Attention 的过程，就是拿着你手中的 Q，去和全场所有人的 K 进行匹配（点积）。匹配度越高，你就越关注对方的 V，把对方的 V 加权融合到自己身上。

硬核计算流程（手撕数学公式）

别怕数学，我们一步步来拆解这个过程。假设输入矩阵是 X。

第一步：线性变换生成 Q, K, V 我们需要三个权重矩阵 W^Q, W^K, W^V（这些是模型训练出来的参数）。

$$Q = X \cdot W^Q$$ $$K = X \cdot W^K$$ $$V = X \cdot W^V$$

第二步：计算注意力分数（Attention Score） 我们要计算 Q 和 K 的相似度。最简单的方法就是点积（Dot Product）。

$$Score = Q \cdot K^T$$

如果两个向量方向一致，点积就大；垂直则为 0。这就量化了'匹配度'。

第三步：缩放（Scale） 这一步是 Google 的独门秘籍。他们把 Score 除以 $\sqrt{d_k}$（通常是 $d_{model}$ 维度的平方根，比如 $\sqrt{64}=8$）。

$$Scaled_Score = \frac{Score}{\sqrt{d_k}}$$

为什么要做这一步？ 如果 $d_k$ 很大，点积结果会变得非常大。一旦进入 Softmax 函数，大的数值会把概率挤向 1 和 0 的两端，导致梯度几乎为 0（梯度消失），模型就学不动了。除以 $\sqrt{d_k}$ 是为了把数值拉回梯度敏感区。

第四步：归一化（Softmax） 将分数转化为概率分布，所有分数加起来等于 1。

$$Attention_Weight = Softmax(Scaled_Score)$$

比如 'it' 对 'animal' 的关注度是 0.8，对 'street' 是 0.1，对其他词是 0.1。

第五步：加权求和（Weighted Sum） 最后，用算出来的权重去乘对应的 V 向量。

$$Z = Attention_Weight \cdot V$$

这就是最终的 Self-Attention 输出。

总结成一个公式，就是 AI 界的 $E=mc^2$：

$$Attention(Q, K, V) = Softmax\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

Python 代码实战：手写 Attention

光说不练假把式，我们用 PyTorch 来写一个函数：

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

def scaled_dot_product_attention(query, key, value, mask=None):
    """
    计算 'Scaled Dot-Product Attention'
    query, key, value 的形状：[batch_size, num_heads, seq_len, d_k]
    """
    d_k = query.size(-1)
    # 获取维度，用于缩放
    
    # 1. 计算 Q * K.T (矩阵乘法)
    # key.transpose(-2, -1) 是为了把最后两个维度交换，以便进行矩阵乘法
    scores = torch.matmul(query, key.transpose(-2,-1)) / math.sqrt(d_k)
    
    # 2. 如果有掩码（Mask），把不需要关注的地方设为负无穷
    # 这样 Softmax 之后就会变成 0
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    
    # 3. Softmax 归一化，得到注意力权重
    p_attn = F.softmax(scores, dim=-1)
    
    # 4. 加权求和：Prob * Value
    return torch.matmul(p_attn, value), p_attn

# 测试一下
d_k = 64
q = torch.randn(2, 8, 10, d_k) # Batch=2, Heads=8, Len=10
k = torch.randn(2, 8, 10, d_k)
v = torch.randn(2, 8, 10, d_k)
output, weights = scaled_dot_product_attention(q, k, v)
print("Output shape:", output.shape) # 应该是 [2, 8, 10, 64]

五、多头注意力：三个臭皮匠，顶个诸葛亮

为什么要'多头'（Multi-Head）？

如果只用一组 Q、K、V，可能只能捕捉到一种关系。但在语言中，词与词的关系是复杂的：

• 语法关系：主谓宾搭配。
• 语义关系：指代消歧（it -> animal）。
• 情感关系：修饰词的强烈程度。

只用一个头，就像只用一只眼睛看世界，虽然能看清，但缺乏立体感。所以，Transformer 搞了 Multi-Head Attention（多头注意力机制）。

这就好比请了一个专家团：

• Head 1：专门负责看语法结构。
• Head 2：专门负责看词义指代。
• Head 3：专门负责看上下文情感。
• …
• Head 8：专门负责看其他的。

最后把大家的意见拼起来（Concat），再经过一次线性变换（Linear），就是最终结果。

水管分流：维度变换的魔法

很多初学者在这里会晕：头多了，计算量岂不是爆炸了？ Google 的工程师非常鸡贼（褒义）。假设模型的标准维度 $d_{model} = 512$，我们用 8 个头 ($h=8$)。那么，每个头处理的维度并不是 512，而是 $512 / 8 = 64$！

这样做的好处是：虽然头多了，但总的计算量并没有显著增加，因为每个头的维度变小了。我们只是把一个大的高维空间，切割成了 8 个独立的子空间，让模型在不同的子空间里学习不同的特征。

看上图这幅生动的 '水管分流图'：

1. 原始向量（最左边的粗管子）被分流成了 8 条细水管（Multi-Head）。
2. 每一条细水管流向一个独立的 处理池（Attention Head），在里面进行独有的搅拌和反应。
3. 处理完后，8 条细水管流出的液体再次 汇聚（Concat），变成一条粗水管。
4. 最后经过一个总阀门（Linear Layer），输出最终结果。

代码实现

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        assert d_model % num_heads == 0 # 确保能整除
        self.d_k = d_model // num_heads
        self.num_heads = num_heads
        self.d_model = d_model
        
        # 定义四个全连接层：Q, K, V 和 Output
        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, query, key, value, mask=None):
        batch_size = query.size(0)
        
        # 1. 线性变换 + 分头 (Reshape)
        # [batch, len, d_model] -> [batch, len, heads, d_k] -> [batch, heads, len, d_k]
        query = self.w_q(query).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        key = self.w_k(key).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        value = self.w_v(value).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        
        # 2. 计算 Attention (这是我们上面写的函数)
        x, self.attn = scaled_dot_product_attention(query, key, value, mask)
        
        # 3. 拼接 (Concat)
        # [batch, heads, len, d_k] -> [batch, len, heads, d_k] -> [batch, len, d_model]
        x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
        
        # 4. 最终线性变换
        return self.w_o(x)

六、那些容易被忽视的细节：位置编码与残差

Positional Encoding：给盲人指路

Transformer 有个大 bug：它没有时间概念。因为它是并行计算的，对于 Self-Attention 来说，'我爱你'和'你爱我'，输入的词向量集合是一样的，它不知道谁排在谁前面。它就像一个脸盲症患者，虽然认识所有人，但不知道大家排队的顺序。

为了解决这个问题，我们需要给每个词加上一个 位置编码（Positional Encoding）。这就像给每个词发一个'座位号'。

Google 的做法非常数学，他们没有用简单的 1, 2, 3, 4（因为这样数值会无限大），也没有用 0~1 的归一化（因为不同长度句子的步长不一样）。他们选用了 正弦（Sin）和余弦（Cos）函数：

$$PE_{(pos, 2i)} = \sin(pos / 10000^{2i/d_{model}})$$ $$PE_{(pos, 2i+1)} = \cos(pos / 10000^{2i/d_{model}})$$

这里的 $pos$ 是词在句子中的位置，$i$ 是维度的索引。

为什么要这么搞？

1. 每个位置唯一：正弦和余弦的组合能保证每个位置的编码都是独一无二的。
2. 相对距离不变：这是最骚的操作。根据三角函数公式 $\sin(\alpha+\beta) = \sin\alpha \cos\beta + \cos\alpha \sin\beta$，位置 $pos+k$ 的编码可以被位置 $pos$ 的编码线性表示。这意味着模型可以轻松学习到'词 A 在词 B 后面第 3 个位置'这样的相对关系，而不仅仅是绝对位置。

Add & Norm：模型的护身符

仔细看架构图，你会发现在每个子层（Self-Attention, FFN）后面，都跟着一个 Add & Norm。

1. Norm（层归一化，Layer Normalization）：不同于 CV 领域常用的 Batch Norm（BN），NLP 中常用 Layer Norm（LN）。
   • BN：把一个 Batch 里所有人的同一门课成绩归一化。这在 NLP 里行不通，因为句子长短不一，且 batch size 较小。
   • LN：把一个人（一个样本）的所有课程成绩归一化。它让每个样本的特征分布保持稳定，加速训练收敛。

Add（残差连接，Residual Connection）：$Output = Input + SubLayer(Input)$

这招是从 ResNet 学来的。它的作用是修一条'高速公路'，让梯度可以直接流到前面的层。

大白话：如果在这一层没学到啥新东西，或者把东西学乱了，至少保留原来的样子别变坏。这让深层网络训练成为可能。

七、解码器：带面具的预言家

Decoder 的结构跟 Encoder 很像，但有两个关键不同点，这决定了它'生成式'的特性。

Masked Self-Attention（掩码自注意力）

Decoder 在训练时，虽然输入了完整的正确答案（Target），比如 'I am a student'。但在预测第 1 个词 'I' 的时候，它绝对不能看到后面的 'am a student'。否则这就是考试作弊！

所以，我们需要用一个 Mask（掩码），把当前位置后面的内容全都遮住。具体操作是生成一个上三角矩阵，把上三角区域（未来信息）设为 $-\infty$。在 Softmax 之后，这些位置的概率就变成了 0。

这就像考试时，监考老师拿一张黑纸挡住你后面还没写的题，逼着你只能根据前面写的内容推断下一步。

Encoder-Decoder Attention（交叉注意力）

这是 Decoder 的第二层注意力，也是连接 Encoder 和 Decoder 的桥梁。

• Query (Q)：来自 Decoder 上一层的输出。（我在找什么？）
• Key (K)：来自 Encoder 的最终输出。（原文里有什么？）
• Value (V)：来自 Encoder 的最终输出。（原文的具体信息）

这个机制的本质是：翻译官（Decoder）每写一个单词，都要抬头看一眼原文（Encoder），看看这句话对应原文的哪一部分，然后把那部分信息抓取过来。

八、拓展方案：Transformer 的进化与变体

Transformer 虽然强大，但不是完美的。计算复杂度 $O(N^2)$ 是它的阿喀琉斯之踵。这几年，大佬们基于它搞出了无数变体。这里给大家介绍三个最硬核的拓展方案。

拓展方案 1：FlashAttention —— 天下武功，唯快不破

痛点：标准 Attention 的计算复杂度是 $O(N^2)$。序列长度翻倍，显存占用爆炸。处理长文本（比如写小说、读财报）时，GPU 显存直接 OOM（Out Of Memory）。而且，显存的读写速度（HBM）远慢于计算核心的速度（SRAM），大量时间浪费在搬运数据上。

解决方案：FlashAttention。这是一种 IO 感知（IO-Aware） 的精确注意力算法。它不改变数学结果，而是改变了计算顺序。它利用了 GPU 内存层级的特点（SRAM 比 HBM 快得多），通过 Tiling（分块计算） 和 Recomputation（重计算） 技术：

1. 把大矩阵切成小块，塞进高速缓存（SRAM）里算。
2. 算完一块再算下一块，避免把整个巨大的 $N \times N$ 矩阵写回显存。
3. 虽然多算了一些步骤（重计算），但因为减少了慢速显存的读写，整体速度反而飞起。

• 效果：速度快 2-4 倍，显存占用降低 5-10 倍。
• 应用：现在的 GPT-4、Llama 3 能处理 128k 甚至 1M 的上下文，全靠它续命。

拓展方案 2：Vision Transformer (ViT) —— 跨界打劫

痛点：CNN（卷积神经网络）长期统治 CV 界，但 CNN 的卷积核只能看局部，看完全局需要堆很深。

解决方案：ViT (Vision Transformer)。Google 说：'谁说 Transformer 只能处理文字？'他们把一张图片（比如 256x256）切成一个个 16x16 的小方块（Patches）。然后，把每个小方块展平，当成一个'单词'（Token），加上位置编码，直接扔进 Transformer Encoder。

看上图的演示：

1. 图片切割：一张风景照被切成了九宫格。
2. 序列化：每个格子变成了一个向量，排成了长队。
3. 输入：这队向量像单词一样进入 Transformer。

• 核心思想：Image Patch Embedding。
• 效果：在海量数据预训练下，性能吊打 ResNet。它证明了 Transformer 架构的普适性——万物皆可 Token 化。

拓展方案 3：Dynamic Sparse Attention (DSA) —— 动态稀疏注意力

痛点：全注意力机制太浪费了。翻译'我今天吃了苹果'中的'苹果'时，真的需要关注'今天'吗？可能并不需要。全量计算包含大量噪音和无效计算。

解决方案：DSA (Dynamic Sparse Attention)。这种技术被 DeepSeek 和智谱 GLM 等模型广泛探索。它动态地筛选出 Top-k 个最重要的 Token 进行计算，其他的直接忽略。就像人类阅读时，只盯着关键词看，自动过滤掉'的'、'了'、'吗'这些虚词。

• 原理：利用一个低成本的评估网络，预测哪些块（Block）是重要的，只计算这些块的 Attention。
• 效果：在不降低模型效果的前提下，大幅降低计算成本，让模型能跑在更便宜的显卡上，甚至适配国产芯片。

九、最终章：完整流程总结

让我们把所有零件拼起来，最后走一遍流程（以机器翻译'我爱 AI' -> 'I love AI'为例）：

1. 输入处理：
   • '我爱 AI' -> 分词 -> Word Embedding -> 加上 Positional Encoding -> 得到向量矩阵 X。
2. Encoder 编码：
   • X 进入 Layer 1。
   • Self-Attention：捕捉'我'和'爱'的关系。
   • Add & Norm。
   • Feed Forward：特征变换。
   • Add & Norm。
   • … 重复 6 次 …
   • 输出包含丰富语义的 Key 和 Value 矩阵（也就是原文的深度记忆）。
3. Decoder 解码：
   • Step 1：输入 <Start> 符。
   • Masked Self-Attention（自己看自己，防止剧透）。
   • Encoder-Decoder Attention：拿着 <Start> 的 Query，去查 Encoder 输出的 Key/Value。
   • Feed Forward -> Linear -> Softmax。
   • 输出概率最高的词：'I'。
   • Step 2：输入 <Start> I。
   • 重复上述过程，输出 'love'。
   • Step 3：输入 <Start> I love。
   • 输出 'AI'。
   • Step 4：输出 <End>，结束。

十、实战：手写一个完整的 Encoder Layer

光说不练假把式。这是一个基于 PyTorch 的微型 Encoder Layer 实现，你可以直接复制运行：

import torch
import torch.nn as nn

class EncoderLayer(nn.Module):
    def __init__(self, d_model, heads, dropout=0.1):
        super(EncoderLayer, self).__init__()
        # 1. 两个层归一化
        self.norm_1 = nn.LayerNorm(d_model)
        self.norm_2 = nn.LayerNorm(d_model)
        
        # 2. 多头注意力
        self.attn = nn.MultiheadAttention(d_model, heads)
        
        # 3. 前馈网络 (通常中间层维度是 d_model 的 4 倍)
        self.ff = nn.Sequential(
            nn.Linear(d_model, 4 * d_model),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(4 * d_model, d_model)
        )
        
        # 4. Dropout 层
        self.dropout_1 = nn.Dropout(dropout)
        self.dropout_2 = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        """
        x: [seq_len, batch_size, d_model]
        """
        # --- 子层 1: Self-Attention ---
        # 这里的实现采用了 Pre-Norm 结构（先 Norm 再 Attention），训练更稳定
        x2 = self.norm_1(x)
        # PyTorch 的 MultiheadAttention 输入是 (seq_len, batch, dim)
        attn_output, _ = self.attn(x2, x2, x2, attn_mask=mask)
        # 残差连接：x + Dropout(SubLayer(x))
        x = x + self.dropout_1(attn_output)
        
        # --- 子层 2: Feed Forward ---
        x2 = self.norm_2(x)
        ff_output = self.ff(x2)
        # 残差连接
        x = x + self.dropout_2(ff_output)
        return x

# --- 测试代码 ---
d_model = 512
heads = 8
seq_len = 20
batch_size = 4

# 模拟输入
x = torch.randn(seq_len, batch_size, d_model)
layer = EncoderLayer(d_model, heads)
out = layer(x)

print(f"输入尺寸：{x.shape}")
print(f"输出尺寸：{out.shape}") # 结果应该是完全一致的，证明维度保持不变

十一、结语

Transformer 不仅仅是一个模型，它更像是一种思维方式的转变——从串行到并行，从局部到全局。

现在的 ChatGPT、文心一言、Stable Diffusion，甚至是 AlphaFold（预测蛋白质结构），骨子里流淌的都是 Transformer 的血液。它不仅统一了 NLP，正在统一 CV，甚至在向科学计算领域扩张。它就是 AI 时代的'通用语'。

这就好比上图展示的，复古的机械（代表传统算法）与未来的光影（代表 Transformer AI）在星辰大海中握手。我们正处在一个时代的转折点上，而你，现在已经掌握了开启这个时代的钥匙。

参考链接：

• Attention Is All You Need (Original Paper)
• The Illustrated Transformer (Jay Alammar)
• PyTorch Official Documentation
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• Vision Transformer (ViT)