【机器学习精通】第11章 | 注意力机制与Transformer：自注意力革命

所谓伊人，在水一方333

554人浏览 · 2026-03-16 07:59:33

所谓伊人，在水一方333 · 2026-03-16 07:59:33 发布

补充：注意力机制被誉为深度学习领域继卷积神经网络之后的又一里程碑式突破，它彻底改变了序列建模的方式。

环境声明

Python 版本：Python 3.12+（建议使用3.10以上版本）
深度学习框架：PyTorch 2.3+
开发工具：PyCharm 或 VS Code
GPU支持：CUDA 11.8+（可选但推荐）
操作系统：Windows / macOS / Linux（通用）

学习目标

完成本章学习后，你将能够：

理解注意力机制的核心原理，包括Query、Key、Value的计算方式
掌握自注意力与多头注意力的数学推导过程
深入理解Transformer的Encoder-Decoder架构设计
区分绝对位置编码与相对位置编码的优劣
了解Vision Transformer在计算机视觉领域的应用
熟悉BERT、GPT系列及Mamba等前沿变体模型
独立使用PyTorch实现简化版Transformer

1. 注意力机制原理

1.1 从人类注意力到机器注意力

人类在阅读长文本时，不会对所有词语给予同等关注，而是根据上下文动态调整注意力焦点。例如，在阅读"猫坐在垫子上，因为它很温暖"时，我们的注意力会集中在"它"与"垫子"之间的关联上。

机器注意力机制正是模拟这一过程，让模型能够：

选择性关注：从大量输入信息中选择性地关注相关部分
动态权重：根据当前任务动态调整不同输入的重要性
长距离依赖：建立序列中任意两个位置之间的直接联系

1.2 Query、Key、Value三要素

注意力机制的核心是三个向量矩阵：

组件	符号	含义	类比
Query	Q	查询向量，表示当前要查询的信息	你在图书馆查找书籍时的检索关键词
Key	K	键向量，表示输入序列中每个位置的标识	图书馆书籍的索引标签
Value	V	值向量，表示输入序列中每个位置的实际内容	书籍本身的内容

1.3 缩放点积注意力公式

注意力机制的核心计算公式如下：

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

公式分解：

$QK^T$ ：计算Query与所有Key的点积，得到相似度分数
$1d_k\frac{1}{\sqrt{d\_k}}$ ：缩放因子，防止点积结果过大导致softmax梯度消失
softmax：将相似度转换为概率分布（注意力权重）
与V相乘：根据权重对Value进行加权求和

1.4 注意力计算示例

假设我们有一个序列长度为3，维度 $d\_k=4$ 的简单例子：

import torch
import torch.nn.functional as F

# 设置随机种子保证可复现
torch.manual_seed(42)

# 定义输入维度
seq_len = 3      # 序列长度
d_k = 4          # 键/查询维度
d_v = 4          # 值维度

# 随机生成Q, K, V矩阵
Q = torch.randn(seq_len, d_k)  # (3, 4)
K = torch.randn(seq_len, d_k)  # (3, 4)
V = torch.randn(seq_len, d_v)  # (3, 4)

# 计算注意力分数
scores = torch.matmul(Q, K.transpose(-2, -1))  # (3, 3)
print("注意力分数矩阵:")
print(scores)

# 应用缩放因子
scores = scores / torch.sqrt(torch.tensor(d_k, dtype=torch.float32))

# 应用softmax获取注意力权重
attn_weights = F.softmax(scores, dim=-1)
print("\n注意力权重矩阵:")
print(attn_weights)

# 加权求和得到输出
output = torch.matmul(attn_weights, V)  # (3, 4)
print("\n注意力输出:")
print(output)

输出结果展示了注意力权重矩阵的行和为1（概率分布特性），每个输出位置都是所有Value的加权组合。

2. 自注意力与多头注意力

2.1 自注意力机制（Self-Attention）

自注意力是注意力机制的一种特殊形式，其中Query、Key、Value都来自同一个输入序列。这使得序列中的每个位置都能关注到序列中的所有位置。

自注意力的优势：

并行计算：不同于RNN的顺序处理，自注意力可以并行计算所有位置的表示
长距离依赖：任意两个位置之间的距离都是 $O (1)$ ，不受序列长度影响
可解释性：注意力权重矩阵直观展示了模型关注的重点

2.2 自注意力完整推导

