1. Transformer 的诞生与核心思想

2017 年 Google 发表《Attention Is All You Need》,Transformer 彻底取代了 RNN/LSTM 的序列建模范式。其核心创新:完全基于注意力机制,消除循环结构,实现高度并行化训练

Transformer 的三大设计支柱:

组件作用解决的问题
Self-Attention捕捉序列内任意位置间的依赖RNN 长距离梯度消失
Positional Encoding注入位置信息注意力本身是排列不变的
Residual + LayerNorm稳定深层网络训练深层网络梯度退化

2. 自注意力机制(Self-Attention)

2.1 公式推导

给定输入序列 $X \in \mathbb{R}^{n \times d}$,通过三个线性投影得到 Query、Key、Value:

$$Q = XW_Q, \quad K = XW_K, \quad V = XW_V$$

注意力计算公式:

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

其中 $\sqrt{d_k}$ 是缩放因子,防止内积值过大导致 softmax 梯度饱和。

为什么除以 $\sqrt{d_k}$? 当 $d_k$ 较大时,$QK^T$ 的方差为 $d_k$(假设 $Q, K$ 各分量独立标准正态),softmax 输入过大时梯度趋近于零。缩放后方差恢复为 1。

2.2 计算复杂度分析

操作时间复杂度空间复杂度
$QK^T$$O(n^2 \cdot d_k)$$O(n^2)$
softmax$O(n^2)$$O(n^2)$
$\times V$$O(n^2 \cdot d_v)$$O(n \cdot d_v)$
总计$O(n^2 \cdot d)$$O(n^2 + n \cdot d)$

序列长度 $n$ 是平方级瓶颈,这直接催生了 Flash Attention、稀疏注意力等优化方向。

3. 多头注意力(Multi-Head Attention)

3.1 原理

单一注意力头只能学习一种关联模式。多头注意力将 $d_{model}$ 维度拆分为 $h$ 个头,每个头独立做注意力,最后拼接:

$$\text{MultiHead}(Q,K,V) = \text{Concat}(\text{head}_1, …, \text{head}_h)W^O$$

$$\text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$$

其中 $W_i^Q \in \mathbb{R}^{d_{model} \times d_k}$,$d_k = d_{model} / h$。

3.2 代码实现

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

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, num_heads=8, dropout=0.1):
        super().__init__()
        assert d_model % num_heads == 0
        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, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.shape
        
        # 投影并分头: [B, n, d_model] -> [B, h, n, d_k]
        Q = self.W_q(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        
        # 注意力分数: [B, h, n, n]
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        attn_weights = F.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)
        
        # 加权求和: [B, h, n, d_k]
        context = torch.matmul(attn_weights, V)
        
        # 合并多头: [B, n, d_model]
        context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
        return self.W_o(context)

4. 前馈网络(Feed-Forward Network)

每个位置独立经过两层线性变换和 ReLU 激活:

$$\text{FFN}(x) = \max(0, xW_1 + b_1)W_2 + b_2$$

FFN 的隐藏维度通常是 $d_{model}$ 的 4 倍(如 $d_{model}=512$ 时 FFN 隐藏层为 2048)。GPT 系列使用 GELU 替代 ReLU:

$$\text{GELU}(x) = x \cdot \Phi(x) \approx 0.5x\left(1 + \tanh\left[\frac{\sqrt{2}}{\pi}(x + 0.044715x^3)\right]\right)$$

FFN 的作用:注意力层负责信息混合(序列内交互),FFN 负责特征变换(逐位置非线性)。两者分工明确。

5. 层归一化与残差连接

5.1 残差连接

每个子层的输出为:

$$\text{output} = \text{LayerNorm}(x + \text{Sublayer}(x))$$

残差连接解决梯度消失,允许训练更深的网络。

5.2 LayerNorm vs RMSNorm

原始 Transformer 使用 LayerNorm:

$$\text{LN}(x) = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} \cdot \gamma + \beta$$

LLaMA 等现代模型改用 RMSNorm,移除均值偏移,计算更快:

$$\text{RMSNorm}(x) = \frac{x}{\sqrt{\frac{1}{d}\sum_{i=1}^d x_i^2 + \epsilon}} \cdot \gamma$$

5.3 Pre-Norm vs Post-Norm

方式公式特点
Post-Norm (原始)$\text{LN}(x + \text{Sublayer}(x))$训练不稳定,需 warmup
Pre-Norm (现代)$x + \text{Sublayer}(\text{LN}(x))$训练稳定,主流选择

