Transformer架构详解：现代AI的基石
Transformer架构详解：现代AI的基石
引言
2017年，Google发表了划时代的论文《Attention is All You Need》，提出了Transformer架构。这一架构彻底改变了自然语言处理领域，并成为GPT、BERT等大语言模型的基础。
为什么需要Transformer？
传统的RNN和LSTM存在几个关键问题：
- 序列依赖：必须按顺序处理，无法并行化
- 长距离依赖：难以捕捉远距离的语义关系
- 梯度消失：序列越长，训练越困难
Transformer通过自注意力机制解决了这些问题。
Transformer的核心组件
1. 自注意力机制（Self-Attention）
自注意力机制允许模型在处理每个位置时，关注输入序列中的所有其他位置。
计算公式：
Attention(Q, K, V) = softmax(QK^T / √d_k) V
其中：
- Q（Query）：查询向量
- K（Key）：键向量
- V（Value）：值向量
- d_k：键向量的维度
2. 多头注意力（Multi-Head Attention）
多头注意力将注意力计算扩展到多个子空间，让模型能够同时关注不同位置的不同表示子空间。
import torch
import torch.nn as nn
class MultiHeadAttention(nn.Module):
def __init__(self, d_model, num_heads):
super().__init__()
self.num_heads = num_heads
self.d_model = d_model
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)
# 计算注意力
scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
attention = torch.softmax(scores, dim=-1)
output = torch.matmul(attention, V)
# 合并多头
output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
return self.W_o(output)
3. 位置编码（Positional Encoding）
由于Transformer没有循环结构，需要额外的位置信息来区分不同位置的词。
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, dtype=torch.float).unsqueeze(1)
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)
self.register_buffer('pe', pe)
def forward(self, x):
return x + self.pe[:, :x.size(1)]
4. 前馈网络（Feed-Forward Network）
每个Transformer层都包含一个前馈网络，通常由两个线性变换和一个激活函数组成。
class FeedForward(nn.Module):
def __init__(self, d_model, d_ff, dropout=0.1):
super().__init__()
self.linear1 = nn.Linear(d_model, d_ff)
self.linear2 = nn.Linear(d_ff, d_model)
self.dropout = nn.Dropout(dropout)
self.relu = nn.ReLU()
def forward(self, x):
return self.linear2(self.dropout(self.relu(self.linear1(x))))
Transformer编码器
编码器由N个相同的层堆叠而成，每层包含：
- 多头自注意力机制
- 前馈网络
- 残差连接和层归一化
class TransformerEncoderLayer(nn.Module):
def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
super().__init__()
self.self_attention = MultiHeadAttention(d_model, num_heads)
self.feed_forward = FeedForward(d_model, d_ff, dropout)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, x, mask=None):
# 自注意力 + 残差 + 归一化
attention_output = self.self_attention(x, x, x, mask)
x = self.norm1(x + self.dropout(attention_output))
# 前馈网络 + 残差 + 归一化
ff_output = self.feed_forward(x)
x = self.norm2(x + self.dropout(ff_output))
return x
Transformer解码器
解码器在编码器的基础上增加了：
- 掩码多头注意力（防止看到未来位置）
- 编码器-解码器注意力（关注编码器输出）
Transformer的优势
- 并行计算：所有位置可以同时计算，大大加速训练
- 长距离依赖：自注意力可以直接建立任意距离的连接
- 可解释性：注意力权重可以可视化，理解模型关注点
- 灵活性：可以处理任意长度的序列
实际应用
- GPT系列：只使用解码器，用于文本生成
- BERT：只使用编码器，用于文本理解
- T5：使用完整的编码器-解码器，用于多种NLP任务
- Vision Transformer（ViT）：将Transformer应用于图像识别
Transformer的演进：从原始架构到LLM时代
原始Transformer（2017）只有编码器-解码器结构，但后来的模型对其进行了大量改进：
关键改进
| 改进 | 提出时间 | 解决的问题 |
|---|---|---|
| RoPE（旋转位置编码） | 2021 | 更好的外推能力 |
| Grouped Query Attention (GQA) | 2023 | 减少KV Cache内存 |
| FlashAttention | 2022/2024 | 加速注意力计算 |
| SwiGLU激活函数 | 2021 | 提升模型性能 |
| RMSNorm | 2019/2023 | 更稳定的归一化 |
| MoE（混合专家） | 2017/2024 | 扩大模型规模 |
RoPE（旋转位置编码）
RoPE是目前主流LLM（LLaMA、GPT-NeoX等）使用的位置编码方案：
import torch
import torch.nn as nn
import math
class RoPE(nn.Module):
"""旋转位置编码（Rotary Position Embedding）"""
def __init__(self, dim, base=10000):
super().__init__()
self.dim = dim
self.base = base
# 预计算频率
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
self.register_buffer('inv_freq', inv_freq)
def forward(self, x, seq_len):
# 生成位置索引
pos = torch.arange(seq_len, device=x.device).float()
freqs = torch.outer(pos, self.inv_freq)
# 构建旋转矩阵
emb = torch.cat((freqs, freqs), dim=-1)
cos, sin = emb.cos(), emb.sin()
return cos, sin
MoE（Mixture of Experts）
MoE是GPT-4、Mixtral等超大模型的核心技术，它让模型在保持巨大参数量的同时，每次推理只激活部分参数：
┌──────────┐
输入Token ──┤ 门控网络 │
└────┬─────┘
│ 选择Top-2专家
┌────────────┼────────────┐
▼ ▼ ▼
┌──────────┐ ┌──────────┐ ┌──────────┐
│ Expert 1 │ │ Expert 2 │ │ Expert 3 │ (其余专家闲置)
└────┬─────┘ └────┬─────┘ └────┬─────┘
│ │ │
└─────────────┼─────────────┘
▼
加权输出
MoE的优势：
- 训练参数量可以是激活参数量的10-100倍
- 推理延迟几乎不增加
- Mixtral 8x7B：470亿总参数，仅激活120亿
注意力机制的优化
原始注意力 vs 优化注意力
原始注意力复杂度：O(n² × d)
- 序列长度n的平方级增长
- 长文本时内存和计算开销巨大
FlashAttention：O(n × d) 有效复杂度
- 利用GPU SRAM分层计算
- 减少HBM（高带宽内存）读写
- 速度提升2-4倍，显存降低50%+
降低KV Cache内存
# 标准Attention vs GQA vs MQA 对比
# 假设：32个头，序列长度2048，hidden_dim=4096
# Multi-Head Attention (MHA) - 原始Transformer
# 32个Q头 + 32个K头 + 32个V头 = 64个头向量
# Grouped Query Attention (GQA)
# 32个Q头 + 8个K头 + 8个V头 = 16个头向量（减少75% KV Cache）
# Multi-Query Attention (MQA)
# 32个Q头 + 1个K头 + 1个V头 = 2个头向量（减少97% KV Cache）
Transformer的数学本质
Transformer的核心操作可以用矩阵运算来理解：
输入嵌入: x ∈ R^(batch × seq_len × d_model)
↓
位置编码 + 嵌入: x' = x + PE(x)
↓
每一层Transformer Block:
┌─────────────────────────────────────┐
│ Self-Attention: │
│ attn = softmax(QK^T/√d) · V │
│ │
│ FFN: │
│ ff = W2 · gelu(W1 · attn + b1) + b2│
│ │
│ 残差 + 归一化: │
│ x = LayerNorm(x + attn) │
│ x = LayerNorm(x + ff) │
└─────────────────────────────────────┘
↓
输出: y ∈ R^(batch × seq_len × d_model)
总结
Transformer架构通过自注意力机制实现了高效的序列建模，成为现代AI的基石。理解Transformer的原理对于深入学习大语言模型至关重要。从2017年的原始论文到现在，Transformer已经演化为LLM时代的核心架构，RoPE、MoE、FlashAttention等创新让它在规模和效率上不断突破极限。