给定输入序列 $X∈Rn×d_modelX \in \mathbb{R}^{n \times d\_{model}}$ ，自注意力的计算过程：

第一步：线性变换

$XW^Q, \quad K = XW^K, \quad V = XW^V$

其中 $WQ,WK∈Rd_model×d_kW^Q, W^K \in \mathbb{R}^{d\_{model} \times d\_k}$ ， $WV∈Rd_model×d_vW^V \in \mathbb{R}^{d\_{model} \times d\_v}$

第二步：计算注意力

$SelfAttention(X)=softmax(XWQ(XWK)Td_k)XWV \text{SelfAttention}(X) = \text{softmax}\left(\frac{XW^Q(XW^K)^T}{\sqrt{d\_k}}\right)XW^V$

2.3 多头注意力（Multi-Head Attention）

单一注意力头可能只关注特定类型的信息。多头注意力通过 $h$ 个并行的注意力头，让模型同时关注不同子空间的信息。

多头注意力公式：

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

其中每个注意力头：

$head_i=Attention(QW_iQ,KW_iK,VW_iV) \text{head}\_i = \text{Attention}(QW\_i^Q, KW\_i^K, VW\_i^V)$

参数维度说明：

参数	维度	说明
$W\_i^Q$	$d_model×d_kd\_{model} \times d\_k$	第i个头的Query投影矩阵
$W\_i^K$	$d_model×d_kd\_{model} \times d\_k$	第i个头的Key投影矩阵
$W\_i^V$	$d_model×d_vd\_{model} \times d\_v$	第i个头的Value投影矩阵
$W^O$	$hd_v×d_modelhd\_v \times d\_{model}$	输出投影矩阵

通常设置 $d\_k = d\_v = d\_{model}/h$ ，保证计算效率。

2.4 多头注意力PyTorch实现

import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    """
    多头注意力机制实现
    """
    def __init__(self, d_model=512, num_heads=8, dropout=0.1):
        super(MultiHeadAttention, self).__init__()
        
        # 确保d_model能被num_heads整除
        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)
        
        self.dropout = nn.Dropout(dropout)
        
    def scaled_dot_product_attention(self, Q, K, V, mask=None):
        """
        计算缩放点积注意力
        Args:
            Q: (batch, num_heads, seq_len, d_k)
            K: (batch, num_heads, seq_len, d_k)
            V: (batch, num_heads, seq_len, d_v)
            mask: 可选的掩码矩阵
        """
        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_weights = torch.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)
        
        output = torch.matmul(attn_weights, V)
        return output, attn_weights
    
    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)
        
        # 1. 线性投影并分头
        # (batch, seq, d_model) -> (batch, seq, num_heads, d_k) -> (batch, num_heads, seq, d_k)
        Q = self.W_q(query).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(key).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(value).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        
        # 2. 计算注意力
        attn_output, attn_weights = self.scaled_dot_product_attention(Q, K, V, mask)
        
        # 3. 拼接多头并投影
        # (batch, num_heads, seq, d_k) -> (batch, seq, num_heads, d_k) -> (batch, seq, d_model)
        attn_output = attn_output.transpose(1, 2).contiguous().view(
            batch_size, -1, self.d_model
        )
        
        output = self.W_o(attn_output)
        return output, attn_weights


# 测试多头注意力
if __name__ == "__main__":
    batch_size = 2
    seq_len = 10
    d_model = 512
    num_heads = 8
    
    mha = MultiHeadAttention(d_model=d_model, num_heads=num_heads)
    x = torch.randn(batch_size, seq_len, d_model)
    
    output, weights = mha(x, x, x)
    print(f"输入形状: {x.shape}")
    print(f"输出形状: {output.shape}")
    print(f"注意力权重形状: {weights.shape}")

3. Transformer架构详解

3.1 Transformer整体架构

Transformer由Google在2017年论文《Attention Is All You Need》中提出，完全基于注意力机制，摒弃了RNN和CNN。

架构组成：