Pre-Norm 是当前 LLM 的标准选择,GPT-2 之后几乎成为共识。

6. 编码器-解码器 vs 解码器-only

6.1 原始 Transformer(Encoder-Decoder)

  • 编码器:双向注意力,用于理解输入序列
  • 解码器:带因果掩码的自注意力 + 交叉注意力
  • 典型应用:机器翻译

6.2 三大架构流派

架构代表模型注意力类型适用任务
Encoder-onlyBERT双向理解类(分类、NER)
Encoder-DecoderT5, BART编码双向 + 解码因果序列到序列(翻译、摘要)
Decoder-onlyGPT, LLaMA因果(单向)生成类(对话、续写)

6.3 GPT 的 Decoder-only 架构

GPT 去除了编码器和交叉注意力,仅保留因果自注意力 + FFN:

class DecoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, num_heads, dropout)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Linear(d_ff, d_model),
        )
        self.norm1 = nn.RMSNorm(d_model)  # 现代用 RMSNorm
        self.norm2 = nn.RMSNorm(d_model)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x, mask=None):
        # Pre-Norm + 因果自注意力
        x = x + self.dropout(self.attn(self.norm1(x), mask))
        # Pre-Norm + FFN
        x = x + self.dropout(self.ffn(self.norm2(x)))
        return x

因果掩码确保位置 $i$ 只能看到位置 $\leq i$:

def causal_mask(seq_len):
    return torch.tril(torch.ones(seq_len, seq_len)).unsqueeze(0).unsqueeze(0)

7. 训练效率分析

7.1 并行性对比

模型训练并行度序列依赖
RNN/LSTM低(时序依赖)
Transformer高(位置独立)

Transformer 的训练效率来自:所有位置可并行计算,GPU 友好的矩阵运算。

7.2 参数量计算

以 GPT-2 Small($d_{model}=768$, $L=12$, $h=12$)为例:

组件参数量公式
Embedding38.6Mvocab_size × d_model
Attention2.36M/层4 × d_model²
FFN4.72M/层2 × d_model × 4d_model
12层总计85.0M
总参数~124M

7.3 FLOPS 估算

前向传播 FLOPS 约为 $2N$($N$ 为参数量),反向传播约 $4N$,总训练 FLOPS:

$$\text{FLOPS}_{\text{train}} \approx 6ND$$

其中 $D$ 为训练 token 数。GPT-3($N=175B$, $D=300B$ tokens)约需 $3.15 \times 10^{23}$ FLOPS。

8. 从 Transformer 到现代 LLM 的演进

Transformer (2017)
    ├── BERT (Encoder-only, 2018)
    ├── GPT 系列 (Decoder-only, 2018-2024)
    │   ├── GPT-2: Pre-Norm, 1.5B
    │   ├── GPT-3: 175B, few-shot
    │   └── GPT-4: MoE, 多模态
    ├── LLaMA 系列 (2023-2024)
    │   ├── RMSNorm, RoPE, SwiGLU
    │   └── LLaMA 3: 405B, 128K 上下文
    └── DeepSeek/Mixtral (MoE, 2024)

现代改进总结

组件原始 Transformer现代 LLM
归一化LayerNorm (Post-Norm)RMSNorm (Pre-Norm)
激活函数ReLUSwiGLU / GELU
位置编码SinusoidalRoPE
注意力Dense AttentionGQA / Flash Attention
FFN标准 FFNSwiGLU(门控 FFN)

SwiGLU 是当前最流行的 FFN 变体:

$$\text{SwiGLU}(x) = \text{Swish}(xW_1) \otimes (xW_2)$$

其中 $\text{Swish}(x) = x \cdot \sigma(\beta x)$,通过门控机制提升表达能力。

9. 总结

Transformer 架构的成功源于三个核心设计:注意力机制实现长距离依赖建模、并行计算带来训练效率、残差连接与归一化支持深层网络训练。从原始 Transformer 到 GPT/LLaMA,架构本质未变,但每个组件都经历了精细化优化——RMSNorm 替代 LayerNorm、RoPE 替代绝对位置编码、SwiGLU 替代 ReLU FFN、GQA 替代标准 MHA。理解这些组件的设计原理,是深入 LLM 技术栈的基础。

加入讨论

这篇文章有姊妹讨论帖在硅基AGI论坛 — 全球首个碳基硅基认知交流平台。

碳基与硅基的智慧碰撞,认知差异创造无限可能。