┌─────────────────────────────────────────────────────────────┐
│                        Transformer                          │
├─────────────────────────────────────────────────────────────┤
│  Encoder Stack (N=6层)        │  Decoder Stack (N=6层)      │
│  ┌─────────────────────┐      │  ┌─────────────────────┐    │
│  │ Multi-Head Attention│      │  │ Masked Multi-Head   │    │
│  │ (Self-Attention)    │      │  │ Attention           │    │
│  ├─────────────────────┤      │  ├─────────────────────┤    │
│  │ Add & Norm          │      │  │ Add & Norm          │    │
│  ├─────────────────────┤      │  ├─────────────────────┤    │
│  │ Feed Forward        │      │  │ Multi-Head Attention│    │
│  │ (全连接层)           │      │  │ (Encoder-Decoder)   │    │
│  ├─────────────────────┤      │  ├─────────────────────┤    │
│  │ Add & Norm          │      │  │ Add & Norm          │    │
│  └─────────────────────┘      │  ├─────────────────────┤    │
│                               │  │ Feed Forward        │    │
│                               │  ├─────────────────────┤    │
│                               │  │ Add & Norm          │    │
│                               │  └─────────────────────┘    │
└─────────────────────────────────────────────────────────────┘

3.2 Encoder详解

Encoder由N个相同的层堆叠而成，每层包含两个子层：

多头自注意力机制：处理输入序列内部的关系
前馈神经网络：对每个位置独立进行非线性变换

每个子层后都有残差连接和层归一化：

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

3.3 Decoder详解

Decoder同样由N个相同的层堆叠，每层包含三个子层：

掩码多头自注意力：防止位置关注后续位置（保持自回归特性）
Encoder-Decoder注意力：Query来自Decoder，Key/Value来自Encoder
前馈神经网络：与Encoder相同

3.4 前馈神经网络（FFN）

每个Encoder和Decoder层都包含一个全连接前馈网络：

$FFN(x)=max⁡(0,xW_1+b_1)W_2+b_2 \text{FFN}(x) = \max(0, xW\_1 + b\_1)W\_2 + b\_2$

这是一个两层的全连接网络，中间使用ReLU激活函数。输入和输出的维度是 $d\_{model}=512$ ，中间层维度 $d\_{ff}=2048$ 。

3.5 残差连接与层归一化

残差连接（Residual Connection）：

解决深层网络梯度消失问题，允许梯度直接回传：

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

层归一化（Layer Normalization）：

对每个样本的所有特征进行归一化，稳定训练过程：

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

其中 $μ\mu$ 和 $σ2\sigma^2$ 是样本的均值和方差， $γ\gamma$ 和 $β\beta$ 是可学习的缩放和平移参数。

4. 位置编码与层归一化

4.1 为什么需要位置编码

Transformer没有RNN的递归结构，也没有CNN的局部感受野，因此模型本身无法感知序列中元素的顺序信息。位置编码（Positional Encoding）为模型提供位置信息。

4.2 绝对位置编码（Original Transformer）

原始Transformer使用正弦和余弦函数生成位置编码：

$PE_(pos,2i)=sin⁡(pos100002i/d_model) PE\_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d\_{model}}}\right)$

$PE_(pos,2i+1)=cos⁡(pos100002i/d_model) PE\_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d\_{model}}}\right)$

其中：

$p os$ ：位置索引
$i$ ：维度索引
$d\_{model}$ ：模型维度

设计优点：

唯一性：每个位置有唯一的编码
相对位置：对于固定的偏移 $k$ ， $PE\_{pos+k}$ 可以表示为 $PE\_{pos}$ 的线性函数
有界性：值域在[-1, 1]之间

4.3 可学习位置编码

BERT等模型采用可学习的位置编码，将位置编码作为可训练参数：

self.pos_embedding = nn.Embedding(max_seq_len, d_model)

优缺点对比：

类型	优点	缺点
正弦位置编码	无需训练，可外推到更长序列	表达能力有限
可学习位置编码	表达能力更强	无法外推到训练时未见过的长度

4.4 相对位置编码

相对位置编码不再为每个绝对位置分配编码，而是编码位置之间的相对距离。代表方法包括：

Transformer-XL：引入相对位置偏置
T5：使用可学习的相对位置嵌入
RoPE（Rotary Position Embedding）：通过旋转矩阵编码位置，被LLaMA等模型采用

RoPE公式：

$qe^{im\theta}$

其中 $m$ 是位置， $θ\theta$ 是预定义的旋转角度。

4.5 Pre-Norm vs Post-Norm

层归一化的位置有两种选择：

Post-Norm（原始Transformer）：

$x_l+1=LayerNorm(x_l+Sublayer(x_l)) x\_{l+1} = \text{LayerNorm}(x\_l + \text{Sublayer}(x\_l))$

Pre-Norm（现代主流）：

$x_l+1=x_l+Sublayer(LayerNorm(x_l)) x\_{l+1} = x\_l + \text{Sublayer}(\text{LayerNorm}(x\_l))$

对比：

特性	Post-Norm	Pre-Norm
训练稳定性	较差，需要学习率预热	更好，训练更稳定
梯度流动	可能梯度消失	更直接的残差连接
性能上限	可能略高	略低但更易训练
现代应用	原始Transformer	GPT、LLaMA等主流模型

4.6 位置编码PyTorch实现

import torch
import torch.nn as nn
import math

class PositionalEncoding(nn.Module):
    """
    正弦位置编码实现
    """
    def __init__(self, d_model, max_seq_len=5000, dropout=0.1):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        
        # 创建位置编码矩阵
        pe = torch.zeros(max_seq_len, d_model)
        position = torch.arange(0, max_seq_len, dtype=torch.float).unsqueeze(1)
        
        # 计算div_term
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * 
            (-math.log(10000.0) / d_model)
        )
        
        # 偶数维度使用sin，奇数维度使用cos
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        
        # 添加batch维度并注册为buffer（不参与训练）
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)
    
    def forward(self, x):
        # x: (seq_len, batch_size, d_model)
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)


class LearnablePositionalEncoding(nn.Module):
    """
    可学习位置编码实现
    """
    def __init__(self, d_model, max_seq_len=5000, dropout=0.1):
        super(LearnablePositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        self.pos_embedding = nn.Embedding(max_seq_len, d_model)
    
    def forward(self, x):
        # x: (batch_size, seq_len, d_model)
        seq_len = x.size(1)
        positions = torch.arange(seq_len, device=x.device).unsqueeze(0)
        x = x + self.pos_embedding(positions)
        return self.dropout(x)

5. Vision Transformer（ViT）

5.1 从NLP到CV的跨越

2020年Google提出Vision Transformer（ViT），首次证明纯Transformer架构在图像分类任务上可以达到甚至超越CNN的性能。

5.2 ViT核心思想

ViT将图像视为序列数据，核心步骤：

图像分块（Patch Embedding）：将图像分割成固定大小的 patches
线性投影：将每个patch映射为向量
添加位置编码：为每个patch添加位置信息
添加类别token：类似BERT的[CLS] token，用于分类
输入Transformer Encoder：标准Transformer编码器处理

5.3 图像分块详解

对于输入图像 $\in \mathbb{R}^{H \times W \times C}$ ：

将图像分割为 $\times P$ 大小的patches
序列长度： $\frac{HW}{P^2}$
每个patch展平后维度： $P2×CP^2 \times C$
通过线性投影映射到 $d\_{model}$ 维度

示例：

224x224的图像，patch大小16x16：

Patch数量： $224/16)^2 = 196$
每个patch原始维度： $16 \times 16 \times 3 = 768$
输入序列长度：196 + 1（类别token）= 197

5.4 ViT与CNN对比

特性	CNN	ViT
归纳偏置	局部性、平移等变性	较少，主要从数据学习
感受野	局部到全局（逐层扩大）	全局（自注意力直接建模）
数据效率	少量数据即可训练	需要大量预训练数据
计算复杂度	$O (n)$ ，n为像素数	$O(n^2)$ ，patch间注意力
可解释性	特征图可视化	注意力图可视化
典型模型	ResNet、EfficientNet	ViT、Swin Transformer

5.5 ViT变体进展

2024-2025年ViT领域的重要进展：

Swin Transformer：引入层次化结构和移位窗口，降低计算复杂度
DeiT（Data-efficient Image Transformer）：使用知识蒸馏，减少数据依赖
MAE（Masked Autoencoder）：自监督预训练方法，提升数据效率
DyT（Dynamic Tanh）：2024年何恺明团队提出，用动态tanh替代LayerNorm，在ViT等模型上取得更好效果

5.6 ViT PyTorch实现

import torch
import torch.nn as nn

class PatchEmbedding(nn.Module):
    """
    图像分块嵌入层
    """
    def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_dim=768):
        super(PatchEmbedding, self).__init__()
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = (img_size // patch_size) ** 2
        
        # 使用卷积实现分块和投影
        self.proj = nn.Conv2d(
            in_channels, embed_dim, 
            kernel_size=patch_size, stride=patch_size
        )
    
    def forward(self, x):
        # x: (batch, channels, height, width)
        x = self.proj(x)  # (batch, embed_dim, H/P, W/P)
        x = x.flatten(2)  # (batch, embed_dim, num_patches)
        x = x.transpose(1, 2)  # (batch, num_patches, embed_dim)
        return x


class VisionTransformer(nn.Module):
    """
    简化版Vision Transformer
    """
    def __init__(
        self, 
        img_size=224, 
        patch_size=16, 
        in_channels=3, 
        num_classes=1000,
        embed_dim=768, 
        depth=12, 
        num_heads=12, 
        mlp_ratio=4,
        dropout=0.1
    ):
        super(VisionTransformer, self).__init__()
        
        # Patch嵌入
        self.patch_embed = PatchEmbedding(
            img_size, patch_size, in_channels, embed_dim
        )
        num_patches = self.patch_embed.num_patches
        
        # 类别token
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        
        # 位置编码
        self.pos_embed = nn.Parameter(
            torch.zeros(1, num_patches + 1, embed_dim)
        )
        self.dropout = nn.Dropout(dropout)
        
        # Transformer编码器
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=embed_dim,
            nhead=num_heads,
            dim_feedforward=int(embed_dim * mlp_ratio),
            dropout=dropout,
            activation='gelu',
            batch_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=depth)
        
        # 分类头
        self.norm = nn.LayerNorm(embed_dim)
        self.head = nn.Linear(embed_dim, num_classes)
        
        # 初始化
        nn.init.normal_(self.cls_token, std=0.02)
        nn.init.normal_(self.pos_embed, std=0.02)
    
    def forward(self, x):
        batch_size = x.shape[0]
        
        # Patch嵌入
        x = self.patch_embed(x)  # (batch, num_patches, embed_dim)
        
        # 添加类别token
        cls_tokens = self.cls_token.expand(batch_size, -1, -1)
        x = torch.cat([cls_tokens, x], dim=1)  # (batch, num_patches+1, embed_dim)
        
        # 添加位置编码
        x = x + self.pos_embed
        x = self.dropout(x)
        
        # Transformer编码
        x = self.transformer(x)
        
        # 取类别token输出进行分类
        x = self.norm(x[:, 0])
        x = self.head(x)
        return x


# 测试ViT
if __name__ == "__main__":
    vit = VisionTransformer(img_size=224, patch_size=16, num_classes=10)
    x = torch.randn(2, 3, 224, 224)
    output = vit(x)
    print(f"输入形状: {x.shape}")
    print(f"输出形状: {output.shape}")

6. Transformer变体与前沿

6.1 BERT（Bidirectional Encoder Representations）

2018年Google提出，基于Transformer Encoder：

核心创新：

双向编码：同时利用左右上下文
预训练任务：
- MLM（Masked Language Model）：随机掩码15%的token进行预测
- NSP（Next Sentence Prediction）：预测两个句子是否连续

应用场景：文本分类、命名实体识别、问答系统等

6.2 GPT系列

OpenAI的生成式预训练模型系列：

模型	年份	参数量	核心特点
GPT-1	2018	1.17亿	证明无监督预训练有效性
GPT-2	2019	15亿	零样本学习能力
GPT-3	2020	1750亿	上下文学习、涌现能力
GPT-4	2023	未公开	多模态、更强推理能力
GPT-4o	2024	未公开	原生多模态、实时交互

GPT架构特点：

仅使用Decoder（自回归生成）
因果掩码（只能看到之前的位置）
从左到右的单向注意力

6.3 LLaMA系列

Meta开源的大语言模型系列：

LLaMA 2（2023）：

7B、13B、70B三个版本
上下文长度4096
采用RoPE位置编码
使用SwiGLU激活函数

LLaMA 3/4（2024-2025）：

更大的训练数据（15万亿token）
更长的上下文窗口
改进的分词器
多模态能力

6.4 Mamba与状态空间模型

2023年底提出的Mamba架构，试图解决Transformer的 $O(n^2)$ 复杂度问题：

核心思想：

基于状态空间模型（State Space Model, SSM）
线性复杂度 $O (n)$
选择性状态空间：根据输入动态调整状态转移

Mamba-2（2024）：

速度提升2-8倍
理论解释更清晰
在语言建模任务上匹敌Transformer

Mamba-3（2025）：

三大核心改进趋近完全体
在视觉任务上展现潜力
可能成为Transformer的有力替代

6.5 其他重要变体

模型	特点	应用场景
T5	Encoder-Decoder统一框架	翻译、摘要
ELECTRA	判别式预训练	高效预训练
ALBERT	参数共享、因式分解	轻量级模型
DeBERTa	解耦注意力机制	理解任务
Mixtral	稀疏混合专家（MoE）	高效推理

6.6 2024-2025研究趋势

效率优化：Mamba、RWKV等线性复杂度架构
长上下文：百万级token上下文建模
多模态融合：文本、图像、音频、视频统一建模
模型压缩：量化、剪枝、知识蒸馏
推理优化：投机解码、KV Cache优化

7. 实战案例：使用PyTorch实现简化版Transformer

7.1 完整Transformer实现

以下是一个完整的、可用于机器翻译任务的简化版Transformer实现：

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


class Transformer(nn.Module):
    """
    完整Transformer模型（用于序列到序列任务）
    """
    def __init__(
        self,
        src_vocab_size,
        tgt_vocab_size,
        d_model=512,
        num_heads=8,
        num_encoder_layers=6,
        num_decoder_layers=6,
        d_ff=2048,
        max_seq_len=5000,
        dropout=0.1
    ):
        super(Transformer, self).__init__()
        
        self.d_model = d_model
        
        # 词嵌入层
        self.src_embedding = nn.Embedding(src_vocab_size, d_model)
        self.tgt_embedding = nn.Embedding(tgt_vocab_size, d_model)
        
        # 位置编码
        self.pos_encoding = PositionalEncoding(d_model, max_seq_len, dropout)
        
        # Encoder和Decoder
        self.encoder = TransformerEncoder(
            d_model, num_heads, num_encoder_layers, d_ff, dropout
        )
        self.decoder = TransformerDecoder(
            d_model, num_heads, num_decoder_layers, d_ff, dropout
        )
        
        # 输出投影
        self.output_projection = nn.Linear(d_model, tgt_vocab_size)
        
        # 初始化参数
        self._init_parameters()
    
    def _init_parameters(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)
    
    def forward(self, src, tgt, src_mask=None, tgt_mask=None):
        # 词嵌入并缩放
        src_emb = self.src_embedding(src) * math.sqrt(self.d_model)
        tgt_emb = self.tgt_embedding(tgt) * math.sqrt(self.d_model)
        
        # 添加位置编码
        src_emb = self.pos_encoding(src_emb.transpose(0, 1)).transpose(0, 1)
        tgt_emb = self.pos_encoding(tgt_emb.transpose(0, 1)).transpose(0, 1)
        
        # Encoder
        memory = self.encoder(src_emb, src_mask)
        
        # Decoder
        output = self.decoder(tgt_emb, memory, src_mask, tgt_mask)
        
        # 投影到词表
        return self.output_projection(output)


class TransformerEncoder(nn.Module):
    """
    Transformer编码器（多层堆叠）
    """
    def __init__(self, d_model, num_heads, num_layers, d_ff, dropout):
        super(TransformerEncoder, self).__init__()
        self.layers = nn.ModuleList([
            EncoderLayer(d_model, num_heads, d_ff, dropout)
            for _ in range(num_layers)
        ])
        self.norm = nn.LayerNorm(d_model)
    
    def forward(self, x, mask=None):
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)


class EncoderLayer(nn.Module):
    """
    单个Encoder层
    """
    def __init__(self, d_model, num_heads, d_ff, dropout):
        super(EncoderLayer, self).__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads, dropout)
        self.feed_forward = PositionwiseFeedForward(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):
        # 自注意力子层（Pre-Norm结构）
        attn_output = self.self_attn(self.norm1(x), mask)
        x = x + self.dropout(attn_output)
        
        # 前馈子层
        ff_output = self.feed_forward(self.norm2(x))
        x = x + self.dropout(ff_output)
        
        return x


class TransformerDecoder(nn.Module):
    """
    Transformer解码器（多层堆叠）
    """
    def __init__(self, d_model, num_heads, num_layers, d_ff, dropout):
        super(TransformerDecoder, self).__init__()
        self.layers = nn.ModuleList([
            DecoderLayer(d_model, num_heads, d_ff, dropout)
            for _ in range(num_layers)
        ])
        self.norm = nn.LayerNorm(d_model)
    
    def forward(self, x, memory, src_mask=None, tgt_mask=None):
        for layer in self.layers:
            x = layer(x, memory, src_mask, tgt_mask)
        return self.norm(x)


class DecoderLayer(nn.Module):
    """
    单个Decoder层
    """
    def __init__(self, d_model, num_heads, d_ff, dropout):
        super(DecoderLayer, self).__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads, dropout)
        self.cross_attn = MultiHeadAttention(d_model, num_heads, dropout)
        self.feed_forward = PositionwiseFeedForward(d_model, d_ff, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x, memory, src_mask=None, tgt_mask=None):
        # 掩码自注意力
        attn_output = self.self_attn(self.norm1(x), tgt_mask)
        x = x + self.dropout(attn_output)
        
        # 交叉注意力
        cross_output = self.cross_attn(
            self.norm2(x), memory=memory, mask=src_mask
        )
        x = x + self.dropout(cross_output)
        
        # 前馈网络
        ff_output = self.feed_forward(self.norm3(x))
        x = x + self.dropout(ff_output)
        
        return x


class MultiHeadAttention(nn.Module):
    """
    多头注意力（优化版）
    """
    def __init__(self, d_model, num_heads, dropout=0.1):
        super(MultiHeadAttention, self).__init__()
        assert d_model % num_heads == 0
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        # 合并Q、K、V的线性投影
        self.qkv_proj = nn.Linear(d_model, 3 * d_model)
        self.out_proj = nn.Linear(d_model, d_model)
        
        self.dropout = nn.Dropout(dropout)
        self.scale = math.sqrt(self.d_k)
    
    def forward(self, query, memory=None, mask=None):
        batch_size = query.size(0)
        
        # 如果是自注意力，memory为None
        if memory is None:
            memory = query
        
        # 线性投影
        q = self.qkv_proj(query)[:, :, :self.d_model]
        kv = self.qkv_proj(memory)[:, :, self.d_model:]
        k, v = kv.chunk(2, dim=-1)
        
        # 分头并调整维度: (batch, heads, seq, d_k)
        q = q.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        k = k.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        v = v.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        
        # 计算注意力
        scores = torch.matmul(q, k.transpose(-2, -1)) / self.scale
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        
        attn = F.softmax(scores, dim=-1)
        attn = self.dropout(attn)
        
        # 加权求和
        context = torch.matmul(attn, v)
        
        # 拼接多头
        context = context.transpose(1, 2).contiguous().view(
            batch_size, -1, self.d_model
        )
        
        return self.out_proj(context)


class PositionwiseFeedForward(nn.Module):
    """
    位置前馈网络
    """
    def __init__(self, d_model, d_ff, dropout=0.1):
        super(PositionwiseFeedForward, self).__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)
        self.activation = nn.GELU()  # 现代Transformer常用GELU
    
    def forward(self, x):
        x = self.linear1(x)
        x = self.activation(x)
        x = self.dropout(x)
        x = self.linear2(x)
        return x


class PositionalEncoding(nn.Module):
    """
    正弦位置编码
    """
    def __init__(self, d_model, max_len=5000, dropout=0.1):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        
        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).transpose(0, 1)
        self.register_buffer('pe', pe)
    
    def forward(self, x):
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)


def generate_square_subsequent_mask(sz):
    """
    生成下三角掩码（用于Decoder的自回归特性）
    """
    mask = torch.triu(torch.ones(sz, sz), diagonal=1)
    mask = mask.masked_fill(mask == 1, float('-inf'))
    return mask


# 测试完整Transformer
if __name__ == "__main__":
    # 超参数
    src_vocab_size = 10000
    tgt_vocab_size = 10000
    batch_size = 2
    src_seq_len = 20
    tgt_seq_len = 15
    
    # 创建模型
    model = Transformer(
        src_vocab_size=src_vocab_size,
        tgt_vocab_size=tgt_vocab_size,
        d_model=512,
        num_heads=8,
        num_encoder_layers=6,
        num_decoder_layers=6,
        d_ff=2048,
        dropout=0.1
    )
    
    # 模拟输入
    src = torch.randint(0, src_vocab_size, (batch_size, src_seq_len))
    tgt = torch.randint(0, tgt_vocab_size, (batch_size, tgt_seq_len))
    
    # 生成掩码
    tgt_mask = generate_square_subsequent_mask(tgt_seq_len)
    
    # 前向传播
    output = model(src, tgt, tgt_mask=tgt_mask)
    
    print(f"源序列形状: {src.shape}")
    print(f"目标序列形状: {tgt.shape}")
    print(f"输出形状: {output.shape}")
    print(f"模型参数量: {sum(p.numel() for p in model.parameters()):,}")

7.2 训练流程示例

def train_transformer(model, train_loader, optimizer, criterion, device):
    """
    Transformer训练函数
    """
    model.train()
    total_loss = 0
    
    for batch_idx, (src, tgt) in enumerate(train_loader):
        src, tgt = src.to(device), tgt.to(device)
        
        # 准备输入和标签（偏移一位用于自回归训练）
        tgt_input = tgt[:, :-1]
        tgt_label = tgt[:, 1:]
        
        # 生成掩码
        tgt_mask = generate_square_subsequent_mask(tgt_input.size(1)).to(device)
        
        # 前向传播
        optimizer.zero_grad()
        output = model(src, tgt_input, tgt_mask=tgt_mask)
        
        # 计算损失
        loss = criterion(output.view(-1, output.size(-1)), tgt_label.view(-1))
        
        # 反向传播
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()
        
        total_loss += loss.item()
    
    return total_loss / len(train_loader)

8. 避坑小贴士

8.1 注意力计算常见错误

错误1：忘记缩放因子

# 错误：没有除以sqrt(d_k)
scores = torch.matmul(Q, K.transpose(-2, -1))  # 数值过大导致softmax梯度消失

# 正确：添加缩放
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)

错误2：掩码应用位置错误

# 错误：在softmax之后应用掩码
attn_weights = F.softmax(scores, dim=-1)
attn_weights = attn_weights.masked_fill(mask == 0, 0)  # 太晚！

# 正确：在softmax之前应用掩码
scores = scores.masked_fill(mask == 0, -1e9)
attn_weights = F.softmax(scores, dim=-1)

8.2 位置编码注意事项

问题：位置编码维度不匹配

确保位置编码的维度与词嵌入维度完全一致，且正确进行广播：

# 确保形状兼容
# pe: (max_len, 1, d_model)
# x: (seq_len, batch_size, d_model)
x = x + self.pe[:x.size(0), :]  # 正确

8.3 训练稳定性建议

学习率预热：Transformer对初始学习率敏感，建议使用warmup策略
梯度裁剪：防止梯度爆炸，设置max_norm=1.0
标签平滑：使用label smoothing（eps=0.1）提升泛化能力
Dropout位置：在注意力权重和残差连接后都应用dropout

8.4 内存优化技巧

技巧	说明	适用场景
梯度检查点	用计算换内存	显存不足时
混合精度训练	FP16/BF16加速	支持Tensor Core的GPU
梯度累积	模拟大批量训练	显存受限
动态批处理	按长度分组	序列长度差异大

9. 本章小结

本章深入探讨了注意力机制与Transformer架构，核心要点总结如下：

核心概念：

注意力机制通过Query、Key、Value计算，实现动态信息筛选
自注意力让序列中每个位置都能关注所有位置
多头注意力并行关注不同子空间信息

Transformer架构：

Encoder-Decoder结构，完全基于注意力机制
位置编码提供序列顺序信息
残差连接和层归一化稳定深层网络训练

前沿发展：

Vision Transformer将Transformer引入计算机视觉
BERT、GPT系列推动NLP进入大模型时代
Mamba等线性复杂度架构挑战Transformer地位

实践要点：

注意缩放因子、掩码应用等实现细节
合理选择Pre-Norm或Post-Norm
使用学习率预热等技巧稳定训练

系列导航

上一篇：[【机器学习精通】第8章 | 循环神经网络与LSTM：序列建模的艺术](file:///d:/python/Python全栈开发/机器学习精通系列/第08章-循环神经网络与LSTM.md)

下一篇：[【机器学习精通】第10章 | 生成模型：VAE与扩散模型](file:///d:/python/Python全栈开发/机器学习精通系列/第10章-生成模型.md)

本文是《机器学习精通》系列教程的第11章，专注深度学习核心技术的系统学习。如有疑问，欢迎在评论区交流讨论。

AtomGit开源社区

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念，把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起，为开发者提供从开发、训练到部署的一站式体验。

更多推荐

告诉大家OpenClaw 社区如此重视 Skill 生态的原因

AtomGit开源社区

Spring AI Alibaba 多模态全家桶：图片理解、图片生成与语音合成实战

实际业务中通常需要指定图片尺寸和生成数量，通过.withN(count) // 生成几张.withWidth(width) // 宽度.withHeight(height) // 高度.build()));DIAGRAM: 图片生成调用时序图——客户端 → Controller → ImageModel.call() → DashScope 提交任务 → 轮询状态 → 返回图片 URL能力核心类模