AIAttnStudio

深入理解 3D Gaussian Splatting (3DGS)：从原理到实践

Thu, 16 Jul 2026 00:00:00 GMT

1. 引言：从 NeRF 到 3DGS#

1.1 什么是新视角合成？#

新视角合成（Novel View Synthesis, NVS） 是计算机视觉与图形学的经典任务：给定一个场景的多张已知视角图像，生成从任意新视角观察该场景的图像。

它在很多领域都有重要应用：

VR/AR 内容生成：构建沉浸式三维场景
数字孪生：工业仿真、城市建模
影视特效：实拍场景的三维重建与编辑
文化遗产：文物、遗址的高保真数字化
自动驾驶：场景重建与仿真

过去几年，神经辐射场（NeRF） 是这个领域最耀眼的明星，但它的训练和渲染速度始终是痛点。直到 2023 年，3D Gaussian Splatting (3DGS) 横空出世，一举解决了”高质量 + 实时”这一对看似矛盾的需求。

1.2 3D Gaussian Splatting 是什么？#

3D Gaussian Splatting (3DGS) 是 Kerbl 等人在 SIGGRAPH 2023 上提出的革命性方法，论文标题为 “3D Gaussian Splatting for Real-Time Radiance Field Rendering”。

它的核心思想非常优雅：

不再用隐式的神经网络表示场景，而是用一组显式的、各向异性的 3D 高斯椭球作为基本渲染单元，并通过可微的光栅化器将它们”泼溅”到图像平面上。

这里的 Splatting（泼溅） 是关键——它指的是将 3D 的高斯椭球投影到 2D 图像平面（投影后仍是高斯分布），再进行 alpha 合成。

1.3 为什么 3DGS 能引起轰动？#

指标	NeRF (原始)	Instant-NGP	3DGS
训练时间	数小时~数天	几分钟	~30-60 分钟
渲染速度	极慢（逐光线采样）	~30 FPS	30+ FPS @ 1080p
显存占用	高	中等	与场景规模线性相关
显式编辑	困难	困难	天然支持（高斯可编辑）

3DGS 在保持高质量重建的同时，把渲染速度提升到了实时级别，并且场景表示是显式的、可编辑的。

2. 为什么选择高斯椭球？#

2.1 表示方法的演进#

在讨论 3DGS 之前，我们先回顾一下场景表示方法的演进：

方法	表示形式	优点	缺点
点云（Point Cloud）	离散点 + 颜色	简单、易获取	无体素、稀疏、不连续
体素网格（Voxel Grid）	规则三维网格	规则、易采样	内存爆炸、分辨率受限
Mesh（网格）	顶点 + 面片	高效渲染、纹理清晰	难以从图像重建
NeRF	MLP（隐式）	连续表示、高质量	训练慢、渲染慢、难编辑
3DGS	3D 高斯椭球	实时、显式、可编辑	场景文件较大

2.2 高斯椭球的直观理解#

为什么是”椭球”而不是”球”？

各向异性（Anisotropy） 是关键。想象一个真实的物体表面：

细长的树枝：沿长度方向延伸，垂直方向较窄 → 需要细长椭球
平坦的墙面：几乎是一个平面 → 需要极扁的椭球（接近圆盘）
球形果实：各方向尺度相当 → 球形高斯即可

如果只能用球形高斯来表示这些结构，就需要数量巨大的小球，效率极低。椭球可以用更少的元素表达更丰富的几何形状，这是 3DGS 高效性的基础。

2.3 与点云的关系#

3DGS 可以看作是一种”升级版”的点云：

点云：每个点是一个 (x, y, z, color) 的零维样本
3DGS：每个”点”是一个有位置、形状（协方差矩阵）、颜色、不透明度的三维概率分布

这种”带体积”和”带朝向”的特性，使得 3DGS 既能表达几何，也能表达外观，并且渲染时不会有”空洞”。

3. 数学基础：3D 高斯椭球的表示#

3.1 单个高斯的参数化#

场景中第 iii 个高斯椭球由以下参数定义：

Gi(x)=e−12(x−μi)TΣi−1(x−μi)G_i(\mathbf{x}) = e^{-\frac{1}{2}(\mathbf{x} - \boldsymbol{\mu}_i)^T \boldsymbol{\Sigma}_i^{-1} (\mathbf{x} - \boldsymbol{\mu}_i)}Gi(x)=e−21(x−μi)TΣi−1(x−μi)

其中：

μi∈R3\boldsymbol{\mu}_i \in \mathbb{R}^3μi∈R3：均值向量（高斯的中心位置）
Σi∈R3×3\boldsymbol{\Sigma}_i \in \mathbb{R}^{3\times 3}Σi∈R3×3：协方差矩阵（决定椭球的形状与朝向）

为了保证协方差矩阵的半正定性（即数学上构成有效的椭球），3DGS 采用了巧妙的分解方式：

Σi=RiSiSiTRiT\boldsymbol{\Sigma}_i = \mathbf{R}_i \mathbf{S}_i \mathbf{S}_i^T \mathbf{R}_i^TΣi=RiSiSiTRiT

其中：

Ri\mathbf{R}_iRi：从四元数 qiq_iqi 构造的旋转矩阵
Si=diag(si,x,si,y,si,z)\mathbf{S}_i = \text{diag}(s_{i,x}, s_{i,y}, s_{i,z})Si=diag(si,x,si,y,si,z)：缩放矩阵（三个轴方向的尺度）

3.2 各向异性与各向同性#

|| 概念 | 协方差矩阵 | 视觉效果 | ||------|------------|----------| || 各向同性（Isotropic） | σ2I\sigma^2 \mathbf{I}σ2I（球形） | 球状云团 | || 正交各向异性（Orthotropic） | diag(σx2,σy2,σz2)\text{diag}(\sigma_x^2, \sigma_y^2, \sigma_z^2)diag(σx2,σy2,σz2) | 沿坐标轴对齐的椭球 | || 完全各向异性（Anisotropic） | RSSTRT\mathbf{R} \mathbf{S} \mathbf{S}^T \mathbf{R}^TRSSTRT（任意朝向） | 任意朝向的椭球（3DGS 使用） |

3DGS 选择了完全各向异性的椭球，这让每个高斯都能精确地贴合局部几何结构。

这里是付费内容

3.3 完整参数列表#

每个高斯椭球 GiG_iGi 的参数总结：

参数	符号	维度	含义	优化方式
位置	μi\boldsymbol{\mu}_iμi	3	高斯的中心	直接优化（梯度下降）
旋转	qiq_iqi	4（四元数）	椭球的朝向	直接优化 + 归一化
缩放	si\mathbf{s}_isi	3	三轴尺度	通过 exp⁡(⋅)\exp(\cdot)exp(⋅) 保证正值
不透明度	αi\alpha_iαi	1	椭球的”密度”	通过 sigmoid(⋅)\text{sigmoid}(\cdot)sigmoid(⋅) 约束到 [0,1][0, 1][0,1]
球谐系数	SHi\mathbf{SH}_iSHi	3×(k+1)23 \times (k+1)^23×(k+1)2	视角相关颜色	直接优化

球谐函数（Spherical Harmonics, SH） 用于表示视角相关的辐射场——这就是为什么 3DGS 能够正确渲染出镜面反射等视角相关效果，而不仅仅是漫反射。

3.4 视角相关颜色的直观理解#

传统点云的颜色是固定的（RGB），但 3DGS 引入球谐后可以表达：

漫反射：低阶球谐即可表达
高光 / 镜面反射：需要高阶球谐（视角变化时颜色会变）

例如，看一个金属球时，正面和侧面看到的颜色亮度不同——这正是 SH 表达的能力。

付费内容

此部分内容需要赞助后解锁

去爱发电赞助

4. 核心创新：可微 3D 高斯光栅化#

4.1 传统光栅化与可微光栅化#

要把 3D 椭球渲染成 2D 图像，需要解决两个问题：

可见性（Visibility）：从当前相机视角，哪些高斯是可见的？
可微性（Differentiability）：渲染结果对每个高斯的参数都是可导的，这样才能用梯度下降优化。

3DGS 的最大贡献之一，是第一个高效的可微 3D 高斯光栅化器。

4.2 椭球的 2D 投影#

给定一个 3D 高斯椭球（μ,Σ\boldsymbol{\mu}, \boldsymbol{\Sigma}μ,Σ）和一个相机（投影矩阵 W\mathbf{W}W），其 2D 投影仍然是高斯分布。

具体来说，2D 投影后的协方差矩阵为：

Σ′=JWΣWTJT\boldsymbol{\Sigma}' = \mathbf{J} \mathbf{W} \boldsymbol{\Sigma} \mathbf{W}^T \mathbf{J}^TΣ′=JWΣWTJT

其中：

W\mathbf{W}W：世界到相机的变换矩阵的 3×33\times 33×3 部分
J\mathbf{J}J：相机投影函数的雅可比矩阵（Jacobian）

直观上：

3D 椭球投影到 2D 图像上仍然是椭球，只是形状和大小随视角变化（近大远小）。

4.3 分块排序（Tile-Based Rasterization）#

3DGS 的渲染器采用分块（Tile）并行策略：

将屏幕划分为若干小块（如 16×1616 \times 1616×16 像素）
视锥剔除（Frustum Culling）：剔除完全不在视锥内的高斯
逐高斯处理：对每个高斯，计算它覆盖哪些 tile
将高斯分配到对应 tile 的列表中（同一高斯可能分配到多个 tile）
对每个 tile 内部的高斯按深度排序
alpha 合成（Alpha Compositing）：在每个 tile 内从前向后合成像素颜色

这种策略既保证了渲染速度，又避免了对所有高斯做全局排序的代价。

4.4 Alpha 合成公式#

最终的像素颜色 CCC 由所有覆盖该像素的高斯按从前到后的顺序合成：

C=∑i=1NTiαici,Ti=∏j=1i−1(1−αj)C = \sum_{i=1}^{N} T_i \alpha_i \mathbf{c}_i, \quad T_i = \prod_{j=1}^{i-1} (1 - \alpha_j)C=i=1∑NTiαici,Ti=j=1∏i−1(1−αj)

其中：

ci\mathbf{c}_ici：第 iii 个高斯的颜色（由 SH 和视角决定）
αi\alpha_iαi：该高斯对当前像素的不透明度贡献
TiT_iTi：累积透射率（到第 iii 个高斯时，光线尚未被遮挡的比例）

注意：这里 TiT_iTi 的计算必须从前往后，但 αi\alpha_iαi 的梯度计算则依赖从前向后的链式求导——这与 NeRF 体渲染的”从后往前”积分有本质区别。

4.5 简化代码示意#

虽然完整的光栅化器需要数百行 CUDA 代码，但核心逻辑可简化为：

1
import torch
2

3
def render_pixel(tile_gaussians, pixel_uv):
4
    """
5
    tile_gaussians: 当前 tile 覆盖的高斯（按深度排序）
6
    pixel_uv: 目标像素坐标
7
    """
8
    color = torch.zeros(3, device=pixel_uv.device)
9
    transmittance = torch.ones(1, device=pixel_uv.device)
10

11
    for gauss in tile_gaussians:
12
        # 2D 高斯在该像素处的不透明度
13
        alpha = compute_2d_alpha(gauss, pixel_uv)
14

15
        # alpha 合成
16
        color += transmittance * alpha * gauss.color
17
        transmittance *= (1.0 - alpha)
18

19
        # 早期终止：透射率足够小
20
        if transmittance.item() < 1e-4:
21
            break
22

23
    return color

展开收起

5. 自适应密度控制（Adaptive Density Control）#

5.1 为什么需要密度控制？#

初始化时，3DGS 通常使用稀疏点云（来自 COLMAP 或随机采样）作为高斯中心。如果高斯数量太少，几何不完整；太多，又会导致冗余和浪费。3DGS 在训练过程中动态增减高斯，这是其高效性的关键之一。

5.2 克隆（Cloning / Densification）#

对于欠重建区域（几何细节缺失），即高斯的位置梯度大但形状扁平：

克隆：在该高斯周围复制一份，沿位置梯度方向略微移动

适合场景：欠重建的小细节（如物体的微小结构）

5.3 分裂（Splitting）#

对于过度重建区域，即一个高斯覆盖了过大的范围：

分裂：将一个大高斯替换为两个稍小的高斯，沿着最大缩放方向

适合场景：大而扁平的高斯需要更细粒度的表达

5.4 周期性修剪（Periodic Pruning）#

每经过若干次迭代，还需要做剪枝：

删除不透明度极低的高斯（透明，几乎不可见）
删除在训练视角下体积过大的高斯（通常是漂浮物）

5.5 密度控制伪代码#

1
def adaptive_density_control(gaussians, optimizers, viewspace_grad, t):
2
    if t % 1000 != 0:
3
        return gaussians
4

5
    # 提取平均位置梯度
6
    avg_grad = viewspace_grad.norm(dim=-1).mean(dim=0)  # (N,)
7

8
    # 判断需要克隆还是分裂
9
    mask_clone = (avg_grad > GRAD_THRESH) & (gaussians.scale.max(dim=1) < SCALE_THRESH)
10
    mask_split = (avg_grad > GRAD_THRESH) & (gaussians.scale.max(dim=1) >= SCALE_THRESH)
11

12
    # 克隆
13
    gaussians.clone(mask_clone)
14
    # 分裂
15
    gaussians.split(mask_split)
16
    # 剪枝
17
    gaussians.prune(opacity < 0.005)
18

19
    # 重置优化器状态
20
    reset_optimizer(optimizers, gaussians)
21
    return gaussians

展开收起

6. 训练流程#

6.1 整体流程图#

3DGS 的训练过程是一个迭代优化：

1
输入图像 + COLMAP 标定
2
        ↓
3
   初始化点云 → 初始化高斯椭球
4
        ↓
5
   ┌────────────────────────────┐
6
   │  迭代优化 (30000 次)        │
7
   │  ├─ 渲染当前视角            │
8
   │  ├─ 计算损失 L1 + SSIM     │
9
   │  ├─ 反向传播                │
10
   │  ├─ Adam 优化器更新参数     │
11
   │  └─ 周期性自适应密度控制    │
12
   └────────────────────────────┘
13
        ↓
14
   输出: 优化后的高斯集合 + 球谐系数
15
        ↓
16
   实时渲染（1080p @ 30+ FPS）

展开收起

6.2 损失函数#

3DGS 使用一个组合损失：

L=(1−λ)L1+λLSSIM\mathcal{L} = (1 - \lambda) \mathcal{L}_1 + \lambda \mathcal{L}_{\text{SSIM}}L=(1−λ)L1+λLSSIM

其中：

L1=∥Irendered−Igt∥1\mathcal{L}_1 = \| I_{\text{rendered}} - I_{\text{gt}} \|_1L1=∥Irendered−Igt∥1：像素级 L1 损失
LSSIM\mathcal{L}_{\text{SSIM}}LSSIM：结构相似性损失（SSIM）
λ\lambdaλ：通常设为 0.2

SSIM 损失 让 3DGS 在保持锐利边缘的同时，获得更好的整体结构相似性。如果只用 L1，结果会偏模糊。

6.3 优化器与学习率#

3DGS 使用 Adam 优化器，但不同参数有不同的学习率：

参数	学习率	说明
位置 μ\boldsymbol{\mu}μ	1.6×10−41.6 \times 10^{-4}1.6×10−4	较慢，避免位置漂移
旋转 qqq	1.0×10−31.0 \times 10^{-3}1.0×10−3	较快，配合归一化
缩放 s\mathbf{s}s	5.0×10−35.0 \times 10^{-3}5.0×10−3	通过 exp⁡(⋅)\exp(\cdot)exp(⋅) 变换
不透明度 α\alphaα	5.0×10−25.0 \times 10^{-2}5.0×10−2	较快，便于剪枝
球谐系数	2.5×10−32.5 \times 10^{-3}2.5×10−3	适度
球谐阶数	每 1000 步提升 1 阶	关键：1→2→3，渐进式

渐进式球谐阶数提升 是个精妙的设计：先用低阶（漫反射）稳定几何，再逐渐引入高阶（视角相关效果），避免训练早期因高阶球谐不稳定而导致的退化。

6.4 典型训练时间#

场景规模	高斯数量	训练时间（单 GPU）
简单物体	~50k	~10 分钟
一般场景	100k~1M	30~60 分钟
大型场景	1M~5M	1~2 小时
城市级别	10M+	数小时

3DGS 在一张 RTX 4090 上即可训练大多数场景，是 NeRF 难以想象的便利。

7. 完整示例：基于 gsplat 库的训练#

7.1 环境准备#

Terminal window

1
pip install gsplat torch torchvision
2
# 或
3
pip install "gsplat[all]"

7.2 核心训练代码#

1
import torch
2
from gsplat import rasterization
3
from torch import nn
4
from torch.optim import Adam
5

6
class GaussianModel(nn.Module):
7
    def __init__(self, points, colors):
8
        super().__init__()
9
        N = points.shape[0]
10
        # 位置
11
        self.means = nn.Parameter(points)
12
        # 缩放（取 log 保证正值）
13
        self.scales = nn.Parameter(torch.full((N, 3), -3.0))
14
        # 旋转（四元数）
15
        self.quats = nn.Parameter(torch.zeros(N, 4))
16
        self.quats.data[:, 0] = 1.0  # 初始化为单位四元数
17
        # 不透明度（取 logit）
18
        self.opacities = nn.Parameter(torch.full((N,), 0.5).logit())
19
        # 球谐系数（这里只用 0 阶，即 DC 颜色）
20
        self.colors = nn.Parameter(colors)
21

22
    def forward(self, camera):
23
        # 调用可微光栅化器
24
        out_img, out_alpha, _ = rasterization(
25
            means=self.means,
26
            quats=self.quats,
27
            scales=torch.exp(self.scales),
28
            opacities=torch.sigmoid(self.opacities),
29
            colors=self.colors,
30
            viewmats=camera.viewmat,
31
            Ks=camera.K,
32
            width=camera.W,
33
            height=camera.H,
34
        )
35
        return out_img

展开收起

7.3 训练循环#

1
import torch.nn.functional as F
2
from torchmetrics.image import StructuralSimilarityIndexMeasure
3

4
ssim = StructuralSimilarityIndexMeasure(data_range=1.0)
5

6
def train_step(model, optimizer, view):
7
    rendered = model(view)        # (1, H, W, 3)
8
    gt = view.image.unsqueeze(0)  # (1, H, W, 3)
9

10
    l1 = F.l1_loss(rendered, gt)
11
    loss_ssim = 1 - ssim(rendered.permute(0, 3, 1, 2), gt.permute(0, 3, 1, 2))
12
    loss = 0.8 * l1 + 0.2 * loss_ssim
13

14
    optimizer.zero_grad()
15
    loss.backward()
16
    optimizer.step()
17
    return loss.item()

展开收起

8. 3DGS vs NeRF：深度对比#

8.1 本质区别#

维度	NeRF	3DGS
表示方式	隐式（MLP）	显式（3D 高斯集合）
渲染方式	沿光线采样	高斯泼溅 + alpha 合成
训练	一次 MLP 训练	逐高斯参数优化
渲染速度	慢（每像素 ~100 次采样）	快（无采样）
显存占用	与图像分辨率相关	与高斯数量相关
场景编辑	困难	天然支持（直接操作高斯）

8.2 性能基准#

在 Mip-NeRF 360 基准测试上的典型结果：

方法	PSNR ↑	SSIM ↑	LPIPS ↓	FPS ↑
NeRF	27.21	0.749	0.229	<1
Instant-NGP	28.18	0.796	0.207	~30
Mip-NeRF 360	29.13	0.844	0.184	<1
3DGS	30.41	0.870	0.149	~120

3DGS 在所有指标上都优于之前的 SOTA，并且渲染速度提升了 100 倍以上。

8.3 各自的适用场景#

场景	推荐方法	理由
静态场景实时渲染	3DGS	速度与质量兼得
大规模 / 城市级重建	3DGS + 分块	高斯可分布式管理
极稀疏视角（< 10 张）	NeRF 变体	3DGS 难以稳定
动态场景 / 4D	4D-GS / DyNeRF	需要时序建模
神经语义编辑	NeRF + Instruct-Nerf2NeRF	隐式表示便于编辑
物理仿真 / 碰撞	3DGS	显式几何可直接用于物理
场景重建 + 物理仿真	PhysGaussian	3DGS 的显式性优势

9. 进阶主题#

9.1 动态场景：4D-GS#

将时间维度引入，每个高斯带时间相关的形变场，可以重建动态场景：

μi(t)=μi+Δθ(μi,t)\boldsymbol{\mu}_i(t) = \boldsymbol{\mu}_i + \boldsymbol{\Delta}_\theta(\boldsymbol{\mu}_i, t)μi(t)=μi+Δθ(μi,t)

代表工作：4D-GS、Deformable 3DGS、Dynamic 3D Gaussians。

9.2 大规模场景：VastGaussian#

面对城市级大场景，单卡训练困难。VastGaussian 提出分区域训练 + 全局合并策略：

将场景分为多个子区域
每个子区域独立训练 3DGS
用”代理高斯”（progressive merging）合并子区域结果

9.3 物理仿真：PhysGaussian#

3DGS 是显式几何，可以直接用于物理仿真。PhysGaussian 实现了：

通过高斯密度场恢复连续体积密度 ρ(x)\rho(\mathbf{x})ρ(x)
在仿真器（如 Material Point Method, MPM）中代入 ρ\rhoρ 进行物理求解
可以模拟切割、碰撞、形变等交互效果

9.4 编辑：GaussianEditor / Instruct-GS2GS#

利用 InstructPix2Pix 或 Segment Anything (SAM) 等 2D 基础模型，引导 3DGS 的局部编辑：

文本指令：“把这个杯子变成红色”
选择性替换部分高斯的颜色 / 球谐系数

9.5 压缩：LightGaussian / CompGS#

3DGS 场景文件通常较大（数百 MB 到 GB）。压缩方向：

剪枝冗余高斯
球哈系数低位量化
矢量量化 + 熵编码
网格化封装

代表方法可在保持 95%+ 质量的同时，将模型缩小 10~25 倍。

10. 关键数学公式汇总#

为方便查阅，下面列出本文涉及的核心公式：

3D 高斯椭球：

协方差分解：

Σi=RiSiSiTRiT\boldsymbol{\Sigma}_i = \mathbf{R}_i \mathbf{S}_i \mathbf{S}_i^T \mathbf{R}_i^TΣi=RiSiSiTRiT

2D 投影协方差：

Σ′=JWΣWTJT\boldsymbol{\Sigma}' = \mathbf{J} \mathbf{W} \boldsymbol{\Sigma} \mathbf{W}^T \mathbf{J}^TΣ′=JWΣWTJT

Alpha 合成：

C=∑i=1NTiαici,Ti=∏j=1i−1(1−αj)C = \sum_{i=1}^{N} T_i \alpha_i \mathbf{c}_i, \quad T_i = \prod_{j=1}^{i-1} (1 - \alpha_j)C=i=1∑NTiαici,Ti=j=1∏i−1(1−αj)

总损失：

L=(1−λ)L1+λLSSIM\mathcal{L} = (1 - \lambda) \mathcal{L}_1 + \lambda \mathcal{L}_{\text{SSIM}}L=(1−λ)L1+λLSSIM

参数激活函数：

Σ=R(q)diag(exp⁡(s))2R(q)T,α=σ(logit)\boldsymbol{\Sigma} = \mathbf{R}(q) \text{diag}(\exp(\mathbf{s}))^2 \mathbf{R}(q)^T, \quad \alpha = \sigma(\text{logit})Σ=R(q)diag(exp(s))2R(q)T,α=σ(logit)

11. 总结与展望#

3DGS 的核心贡献#

显式表示：用一组 3D 高斯椭球作为基本渲染单元，每个椭球带有位置、形状、颜色、视角相关辐射
可微光栅化器：第一个支持高效反向传播的 3D 高斯光栅化器
自适应密度控制：训练过程中动态增删高斯，使场景表示既紧凑又完整
实时高质量：在 RTX 4090 上 1080p 实时渲染，PSNR 优于 NeRF

当前挑战#

场景文件较大：动辄数百 MB
稀疏视角不稳定：训练图像太少时容易过拟合
动态 / 4D 建模：仍处于早期阶段
编辑可控性：相比 NeRF 更可控，但工业级工具链尚不成熟

未来方向#

方向	代表工作 / 思路
动态 4D 重建	4D-GS、Dynamic 3DGS
大规模场景	VastGaussian、城市级分块训练
物理交互	PhysGaussian、GaussianSim
AI 编辑	GaussianEditor、Instruct-GS2GS
模型压缩	LightGaussian、CompGS、剪枝量化
端到端应用	自动驾驶仿真、机器人具身、重光照、姿态估计

3D Gaussian Splatting 不仅是一项渲染技术，它正在成为新一代三维内容创作的”基础设施”——像 LLM 之于文本、Stable Diffusion 之于图像一样，重新定义我们生成、编辑、交互三维场景的方式。

参考资料#

Kerbl, B., Kopanas, G., Leimkühler, T., & Drettakis, G. (2023). “3D Gaussian Splatting for Real-Time Radiance Field Rendering.” ACM Transactions on Graphics (SIGGRAPH).
Mildenhall, B., et al. (2020). “NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis.” ECCV.
Müller, T., et al. (2022). “Instant Neural Graphics Primitives with a Multiresolution Hash Encoding.” ACM Transactions on Graphics (SIGGRAPH).
Barron, J. T., et al. (2022). “Mip-NeRF 360: Unbounded Anti-Aliased Neural Radiance Fields.” CVPR.
Wu, G., et al. (2024). “4D Gaussian Splatting for Real-Time Dynamic Scene Rendering.” CVPR.
Lin, J., et al. (2024). “VastGaussian: Vast 3D Gaussians for Large Scene Reconstruction.” CVPR.
Xie, T., et al. (2024). “PhysGaussian: Physics-Integrated 3D Gaussians for Generative Dynamics.” CVPR.
Ye, M., et al. (2024). “GaussianEditor: Editing 3D Gaussians with Text Prompts.” CVPR.
Fan, Z., et al. (2024). “LightGaussian: Unbounded 3D Gaussian Compression with 15x Reduction.” CVPR.
Kerbl, B., et al. (2024). “Hierarchical 3D Gaussian Splatting (HGS).” ACM TOG.
Zwicker, M., Pfister, H., van Baar, J., & Gross, M. (2001). “EWA Splatting.” IEEE TVCG — 3DGS 的数学前身。
“gsplat: An Open-Source Library for Gaussian Splatting.” GitHub: nerfstudio-project/gsplat.

深度解析 ActionMesh：Temporal 3D Diffusion 实现动画 3D 网格生成

Thu, 16 Jul 2026 00:00:00 GMT

1. 引言：从视频到动画 3D 网格的挑战#

1.1 动画 3D 内容的应用价值#

自动从简单输入生成动画 3D 物体在以下领域至关重要：

应用领域	具体场景
电子游戏	角色动画、NPC 行为
动画电影与广告	3D 资产快速生成
增强/虚拟现实	动态物体交互
机器人仿真	运动模仿与迁移

1.2 现有方法的三大局限#

局限	问题描述	影响
模态受限	只能处理特定输入（如单目视频）和特定类别（如双足动物）	泛化能力差
速度慢	需要 30~45 分钟的逐场景优化	无法快速迭代
质量不足	输出不符合生产标准	无法直接用于工业流程

1.3 ActionMesh 的核心思想#

ActionMesh（CVPR 2026）的核心洞察：

将现有 3D 扩散模型扩展加入时间轴（Temporal Axis），实现”Temporal 3D Diffusion”，从而将 3D 生成与动画预测解耦。

关键创新：

两阶段管线：先生成时序一致的独立 3D 形状，再通过时序 3D 自编码器转换为固定拓扑的动画
无需 rig：不依赖骨骼绑定系统，复杂物体也能直接动画化
拓扑一致：所有帧共享相同网格拓扑，便于纹理映射和重定向
10 倍加速：3 分钟完成 16 帧视频的动画网格生成

2. 核心问题定义与术语#

2.1 三种 3D 表示#

术语	定义	数学表示
3D Mesh	三角网格	M=(V,F)M = (V, F)M=(V,F)，其中 V∈RNv×3V \in \mathbb{R}^{N_v \times 3}V∈RNv×3 为顶点位置，F∈{1,…,Nv}Nf×3F \in \{1, \ldots, N_v\}^{N_f \times 3}F∈{1,…,Nv}Nf×3 为面连接
4D Mesh	时序变化但拓扑独立的网格序列	{(Vk,Fk)}k=1N\{(V_k, F_k)\}_{k=1}^{N}{(Vk,Fk)}k=1N
Animated 3D Mesh	拓扑一致的动画网格	{(Vk,F)}k=1N\{(V_k, F)\}_{k=1}^{N}{(Vk,F)}k=1N

2.2 任务目标#

给定输入视频 {Ik}k=1N\{I_k\}_{k=1}^{N}{Ik}k=1N（其中 Ik∈RH×W×3I_k \in \mathbb{R}^{H \times W \times 3}Ik∈RH×W×3），预测与视频对应的 Animated 3D Mesh。

具体目标：预测参考网格 M=(V,F)M = (V, F)M=(V,F) 及其在各帧的顶点更新 {Vk}\{V_k\}{Vk}，使得 VkV_kVk 表示 MMM 在时间步 tkt_ktk 与视频 IkI_kIk 匹配的新顶点位置。

3. 技术背景：3DShape2VecSet#

3.1 VecSet 架构概述#

ActionMesh 基于 3DShape2VecSet 的潜在扩散框架，该框架包含两阶段：

1
┌─────────────────────────────────────────┐
2
│ Stage 1: 3D Variational Autoencoder     │
3
│                                         │
4
│ Encoder E3D ──→ Latent z ──→ Decoder D3D │
5
│                                         │
6
│  • 编码器：点集 → 潜在向量集              │
7
│  • 解码器：潜在向量 → SDF/占用率          │
8
│  • Marching Cubes → 3D Mesh            │
9
└─────────────────────────────────────────┘
10
               ↓
11
┌─────────────────────────────────────────┐
12
│ Stage 2: 3D Latent Diffusion Model      │
13
│                                         │
14
│  Conditioning (DINOv2) + 3D Latent     │
15
│         ↓                               │
16
│  DiT (Decoder-only Transformer)         │
17
│         ↓                               │
18
│  Denoised 3D Latent                     │
19
└─────────────────────────────────────────┘

展开收起

3.2 VecSet 核心机制#

点集采样：在网格表面采样稠密点 PPP
Query 向量：学习或从 PPP 下采样的稀疏查询向量 QQQ
交叉注意力：用 PPP 作为上下文编码 QQQ
自注意力：多层自注意力处理生成低维潜在向量 zzz
解码：解码器处理 zzz 并通过交叉注意力预测查询点的 SDF/占用率

3.3 ActionMesh 的基础 Backbone#

ActionMesh 采用 TripoSG 作为 backbone（也可替换为 Craftsman）：

特性	TripoSG
生成范式	Rectified Flow
时间步嵌入	Fourier 嵌入后拼接为额外 token
Conditioning	Frozen DINOv2
开源	是（MIT）

4. Stage I：Temporal 3D Diffusion#

4.1 朴素方法的失败#

朴素方法：对视频每一帧独立应用图像转 3D 生成器。

失败原因：

3D 方向不一致
几何细节闪烁
缺乏跨帧一致性机制

4.2 核心创新：膨胀注意力（Inflated Attention）#

灵感来源：MVDream（多视角生成）→ 扩展为时序维度

核心公式：

infattn(X)=reshape−1(selfattn(reshape(X)))\text{infattn}(X) = \text{reshape}^{-1}(\text{selfattn}(\text{reshape}(X)))infattn(X)=reshape−1(selfattn(reshape(X)))

其中：

X∈RN×T×DX \in \mathbb{R}^{N \times T \times D}X∈RN×T×D：TTT 个 token，维度 DDD，对应 NNN 帧
reshape\text{reshape}reshape：将前两维展平为 1×NT1 \times NT1×NT
selfattn\text{selfattn}selfattn：标准自注意力
reshape−1\text{reshape}^{-1}reshape−1：恢复为 N×T×DN \times T \times DN×T×D

关键优势：

Token 现在可以关注所有帧的 token
充分利用已预训练的层

4.3 计算优化#

为降低 (NT)2(NT)^2(NT)2 的复杂度开销，使用 FlashAttention-2 实现高效计算。

4.4 旋转位置编码（RoPE）#

问题：膨胀注意力后观察到相邻帧间仍有小幅抖动。

解决方案：在膨胀注意力层内注入相对帧位置信息（通过旋转位置编码）。

效果：生成更平滑的帧间运动。

4.5 掩码生成（Masked Generation）#

动机：希望模型能够从已知 3D 网格开始生成（如 3D+视频 → 动画任务）。

方法：将模型改造为掩码生成模型，部分 3D 潜在向量已知，仅需生成剩余的”掩码”潜在向量。

实现细节：

训练时：随机采样 NSN_SNS 个潜在向量，保持无噪声
无噪声潜在向量的 Flow Matching 步设为 0（比 CAT3D 的二元掩码注入更自然）
推理时：将干净潜在向量复制到噪声潜在序列中

4.6 Stage I 推理流程#

1
1. 使用图像转 3D 生成器（如 TripoSG）生成首帧 3D 网格
2
2. 对视频每帧编码为 3D 潜在向量
3
3. 将首帧潜在向量设为已知（flow step = 0）
4
4. 运行 Temporal 3D Diffusion 去噪，生成时序一致的 4D 网格

5. Stage II：Temporal 3D Autoencoder#

5.1 为什么需要 Stage II？#

Stage I 输出的是 4D Mesh（每帧拓扑独立），但实际应用需要：

拓扑一致的 Animated 3D Mesh（所有帧共享相同网格拓扑）

5.2 核心思想#

预测时变顶点变形场 δk\delta_kδk：

(V+δk,F)≈(Vk,Fk)(V + \delta_k, F) \approx (V_k, F_k)(V+δk,F)≈(Vk,Fk)

即找到参考网格的顶点偏移，使得变形势与每帧的目标表面匹配。

5.3 架构设计#

1
输入：时序独立网格序列 {(V_k, F_k)}
2
        ↓
3
  ┌─────────────────────────────────────────┐
4
  │ Step 1: 独立编码（冻结 3D 编码器）      │
5
  │ 3D Encoder E3D → Shape Latents Z = {z_k} │
6
  └─────────────────────────────────────────┘
7
               ↓
8
  ┌─────────────────────────────────────────┐
9
  │ Step 2: 时序解码                         │
10
  │ Temporal 3D Decoder D4D                 │
11
  │  • 处理完整潜在序列                       │
12
  │  • 预测每帧的参考网格变形场               │
13
  │  • 膨胀自注意力 + RoPE                   │
14
  └─────────────────────────────────────────┘
15
               ↓
16
        Animated 3D Mesh {(V + δ_k, F)}

展开收起

5.4 时序解码器详解#

解码器核心机制：

自注意力处理：处理所有潜在 token
帧间跳转预测：对任意两帧 ti,tjt_i, t_jti,tj，输出查询点从 tit_iti 到 tjt_jtj 的位移
位置编码：将 ti,tjt_i, t_jti,tj 的 Fourier 嵌入拼接为额外 token
查询点增强：推理时使用参考网格顶点位置 + 顶点法向量

局部几何信息（法向量）有助于区分空间接近但拓扑距离远的点。

5.5 平移等变性#

该自编码器具有平移等变性：将点云序列转换为变形场序列，这是拓扑一致动画的关键保证。

6. 应用场景#

6.1 Video → 4D（主任务）#

1
输入：单目视频 {I_k}
2
        ↓
3
  Stage I (Temporal 3D Diffusion)
4
        ↓
5
  时序一致的独立网格序列 ((V_k, F_k))
6
        ↓
7
  Stage II (Temporal 3D Autoencoder)
8
        ↓
9
  Animated 3D Mesh ((V + δ_k, F))

6.2 3D + Text → Animation#

1
输入：3D 网格 M + 动画文本描述
2
        ↓
3
  1. 渲染 M 为正面视角图像 I_1
4
  2. 使用视频生成模型从文本生成视频 {I_k}
5
  3. 使用 ActionMesh ({3D + Video} → 4D) 生成动画

6.3 Image + Text → 4D#

1
输入：物体图像 + 动画文本描述
2
        ↓
3
  1. 使用图像转 3D 模型（TripoSG）生成初始网格
4
  2. 继续 {3D + Text} → Animation 流程

6.4 Text → 4D#

1
输入：文本描述
2
        ↓
3
  方案 A：文本 → 视频 → ActionMesh
4
  方案 B：文本 → 图像 → Image + Text → 4D（实际使用）

6.5 Motion Transfer / Retargeting#

意外发现：尽管未针对重定向训练，模型能够将视频中物体 A 的运动迁移到不同网格 B 上。

条件：物体 A 和 B 之间能够建立语义对应（如”飞鸟” → “龙”）。

6.6 Animation Extrapolation#

利用自回归建模，ActionMesh 支持动画外推：

1
长视频 → 分块 → 递归处理 → 连贯的长动画

7. 实验结果与分析#

7.1 Objaverse 定量评估#

评估指标：

CD-3D：逐帧 3D 重建精度（Chamfer Distance）
CD-4D：4D 重建质量（全局 ICP 对齐后的 Chamfer Distance）
CD-M：运动保真度（基于首帧对应点的双向 Chamfer Distance）

方法	时间	CD-3D ↓	CD-4D ↓	CD-M ↓
LIM	15 min	0.095	0.127	0.258
DreamMesh4D	35 min	0.095	0.140	0.247
V2M4	35 min	0.063	0.223	0.500
ActionMesh	3 min	0.050	0.069	0.137

ActionMesh 在所有指标上显著超越 SOTA，同时快 10 倍！

7.2 Consistent4D 定性对比#

方法	观察
LIM / DreamMesh4D	几何质量低、细节软、明显伪影
V2M4 / ShapeGen4D	细节锐利但有伪影和时序漂移
ActionMesh	高质量网格、强时序一致性、强运动保真度

7.3 消融实验#

消融类型	CD-3D ↓	CD-4D ↓	CD-M ↓
完整模型	0.050	0.069	0.137
无 Stage II	0.050	0.069	—
无 Stage I & II	0.050	0.187	—
Craftsman Backbone	0.072	0.117	0.216

关键发现：

Stage I 是关键：单独 Stage I 就能实现精确的 4D 重建
Stage II 保持质量：添加 Stage II 不损害 3D 重建质量
Backbone 可替换：Craftsman 也能达到竞争性能

7.4 ActionBench 基准#

ActionMesh 团队发布了 ActionBench：128 个配对的视频 ↔ 动画点云样本。

方法	CD-3D ↓	CD-4D ↓	CD-M ↓
TRELLIS	0.065	0.181	—
TripoSG	0.056	0.184	—
DreamMesh4D	0.104	0.152	0.265
LIM	0.089	0.126	0.243
V2M4	0.068	0.340	0.616
ShapeGen4D	0.056	0.170	0.348
ActionMesh (fast)	0.054	0.089	0.156
ActionMesh	0.054	0.085	0.153

8. 与现有方法的对比#

8.1 Video-to-4D 方法分类#

类别	代表方法	特点
优化类	DreamMesh4D, V2M4, LIM	高质量但需逐场景优化（30~45 分钟）
前馈类	L4GM, 4DGT	快速但不生成拓扑一致的网格
混合类	ShapeGen4D	无显式拓扑一致性保证
ActionMesh	前馈 + 拓扑一致	快速 + 高质量 + 拓扑一致

8.2 核心优势总结#

优势	ActionMesh	其他方法
推理速度	3 分钟	15~45 分钟
拓扑一致性	✅ 所有帧共享拓扑	❌ 大多不支持
Rig-free	✅ 无需骨骼绑定	❌ 人形/双足专用
输入多样性	视频/文本/图像/3D	通常单一模态
直接可导出	✅ GLB 格式	需额外处理

9. 局限性与未来方向#

9.1 当前局限#

局限	描述	示例
拓扑变化	假设固定连接性，无法建模拓扑变化	昆虫展开翅膀
强遮挡	遮挡区域重建困难	物体自遮挡

9.2 未来方向#

方向	描述
拓扑感知潜在更新	实例化、融合或移除局部部件，无需手动编辑连接性
更强的遮挡补全	改进参考帧的缺失区域补全能力
大规模视频预训练	从海量视频数据直接学习几何运动先验

10. 关键数学公式汇总#

Animated 3D Mesh 定义：

{(Vk,F)}k=1N其中 F 对所有帧恒定\{(V_k, F)\}_{k=1}^{N} \quad \text{其中} \ F \text{ 对所有帧恒定}{(Vk,F)}k=1N其中 F 对所有帧恒定

膨胀注意力：

infattn(X)=reshape−1(selfattn(reshape(X)))\text{infattn}(X) = \text{reshape}^{-1}(\text{selfattn}(\text{reshape}(X)))infattn(X)=reshape−1(selfattn(reshape(X)))

4D 重建目标：

(V+δk,F)≈(Vk,Fk)(V + \delta_k, F) \approx (V_k, F_k)(V+δk,F)≈(Vk,Fk)

时序解码器查询：

给定 ti,tj, 输出查询点从 ti 到 tj 的位移\text{给定} \ t_i, t_j, \ \text{输出查询点从} \ t_i \ \text{到} \ t_j \ \text{的位移}给定 ti,tj, 输出查询点从 ti 到 tj 的位移

11. 代码使用示例#

11.1 基础安装#

Terminal window

1
git clone https://github.com/facebookresearch/actionmesh.git
2
cd actionmesh
3
git submodule update --init --recursive
4
pip install -r requirements.txt
5
pip install -e .

11.2 Video → Animated Mesh#

Terminal window

1
# 标准模式（约 3 分钟）
2
python inference/video_to_animated_mesh.py --input assets/examples/davis_camel
3

4
# 快速模式（用于 Hugging Face Demo）
5
python inference/video_to_animated_mesh.py --input assets/examples/davis_camel --fast

11.3 输出说明#

输出	描述
`mesh_000.glb` ~ `mesh_XXX.glb`	每帧的 GLB 网格文件
`animated_mesh.glb`	带嵌入动画的完整网格（Blender 可直接导入）
`output.mp4`	渲染后的动画视频

12. 总结与展望#

ActionMesh 的核心贡献#

Temporal 3D Diffusion：首次将 3D 扩散模型扩展为时序版本
Temporal 3D Autoencoder：将独立形状序列转换为拓扑一致的动画
掩码生成机制：支持从已知 3D 网格开始的动画生成
两阶段解耦设计：3D 生成与动画预测分离，模块化且可扩展
10 倍加速：3 分钟完成动画网格生成，质量超越 15~45 分钟的优化方法

在 3D 生成版图中的位置#

1
3D 生成（静态）
2
├── LRM / TripoSG / Hunyuan3D
3
│   └── 单帧图像 → 3D 网格
4
│
5
└── ActionMesh (CVPR 2026)
6
    └── Temporal 3D Diffusion
7
        ├── Stage I: 时序一致的 4D 潜在向量
8
        └── Stage II: 拓扑一致的动画网格

ActionMesh 证明了”Temporal 3D Diffusion”是连接大规模视频语料与网格原生推理的有效路径。 它不仅实现了前所未有的速度和品质，更为零样本视频驱动的 4D 理解和生成开辟了新方向。

参考资料#

Sabathier, R., Novotny, D., Mitra, N. J., & Monnier, T. (2026). “ActionMesh: Animated 3D Mesh Generation with Temporal 3D Diffusion.” CVPR 2026.
“facebookresearch/actionmesh.” GitHub: facebookresearch/actionmesh. https://github.com/facebookresearch/actionmesh
“facebook/ActionMesh.” Hugging Face: facebook/ActionMesh. https://huggingface.co/facebook/ActionMesh
Zhang, R., et al. (2023). “3DShape2VecSet: A Neural Field for 3D Shape Generation and Manipulation.” NeurIPS 2023.
Li, Y., et al. (2025). “TripoSG: High-Fidelity 3D Shape Synthesis using Large-Scale Rectified Flow Models.” IEEE TPAMI 2025.
Cao, W., et al. (2024). “Motion2VecSets: 4D Latent Vector Set Diffusion for Non-rigid Shape Reconstruction and Tracking.” CVPR 2024.
Jiang, Y., et al. (2024). “Consistent4D: Consistent 360° Dynamic Object Generation from Monocular Video.” ICLR 2024.
Li, Z., et al. (2024). “DreamMesh4D: Video-to-4D Generation with Sparse-Controlled Gaussian-Mesh Hybrid Representation.” NeurIPS 2024.
Chen, J., et al. (2025). “V2M4: 4D Mesh Animation Reconstruction from a Single Monocular Video.” ICCV 2025.
Guo, Y., et al. (2025). “LIM: Learning Implicit Field for 4D Shape Generation.” arXiv 2025.

深入理解 AI Agent 架构：从单 Agent 到多 Agent 协作

Thu, 16 Jul 2026 00:00:00 GMT

1. 引言：从 LLM 到 Agent#

1.1 什么是 AI Agent？#

AI Agent 是能够自主感知环境、做出决策、执行行动的人工智能系统。与传统 LLM 不同，Agent 不只是被动回答问题，而是能够：

感知 (Perception)：观察环境状态、接收用户输入、读取工具结果
推理 (Reasoning)：分析信息、规划任务、做出决策
行动 (Action)：执行工具调用、操作环境、生成响应

Agent = LLM + 工具 + 记忆 + 规划

可以理解为：Agent 是一个有”目标感”的 LLM，它不仅能思考，还能主动采取行动。

1.2 为什么需要 Agent？#

1.2.1 传统 LLM 的局限#

1
孤立 LLM:
2
  静态知识 (训练截止)
3
  + 无工具 (不能执行)
4
  + 无记忆 (每次独立)
5
  + 无目标 (被动响应)
6
  = 强大的聊天机器人

1.2.2 Agent 的增强#

1
增强 Agent:
2
  动态知识 (RAG 检索)
3
  + 多工具 (浏览器/代码/API)
4
  + 长记忆 (向量化)
5
  + 自主目标 (规划执行)
6
  = 通用问题解决者

1.3 Agent 的发展历程#

1
2022.10  ReAct (Yao et al.)
2
2023.03  AutoGPT / BabyAGI 引领自主 Agent 浪潮
3
2023.06  OpenAI Function Calling
4
2023.11  LangChain 推出 Agent 框架
5
2024.01  LangGraph 发布
6
2024.05  Anthropic Computer Use
7
2024.06  AutoGen (Microsoft)
8
2024.10  Claude 3.5 Sonnet Computer Use
9
2025.01  Anthropic Prompt Caching
10
2025+    Agent as a Service (Agent 服务化)

1.4 应用场景#

场景	Agent 任务
软件开发	代码生成、Bug 修复、PR Review
研究分析	信息检索、数据分析、报告撰写
客户支持	自动问答、工单处理、问题升级
数据分析	查询生成、可视化、洞察发现
工作流自动化	多步骤任务、跨系统协调

2. Agent 核心架构#

2.1 Agent 三要素#

经典的 Agent 架构包含三个核心组件：

1
         ┌─────────────────────────────────────┐
2
         │           AI Agent                  │
3
         │                                      │
4
         │   ┌────────────┐  ┌────────────┐   │
5
         │   │ Perception │  │  Memory    │   │
6
         │   └─────┬──────┘  └─────▲──────┘   │
7
         │         │               │          │
8
         │         ▼               │          │
9
         │   ┌─────────────────────┴────────┐ │
10
         │   │        Reasoning             │ │
11
         │   │   (LLM + Planning)          │ │
12
         │   └──────────────┬───────────────┘ │
13
         │                  │                 │
14
         │                  ▼                 │
15
         │         ┌──────────────┐          │
16
         │         │    Action    │          │
17
         │         └──────┬───────┘          │
18
         └────────────────┼──────────────────┘
19
                          ▼
20
                    [Environment]

展开收起

2.2 Perception 模块#

Perception 是 Agent 的”感官”，负责接收和处理输入：

1
class Perception:
2
    """Agent 感知模块"""
3
    def __init__(self):
4
        self.input_processors = []
5

6
    def add_processor(self, processor):
7
        """注册输入处理器"""
8
        self.input_processors.append(processor)
9

10
    def perceive(self, raw_input):
11
        """处理输入"""
12
        processed = raw_input
13

14
        for processor in self.input_processors:
15
            processed = processor(processed)
16

17
        return processed
18

19
# 示例: 多模态感知
20
perception = Perception()
21

22
# 文本输入
23
perception.add_processor(lambda x: x.strip())
24

25
# 包含图像
26
def image_processor(input):
27
    if "<image>" in input:
28
        # 调用视觉模型描述图像
29
        return describe_image(input)
30
    return input
31
perception.add_processor(image_processor)

展开收起

感知来源包括：

用户输入：文本、语音、图像
环境状态：数据库内容、文件系统、API 返回值
工具结果：搜索响应、计算结果、API 响应
其他 Agent：多 Agent 系统中的消息

2.3 Memory 模块#

Memory 是 Agent 的”大脑”，负责存储和检索信息：

1
class Memory:
2
    """Agent 记忆系统"""
3
    def __init__(self):
4
        self.short_term = []      # 短期记忆 (当前任务)
5
        self.long_term = {}        # 长期记忆 (向量化)
6
        self.episodic = []         # 情节记忆 (历史行动)
7
        self.semantic = {}         # 语义记忆 (事实知识)
8

9
    def store(self, content, memory_type='short_term'):
10
        """存储记忆"""
11
        if memory_type == 'short_term':
12
            self.short_term.append(content)
13
        elif memory_type == 'long_term':
14
            vector = self.embed(content)
15
            self.long_term[vector] = content
16
        # ...
17

18
    def retrieve(self, query, top_k=5):
19
        """检索相关记忆"""
20
        # 使用 embedding 相似度
21
        query_vector = self.embed(query)
22
        relevant = self.similarity_search(query_vector, top_k)
23
        return relevant

展开收起

Memory 的分类参见第 5 节。

2.4 Reasoning 模块#

Reasoning 是 Agent 的”思维”，负责规划和决策：

1
class Reasoning:
2
    """Agent 推理模块"""
3
    def __init__(self, llm):
4
        self.llm = llm
5
        self.planner = Planner(llm)
6
        self.reflector = Reflector(llm)
7

8
    def reason(self, state, goal):
9
        """推理下一步行动"""
10
        # 1. 规划：分解目标
11
        plan = self.planner.create_plan(state, goal)
12

13
        # 2. 执行当前步骤
14
        step = plan.current_step()
15

16
        # 3. 反思：如果失败，调整
17
        if self._is_stuck(state):
18
            plan = self.reflector.reflect(state, plan)
19
            step = plan.current_step()
20

21
        return step
22

23
class Planner:
24
    """规划器"""
25
    def create_plan(self, state, goal):
26
        prompt = f"""
27
        目标: {goal}
28
        当前状态: {state}
29

30
        请创建一个详细的执行计划：
31
        1. [步骤1]
32
        2. [步骤2]
33
        ...
34
        """
35
        plan_text = self.llm.generate(prompt)
36
        return self._parse_plan(plan_text)

展开收起

2.5 Action 模块#

Action 是 Agent 的”手脚”，负责执行决策：

1
class Action:
2
    """Agent 行动模块"""
3
    def __init__(self, tools):
4
        self.tools = tools    # 可用工具集合
5

6
    def execute(self, action):
7
        """执行行动"""
8
        tool_name = action.tool
9
        args = action.arguments
10

11
        tool = self.tools.get(tool_name)
12
        if not tool:
13
            raise ToolNotFound(tool_name)
14

15
        result = tool.execute(**args)
16
        return result
17

18
class Tool:
19
    """工具"""
20
    def __init__(self, name, func, schema):
21
        self.name = name
22
        self.func = func
23
        self.schema = schema
24

25
    def execute(self, **kwargs):
26
        try:
27
            result = self.func(**kwargs)
28
            return ActionResult(success=True, data=result)
29
        except Exception as e:
30
            return ActionResult(success=False, error=str(e))

展开收起

3. Agent 设计模式#

3.1 模式 1: ReAct#

ReAct (Reasoning + Acting) 是最经典的 Agent 模式，由 Yao et al. (2022) 提出：

1
循环:
2
  Thought → Action → Observation → Thought → ...

1
class ReActAgent:
2
    """ReAct Agent"""
3
    def __init__(self, llm, tools, max_steps=10):
4
        self.llm = llm
5
        self.tools = tools
6
        self.max_steps = max_steps
7

8
    def run(self, question):
9
        history = []
10

11
        for step in range(self.max_steps):
12
            # 1. Thought + Action
13
            prompt = self._build_prompt(question, history)
14
            response = self.llm.generate(prompt)
15

16
            thought, action = self._parse_response(response)
17

18
            # 2. 完成检测
19
            if self._is_final_answer(response):
20
                return self._extract_answer(response)
21

22
            # 3. 执行 Action
23
            observation = self._execute_tool(action)
24

25
            # 4. 更新历史
26
            history.append({
27
                'thought': thought,
28
                'action': action,
29
                'observation': observation
30
            })
31

32
        return "达到最大步数"
33

34
# ReAct 循环示意:
35
# Thought 1: 需要查询北京天气
36
# Action 1: get_weather(city="北京")
37
# Observation 1: 25°C, 晴
38
#
39
# Thought 2: 现在查空气质量
40
# Action 2: get_air_quality(city="北京")
41
# Observation 2: AQI 85
42
#
43
# Thought 3: 可以给出建议
44
# Final Answer: 北京今天 25°C, 晴朗, 空气良好...

展开收起

3.2 模式 2: Plan-and-Execute#

Plan-and-Execute 先完整规划，再逐步执行：

1
规划阶段: 一次生成完整计划
2
执行阶段: 按计划逐步执行

1
class PlanExecuteAgent:
2
    """Plan-and-Execute Agent"""
3

4
    def run(self, question):
5
        # 1. 规划阶段
6
        plan = self.create_plan(question)
7

8
        # 2. 执行阶段
9
        results = []
10
        for step in plan:
11
            try:
12
                result = self.execute_step(step)
13
                results.append(result)
14
            except Exception as e:
15
                # 失败：重规划
16
                plan = self.replan(question, plan, results, e)
17
                return self.run(question)  # 重试
18

19
        return self.synthesize(results)
20

21
    def create_plan(self, question):
22
        """创建完整计划"""
23
        prompt = f"""
24
        任务: {question}
25

26
        请创建一个详细的执行计划。每个步骤应该明确指定:
27
        - 工具
28
        - 参数
29
        - 预期输出
30

31
        计划:
32
        """
33
        plan_text = self.llm.generate(prompt)
34
        return self._parse_plan(plan_text)
35

36
    def replan(self, question, plan, results, error):
37
        """失败后重新规划"""
38
        prompt = f"""
39
        原任务: {question}
40
        原计划: {plan}
41
        已执行: {results}
42
        错误: {error}
43

44
        请重新规划。跳过失败步骤或找到替代方案。
45

46
        新计划:
47
        """
48
        new_plan_text = self.llm.generate(prompt)
49
        return self._parse_plan(new_plan_text)

展开收起

优势：

全局可见的计划
执行阶段 LLM 调用少
适合复杂多步任务

3.3 模式 3: Reflexion#

Reflexion 在 ReAct 基础上加入反思机制：

1
class ReflexionAgent(ReActAgent):
2
    """带反思的 Agent"""
3

4
    def __init__(self, llm, tools, max_reflections=3):
5
        super().__init__(llm, tools)
6
        self.max_reflections = max_reflections
7
        self.reflections = []
8

9
    def run(self, question):
10
        for attempt in range(self.max_reflections + 1):
11
            # 执行 ReAct
12
            result = super().run(question)
13

14
            # 评估
15
            if self.is_successful(result):
16
                return result
17

18
            if attempt < self.max_reflections:
19
                # 反思
20
                reflection = self.reflect(question, result)
21
                self.reflections.append(reflection)
22

23
        return result
24

25
    def reflect(self, question, result):
26
        """反思失败原因"""
27
        prompt = f"""
28
        问题: {question}
29
        尝试结果: {result}
30

31
        反思以下:
32
        1. 为什么失败？
33
        2. 下次可以采取什么不同的策略？
34
        3. 你学到了什么？
35

36
        反思:
37
        """
38
        return self.llm.generate(prompt)

展开收起

反思范式：

行动反思：行动后立即反思
失败反思：失败后反思
策略反思：更新整体策略

3.4 模式 4: ReWOO#

ReWOO (Reasoning WithOut Observation) 通过预先生成所有任务减少冗余：

1
class ReWOOAgent:
2
    """ReWOO Agent - 减少观察次数"""
3

4
    def run(self, question):
5
        # 1. 一次性生成所有步骤
6
        plan = self.generate_full_plan(question)
7

8
        # 2. 执行所有工具调用（无需中间反馈）
9
        evidence = {}
10
        for step in plan:
11
            result = self.execute_tool(step)
12
            evidence[step.id] = result
13

14
        # 3. 综合所有证据
15
        return self.synthesize(question, plan, evidence)
16

17
    def generate_full_plan(self, question):
18
        """生成完整计划，包括变量替换"""
19
        # 计划结构:
20
        # Step 1: Search("Beijing weather") -> #E1
21
        # Step 2: Search("Shanghai weather") -> #E2
22
        # Step 3: LLM(#E1, #E2, "compare weathers") -> #E3
23
        return self.llm.generate(prompt_for_plan(question))

展开收起

优势：

减少 LLM 调用次数
适合并行工具调用
降低成本

3.5 模式 5: Tree of Thoughts#

Tree of Thoughts (ToT) 把推理组织成树搜索：

1
class ToTAgent:
2
    """Tree of Thoughts Agent"""
3

4
    def __init__(self, llm, branching_factor=3, max_depth=5):
5
        self.llm = llm
6
        self.branching_factor = branching_factor
7
        self.max_depth = max_depth
8

9
    def run(self, question):
10
        root = ThoughtNode(question)
11
        tree = ThoughtTree(root)
12

13
        for depth in range(self.max_depth):
14
            # 1. 扩展当前层
15
            current_nodes = tree.get_nodes_at_depth(depth)
16
            for node in current_nodes:
17
                # 生成多个候选
18
                thoughts = self.generate_thoughts(node.state, n=self.branching_factor)
19
                for thought in thoughts:
20
                    child = ThoughtNode(thought, parent=node)
21

22
                    # 评估
23
                    child.score = self.evaluate(thought)
24

25
                    tree.add(child)
26

27
            # 2. 剪枝
28
            tree.prune(keep_top_n=self.branching_factor)
29

30
        return tree.get_best_path()
31

32
    def generate_thoughts(self, state, n):
33
        """生成 n 个下一步思考"""
34
        prompt = f"当前: {state}\n生成 {n} 个可能的下一步思考："
35
        response = self.llm.generate(prompt)
36
        return self._parse_thoughts(response, n)
37

38
    def evaluate(self, thought):
39
        """评估思考的质量"""
40
        prompt = f"评估以下思考（1-10）：{thought}"
41
        return int(self.llm.generate(prompt))

展开收起

3.6 模式对比#

模式	LLM 调用	适合场景	复杂度
ReAct	高 (每步)	简单任务	低
Plan-Execute	低	复杂多步	中
Reflexion	中	反复尝试	中
ReWOO	很低	独立任务	低
ToT	极高	复杂搜索	高

4. Planning 模块深入#

4.1 任务分解#

1
class TaskDecomposer:
2
    """任务分解器"""
3

4
    def decompose(self, task):
5
        """将复杂任务分解为子任务"""
6
        prompt = f"""
7
        任务: {task}
8

9
        请将这个任务分解为多个子任务。每个子任务应该：
10
        1. 具体明确
11
        2. 可独立执行
12
        3. 有明确的成功标准
13

14
        子任务列表（JSON 数组）:
15
        """
16
        response = self.llm.generate(prompt)
17
        return self._parse_tasks(response)
18

19
# 示例分解:
20
# 原始任务: "研究某个 AI 主题"
21
# 分解为:
22
# 1. 搜索 arXiv 相关论文
23
# 2. 阅读最重要的 3 篇
24
# 3. 对比方法
25
# 4. 总结发现

展开收起

4.2 计划表示#

1
from dataclasses import dataclass
2
from typing import List, Optional
3

4
@dataclass
5
class PlanStep:
6
    """计划步骤"""
7
    id: int
8
    description: str
9
    tool: Optional[str] = None    # 可选的工具
10
    args: dict = None             # 工具参数
11
    depends_on: List[int] = None  # 依赖的前置步骤
12
    status: str = "pending"       # pending/running/done/failed
13
    result: any = None
14

15
@dataclass
16
class Plan:
17
    """完整计划"""
18
    goal: str
19
    steps: List[PlanStep]
20

21
    def ready_steps(self):
22
        """获取所有可执行的步骤（依赖已满足）"""
23
        done_ids = {s.id for s in self.steps if s.status == 'done'}
24
        return [
25
            s for s in self.steps
26
            if s.status == 'pending' and
27
            all(dep in done_ids for dep in (s.depends_on or []))
28
        ]
29

30
    def is_complete(self):
31
        return all(s.status == 'done' for s in self.steps)
32

33
    def progress(self):
34
        done = sum(1 for s in self.steps if s.status == 'done')
35
        return done / len(self.steps)

展开收起

4.3 动态重规划#

1
class DynamicPlanner:
2
    """动态规划器"""
3

4
    def run_with_replanning(self, plan, executor):
5
        """带重规划的执行"""
6
        while not plan.is_complete():
7
            ready = plan.ready_steps()
8

9
            # 并行执行所有就绪步骤
10
            futures = []
11
            for step in ready:
12
                step.status = 'running'
13
                futures.append(executor.execute_async(step))
14

15
            # 等待结果
16
            for step, future in zip(ready, futures):
17
                try:
18
                    result = future.result()
19
                    step.status = 'done'
20
                    step.result = result
21
                except Exception as e:
22
                    step.status = 'failed'
23
                    step.result = str(e)
24

25
            # 检查是否需要重规划
26
            failed = [s for s in plan.steps if s.status == 'failed']
27
            if failed:
28
                plan = self.replan(plan, failed)
29

30
        return plan

展开收起

4.4 分层规划#

1
class HierarchicalPlanner:
2
    """分层规划器"""
3

4
    def plan(self, goal):
5
        """分层规划"""
6
        # Level 1: 战略规划
7
        strategy = self.strategic_plan(goal)
8

9
        # Level 2: 战术规划
10
        tactics = [self.tactical_plan(subgoal) for subgoal in strategy]
11

12
        # Level 3: 操作规划
13
        operations = []
14
        for tactic in tactics:
15
            ops = self.operational_plan(tactic)
16
            operations.extend(ops)
17

18
        return Plan(strategy, tactics, operations)
19

20
    def strategic_plan(self, goal):
21
        """高层次规划"""
22
        prompt = f"目标: {goal}\n拆分为 3-5 个主要阶段："
23
        return self.llm.generate(prompt)
24

25
    def tactical_plan(self, subgoal):
26
        """子目标规划"""
27
        prompt = f"子目标: {subgoal}\n具体行动步骤："
28
        return self.llm.generate(prompt)
29

30
    def operational_plan(self, action):
31
        """操作级规划"""
32
        # 确定具体的工具调用
33
        # ...

展开收起

5. Memory 系统设计#

5.1 Memory 分类#

1
┌─────────────────────────────────────────────────────┐
2
│                  Agent Memory                        │
3
├─────────────────────────────────────────────────────┤
4
│                                                      │
5
│  ┌──────────────────────────────────────────────┐    │
6
│  │         Short-Term Memory                    │    │
7
│  │  - 当前对话上下文                            │    │
8
│  │  - 最近 N 轮对话                            │    │
9
│  │  - 滑动窗口                                 │    │
10
│  └──────────────────────────────────────────────┘    │
11
│                                                      │
12
│  ┌──────────────────────────────────────────────┐    │
13
│  │         Long-Term Memory (向量存储)          │    │
14
│  │  - 历史对话摘要                              │    │
15
│  │  - 用户偏好                                  │    │
16
│  │  - 事实知识                                  │    │
17
│  └──────────────────────────────────────────────┘    │
18
│                                                      │
19
│  ┌──────────────────────────────────────────────┐    │
20
│  │         Episodic Memory                      │    │
21
│  │  - 具体的过去事件                            │    │
22
│  │  - "上次我尝试 X 然后..."                     │    │
23
│  │  - 成功/失败模式                             │    │
24
│  └──────────────────────────────────────────────┘    │
25
│                                                      │
26
│  ┌──────────────────────────────────────────────┐    │
27
│  │         Semantic Memory                      │    │
28
│  │  - 概念关系                                  │    │
29
│  │  - 知识图谱                                  │    │
30
│  │  - 领域知识                                  │    │
31
│  └──────────────────────────────────────────────┘    │
32
│                                                      │
33
└─────────────────────────────────────────────────────┘

展开收起

5.2 短期记忆实现#

1
class ShortTermMemory:
2
    """短期记忆"""
3

4
    def __init__(self, max_tokens=4000):
5
        self.messages = []
6
        self.max_tokens = max_tokens
7

8
    def add(self, role, content):
9
        """添加消息"""
10
        msg = {'role': role, 'content': content}
11
        self.messages.append(msg)
12
        self._truncate()
13

14
    def get_messages(self):
15
        return self.messages
16

17
    def _truncate(self):
18
        """超出 token 限制时截断"""
19
        while self._count_tokens() > self.max_tokens:
20
            # 保留 system prompt
21
            if len(self.messages) > 2:
22
                self.messages.pop(1)  # 移除第二条 (最早用户消息)
23

24
    def _count_tokens(self):
25
        return sum(len(m['content']) // 4 for m in self.messages)

展开收起

5.3 长期记忆实现#

1
import numpy as np
2
from typing import List, Tuple
3

4
class LongTermMemory:
5
    """长期记忆 - 向量存储"""
6

7
    def __init__(self, embedder, vector_store):
8
        self.embedder = embedder  # embedding 模型
9
        self.vector_store = vector_store  # 向量数据库
10

11
    def store(self, content, metadata=None):
12
        """存储记忆"""
13
        embedding = self.embedder.embed(content)
14
        self.vector_store.add(
15
            vector=embedding,
16
            content=content,
17
            metadata=metadata or {}
18
        )
19

20
    def retrieve(self, query, top_k=5):
21
        """检索相关记忆"""
22
        query_embedding = self.embedder.embed(query)
23
        results = self.vector_store.search(
24
            vector=query_embedding,
25
            top_k=top_k
26
        )
27
        return results
28

29
    def retrieve_relevant_for_context(self, query):
30
        """检索并格式化为上下文"""
31
        results = self.retrieve(query)
32
        context = "\n".join([
33
            f"[{i}] {r['content']}"
34
            for i, r in enumerate(results)
35
        ])
36
        return context

展开收起

5.4 记忆压缩#

1
class MemoryCompressor:
2
    """记忆压缩器"""
3

4
    def __init__(self, llm):
5
        self.llm = llm
6

7
    def summarize_conversation(self, messages):
8
        """总结对话"""
9
        conversation = "\n".join([
10
            f"{m['role']}: {m['content']}"
11
            for m in messages
12
        ])
13

14
        prompt = f"""
15
        请总结以下对话的关键信息:
16
        {conversation}
17

18
        总结应包含:
19
        1. 用户的主要需求
20
        2. 已完成的工作
21
        3. 重要的决定或发现
22
        4. 待解决的问题
23

24
        总结:
25
        """
26
        return self.llm.generate(prompt)
27

28
    def compress_old_memories(self, memories):
29
        """压缩旧记忆"""
30
        # 相似记忆去重
31
        unique = self._deduplicate(memories)
32

33
        # 过长记忆总结
34
        compressed = []
35
        for memory in unique:
36
            if len(memory) > 500:
37
                summarized = self.summarize_conversation([{'content': memory}])
38
                compressed.append(summarized)
39
            else:
40
                compressed.append(memory)
41

42
        return compressed

展开收起

5.5 Memory-Augmented Agent#

1
class MemoryAugmentedAgent:
2
    """带记忆的 Agent"""
3

4
    def __init__(self, llm, tools, memory):
5
        self.llm = llm
6
        self.tools = tools
7
        self.memory = memory
8

9
    def run(self, user_input):
10
        # 1. 检索相关记忆
11
        relevant_memories = self.memory.retrieve_relevant_for_context(user_input)
12

13
        # 2. 增强 prompt
14
        augmented_prompt = f"""
15
        相关历史:
16
        {relevant_memories}
17

18
        当前请求:
19
        {user_input}
20
        """
21

22
        # 3. 生成响应（可调用工具）
23
        response = self._generate(augmented_prompt)
24

25
        # 4. 存储到记忆
26
        self.memory.store(
27
            f"用户: {user_input}\n助手: {response}"
28
        )
29

30
        return response

展开收起

6. Tool Use#

6.1 Tool 抽象层#

1
from typing import Callable, Dict, Any, List
2
from dataclasses import dataclass
3
import json
4

5
@dataclass
6
class Tool:
7
    """工具"""
8
    name: str
9
    description: str
10
    func: Callable
11
    schema: Dict[str, Any]
12

13
    def execute(self, **kwargs) -> str:
14
        try:
15
            result = self.func(**kwargs)
16
            return json.dumps({"success": True, "data": result})
17
        except Exception as e:
18
            return json.dumps({"success": False, "error": str(e)})
19

20
    def to_openai_format(self) -> Dict:
21
        return {
22
            "type": "function",
23
            "function": {
24
                "name": self.name,
25
                "description": self.description,
26
                "parameters": self.schema
27
            }
28
        }
29

30
class ToolRegistry:
31
    """工具注册表"""
32

33
    def __init__(self):
34
        self.tools: Dict[str, Tool] = {}
35

36
    def register(self, name=None, description=None):
37
        def decorator(func):
38
            tool_name = name or func.__name__
39
            tool_desc = description or (func.__doc__ or "")
40
            schema = self._infer_schema(func)
41
            tool = Tool(tool_name, tool_desc, func, schema)
42
            self.tools[tool_name] = tool
43
            return func
44
        return decorator
45

46
    def _infer_schema(self, func):
47
        import inspect
48
        sig = inspect.signature(func)
49
        properties = {}
50
        required = []
51

52
        for name, param in sig.parameters.items():
53
            properties[name] = {"type": "string"}
54
            if param.default is param.empty:
55
                required.append(name)
56

57
        return {
58
            "type": "object",
59
            "properties": properties,
60
            "required": required
61
        }
62

63
    def get(self, name: str) -> Tool:
64
        return self.tools.get(name)
65

66
    def to_openai_tools(self) -> List[Dict]:
67
        return [t.to_openai_format() for t in self.tools.values()]

展开收起

6.2 Tool Selection#

1
class ToolSelector:
2
    """智能工具选择器"""
3

4
    def __init__(self, all_tools, embedding_model):
5
        self.all_tools = all_tools
6
        self.embedding_model = embedding_model
7
        self.tool_embeddings = self._embed_tools()
8

9
    def _embed_tools(self):
10
        embeddings = {}
11
        for tool in self.all_tools:
12
            tool_text = f"{tool.name}: {tool.description}"
13
            embeddings[tool.name] = self.embedding_model.embed(tool_text)
14
        return embeddings
15

16
    def select(self, query, top_k=5):
17
        """根据查询选择相关工具"""
18
        query_emb = self.embedding_model.embed(query)
19

20
        similarities = {}
21
        for name, tool_emb in self.tool_embeddings.items():
22
            sim = self._cosine_similarity(query_emb, tool_emb)
23
            similarities[name] = sim
24

25
        # 返回 top-k
26
        sorted_tools = sorted(
27
            similarities.items(),
28
            key=lambda x: x[1],
29
            reverse=True
30
        )
31

32
        return [name for name, _ in sorted_tools[:top_k]]
33

34
    def _cosine_similarity(self, a, b):
35
        return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

展开收起

6.3 Tool 错误处理#

1
class RobustToolExecutor:
2
    """健壮的工具执行器"""
3

4
    def __init__(self, tools, max_retries=3):
5
        self.tools = tools
6
        self.max_retries = max_retries
7

8
    async def execute(self, name, args):
9
        """异步执行，带重试"""
10
        last_error = None
11

12
        for attempt in range(self.max_retries):
13
            try:
14
                tool = self.tools.get(name)
15
                if not tool:
16
                    return {"error": f"未知工具: {name}"}
17

18
                result = tool.execute(**args)
19

20
                if isinstance(result, dict) and result.get('success'):
21
                    return result
22
                else:
23
                    raise Exception(result.get('error', 'Unknown error'))
24

25
            except Exception as e:
26
                last_error = e
27

28
                # 重试延迟
29
                if attempt < self.max_retries - 1:
30
                    delay = 2 ** attempt  # 指数退避
31
                    await asyncio.sleep(delay)
32

33
        return {"error": f"执行失败 (尝试 {self.max_retries} 次): {last_error}"}

展开收起

7. Multi-Agent 系统#

7.1 为什么需要 Multi-Agent？#

单 Agent 在复杂任务中存在局限：

上下文限制：难以处理庞大信息
能力单一：一个 Agent 不可能精通所有
协作困难：单 Agent 缺乏不同视角
专业不足：通用模型在专业领域不够精准

Multi-Agent 通过分工协作解决这些问题：

专业分工：每个 Agent 专注于特定领域
并行处理：多个 Agent 同时工作
交叉验证：不同 Agent 可以相互校验
可扩展性：根据需要增加新 Agent

7.2 Multi-Agent 架构#

7.2.1 Supervisor 架构#

一个中心 Supervisor 协调多个 Worker Agents：

1
                Supervisor Agent
2
                       │
3
        ┌──────────────┼──────────────┐
4
        ▼              ▼              ▼
5
   Researcher     Coder         Reviewer
6
     Agent        Agent           Agent

1
class SupervisorAgent:
2
    """监督者 Agent"""
3

4
    def __init__(self, llm, worker_agents):
5
        self.llm = llm
6
        self.worker_agents = worker_agents
7

8
    def run(self, task):
9
        while not self._is_complete(task):
10
            # 决定下一步分配给谁
11
            decision = self._decide_next_agent(task)
12
            agent_name = decision['agent']
13
            sub_task = decision['sub_task']
14

15
            # 调用 Worker Agent
16
            if agent_name == "FINISH":
17
                break
18

19
            worker = self.worker_agents[agent_name]
20
            result = worker.run(sub_task)
21

22
            # 更新任务状态
23
            task.add_result(agent_name, result)
24

25
        return task.synthesize()
26

27
# 定义 Worker Agents
28
class ResearchAgent:
29
    def run(self, task):
30
        # 研究任务
31
        return search_and_summarize(task)
32

33
class CoderAgent:
34
    def run(self, task):
35
        # 编程任务
36
        return write_code(task)
37

38
class ReviewerAgent:
39
    def run(self, task):
40
        # 评审任务
41
        return review(task)
42

43
# 创建系统
44
supervisor = SupervisorAgent(
45
    llm=llm,
46
    worker_agents={
47
        'researcher': ResearchAgent(llm),
48
        'coder': CoderAgent(llm),
49
        'reviewer': ReviewerAgent(llm)
50
    }
51
)

展开收起

7.2.2 对等架构#

各 Agent 完全对等，点对点通信：

1
class PeerToPeerAgent:
2
    """点对点 Agent"""
3

4
    def __init__(self, name, llm, tools, peers=None):
5
        self.name = name
6
        self.llm = llm
7
        self.tools = tools
8
        self.peers = peers or {}  # 其他 Agent
9
        self.mailbox = []        # 消息队列
10

11
    def send_message(self, peer_name, message):
12
        """发送消息给其他 Agent"""
13
        peer = self.peers.get(peer_name)
14
        if peer:
15
            peer.receive_message(self.name, message)
16

17
    def receive_message(self, sender, message):
18
        """接收消息"""
19
        self.mailbox.append({'from': sender, 'content': message})
20

21
    def run(self, task):
22
        # 可以询问其他 Agent
23
        context = self._format_context()
24
        prompt = f"{context}\n\n任务: {task}\n\n响应:"
25
        response = self.llm.generate(prompt)
26

27
        # 决定是否需要其他人协助
28
        if self._needs_help(response):
29
            helper = self._choose_helper()
30
            self.send_message(helper, "需要帮助")
31

32
        return response

展开收起

7.2.3 层级架构#

按层级组织 Agent：

1
顶层: Planner Agent (战略规划)
2
   │
3
   ├─ 中层: 多个 Specialist Agent (领域专家)
4
   │
5
   └─ 底层: 多个 Executor Agent (执行单元)

7.3 Agent 通信协议#

1
class AgentMessage:
2
    """Agent 间消息"""
3
    def __init__(self, sender, receiver, content,
4
                 msg_type='text', metadata=None):
5
        self.sender = sender
6
        self.receiver = receiver
7
        self.content = content
8
        self.msg_type = msg_type  # text/request/response/error
9
        self.metadata = metadata or {}
10
        self.timestamp = time.time()
11

12
class AgentCommunicationBus:
13
    """Agent 通信总线"""
14

15
    def __init__(self):
16
        self.agents = {}
17
        self.message_log = []
18

19
    def register(self, agent):
20
        self.agents[agent.name] = agent
21

22
    def send(self, message):
23
        """发送消息"""
24
        # 记录
25
        self.message_log.append(message)
26

27
        # 路由
28
        receiver = self.agents.get(message.receiver)
29
        if receiver:
30
            receiver.receive_message(message)
31

32
    def broadcast(self, sender, content):
33
        """广播"""
34
        for name, agent in self.agents.items():
35
            if name != sender:
36
                self.send(AgentMessage(sender, name, content))

展开收起

7.4 协作模式#

7.4.1 投票 (Voting)#

1
class VotingAgent:
2
    """投票 Agent"""
3

4
    def __init__(self, agents):
5
        self.agents = agents
6

7
    def run(self, task):
8
        # 每个 Agent 独立给出答案
9
        answers = []
10
        for agent in self.agents:
11
            answer = agent.run(task)
12
            answers.append(answer)
13

14
        # 投票
15
        from collections import Counter
16
        votes = Counter(answers)
17
        most_common = votes.most_common(1)[0][0]
18

19
        return most_common, votes

展开收起

7.4.2 辩论 (Debate)#

1
class DebateAgent:
2
    """辩论 Agent"""
3

4
    def __init__(self, agents, rounds=3):
5
        self.agents = agents
6
        self.rounds = rounds
7

8
    def run(self, task):
9
        opinions = [agent.run(task) for agent in self.agents]
10

11
        for round in range(self.rounds):
12
            new_opinions = []
13
            for i, agent in enumerate(self.agents):
14
                # 让每个 Agent 看到其他 Agent 的意见
15
                others = [op for j, op in enumerate(opinions) if j != i]
16
                context = "\n".join([f"Agent {j}: {op}" for j, op in enumerate(others)])
17

18
                prompt = f"""
19
                任务: {task}
20

21
                其他 Agent 的意见:
22
                {context}
23

24
                你的新意见:
25
                """
26
                new_opinion = agent.llm.generate(prompt)
27
                new_opinions.append(new_opinion)
28

29
            opinions = new_opinions
30

31
        # 综合最终答案
32
        final = self.synthesize(opinions)
33
        return final

展开收起

7.4.3 流水线 (Pipeline)#

1
class PipelineAgent:
2
    """流水线 Agent"""
3

4
    def __init__(self, agents):
5
        self.agents = agents
6

7
    def run(self, task):
8
        result = task
9
        for agent in self.agents:
10
            result = agent.run(result)
11
        return result
12

13
# 使用：研究 -> 写作 -> 审查 -> 发布
14
pipeline = PipelineAgent([
15
    ResearchAgent(llm),
16
    WriterAgent(llm),
17
    ReviewerAgent(llm),
18
    PublisherAgent(llm)
19
])

展开收起

7.5 AutoGen 框架示例#

1
# AutoGen 风格的多 Agent 协作
2
from autogen import AssistantAgent, UserProxyAgent, GroupChat, GroupChatManager
3

4
# 定义 Agent
5
researcher = AssistantAgent(
6
    name="researcher",
7
    llm_config={"model": "gpt-4"},
8
    system_message="你是一个研究员。负责查找和分析信息。"
9
)
10

11
coder = AssistantAgent(
12
    name="coder",
13
    llm_config={"model": "gpt-4"},
14
    system_message="你是一个程序员。负责编写代码。"
15
)
16

17
user_proxy = UserProxyAgent(
18
    name="user_proxy",
19
    human_input_mode="TERMINATE"
20
)
21

22
# 创建群聊
23
groupchat = GroupChat(
24
    agents=[user_proxy, researcher, coder],
25
    messages=[],
26
    max_round=10
27
)
28

29
manager = GroupChatManager(groupchat=groupchat, llm_config={"model": "gpt-4"})
30

31
# 启动
32
user_proxy.initiate_chat(
33
    manager,
34
    message="开发一个计算斐波那契数列的 Python 程序"
35
)

展开收起

8. LangGraph 框架#

8.1 LangGraph 概述#

LangGraph 是 LangChain 推出的图状 Agent 框架：

1
# LangGraph 用 StateGraph 定义 Agent
2
from langgraph.graph import StateGraph, END
3

4
# 定义 State
5
class AgentState(TypedDict):
6
    messages: list
7
    next_step: str
8

9
# 定义节点
10
def research_step(state):
11
    # 研究节点
12
    return {"messages": state["messages"] + ["research done"]}
13

14
def code_step(state):
15
    # 编码节点
16
    return {"messages": state["messages"] + ["code done"]}
17

18
def review_step(state):
19
    # 审查节点
20
    return {"messages": state["messages"] + ["review done"]}
21

22
# 定义图
23
workflow = StateGraph(AgentState)
24

25
workflow.add_node("research", research_step)
26
workflow.add_node("code", code_step)
27
workflow.add_node("review", review_step)
28

29
# 定义边
30
workflow.add_edge("research", "code")
31
workflow.add_edge("code", "review")
32
workflow.add_edge("review", END)
33

34
# 入口点
35
workflow.set_entry_point("research")
36

37
# 编译
38
app = workflow.compile()
39

40
# 执行
41
result = app.invoke({"messages": []})

展开收起

8.2 条件分支#

1
def should_continue(state):
2
    """决定下一步"""
3
    last_message = state["messages"][-1]
4

5
    if "需要更多信息" in last_message:
6
        return "research"
7
    elif "需要修复" in last_message:
8
        return "code"
9
    elif "完成" in last_message:
10
        return END
11
    else:
12
        return "review"
13

14
# 添加条件边
15
workflow.add_conditional_edges(
16
    "review",
17
    should_continue,
18
    {
19
        "research": "research",
20
        "code": "code",
21
        END: END
22
    }
23
)

展开收起

8.3 循环工作流#

1
class IterativeState(TypedDict):
2
    code: str
3
    tests: str
4
    iteration: int
5
    passed: bool
6

7
def generate_code(state):
8
    if state["iteration"] == 0:
9
        code = generate_initial(state["tests"])
10
    else:
11
        code = fix_code(state["code"], state["test_results"])
12
    return {"code": code, "iteration": state["iteration"] + 1}
13

14
def run_tests(state):
15
    passed = execute_tests(state["code"], state["tests"])
16
    return {"passed": passed}
17

18
def should_continue(state):
19
    if state["passed"] or state["iteration"] >= 5:
20
        return "end"
21
    return "retry"
22

23
workflow = StateGraph(IterativeState)
24
workflow.add_node("code", generate_code)
25
workflow.add_node("test", run_tests)
26

27
workflow.add_edge("code", "test")
28
workflow.add_conditional_edges(
29
    "test",
30
    should_continue,
31
    {"end": END, "retry": "code"}
32
)
33
workflow.set_entry_point("code")
34

35
app = workflow.compile()

展开收起

9. Agent 评估#

9.1 评估维度#

1
class AgentEvaluator:
2
    """Agent 评估器"""
3

4
    METRICS = {
5
        'success_rate': '任务完成率',
6
        'efficiency': '效率（步数/时间）',
7
        'tool_accuracy': '工具调用准确率',
8
        'response_quality': '响应质量',
9
        'robustness': '鲁棒性',
10
        'cost': '成本（token/金钱）'
11
    }
12

13
    def evaluate(self, agent, test_cases):
14
        results = {
15
            'success': 0,
16
            'total': 0,
17
            'steps': [],
18
            'tokens': 0,
19
            'time': 0,
20
        }
21

22
        for case in test_cases:
23
            start_time = time.time()
24

25
            result = agent.run(case['input'])
26

27
            results['time'] += time.time() - start_time
28
            results['total'] += 1
29

30
            # 成功评估
31
            if self._check_success(result, case['expected']):
32
                results['success'] += 1
33

34
        return {
35
            'success_rate': results['success'] / results['total'],
36
            'avg_time': results['time'] / results['total'],
37
            'avg_steps': sum(results['steps']) / results['total'],
38
            'total_cost': results['tokens']
39
        }

展开收起

9.2 评估基准#

1
AGENT_BENCHMARKS = {
2
    'HotpotQA': '多跳问答',
3
    'FEVER': '事实核查',
4
    'SWE-bench': '软件工程',
5
    'HumanEval': '代码生成',
6
    'WebShop': '网页购物',
7
    'AlfWorld': '文本游戏',
8
    'Mind2Web': '网页操作',
9
    'GAIA': '通用助手',
10
    'AgentBench': '综合 Agent 评测',
11
}
12

13
class BenchmarkRunner:
14
    """基准测试运行器"""
15

16
    def run_benchmark(self, agent, benchmark_name):
17
        cases = self.load_benchmark(benchmark_name)
18
        return self.evaluate(agent, cases)

展开收起

9.3 失败分析#

1
class FailureAnalyzer:
2
    """失败分析器"""
3

4
    FAILURE_TYPES = [
5
        'planning_failure',     # 规划失败
6
        'tool_selection',      # 工具选错
7
        'tool_execution',      # 工具执行失败
8
        'reasoning_loop',      # 推理循环
9
        'context_overflow',    # 上下文溢出
10
        'hallucination',       # 幻觉
11
        'timeout',             # 超时
12
    ]
13

14
    def analyze(self, run_history):
15
        """分析失败原因"""
16
        for step in run_history:
17
            step['failure_type'] = self._classify_failure(step)
18

19
        distribution = Counter(
20
            s['failure_type'] for s in run_history
21
            if s.get('failure_type')
22
        )
23

24
        return {
25
            'distribution': distribution,
26
            'recommendations': self._generate_recommendations(distribution)
27
        }
28

29
    def _generate_recommendations(self, distribution):
30
        """生成改进建议"""
31
        recs = []
32

33
        if distribution.get('planning_failure', 0) > 0.2:
34
            recs.append("改进规划策略：使用 Plan-and-Execute")
35

36
        if distribution.get('tool_selection', 0) > 0.2:
37
            recs.append("改进工具描述：使 description 更清晰")
38

39
        if distribution.get('reasoning_loop', 0) > 0.1:
40
            recs.append("加入循环检测机制")
41

42
        return recs

展开收起

10. 生产实践#

10.1 性能优化#

10.1.1 Token 优化#

1
class TokenOptimizedAgent:
2
    """Token 优化的 Agent"""
3

4
    def __init__(self, llm, tools):
5
        self.llm = llm
6
        self.tools = tools
7

8
    def compress_prompt(self, prompt, max_tokens=4000):
9
        """压缩 prompt"""
10
        if self._count_tokens(prompt) > max_tokens:
11
            # 策略 1: 删除旧的工具结果
12
            prompt = self._remove_old_observations(prompt)
13

14
            # 策略 2: 总结长对话
15
            if self._count_tokens(prompt) > max_tokens:
16
                prompt = self._summarize_old_messages(prompt)
17

18
        return prompt
19

20
    def cache_system_prompt(self):
21
        """缓存 system prompt (如果模型支持)"""
22
        # Claude: cache_control
23
        # OpenAI: 自动部分优化
24
        ...
25

26
    def parallel_tool_calls(self):
27
        """并行工具调用"""
28
        # 多个独立工具同时执行
29
        ...

展开收起

10.1.2 模型路由#

1
class ModelRouter:
2
    """模型路由器"""
3

4
    def __init__(self):
5
        self.cheap_model = "gpt-3.5-turbo"
6
        self.expensive_model = "gpt-4"
7

8
    def choose_model(self, task_complexity):
9
        """根据复杂度选择模型"""
10
        if task_complexity < 0.3:
11
            return self.cheap_model
12
        else:
13
            return self.expensive_model
14

15
# 在 Agent 中使用
16
class CostOptimizedAgent:
17
    def __init__(self):
18
        self.router = ModelRouter()
19

20
    def run(self, task):
21
        # 简单任务用便宜模型
22
        if self._is_simple(task):
23
            model = self.router.cheap_model
24
        else:
25
            model = self.router.expensive_model
26

27
        return self._generate(model, task)

展开收起

10.2 可靠性#

1
class ReliableAgent:
2
    """可靠的 Agent"""
3

4
    def __init__(self, llm, tools, max_retries=3):
5
        self.llm = llm
6
        self.tools = tools
7
        self.max_retries = max_retries
8

9
    def run(self, task):
10
        for attempt in range(self.max_retries):
11
            try:
12
                return self._try_run(task)
13
            except Exception as e:
14
                if attempt < self.max_retries - 1:
15
                    # 记录错误，调整策略
16
                    self._handle_error(e)
17
                else:
18
                    raise
19

20
    def _handle_error(self, error):
21
        """错误处理"""
22
        if isinstance(error, ToolExecutionError):
23
            # 工具错误：重试或换工具
24
            self._fallback_strategy(error)
25
        elif isinstance(error, RateLimitError):
26
            # 限流：等待
27
            time.sleep(60)
28
        # ...

展开收起

10.3 监控#

1
class AgentMonitor:
2
    """Agent 监控器"""
3

4
    def __init__(self):
5
        self.metrics = []
6

7
    def record_step(self, step_data):
8
        """记录每步"""
9
        self.metrics.append({
10
            'timestamp': time.time(),
11
            'step_type': step_data['type'],
12
            'duration': step_data['duration'],
13
            'tokens': step_data['tokens'],
14
            'success': step_data['success']
15
        })
16

17
    def get_dashboard(self):
18
        """获取仪表盘数据"""
19
        return {
20
            'total_runs': len(self.metrics),
21
            'success_rate': self._calc_success_rate(),
22
            'avg_duration': self._avg('duration'),
23
            'p95_duration': self._percentile('duration', 95),
24
            'avg_tokens': self._avg('tokens'),
25
            'cost_estimate': self._estimate_cost()
26
        }

展开收起

10.4 安全性#

1
class SecureAgent:
2
    """安全的 Agent"""
3

4
    SENSITIVE_TOOLS = {'make_payment', 'delete_data', 'send_email'}
5

6
    def execute_tool(self, name, args):
7
        tool = self.tools.get(name)
8

9
        # 敏感操作需要确认
10
        if name in self.SENSITIVE_TOOLS:
11
            if not self._user_confirmation(name, args):
12
                return "用户拒绝"
13

14
        # 输入消毒
15
        sanitized_args = self._sanitize(name, args)
16

17
        # 执行
18
        return tool.execute(**sanitized_args)
19

20
    def _user_confirmation(self, action, args):
21
        """用户确认"""
22
        print(f"\n⚠️  确认执行: {action}({args})")
23
        return input("(yes/no): ").lower() == 'yes'
24

25
    def _sanitize(self, name, args):
26
        """输入消毒"""
27
        # SQL 注入检测
28
        # 路径遍历检测
29
        # 命令注入检测
30
        # ...
31
        return args

展开收起

11. 实战案例#

11.1 研究助手 Agent#

1
class ResearchAssistantAgent:
2
    """研究助手 Agent"""
3

4
    def __init__(self):
5
        self.tools = ToolRegistry()
6
        self._register_tools()
7
        self.memory = LongTermMemory(embedder, vector_store)
8

9
    def _register_tools(self):
10
        @self.tools.register(description="搜索 arXiv 论文")
11
        def search_arxiv(query, max_results=5):
12
            # ...
13
            return results
14

15
        @self.tools.register(description="阅读论文内容")
16
        def read_paper(paper_id):
17
            # ...
18
            return content
19

20
        @self.tools.register(description="总结段落")
21
        def summarize(text):
22
            return self.llm.summarize(text)
23

24
        @self.tools.register(description="生成引用")
25
        def cite(paper_id):
26
            return format_citation(paper_id)
27

28
    def research(self, question):
29
        plan = self._create_research_plan(question)
30

31
        report = []
32
        for step in plan:
33
            if step.tool == "search_arxiv":
34
                papers = self.tools.get("search_arxiv").execute(**step.args)
35
                # 选最相关的几篇
36
                for paper in papers[:3]:
37
                    content = self.tools.get("read_paper").execute(paper.id)
38
                    summary = self.tools.get("summarize").execute(content)
39
                    report.append({
40
                        'paper': paper,
41
                        'summary': summary,
42
                        'citation': self.tools.get("cite").execute(paper.id)
43
                    })
44

45
        return self._write_report(question, report)

展开收起

11.2 编码 Agent#

1
class CodingAgent:
2
    """编码 Agent - 类似 SWE-Agent"""
3

4
    def __init__(self, repo_path):
5
        self.tools = ToolRegistry()
6
        self._register_coding_tools()
7
        self.repo_path = repo_path
8

9
    def _register_coding_tools(self):
10
        @self.tools.register(description="读取文件")
11
        def read_file(path):
12
            with open(f"{self.repo_path}/{path}", 'r') as f:
13
                return f.read()
14

15
        @self.tools.register(description="编辑文件")
16
        def edit_file(path, old_string, new_string):
17
            # 应用编辑
18
            ...
19

20
        @self.tools.register(description="运行命令")
21
        def run_command(cmd):
22
            return subprocess.run(cmd, shell=True, capture_output=True)
23

24
        @self.tools.register(description="搜索代码")
25
        def search_code(pattern):
26
            return grep(self.repo_path, pattern)
27

28
    def solve_issue(self, issue_text):
29
        # 1. 理解问题
30
        # 2. 查看相关文件
31
        # 3. 设计修复
32
        # 4. 应用修改
33
        # 5. 运行测试
34
        # 6. 迭代到通过
35
        ...

展开收起

11.3 数据分析 Agent#

1
class DataAnalysisAgent:
2
    """数据分析 Agent"""
3

4
    def analyze(self, dataset_path, question):
5
        # 1. 理解数据集
6
        schema = self._read_schema(dataset_path)
7

8
        # 2. 规划分析步骤
9
        plan = self._plan_analysis(question, schema)
10

11
        # 3. 执行
12
        results = {}
13
        for step in plan:
14
            if step.tool == "execute_sql":
15
                results[step.id] = self.execute_sql(step.sql)
16
            elif step.tool == "create_chart":
17
                results[step.id] = self.create_chart(results[step.deps])
18
            elif step.tool == "interpret":
19
                results[step.id] = self.interpret(results[step.deps])
20

21
        # 4. 生成报告
22
        return self.generate_report(question, results)

展开收起

12. 框架对比#

12.1 主流 Agent 框架#

框架	特点	适合场景
LangChain	通用 LLM 框架	多用途
LangGraph	图状 Agent	复杂工作流
AutoGen	多 Agent 协作	对话式
CrewAI	角色扮演多 Agent	协作任务
LangSmith	可观测性	调试优化
Semantic Kernel	企业级	.NET 生态
SmolAgents	轻量级	简单应用

12.2 选型建议#

1
FRAMEWORK_SELECTION = """
2
选择 LangChain / LangGraph 当:
3
  - 需要 RAG、链式调用等
4
  - 工作流是图/树结构
5
  - 需要细粒度控制
6

7
选择 AutoGen 当:
8
  - 主要需求是多 Agent 对话
9
  - 人类在环 (human-in-the-loop)
10
  - 灵活的对话管理
11

12
选择 CrewAI 当:
13
  - 角色清晰的多 Agent 协作
14
  - 任务流程标准化
15
  - 希望快速搭建
16

17
选择 SmolAgents 当:
18
  - 简单任务
19
  - 学习 Agent 原理
20
  - 快速原型
21

22
自研 当:
23
  - 完全定制需求
24
  - 学习/教育目的
25
  - 已有基础设施
26
"""

展开收起

13. 未来发展方向#

13.1 当前挑战#

1
┌─────────────────────────────────────────────────────┐
2
│              Agent 的核心挑战                         │
3
├─────────────────────────────────────────────────────┤
4
│                                                      │
5
│  1. 可靠性                                            │
6
│     - 长链任务容易出错                                │
7
│     - 需要自我纠错                                    │
8
│                                                      │
9
│  2. 效率                                              │
10
│     - Token 消耗巨大                                  │
11
│     - 需要优化策略                                    │
12
│                                                      │
13
│  3. 评估                                              │
14
│     - 难以量化性能                                    │
15
│     - 需要更好基准                                    │
16
│                                                      │
17
│  4. 安全                                              │
18
│     - 自主行动的边界                                  │
19
│     - 防止滥用                                        │
20
│                                                      │
21
│  5. 协同                                              │
22
│     - 多 Agent 协调复杂                              │
23
│     - 通信协议尚未成熟                                │
24
│                                                      │
25
└─────────────────────────────────────────────────────┘

展开收起

13.2 未来方向#

更智能的规划：从被动响应到主动规划
更强记忆：跨会话、跨任务的长期记忆
更好的工具学习：自动发现和使用新工具
多模态融合：视觉、听觉、触觉的统一
自演化 Agent：根据使用自动改进

13.3 Agent OS 愿景#

1
未来的 Agent OS:
2
┌────────────────────────────────────┐
3
│        Agent Operating System     │
4
├────────────────────────────────────┤
5
│ - 标准化工具调用                    │
6
│ - 统一的记忆接口                    │
7
│ - 跨平台 Agent 运行时              │
8
│ - Agent 网络通信                    │
9
│ - 安全沙箱                          │
10
└────────────────────────────────────┘
11
         │           │
12
         ▼           ▼
13
    ┌─────────┐ ┌─────────┐
14
    │ Agent A │ │ Agent B │
15
    └─────────┘ └─────────┘

展开收起

14. 核心概念总结#

14.1 Agent 的核心定义#

1
AGENT_EQUATION = """
2
Agent = Perception + Memory + Reasoning + Action + Tools
3

4
其中:
5
- Perception: 感知环境
6
- Memory: 记忆系统
7
- Reasoning: 推理决策
8
- Action: 行动执行
9
- Tools: 工具集
10

11
Agent 的核心特征:
12
1. 自主性 (Autonomy)
13
2. 反应性 (Reactivity)
14
3. 主动性 (Proactivity)
15
4. 社交性 (Social Ability)
16
"""

展开收起

14.2 关键架构模式#

1
ARCHITECTURE_PATTERNS = """
2
1. 单 Agent
3
   - ReAct
4
   - Plan-Execute
5
   - Reflexion
6
   - ReWOO
7

8
2. Multi-Agent
9
   - Supervisor
10
   - Hierarchical
11
   - Peer-to-Peer
12
   - Collaborative
13

14
3. 设计原则
15
   - 模块化
16
   - 可观测
17
   - 容错
18
   - 可扩展
19
"""

展开收起

14.3 数学化表达#

Agent 决策过程：

at=π(st,Mt,T)a_t = \pi(s_t, M_t, T)at=π(st,Mt,T)

其中 ata_tat 是行动，sts_tst 是状态，MtM_tMt 是记忆，TTT 是工具集。

Plan-Execute 计划评估：

Plan Quality=f(completeness,efficiency,robustness)\text{Plan Quality} = f(\text{completeness}, \text{efficiency}, \text{robustness})Plan Quality=f(completeness,efficiency,robustness)

记忆检索相似度：

sim(q,m)=vq⋅vm∣∣vq∣∣⋅∣∣vm∣∣\text{sim}(q, m) = \frac{v_q \cdot v_m}{||v_q|| \cdot ||v_m||}sim(q,m)=∣∣vq∣∣⋅∣∣vm∣∣vq⋅vm

15. 实战建议#

15.1 设计 Agent 的最佳实践#

1
BEST_PRACTICES = """
2
1. 从简单开始
3
   - 先实现 ReAct
4
   - 逐步添加复杂功能
5
   - 测试每个组件
6

7
2. 工具设计
8
   - 单一职责
9
   - 清晰描述
10
   - 错误处理
11

12
3. 记忆策略
13
   - 短期+长期结合
14
   - 定期压缩
15
   - 重要信息优先
16

17
4. 规划策略
18
   - 简单任务不需要复杂规划
19
   - 复杂任务使用 Plan-Execute
20
   - 加入容错
21

22
5. 评估驱动
23
   - 先定义评估标准
24
   - 持续监控
25
   - 数据驱动改进
26
"""

展开收起

15.2 常见反模式#

1
ANTI_PATTERNS = """
2
1. 过度复杂
3
   - 一开始就设计 10 个 Agent
4
   - 应该从单 Agent 开始
5

6
2. 缺乏测试
7
   - 没有评估标准
8
   - 难以改进
9

10
3. 忽视成本
11
   - 不监控 Token 使用
12
   - 应该控制成本
13

14
4. 没有 fallback
15
   - 单点失败
16
   - 应该设计降级方案
17

18
5. 模糊目标
19
   - Agent 不知道目标
20
   - 应该明确任务分解
21
"""

展开收起

15.3 学习路径#

1
LEARNING_ROADMAP = """
2
1. 基础 (1-2 周)
3
   - 理解 LLM 调用
4
   - 学会 Function Calling
5
   - 实现 ReAct Agent
6

7
2. 进阶 (2-4 周)
8
   - 加入 Memory 系统
9
   - 学习 Plan-Execute
10
   - 实现复杂 Agent
11

12
3. 多 Agent (4-8 周)
13
   - 学习 Multi-Agent 概念
14
   - 掌握 1-2 个框架（LangGraph/AutoGen）
15
   - 实现协作系统
16

17
4. 专家 (3 个月+)
18
   - 深入研究
19
   - 性能优化
20
   - 生产部署
21
   - 安全考虑
22
"""

展开收起

16. 总结#

Agent 架构的核心#

AI Agent 不只是”会调用工具的 LLM”，而是一个完整的智能系统：

感知：理解环境
记忆：积累知识
推理：规划决策
行动：执行任务

Agent 的发展趋势#

阶段	时间	特征
LLM	2020-2022	文本生成
Tool-LLM	2023	工具调用
ReAct Agent	2022-2023	思考+行动
Multi-Agent	2023-2024	多代理协作
Agent OS	2025+	平台化

结语#

AI Agent 是 LLM 走向通用人工智能的关键一步。掌握 Agent 架构设计，理解不同模式、Memory 系统、Multi-Agent 协作，将为你打开 AI 应用开发的新天地。

无论技术如何演进，Perception → Reasoning → Action 的循环思想始终是 Agent 系统设计的核心范式。

参考资料#

Yao, S., et al. (2022). “ReAct: Synergizing Reasoning and Acting in Language Models.” ICLR.
Shinn, N., et al. (2023). “Reflexion: Language Agents with Verbal Reinforcement Learning.” NeurIPS.
Wei, J., et al. (2022). “Chain-of-Thought Prompting Elicits Reasoning in Large Language Models.” NeurIPS.
Liu, J., et al. (2023). “ReWOO: Decoupling Reasoning from Observations for Efficient Augmented Language Models.” arXiv.
Schick, T., et al. (2023). “Toolformer: Language Models Can Teach Yourself to Use Tools.” NeurIPS.
Wang, L., et al. (2024). “A Survey on Large Language Model based Autonomous Agents.” arXiv.
Xi, Z., et al. (2023). “The Rise and Potential of Large Language Model Based Agents.” arXiv.
Wu, Q., et al. (2023). “AutoGen: Enabling Next-Gen LLM Applications via Multi-Agent Conversation.” arXiv.
LangChain. (2024). “LangGraph: Multi-Agent Workflows.” Documentation.
Anthropic. (2024). “Building Effective Agents with Claude.” Engineering Blog.
Microsoft. (2024). “AutoGen: Enabling Next-Gen LLM Applications.” arXiv.
Patil, S., et al. (2024). “Function Calling Best Practices.” OpenAI Cookbook.
Liu, F., et al. (2024). “AgentBench: Evaluating LLMs as Agents.” ICLR.

深入理解 Chain-of-Thought (CoT)：让大模型学会"思考"

Thu, 16 Jul 2026 00:00:00 GMT

1. Chain-of-Thought 的引入#

1.1 大模型为什么不会”思考”？#

长期以来，深度学习模型在处理需要多步骤推理的任务时表现不佳：

简单的算术题：模型可能给出错误答案
复杂逻辑推理：模型常常”跳步”推理
多约束优化：模型难以分解问题

核心问题：

大语言模型本质上是”下一个 token 预测器”——它们直接生成答案，而不展示思考过程。

这就像一个学生不展示任何计算过程，直接写出最终答案——你无法判断他是真的理解了，还是在瞎猜。

1.2 CoT 的核心思想#

Chain-of-Thought (CoT，思维链) 是 Wei et al. (2022) 在 Google Brain 提出的突破性技术：

不要直接问模型答案，而是让模型”一步一步地思考”——展示推理过程，答案自然浮现。

直觉类比：

1
普通提问：
2
  "12 × 15 等于多少？"
3
  模型直接回答: "180"  ← 可能是对的，也可能是错的
4

5
CoT 提问：
6
  "12 × 15 等于多少？请一步一步计算"
7
  模型回答: "12 × 15 = 12 × 10 + 12 × 5 = 120 + 60 = 180"
8
  ← 推理过程透明，更可能正确

1.3 CoT 发展简史#

1
2020: GPT-3 时代
2
  │  大模型展现强大的 few-shot 能力
3
  │  但推理能力仍然有限
4
  │
5
2022.1: Chain-of-Thought Prompting (Wei et al.)
6
  │  Google Brain 首次提出 CoT 概念
7
  │  在算术、常识、符号推理任务上取得突破
8
  │
9
2022.3: Zero-Shot CoT (Kojima et al.)
10
  │  仅用一句"让我们一步一步思考"就能激发推理
11
  │
12
2022.4: Self-Consistency (Wang et al.)
13
  │  多个推理路径投票，提高鲁棒性
14
  │
15
2022.9: Tree of Thoughts (Yao et al.)
16
  │  将推理组织成树形搜索
17
  │
18
2022.10: ReAct (Yao et al.)
19
  │  推理 + 行动，与工具结合
20
  │
21
2023: Auto-CoT, Active-Prompt
22
  │  自动选择最有信息量的演示样例
23
  │
24
2023-2024: CoT + 工具使用
25
  │  Toolformer, PAL, Program-aided Reasoning
26
  │
27
2024-2025: 多模态 CoT
28
  │  视觉 + 语言的协同推理
29
  │
30
2025+: 推理模型崛起
31
  │  OpenAI o1, DeepSeek-R1, Claude 3.7
32
  │  CoT 训练进模型权重

展开收起

1.4 CoT 的革命性意义#

CoT 是大模型推理能力的第一次重大飞跃：

维度	直接回答	CoT
准确率	中等	显著提升
可解释性	黑盒	推理过程透明
错误定位	难以调试	可以定位错误步骤
复杂问题	容易失败	可以分解处理
推理深度	浅层	多步骤深度推理

2. CoT 的核心原理#

2.1 为什么 CoT 有效？#

2.1.1 认知负荷理论#

CoT 通过分解复杂任务来降低认知负荷：

1
直接回答：
2
  [困难推理 + 直接输出]
3
  └─────── 高认知负荷 ──────┐
4
                          ▼
5
                    可能失败
6

7
CoT 分解：
8
  步骤1: 提取关键信息  ── 低负荷
9
  步骤2: 应用规则    ── 低负荷
10
  步骤3: 推导中间结果 ── 低负荷
11
  步骤4: 得出最终答案 ── 低负荷
12
  └─────── 总负荷可管理 ──────┐
13
                          ▼
14
                    高成功率

2.1.2 显式中间表示#

CoT 在回答过程中生成显式的中间 token，这些 token：

为后续推理提供”工作记忆”
引导模型走向正确答案
让复杂的”思维”成为可能的”思考路径”

2.1.3 计算假设#

Transformer 中的 CoT 实际上相当于增加了计算深度：

想象一个 L 层 Transformer：

直接回答：使用 L 层计算直接预测答案

CoT：使用 L 层生成 k 个推理 token，再用 k 层计算预测答案

有效计算量从 L 增加到 L + k！

2.2 CoT 的形式化定义#

给定一个问题 qqq 和 few-shot 演示 D={(x1,c1,y1),…,(xn,cn,yn)}D = \{(x_1, c_1, y_1), \ldots, (x_n, c_n, y_n)\}D={(x1,c1,y1),…,(xn,cn,yn)}，其中 cic_ici 是思维链，yiy_iyi 是最终答案。

CoT 模型生成的输出：

CoT(q)=LLM(q,D)=(cq,yq)\text{CoT}(q) = \text{LLM}(q, D) = (c_q, y_q)CoT(q)=LLM(q,D)=(cq,yq)

其中 cq=(z1,z2,…,zk)c_q = (z_1, z_2, \ldots, z_k)cq=(z1,z2,…,zk) 是推理步骤，yqy_qyq 是最终答案。

2.3 CoT 的理论基础#

2.3.1 涌现能力 (Emergent Ability)#

CoT 是典型的涌现能力：

小模型 (< 10B)：CoT 几乎无效果，甚至有害
中等模型 (10B-50B)：CoT 开始展现优势
大模型 (> 50B)：CoT 带来巨大提升

1
性能提升
2
│
3
│                ●  大模型 + CoT
4
│
5
│         ●  大模型 (无 CoT)
6
│
7
│    ●  小模型 + CoT  ≈  小模型 (无 CoT)
8
│
9
└──────────────────────────────────────
10
  小模型         中模型         大模型

2.3.2 CoT 的数学解释#

从理论上讲，CoT 可以被理解为：

隐式计算：每个中间 token 都是一次”软计算”
信息瓶颈：通过中间表示将复杂问题分解
递归分解：每一步将问题降维

1
# 简化的 CoT 计算模型
2
def cot_recursive_decomposition(problem, max_depth=5):
3
    """
4
    CoT 作为递归问题分解
5
    """
6
    if max_depth == 0 or is_simple(problem):
7
        return solve_directly(problem)
8

9
    # 分解问题
10
    sub_problems = decompose(problem)
11

12
    # 递归求解每个子问题
13
    sub_answers = []
14
    for sub_p in sub_problems:
15
        sub_answer = cot_recursive_decomposition(sub_p, max_depth - 1)
16
        sub_answers.append(sub_answer)
17

18
    # 组合答案
19
    final_answer = combine(sub_answers)
20
    return final_answer

展开收起

3. CoT 的基础方法#

3.1 Few-Shot CoT (原始方法)#

Few-Shot CoT 是 Wei et al. 提出的原始方法，通过手工设计的演示激发 CoT：

1
# Few-Shot CoT Prompt 示例
2

3
prompt = """
4
问题：小明有 5 个苹果。他又买了 3 袋苹果，每袋 4 个。现在小明一共有多少个苹果？
5

6
推理：小明原来有 5 个苹果。他买了 3 袋，每袋 4 个，所以新买了 3 × 4 = 12 个苹果。总共：5 + 12 = 17 个苹果。
7

8
答案：17 个苹果
9

10
问题：一本书有 240 页。小红第一天读了 1/3，第二天读了剩下的 1/4。第三天她需要读多少页才能读完？
11

12
推理：小红第一天读了 240 × 1/3 = 80 页。剩下 240 - 80 = 160 页。第二天她读了剩下 160 × 1/4 = 40 页。最后剩下 160 - 40 = 120 页。
13

14
答案：120 页
15

16
问题：{用户问题}
17

18
推理：
19
"""

展开收起

关键要素：

多个示例：3-8 个 few-shot 示例
推理过程：每个示例都有完整的思维链
格式一致：保持答案格式的一致性

3.2 Zero-Shot CoT (零样本)#

Zero-Shot CoT 是 Kojima et al. (2022) 的重大发现：

只需要在 prompt 中加一句”让我们一步一步思考”，模型就会自动进行 CoT 推理！

1
# Zero-Shot CoT
2

3
prompt_direct = "12 × 15 等于多少？"
4
# 输出: "180" (可能错误)
5

6
prompt_cot = "12 × 15 等于多少？让我们一步一步思考。"
7
# 输出: "12 × 15 = 12 × 10 + 12 × 5 = 120 + 60 = 180" (更准确)

经典触发短语：

中文：“让我们一步一步思考”
英文：“Let’s think step by step”
“Take a deep breath and work through this step by step”
“Think through this carefully”

3.3 对比：Few-Shot vs Zero-Shot#

维度	Few-Shot CoT	Zero-Shot CoT
是否需要演示	是（3-8 个）	否
设计成本	高	低
推理质量	高（取决于演示质量）	中等
适用场景	生产环境	快速原型
灵活性	低（受演示影响）	高

3.4 Auto-CoT (自动化)#

Auto-CoT 自动选择或生成演示样例：

1
class AutoCoT:
2
    """
3
    自动思维链：自动选择/生成多样化的演示
4
    """
5
    def __init__(self, llm, dataset):
6
        self.llm = llm
7
        self.dataset = dataset
8

9
    def sample_diverse_questions(self, n=8):
10
        """
11
        使用聚类选择多样化的代表性问题
12
        """
13
        # 1. 编码所有问题
14
        embeddings = self.embed_questions(self.dataset.questions)
15

16
        # 2. K-means 聚类
17
        clusters = self.kmeans(embeddings, n_clusters=n)
18

19
        # 3. 从每个聚类中心选择代表问题
20
        diverse_questions = []
21
        for cluster_id in range(n):
22
            center = clusters[cluster_id].center
23
            closest = self.find_closest_question(center, embeddings)
24
            diverse_questions.append(closest)
25

26
        return diverse_questions
27

28
    def generate_cot_for_question(self, question):
29
        """使用 Zero-Shot CoT 生成推理链"""
30
        prompt = f"{question}\n让我们一步一步思考。"
31
        reasoning = self.llm.generate(prompt)
32
        return reasoning
33

34
    def build_prompt(self, user_question):
35
        """构建 Auto-CoT prompt"""
36
        # 选择多样化问题
37
        diverse_qs = self.sample_diverse_questions()
38

39
        # 为每个问题生成 CoT
40
        demonstrations = []
41
        for q in diverse_qs:
42
            cot = self.generate_cot_for_question(q)
43
            demonstrations.append({
44
                'question': q,
45
                'reasoning': cot
46
            })
47

48
        # 构建最终 prompt
49
        prompt = self.format_demonstrations(demonstrations)
50
        prompt += f"\n问题：{user_question}\n推理："
51

52
        return prompt

展开收起

3.5 Self-Consistency (自一致性)#

Self-Consistency 通过多次采样 + 投票提高 CoT 的鲁棒性：

1
class SelfConsistency:
2
    """
3
    自一致性：从多个推理路径中投票
4
    """
5
    def __init__(self, llm, num_samples=10, temperature=0.7):
6
        self.llm = llm
7
        self.num_samples = num_samples
8
        self.temperature = temperature
9

10
    def generate(self, question):
11
        """生成多个推理路径并投票"""
12
        # 1. 生成多个 CoT 推理
13
        reasoning_paths = []
14
        for i in range(self.num_samples):
15
            prompt = self._build_prompt(question)
16
            response = self.llm.generate(
17
                prompt,
18
                temperature=self.temperature,
19
                do_sample=True
20
            )
21
            reasoning_paths.append(self._extract_answer(response))
22

23
        # 2. 多数投票
24
        answer_counts = Counter(reasoning_paths)
25
        final_answer = answer_counts.most_common(1)[0][0]
26

27
        return {
28
            'answer': final_answer,
29
            'reasoning_paths': reasoning_paths,
30
            'confidence': answer_counts[final_answer] / self.num_samples
31
        }
32

33
    def _extract_answer(self, response):
34
        """从响应中提取最终答案"""
35
        # 多种答案提取策略
36
        patterns = [
37
            r'答案是[：:]\s*(.+?)(?:\n|$)',
38
            r'答案[：:]\s*(.+?)(?:\n|$)',
39
            r'####\s*(.+?)(?:\n|$)',
40
        ]
41
        for pattern in patterns:
42
            match = re.search(pattern, response)
43
            if match:
44
                return match.group(1).strip()
45
        return response.strip()

展开收起

自一致性的优势：

不依赖单一推理路径：即使某个推理出错，多数投票仍能选出正确答案
置信度估计：可以给出答案的可信度
与 Temperature 配合：高 temperature 增加多样性

1
推理路径 1:  A → B → C → 答案 X
2
推理路径 2:  A → D → E → 答案 Y  ← 多数
3
推理路径 3:  A → F → G → 答案 Y  ← 多数
4
推理路径 4:  A → H → I → 答案 Y  ← 多数
5
推理路径 5:  A → J → K → 答案 X
6

7
多数投票结果: Y (3 票)

4. 高级 CoT 方法#

4.1 Tree of Thoughts (ToT)#

Tree of Thoughts (ToT) 将 CoT 扩展为树形搜索：

1
class TreeOfThoughts:
2
    """
3
    树形思维：探索多个推理路径
4
    """
5
    def __init__(self, llm, branching_factor=3, max_depth=5):
6
        self.llm = llm
7
        self.branching_factor = branching_factor
8
        self.max_depth = max_depth
9
        self.tree = ThoughtTree()
10

11
    def generate_thoughts(self, state):
12
        """在当前状态下生成多个可能的下一步"""
13
        prompt = f"""
14
        当前状态：{state}
15

16
        请提出 {self.branching_factor} 个可能的下一步推理：
17
        """
18
        response = self.llm.generate(prompt)
19
        thoughts = self.parse_thoughts(response)
20
        return thoughts
21

22
    def evaluate_state(self, state):
23
        """评估当前状态的优劣"""
24
        prompt = f"""
25
        状态：{state}
26

27
        这个状态距离最终答案有多远？评估（1-10）：
28
        """
29
        score = int(self.llm.generate(prompt).strip())
30
        return score
31

32
    def search(self, problem):
33
        """使用 BFS/DFS 搜索最佳推理路径"""
34
        # 初始化
35
        root = ThoughtNode(state=problem, depth=0)
36
        self.tree.add(root)
37

38
        for depth in range(self.max_depth):
39
            current_nodes = self.tree.get_nodes_at_depth(depth)
40
            new_nodes = []
41

42
            for node in current_nodes:
43
                # 生成多个下一步
44
                thoughts = self.generate_thoughts(node.state)
45

46
                for thought in thoughts:
47
                    new_state = node.state + " → " + thought
48
                    score = self.evaluate_state(new_state)
49

50
                    new_node = ThoughtNode(
51
                        state=new_state,
52
                        parent=node,
53
                        depth=depth + 1,
54
                        score=score
55
                    )
56
                    new_nodes.append(new_node)
57

58
            # 剪枝：保留 top-k 最佳节点
59
            new_nodes.sort(key=lambda x: x.score, reverse=True)
60
            top_nodes = new_nodes[:self.branching_factor]
61

62
            for node in top_nodes:
63
                self.tree.add(node)
64

65
        # 选择最佳路径
66
        best_path = self.tree.get_best_path()
67
        return best_path

展开收起

ToT 适用场景：

数学证明（需要回溯）
战略游戏（需要规划）
谜题求解（需要探索）

4.2 Graph of Thoughts (GoT)#

Graph of Thoughts 将推理建模为图结构，允许合并和重用推理路径：

1
class GraphOfThoughts:
2
    """
3
    图思维：在推理图中合并、重用
4
    """
5
    def __init__(self, llm):
6
        self.llm = llm
7
        self.graph = ThoughtGraph()
8

9
    def transform(self, nodes, operation):
10
        """
11
        图操作：生成、聚合、改进、验证
12
        """
13
        if operation == 'aggregate':
14
            # 合并多个想法
15
            prompt = f"合并这些想法：{nodes}"
16
            new_node = self.llm.generate(prompt)
17
            return [new_node]
18

19
        elif operation == 'improve':
20
            # 改进单个想法
21
            prompt = f"改进这个想法：{nodes[0]}"
22
            return [self.llm.generate(prompt)]
23

24
        elif operation == 'generate':
25
            # 生成新想法
26
            prompt = f"基于：{nodes}\n生成新想法："
27
            return [self.llm.generate(prompt) for _ in range(3)]
28

29
    def solve(self, problem):
30
        """通过图操作解决"""
31
        # 添加初始节点
32
        initial = self.graph.add(problem)
33

34
        # 生成多个思路
35
        thoughts = self.transform([initial], 'generate')
36

37
        # 聚合思路
38
        merged = self.transform(thoughts, 'aggregate')
39

40
        # 改进思路
41
        refined = self.transform(merged, 'improve')
42

43
        return refined[0]

展开收起

4.3 ReAct (Reasoning + Acting)#

ReAct 将 CoT 与工具使用结合：

1
class ReActAgent:
2
    """
3
    ReAct: 推理 + 行动
4
    """
5
    def __init__(self, llm, tools):
6
        self.llm = llm
7
        self.tools = tools
8
        self.memory = []
9

10
    def think(self, observation):
11
        """思考下一步行动"""
12
        prompt = self._build_prompt(observation)
13
        thought = self.llm.generate(prompt)
14
        return thought
15

16
    def act(self, thought):
17
        """根据思考执行行动"""
18
        # 解析行动
19
        action_name, action_input = self.parse_action(thought)
20

21
        if action_name in self.tools:
22
            result = self.tools[action_name](action_input)
23
            return f"Observation: {result}"
24
        else:
25
            return f"Unknown action: {action_name}"
26

27
    def step(self, question):
28
        """执行一步 ReAct"""
29
        # 思考
30
        thought = self.think(question)
31
        self.memory.append(f"Thought: {thought}")
32

33
        # 行动
34
        if "Final Answer" in thought:
35
            return self.extract_answer(thought)
36

37
        observation = self.act(thought)
38
        self.memory.append(observation)
39

40
        return None
41

42
    def solve(self, question, max_steps=10):
43
        for step in range(max_steps):
44
            result = self.step(question)
45
            if result:
46
                return result
47
        return "Failed to find answer"

展开收起

ReAct 的 prompt 模板：

1
问题：{question}
2

3
Thought 1: 我需要查找 X 信息
4
Action 1: Search[X]
5
Observation 1: ...
6

7
Thought 2: 基于观察，我需要进一步查询 Y
8
Action 2: Lookup[Y]
9
Observation 2: ...
10

11
Thought 3: 现在我有了足够信息，可以给出答案
12
Action 3: Finish[答案]

4.4 Self-Refine#

Self-Refine 让模型自我评估并改进：

1
class SelfRefine:
2
    """
3
    自我精炼：迭代评估和改进
4
    """
5
    def __init__(self, llm, max_iterations=3):
6
        self.llm = llm
7
        self.max_iterations = max_iterations
8

9
    def generate_initial(self, question):
10
        """初始生成"""
11
        return self.llm.generate(question + "\n让我们一步一步思考。")
12

13
    def critique(self, response):
14
        """自我评估"""
15
        prompt = f"""
16
        以下是对问题的回答：
17
        {response}
18

19
        请评估这个回答，指出：
20
        1. 推理中是否有错误？
21
        2. 是否遗漏了重要步骤？
22
        3. 如何改进？
23

24
        评估：
25
        """
26
        return self.llm.generate(prompt)
27

28
    def improve(self, response, critique):
29
        """根据评估改进"""
30
        prompt = f"""
31
        原回答：{response}
32
        评估：{critique}
33

34
        请根据评估改进原回答：
35
        """
36
        return self.llm.generate(prompt)
37

38
    def solve(self, question):
39
        """通过自我精炼求解"""
40
        response = self.generate_initial(question)
41

42
        for i in range(self.max_iterations):
43
            critique = self.critique(response)
44

45
            if self._is_satisfied(critique):
46
                break
47

48
            response = self.improve(response, critique)
49

50
        return response

展开收起

4.5 各种 CoT 方法对比#

方法	复杂度	准确性	适用场景	计算成本
Direct	低	低	简单问题	1×
Zero-Shot CoT	低	中	一般推理	1.2×
Few-Shot CoT	中	高	生产环境	1.5×
Self-Consistency	中	高	关键任务	10×
ToT	高	很高	复杂规划	10-100×
ReAct	高	高	工具使用	5-20×
Self-Refine	中	高	优化质量	3-5×

5. CoT 的工程实现#

5.1 Prompt 工程最佳实践#

1
class CoTPromptBuilder:
2
    """
3
    高质量的 CoT Prompt 构建器
4
    """
5
    def __init__(self):
6
        self.templates = {}
7

8
    def build_math_cot_prompt(self, question):
9
        """数学推理的 CoT prompt"""
10
        template = """
11
你是一位数学专家。请通过清晰的步骤解决问题。
12

13
问题: {question}
14

15
让我们一步一步思考:
16
- 第一步: 提取已知信息
17
- 第二步: 确定解题方法
18
- 第三步: 应用公式或定理
19
- 第四步: 计算中间结果
20
- 第五步: 验证答案
21

22
推理:
23
"""
24
        return template.format(question=question)
25

26
    def build_logic_cot_prompt(self, question):
27
        """逻辑推理的 CoT prompt"""
28
        template = """
29
请分析这个逻辑问题，展示完整推理过程。
30

31
问题: {question}
32

33
推理框架:
34
1. 识别前提
35
2. 列出约束条件
36
3. 推导可能结论
37
4. 排除矛盾选项
38
5. 得出最终答案
39

40
推理:
41
"""
42
        return template.format(question=question)
43

44
    def build_code_cot_prompt(self, question):
45
        """代码生成的 CoT prompt"""
46
        template = """
47
请编写解决以下问题的代码，并解释思路。
48

49
问题: {question}
50

51
步骤:
52
1. 理解问题需求
53
2. 设计算法
54
3. 编写代码
55
4. 测试边界条件
56

57
代码与解释:
58
"""
59
        return template.format(question=question)

展开收起

5.2 答案提取#

1
class AnswerExtractor:
2
    """
3
    从 CoT 响应中提取结构化答案
4
    """
5
    def __init__(self):
6
        self.patterns = [
7
            # 中文模式
8
            r'答案是[：:]?\s*(.+?)(?:\n|$)',
9
            r'最终答案[：:]?\s*(.+?)(?:\n|$)',
10
            r'结论[：:]?\s*(.+?)(?:\n|$)',
11
            r'所以[，,]?\s*(.+?)(?:\n|$)',
12

13
            # 英文模式
14
            r'[Tt]he answer is\s*[:：]?\s*(.+?)(?:\n|$)',
15
            r'[Ff]inal [Aa]nswer\s*[:：]?\s*(.+?)(?:\n|$)',
16
            r'[Tt]herefore[，,]?\s*(.+?)(?:\n|$)',
17
            r'####\s*(.+?)(?:\n|$)',
18

19
            # LaTeX boxed
20
            r'\\boxed\{(.+?)\}',
21
        ]
22

23
    def extract(self, response):
24
        """从响应中提取答案"""
25
        for pattern in self.patterns:
26
            match = re.search(pattern, response)
27
            if match:
28
                return match.group(1).strip()
29

30
        # 如果没有匹配，尝试使用 LLM 提取
31
        return self._extract_with_llm(response)
32

33
    def extract_with_confidence(self, response):
34
        """提取答案并给出置信度"""
35
        candidates = []
36
        for pattern in self.patterns:
37
            matches = re.findall(pattern, response)
38
            for match in matches:
39
                candidates.append(match.strip())
40

41
        if not candidates:
42
            return None, 0.0
43

44
        # 选择最后一个或最常出现的
45
        from collections import Counter
46
        counts = Counter(candidates)
47
        most_common = counts.most_common(1)[0]
48

49
        confidence = most_common[1] / len(candidates) if candidates else 0
50
        return most_common[0], confidence

展开收起

5.3 推理评估#

1
class CoTEvaluator:
2
    """
3
    CoT 推理质量评估器
4
    """
5
    def __init__(self, llm):
6
        self.llm = llm
7

8
    def evaluate_correctness(self, question, reasoning, answer, expected):
9
        """评估答案正确性"""
10
        # 1. 精确匹配
11
        if str(answer).strip() == str(expected).strip():
12
            return 1.0
13

14
        # 2. 数值比较（对数学题）
15
        try:
16
            if abs(float(answer) - float(expected)) < 1e-6:
17
                return 1.0
18
        except (ValueError, TypeError):
19
            pass
20

21
        # 3. 语义相似度（使用 LLM 判断）
22
        prompt = f"""
23
        预期答案：{expected}
24
        实际答案：{answer}
25
        这两个答案在语义上是否等价？
26
        回答：是/否
27
        """
28
        response = self.llm.generate(prompt).strip()
29
        return 1.0 if "是" in response or "yes" in response.lower() else 0.0
30

31
    def evaluate_reasoning_quality(self, reasoning):
32
        """评估推理质量"""
33
        criteria = {
34
            'logical_consistency': '推理步骤是否逻辑一致？',
35
            'step_clarity': '每一步是否清晰？',
36
            'completeness': '是否包含所有必要步骤？',
37
            'efficiency': '是否高效，没有冗余步骤？'
38
        }
39

40
        scores = {}
41
        for criterion, question in criteria.items():
42
            prompt = f"推理过程：{reasoning}\n\n{question}（评分 1-10）"
43
            score = self.llm.generate(prompt)
44
            try:
45
                scores[criterion] = float(score.strip().split()[0]) / 10
46
            except (ValueError, IndexError):
47
                scores[criterion] = 0.5
48

49
        return sum(scores.values()) / len(scores)

展开收起

6. CoT 与工具的结合#

6.1 PAL (Program-Aided Language Models)#

PAL 让 LLM 生成代码来辅助推理：

1
class PAL:
2
    """
3
    Program-Aided Reasoning
4
    """
5
    def __init__(self, llm, code_executor):
6
        self.llm = llm
7
        self.executor = code_executor
8

9
    def solve(self, question):
10
        """通过生成代码来求解"""
11
        prompt = f"""
12
        问题：{question}
13

14
        请生成 Python 代码来解决这个问题。代码应该包含推理过程。
15

16
        代码：
17
        ```python
18
        """
19
        code = self.llm.generate(prompt)
20

21
        # 执行代码
22
        result = self.executor.execute(code)
23
        return result

展开收起

PAL 示例：

1
# PAL 解决数学问题
2

3
question = "一家餐厅周一卖出 15 个汉堡，周二比周一多卖出 1/3，周三是周二的 80%。三天共卖出多少汉堡？"
4

5
# LLM 生成代码
6
generated_code = """
7
# 周一
8
monday = 15
9
# 周二: 周一 + 周一的 1/3
10
tuesday = monday + monday * (1/3)
11
# 周三: 周二的 80%
12
wednesday = tuesday * 0.8
13
# 总数
14
total = monday + tuesday + wednesday
15
print(f"周一: {monday}, 周二: {tuesday}, 周三: {wednesday}, 总数: {total}")
16
"""
17

18
# 执行得到答案
19
# 输出: 周一: 15, 周二: 20.0, 周三: 16.0, 总数: 51.0

展开收起

6.2 Toolformer#

Toolformer 让模型学会自动调用工具：

1
class ToolformerAgent:
2
    """
3
    自主工具调用
4
    """
5
    def __init__(self, llm, tools):
6
        self.llm = llm
7
        self.tools = tools
8
        self.tool_descriptions = self._format_tools(tools)
9

10
    def decide_tool_use(self, context, question):
11
        """决定是否使用工具"""
12
        prompt = f"""
13
        可用工具:
14
        {self.tool_descriptions}
15

16
        当前问题: {question}
17

18
        请判断是否需要使用工具。如果需要，调用工具；如果不需要，直接回答。
19

20
        决定：
21
        """
22
        decision = self.llm.generate(prompt)
23

24
        if "[Tool:" in decision:
25
            tool_call = self.parse_tool_call(decision)
26
            return self.execute_tool(tool_call)
27
        else:
28
            return decision
29

30
    def execute_tool(self, tool_call):
31
        """执行工具调用"""
32
        tool_name = tool_call['name']
33
        tool_input = tool_call['input']
34

35
        if tool_name in self.tools:
36
            return self.tools[tool_name](tool_input)
37
        return None

展开收起

6.3 RAG + CoT#

将检索增强生成 (RAG) 与 CoT 结合：

1
class RAGWithCoT:
2
    """
3
    检索增强的思维链
4
    """
5
    def __init__(self, llm, retriever):
6
        self.llm = llm
7
        self.retriever = retriever
8

9
    def solve(self, question):
10
        # 1. 检索相关文档
11
        docs = self.retriever.retrieve(question, top_k=5)
12
        context = "\n".join([d.content for d in docs])
13

14
        # 2. 基于检索的 CoT
15
        prompt = f"""
16
        参考资料：
17
        {context}
18

19
        问题：{question}
20

21
        请基于参考资料，一步一步推理：
22
        1. 提取关键信息
23
        2. 整合多个资料
24
        3. 推导答案
25
        4. 给出最终答案
26

27
        推理：
28
        """
29
        return self.llm.generate(prompt)

展开收起

7. CoT 的训练方法#

7.1 CoT 数据生成#

1
class CoTDataGenerator:
2
    """
3
    自动生成 CoT 训练数据
4
    """
5
    def __init__(self, llm, dataset):
6
        self.llm = llm
7
        self.dataset = dataset
8

9
    def generate_cot_for_problem(self, problem, solution):
10
        """为现有问题生成 CoT"""
11
        prompt = f"""
12
        问题：{problem}
13
        答案：{solution}
14

15
        请生成到达这个答案的详细推理步骤：
16
        """
17
        cot = self.llm.generate(prompt)
18
        return cot
19

20
    def bootstrap_cot_data(self, num_samples=1000):
21
        """通过 bootstrapping 生成 CoT 数据"""
22
        cot_dataset = []
23

24
        for item in self.dataset.sample(num_samples):
25
            # 1. 使用 Zero-Shot CoT 生成初始推理
26
            initial_cot = self.generate_cot_for_problem(
27
                item.question, item.answer
28
            )
29

30
            # 2. 验证推理是否正确
31
            is_correct = self.verify_cot(
32
                item.question, initial_cot, item.answer
33
            )
34

35
            if is_correct:
36
                cot_dataset.append({
37
                    'question': item.question,
38
                    'cot': initial_cot,
39
                    'answer': item.answer
40
                })
41

42
        return cot_dataset
43

44
    def verify_cot(self, question, cot, expected_answer):
45
        """验证 CoT 是否正确"""
46
        # 执行 CoT 中的推理步骤
47
        # 或使用 LLM 验证
48
        prompt = f"""
49
        问题：{question}
50
        推理：{cot}
51
        期望答案：{expected_answer}
52

53
        这个推理是否正确地推导出期望答案？
54
        """
55
        response = self.llm.generate(prompt)
56
        return "正确" in response or "correct" in response.lower()

展开收起

7.2 CoT 微调#

1
class CoTFinetune:
2
    """
3
    CoT 微调训练
4
    """
5
    def __init__(self, model, tokenizer, cot_dataset):
6
        self.model = model
7
        self.tokenizer = tokenizer
8
        self.dataset = cot_dataset
9

10
    def prepare_training_data(self):
11
        """准备训练数据"""
12
        formatted_data = []
13
        for item in self.dataset:
14
            # 格式: 问题 + CoT + 答案
15
            text = f"""
16
            问题：{item['question']}
17
            推理：{item['cot']}
18
            答案：{item['answer']}
19
            """
20
            formatted_data.append({
21
                'text': text,
22
                'input_ids': self.tokenizer.encode(text)
23
            })
24
        return formatted_data
25

26
    def train(self, epochs=3, batch_size=8, lr=2e-5):
27
        """训练循环"""
28
        optimizer = torch.optim.AdamW(
29
            self.model.parameters(),
30
            lr=lr,
31
            weight_decay=0.01
32
        )
33

34
        for epoch in range(epochs):
35
            for batch in self.dataloader(batch_size):
36
                outputs = self.model(
37
                    input_ids=batch['input_ids'],
38
                    labels=batch['input_ids']
39
                )
40
                loss = outputs.loss
41

42
                loss.backward()
43
                optimizer.step()
44
                optimizer.zero_grad()

展开收起

7.3 强化学习与 CoT#

1
class CoTRLHF:
2
    """
3
    使用 RL 优化 CoT 推理
4
    """
5
    def __init__(self, model, reward_model):
6
        self.model = model
7
        self.reward_model = reward_model
8

9
    def compute_reward(self, question, cot, answer):
10
        """计算 CoT 推理的奖励"""
11
        # 1. 答案正确性奖励
12
        correctness_reward = self.reward_model.score(question, answer)
13

14
        # 2. 推理质量奖励
15
        quality_reward = self.evaluate_reasoning(cot)
16

17
        # 3. 格式奖励
18
        format_reward = self.evaluate_format(cot)
19

20
        return (
21
            0.6 * correctness_reward +
22
            0.3 * quality_reward +
23
            0.1 * format_reward
24
        )
25

26
    def train_step(self, batch):
27
        """单步训练"""
28
        # 1. 模型生成 CoT
29
        generated_cots = self.model.generate(batch['questions'])
30

31
        # 2. 计算奖励
32
        rewards = []
33
        for q, cot in zip(batch['questions'], generated_cots):
34
            answer = self.extract_answer(cot)
35
            reward = self.compute_reward(q, cot, answer)
36
            rewards.append(reward)
37

38
        # 3. PPO 更新
39
        loss = self.ppo_loss(batch, generated_cots, rewards)
40
        loss.backward()
41
        # ... optimizer step

展开收起

8. CoT 在不同任务中的应用#

8.1 数学推理#

1
class MathCoT:
2
    """
3
    数学推理的 CoT 优化
4
    """
5
    def __init__(self, llm):
6
        self.llm = llm
7

8
    def build_math_prompt(self, problem):
9
        """数学专用 CoT prompt"""
10
        template = """
11
解决以下数学问题。展示所有计算步骤。
12

13
问题: {problem}
14

15
让我们一步步解决:
16
1. 识别已知量和未知量
17
2. 确定适用的公式或定理
18
3. 逐步计算
19
4. 验证答案是否合理
20

21
解:
22
"""
23
        return template.format(problem=problem)
24

25
    def solve(self, problem):
26
        prompt = self.build_math_prompt(problem)
27
        return self.llm.generate(prompt)

展开收起

数学 CoT 的关键：

逐步计算：每一步都要明确写出
公式引用：使用准确的数学公式
单位一致：保持单位的一致性
答案验证：通过代入法或估算验证

8.2 常识推理#

1
class CommonsenseCoT:
2
    """
3
    常识推理 CoT
4
    """
5
    def __init__(self, llm):
6
        self.llm = llm
7

8
    def solve(self, question):
9
        prompt = f"""
10
        问题：{question}
11

12
        请基于常识进行推理。思考：
13
        - 这个情境下的人/物通常会如何反应？
14
        - 哪些常识知识可以应用？
15
        - 推理过程：
16

17
        推理：
18
        """
19
        return self.llm.generate(prompt)

展开收起

8.3 多步推理 (Multi-Hop QA)#

1
class MultiHopCoT:
2
    """
3
    多跳问答的 CoT
4
    """
5
    def __init__(self, llm, retriever):
6
        self.llm = llm
7
        self.retriever = retriever
8

9
    def solve(self, question):
10
        prompt = f"""
11
        问题：{question}
12

13
        这是一个需要多步推理的问题。请：
14
        1. 将问题分解为子问题
15
        2. 依次解决每个子问题
16
        3. 整合答案
17

18
        分解：
19
        - 子问题 1:
20
        - 子问题 2:
21
        - 子问题 3:
22

23
        推理：
24
        """
25
        return self.llm.generate(prompt)

展开收起

8.4 代码生成#

1
class CodeCoT:
2
    """
3
    代码生成的 CoT
4
    """
5
    def __init__(self, llm):
6
        self.llm = llm
7

8
    def solve(self, problem):
9
        prompt = f"""
10
        编程问题：{problem}
11

12
        请按以下步骤解决：
13
        1. 理解问题需求
14
        2. 设计算法思路
15
        3. 分析时间/空间复杂度
16
        4. 编写代码
17
        5. 测试边界情况
18

19
        思路：
20
        ```python
21
        # 算法思路

展开收起

代码：

1
# 实现

""" return self.llm.generate(prompt)

1
## 9. CoT 的评估基准
2

3
### 9.1 主要基准
4

5
| 基准 | 类型 | 难度 | 评估指标 |
6
|------|------|------|----------|
7
| GSM8K | 数学 | 中 | 准确率 |
8
| MATH | 数学 | 高 | 准确率 |
9
| AQuA-RAT | 数学 | 中 | 准确率 |
10
| StrategyQA | 常识 | 中 | 准确率 |
11
| HotpotQA | 多跳QA | 高 | F1, EM |
12
| HumanEval | 代码 | 中 | Pass@1 |
13
| LogiQA | 逻辑 | 中 | 准确率 |
14

15
### 9.2 GSM8K 评估示例
16

17
```python
18
class GSM8KEvaluator:
19
    """
20
    GSM8K 数学基准评估
21
    """
22
    def __init__(self, llm):
23
        self.llm = llm
24

25
    def evaluate(self, num_samples=100):
26
        correct = 0
27
        total = 0
28

29
        for problem, expected_answer in self.load_gsm8k(num_samples):
30
            prompt = f"""
31
            问题：{problem.question}
32

33
            让我们一步一步思考。
34
            """
35
            response = self.llm.generate(prompt)
36
            predicted_answer = self.extract_number(response)
37

38
            if abs(float(predicted_answer) - float(expected_answer)) < 1e-3:
39
                correct += 1
40
            total += 1
41

42
        return {
43
            'accuracy': correct / total,
44
            'correct': correct,
45
            'total': total
46
        }

展开收起

9.3 不同方法的基准结果#

基准	Direct	CoT	Self-Consistency	ToT
GSM8K	12.5%	56.5%	74.4%	-
MATH	5.8%	32.5%	50.3%	-
StrategyQA	53.3%	73.0%	81.7%	84.6%
HotpotQA	28.1%	39.7%	47.0%	-

10. CoT 的局限与挑战#

10.1 主要局限#

局限	描述	缓解方法
幻觉推理	模型可能生成看似合理但错误的推理	Self-Consistency, 验证步骤
成本增加	生成更多 token，增加推理成本	选择性 CoT
不稳定的触发	不同 prompt 触发效果差异大	Auto-CoT
不可控的深度	模型可能推理过浅或过深	指定推理深度
难以复杂任务	超复杂任务仍无法处理	分层 CoT

10.2 CoT 的失败模式#

1
class CoTFailureAnalyzer:
2
    """
3
    分析 CoT 失败的原因
4
    """
5
    def __init__(self):
6
        self.failure_types = {
7
            'arithmetic_error': '计算错误',
8
            'logical_fallacy': '逻辑谬误',
9
            'misunderstood_problem': '误解问题',
10
            'incomplete_reasoning': '推理不完整',
11
            'hallucination': '幻觉信息',
12
            'wrong_approach': '方法错误'
13
        }
14

15
    def analyze(self, question, reasoning, expected, predicted):
16
        """分析失败类型"""
17
        analysis = {
18
            'question': question,
19
            'reasoning': reasoning,
20
            'expected': expected,
21
            'predicted': predicted,
22
            'failure_types': []
23
        }
24

25
        # 检测各类失败
26
        if self.has_arithmetic_error(reasoning):
27
            analysis['failure_types'].append('arithmetic_error')
28

29
        if self.has_logical_fallacy(reasoning):
30
            analysis['failure_types'].append('logical_fallacy')
31

32
        # ... 更多检查
33

34
        return analysis

展开收起

10.3 改进方向#

验证机制：在每一步推理后添加验证
混合方法：CoT + 工具 + 搜索
过程奖励模型：训练专门的”推理质量评估器”
强化学习：通过 RL 优化 CoT 策略
多模态 CoT：结合视觉、听觉等多模态信息

11. CoT 与推理模型 (Reasoning Models)#

11.1 OpenAI o1 系列#

OpenAI o1 是 CoT 训练进模型权重的代表：

1
传统模型 + CoT Prompt:
2
  Input: "问题 + Let's think step by step"
3
  Model: 生成推理 token
4
  Output: 最终答案
5

6
OpenAI o1:
7
  Input: "问题"
8
  Model: 内部自动进行长 CoT 推理 (上千 token)
9
  Output: 最终答案 (CoT 隐藏在内部)

o1 的特点：

内部进行长链推理（数千 token）
不向用户展示推理过程
在数学、编程、科学问题上达到博士级水平

11.2 DeepSeek-R1#

DeepSeek-R1 是开源推理模型的代表：

特性	描述
训练方法	纯强化学习 + GRPO
推理能力	数学竞赛级别
开源	是（完全开源）
推理 token	数千到数万

11.3 推理模型的训练范式#

1
class ReasoningModelTrainer:
2
    """
3
    推理模型的训练流程
4
    """
5
    def __init__(self, base_model):
6
        self.model = base_model
7

8
    def phase_1_cot_sft(self, cot_dataset):
9
        """阶段 1: CoT 监督微调"""
10
        print("阶段 1: CoT SFT")
11
        # 让模型学会生成 CoT
12
        for batch in cot_dataset:
13
            loss = self.compute_sft_loss(batch)
14
            self.update(loss)
15

16
    def phase_2_rlhf(self, preference_data):
17
        """阶段 2: 人类反馈强化学习"""
18
        print("阶段 2: RLHF")
19
        # 训练奖励模型
20
        reward_model = self.train_reward_model(preference_data)
21

22
        # PPO 优化
23
        for batch in preference_data:
24
            response = self.model.generate(batch.question)
25
            reward = reward_model.score(batch.question, response)
26
            loss = self.compute_ppo_loss(batch, response, reward)
27
            self.update(loss)
28

29
    def phase_3_distillation(self, teacher_model):
30
        """阶段 3: 蒸馏到小模型"""
31
        print("阶段 3: 蒸馏")
32
        # 用大推理模型生成训练数据
33
        distilled_data = teacher_model.generate_cot_dataset()
34

35
        # 训练小模型
36
        self.phase_1_cot_sft(distilled_data)

展开收起

12. 核心概念总结#

12.1 CoT 的关键 Prompt 模式#

1
# 模式 1: Zero-Shot CoT
2
"Let's think step by step."
3

4
# 模式 2: 中文 Zero-Shot
5
"让我们一步一步思考。"
6

7
# 模式 3: Few-Shot CoT
8
"""
9
示例 1: 问题 → 推理 → 答案
10
示例 2: 问题 → 推理 → 答案
11
问题: {用户问题}
12
推理:
13
"""
14

15
# 模式 4: Self-Consistency
16
"""
17
[生成多个推理路径]
18
[投票选择最常见答案]
19
"""
20

21
# 模式 5: ToT
22
"""
23
[维护推理树]
24
[评估和搜索最佳路径]
25
"""

展开收起

12.2 CoT 的核心数学表达#

CoT 输出形式：

CoT(q)=(z1,z2,…,zk,y)\text{CoT}(q) = (z_1, z_2, \ldots, z_k, y)CoT(q)=(z1,z2,…,zk,y)

其中 ziz_izi 是推理步骤，yyy 是最终答案。

Self-Consistency 投票：

y^=arg⁡max⁡y∑i=1N1(yi=y)\hat{y} = \arg\max_y \sum_{i=1}^{N} \mathbb{1}(y_i = y)y^=argymaxi=1∑N1(yi=y)

CoT 的有效计算深度：

Leff=Lmodel+kL_{\text{eff}} = L_{\text{model}} + kLeff=Lmodel+k

其中 kkk 是 CoT 生成的 token 数。

ToT 评分函数：

V(s)=LLMeval(s)∈[1,10]V(s) = \text{LLM}_{\text{eval}}(s) \in [1, 10]V(s)=LLMeval(s)∈[1,10]

13. 实践建议#

13.1 选择合适的 CoT 方法#

1
def choose_cot_method(task_type, budget, accuracy_need):
2
    """
3
    根据任务选择合适的 CoT 方法
4
    """
5
    if budget == 'low' and accuracy_need == 'low':
6
        return 'direct'  # 直接回答
7

8
    if budget == 'low' and accuracy_need == 'medium':
9
        return 'zero_shot_cot'  # Zero-Shot CoT
10

11
    if budget == 'medium' and accuracy_need == 'high':
12
        return 'few_shot_cot'  # Few-Shot CoT
13

14
    if budget == 'high' and accuracy_need == 'high':
15
        return 'self_consistency'  # Self-Consistency
16

17
    if task_type == 'planning' or task_type == 'search':
18
        return 'tot'  # Tree of Thoughts
19

20
    if task_type == 'tool_use':
21
        return 'react'  # ReAct
22

23
    return 'few_shot_cot'  # 默认

展开收起

13.2 CoT 设计的最佳实践#

演示质量 > 数量：精心设计的 3 个示例胜过随意的 10 个
推理要清晰：每一步都要明确
格式要一致：保持答案格式的一致性
验证是必需的：对关键任务使用 Self-Consistency
成本要权衡：CoT 增加推理成本，需评估性价比

13.3 常见错误及避免#

常见错误	解决方案
演示过于复杂	使用简单、清晰的演示
推理跳步	强调”每一步都要明确”
答案格式不统一	强制使用特定格式（如 #### ）
演示与问题不匹配	选择与目标问题相似的演示
忽略验证	添加答案验证步骤

14. 总结与展望#

CoT 的核心价值#

激活大模型潜能：CoT 让大模型的推理能力”涌现”
可解释性：推理过程透明，可调试
灵活性：适用于多种推理任务
可扩展性：可以与各种技术结合

CoT 的未来方向#

训练即推理：将 CoT 能力训练进模型权重（OpenAI o1）
多模态 CoT：结合视觉、听觉等多模态信息
自适应 CoT：根据问题难度自动调整推理深度
协同 CoT：多个模型协同推理
可验证 CoT：自动验证推理正确性

学习路径建议#

基础：理解 Few-Shot CoT 和 Zero-Shot CoT
进阶：掌握 Self-Consistency 和 ToT
实践：在实际项目中应用 CoT
深入：研究 CoT 的理论解释
前沿：关注推理模型的发展

CoT 是大模型推理能力的重要里程碑。从简单的 prompt 技巧到 OpenAI o1 这样的推理模型，CoT 思想已经深刻地改变了 AI 领域。掌握 CoT 是理解现代大模型推理能力的关键。

参考资料#

Wei, J., et al. (2022). “Chain-of-Thought Prompting Elicits Reasoning in Large Language Models.” NeurIPS.
Kojima, T., et al. (2022). “Large Language Models are Zero-Shot Reasoners.” NeurIPS.
Wang, X., et al. (2022). “Self-Consistency Improves Chain of Thought Reasoning in Language Models.” ICLR.
Yao, S., et al. (2023). “Tree of Thoughts: Deliberate Problem Solving with Large Language Models.” NeurIPS.
Yao, S., et al. (2022). “ReAct: Synergizing Reasoning and Acting in Language Models.” ICLR.
Zhang, Z., et al. (2022). “Automatic Chain of Thought Prompting in Large Language Models.” arXiv.
Besta, M., et al. (2023). “Graph of Thoughts: Solving Elaborate Problems with Large Language Models.” AAAI.
Madaan, A., et al. (2023). “Self-Refine: Iterative Refinement with Self-Feedback.” NeurIPS.
Gao, L., et al. (2022). “PAL: Program-aided Language Models.” ICML.
Schick, T., et al. (2023). “Toolformer: Language Models Can Teach Themselves to Use Tools.” NeurIPS.
Wei, J., et al. (2022). “Emergent Abilities of Large Language Models.” TMLR.
OpenAI. (2024). “Learning to Reason with LLMs.” OpenAI Blog.
DeepSeek-AI. (2025). “DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning.” arXiv.

CLIP 深度剖析：对比语言-图像预训练的完整技术栈

Thu, 16 Jul 2026 00:00:00 GMT

1. CLIP 是什么：打破数据标注的瓶颈#

1.1 传统视觉训练的困境#

在 CLIP 出现之前，视觉模型的训练严重依赖人工标注数据：

1
ImageNet 分类: 1.2M 图像, 1,000 个类别, 人工标注
2
COCO 检测: 330K 图像, 80 个物体类别, 每张图多人标注
3

4
问题:
5
  1. 标注成本极高——ImageNet 用了 2.5 万名众包工人
6
  2. 类别固定——无法泛化到 ImageNet 之外的任何类别
7
  3. 任务单一——分类模型无法直接用于检测、分割、OCR
8
  4. 分布偏置——互联网图片 ≠ ImageNet 风格

1.2 CLIP 的核心洞察#

互联网上天然存在大量的图像-文本配对数据**（网页 alt 标签、标题描述、用户上传说明等）。与其让人工给图像打标签，不如直接用自然语言 supervision——让模型学会”图像和什么文字描述最匹配”。**

这带来了一个根本性的范式转变：

1
传统:  图像 → (类别标签) → 预测类别      # 需要人工标注
2
CLIP:  图像 + 文本 → 匹配程度预测          # 用自然语言 supervision
3

4
文本本身就是标签——不需要预先定义类别。
5
新的类别 = 新的文本描述 = 零成本。

1.3 CLIP 的能力#

CLIP 训练完成后，可以在零样本（zero-shot）条件下完成从未见过的分类任务：

1
任务: 分类一张猫的照片
2

3
传统做法:
4
  - 需要在 ImageNet-1K 的 1000 个类上微调
5
  - 如果类别不是"猫"，模型完全无法识别
6

7
CLIP 零样本:
8
  - 类别描述: ["a photo of a cat", "a photo of a dog", "a photo of a bird"]
9
  - 把这些文本编码成向量
10
  - 把图像编码成向量
11
  - 计算余弦相似度 → 最高者胜出
12
  - 无需任何微调，即插即用

1.4 论文信息#

1
论文: "Learning Transferable Visual Models From Natural Language Supervision"
2
作者: Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh,
3
      Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell,
4
      Pamela Mishkin, Jack Clark, Gretchen Krueger, Ilya Sutskever
5
单位: OpenAI
6
发表于: ICML 2021
7
引用: > 20,000 次（截至 2024）★★★★★
8
开源: https://github.com/openai/CLIP

1.5 一句话概括 CLIP#

CLIP 通过一个双塔对比学习框架，用自然语言作为监督信号，让视觉编码器和文本编码器在同一个特征空间中对齐——从而实现零样本图像分类、开放词汇检测、多模态理解等下游任务，并成为几乎所有 VLM 和 VLA 的视觉编码器基础。

2. 数学框架：对比学习目标#

2.1 双塔架构#

CLIP 由两个编码器组成：

1
┌──────────────────────────────────────────────────────────────┐
2
│                                                              │
3
│   [图像] ──→ 视觉编码器 ──→ I(x) ∈ ℝ^D  (视觉特征向量)      │
4
│                    ViT 或 ResNet                             │
5
│                                                              │
6
│   [文本] ──→ 文本编码器 ──→ T(t) ∈ ℝ^D  (文本特征向量)      │
7
│                    Transformer                               │
8
│                                                              │
9
│   对比: 最大化匹配的 I·T，余弦相似度矩阵                     │
10
│                                                              │
11
└──────────────────────────────────────────────────────────────┘

2.2 InfoNCE 对比损失#

给定一个 batch 的图文对 (xiI,xiT)(x_i^I, x_i^T)(xiI,xiT)，CLIP 最大化正样本对的相似度，最小化负样本对的相似度。

余弦相似度矩阵：

S=Ii⋅Tj∥Ii∥∥Tj∥S = \frac{I_i \cdot T_j}{\|I_i\| \|T_j\|}S=∥Ii∥∥Tj∥Ii⋅Tj

其中 IiI_iIi 是第 iii 张图像的归一化特征，TjT_jTj 是第 jjj 个文本的归一化特征。

对称的对比损失（两方向都要优化）：

L=12(Li+Lt)\mathcal{L} = \frac{1}{2} \left( \mathcal{L}_i + \mathcal{L}_t \right)L=21(Li+Lt)

图像→文本损失（Image-to-Text）：

Li=−1N∑i=1Nlog⁡exp⁡(Si,i/τ)∑j=1Nexp⁡(Si,j/τ)\mathcal{L}_i = -\frac{1}{N} \sum_{i=1}^{N} \log \frac{\exp(S_{i,i}/\tau)}{\sum_{j=1}^{N} \exp(S_{i,j}/\tau)}Li=−N1i=1∑Nlog∑j=1Nexp(Si,j/τ)exp(Si,i/τ)

文本→图像损失（Text-to-Image）：

Lt=−1N∑j=1Nlog⁡exp⁡(Sj,j/τ)∑i=1Nexp⁡(Si,j/τ)\mathcal{L}_t = -\frac{1}{N} \sum_{j=1}^{N} \log \frac{\exp(S_{j,j}/\tau)}{\sum_{i=1}^{N} \exp(S_{i,j}/\tau)}Lt=−N1j=1∑Nlog∑i=1Nexp(Si,j/τ)exp(Sj,j/τ)

其中 τ\tauτ 是温度系数（可学习的参数），控制 softmax 的锐度。

2.3 损失的直觉解释#

1
def info_nce_loss(image_embeds, text_embeds, temperature):
2
    """
3
    InfoNCE 对比损失（简化版）。
4

5
    参数:
6
        image_embeds: (B, D) 图像特征
7
        text_embeds:  (B, D) 文本特征
8
        temperature:  温度系数 τ
9
    """
10
    # 归一化
11
    image_embeds = F.normalize(image_embeds, dim=-1)
12
    text_embeds  = F.normalize(text_embeds, dim=-1)
13

14
    # 计算余弦相似度矩阵 (B, B)
15
    logits = torch.matmul(image_embeds, text_embeds.T) / temperature
16

17
    # labels: 对角线上是正样本
18
    batch_size = image_embeds.shape[0]
19
    labels = torch.arange(batch_size, device=image_embeds.device)
20

21
    # 两方向的对称损失
22
    loss_i2t = F.cross_entropy(logits, labels)       # 图像→文本
23
    loss_t2i = F.cross_entropy(logits.T, labels)      # 文本→图像
24

25
    loss = (loss_i2t + loss_t2i) / 2
26
    return loss

展开收起

1
直观理解:
2

3
对角线 (正样本):   S_ii → 最大 (希望)
4
非对角线 (负样本): S_ij (i≠j) → 最小 (希望)
5

6
每个图像 i 应该与对应的文本 j 最匹配
7
每个文本 j 应该与对应的图像 i 最匹配
8

9
这就是"对比"——把正确配对从 N 个配对中对比出来。

2.4 温度系数的意义#

1
def temperature_effect():
2
    """
3
    温度系数 τ 对对比损失的影响。
4
    τ 是可学习的参数，初始化通常在 0.01~0.1 之间。
5
    """
6
    configs = {
7
        "τ 极小 (0.01)": "softmax 极度尖锐 → 只关注最难的负样本",
8
        "τ 适中 (0.07)": "平衡——正确配对突出，负样本被适度压制",
9
        "τ 过大 (1.0)":  "softmax 趋于均匀 → 无法区分正负样本",
10
    }
11
    return configs
12

13
    # CLIP 原始实现中 τ 被建模为 log(scaled_logits) 的可学习参数
14
    # logit_scale = nn.Parameter(torch.ones([]) * np.log(1/0.07))
15
    # temperature = 1 / torch.exp(logit_scale) ≈ 0.07

展开收起

3. 视觉编码器：Vision Transformer#

3.1 图像分块#

CLIP 的视觉编码器使用 ViT（Vision Transformer） 而非 ResNet。图像被切分成固定大小的 patch：

1
输入图像: 224 × 224 × 3
2

3
patch_size = 16 × 16
4

5
图像分块:
6
  - 224 / 16 = 14
7
  - 总 patch 数 = 14 × 14 = 196
8
  - 每个 patch = 16 × 16 × 3 = 768 维
9

10
加上 [CLS] token → 197 个视觉 token

1
class ImagePatcher(nn.Module):
2
    """
3
    图像分块：把图像切成 patch 并线性投影。
4
    """
5
    def __init__(self, image_size=224, patch_size=16, in_channels=3, embed_dim=768):
6
        super().__init__()
7
        self.image_size = image_size
8
        self.patch_size = patch_size
9
        self.num_patches = (image_size // patch_size) ** 2  # 196
10

11
        # 线性投影层（等价于 2D conv with kernel_size=patch_size）
12
        self.proj = nn.Conv2d(in_channels, embed_dim,
13
                               kernel_size=patch_size, stride=patch_size)
14

15
    def forward(self, x):
16
        """
17
        参数:
18
            x: (B, 3, 224, 224) 图像 batch
19
        返回:
20
            patches: (B, num_patches, embed_dim)  每个 patch 的嵌入
21
        """
22
        # Conv2d: (B, 3, 224, 224) → (B, embed_dim, 14, 14)
23
        x = self.proj(x)
24
        # 展平空间维度: (B, embed_dim, 14, 14) → (B, embed_dim, 196)
25
        x = x.flatten(2)
26
        # 调换维度: (B, embed_dim, 196) → (B, 196, embed_dim)
27
        x = x.transpose(1, 2)
28
        return x

展开收起

3.2 ViT Transformer 编码器#

1
class ViTEncoder(nn.Module):
2
    """
3
    CLIP 的 ViT 视觉编码器。
4
    """
5
    def __init__(self, image_size=224, patch_size=16, embed_dim=768,
6
                 num_layers=12, num_heads=12, mlp_ratio=4.0):
7
        super().__init__()
8
        self.patch_embed = ImagePatcher(image_size, patch_size, 3, embed_dim)
9
        self.num_patches = self.patch_embed.num_patches
10

11
        # [CLS] token
12
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
13

14
        # 位置编码（可学习）
15
        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + 1, embed_dim))
16

17
        # Transformer 编码器
18
        encoder_layer = nn.TransformerEncoderLayer(
19
            d_model=embed_dim,
20
            nhead=num_heads,
21
            dim_feedforward=int(embed_dim * mlp_ratio),
22
            activation="gelu",
23
            batch_first=True,
24
            norm_first=True,  # Pre-LN (CLIP 使用)
25
        )
26
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)
27

28
        # LayerNorm
29
        self.ln = nn.LayerNorm(embed_dim)
30

31
    def forward(self, x):
32
        """
33
        参数:
34
            x: (B, 3, 224, 224)
35
        返回:
36
            pooled: (B, embed_dim) [CLS] token 的最终表示
37
            sequence: (B, num_patches+1, embed_dim) 完整序列
38
        """
39
        B = x.shape[0]
40

41
        # 1) 分块 + 投影
42
        patches = self.patch_embed(x)  # (B, 196, embed_dim)
43

44
        # 2) 添加 [CLS] token
45
        cls_tokens = self.cls_token.expand(B, -1, -1)  # (B, 1, embed_dim)
46
        x = torch.cat([cls_tokens, patches], dim=1)   # (B, 197, embed_dim)
47

48
        # 3) 添加位置编码
49
        x = x + self.pos_embed  # (B, 197, embed_dim)
50

51
        # 4) Transformer 编码
52
        x = self.transformer(x)  # (B, 197, embed_dim)
53

54
        # 5) LayerNorm
55
        x = self.ln(x)
56

57
        # [CLS] token
58
        pooled = x[:, 0, :]       # (B, embed_dim)
59
        sequence = x              # (B, 197, embed_dim)
60

61
        return pooled, sequence

展开收起

3.3 CLIP 使用的两种视觉编码器#

编码器	规模	patch size	层数	参数量	特点
ViT-B/32	Base	32×32	12	86M	CLIP 默认, 速度最快
ViT-B/16	Base	16×16	12	86M	更高分辨率
ViT-L/14	Large	14×14	24	304M	CLIP 最强, 最常用
ViT-L/14@336	Large	14×14	24	304M	微调版, 336px
RN50	ResNet-50	-	-	38M	论文消融实验

4. 文本编码器：Text Transformer#

4.1 文本预处理#

1
class TextTokenizer:
2
    """
3
    CLIP 使用 BPE (Byte Pair Encoding) 分词器。
4
    与 GPT-2 / RoBERTa 相同。
5
    """
6
    def __init__(self, vocab_size=50257):
7
        self.tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
8
        self.tokenizer.add_special_tokens(["<|startoftext|>", "<|endoftext|>"])
9

10
    def encode(self, texts):
11
        """
12
        参数:
13
            texts: str 或 list[str]
14
        返回:
15
            input_ids: (B, L) token IDs
16
            attention_mask: (B, L) 注意力掩码
17
        """
18
        if isinstance(texts, str):
19
            texts = [texts]
20

21
        encoded = self.tokenizer(
22
            texts,
23
            padding=True,
24
            truncation=True,
25
            max_length=77,     # CLIP 固定文本长度
26
            return_tensors="pt",
27
        )
28
        return encoded["input_ids"], encoded["attention_mask"]

展开收起

4.2 Transformer 文本编码器#

1
class TextTransformerEncoder(nn.Module):
2
    """
3
    CLIP 的文本编码器。
4
    Transformer Encoder 处理文本 token 序列。
5
    """
6
    def __init__(self, vocab_size=50257, embed_dim=512,
7
                 context_length=77, num_layers=12, num_heads=8):
8
        super().__init__()
9
        self.context_length = context_length
10

11
        # Token 嵌入
12
        self.token_embedding = nn.Embedding(vocab_size, embed_dim)
13

14
        # 位置编码（可学习，固定长度）
15
        self.pos_embedding = nn.Parameter(
16
            torch.zeros(1, context_length, embed_dim)  # (1, 77, 512)
17
        )
18

19
        # Transformer 编码器
20
        encoder_layer = nn.TransformerEncoderLayer(
21
            d_model=embed_dim,
22
            nhead=num_heads,
23
            dim_feedforward=embed_dim * 4,
24
            activation="gelu",
25
            batch_first=True,
26
            norm_first=True,
27
        )
28
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)
29

30
        # LayerNorm
31
        self.ln = nn.LayerNorm(embed_dim)
32

33
    def forward(self, input_ids, attention_mask=None):
34
        """
35
        参数:
36
            input_ids:     (B, L) token IDs
37
            attention_mask: (B, L) 可选
38
        返回:
39
            pooled: (B, embed_dim) [EOS] token 的表示
40
            sequence: (B, L, embed_dim) 完整序列
41
        """
42
        B, L = input_ids.shape
43

44
        # 1) Token 嵌入
45
        x = self.token_embedding(input_ids)  # (B, L, embed_dim)
46

47
        # 2) 添加位置编码
48
        x = x + self.pos_embedding[:, :L, :]  # (B, L, embed_dim)
49

50
        # 3) Transformer 编码
51
        x = self.transformer(x)  # (B, L, embed_dim)
52

53
        # 4) LayerNorm
54
        x = self.ln(x)
55

56
        # 5) 取 [EOS] token（通常是文本序列最后一个 token）作为句子表示
57
        #    或者用 attention pooling
58
        pooled = x[:, -1, :]  # (B, embed_dim)
59

60
        return pooled, x

展开收起

4.3 文本表示的选择#

CLIP 原文使用 [EOS] token 的输出作为整个文本的表示，而非 attention pooling：

1
方法 1: [EOS] token → 简单, 最常用
2
  "hello world [EOS]"
3
                  ↑
4
              取这个位置的特征
5

6
方法 2: Attention Pooling → 更灵活
7
  attention_weights = softmax(w · tanh(W · h_i^T))
8
  text_repr = Σ attention_weights_i · h_i
9

10
方法 3: Mean Pooling → 无参数
11
  text_repr = mean(h_1, h_2, ..., h_L)

CLIP 实验发现 [EOS] token 效果最好，原因可能是：Transformer 的 [EOS] 位置已经”看过”了完整序列。

5. 训练细节：从零开始的工程实践#

5.1 训练数据#

CLIP 使用 WIT-400M 数据集（WebImageText）：

1
WIT-400M:
2
  - 400,000,000 图像-文本对
3
  - 来源: 从互联网上抓取的图像+alt 文本/title/caption
4
  - 文本平均长度: ~12 个词
5
  - 覆盖类别: 自然场景、物体、人物、动作、场景描述等
6

7
与 ImageNet 的对比:
8
  ImageNet:  1.2M 图像, 1,000 类, 人工标注
9
  WIT-400M: 400M 图像, 无固定类别, 自然语言描述

数据预处理：

1
class WITDataset(Dataset):
2
    """
3
    CLIP 的训练数据加载器。
4
    """
5
    def __init__(self, image_paths, texts, image_size=224):
6
        self.image_paths = image_paths
7
        self.texts = texts
8
        self.image_size = image_size
9
        self.img_transform = transforms.Compose([
10
            transforms.Resize(image_size, interpolation=transforms.InterpolationMode.BICUBIC),
11
            transforms.CenterCrop(image_size),
12
            transforms.ToTensor(),
13
            transforms.Normalize(
14
                mean=[0.48145466, 0.4578275, 0.40821073],  # ImageNet RGB 统计
15
                std=[0.26862954, 0.26130258, 0.27577711],
16
            ),
17
        ])
18

19
    def __getitem__(self, idx):
20
        # 加载图像
21
        image = Image.open(self.image_paths[idx]).convert("RGB")
22
        image = self.img_transform(image)
23

24
        # 获取文本
25
        text = self.texts[idx][:77]  # 截断到 77 tokens
26

27
        return {"image": image, "text": text}

展开收起

5.2 训练超参数#

1
@dataclass
2
class CLIPTrainingConfig:
3
    """CLIP 训练配置。"""
4
    # 模型
5
    vision_model: str = "ViT-L/14"       # CLIP 默认使用 ViT-L/14
6
    embed_dim: int = 768                 # 视觉/文本特征的维度
7

8
    # 训练
9
    batch_size: int = 32768               # CLIP 原始实现用 32768（很大！）
10
    num_epochs: int = 32
11
    learning_rate: float = 5e-4            # AdamW 学习率
12
    weight_decay: float = 0.1
13
    warmup_steps: int = 5000
14

15
    # 优化器
16
    optimizer: str = "AdamW"
17
    beta1: float = 0.9
18
    beta2: float = 0.98
19
    eps: float = 1e-6
20

21
    # 损失
22
    temperature_init: float = 0.07        # 温度系数初始化
23
    label_smoothing: float = 0.0
24

25
    # 正则化
26
    gradient_clip: float = 1.0
27
    mixed_precision: bool = True         # FP16 混合精度
28

29
    # 数据
30
    image_size: int = 224
31
    context_length: int = 77
32
    max_text_length: int = 77

展开收起

注意：CLIP 原始实现的 batch_size 高达 32768！这需要分布式训练（多机多卡）。社区复现通常用 256-1024。

5.3 训练循环#

1
def train_clip(model, dataloader, config):
2
    """
3
    CLIP 训练循环（简化版）。
4
    """
5
    optimizer = torch.optim.AdamW(
6
        model.parameters(),
7
        lr=config.learning_rate,
8
        weight_decay=config.weight_decay,
9
        betas=(config.beta1, config.beta2),
10
        eps=config.eps,
11
    )
12

13
    # 学习率调度器
14
    def lr_lambda(step):
15
        if step < config.warmup_steps:
16
            return step / config.warmup_steps
17
        else:
18
            progress = (step - config.warmup_steps) / (config.total_steps - config.warmup_steps)
19
            return 0.5 * (1 + np.cos(np.pi * progress))
20

21
    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
22

23
    # 可学习的温度系数
24
    logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / config.temperature_init))
25
    logit_scale = logit_scale.to(device)
26

27
    model.train()
28
    for epoch in range(config.num_epochs):
29
        epoch_loss = 0.0
30
        for batch in tqdm(dataloader):
31
            images = batch["image"].to(device)
32
            texts  = batch["text"]
33

34
            # 编码
35
            image_features = model.visual(images)    # (B, D)
36
            text_features  = model.text(texts)       # (B, D)
37

38
            # L2 归一化
39
            image_features = F.normalize(image_features, dim=-1)
40
            text_features  = F.normalize(text_features, dim=-1)
41

42
            # 对比 logits
43
            temperature = (1 / torch.exp(logit_scale))
44
            logits_per_image = (image_features @ text_features.T) / temperature  # (B, B)
45
            logits_per_text  = logits_per_image.T                                        # (B, B)
46

47
            # Labels: 对角线是正样本
48
            batch_size = images.shape[0]
49
            labels = torch.arange(batch_size, device=device)
50

51
            # 对称交叉熵损失
52
            loss_i2t = F.cross_entropy(logits_per_image, labels)
53
            loss_t2i = F.cross_entropy(logits_per_text, labels)
54
            loss = (loss_i2t + loss_t2i) / 2
55

56
            # 反向传播
57
            optimizer.zero_grad()
58
            loss.backward()
59

60
            # 梯度裁剪
61
            torch.nn.utils.clip_grad_norm_(model.parameters(), config.gradient_clip)
62

63
            optimizer.step()
64
            scheduler.step()
65

66
            epoch_loss += loss.item()
67

68
        # 打印
69
        print(f"Epoch {epoch}: loss={epoch_loss / len(dataloader):.4f}, "
70
              f"temp={temperature.item():.4f}")

展开收起

6. 零样本分类：CLIP 的杀手锏#

6.1 零样本分类流程#

1
@torch.no_grad()
2
def zero_shot_classify(model, image, class_names, device):
3
    """
4
    CLIP 零样本分类。
5

6
    参数:
7
        model:       CLIP 模型
8
        image:       (C, H, W) 图像 tensor
9
        class_names: list[str] 类别名称，如 ["cat", "dog", "bird"]
10
        device:      计算设备
11
    返回:
12
        预测的类别索引和概率
13
    """
14
    model.eval()
15

16
    # 1) 构建文本提示（class name → sentence）
17
    #    CLIP 原论文发现"a photo of a {class}"效果最好
18
    prompts = [f"a photo of a {c}." for c in class_names]
19

20
    # 2) 编码文本
21
    text_tokens = model.tokenizer(prompts).to(device)
22
    text_features = model.encode_text(text_tokens)
23
    text_features = F.normalize(text_features, dim=-1)  # (num_classes, D)
24

25
    # 3) 编码图像
26
    if len(image.shape) == 3:
27
        image = image.unsqueeze(0)  # (1, C, H, W)
28
    image = image.to(device)
29
    image_features = model.encode_image(image)
30
    image_features = F.normalize(image_features, dim=-1)  # (1, D)
31

32
    # 4) 计算余弦相似度
33
    similarity = (image_features @ text_features.T)  # (1, num_classes)
34
    probs = F.softmax(similarity * model.logit_scale, dim=-1)  # (1, num_classes)
35

36
    # 5) 返回结果
37
    probs = probs[0].cpu().numpy()
38
    pred_idx = probs.argmax()
39
    return pred_idx, class_names[pred_idx], probs[pred_idx], probs

展开收起

6.2 提示工程（Prompt Engineering）#

CLIP 对文本提示的措辞非常敏感：

1
def prompt_ensemble(model, class_names, image, device):
2
    """
3
    提示集成：使用多个提示模板提升零样本分类精度。
4
    CLIP 原论文发现这一招效果显著。
5
    """
6
    # 多种提示模板
7
    prompt_templates = [
8
        "a photo of a {}.",
9
        "a blurry photo of a {}.",
10
        "a low resolution photo of a {}.",
11
        "a bright photo of a {}.",
12
        "a dark photo of a {}.",
13
        "a photo of a small {}.",
14
        "a photo of a large {}.",
15
        "a photo of the {}.",
16
        "a photo of a {} in the scene.",
17
        "a photo of the {} in the scene.",
18
    ]
19

20
    all_text_features = []
21

22
    for template in prompt_templates:
23
        prompts = [template.format(c) for c in class_names]
24
        text_tokens = model.tokenizer(prompts).to(device)
25
        text_features = model.encode_text(text_tokens)
26
        text_features = F.normalize(text_features, dim=-1)
27
        all_text_features.append(text_features)
28

29
    # 平均所有提示的特征
30
    text_features_avg = torch.stack(all_text_features, dim=0).mean(dim=0)  # (num_classes, D)
31
    text_features_avg = F.normalize(text_features_avg, dim=-1)
32

33
    # 图像编码
34
    image_features = model.encode_image(image.to(device).unsqueeze(0))
35
    image_features = F.normalize(image_features, dim=-1)
36

37
    # 相似度
38
    similarity = image_features @ text_features_avg.T
39
    probs = F.softmax(similarity * model.logit_scale, dim=-1)
40

41
    return probs[0].cpu().numpy()

展开收起

提示工程的效果（来自 CLIP 论文）：

1
ImageNet 零样本分类准确率对比:
2

3
单一提示 "a photo of a {}":           56.3%
4
集成 80 种提示:                         59.1%  (+2.8%)
5
+ 描述性提示 (如 "a blurry photo..."):  +2.0-4.0%
6
+ 类别同义词 ("cat, kitty"):            +1.0-2.0%

6.3 零样本迁移的理论解释#

1
def why_zero_shot_works():
2
    """
3
    CLIP 零样本分类有效的原因。
4
    """
5
    explanations = {
6
        "语义丰富": "文本编码器能理解语义——'golden retriever' 和 'labrador' 都指向同一种狗",
7
        "分布匹配": "训练数据是互联网图文对，测试分布与训练分布相似",
8
        "语言先验": "语言模型带来了世界知识——知道什么是'斑马'的描述",
9
        "特征泛化": "ViT 学到了通用的视觉表示，适用于任何视觉概念",
10
        "开放词汇": "不需要固定类别，任意文本都可以作为查询",
11
    }
12
    return explanations

7. 扩展：ALIGN / EVA-CLIP / SigLIP#

7.1 ALIGN：更大数据 + 更强骨干#

1
ALIGN (Jia et al., 2021):
2
  - 数据: 1.8B 图像-文本对（比 CLIP 多 4.5 倍）
3
  - 骨干: EfficientNet-L2 (更大的视觉编码器)
4
  - 结果: 在 ImageNet 上零样本 85.5% (CLIP: 76.2%)
5

6
核心发现: 数据规模 > 模型架构
7
  400M 图像 + ViT-L = 76.2%
8
  1.8B 图像 + ViT-L = 85.5% (+9.3%)
9

10
→ 更多的图文数据比更大的模型更重要

7.2 EVA-CLIP：打补丁的 CLIP#

1
EVA-CLIP (Sun et al., 2023):
2
  - 数据: LAION-5B (50 亿图文对)
3
  - 骨干: EVA-02 (视觉编码器改进版)
4
  - 关键改进:
5
    1. CLIP 预训练 → EVA 预训练 → CLIP 对比微调（两阶段）
6
    2. 更大的 ViT-G (2B 参数)
7
    3. 渐进式训练策略
8

9
结果:
10
  - EVA-CLIP-18B: ImageNet 零样本 82.0%
11
  - EVA-CLIP-18B + 融合: 86.5%

7.3 SigLIP：sigmoid 损失替代 softmax#

SigLIP（Google, 2024）用成对 sigmoid 损失替代 CLIP 的 softmax InfoNCE：

L=−1N2∑i=1N∑j=1N[yij⋅log⁡(σ(Sij/τ))+(1−yij)⋅log⁡(1−σ(Sij/τ))]\mathcal{L} = -\frac{1}{N^2} \sum_{i=1}^{N} \sum_{j=1}^{N} \left[ y_{ij} \cdot \log(\sigma(S_{ij}/\tau)) + (1 - y_{ij}) \cdot \log(1 - \sigma(S_{ij}/\tau)) \right]L=−N21i=1∑Nj=1∑N[yij⋅log(σ(Sij/τ))+(1−yij)⋅log(1−σ(Sij/τ))]

其中 yij=1y_{ij} = 1yij=1 当且仅当 i=ji = ji=j（正样本对）。

1
def siglip_loss(image_features, text_features, temperature):
2
    """
3
    SigLIP 的 Sigmoid 对比损失。
4

5
    特点:
6
      - 每个样本独立优化（不需要 global softmax）
7
      - 可以处理更大的 batch size
8
      - 避免正样本被其他正样本"稀释"
9
    """
10
    image_features = F.normalize(image_features, dim=-1)
11
    text_features  = F.normalize(text_features, dim=-1)
12

13
    logits = image_features @ text_features.T / temperature  # (B, B)
14

15
    # 正样本标签
16
    B = image_features.shape[0]
17
    labels = torch.eye(B, device=image_features.device)
18

19
    # Sigmoid 交叉熵（而非 softmax）
20
    loss = F.binary_cross_entropy_with_logits(logits, labels)
21

22
    return loss

展开收起

SigLIP vs CLIP:

维度	CLIP (softmax)	SigLIP (sigmoid)
损失函数	InfoNCE (softmax)	Pairwise sigmoid
Batch 内交互	归一化（全局 softmax）	独立（成对）
最优温度	~0.07	~0.1
大 batch 依赖	高（需要大 batch 提供足够负样本）	低
精度	76.2%	78.4% (ImageNet zero-shot)

8. 完整实现#

1
import torch
2
import torch.nn as nn
3
import torch.nn.functional as F
4
from dataclasses import dataclass
5

6

7
@dataclass
8
class CLIPConfig:
9
    """CLIP 模型配置。"""
10
    # 视觉编码器
11
    image_size: int = 224
12
    patch_size: int = 14
13
    vision_embed_dim: int = 768      # ViT-L/14
14
    vision_layers: int = 24
15
    vision_heads: int = 16
16

17
    # 文本编码器
18
    vocab_size: int = 49408
19
    text_embed_dim: int = 512
20
    text_layers: int = 12
21
    text_heads: int = 8
22
    context_length: int = 77
23

24
    # 融合
25
    embed_dim: int = 768             # 视觉和文本特征最终维度
26

27

28
class ImagePatcher(nn.Module):
29
    """图像分块层。"""
30
    def __init__(self, patch_size=14, in_channels=3, embed_dim=768):
31
        super().__init__()
32
        self.proj = nn.Conv2d(in_channels, embed_dim,
33
                               kernel_size=patch_size, stride=patch_size)
34

35
    def forward(self, x):
36
        x = self.proj(x)              # (B, embed_dim, H/P, W/P)
37
        x = x.flatten(2).transpose(1, 2)  # (B, num_patches, embed_dim)
38
        return x
39

40

41
class CLIPVisionEncoder(nn.Module):
42
    """CLIP 的 ViT 视觉编码器。"""
43
    def __init__(self, config: CLIPConfig):
44
        super().__init__()
45
        self.config = config
46
        self.num_patches = (config.image_size // config.patch_size) ** 2
47

48
        self.patch_embed = ImagePatcher(config.patch_size, 3, config.vision_embed_dim)
49

50
        # [CLS] token
51
        self.cls_token = nn.Parameter(torch.zeros(1, 1, config.vision_embed_dim))
52

53
        # 位置编码
54
        self.pos_embed = nn.Parameter(
55
            torch.zeros(1, self.num_patches + 1, config.vision_embed_dim)
56
        )
57

58
        # Transformer 编码器
59
        encoder_layer = nn.TransformerEncoderLayer(
60
            d_model=config.vision_embed_dim,
61
            nhead=config.vision_heads,
62
            dim_feedforward=config.vision_embed_dim * 4,
63
            activation="gelu",
64
            batch_first=True,
65
            norm_first=True,
66
        )
67
        self.transformer = nn.TransformerEncoder(encoder_layer, config.vision_layers)
68

69
        self.ln_post = nn.LayerNorm(config.vision_embed_dim)
70

71
    def forward(self, x):
72
        B = x.shape[0]
73

74
        # 分块
75
        x = self.patch_embed(x)  # (B, num_patches, embed_dim)
76

77
        # [CLS] token
78
        cls_tokens = self.cls_token.expand(B, -1, -1)
79
        x = torch.cat([cls_tokens, x], dim=1)
80

81
        # 位置编码
82
        x = x + self.pos_embed
83

84
        # Transformer
85
        x = self.transformer(x)
86
        x = self.ln_post(x)
87

88
        return x
89

90

91
class CLIPTextEncoder(nn.Module):
92
    """CLIP 的文本编码器。"""
93
    def __init__(self, config: CLIPConfig):
94
        super().__init__()
95
        self.config = config
96

97
        # Token 嵌入
98
        self.token_embedding = nn.Embedding(config.vocab_size, config.text_embed_dim)
99

100
        # 位置编码
101
        self.pos_embedding = nn.Parameter(
102
            torch.zeros(1, config.context_length, config.text_embed_dim)
103
        )
104

105
        # Transformer
106
        encoder_layer = nn.TransformerEncoderLayer(
107
            d_model=config.text_embed_dim,
108
            nhead=config.text_heads,
109
            dim_feedforward=config.text_embed_dim * 4,
110
            activation="gelu",
111
            batch_first=True,
112
            norm_first=True,
113
        )
114
        self.transformer = nn.TransformerEncoder(encoder_layer, config.text_layers)
115

116
        self.ln_final = nn.LayerNorm(config.text_embed_dim)
117

118
    def forward(self, input_ids):
119
        B, L = input_ids.shape
120

121
        # Token 嵌入 + 位置编码
122
        x = self.token_embedding(input_ids) + self.pos_embedding[:, :L, :]
123

124
        # Transformer
125
        x = self.transformer(x)
126
        x = self.ln_final(x)
127

128
        # 取 [EOS] token 作为句子表示
129
        return x
130

131

132
class CLIP(nn.Module):
133
    """
134
    CLIP 完整模型。
135
    论文: "Learning Transferable Visual Models From Natural Language Supervision"
136
    """
137

138
    def __init__(self, config: CLIPConfig):
139
        super().__init__()
140
        self.config = config
141

142
        # 编码器
143
        self.visual = CLIPVisionEncoder(config)
144
        self.transformer = CLIPTextEncoder(config)
145

146
        # 投影层：统一特征维度
147
        self.visual_projection = nn.Linear(config.vision_embed_dim, config.embed_dim)
148
        self.text_projection   = nn.Linear(config.text_embed_dim, config.embed_dim)
149

150
        # 可学习的温度（logit scale）
151
        self.logit_scale = nn.Parameter(torch.ones([]) * 2.659)
152

153
        # 分词器
154
        from transformers import GPT2Tokenizer
155
        self.tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
156
        self.tokenizer.add_special_tokens(["<|startoftext|>", "<|endoftext|>"])
157
        self.tokenizer.pad_token = self.tokenizer.eos_token
158

159
    def encode_image(self, image):
160
        """编码图像。"""
161
        x = self.visual(image)                # (B, num_patches+1, vision_embed_dim)
162
        x = x[:, 0, :]                        # [CLS] token: (B, vision_embed_dim)
163
        x = self.visual_projection(x)          # (B, embed_dim)
164
        return F.normalize(x, dim=-1)
165

166
    def encode_text(self, text):
167
        """
168
        编码文本。
169

170
        参数:
171
            text: str 或 list[str] 或 token IDs
172
        """
173
        # Tokenize
174
        if isinstance(text, str):
175
            text = [text]
176

177
        if isinstance(text, list) and isinstance(text[0], str):
178
            encoded = self.tokenizer(
179
                text,
180
                padding=True,
181
                truncation=True,
182
                max_length=self.config.context_length,
183
                return_tensors="pt",
184
            )
185
            input_ids = encoded["input_ids"].to(next(self.parameters()).device)
186
        else:
187
            input_ids = text.to(next(self.parameters()).device)
188

189
        x = self.transformer(input_ids)      # (B, L, text_embed_dim)
190
        x = x[torch.arange(x.shape[0]), input_ids.argmax(dim=-1)]  # [EOS] token
191
        x = self.text_projection(x)           # (B, embed_dim)
192
        return F.normalize(x, dim=-1)
193

194
    def forward(self, image, text, return_loss=False):
195
        """
196
        完整前向。
197

198
        参数:
199
            image:       (B, 3, H, W) 图像
200
            text:        list[str] 文本列表
201
            return_loss: 是否返回 InfoNCE 损失
202
        """
203
        # 编码
204
        image_features = self.encode_image(image)   # (B, embed_dim)
205
        text_features  = self.encode_text(text)     # (B, embed_dim)
206

207
        if return_loss:
208
            # 计算对比损失
209
            logit_scale = self.logit_scale.exp()
210
            logits_per_image = image_features @ text_features.T * logit_scale
211
            logits_per_text  = logits_per_image.T
212

213
            B = image.shape[0]
214
            labels = torch.arange(B, device=image.device)
215

216
            loss_i2t = F.cross_entropy(logits_per_image, labels)
217
            loss_t2i = F.cross_entropy(logits_per_text, labels)
218
            loss = (loss_i2t + loss_t2i) / 2
219
            return {"loss": loss, "image_features": image_features, "text_features": text_features}
220

221
        return {"image_features": image_features, "text_features": text_features}
222

223
    @torch.no_grad()
224
    def zero_shot_classify(self, image, class_names, prompt_template="a photo of a {}."):
225
        """
226
        零样本分类。
227

228
        参数:
229
            image:          (3, H, W) 图像 tensor
230
            class_names:    list[str] 类别名称
231
            prompt_template: str 带占位符的模板
232
        """
233
        # 构建提示
234
        prompts = [prompt_template.format(c) for c in class_names]
235

236
        # 编码
237
        image_features = self.encode_image(image.unsqueeze(0))
238
        text_features  = self.encode_text(prompts)
239

240
        # 相似度
241
        logit_scale = self.logit_scale.exp()
242
        similarity = (image_features @ text_features.T) * logit_scale
243
        probs = F.softmax(similarity, dim=-1)[0]
244

245
        return probs.cpu().numpy()

展开收起

9. 应用场景#

9.1 开放词汇检测与分割#

CLIP 的特征可以泛化到任意视觉概念，因此被广泛用于：

1
def clip_applications():
2
    """
3
    CLIP 的主要应用场景。
4
    """
5
    apps = {
6
        "开放词汇检测 (OVOD)": {
7
            "代表": "GLIP, MDETR, OWLVIT",
8
            "方法": "用 CLIP 特征替换检测器的类别嵌入",
9
            "效果": "检测任意文本描述的物体",
10
        },
11
        "零样本分割 (ZS3)": {
12
            "代表": "Lseg, ZeroSeg, SA-V",
13
            "方法": "CLIP 特征 + 分割模型",
14
            "效果": "用文本描述来分割图像区域",
15
        },
16
        "图像检索": {
17
            "代表": "LAION-5B 检索",
18
            "方法": "图像-图像 / 文本-图像相似度搜索",
19
            "效果": "给定图像找相似图 / 给定文本找图像",
20
        },
21
        "图像生成条件": {
22
            "代表": "Stable Diffusion, DALL-E 3",
23
            "方法": "CLIP text encoder 作为生成模型的条件",
24
            "效果": "CLIP 理解文本 → 指导图像生成",
25
        },
26
        "分布外检测": {
27
            "代表": "CLIP Score, Winoground",
28
            "方法": "用 CLIP 相似度衡量图像-文本一致性",
29
            "效果": "判断模型是否正确理解图像内容",
30
        },
31
    }
32
    return apps

展开收起

9.2 CLIP 作为视觉特征提取器#

1
def extract_visual_features(model, image_paths):
2
    """
3
    用 CLIP 提取视觉特征，用于下游任务。
4
    """
5
    features = []
6
    for path in image_paths:
7
        image = Image.open(path).convert("RGB")
8
        image_tensor = preprocess(image).unsqueeze(0)
9

10
        with torch.no_grad():
11
            feat = model.encode_image(image_tensor)
12

13
        features.append(feat)
14

15
    return torch.cat(features, dim=0)  # (N, embed_dim)
16

17

18
def extract_text_features(model, texts):
19
    """
20
    用 CLIP 提取文本特征。
21
    """
22
    with torch.no_grad():
23
        features = model.encode_text(texts)
24
    return features

展开收起

10. 局限性#

1
def clip_limitations():
2
    """
3
    CLIP 的主要局限性。
4
    """
5
    limitations = {
6
        "分布偏置": {
7
            "问题": "训练数据偏向英语和西方文化，对其他语言和文化适应性差",
8
            "数据": "WIT-400M 主要是英语图文对",
9
        },
10
        "长尾分布": {
11
            "问题": "对常见物体分类好，对稀有/抽象概念分类差",
12
            "原因": "训练数据中稀有概念出现少",
13
        },
14
        "文字识别": {
15
            "问题": "CLIP 不擅长 OCR，文字区域的特征不佳",
16
            "改进": "需要额外的 OCR 预训练（如 CLIP-OCR）",
17
        },
18
        "空间关系": {
19
            "问题": "CLIP 不擅长理解物体之间的空间关系",
20
            "示例": "'红在蓝左边' vs '红在蓝右边' 难以区分",
21
            "原因": "图文对数据中空间描述较少",
22
        },
23
        "计数": {
24
            "问题": "CLIP 不擅长精确计数",
25
            "示例": "'三只猫' vs '两只猫' 容易混淆",
26
            "原因": "自然语言中精确数量描述少",
27
        },
28
        "对抗鲁棒性": {
29
            "问题": "对对抗噪声的鲁棒性较低",
30
            "原因": "对比学习没有显式防御机制",
31
        },
32
        "组合推理": {
33
            "问题": "不擅长理解属性的组合（颜色+形状+材质）",
34
            "示例": "'蓝色的圆形' vs '红色的方形'",
35
            "改进": "需要组合推理能力的改进（LEO, COMET）",
36
        },
37
    }
38
    return limitations

展开收起

11. 总结#

11.1 核心要点#

维度	关键要点
核心创新	用自然语言作为监督信号，打破类别标签的限制
模型架构	双塔：ViT 视觉编码器 + Transformer 文本编码器
损失函数	对称 InfoNCE（两方向的对比损失）
训练数据	WIT-400M（4 亿图文对）
零样本	任何文本描述都可以作为分类标签
提示工程	多种提示模板集成可显著提升精度
扩展方向	ALIGN（更大数据）/ EVA-CLIP（两阶段训练）/ SigLIP（sigmoid 损失）
下游应用	开放词汇检测、分割、图像检索、生成条件、图文匹配

11.2 一句话总结#

CLIP 的核心思想简单而深刻：与其让人类给图像打上固定的类别标签，不如直接用自然语言作为监督——让视觉编码器和文本编码器在同一个特征空间里”对齐”。通过 4 亿图文对的对比学习，CLIP 学到了通用的视觉-语义表示，实现了真正的零样本迁移，成为几乎所有多模态模型（VLM、VLA、图像生成）的视觉编码器基础。

11.3 推荐资源#

1
论文:
2
  - CLIP (Radford et al., 2021): "Learning Transferable Visual Models From Natural Language Supervision"
3
  - ALIGN (Jia et al., 2021): "Scaling Up Visual and Vision-Language Representation Learning"
4
  - EVA-CLIP (Sun et al., 2023): "EVA-CLIP: Improved Training of CLIP on Large-scale Web Data"
5
  - SigLIP (Zhai et al., 2024): "Sigmoid Loss for Language-Image Pre-Training"
6
  - OpenCLIP (Ilharco et al., 2021): "OpenCLIP"
7

8
代码:
9
  - openai/CLIP (官方实现)
10
  - mlfoundations/open_clip (OpenCLIP, 支持更多模型)
11
  - facebookresearch/CLIP (PyTorch 官方重实现)

Consistency Models 深度解析：一步生成扩散模型

Thu, 16 Jul 2026 00:00:00 GMT

1. 从多步采样到一步生成#

1.1 扩散模型的核心瓶颈#

DDPM 需要 100-1000 步才能生成一张高质量图像。即使是加速后的 DDIM 也需要 20-50 步。每次生成都意味着：

1
DDPM 多步采样的代价:
2

3
  模型前向: 1000 步 × N 层 × 复杂度 = 高昂计算成本
4
  延迟: 每步 ~50ms → 1000步 ≈ 50秒 → 无法实时
5
  显存: 所有中间激活值需要保存
6
  能耗: 移动端 / 边缘部署几乎不可能
7

8
实际需求:
9
  ✓ "我要生成一张猫的照片" → 等待 50 秒？
10
  ✗ "我要实时生成视频帧" → 根本不可能
11
  ✗ "在手机上运行 SD" → 显存和延迟双重爆炸

1.2 加速采样的三条技术路线#

1
┌──────────────────────────────────────────────────────────────┐
2
│  扩散模型加速采样的三条路线                                   │
3
├──────────────────────────────────────────────────────────────┤
4
│                                                              │
5
│  路线 1: 更少的采样步数 (DDIM / PNDM)                      │
6
│  ────────────────────────────────────────────────────────   │
7
│  思想: 跳过部分中间步，通过 ODE 插值重建                     │
8
│  效果: 50-100 步仍可接受，20 步开始明显质量下降             │
9
│  缺点: 仍有固定步数，延迟有下限                              │
10
│                                                              │
11
│  路线 2: 预测更少的变量 (Latent Consistency Models / LCM)   │
12
│  ────────────────────────────────────────────────────────   │
13
│  思想: 预测 x_0 或一致性函数，减少步数                       │
14
│  效果: 2-8 步生成，质量接近全步                              │
15
│  缺点: 需要额外引导，对某些模型结构有依赖                    │
16
│                                                              │
17
│  路线 3: 一步直接生成 (Consistency Models)                 │
18
│  ────────────────────────────────────────────────────────   │
19
│  思想: 训练一个模型，直接从噪声一步映射到数据                 │
20
│  效果: 1 步生成，质量通过蒸馏提升                           │
21
│  缺点: 需要从预训练扩散模型蒸馏或从零训练                   │
22
│                                                              │
23
└──────────────────────────────────────────────────────────────┘

展开收起

1.3 Consistency Models 的核心洞察#

Song 等人 (OpenAI) 的核心洞察是：DDPM 的前向过程和反向过程本质上是一条 ODE 轨迹——任意时间步的 x 都沿着这条轨迹流向数据分布或噪声分布。既然如此，我们可以直接学习这条轨迹的”一致性函数”：轨迹上的任意点，都应该映射到轨迹的起点（即数据分布）。这就是 Consistency Models 的核心思想——训练一个函数 f，使得 f(x_t) = x_0 对所有 t ∈ [0, T] 成立。

1.4 论文信息#

1
论文: "Consistency Models"
2
作者: Yang Song, Prafulla Dhariwal, Mark Chen, Ilya Sutskever
3
单位: OpenAI
4
发表于: ICML 2023
5
引用: > 2,000 次
6
开源: https://github.com/openai/consistency_models

2. 核心思想：ODE 轨迹与一致性映射#

2.1 扩散模型的 ODE 视角#

DDPM 的前向过程 q(x_t | x_0) 可以被等价的 ODE（随机微分方程）描述：

1
def diffusion_ode():
2
    """
3
    扩散模型的随机微分方程 (SDE) 视角。
4

5
    关键洞察:
6
    前向过程和反向过程都在同一条轨迹上移动。
7
    """
8

9
    return {
10
        "前向 SDE (添加噪声)": {
11
            "方程": "dx = σ_max^t · dW_t",
12
            "解释": " Wiener 过程 W_t 驱动 x 趋向噪声",
13
            "边界": "x_0 → x_T (T → ∞ 时，x_T ≈ N(0, σ_max²))",
14
        },
15
        "反向 ODE (生成采样)": {
16
            "方程": "dx = [x - ∇_x log p_θ(x, t)] dt",
17
            "解释": "得分函数 ∇ log p 指向数据密度梯度",
18
            "关键": "这是确定性路径，从噪声流向数据",
19
        },
20
        "轨迹一致性": {
21
            "观察": "反向 ODE 定义了一条从 x_T 到 x_0 的路径",
22
            "性质": "轨迹上任意点 x_t 都是同一个扩散过程的"快照"",
23
            "需求": "我们希望一个函数 f 满足: f(x_t) = x_0 对所有 t",
24
        },
25
    }
26

27

28
def consistency_mapping():
29
    """
30
    一致性映射。
31

32
    目标: 学习一个函数 f_θ
33
    使得: f_θ(x_t) = x_0 对所有 t ∈ [0, T]
34
    其中 x_0 和 x_t 是同一轨迹上的点。
35

36
    这意味着:
37
    - f_θ 沿着 ODE 轨迹是"一致的"
38
    - 轨迹上的所有点映射到同一点 (轨迹起点)
39
    - 一旦学会，从任意中间步 x_t 一步到 x_0
40
    """
41

42
    return {
43
        "数学定义": "f_θ(x_t) = x_0",
44
        "物理含义": "轨迹上任意时刻的状态都"记住"了起点",
45
        "推理优势": "从任意 x_t，直接一步到达 x_0",
46
        "训练挑战": "x_0 是未知的 (我们只有噪声样本)",
47
    }

展开收起

2.2 一致性函数的数学推导#

1
def consistency_function_derivation():
2
    """
3
    一致性函数的数学推导。
4

5
    核心问题: 如何训练 f_θ 使得 f_θ(x_t) ≈ x_0？
6

7
    方法 1: 自监督学习
8
    ─────────────────────────────────────────────
9
    1. 采样 x_0 ~ p_data (真实数据)
10
    2. 采样 t_1, t_2 ~ Uniform(ε, T), 其中 t_1 < t_2
11
    3. 通过前向过程: x_{t_1} = α_{t_1} x_0 + σ_{t_1} ε_1
12
                      x_{t_2} = α_{t_2} x_0 + σ_{t_2} ε_2
13
    4. 损失: |f_θ(x_{t_1}) - x_{t_2}|²
14

15
    关键观察:
16
    - 如果 t_2 → 0，则 x_{t_2} ≈ x_0
17
    - f_θ(x_{t_1}) 应该接近 x_{t_2} (都接近 x_0)
18
    - 不需要知道真实的 x_0！
19

20
    方法 2: 蒸馏 (从预训练扩散模型)
21
    ─────────────────────────────────────────────
22
    1. 使用预训练的 DDPM 作为教师
23
    2. 采样 x_{t+ε}，用教师模型计算 x_0^teacher
24
    3. 损失: |f_θ(x_t) - x_0^teacher|²
25

26
    方法 3: TTCT (Two-Track Consistency Training)
27
    ─────────────────────────────────────────────
28
    结合方法 1 和 2，用更强的教师信号
29
    """
30

31
    return {
32
        "TTCT 损失": {
33
            "公式": "L = λ_1 · |f_θ(x_t) - f_φ(x_{t+ε})|² + λ_2 · |f_θ(x_t) - x_0^teacher|²",
34
            "第一项": "自监督一致性 (不依赖教师)",
35
            "第二项": "蒸馏一致性 (利用预训练知识)",
36
            "优点": "既能从零训练，又保留蒸馏的效率",
37
        },
38
    }

展开收起

2.3 EDM 框架整合#

Consistency Models 被设计为可以直接在 EDM（Elucidating the Design Space of Diffusion Models）框架上工作：

1
def edm_integration():
2
    """
3
    Consistency Models 与 EDM 的整合。
4

5
    EDM 的核心设计:
6
    - 使用网络输出 c_out = σ_data · F_θ(x, t)
7
    - 其中 F_θ 是 UNet，σ_data 是数据标准差
8
    - 采样使用 EDM 的 ODE 求解器
9

10
    Consistency Model 的 EDM 版本:
11
    - f_θ(x_t, t) 直接输出 x_0
12
    - 网络: F_θ(x, t) 输出 c_skip(t)·x + c_out(t)·F_θ(x, t)
13
    - 其中 c_skip 和 c_out 是时间依赖的跳跃连接权重
14

15
    优势:
16
    - 与 SD 的 EDM 实现兼容
17
    - 可以直接加载 SD 的预训练权重进行蒸馏
18
    """
19

20
    return {
21
        "网络输出": {
22
            "标准 EDM": "c_out = σ_data · F_θ(x, t)",
23
            "Consistency Model": "x_0 = c_skip(t)·x + c_out(t)·F_θ(x, t)",
24
        },
25
        "时间依赖权重": {
26
            "c_skip(t)": "控制输入 x 的残差比例",
27
            "c_out(t)": "控制网络输出的比例",
28
            "设计": "与 EDM 相同，通过 MLP 从 t 计算",
29
        },
30
    }

展开收起

3. 训练方法：TTCT 双轨训练#

3.1 TTCT 详解#

TTCT（Two-Track Consistency Training）是 Consistency Models 提出的双轨训练方法：

1
TTCT 训练流程:
2

3
  轨道 1: 自监督一致性
4
  ───────────────────────────────────────────────────────
5
  x_0 ~ p_data
6
      ↓
7
  t_1, t_2 ~ Uniform(ε, T), t_1 < t_2
8
      ↓
9
  x_{t_1} = α_{t_1} x_0 + σ_{t_1} ε_1
10
  x_{t_2} = α_{t_2} x_0 + σ_{t_2} ε_2
11
      ↓
12
  损失: L_1 = |f_θ(x_{t_1}) - x_{t_2}|²
13

14
  轨道 2: 蒸馏一致性 (从预训练教师)
15
  ───────────────────────────────────────────────────────
16
  x_ε ~ N(0, σ²)
17
      ↓
18
  用预训练教师模型计算 x_0^teacher
19
      ↓
20
  损失: L_2 = |f_θ(x_ε) - x_0^teacher|²
21

22
  总损失: L = λ_1 · L_1 + λ_2 · L_2

展开收起

1
def ttct_loss():
2
    """
3
    TTCT 损失函数实现。
4
    """
5
    def compute_ttct_loss(student_model, teacher_model, p_data,
6
                          optimizer, config):
7
        """
8
        参数:
9
            student_model: 要训练的 Consistency Model
10
            teacher_model: 预训练的 DDPM (用于蒸馏)
11
            p_data: 真实数据分布
12
            config: 训练配置
13
        """
14
        # 采样一个 batch
15
        x_0 = sample_batch(p_data, batch_size=config.batch_size)
16

17
        # ============ 轨道 1: 自监督一致性 ============
18
        # 采样两个时间步
19
        epsilon = config.epsilon  # 最小时间步
20
        T = config.T             # 最大时间步
21

22
        t_1 = torch.rand(batch_size) * (T - epsilon) + epsilon
23
        t_2 = torch.rand(batch_size) * (T - t_1) + t_1  # t_2 > t_1
24

25
        # 加噪
26
        noise_1 = torch.randn_like(x_0)
27
        noise_2 = torch.randn_like(x_0)
28

29
        x_t1 = alpha(t_1).view(-1, 1, 1, 1) * x_0 + \
30
               sigma(t_1).view(-1, 1, 1, 1) * noise_1
31

32
        x_t2 = alpha(t_2).view(-1, 1, 1, 1) * x_0 + \
33
               sigma(t_2).view(-1, 1, 1, 1) * noise_2
34

35
        # 学生模型预测
36
        pred_t1 = student_model(x_t1, t_1)
37

38
        # 自监督损失: f_θ(x_{t_1}) 应该接近 x_{t_2}
39
        # (因为两者都接近 x_0)
40
        loss_1 = F.mse_loss(pred_t1, x_t2.detach())
41

42
        # ============ 轨道 2: 蒸馏一致性 ============
43
        # 采样噪声
44
        x_epsilon = torch.randn_like(x_0) * sigma(epsilon)
45

46
        # 教师模型预测 (全步 DDPM)
47
        with torch.no_grad():
48
            x_0_teacher = teacher_model(x_epsilon, epsilon)
49

50
        # 学生模型预测
51
        pred_eps = student_model(x_epsilon, epsilon)
52

53
        # 蒸馏损失
54
        loss_2 = F.mse_loss(pred_eps, x_0_teacher.detach())
55

56
        # ============ 总损失 ============
57
        loss = config.lambda_1 * loss_1 + config.lambda_2 * loss_2
58

59
        return loss, {"self_supervised": loss_1, "distillation": loss_2}
60

61
    return compute_ttct_loss

展开收起

3.2 从零训练 vs 蒸馏#

1
def training_modes():
2
    """
3
    Consistency Models 的两种训练模式。
4
    """
5
    return {
6
        "Training-from-Scratch": {
7
            "方法": "不使用预训练教师，只用自监督一致性",
8
            "损失": "L = |f_θ(x_{t_1}) - x_{t_2}|²",
9
            "优点": "不需要预训练扩散模型",
10
            "缺点": "需要更多数据和训练时间",
11
            "适用": "从头构建新模型",
12
            "收敛": "慢，但最终质量可达",
13
        },
14
        "SCT (Sampling Consistency Training)": {
15
            "方法": "从零训练 + 使用单步学生作为教师",
16
            "损失": "L = |f_θ(x_t) - f_φ(x_{t+Δ})|²",
17
            "优点": "渐进提升，不需要预训练教师",
18
            "缺点": "需要迭代训练教师-学生",
19
            "适用": "持续改进模型",
20
        },
21
        "CFG (Consistency Fine-tuning)": {
22
            "方法": "在预训练 CM 基础上微调，加入 CFG",
23
            "损失": "L = |f_θ(x_t, c) - x_0^teacher(x_t, c)|²",
24
            "优点": "继承预训练 CM + 支持条件生成",
25
            "缺点": "需要额外的条件编码器",
26
            "适用": "快速适配下游任务",
27
        },
28
        "CD (Consistency Distillation)": {
29
            "方法": "从预训练扩散模型蒸馏",
30
            "损失": "L = |f_θ(x_t) - x_0^teacher(x_t)|²",
31
            "优点": "快速收敛，质量高",
32
            "缺点": "需要预训练扩散模型",
33
            "适用": "已有 SD 模型，想加速采样",
34
        },
35
    }

展开收起

4. 网络架构#

4.1 与 EDM UNet 的兼容性#

Consistency Models 直接复用 EDM 的 UNet 架构，只需修改输出头：

1
class ConsistencyModelEDM(nn.Module):
2
    """
3
    基于 EDM UNet 的 Consistency Model。
4

5
    核心修改:
6
    - 网络输出改为预测 x_0 (而非噪声或 v-prediction)
7
    - 添加时间依赖的 c_skip 和 c_out 跳跃连接
8
    """
9

10
    def __init__(self, unet_backbone):
11
        super().__init__()
12
        self.unet = unet_backbone
13

14
        # 时间嵌入 → 跳跃连接权重
15
        self.c_skip_net = nn.Sequential(
16
            nn.Linear(256, 1),
17
            nn.Sigmoid(),
18
        )
19
        self.c_out_net = nn.Sequential(
20
            nn.Linear(256, 1),
21
            nn.Sigmoid(),
22
        )
23

24
    def forward(self, x, t):
25
        """
26
        参数:
27
            x: (B, C, H, W) 输入 (可以是任意时间步的 x_t)
28
            t: (B,) 时间步 (0 到 T)
29
        返回:
30
            x_0: (B, C, H, W) 预测的原始数据
31
        """
32
        # 时间嵌入
33
        t_emb = get_timestep_embedding(t, 128)
34
        t_emb = self.unet.time_embed(t_emb)  # (B, 256)
35

36
        # EDM UNet 前向
37
        h = self.unet.prep(x)
38
        h = self.unet.downsample(h)
39
        h = self.unet.attention(h, t_emb)
40
        h = self.unet.upsample(h)
41

42
        # EDM 输出: F_θ(x, t)
43
        f_theta = self.unet.final(h)  # (B, C, H, W)
44

45
        # 时间依赖跳跃连接
46
        c_skip = self.c_skip_net(t_emb)        # (B, 1) ∈ (0, 1)
47
        c_out = self.c_out_net(t_emb)          # (B, 1) ∈ (0, 1)
48

49
        # Consistency Model 输出: x_0 = c_skip · x + c_out · F_θ
50
        x_0 = c_skip.view(-1, 1, 1, 1) * x + \
51
              c_out.view(-1, 1, 1, 1) * f_theta
52

53
        return x_0

展开收起

4.2 边界条件处理#

1
def boundary_conditions():
2
    """
3
    Consistency Models 的边界条件。
4

5
    t = 0 时:
6
    - x_t = x_0 (无噪声)
7
    - f_θ(x_0, 0) 应该 = x_0 (恒等映射)
8

9
    t = T 时:
10
    - x_T ≈ 纯噪声
11
    - f_θ(x_T, T) 应该 ≈ x_0 (仍然预测数据)
12
    """
13

14
    return {
15
        "t = 0 (边界)": {
16
            "条件": "f_θ(x, 0) = x",
17
            "实现": "强制 c_skip(0) = 1, c_out(0) = 0",
18
            "原因": "此时 x 就是真实数据，不需要预测",
19
        },
20
        "t = T (边界)": {
21
            "条件": "f_θ(x, T) = E[x_0 | x_T = x]",
22
            "实现": "网络在高噪声时预测数据分布的均值",
23
            "原因": "从纯噪声中重建数据",
24
        },
25
        "中间 t": {
26
            "条件": "f_θ(x_t, t) 应该介于 x_t 和 x_0 之间",
27
            "实现": "c_skip(t) 和 c_out(t) 通过 MLP 学习",
28
            "观察": "t 越大，c_skip 越小（x_t 越偏离 x_0）",
29
        },
30
    }

展开收起

5. 推理：多步生成与 TTA#

5.1 基本推理流程#

1
@torch.no_grad()
2
def cm_sample(model, x_T, num_steps=1, guide_scale=1.0):
3
    """
4
    Consistency Model 推理。
5

6
    参数:
7
        model: 训练好的 Consistency Model
8
        x_T: (B, C, H, W) 初始噪声
9
        num_steps: 采样步数 (1 = 一步生成)
10
        guide_scale: CFG 引导强度
11
    返回:
12
        x_0: (B, C, H, W) 生成的图像
13
    """
14

15
    if num_steps == 1:
16
        # ============ 一步生成 (核心) ============
17
        t = torch.full((x_T.shape[0],), 1.0, device=x_T.device)
18

19
        if guide_scale == 1.0:
20
            # 无条件生成
21
            x_0 = model(x_T, t)
22
        else:
23
            # CFG 引导
24
            # 需要条件和无条件输入
25
            x_0_cond = model(x_T, t, condition)      # 条件输入
26
            x_0_uncond = model(x_T, t, uncondition)  # 无条件输入
27
            x_0 = x_0_uncond + guide_scale * (x_0_cond - x_0_uncond)
28

29
        return x_0
30

31
    else:
32
        # ============ 多步生成 (TTA 增强) ============
33
        return cm_sample_tta(model, x_T, num_steps)
34

35

36
@torch.no_grad()
37
def cm_sample_tta(model, x_T, num_steps=4):
38
    """
39
    TTST (Test-Time Self-Taught) 多步增强。
40

41
    核心思想:
42
    - 多步生成比单步更稳定
43
    - 但步数仍然远少于 DDPM
44
    - 通过 self-distillation 持续改进
45
    """
46

47
    # 初始化
48
    x = x_T.clone()
49
    T = 1.0
50

51
    # 时间步调度 (从大到小)
52
    timesteps = torch.linspace(T, 0, num_steps + 1).tolist()
53

54
    for i in range(num_steps):
55
        t = timesteps[i]
56
        t_next = timesteps[i + 1]
57

58
        # 单步预测
59
        x_0_pred = model(x, torch.full((x.shape[0],), t, device=x.device))
60

61
        # 如果 t_next = 0，则这就是最终输出
62
        if t_next == 0:
63
            return x_0_pred
64

65
        # 否则，用预测的 x_0 重建下一步
66
        # x_{t_next} = α_{t_next} · x_0 + σ_{t_next} · noise
67
        # 其中 noise = (x - α_t · x_0) / σ_t
68
        alpha_t = alpha(torch.tensor(t))
69
        sigma_t = sigma(torch.tensor(t))
70
        alpha_next = alpha(torch.tensor(t_next))
71
        sigma_next = sigma(torch.tensor(t_next))
72

73
        noise = (x - alpha_t * x_0_pred) / sigma_t.clamp(min=1e-6)
74
        x = alpha_next * x_0_pred + sigma_next * noise
75

76
    return x

展开收起

5.2 TTST 推理详解#

1
def ttst_explained():
2
    """
3
    TTST (Test-Time Self-Taught) 推理详解。
4

5
    目标: 用多步推理提升单步 CM 的质量。
6

7
    原理:
8
    1. 从噪声 x_T 开始
9
    2. 用 CM 预测 x_0
10
    3. 用预测的 x_0 重建中间步 x_{t-Δ}
11
    4. 再从 x_{t-Δ} 预测 x_0
12
    5. 重复 N 次，质量逐步提升
13

14
    为什么有效:
15
    - 每一步都在"纠正"上一步的误差
16
    - 与 ODE 求解器的思想类似
17
    - 但不需要梯度，计算量小
18
    """
19

20
    return {
21
        "单步 TTST": {
22
            "步骤": "x_T → f_θ(x_T) = x_0",
23
            "质量": "★★★☆☆",
24
            "速度": "★★★★★",
25
        },
26
        "4步 TTST": {
27
            "步骤": "x_T → x_{0.75} → x_{0.5} → x_{0.25} → x_0",
28
            "质量": "★★★★☆",
29
            "速度": "★★★★☆",
30
        },
31
        "8步 TTST": {
32
            "步骤": "更细粒度的多步重建",
33
            "质量": "★★★★★",
34
            "速度": "★★★☆☆",
35
        },
36
        "对比 DDPM": {
37
            "DDPM 100步": "质量 ★★★★★, 速度 ★☆☆☆☆",
38
            "CM 1步": "质量 ★★★☆☆, 速度 ★★★★★",
39
            "CM 4步 TTST": "质量 ★★★★☆, 速度 ★★★★☆",
40
        },
41
    }

展开收起

6. CFG 条件一致性#

6.1 为 Consistency Model 添加条件控制#

原始 Consistency Models 是无条件模型。通过 Consistency Fine-tuning（CFG），可以支持条件控制：

1
def cfg_consistency_model():
2
    """
3
    CFG Consistency Model。
4

5
    方法: 微调预训练 CM，添加条件信息。
6
    损失: L = |f_θ(x_t, c) - x_0^teacher(x_t, c)|²
7
    """
8

9
    class CFGConsistencyModel(nn.Module):
10
        def __init__(self, base_cm, text_encoder):
11
            super().__init__()
12
            self.base_cm = base_cm
13
            self.text_encoder = text_encoder
14

15
            # 额外的条件 embedding 投影
16
            self.cfg_head = nn.Sequential(
17
                nn.Linear(768 + 768, 256),  # cond + uncond → time_embed dim
18
                nn.SiLU(),
19
                nn.Linear(256, 256),
20
            )
21

22
        def forward(self, x, t, cond_embed=None):
23
            """
24
            参数:
25
                x: (B, C, H, W) 输入
26
                t: (B,) 时间步
27
                cond_embed: (B, D) 条件 embedding (如文本)
28
            """
29
            # 文本编码
30
            if cond_embed is None:
31
                cond_embed = torch.zeros(x.shape[0], 768, device=x.device)
32

33
            # 无条件 (全零 embedding)
34
            uncond_embed = torch.zeros_like(cond_embed)
35

36
            # 条件预测
37
            x_cond = self.base_cm(x, t, cond_embed)
38

39
            # 无条件预测
40
            x_uncond = self.base_cm(x, t, uncond_embed)
41

42
            # CFG 合并
43
            cfg_scale = 7.5
44
            x_0 = x_uncond + cfg_scale * (x_cond - x_uncond)
45

46
            return x_0
47

48
    return CFGConsistencyModel
49

50

51
def cfg_training():
52
    """
53
    CFG Consistency Model 训练。
54
    """
55

56
    def train_cfg_cm(student_cm, teacher_cm, p_data, p_text, optimizer, config):
57
        """
58
        参数:
59
            student_cm: 要训练的 CFG CM
60
            teacher_cm: 预训练的 CFG 扩散模型
61
            p_data: 图像分布
62
            p_text: 文本分布
63
        """
64
        # 采样
65
        x_0 = sample_batch(p_data)
66
        text = sample_batch(p_text)
67

68
        # 文本编码
69
        cond_embed = student_cm.text_encoder(text)
70

71
        # 加噪
72
        t = torch.rand(x_0.shape[0]) * (config.T - config.epsilon) + config.epsilon
73
        noise = torch.randn_like(x_0)
74
        x_t = alpha(t).view(-1, 1, 1, 1) * x_0 + sigma(t).view(-1, 1, 1, 1) * noise
75

76
        # 学生预测
77
        pred = student_cm(x_t, t, cond_embed)
78

79
        # 教师预测 (使用相同的条件)
80
        with torch.no_grad():
81
            target = teacher_cm(x_t, t, cond_embed)
82

83
        # CFG 损失
84
        loss = F.mse_loss(pred, target)
85

86
        return loss

展开收起

7. 完整 PyTorch 实现#

7.1 Consistency Model 完整训练#

1
import torch
2
import torch.nn as nn
3
import torch.nn.functional as F
4
import math
5
from typing import Optional
6

7

8
def alpha_sigma(t, sigma_data=0.5, p_mean=-1.2, p_std=1.2):
9
    """
10
    EDM 风格的 alpha 和 sigma 调度。
11
    """
12
    # t 是 [0, 1] 归一化的时间步
13
    # sigma(t) 使用指数调度
14
    sigma = sigma_data ** (1 + 0.5 * (t - 0.5))  # 非线性调度
15

16
    # alpha 从 t 计算
17
    # 使用 sigmoid 中心化
18
    c = -2.5 * (t - 0.5)
19
    alpha = 1 / (1 + math.e ** c)
20

21
    return alpha, sigma
22

23

24
class ConsistencyModel(nn.Module):
25
    """
26
    Consistency Model 完整实现。
27

28
    核心设计:
29
    - 输入任意时间步的 x_t，输出 x_0
30
    - 使用 EDM UNet 作为 backbone
31
    - 时间依赖的 c_skip 和 c_out 跳跃连接
32
    """
33

34
    def __init__(self, unet: nn.Module, sigma_data=0.5):
35
        super().__init__()
36
        self.unet = unet
37
        self.sigma_data = sigma_data
38

39
        # 时间嵌入维度 (与 UNet 匹配)
40
        time_embed_dim = 256
41

42
        # c_skip(t) 和 c_out(t) 的 MLP
43
        self.c_skip_net = nn.Sequential(
44
            nn.Linear(time_embed_dim, 256),
45
            nn.SiLU(),
46
            nn.Linear(256, 1),
47
        )
48
        self.c_out_net = nn.Sequential(
49
            nn.Linear(time_embed_dim, 256),
50
            nn.SiLU(),
51
            nn.Linear(256, 1),
52
        )
53

54
        # 时间步嵌入
55
        self.time_embed = nn.Sequential(
56
            nn.Linear(128, time_embed_dim),
57
            nn.SiLU(),
58
            nn.Linear(time_embed_dim, time_embed_dim),
59
        )
60

61
    def get_time_embedding(self, t):
62
        """获取时间步的正弦位置嵌入。"""
63
        half_dim = 64
64
        emb = math.log(10000) / (half_dim - 1)
65
        emb = torch.exp(torch.arange(half_dim, device=t.device) * -emb)
66
        emb = t[:, None] * emb[None, :]
67
        emb = torch.cat([emb.sin(), emb.cos()], dim=-1)
68
        return self.time_embed(emb)
69

70
    def forward(self, x, t):
71
        """
72
        前向传播。
73

74
        参数:
75
            x: (B, C, H, W) 输入 (任意时间步的 x_t)
76
            t: (B,) 时间步 (0 到 1)
77
        返回:
78
            x_0: (B, C, H, W) 预测的原始数据
79
        """
80
        # 时间嵌入
81
        t_emb = self.get_time_embedding(t)  # (B, 256)
82

83
        # EDM UNet 前向
84
        h = self.unet.prep(x)
85
        h = self.unet.downsample(h, t_emb)
86
        h = self.unet.attention(h, t_emb)
87
        h = self.unet.upsample(h, t_emb)
88

89
        # 网络输出 F_θ(x, t)
90
        f_theta = self.unet.final(h)
91

92
        # 时间依赖跳跃连接
93
        c_skip = self.c_skip_net(t_emb)  # (B, 1)
94
        c_out = self.c_out_net(t_emb)    # (B, 1)
95

96
        # 边界条件: t=0 时，c_skip=1, c_out=0
97
        # 边界条件: t=1 时，c_skip=0, c_out=1
98
        # 通过 Sigmoid 实现自动插值
99
        c_skip = torch.sigmoid(c_skip)
100
        c_out = torch.sigmoid(c_out)
101

102
        # Consistency Model 输出
103
        x_0 = c_skip.view(-1, 1, 1, 1) * x + c_out.view(-1, 1, 1, 1) * f_theta
104

105
        return x_0
106

107

108
class ConsistencyModelTrainer:
109
    """Consistency Model 训练器。"""
110

111
    def __init__(self, model: ConsistencyModel, teacher_model: Optional[nn.Module] = None,
112
                 device="cuda", lambda_1=1.0, lambda_2=1.0):
113
        self.model = model.to(device)
114
        self.teacher = teacher_model.to(device) if teacher_model else None
115
        self.device = device
116
        self.lambda_1 = lambda_1
117
        self.lambda_2 = lambda_2
118
        self.step = 0
119

120
    def train_step(self, x_0: torch.Tensor, cfg_emb: Optional[torch.Tensor] = None):
121
        """
122
        单步训练。
123

124
        参数:
125
            x_0: (B, C, H, W) 真实数据
126
            cfg_emb: (B, D) 条件 embedding (可选)
127
        """
128
        B = x_0.shape[0]
129

130
        # ============ 轨道 1: 自监督一致性 ============
131
        epsilon = 1e-3
132
        T = 1.0
133

134
        # 采样两个时间步 (t_1 < t_2)
135
        t_1 = torch.rand(B, device=self.device) * (T - epsilon) + epsilon
136
        t_2 = torch.rand(B, device=self.device) * (T - t_1) + t_1
137

138
        # 加噪
139
        noise_1 = torch.randn_like(x_0)
140
        noise_2 = torch.randn_like(x_0)
141

142
        alpha_1, sigma_1 = alpha_sigma(t_1)
143
        alpha_2, sigma_2 = alpha_sigma(t_2)
144

145
        x_t1 = alpha_1.view(-1, 1, 1, 1) * x_0 + sigma_1.view(-1, 1, 1, 1) * noise_1
146
        x_t2 = alpha_2.view(-1, 1, 1, 1) * x_0 + sigma_2.view(-1, 1, 1, 1) * noise_2
147

148
        # 学生预测
149
        pred_1 = self.model(x_t1, t_1)
150

151
        # 自监督损失: f_θ(x_{t_1}) ≈ x_{t_2}
152
        loss_1 = F.mse_loss(pred_1, x_t2.detach())
153

154
        # ============ 轨道 2: 蒸馏一致性 ============
155
        if self.teacher is not None:
156
            # 用教师模型计算目标
157
            t_eps = torch.full((B,), epsilon, device=self.device)
158
            alpha_eps, sigma_eps = alpha_sigma(t_eps)
159
            x_eps = alpha_eps * x_0 + sigma_eps * torch.randn_like(x_0)
160

161
            with torch.no_grad():
162
                # 教师预测 x_0
163
                target = self.teacher(x_eps, t_eps)
164

165
            # 学生预测
166
            pred_eps = self.model(x_eps, t_eps)
167

168
            # 蒸馏损失
169
            loss_2 = F.mse_loss(pred_eps, target.detach())
170

171
            total_loss = self.lambda_1 * loss_1 + self.lambda_2 * loss_2
172
        else:
173
            total_loss = loss_1
174
            loss_2 = torch.tensor(0.0)
175

176
        # 反向传播
177
        total_loss.backward()
178

179
        return {"total": total_loss.item(), "self_sup": loss_1.item(), "distill": loss_2.item()}
180

181
    @torch.no_grad()
182
    def sample(self, batch_size=1, steps=1, guide_scale=1.0):
183
        """
184
        推理采样。
185
        """
186
        # 初始化噪声
187
        x = torch.randn(batch_size, 3, 32, 32, device=self.device)
188

189
        if steps == 1:
190
            # 一步生成
191
            t = torch.ones(batch_size, device=self.device)
192
            x_0 = self.model(x, t)
193
        else:
194
            # TTST 多步生成
195
            timesteps = torch.linspace(1.0, 0.0, steps + 1, device=self.device)
196
            for i in range(steps):
197
                t = timesteps[i]
198
                t_next = timesteps[i + 1]
199

200
                pred = self.model(x, t.expand(batch_size))
201

202
                if t_next == 0:
203
                    return pred
204

205
                # 重建下一步
206
                alpha_t, sigma_t = alpha_sigma(t)
207
                alpha_next, sigma_next = alpha_sigma(t_next)
208

209
                noise = (x - alpha_t * pred) / sigma_t.clamp(min=1e-6)
210
                x = alpha_next * pred + sigma_next * noise
211

212
            x_0 = x
213

214
        return x_0

展开收起

8. LCM：Latent Consistency Models#

8.1 LCM 的核心思想#

LCM（Latent Consistency Models, Luo et al., 2023）是 Consistency Models 在 Stable Diffusion 上的应用：

1
LCM vs 标准 CM:
2

3
  Consistency Model:
4
    - 在像素空间操作
5
    - 需要处理高分辨率 (512×512×3)
6
    - 计算量大
7

8
  LCM (Latent Consistency Model):
9
    - 在潜空间操作 (使用 SD 的 VAE)
10
    - 处理低维潜变量 (64×64×4)
11
    - 计算量大幅降低
12
    - 可以直接加载 SD 权重进行蒸馏
13

14
  核心公式:
15
    x_0 = c_skip(t) · x + c_out(t) · ε_θ(x, t)
16

17
  LCM 的特殊处理:
18
    - 使用 SD 的 UNet 作为 backbone
19
    - 直接蒸馏 ε_θ 的预测
20
    - 输出 x_0 而非噪声 ε

展开收起

8.2 LCM 的训练与推理#

1
def lcm_training():
2
    """
3
    LCM 训练 (与 Consistency Model 相同，但操作在潜空间)。
4
    """
5

6
    def train_lcm(sd_unet, vae, text_encoder, dataloader, optimizer):
7
        """
8
        LCM 训练。
9
        """
10
        for batch in dataloader:
11
            images, texts = batch
12

13
            # VAE 编码到潜空间
14
            with torch.no_grad():
15
                z_0 = vae.encode(images)
16

17
            # 采样时间步
18
            t = torch.rand(B) * 0.998 + 0.002  # 避免 t=0, t=1
19

20
            # 加噪
21
            noise = torch.randn_like(z_0)
22
            z_t = alpha(t) * z_0 + sigma(t) * noise
23

24
            # 文本编码
25
            context = text_encoder(texts)
26

27
            # LCM 预测 x_0
28
            # 使用 SD UNet，但输出头替换为 CM 的 c_skip/c_out
29
            epsilon_pred = sd_unet(z_t, t, context)
30
            x_0_pred = (z_t - sigma(t) * epsilon_pred) / alpha(t)
31

32
            # 教师目标
33
            with torch.no_grad():
34
                target = x_0_pred  # 可以用更好的教师
35

36
            loss = F.mse_loss(x_0_pred, target)
37

38
            optimizer.zero_grad()
39
            loss.backward()
40
            optimizer.step()
41

42
    return train_lcm
43

44

45
@torch.no_grad()
46
def lcm_sample(lcm_model, text_encoder, prompt,
47
               num_steps=4, guidance_scale=7.0):
48
    """
49
    LCM 推理 — 4 步生成 SD 图像。
50
    """
51
    # 文本编码
52
    context = text_encoder([prompt])
53
    uncond = text_encoder([""])
54

55
    # 初始化噪声
56
    z = torch.randn(1, 4, 64, 64, device=device)
57

58
    # 调度
59
    timesteps = [0.9, 0.7, 0.4, 0.0]
60

61
    for i, t in enumerate(tqdm(timesteps)):
62
        t_tensor = torch.tensor([t], device=device)
63

64
        # CFG
65
        z_cond = lcm_model(z, t_tensor, context)
66
        z_uncond = lcm_model(z, t_tensor, uncond)
67
        z_0 = z_uncond + guidance_scale * (z_cond - z_uncond)
68

69
        if t == 0.0:
70
            break
71

72
        # 重建下一步
73
        alpha_t = alpha_of_t(t)
74
        alpha_next = alpha_of_t(timesteps[i + 1])
75
        sigma_t = sigma_of_t(t)
76
        sigma_next = sigma_of_t(timesteps[i + 1])
77

78
        noise = (z - alpha_t * z_0) / sigma_t.clamp(min=1e-6)
79
        z = alpha_next * z_0 + sigma_next * noise
80

81
    # VAE 解码
82
    image = vae.decode(z)
83

84
    return image

展开收起

9. 与其他方法的对比#

9.1 加速方法全面对比#

1
def comparison_all():
2
    """
3
    扩散模型加速方法的全面对比。
4
    """
5
    return {
6
        "DDIM": {
7
            "步数": "20-50 步",
8
            "质量": "★★★★☆ (20步) - ★★★★★ (50步)",
9
            "速度": "★★★☆☆",
10
            "方法": "确定性 ODE 采样",
11
            "优点": "无需训练，通用",
12
            "缺点": "仍有固定步数延迟",
13
        },
14
        "LCM": {
15
            "步数": "2-8 步",
16
            "质量": "★★★★☆",
17
            "速度": "★★★★★",
18
            "方法": "潜空间一致性模型",
19
            "优点": "步数极少，兼容 SD",
20
            "缺点": "需要蒸馏训练",
21
        },
22
        "Consistency Models": {
23
            "步数": "1 步 (可达 ★★★★☆)",
24
            "质量": "★★★★☆ (TTST) - ★★★☆☆ (单步)",
25
            "速度": "★★★★★",
26
            "方法": "ODE 轨迹一致性蒸馏",
27
            "优点": "一步生成，理论优美",
28
            "缺点": "单步质量有限，需 TTST 增强",
29
        },
30
        "LCM + TTST": {
31
            "步数": "4-8 步",
32
            "质量": "★★★★★",
33
            "速度": "★★★★☆",
34
            "方法": "LCM + 多步自改进",
35
            "优点": "质量和速度的最佳平衡",
36
            "缺点": "比单步慢",
37
        },
38
        "对抗蒸馏": {
39
            "步数": "1 步",
40
            "质量": "★★★★☆",
41
            "速度": "★★★★★",
42
            "方法": "GAN + 扩散模型蒸馏",
43
            "优点": "一步高质量",
44
            "缺点": "训练不稳定，mode collapse",
45
        },
46
    }

展开收起

9.2 Consistency Models vs DDPM#

1
def cm_vs_ddpm():
2
    """
3
    Consistency Models vs DDPM 对比。
4
    """
5
    return {
6
        "采样": {
7
            "DDPM": "100-1000 步随机采样",
8
            "CM": "1-N 步确定性/随机采样",
9
        },
10
        "损失": {
11
            "DDPM": "噪声预测 |ε - ε_θ|²",
12
            "CM": "一致性 |f_θ(x_t) - f_θ(x_{t+Δ})|²",
13
        },
14
        "训练": {
15
            "DDPM": "单阶段，从零训练",
16
            "CM": "双轨 (自监督 + 蒸馏)，或从零训练",
17
        },
18
        "条件": {
19
            "DDPM": "天然支持 (cross-attention)",
20
            "CM": "需要 CFG 微调",
21
        },
22
        "理论基础": {
23
            "DDPM": "变分推断 + ELBO",
24
            "CM": "ODE 轨迹一致性 + 最优传输",
25
        },
26
        "生成模式": {
27
            "DDPM": "自回归多步",
28
            "CM": "直接映射 (类似 GAN)",
29
        },
30
    }

展开收起

10. 总结#

10.1 核心要点#

维度	关键要点
核心思想	ODE 轨迹一致性：轨迹上任意点映射到轨迹起点 x_0
数学框架	f_θ(x_t) = x_0 对所有 t ∈ [0, T]
训练目标	TTCT = 自监督一致性 + 蒸馏一致性
边界条件	t=0 时 f_θ(x, 0) = x，t=T 时预测数据均值
推理方式	单步生成或 TTST 多步增强
LCM	Consistency Models + 潜空间 + SD 权重
条件控制	CFG Consistency Fine-tuning
与 DDPM 关系	CM 压缩了 DDPM 的多步采样为单步映射

10.2 采样速度对比#

1
生成一张 512×512 图像所需的前向次数:
2

3
DDPM (100步)    ████████████████████████████████████ 100步  ~50秒
4
DDPM (1000步)   ██████████████████████████████████████████████ 1000步  ~500秒
5
DDIM (50步)     ██████████████████                    50步   ~25秒
6
DDIM (20步)     ████████                              20步   ~10秒
7
LCM (4步)       █                                      4步   ~2秒
8
CM (1步)        █                                      1步   ~0.5秒
9
CM (TTST 4步)  ████                                   4步   ~2秒
10
GAN (1步)      █                                      1步   ~0.5秒
11

12
质量对比:
13
GAN (1步)      ████████████████████████████████  ~85% 质量
14
CM (1步)       ██████████████████████████  ~80% 质量
15
LCM (4步)      █████████████████████████████████████  ~95% 质量
16
CM (TTST 8步) ██████████████████████████████████████  ~97% 质量
17
DDIM (50步)   ████████████████████████████████████████  ~99% 质量
18
DDPM (100步)  ███████████████████████████████████████████████  100% 质量

展开收起

10.3 一句话总结#

Consistency Models 的核心创新是将扩散模型的”多步去噪”视为一条 ODE 轨迹，并训练一个”一致性函数”f_θ，使得轨迹上的任意点都能一步映射到数据分布的起点 x_0——这通过 TTCT 双轨训练（自监督一致性 + 蒸馏一致性）和 EDM 兼容的 UNet 架构实现。CM 可以从零训练也可以从预训练扩散模型蒸馏，支持 1 步到多步推理；在潜空间的变体 LCM 则与 Stable Diffusion 无缝结合，实现了 4 步生成高质量图像，代表了扩散模型从”慢速多步”向”快速一步”演进的重要技术路线。

10.4 推荐资源#

1
论文:
2
  - Consistency Models (Song et al., 2023): "Consistency Models"
3
    ICML 2023, OpenAI
4
  - LCM (Luo et al., 2023): "Latent Consistency Models: Synthesizing
5
    Convergence-Offering Latent Space for Image Generation"
6
  - EDM (Karras et al., 2022): "Elucidating the Design Space of Diffusion Models"
7

8
代码:
9
  - openai/consistency_models: https://github.com/openai/consistency_models
10
  - openai/scoreODE: https://github.com/openai/scoreODE
11
  - Latent-Consistency-Models/sd-webui-lcm: LCM for WebUI
12

13
相关:
14
  - DDIM (Song et al., 2021): DDIM 采样加速
15
  - PF-ODE (Lu et al., 2022): Rectified Flow，另一种 ODE 加速
16
  - flow-matching: (Lipman et al., 2022) 统一扩散和流匹配

展开收起

ControlNet 深度剖析：条件控制扩散模型的结构化控制系统

Thu, 16 Jul 2026 00:00:00 GMT

1. ControlNet 的核心问题#

1.1 扩散模型缺乏可控性#

尽管 Stable Diffusion 等扩散模型能生成高质量图像，但”无条件生成”在实际应用中远远不够：

1
实际应用的需求:
2
  ✓ "生成一只猫" → 只需文本
3
  ✗ "生成一只坐着的猫" → 需要姿态控制
4
  ✗ "生成一只猫，放在我的书桌上" → 需要空间布局
5
  ✗ "按照我画的草图生成完整图像" → 需要边缘/轮廓控制
6
  ✗ "用我给的深度图作为参考" → 需要深度条件

核心矛盾：文本描述的是语义内容，但无法精确控制空间布局、物体姿态、边缘轮廓等结构性信息。

1.2 现有的四种失败方案#

ControlNet 论文（Adding Conditional Control to Text-to-Image Diffusion Models, Zhang et al., ICCV 2023）系统分析了已有的失败方案：

1
方案 1: 在 UNet 中加入额外输入通道
2
  思想: 把条件图像拼接到潜变量通道
3
  问题: ❌ 需要重新训练整个 UNet，显存爆炸
4
  参数量: +86M 额外参数需要从头训练
5

6
方案 2: 只微调额外输入层
7
  思想: 新增输入分支，只训练新层
8
  问题: ❌ 新层参数随机初始化，微调效果差
9
  原因: 随机初始化的层在前向传播中产生噪声，破坏预训练 UNet 的特征
10

11
方案 3: 注意力机制
12
  思想: 用 cross-attention 注入条件
13
  问题: ❌ 文本注意力不够强，结构条件信息丢失
14
  原因: 文本 token 数量有限，无法精细编码空间信息
15

16
方案 4: 分类器引导
17
  思想: 训练一个条件分类器，引导生成
18
  问题: ❌ 需要额外的分类器，数据标注成本高
19
  原因: 分类器的质量直接限制生成效果

展开收起

1.3 ControlNet 的核心洞察#

把预训练的 UNet 视为强大的”主脑”，在旁边训练一个可拆卸的”小脑”——这个”小脑”学会理解额外的条件输入（边缘、姿态、深度等），并通过渐进式零卷积初始化保证训练稳定性——小脑产生的信号从零开始逐步增长，不破坏主脑的预训练知识。

1.4 论文信息#

1
论文: "Adding Conditional Control to Text-to-Image Diffusion Models"
2
作者:Lvmin Zhang, Maneesh Agrawala
3
单位: Stanford University
4
发表于: ICCV 2023
5
引用: > 5,000 次 ★★★★★
6
开源: https://github.com/lllyasviel/ControlNet
7
在线体验: https://modelscope.cn/studios/iic/ControlNet

1.5 一句话概括 ControlNet#

ControlNet 的核心是”复制并冻结 UNet 编码器 + 训练一个可拆卸的副本”——冻结的编码器保证预训练知识不被破坏，副本在条件输入的监督下学会精确的结构控制，然后通过零卷积初始化的跳跃连接把控制信号渐进注入 UNet 解码器，实现对图像生成的结构级控制，同时保留文本到图像的质量。

2. 架构设计：复制 + 冻结 + 零卷积#

2.1 整体架构#

1
┌────────────────────────────────────────────────────────────────┐
2
│                                                                │
3
│   [条件图像 c] → 条件编码器                                     │
4
│      (Canny边缘/深度图/姿态)    ↓                              │
5
│                              可训练的副本分支                  │
6
│                              (ControlNet Block)               │
7
│                                        ↓                       │
8
│   [噪声潜变量 z_t] ───────────────────┐                        │
9
│           │                          ↓                         │
10
│   [时间步 t] ─→ UNet 编码器 ─→ [冻结 ❄️] ──────────────────┐ │
11
│                              ↓                               │ │
12
│                          UNet 解码器                         │ │
13
│                              ↑                               │ │
14
│                      跳跃连接 ← ← ← ← ← ← ← ← ← ← ← ← ← ← ┘ │
15
│                      (零卷积 ZeroConv)                        │
16
│                              ↑                                │
17
│                   ControlNet 副本编码器 ← 条件图像 c          │
18
│                   (训练 🔥)                                   │
19
│                                                                │
20
│   [文本条件] ──→ CLIP Text Encoder ──→ 交叉注意力注入        │
21
│                                                                │
22
└────────────────────────────────────────────────────────────────┘
23

24
关键设计:
25
  1. UNet 编码器完整复制 → 冻结 → 预训练权重
26
  2. 条件编码器: 与 UNet 编码器结构相同的副本 → 训练
27
  3. 零卷积: 跳跃连接处的 1×1 卷积，初始为零 → 渐进增长

展开收起

2.2 核心组件：ControlNet Block#

1
class ControlNetBlock(nn.Module):
2
    """
3
    ControlNet 的基本模块。
4
    与 UNet 的 ResBlock 结构相同，但包含跳跃连接的零卷积。
5
    """
6
    def __init__(self, in_channels, out_channels, context_dim=None):
7
        super().__init__()
8

9
        # 主体: ResBlock
10
        self.norm1 = nn.GroupNorm(32, in_channels)
11
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
12

13
        if context_dim is not None:
14
            # 如果需要，添加交叉注意力（ControlNet 不需要，这里预留给扩展）
15
            self.has_attn = False  # ControlNet 主体不用 attention
16

17
        self.norm2 = nn.GroupNorm(32, out_channels)
18
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
19

20
        # 残差连接
21
        if in_channels != out_channels:
22
            self.shortcut = nn.Conv2d(in_channels, out_channels, 1)
23
        else:
24
            self.shortcut = nn.Identity()
25

26
        # ===== 关键: 跳跃连接的零卷积 =====
27
        # 初始为零，训练时渐进增长
28
        self.zero_conv = ZeroConv2d(in_channels, out_channels)
29

30

31
class ZeroConv2d(nn.Module):
32
    """
33
    零卷积层。
34
    权重初始化为零 → 前向传播输出为零 → 不破坏预训练网络。
35
    训练时通过梯度更新逐渐增长。
36
    """
37
    def __init__(self, in_channels, out_channels):
38
        super().__init__()
39
        self.conv = nn.Conv2d(in_channels, out_channels, 1, padding=0)
40

41
        # 关键: 权重初始化为零
42
        nn.init.constant_(self.conv.weight, 0.0)
43
        nn.init.constant_(self.conv.bias, 0.0)
44

45
    def forward(self, x):
46
        return self.conv(x)

展开收起

2.3 零卷积初始化的数学保证#

1
def zero_conv_guarantee():
2
    """
3
    零卷积的数学保证。
4
    """
5
    analysis = {
6
        "初始化时": {
7
            "权重": "W = zero",
8
            "偏置": "b = zero",
9
            "输出": "y = W·x + b = zero",
10
            "效果": "ControlNet 对主网络的影响 = 0",
11
            "主网络行为": "完全等价于原始 SD，生成质量不变",
12
        },
13
        "训练过程中": {
14
            "梯度": "∂L/∂W ≠ 0 → W 开始更新",
15
            "更新方向": "由条件损失驱动，向条件方向调整",
16
            "增长速度": "初期慢（因为 W 接近零，梯度被放大）→ 后期稳定",
17
            "避免破坏": "W 从零增长，不会产生大的破坏性跳变",
18
        },
19
        "收敛时": {
20
            "权重": "W → W* (最优)",
21
            "效果": "ControlNet 学会了精确的条件控制",
22
            "主网络": "仍然冻结，预训练知识完整保留",
23
        },
24
    }
25
    return analysis

展开收起

2.4 完整的 ControlNet 架构#

1
class ControlNet(nn.Module):
2
    """
3
    完整的 ControlNet。
4
    复制 SD 的 UNet 编码器为可训练的副本。
5
    """
6

7
    def __init__(self, sd_unet):
8
        super().__init__()
9
        self.sd_unet = sd_unet
10

11
        # 冻结 SD UNet
12
        for param in sd_unet.parameters():
13
            param.requires_grad_(False)
14

15
        # 复制编码器结构为 ControlNet
16
        self.control_scales = [1.0] * 12  # 12 层控制强度
17

18
        # 每一层对应一个可训练的 ControlNet Block
19
        self.control_blocks = nn.ModuleList([
20
            self._make_control_block(ch) for ch in sd_unet.channel_list
21
        ])
22

23
        # 零卷积跳跃连接
24
        self.zero_convs = nn.ModuleList([
25
            ZeroConv2d(sd_unet.channel_list[i], sd_unet.channel_list[i])
26
            for i in range(len(sd_unet.channel_list))
27
        ])
28

29
        # 额外的条件图像处理
30
        self.condition_encoder = nn.Sequential(
31
            nn.Conv2d(3, 16, 3, padding=1),  # 假设条件图为 RGB
32
            nn.SiLU(),
33
            nn.Conv2d(16, 32, 3, stride=2, padding=1),
34
            nn.SiLU(),
35
            nn.Conv2d(32, 64, 3, stride=2, padding=1),
36
            nn.SiLU(),
37
            nn.Conv2d(64, 64, 3, padding=1),
38
        )
39

40
    def _make_control_block(self, channels):
41
        """创建一个 ControlNet Block。"""
42
        return nn.Sequential(
43
            ResBlock(channels, channels),
44
            ResBlock(channels, channels),
45
        )
46

47
    def forward(self, x_t, t, cond_image, context):
48
        """
49
        前向传播。
50

51
        参数:
52
            x_t:       (B, 4, H, W) 噪声潜变量
53
            t:         (B,) 时间步
54
            cond_image: (B, 3, H, W) 条件图像（边缘/深度等）
55
            context:   (B, seq_len, d) 文本 embedding
56
        """
57
        # 1) 编码条件图像
58
        cond = self.condition_encoder(cond_image)  # (B, 64, H/4, W/4)
59

60
        # 2) 获取 SD UNet 编码器特征
61
        #    (冻结的编码器用于提供特征，ControlNet 副本用于学习条件)
62
        features = self.sd_unet.get_encoder_features(x_t, t)
63

64
        # 3) 条件注入
65
        controlled_features = []
66
        for i, (feat, ctrl_block, zero_conv) in enumerate(zip(
67
            features, self.control_blocks, self.zero_convs
68
        )):
69
            # ControlNet 副本处理特征
70
            ctrl_feat = ctrl_block(feat + cond)  # 条件信息注入
71

72
            # 零卷积跳跃连接
73
            # 初始时 zero_conv 输出为 0 → 不影响主网络
74
            # 训练后: 输出控制信号
75
            skip = zero_conv(ctrl_feat)
76
            controlled_features.append(skip)
77

78
        # 4) SD UNet 解码器
79
        output = self.sd_unet.decode_with_skip(
80
            x_t, t, context, skip_connections=controlled_features
81
        )
82

83
        return output
84

85
    def unfix_layers(self):
86
        """
87
        解冻部分 SD UNet 层（可选的高级用法）。
88
        渐进解冻: 先训练零卷积，再解冻更多层。
89
        """
90
        pass  # ControlNet 主体设计为冻结 SD

展开收起

3. 训练策略：渐进式解冻#

3.1 零卷积的渐进式解冻#

ControlNet 的训练分为两个阶段：固定阶段和解冻阶段：

1
┌──────────────────────────────────────────────────────────────┐
2
│  阶段 1: 固定阶段 (Locked Stage)                              │
3
│                                                              │
4
│  SD UNet:  冻结 ❄️ (权重不变)                              │
5
│  ZeroConv: 冻结 🔒 (梯度计算，但权重不更新)                  │
6
│                                                              │
7
│  效果:                                                       │
8
│  - 零卷积从零开始，输出 = 0                                   │
9
│  - 只有 ControlNet Block 更新                               │
10
│  - 解码器收到的跳跃连接 = 0                                  │
11
│  - 主网络完全保持预训练状态                                  │
12
│                                                              │
13
├──────────────────────────────────────────────────────────────┤
14
│  阶段 2: 解冻阶段 (Unlocked Stage)                          │
15
│                                                              │
16
│  ZeroConv: 解冻 🔥 (零卷积开始学习)                         │
17
│                                                              │
18
│  效果:                                                       │
19
│  - 零卷积从零开始增长                                        │
20
│  - ControlNet 的控制信号逐渐注入                             │
21
│  - 避免突然的大跳变                                          │
22
│                                                              │
23
│  为什么需要解冻 ZeroConv?                                    │
24
│  - ZeroConv 需要学习最优的控制强度缩放                       │
25
│  - 某些层可能需要更强的控制，某些层需要更弱                   │
26
│  - 解冻后 ZeroConv 可以独立调整每层的贡献                    │
27
└──────────────────────────────────────────────────────────────┘

展开收起

3.2 ControlNet 训练循环#

1
def train_controlnet(controlnet, sd_unet, vae, text_encoder, dataloader,
2
                      optimizer, device, cfg_drop_prob=0.1):
3
    """
4
    ControlNet 训练循环。
5
    """
6
    controlnet.train()
7
    sd_unet.eval()  # 冻结
8

9
    for batch in tqdm(dataloader):
10
        images, cond_images, texts = (
11
            batch["image"].to(device),
12
            batch["condition"].to(device),  # Canny边缘/深度图/姿态
13
            batch["text"],
14
        )
15

16
        # ============ 1) 编码图像到潜空间 ============
17
        with torch.no_grad():
18
            z = vae.encode(images)  # (B, 4, 64, 64)
19

20
        # ============ 2) 采样时间步 ============
21
        batch_size = z.shape[0]
22
        t = torch.randint(0, 1000, (batch_size,), device=device)
23

24
        # ============ 3) 加噪 ============
25
        noise = torch.randn_like(z)
26
        alpha_bar = get_diffusion_schedule(t)  # 预设噪声调度
27
        z_t = alpha_bar.sqrt().view(-1, 1, 1, 1) * z + \
28
              (1 - alpha_bar).sqrt().view(-1, 1, 1, 1) * noise
29

30
        # ============ 4) 文本编码 + CFG Drop ============
31
        context = text_encoder.encode(texts)
32
        uncond_context = text_encoder.encode([""] * batch_size)
33

34
        # 随机丢弃条件（Classifier-Free Guidance）
35
        drop_mask = torch.rand(batch_size, device=device) < cfg_drop_prob
36
        context = context.masked_fill(drop_mask.unsqueeze(1), 0.0)
37

38
        # ============ 5) SD UNet 前向（冻结，提供特征） ============
39
        with torch.no_grad():
40
            # SD 的无条件预测
41
            noise_pred_uncond = sd_unet(z_t, t, uncond_context)
42

43
        # ============ 6) ControlNet 前向（训练） ============
44
        noise_pred_cond = controlnet(z_t, t, cond_images, context)
45

46
        # ============ 7) CFG 合并 ============
47
        # 训练时也用 CFG: 预测的噪声 = uncond + w * (cond - uncond)
48
        guidance_scale = 7.5
49
        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
50

51
        # ============ 8) 计算损失 ============
52
        loss = F.mse_loss(noise_pred, noise)
53

54
        # ============ 9) 反向传播 ============
55
        optimizer.zero_grad()
56
        loss.backward()
57
        torch.nn.utils.clip_grad_norm_(controlnet.parameters(), max_norm=1.0)
58
        optimizer.step()
59

60
    return loss.item()

展开收起

3.3 条件预处理#

不同的控制类型需要不同的预处理方法：

1
class ConditionPreprocessor:
2
    """
3
    条件图像预处理。
4
    把各种类型的条件图像标准化到模型需要的格式。
5
    """
6
    @staticmethod
7
    def canny_edge(image, low_threshold=100, high_threshold=200):
8
        """
9
        Canny 边缘检测 → 二值边缘图。
10
        """
11
        import cv2
12
        gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
13
        edges = cv2.Canny(gray, low_threshold, high_threshold)
14
        # 转为 RGB (3通道，与 SD 输入对齐)
15
        edges_rgb = np.stack([edges, edges, edges], axis=-1)
16
        return edges_rgb / 255.0
17

18
    @staticmethod
19
    def depth_map(image, model="MiDaS"):
20
        """
21
        单目深度估计 → 深度图。
22
        """
23
        # 使用 MiDaS 或 ZoeDepth 估计深度
24
        # 返回归一化的深度图
25
        pass
26

27
    @staticmethod
28
    def human_pose(image, model="OpenPose"):
29
        """
30
        人体姿态估计 → 关键点骨架图。
31
        """
32
        # 使用 OpenPose 或 DWPose 检测 17/25 个关键点
33
        # 绘制骨架连接
34
        pass
35

36
    @staticmethod
37
    def normal_map(image, model="MiDaS"):
38
        """
39
        表面法线估计 → 法线图。
40
        """
41
        # 估计表面法线，用于几何控制
42
        pass
43

44
    @staticmethod
45
    def semantic_seg(image, model="U2Net"):
46
        """
47
        语义分割 → 分割掩码图。
48
        """
49
        # 像素级类别分割
50
        pass
51

52

53
def normalize_condition(cond_image, target_size=(512, 512)):
54
    """
55
    标准化条件图像到 [0, 1] 范围，并 resize 到目标尺寸。
56
    """
57
    import torchvision.transforms as T
58

59
    transform = T.Compose([
60
        T.ToTensor(),
61
        T.Resize(target_size),
62
        T.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
63
    ])
64

65
    return transform(cond_image)

展开收起

4. 八种控制类型#

4.1 控制类型总览#

ControlNet 官方支持八种条件控制：

1
def control_types():
2
    """
3
    ControlNet 支持的八种控制类型。
4
    """
5
    return {
6
        "Canny Edge": {
7
            "输入": "任意图像",
8
            "预处理": "Canny 边缘检测",
9
            "效果": "精确的边缘轮廓控制",
10
            "适用": "草图到图像、线稿上色",
11
            "预训练模型大小": "~1.4GB",
12
        },
13
        "Depth": {
14
            "输入": "任意图像",
15
            "预处理": "单目深度估计 (MiDaS)",
16
            "效果": "空间深度结构控制",
17
            "适用": "深度感知生成、3D 场景构建",
18
            "预训练模型大小": "~1.4GB",
19
        },
20
        "Normal Map": {
21
            "输入": "任意图像",
22
            "预处理": "表面法线估计",
23
            "效果": "物体表面几何控制",
24
            "适用": "纹理映射、几何细节",
25
            "预训练模型大小": "~1.4GB",
26
        },
27
        "OpenPose": {
28
            "输入": "人物图像",
29
            "预处理": "人体姿态估计 (17 关键点)",
30
            "效果": "人物姿态控制",
31
            "适用": "角色动作生成、舞蹈/运动",
32
            "预训练模型大小": "~1.4GB",
33
        },
34
        "Semantic Segmentation": {
35
            "输入": "任意图像",
36
            "预处理": "语义分割 (ADE20K)",
37
            "效果": "物体类别和位置控制",
38
            "适用": "场景布局、物体放置",
39
            "预训练模型大小": "~1.4GB",
40
        },
41
        "M-LSD": {
42
            "输入": "任意图像",
43
            "预处理": "直线检测",
44
            "效果": "建筑/室内线条结构控制",
45
            "适用": "建筑草图、室内设计",
46
            "预训练模型大小": "~1.4GB",
47
        },
48
        "HED Soft Edge": {
49
            "输入": "任意图像",
50
            "预处理": "HED 边缘检测 (软边缘)",
51
            "效果": "柔和的边缘轮廓，比 Canny 更自然",
52
            "适用": "艺术风格边缘、风格迁移",
53
            "预训练模型大小": "~1.4GB",
54
        },
55
        "Scribble": {
56
            "输入": "手绘涂鸦",
57
            "预处理": "涂鸦转二值图",
58
            "效果": "手绘草图控制",
59
            "适用": "从涂鸦生成完整图像",
60
            "预训练模型大小": "~1.4GB",
61
        },
62
    }

展开收起

4.2 Canny Edge 控制详解#

Canny 是最常用的控制类型之一，适用于从线稿生成图像：

1
def canny_control_workflow():
2
    """
3
    Canny Edge 控制的工作流程。
4
    """
5
    return {
6
        "输入图像": "任意照片或生成图",
7
        "↓": "",
8
        "Canny 边缘检测": "低阈值 100, 高阈值 200",
9
        "↓": "得到精确的边缘轮廓图 (512×512×3)",
10
        "ControlNet (Canny)": "学习"边缘 → 图像内容"的映射",
11
        "↓": "输出条件噪声预测",
12
        "SD UNet + CFG": "结合文本条件生成图像",
13
        "↓": "",
14
        "输出图像": "保留原始边缘结构的生成图",
15
    }
16

17

18
def canny_training_data():
19
    """
20
    Canny ControlNet 的训练数据构建。
21
    """
22
    return {
23
        "数据来源": "LAION-5B / COCO / ADE20K",
24
        "样本数": "约 300 万对",
25
        "条件图像": "原始图像经过 Canny 边缘检测",
26
        "目标图像": "原始图像本身",
27
        "文本": "图像的描述 (BLIP-2 生成)",
28
        "训练技巧": "使用 ImageNet 预训练 Canny 检测器",
29
    }

展开收起

4.3 OpenPose 姿态控制详解#

OpenPose 是 ControlNet 在人物生成中最受欢迎的控制类型：

1
OpenPose 骨架:
2
  ┌─────────────────────────────────────────────────┐
3
  │                                                 │
4
  │   [0] 鼻子                                      │
5
  │        │                                        │
6
  │   [1]─[2]─[3]─[4]─[15]─[14]─[16]─[17]        │
7
  │        │                                        │
8
  │   [5]  [6]  [7]  [8]                           │
9
  │        │                                        │
10
  │   [9]  [10] [11]                                │
11
  │        │                                        │
12
  │   [12] [13]                                     │
13
  │                                                 │
14
  │  17 个关键点:                                    │
15
  │  0: 鼻子   1-2: 眼睛   3-4: 耳朵               │
16
  │  5-8: 肩膀、手肘、手腕 (右)                     │
17
  │  9-12: 髋、膝盖、脚踝 (右)                      │
18
  │  13-16: 肩膀、手肘、手腕 (左)                   │
19
  │  17-20: 髋、膝盖、脚踝 (左)                     │
20
  │                                                 │
21
  └─────────────────────────────────────────────────┘
22

23
关键点连线 (骨骼):
24
  头部: 0-1, 0-2, 1-3, 2-4     (脸)
25
  躯干: 5-6, 6-7, 7-8          (右臂)
26
  躯干: 5-11, 11-12, 12-13     (右腿)
27
  躯干: 5-9, 9-10, 10-11       (左臂)
28
  躯干: 5-14, 14-15, 15-16     (左腿)

展开收起

1
class OpenPoseVisualizer:
2
    """
3
    OpenPose 关键点可视化。
4
    """
5
    def __init__(self, keypoint_confidence=0.3):
6
        self.confidence = keypoint_confidence
7

8
        # 骨骼连接
9
        self.pairs = [
10
            (0, 1), (0, 2), (1, 3), (2, 4),  # 脸
11
            (5, 6), (6, 7), (7, 8),           # 右臂
12
            (5, 9), (9, 10), (10, 11),        # 左臂
13
            (5, 11), (11, 12), (12, 13),      # 右腿
14
            (5, 14), (14, 15), (15, 16),      # 左腿
15
        ]
16

17
        # 颜色
18
        self.keypoint_color = (255, 0, 85)    # 红色
19
        self.limb_color = (255, 85, 0)        # 橙色
20

21
    def draw_pose(self, keypoints, image_size):
22
        """
23
        在空白图上绘制姿态骨架。
24

25
        参数:
26
            keypoints: (17, 3) [x, y, confidence] 关键点
27
            image_size: (H, W) 输出图像大小
28
        """
29
        canvas = np.zeros((image_size[0], image_size[1], 3), dtype=np.uint8)
30

31
        # 绘制关键点
32
        for kp in keypoints:
33
            x, y, conf = kp
34
            if conf > self.confidence:
35
                cv2.circle(canvas, (int(x), int(y)), 4, self.keypoint_color, -1)
36

37
        # 绘制骨骼连线
38
        for i, j in self.pairs:
39
            if keypoints[i][2] > self.confidence and keypoints[j][2] > self.confidence:
40
                x1, y1 = int(keypoints[i][0]), int(keypoints[i][1])
41
                x2, y2 = int(keypoints[j][0]), int(keypoints[j][1])
42
                cv2.line(canvas, (x1, y1), (x2, y2), self.limb_color, 2)
43

44
        return canvas

展开收起

4.4 多控制组合#

1
def multi_control_inference(controlnet_models, sd_unet, vae, text_encoder,
2
                             z_t, t, conditions, texts, guidance_scale=7.5):
3
    """
4
    多控制组合推理。
5
    多个 ControlNet 可以组合使用。
6
    """
7
    noise_preds = []
8

9
    for ctrl_model, cond_img in zip(controlnet_models, conditions):
10
        # 每个 ControlNet 预测
11
        ctrl_pred = ctrl_model(z_t, t, cond_img, text_encoder.encode(texts))
12
        noise_preds.append(ctrl_pred)
13

14
    # 平均所有 ControlNet 的预测
15
    noise_pred_cond = torch.stack(noise_preds).mean(dim=0)
16

17
    # 无条件预测
18
    noise_pred_uncond = sd_unet(z_t, t, text_encoder.encode([""]))
19

20
    # CFG 合并
21
    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
22

23
    return noise_pred

展开收起

5. 推理：精确控制生成#

5.1 ControlNet 推理流程#

1
@torch.no_grad()
2
def controlnet_inference(controlnet, sd_unet, vae, text_encoder,
3
                          prompt, cond_image, num_steps=50, guidance_scale=7.5):
4
    """
5
    ControlNet 推理。
6

7
    参数:
8
        prompt:       str 文本描述
9
        cond_image:   PIL.Image 条件图像（边缘/深度/姿态等）
10
        num_steps:    推理步数 (DDIM)
11
        guidance_scale: CFG 强度
12
    """
13
    device = next(controlnet.parameters()).device
14

15
    # ============ 1) 预处理条件图像 ============
16
    cond_tensor = preprocess_condition(cond_image)  # (1, 3, 512, 512)
17
    cond_tensor = cond_tensor.to(device)
18

19
    # ============ 2) 编码文本 ============
20
    context = text_encoder.encode([prompt])     # 条件文本
21
    uncond_context = text_encoder.encode([""])   # 无条件文本
22

23
    # ============ 3) 初始化噪声 ============
24
    z_T = torch.randn(1, 4, 64, 64, device=device)
25

26
    # ============ 4) DDIM 采样 ============
27
    timesteps = get_ddim_timesteps(num_steps)  # 递减时间步
28

29
    for i, t in enumerate(tqdm(timesteps)):
30
        t_tensor = torch.tensor([t], device=device)
31

32
        # 条件预测 (ControlNet)
33
        ctrl_pred = controlnet(z_T, t_tensor, cond_tensor, context)
34

35
        # 无条件预测 (SD)
36
        uncond_pred = sd_unet(z_T, t_tensor, uncond_context)
37

38
        # CFG 合并
39
        noise_pred = uncond_pred + guidance_scale * (ctrl_pred - uncond_pred)
40

41
        # DDIM 去噪步
42
        z_T = ddim_step(z_T, noise_pred, t, timesteps[i - 1] if i > 0 else 0)
43

44
    # ============ 5) VAE 解码 ============
45
    image = vae.decode(z_T)
46

47
    return image

展开收起

5.2 ControlNet 的控制强度#

1
def control_strength():
2
    """
3
    控制强度参数 (control_scale) 的效果。
4
    """
5
    results = {
6
        "0.0 (无控制)": {
7
            "效果": "完全由文本决定，忽略条件图像",
8
            "生成": "与普通 SD 无异",
9
        },
10
        "0.3-0.5 (弱控制)": {
11
            "效果": "条件提供参考，但允许灵活变化",
12
            "生成": "保留条件的部分结构，加入创意",
13
        },
14
        "0.7-0.9 (标准控制)": {
15
            "效果": "条件主导结构，细节和风格由文本决定",
16
            "生成": "精确的结构 + 丰富的细节",
17
        },
18
        "1.0 (强控制)": {
19
            "效果": "条件主导一切，文本仅影响颜色/风格",
20
            "生成": "结构与条件高度一致",
21
        },
22
        "1.2+ (过强)": {
23
            "效果": "可能产生 artifacts 或过度平滑",
24
            "生成": "结构过拟合，缺乏多样性",
25
        },
26
    }
27
    return results
28

29

30
def set_control_scale(controlnet, scale=0.8):
31
    """
32
    设置 ControlNet 的控制强度。
33
    推理时可以动态调整。
34
    """
35
    for module in controlnet.modules():
36
        if isinstance(module, ZeroConv2d):
37
            # ZeroConv 权重乘以 scale
38
            with torch.no_grad():
39
                module.conv.weight.mul_(scale)
40
    return controlnet

展开收起

6. ControlNet-XS：效率优化#

6.1 ControlNet 的效率问题#

原始 ControlNet 的问题是：需要完整复制 UNet 编码器，导致参数量大（与 SD UNet 相当）：

1
原始 ControlNet:
2
  - SD UNet: ~860M 参数
3
  - ControlNet 副本: ~860M 参数 (相同大小！)
4
  - 总计: ~1.7B 参数
5
  - 问题: 部署成本高，推理慢

6.2 ControlNet-XS 的解法#

ControlNet-XS（ControlNet-XS: Lightweight ControlNet for Faster and Better Image Editing, Patil et al., 2024）提出了两个关键改进：

1
改进 1: 只复制编码器的后半部分
2
  - 原始 ControlNet: 复制全部 12 层编码器
3
  - ControlNet-XS:    只复制后 4 层 (下采样最深的部分)
4
  - 参数量: ~860M → ~300M
5

6
  原理:
7
  - 高层语义特征 (深层) 才需要条件信息
8
  - 低层纹理特征 (浅层) 对条件控制贡献小
9

10
改进 2: 更轻量的条件注入
11
  - 不再完整复制编码器结构
12
  - 使用浅层的旁路连接
13
  - 通过 tiny 的调节器（modulator）注入条件

1
class ControlNetXS(nn.Module):
2
    """
3
    ControlNet-XS: 轻量级 ControlNet。
4
    只复制编码器的后半部分。
5
    """
6

7
    def __init__(self, sd_unet, num_layers_to_copy=4):
8
        super().__init__()
9
        self.sd_unet = sd_unet
10
        self.num_layers = num_layers_to_copy  # 只需复制 4 层而非 12 层
11

12
        # 只复制编码器的深层部分
13
        self.control_blocks = nn.ModuleList([
14
            self._make_control_block(ch)
15
            for ch in sd_unet.channel_list[-num_layers_to_copy:]
16
        ])
17

18
        # 轻量级条件编码器
19
        self.cond_encoder = nn.Sequential(
20
            nn.Conv2d(3, 32, 3, padding=1),
21
            nn.SiLU(),
22
            nn.Conv2d(32, 64, 3, stride=2, padding=1),
23
            nn.SiLU(),
24
            nn.Conv2d(64, 64, 3, padding=1),
25
        )
26

27
        # 轻量级跳跃连接
28
        self.zero_convs = nn.ModuleList([
29
            ZeroConv2d(sd_unet.channel_list[-num_layers_to_copy + i],
30
                       sd_unet.channel_list[-num_layers_to_copy + i])
31
            for i in range(num_layers_to_copy)
32
        ])
33

34
    def forward(self, x_t, t, cond_image, context):
35
        # 条件编码
36
        cond = self.cond_encoder(cond_image)
37

38
        # 只在深层注入控制
39
        features = self.sd_unet.get_deep_features(x_t, t)  # 只取后 N 层
40

41
        controlled_skips = []
42
        for i, (feat, ctrl_block, zero_conv) in enumerate(zip(
43
            features, self.control_blocks, self.zero_convs
44
        )):
45
            ctrl_feat = ctrl_block(feat + cond)
46
            controlled_skips.append(zero_conv(ctrl_feat))
47

48
        # UNet 解码（与原始相同）
49
        output = self.sd_unet.decode_with_skip(x_t, t, context, controlled_skips)
50
        return output

展开收起

6.3 ControlNet-XS vs 原始 ControlNet#

1
def controlnet_comparison():
2
    """
3
    ControlNet vs ControlNet-XS 对比。
4
    """
5
    return {
6
        "参数量": {
7
            "ControlNet": "~860M (完整复制)",
8
            "ControlNet-XS": "~300M (部分复制)",
9
            "节省": "65%",
10
        },
11
        "推理速度": {
12
            "ControlNet": "1x (慢)",
13
            "ControlNet-XS": "~2x (快)",
14
        },
15
        "生成质量": {
16
            "ControlNet": "最佳",
17
            "ControlNet-XS": "略有下降，但在可接受范围",
18
            "差距": "FID 差值 < 2%",
19
        },
20
        "适用场景": {
21
            "ControlNet": "追求最高质量",
22
            "ControlNet-XS": "边缘部署、实时应用",
23
        },
24
    }

展开收起

7. ControlLoRA：更轻量的方案#

7.1 ControlLoRA 的核心思路#

ControlLoRA（2024）用 LoRA 低秩适配 替代完整的 ControlNet 副本：

1
原始 ControlNet:
2
  - 完整复制 UNet 编码器: ~860M 参数
3
  - 需要完整微调
4

5
ControlLoRA:
6
  - 只微调编码器的 Low-rank 分解矩阵: ~5-20M 参数
7
  - 冻结原始权重，通过低秩矩阵注入控制
8
  - 训练更高效，部署更轻量

1
class ControlLoRA(nn.Module):
2
    """
3
    ControlLoRA: 用 LoRA 替代完整复制。
4
    """
5
    def __init__(self, sd_unet, rank=16, alpha=16):
6
        super().__init__()
7
        self.sd_unet = sd_unet
8
        self.rank = rank
9
        self.alpha = alpha
10

11
        # 为每层添加 LoRA
12
        self.lora_layers = nn.ModuleDict()
13
        for name, param in sd_unet.named_parameters():
14
            if "conv" in name and "weight" in name:
15
                # 对卷积层添加 LoRA
16
                self.lora_layers[name] = LoRALinear(param, rank, alpha)
17

18
        # 冻结原始权重
19
        for param in sd_unet.parameters():
20
            param.requires_grad_(False)
21

22
        # 条件处理
23
        self.cond_encoder = CondEncoder(rank)
24

25
    def forward(self, x_t, t, cond_image, context):
26
        # 处理条件
27
        cond = self.cond_encoder(cond_image)  # 低秩条件表示
28

29
        # 用 LoRA 注入条件
30
        for name, lora_layer in self.lora_layers.items():
31
            original = self.sd_unet.state_dict()[name]
32
            lora_weight = lora_layer(cond)  # 基于条件的低秩调整
33
            self.sd_unet.set_param(name, original + lora_weight)
34

35
        # 前向
36
        output = self.sd_unet(x_t, t, context)
37

38
        # 恢复原始权重
39
        for name in self.lora_layers:
40
            original = self.sd_unet.state_dict()[name]
41
            self.sd_unet.set_param(name, original)
42

43
        return output

展开收起

7.2 ControlNet vs ControlNet-XS vs ControlLoRA#

1
def control_variants():
2
    """
3
    三种控制方案对比。
4
    """
5
    return {
6
        "ControlNet": {
7
            "参数量": "~860M",
8
            "训练方式": "全量微调",
9
            "质量": "最高",
10
            "速度": "慢",
11
            "适用": "追求最佳效果",
12
        },
13
        "ControlNet-XS": {
14
            "参数量": "~300M",
15
            "训练方式": "部分层微调",
16
            "质量": "略低 (~2% FID差距)",
17
            "速度": "中等",
18
            "适用": "效率敏感的部署",
19
        },
20
        "ControlLoRA": {
21
            "参数量": "~10-20M",
22
            "训练方式": "LoRA 微调",
23
            "质量": "低于 XS",
24
            "速度": "快",
25
            "适用": "轻量部署、个性化控制",
26
        },
27
    }

展开收起

8. T2I-Adapter：另一种条件控制方案#

8.1 T2I-Adapter 的设计#

T2I-Adapter（Mou et al., 2024）与 ControlNet 不同：不是复制 UNet，而是训练一个轻量适配器：

1
ControlNet:
2
  - 完整复制 UNet 编码器
3
  - 每个控制类型需要一个完整的 ControlNet
4
  - 参数量大
5

6
T2I-Adapter:
7
  - 不复制 UNet
8
  - 训练一个轻量适配器 (约 80M 参数)
9
  - 通过特征注入影响 UNet
10
  - 可共享基础 UNet，多个 Adapter 组合使用

1
class T2IAdapter(nn.Module):
2
    """
3
    T2I-Adapter: 轻量级条件适配器。
4
    """
5
    def __init__(self, channels=[320, 640, 1280, 1280], num_blocks=4):
6
        super().__init__()
7

8
        # 条件图像编码器
9
        self.cond_encoder = nn.Sequential(
10
            nn.Conv2d(3, 64, 3, padding=1),
11
            *[ResBlock(64, 64) for _ in range(num_blocks)],
12
            nn.Conv2d(64, channels[-1], 1),
13
        )
14

15
        # 适配块: 与 UNet 各层对齐
16
        self.adapters = nn.ModuleList([
17
            AdapterBlock(in_ch, out_ch)
18
            for in_ch, out_ch in zip(channels[:-1], channels[1:])
19
        ])
20

21
    def forward(self, cond_image, unet_features):
22
        """
23
        注入条件信息到 UNet 特征。
24

25
        参数:
26
            cond_image:     (B, 3, H, W) 条件图像
27
            unet_features:   List[Tensor] UNet 各层特征
28
        """
29
        # 编码条件
30
        cond = self.cond_encoder(cond_image)
31

32
        # 逐层适配
33
        adapted_features = []
34
        for feat, adapter in zip(unet_features, self.adapters):
35
            adapted = adapter(feat, cond)
36
            adapted_features.append(adapted)
37

38
        return adapted_features

展开收起

8.2 ControlNet vs T2I-Adapter#

1
def controlnet_vs_adapter():
2
    """
3
    ControlNet vs T2I-Adapter 对比。
4
    """
5
    return {
6
        "架构": {
7
            "ControlNet": "完整复制 UNet 编码器",
8
            "T2I-Adapter": "独立轻量适配器",
9
        },
10
        "参数量": {
11
            "ControlNet": "~860M",
12
            "T2I-Adapter": "~80M",
13
        },
14
        "可组合性": {
15
            "ControlNet": "每个控制类型独立模型",
16
            "T2I-Adapter": "多个 Adapter 可叠加",
17
        },
18
        "质量": {
19
            "ControlNet": "更高（更多参数）",
20
            "T2I-Adapter": "略低，但足够好",
21
        },
22
        "灵活性": {
23
            "ControlNet": "固定与 SD 配合",
24
            "T2I-Adapter": "更灵活的注入点",
25
        },
26
    }

展开收起

9. 完整实现：简化版 ControlNet#

1
import torch
2
import torch.nn as nn
3
import torch.nn.functional as F
4
from dataclasses import dataclass
5
import math
6

7

8
# ============ 基础模块 ============
9
class ResBlock(nn.Module):
10
    def __init__(self, channels):
11
        super().__init__()
12
        self.norm1 = nn.GroupNorm(32, channels)
13
        self.conv1 = nn.Conv2d(channels, channels, 3, padding=1)
14
        self.norm2 = nn.GroupNorm(32, channels)
15
        self.conv2 = nn.Conv2d(channels, channels, 3, padding=1)
16

17
    def forward(self, x):
18
        h = F.silu(self.norm1(x))
19
        h = self.conv1(h)
20
        h = F.silu(self.norm2(h))
21
        h = self.conv2(h)
22
        return h + x
23

24

25
class ZeroConv2d(nn.Module):
26
    """零卷积: 初始为零，渐进增长。"""
27
    def __init__(self, in_ch, out_ch):
28
        super().__init__()
29
        self.conv = nn.Conv2d(in_ch, out_ch, 1)
30
        nn.init.constant_(self.conv.weight, 0.0)
31
        nn.init.constant_(self.conv.bias, 0.0)
32

33
    def forward(self, x):
34
        return self.conv(x)
35

36

37
# ============ 条件编码器 ============
38
class ConditionEncoder(nn.Module):
39
    """条件图像编码器。"""
40
    def __init__(self, in_ch=3, mid_ch=64, out_ch=320):
41
        super().__init__()
42
        self.encoder = nn.Sequential(
43
            nn.Conv2d(in_ch, mid_ch, 3, padding=1),
44
            nn.SiLU(),
45
            nn.Conv2d(mid_ch, mid_ch, 3, stride=2, padding=1),  # 256
46
            nn.SiLU(),
47
            nn.Conv2d(mid_ch, mid_ch * 2, 3, padding=1),
48
            nn.SiLU(),
49
            nn.Conv2d(mid_ch * 2, mid_ch * 2, 3, stride=2, padding=1),  # 128
50
            nn.SiLU(),
51
            nn.Conv2d(mid_ch * 2, out_ch, 3, padding=1),
52
        )
53

54
    def forward(self, x):
55
        return self.encoder(x)
56

57

58
# ============ ControlNet Block ============
59
class ControlNetBlock(nn.Module):
60
    """ControlNet 基本模块。"""
61
    def __init__(self, in_ch, out_ch):
62
        super().__init__()
63
        self.res1 = ResBlock(in_ch)
64
        self.res2 = ResBlock(out_ch)
65
        self.zero_conv = ZeroConv2d(out_ch, out_ch)
66
        if in_ch != out_ch:
67
            self.shortcut = nn.Conv2d(in_ch, out_ch, 1)
68
        else:
69
            self.shortcut = nn.Identity()
70

71
    def forward(self, x, cond):
72
        h = self.res1(x)
73
        h = h + cond  # 条件注入
74
        h = self.res2(h)
75
        return self.zero_conv(h)
76

77

78
# ============ ControlNet ============
79
class ControlNet(nn.Module):
80
    """
81
    简化版 ControlNet。
82
    """
83

84
    def __init__(self, sd_unet):
85
        super().__init__()
86
        self.sd_unet = sd_unet
87
        self.channels = sd_unet.channel_list  # [320, 640, 1280, 1280]
88

89
        # 冻结 SD UNet
90
        for param in sd_unet.parameters():
91
            param.requires_grad_(False)
92

93
        # 条件编码器
94
        self.cond_encoder = ConditionEncoder(out_ch=self.channels[0])
95

96
        # ControlNet Blocks（与 UNet 编码器结构对应）
97
        self.control_blocks = nn.ModuleList([
98
            nn.ModuleList([
99
                ControlNetBlock(self.channels[i], self.channels[i + 1])
100
                for _ in range(2)  # 每层 2 个 block
101
            ])
102
            for i in range(len(self.channels) - 1)
103
        ])
104

105
        # 下采样
106
        self.downsample = nn.ModuleList([
107
            nn.Conv2d(self.channels[i], self.channels[i], 3, stride=2, padding=1)
108
            for i in range(len(self.channels) - 1)
109
        ])
110

111
    def forward(self, x_t, t, cond_image, context):
112
        """
113
        参数:
114
            x_t:       (B, 4, 64, 64) 噪声潜变量
115
            t:         (B,) 时间步
116
            cond_image: (B, 3, 512, 512) 条件图像
117
            context:   (B, seq_len, d) 文本 embedding
118
        """
119
        # 编码条件
120
        cond = self.cond_encoder(cond_image)  # (B, 320, 64, 64)
121

122
        # 获取 SD UNet 的特征
123
        unet_features = self.sd_unet.get_encoder_features(x_t, t)
124

125
        # 通过 ControlNet 处理并返回跳跃连接
126
        skips = []
127
        h = x_t
128
        for i, (feat, blocks, down) in enumerate(zip(
129
            unet_features, self.control_blocks, self.downsample
130
        )):
131
            for block in blocks:
132
                feat = block(feat, cond if i == 0 else torch.zeros_like(cond))
133
            skips.append(feat)
134
            if i < len(self.downsample) - 1:
135
                h = down(h)
136

137
        # SD UNet 解码
138
        output = self.sd_unet.decode_with_skips(x_t, t, context, skips)
139
        return output
140

141

142
# ============ SD UNet (简化版) ============
143
class SDUNet(nn.Module):
144
    """简化版 Stable Diffusion UNet。"""
145
    def __init__(self):
146
        super().__init__()
147
        self.channel_list = [320, 640, 1280, 1280]
148

149
        # 简化实现
150
        self.encoder_blocks = nn.ModuleList([
151
            nn.ModuleList([ResBlock(ch) for _ in range(2)])
152
            for ch in self.channel_list
153
        ])
154
        self.downsample = nn.ModuleList([
155
            nn.Conv2d(ch, ch, 3, stride=2, padding=1)
156
            for ch in self.channel_list[:-1]
157
        ])
158

159
        self.bottleneck = nn.ModuleList([ResBlock(1280) for _ in range(2)])
160

161
        self.decoder_blocks = nn.ModuleList([
162
            nn.ModuleList([ResBlock(ch) for _ in range(2)])
163
            for ch in reversed(self.channel_list)
164
        ])
165
        self.upsample = nn.ModuleList([
166
            nn.ConvTranspose2d(ch, ch, 4, stride=2, padding=1)
167
            for ch in reversed(self.channel_list[:-1])
168
        ])
169

170
        self.out_conv = nn.Conv2d(320, 4, 3, padding=1)
171
        self.time_embed = TimeEmbedding(320)
172

173
    def get_encoder_features(self, x, t):
174
        """返回编码器各层特征（用于 ControlNet）。"""
175
        features = []
176
        t_emb = self.time_embed(t)
177
        h = x
178

179
        for i, (blocks, down) in enumerate(zip(self.encoder_blocks, self.downsample)):
180
            for block in blocks:
181
                h = block(h)
182
            features.append(h)
183
            h = down(h)
184

185
        return features
186

187
    def decode_with_skips(self, x, t, context, skips):
188
        """用 ControlNet 的跳跃连接解码。"""
189
        t_emb = self.time_embed(t)
190
        h = x
191

192
        for i, (blocks, up) in enumerate(zip(
193
            self.decoder_blocks, self.upsample
194
        )):
195
            # 拼接跳跃连接（ControlNet 提供）
196
            if i < len(skips):
197
                h = torch.cat([h, skips[-(i + 1)]], dim=1)
198

199
            for block in blocks:
200
                h = block(h)
201

202
            h = up(h)
203

204
        return self.out_conv(h)
205

206
    def forward(self, x, t, context):
207
        features = self.get_encoder_features(x, t)
208
        return self.decode_with_skips(x, t, context, [torch.zeros_like(x)] * 4)
209

210

211
class TimeEmbedding(nn.Module):
212
    def __init__(self, dim):
213
        super().__init__()
214
        self.mlp = nn.Sequential(
215
            nn.Linear(128, dim * 4),
216
            nn.SiLU(),
217
            nn.Linear(dim * 4, dim),
218
        )
219

220
    def forward(self, t):
221
        half = 64
222
        emb = math.log(10000) / half
223
        emb = torch.exp(torch.arange(half, device=t.device) * -emb)
224
        emb = t[:, None] * emb[None, :]
225
        emb = torch.cat([emb.sin(), emb.cos()], dim=-1)
226
        return self.mlp(emb)

展开收起

10. 应用场景#

10.1 ControlNet 的典型应用#

1
def controlnet_applications():
2
    """
3
    ControlNet 的典型应用场景。
4
    """
5
    return {
6
        "线稿上色": {
7
            "输入": "黑白线稿",
8
            "条件": "Canny / HED 边缘检测",
9
            "输出": "保留线稿结构的彩色图像",
10
            "模型": "ControlNet Canny / HED",
11
        },
12
        "姿态控制": {
13
            "输入": "姿态骨架图",
14
            "条件": "OpenPose 关键点",
15
            "输出": "符合姿态的人物图像",
16
            "模型": "ControlNet OpenPose",
17
        },
18
        "深度感知生成": {
19
            "输入": "深度图",
20
            "条件": "MiDaS 深度估计",
21
            "输出": "保持深度结构的三维图像",
22
            "模型": "ControlNet Depth",
23
        },
24
        "建筑/室内设计": {
25
            "输入": "建筑线稿",
26
            "条件": "M-LSD 直线检测",
27
            "输出": "保留建筑结构的渲染图",
28
            "模型": "ControlNet M-LSD",
29
        },
30
        "分割引导生成": {
31
            "输入": "语义分割图",
32
            "条件": "ADE20K 分割",
33
            "输出": "符合分割布局的图像",
34
            "模型": "ControlNet Seg",
35
        },
36
        "风格迁移": {
37
            "输入": "参考图像",
38
            "条件": "法线图 / 深度图",
39
            "输出": "保持结构、改变风格",
40
            "模型": "ControlNet Normal / Depth",
41
        },
42
    }

展开收起

10.2 组合多个 ControlNet#

1
def combine_multiple_controlnets():
2
    """
3
    组合多个 ControlNet 实现更精细的控制。
4
    """
5
    return {
6
        "场景": "生成一个坐在室内的女性",
7
        "控制1": "OpenPose (姿态) + ControlNet OpenPose",
8
        "控制2": "深度图 (空间) + ControlNet Depth",
9
        "控制3": "语义分割 (布局) + ControlNet Seg",
10
        "文本": "a beautiful woman sitting in a modern living room",
11
        "方法": "多个 ControlNet 的噪声预测求平均",
12
        "效果": "精确的姿态 + 空间结构 + 场景布局",
13
    }

11. 总结#

11.1 核心要点#

维度	关键要点
核心问题	扩散模型缺乏空间/结构级可控性
解决方案	复制 UNet 编码器为可训练副本
核心机制	零卷积初始化保证训练稳定，不破坏预训练权重
训练策略	先冻结零卷积训练 ControlNet Block，再解冻零卷积
八种控制	Canny/深度/姿态/分割/M-LSD/HED/Scribble/Normal
效率优化	ControlNet-XS（部分复制）/ ControlLoRA（LoRA 替代）
开源影响	成为 SD 最重要的控制扩展，催生 ComfyUI 等生态

11.2 ControlNet vs 其他条件控制方案#

1
         ControlNet    T2I-Adapter   分类器引导   额外通道
2
架构     复制编码器    轻量适配器     训练分类器    修改输入
3
参数量   ~860M         ~80M          ~100M        ~86M
4
训练难度 低（零卷积）   中等          高           高
5
控制精度 高            中            中           中
6
兼容性   最佳          好            差           差
7
开源生态 最活跃        一般          无           一般

11.3 一句话总结#

ControlNet 的核心创新是”复制 + 冻结 + 零卷积”——把预训练的 SD UNet 编码器完整复制一份作为可训练副本，冻结原始权重保证预训练知识不丢失，通过零卷积初始化的跳跃连接把控制信号从零开始渐进注入解码器。这使得任意条件（边缘、姿态、深度、分割等）都能以插件形式精确控制图像生成结构，同时完全保留 Stable Diffusion 的文本生成质量——成为 AI 图像生成从”创意工具”升级为”可控设计工具”的关键技术。

11.4 推荐资源#

1
论文:
2
  - ControlNet (Zhang & Agrawala, 2023): "Adding Conditional Control to Text-to-Image Diffusion Models"
3
  - ControlNet-XS (Patil et al., 2024): "ControlNet-XS: Lightweight ControlNet for Faster and Better Image Editing"
4
  - ControlLoRA (2024): LoRA-based ControlNet
5
  - T2I-Adapter (Mou et al., 2024): "T2I-Adapter: Adapting Diffusion Models for Text-to-Image Generation"
6

7
代码:
8
  - lllyasviel/ControlNet (官方实现)
9
  - lllyasviel/sd-webui-controlnet (AUTOMATIC1111 WebUI 插件)
10
  - comfyanonymous/ComfyUI (节点式 ControlNet 工作流)
11

12
预训练模型:
13
  - HuggingFace: lllyasviel/sd-controlnet-canny
14
  - HuggingFace: lllyasviel/sd-controlnet-depth
15
  - HuggingFace: lllyasviel/sd-controlnet-openpose
16
  - ModelScope: damo/control_v11p_sd15_canny

展开收起

数据构建深度解析：LLM 训练数据的艺术与科学

Thu, 16 Jul 2026 00:00:00 GMT

1. 引言：数据是 LLM 的灵魂#

1.1 数据决定模型上限#

在 LLM 领域，有一个被反复验证的铁律：

“Garbage in, garbage out.” — 数据的质量决定了模型能力的上限，训练只是无限逼近这个上限。

1
数据 vs 模型能力的相关性：
2

3
  GPT-3 (300B tokens) → GPT-4 (?)
4
    → 能力跃升的核心原因之一是数据质量的提升
5
    → 不仅仅是参数规模的扩大
6

7
  LLaMA 2 (2T tokens) vs LLaMA 3 (15T tokens)
8
    → LLaMA 3 用更多但更高质量的数据
9
    → 在更少参数下超越了 LLaMA 2
10

11
  Phi-3 (3.8B 参数)
12
    → 仅用 3.3T tokens 训练
13
    → 超越大它 10 倍的模型
14
    → 秘诀：高质量"教科书级"数据

1.2 LLM 数据的三大类型#

LLM 训练涉及三类核心数据，它们的角色和构建方法截然不同：

1
┌─────────────────────────────────────────────────────────────────┐
2
│                    LLM 训练数据的三种类型                         │
3
├─────────────────────────────────────────────────────────────────┤
4
│                                                                  │
5
│  类型 1: 预训练数据 (Pre-training Data)                         │
6
│    用途: 学习语言模型的基础能力（语法、语义、世界知识）           │
7
│    规模: trillions of tokens（数万亿）                          │
8
│    来源: 网页、书籍、代码、论文、百科等                          │
9
│    构建: 爬取 → 去重 → 质量过滤 → 格式转换                      │
10
│                                                                  │
11
│  类型 2: SFT 数据 (Supervised Fine-Tuning Data)                 │
12
│    用途: 学习指令遵循、对话格式、特定任务能力                     │
13
│    规模: thousands to millions（数千到数百万）                  │
14
│    来源: 人类标注、LLM 生成、公开数据集                           │
15
│    构建: 指令设计 → 回答生成 → 质量筛选                         │
16
│                                                                  │
17
│  类型 3: 偏好数据 (Preference Data)                              │
18
│    用途: 学习人类偏好（RLHF/DPO/KTO 对齐）                       │
19
│    规模: thousands to hundreds of thousands（数千到数十万）       │
20
│    来源: 人类标注、LLM 评判                                     │
21
│    构建: 候选生成 → 偏好标注 → 质量过滤                         │
22
│                                                                  │
23
└─────────────────────────────────────────────────────────────────┘

展开收起

1.3 本系列文章关联#

文章	关联
后训练深度解析	数据在后训练流水线中的使用方式
偏好对齐深度解析	偏好数据的具体应用
知识蒸馏深度解析	蒸馏数据的特殊构建方法
DeepSeek-R1	R1 的数据构建实践

2. 预训练数据构建#

2.1 数据来源#

预训练数据的来源决定了模型的知识面和能力边界：

1
class PretrainingDataSources:
2
    """
3
    预训练数据来源及其特点
4
    """
5

6
    SOURCES = {
7
        # 网页数据：最大、最多样、但噪声最多
8
        "web": {
9
            "examples": ["CommonCrawl", "WebText", "CCNet", "C4"],
10
            "size": "~trillions of tokens",
11
            "pros": ["多样性高", "覆盖广", "成本相对低"],
12
            "cons": ["噪声多", "隐私问题", "质量参差不齐"],
13
        },
14

15
        # 书籍：高质量、连贯、但覆盖有限
16
        "books": {
17
            "examples": ["BookCorpus", "Books3", "Gutenberg", "Libgen"],
18
            "size": "~hundreds of billions",
19
            "pros": ["质量高", "长程依赖", "结构化"],
20
            "cons": ["版权问题", "覆盖有限", "可能过时"],
21
        },
22

23
        # 代码：GPT-4/Codex 能力的关键
24
        "code": {
25
            "examples": ["GitHub", "The Stack", "Python data"],
26
            "size": "~hundreds of billions",
27
            "pros": ["结构化", "逻辑性强", "可执行验证"],
28
            "cons": ["语言有限", "注释质量差", "版权问题"],
29
        },
30

31
        # 学术论文：知识和推理能力
32
        "papers": {
33
            "examples": ["arXiv", "Semantic Scholar"],
34
            "size": "~hundreds of billions",
35
            "pros": ["知识密集", "推理严谨", "引用丰富"],
36
            "cons": ["格式复杂", "数学公式多", "需要处理"],
37
        },
38

39
        # 百科：结构化知识
40
        "encyclopedia": {
41
            "examples": ["Wikipedia", "Wikidata"],
42
            "size": "~billions",
43
            "pros": ["知识准确", "结构化", "引用丰富"],
44
            "cons": ["覆盖面有限", "可能有偏见"],
45
        },
46
    }

展开收起

2.2 网页数据流水线：CCNet#

CommonCrawl 是预训练语料的最大来源，CCNet 是处理它的经典流水线：

1
class CCNetPipeline:
2
    """
3
    CCNet: CommonCrawl Net
4

5
    处理 CommonCrawl 网页数据的完整流水线
6
    论文: Wenzek et al., 2020
7
    """
8

9
    def __init__(self):
10
        self.language = "en"
11
        self.min_quality_score = 0.7
12
        self.min_length = 100
13

14
    def process_common_crawl(self, raw_warc_files):
15
        """
16
        完整流水线：
17
        1. 下载 CommonCrawl WARC 文件
18
        2. 提取文本（Text Extraction）
19
        3. 语言识别（Language Identification）
20
        4. 质量过滤（Quality Filtering）
21
        5. 去重（Deduplication）
22
        """
23
        all_documents = []
24

25
        for warc_file in tqdm(raw_warc_files):
26
            # Step 1: 下载并解压 WARC 文件
27
            warc_records = self.parse_warc(warc_file)
28

29
            # Step 2: 文本提取
30
            for record in warc_records:
31
                text = self.extract_text(record)
32

33
                # Step 3: 语言识别
34
                if self.is_target_language(text, self.language):
35
                    all_documents.append(text)
36

37
        # Step 4: 质量过滤
38
        filtered = self.quality_filter(all_documents)
39

40
        # Step 5: 去重
41
        deduped = self.deduplicate(filtered)
42

43
        return deduped
44

45
    def extract_text(self, warc_record):
46
        """
47
        从 HTML 中提取文本
48

49
        使用 jusText 或类似的启发式提取器
50
        移除导航栏、页脚、广告等噪音内容
51
        """
52
        html = warc_record.content
53

54
        # 方法 1: jusText
55
        # 基于马尔可夫链的文本提取
56
        # 区分"正文内容"和"模板内容"
57

58
        # 方法 2: trafilatura
59
        # 更准确的提取，保留表格和列表
60

61
        # 方法 3: boilerplate removal
62
        # 基于规则的噪音移除
63

64
        text = trafilatura.extract(html)
65

66
        return text
67

68
    def is_target_language(self, text, target_lang):
69
        """
70
        语言识别
71

72
        使用 fastText 的 LID (Language Identification) 模型
73
        在 Wikipedia 上预训练，支持 176 种语言
74
        """
75
        # 加载预训练的 fastText 模型
76
        lid_model = fasttext.load_model('lid.176.bin')
77

78
        # 预测语言和置信度
79
        predictions = lid_model.predict(text, k=1)
80

81
        lang = predictions[0][0].replace('__label__', '')
82
        confidence = predictions[1][0]
83

84
        return lang == target_lang and confidence > 0.8
85

86
    def quality_filter(self, documents):
87
        """
88
        质量过滤
89

90
        使用多个质量指标组合：
91
        1. 语言模型困惑度（LM Perplexity）
92
        2. 单词/字符比率
93
        3. 特殊字符密度
94
        4. 停用词比率
95
        """
96
        scored_docs = []
97

98
        for doc in documents:
99
            score = 0.0
100

101
            # 指标 1: 句子完整度
102
            # 好的文档应该有完整的句子
103
            sentences = sent_tokenize(doc)
104
            complete_ratio = sum(1 for s in sentences if s.endswith(('.','!','?'))) / len(sentences)
105
            score += 0.3 * complete_ratio
106

107
            # 指标 2: 词/字符比率
108
            # 合理范围：4-10 chars/word
109
            words = doc.split()
110
            if len(words) > 0:
111
                avg_word_len = len(doc) / len(words)
112
                if 3 < avg_word_len < 15:
113
                    score += 0.2
114
                else:
115
                    score += 0.1 * (10 - abs(avg_word_len - 6)) / 10
116

117
            # 指标 3: 停用词比率
118
            stopwords = set(stopwords_list)
119
            stopword_ratio = sum(1 for w in words if w.lower() in stopwords) / len(words)
120
            if 0.3 < stopword_ratio < 0.7:
121
                score += 0.2
122

123
            # 指标 4: 重复度
124
            # 高重复度可能是模板/垃圾内容
125
            repeat_ratio = self.compute_repeat_ratio(doc)
126
            if repeat_ratio < 0.2:
127
                score += 0.3 * (1 - repeat_ratio)
128

129
            if score >= self.min_quality_score:
130
                scored_docs.append((score, doc))
131

132
        # 按分数排序
133
        scored_docs.sort(reverse=True)
134

135
        return [doc for _, doc in scored_docs]
136

137
    def compute_repeat_ratio(self, text, n=5):
138
        """
139
        计算 n-gram 重复比率
140

141
        重复比率 = 1 - (唯一 n-gram 数 / 总 n-gram 数)
142
        """
143
        words = text.split()
144
        if len(words) < n:
145
            return 0.0
146

147
        ngrams = [tuple(words[i:i+n]) for i in range(len(words)-n+1)]
148
        unique_ngrams = len(set(ngrams))
149
        total_ngrams = len(ngrams)
150

151
        return 1 - unique_ngrams / total_ngrams

展开收起

2.3 RefineWeb: 高质量网页数据#

RefineWeb（Faoorosh et al., 2023）提出了更精细的过滤策略：

1
class RefineWebPipeline:
2
    """
3
    RefineWeb: Refined Web Dataset
4

5
    相比 CCNet 的改进：
6
    1. 更激进的模板过滤
7
    2. 使用更强的质量分类器
8
    3. 保留文档级去重
9
    """
10

11
    def __init__(self):
12
        self.quality_classifier = None  # 用高质量数据训练的分类器
13

14
    def build_quality_classifier(self, positive_samples, negative_samples):
15
        """
16
        训练质量分类器
17

18
        positive: Wikipedia, Books, 人工筛选的高质量网页
19
        negative: 广告、导航栏、低质量网页
20
        """
21
        from sklearn.linear_model import LogisticRegression
22

23
        def extract_features(text):
24
            """提取质量相关特征"""
25
            features = {}
26

27
            # 词汇丰富度
28
            words = text.split()
29
            unique_words = set(words)
30
            features['vocab_richness'] = len(unique_words) / (len(words) + 1)
31

32
            # 平均句子长度
33
            sentences = sent_tokenize(text)
34
            features['avg_sent_len'] = np.mean([len(s.split()) for s in sentences])
35

36
            # 标点符号密度
37
            punct_count = sum(1 for c in text if c in '.,!?;:')
38
            features['punct_density'] = punct_count / (len(text) + 1)
39

40
            # 大写字母比率
41
            upper_count = sum(1 for c in text if c.isupper())
42
            features['upper_ratio'] = upper_count / (len(text) + 1)
43

44
            # URL 密度（低质量内容 URL 多）
45
            url_count = len(re.findall(r'http[s]?://', text))
46
            features['url_density'] = url_count / (len(text) + 1)
47

48
            return list(features.values())
49

50
        # 训练
51
        X_pos = [extract_features(s) for s in positive_samples]
52
        X_neg = [extract_features(s) for s in negative_samples]
53

54
        X = X_pos + X_neg
55
        y = [1] * len(X_pos) + [0] * len(X_neg)
56

57
        self.quality_classifier = LogisticRegression().fit(X, y)
58

59
    def heuristic_filters(self, doc):
60
        """
61
        启发式过滤器
62
        RefineWeb 使用的 10 个启发式规则
63
        """
64
        words = doc.split()
65

66
        # 规则 1: 最小长度
67
        if len(words) < 50:
68
            return False
69

70
        # 规则 2: 最大长度（避免文档过长）
71
        if len(words) > 100000:
72
            return False
73

74
        # 规则 3: 最小平均词长
75
        avg_word_len = len(doc) / len(words)
76
        if avg_word_len < 2.5 or avg_word_len > 15:
77
            return False
78

79
        # 规则 4: 最大特殊字符密度
80
        special_chars = sum(1 for c in doc if not c.isalnum() and not c.isspace())
81
        if special_chars / len(doc) > 0.3:
82
            return False
83

84
        # 规则 5: 最小句子数
85
        sentences = sent_tokenize(doc)
86
        if len(sentences) < 3:
87
            return False
88

89
        # 规则 6: 检查"段落"结构
90
        # 好的文档应该有多个段落
91
        paragraphs = doc.split('\n\n')
92
        if len(paragraphs) < 2:
93
            return False
94

95
        # 规则 7: 检查"锚文本"密度
96
        # 锚文本（链接文字）过多可能是导航/列表页
97
        anchor_ratio = doc.count('<a ') / (len(words) + 1)
98
        if anchor_ratio > 0.3:
99
            return False
100

101
        # 规则 8: 检查"call to action"
102
        cta_patterns = ['click here', 'sign up', 'subscribe', 'buy now', 'learn more']
103
        if any(p in doc.lower() for p in cta_patterns):
104
            return False
105

106
        # 规则 9: 检查"contact us"
107
        if 'contact' in doc.lower() and 'email' in doc.lower():
108
            return False  # 可能是联系页面
109

110
        # 规则 10: 检查"terms of service"
111
        if 'terms of service' in doc.lower() or 'privacy policy' in doc.lower():
112
            return False  # 法律页面
113

114
        return True

展开收起

2.4 去重策略#

去重是预训练数据处理中最关键的步骤之一：

1
class DeduplicationStrategies:
2
    """
3
    多层次去重策略
4
    """
5

6
    def exact_deduplication(self, documents):
7
        """
8
        精确匹配去重（MinHash）
9

10
        时间复杂度: O(n)
11
        空间复杂度: O(n)
12
        """
13
        from datasketch import MinHash, MinHashEncoder
14

15
        minhashes = []
16
        for doc in documents:
17
            m = MinHash(num_perm=128)
18
            for word in doc.split():
19
                m.update(word.encode('utf8'))
20
            minhashes.append(m)
21

22
        # 编码并去重
23
        encoder = MinHashEncoder(num_perm=128)
24
        encoded = [encoder.encode(m) for m in minhashes]
25

26
        unique_idx = []
27
        seen = set()
28
        for i, e in enumerate(encoded):
29
            if e not in seen:
30
                seen.add(e)
31
                unique_idx.append(i)
32

33
        return [documents[i] for i in unique_idx]
34

35
    def near_deduplication(self, documents, threshold=0.8):
36
        """
37
        近似去重（SimHash 或 MinHash Jaccard）
38

39
        用于检测高度相似的文档
40
        threshold: Jaccard 相似度阈值
41
        """
42
        from datasketch import MinHash, MinHashLSH
43

44
        # 构建 MinHash 索引
45
        lsh = MinHashLSH(threshold=threshold, num_perm=128)
46

47
        for i, doc in enumerate(documents):
48
            m = MinHash(num_perm=128)
49
            for word in doc.split():
50
                m.update(word.encode('utf8'))
51
            lsh.insert(f"doc_{i}", m)
52

53
        # 查找并移除重复
54
        unique_docs = []
55
        duplicate_indices = set()
56

57
        for i, doc in enumerate(documents):
58
            if i in duplicate_indices:
59
                continue
60

61
            unique_docs.append(doc)
62

63
            # 查找相似文档
64
            m = MinHash(num_perm=128)
65
            for word in doc.split():
66
                m.update(word.encode('utf8'))
67

68
            duplicates = lsh.query(m)
69
            for dup_id in duplicates:
70
                dup_idx = int(dup_id.split('_')[1])
71
                if dup_idx != i:
72
                    duplicate_indices.add(dup_idx)
73

74
        return unique_docs
75

76
    def sentence_deduplication(self, text, min_occurrence=10):
77
        """
78
        句子级去重
79

80
        移除在整个语料中出现次数过多的句子
81
        这些通常是模板文本或版权声明
82
        """
83
        from collections import Counter
84

85
        sentences = sent_tokenize(text)
86

87
        # 统计句子出现次数
88
        sentence_counts = Counter(sentences)
89

90
        # 过滤出现过多的句子
91
        filtered = [
92
            s for s in sentences
93
            if sentence_counts[s] <= min_occurrence
94
        ]
95

96
        return ' '.join(filtered)
97

98
    def document_level_deduplication(self, documents):
99
        """
100
        文档级去重（基于内容哈希）
101

102
        方法：
103
        1. 提取文档的 k-mer
104
        2. 计算 MinHash
105
        3. 用 LSH 找相似文档
106
        4. 保留最长/质量最高的版本
107
        """
108
        from datasketch import MinHash, MinHashLSH
109

110
        lsh = MinHashLSH(threshold=0.9, num_perm=256)
111

112
        # 构建索引
113
        indexed = []
114
        for i, doc in enumerate(documents):
115
            m = self._compute_minhash(doc)
116
            key = f"doc_{i}"
117
            lsh.insert(key, m)
118
            indexed.append((i, doc, m))
119

120
        # 聚类并保留代表
121
        clusters = defaultdict(list)
122

123
        for i, doc, m in indexed:
124
            # 查找相似文档
125
            candidates = lsh.query(m)
126
            if candidates:
127
                cluster_id = int(candidates[0].split('_')[1])
128
                clusters[cluster_id].append((i, doc))
129
            else:
130
                clusters[i].append((i, doc))
131

132
        # 从每个 cluster 选代表（最长/最高质量）
133
        representatives = []
134
        for cluster in clusters.values():
135
            best = max(cluster, key=lambda x: len(x[1].split()))
136
            representatives.append(best[1])
137

138
        return representatives
139

140
    def _compute_minhash(self, text, k=5):
141
        """
142
        计算文本的 MinHash
143
        """
144
        from datasketch import MinHash
145

146
        words = text.split()
147
        m = MinHash(num_perm=256)
148

149
        for i in range(len(words) - k + 1):
150
            kmer = ' '.join(words[i:i+k])
151
            m.update(kmer.encode('utf8'))
152

153
        return m

展开收起

3. SFT 数据构建#

3.1 SFT 数据的核心要素#

SFT 数据与预训练数据有本质区别：它需要包含**指令（Input）和回答（Output）**的配对。

1
class SFTDataCore:
2
    """
3
    SFT 数据的核心要素
4
    """
5

6
    @dataclass
7
    class SFTSample:
8
        """
9
        SFT 样本的三个核心要素
10
        """
11
        instruction: str   # 指令/问题/任务描述
12
        input: str        # 可选的输入上下文
13
        output: str       # 期望的回答
14

15
        def format(self, template="chatml"):
16
            """
17
            格式化为训练格式
18
            """
19
            if template == "chatml":
20
                return f"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n{self.instruction}\n{self.input}<|im_end|>\n<|im_start|>assistant\n{self.output}<|im_end|>"
21
            elif template == "alpaca":
22
                return f"Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n### Instruction:\n{self.instruction}\n\n### Input:\n{self.input}\n\n### Response:\n{self.output}"
23
            elif template == "llama3":
24
                return f"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\nYou are a helpful assistant.<|eot_id|>\n<|start_header_id|>user<|end_header_id|>\n\n{self.instruction}\n{self.input}<|eot_id|>\n<|start_header_id|>assistant<|end_header_id|>\n\n{self.output}<|eot_id|>"

展开收起

3.2 指令模板设计#

指令模板的设计直接影响模型学到的能力：

1
class InstructionTemplates:
2
    """
3
    指令模板设计
4
    """
5

6
    TEMPLATES = {
7
        # 开放式问答
8
        "qa_open": [
9
            "{instruction}",
10
            "请回答以下问题：{instruction}",
11
            "Question: {instruction}\nAnswer:",
12
        ],
13

14
        # 带上下文的问答
15
        "qa_context": [
16
            "根据以下信息回答问题：\n{input}\n\n问题：{instruction}",
17
            "Context: {input}\n\nQuestion: {instruction}\n\nAnswer:",
18
        ],
19

20
        # 总结任务
21
        "summarize": [
22
            "请总结以下文本：\n{input}",
23
            "Summarize this: {input}",
24
        ],
25

26
        # 翻译任务
27
        "translate": [
28
            "请将以下文本翻译成{target_lang}：\n{input}",
29
            "Translate to {target_lang}: {input}",
30
        ],
31

32
        # 代码任务
33
        "code": [
34
            "请用{language}编写代码：\n{instruction}",
35
            "Write a {language} function that {instruction}",
36
        ],
37

38
        # 数学任务
39
        "math": [
40
            "求解：{instruction}",
41
            "Calculate: {instruction}",
42
            "Given that {input}, {instruction}",
43
        ],
44

45
        # 分析任务
46
        "analysis": [
47
            "分析以下内容：\n{input}",
48
            "请分析{instruction}，参考：\n{input}",
49
        ],
50
    }
51

52
    def generate_instruction_variants(self, base_instruction, num_variants=5):
53
        """
54
        生成指令变体
55

56
        使用 LLM 重写指令，增加多样性
57
        """
58
        prompt = f"""
59
        请将以下指令改写为 {num_variants} 种不同的表述方式。
60
        保持原意，但改变措辞和句式。
61

62
        原始指令：{base_instruction}
63

64
        改写要求：
65
        1. 使用不同的开头词
66
        2. 使用不同的句式结构
67
        3. 保持相同的任务要求
68
        4. 不要改变任务的核心目标
69
        """
70

71
        response = llm.generate(prompt)
72
        variants = extract_variants(response)
73

74
        return variants

展开收起

3.3 LLM 生成 SFT 数据#

1
class LLMDataGenerator:
2
    """
3
    使用 LLM 生成 SFT 数据
4
    """
5

6
    def __init__(self, generator_model, judge_model=None):
7
        self.generator = generator_model
8
        self.judge = judge_model
9

10
    def generate_from_prompts(self, prompts, task_type="general"):
11
        """
12
        从提示生成 SFT 数据
13
        """
14
        data = []
15

16
        for prompt in tqdm(prompts):
17
            response = self.generator.generate(prompt)
18

19
            # 质量检查
20
            if self.judge:
21
                quality = self.judge.score(prompt, response)
22
                if quality < 0.7:
23
                    continue
24

25
            data.append({
26
                "instruction": self.extract_instruction(prompt),
27
                "input": self.extract_input(prompt),
28
                "output": response,
29
                "source": "llm_generated",
30
            })
31

32
        return data
33

34
    def self_generate(self, model, instructions, num_samples_per_inst=1):
35
        """
36
        Self-Generation: 用模型自己生成回答
37

38
        典型流程：
39
        1. 准备多样化的指令集
40
        2. 用模型生成多个回答（不同 temperature）
41
        3. 过滤低质量回答
42
        4. 得到 (instruction, response) 对
43
        """
44
        data = []
45

46
        for instruction in tqdm(instructions):
47
            for _ in range(num_samples_per_inst):
48
                # 随机 temperature
49
                temp = random.choice([0.5, 0.7, 0.9, 1.1])
50

51
                response = model.generate(
52
                    instruction,
53
                    temperature=temp,
54
                    max_tokens=2048,
55
                    do_sample=True,
56
                )
57

58
                # 质量检查
59
                if self._passes_quality_check(response):
60
                    data.append({
61
                        "instruction": instruction,
62
                        "input": "",
63
                        "output": response,
64
                        "temperature": temp,
65
                    })
66

67
        return data
68

69
    def _passes_quality_check(self, response):
70
        """
71
        简单的质量检查
72
        """
73
        # 长度检查
74
        if len(response.split()) < 20:
75
            return False
76

77
        # 完整性检查
78
        if not response.strip().endswith(('.', '!', '?', '"', "'")):
79
            # 不完整可能表示截断
80
            if len(response) < 200:
81
                return False
82

83
        # 重复检查
84
        if self._has_excessive_repetition(response):
85
            return False
86

87
        return True
88

89
    def _has_excessive_repetition(self, text, n=3):
90
        """
91
        检查过度重复
92
        """
93
        words = text.lower().split()
94
        if len(words) < 10:
95
            return False
96

97
        # 检查连续 n 个词是否重复
98
        for i in range(len(words) - n * 2):
99
            window = ' '.join(words[i:i+n])
100
            next_window = ' '.join(words[i+n:i+2*n])
101
            if window == next_window:
102
                return True
103

104
        return False

展开收起

3.4 人类标注流程#

1
class HumanAnnotationWorkflow:
2
    """
3
    人类标注工作流
4
    """
5

6
    def __init__(self, annotation_platform):
7
        self.platform = annotation_platform
8

9
    def setup_task(self, task_config):
10
        """
11
        配置标注任务
12
        """
13
        task = {
14
            "name": task_config["name"],
15
            "description": task_config["description"],
16

17
            # 输入字段
18
            "input_fields": [
19
                {"name": "instruction", "type": "text", "required": True},
20
                {"name": "context", "type": "text", "required": False},
21
            ],
22

23
            # 输出字段
24
            "output_fields": [
25
                {"name": "response", "type": "textarea", "required": True},
26
                {"name": "quality_score", "type": "rating", "min": 1, "max": 5},
27
                {"name": "issues", "type": "multiselect",
28
                 "options": ["incomplete", "incorrect", "harmful", "off-topic"]},
29
            ],
30

31
            # 质量控制
32
            "qualification_requirements": [
33
                {"type": "approval_rate", "threshold": 0.85},
34
                {"type": "completed_tasks", "threshold": 100},
35
            ],
36

37
            # 奖励设置
38
            "reward": task_config.get("reward_per_task", 0.10),
39
        }
40

41
        return self.platform.create_task(task)
42

43
    def annotate_instructions(self, instructions, num_workers=10):
44
        """
45
        标注指令-回答对
46

47
        典型流程：
48
        1. 分发任务给标注者
49
        2. 每个任务多人标注
50
        3. 收集并合并结果
51
        """
52
        # 分配任务
53
        tasks = self._distribute_tasks(instructions, num_workers)
54

55
        # 收集结果
56
        all_results = []
57

58
        for worker_id, worker_tasks in tasks.items():
59
            results = self.platform.get_worker_results(worker_id)
60
            all_results.extend(results)
61

62
        # 合并重复任务（多人标注同一任务）
63
        merged = self._merge_annotations(all_results)
64

65
        # 过滤低质量标注
66
        filtered = self._filter_by_agreement(merged)
67

68
        return filtered
69

70
    def _distribute_tasks(self, instructions, num_workers):
71
        """
72
        分发任务，确保每个任务被多人标注
73
        """
74
        import random
75

76
        tasks_per_worker = len(instructions) // num_workers
77

78
        tasks = {}
79
        for i in range(num_workers):
80
            start = i * tasks_per_worker
81
            end = start + tasks_per_worker if i < num_workers - 1 else len(instructions)
82
            tasks[f"worker_{i}"] = instructions[start:end]
83

84
        return tasks
85

86
    def _merge_annotations(self, annotations):
87
        """
88
        合并同一任务的多个标注
89

90
        策略：
91
        - 如果多数一致，选择多数结果
92
        - 如果一致，选择质量最高的
93
        """
94
        from collections import defaultdict
95

96
        by_instruction = defaultdict(list)
97

98
        for ann in annotations:
99
            key = ann["instruction"]
100
            by_instruction[key].append(ann)
101

102
        merged = []
103

104
        for instruction, anns in by_instruction.items():
105
            if len(anns) == 1:
106
                merged.append(anns[0])
107
                continue
108

109
            # 检查一致性
110
            responses = [a["response"] for a in anns]
111
            response_counts = Counter(responses)
112

113
            if response_counts.most_common(1)[0][1] >= len(anns) * 0.5:
114
                # 多数一致，选择多数结果
115
                best_response = response_counts.most_common(1)[0][0]
116
            else:
117
                # 不一致，选择质量最高的
118
                best_ann = max(anns, key=lambda a: a["quality_score"])
119
                best_response = best_ann["response"]
120

121
            merged.append({
122
                "instruction": instruction,
123
                "output": best_response,
124
                "agreement": response_counts.most_common(1)[0][1] / len(anns),
125
                "num_annotations": len(anns),
126
            })
127

128
        return merged
129

130
    def _filter_by_agreement(self, annotations, min_agreement=0.5):
131
        """
132
        按一致性过滤
133
        """
134
        return [a for a in annotations if a.get("agreement", 1.0) >= min_agreement]

展开收起

4. 偏好数据构建#

4.1 偏好数据的格式#

偏好数据是对齐训练的核心，其基本格式为：

1
class PreferenceDataFormat:
2
    """
3
    偏好数据的三种格式
4
    """
5

6
    @dataclass
7
    class PreferencePair:
8
        """
9
        标准偏好对
10
        """
11
        prompt: str
12
        chosen: str      # 被选中的回答
13
        rejected: str    # 被拒绝的回答
14

15
    @dataclass
16
    class PreferenceTriple:
17
        """
18
        偏好三元组（包含多个候选）
19
        """
20
        prompt: str
21
        responses: List[str]           # 候选回答列表
22
        ranking: List[int]            # 排序（0 = 最好）
23

24
    @dataclass
25
    class BinaryPreference:
26
        """
27
        二元偏好（正/负样本）
28
        """
29
        prompt: str
30
        response: str
31
        label: bool  # True = desirable, False = undesirable

展开收起

4.2 偏好标注流水线#

1
class PreferenceAnnotationPipeline:
2
    """
3
    偏好标注流水线
4
    """
5

6
    def __init__(self, annotation_platform, judge_model=None):
7
        self.platform = annotation_platform
8
        self.judge = judge_model
9

10
    def generate_candidates(self, prompts, generator, num_candidates=4):
11
        """
12
        为每个 prompt 生成多个候选回答
13
        """
14
        candidates_data = []
15

16
        for prompt in tqdm(prompts):
17
            candidates = []
18

19
            for temp in [0.3, 0.5, 0.7, 0.9]:
20
                response = generator.generate(
21
                    prompt,
22
                    temperature=temp,
23
                    max_tokens=1024,
24
                )
25
                candidates.append(response)
26

27
            candidates_data.append({
28
                "prompt": prompt,
29
                "candidates": candidates,
30
            })
31

32
        return candidates_data
33

34
    def annotation_task(self, candidates_data, num_annotators=3):
35
        """
36
        偏好标注任务
37

38
        标注者需要：
39
        1. 阅读 prompt 和多个候选回答
40
        2. 选择最好的一个（或排序）
41
        3. 标记理由（可选）
42
        """
43
        task_config = {
44
            "name": "Preference Annotation",
45
            "instruction": """
46
            请阅读以下问题及其多个回答，选择最好的一个。
47

48
            评分标准：
49
            - 准确性：回答是否正确
50
            - 完整性：回答是否全面
51
            - 清晰度：回答是否易于理解
52
            - 有用性：回答是否有帮助
53

54
            请为每个回答打分（1-5）并选出最佳。
55
            """,
56
            "input_fields": [
57
                {"name": "prompt", "type": "display"},  # 只读
58
                {"name": "candidate_1", "type": "display"},
59
                {"name": "candidate_2", "type": "display"},
60
                {"name": "candidate_3", "type": "display"},
61
                {"name": "candidate_4", "type": "display"},
62
            ],
63
            "output_fields": [
64
                {"name": "scores", "type": "rating_group", "count": 4, "min": 1, "max": 5},
65
                {"name": "best", "type": "select", "options": ["1", "2", "3", "4"]},
66
                {"name": "reason", "type": "textarea"},
67
            ],
68
        }
69

70
        return task_config
71

72
    def build_preference_pairs(self, annotations):
73
        """
74
        从标注结果构建偏好对
75

76
        策略：两两比较构建 pair
77
        如果 A 得分 > B 得分，则 (prompt, A, B) 是一个正例
78
        """
79
        pairs = []
80

81
        for item in annotations:
82
            prompt = item["prompt"]
83
            candidates = item["candidates"]
84
            scores = item["scores"]  # [score_1, score_2, ...]
85

86
            # 构建所有两两比较
87
            for i in range(len(candidates)):
88
                for j in range(i+1, len(candidates)):
89
                    if scores[i] > scores[j]:
90
                        pairs.append({
91
                            "prompt": prompt,
92
                            "chosen": candidates[i],
93
                            "rejected": candidates[j],
94
                        })
95
                    elif scores[j] > scores[i]:
96
                        pairs.append({
97
                            "prompt": prompt,
98
                            "chosen": candidates[j],
99
                            "rejected": candidates[i],
100
                        })
101

102
        return pairs

展开收起

4.3 LLM-as-Judge#

使用 LLM 作为评判者可以大幅降低标注成本：

1
class LLMAsJudge:
2
    """
3
    LLM-as-Judge: 用 LLM 进行偏好判断
4
    """
5

6
    def __init__(self, judge_model):
7
        self.judge = judge_model
8

9
    def judge_preference(self, prompt, response_a, response_b):
10
        """
11
        判断哪个回答更好
12
        """
13
        judge_prompt = f"""
14
        You are an expert evaluator. Given a question and two responses,
15
        determine which response is better.
16

17
        Question: {prompt}
18

19
        Response A:
20
        {response_a}
21

22
        Response B:
23
        {response_b}
24

25
        Consider:
26
        - Accuracy and correctness
27
        - Completeness and thoroughness
28
        - Clarity and readability
29
        - Helpfulness and relevance
30

31
        Output your decision in JSON format:
32
        {{"winner": "A" or "B", "reason": "brief explanation"}}
33
        """
34

35
        result = self.judge.generate(judge_prompt)
36
        parsed = json.loads(result)
37

38
        return parsed
39

40
    def score_and_rank(self, prompt, responses):
41
        """
42
        给多个回答打分并排序
43
        """
44
        judge_prompt = f"""
45
        Evaluate and rank the following responses to the question.
46

47
        Question: {prompt}
48

49
        {''.join([f'Response {i+1}: {r}\n\n' for i, r in enumerate(responses)])}
50

51
        Score each response (1-10) and provide rankings.
52
        Output in JSON format:
53
        {{
54
            "rankings": [3, 1, 2, 4],  # Response indices in rank order
55
            "scores": [7.5, 9.0, 8.5, 6.0],
56
            "justifications": ["..."]
57
        }}
58
        """
59

60
        result = self.judge.generate(judge_prompt)
61
        parsed = json.loads(result)
62

63
        return parsed
64

65
    def batch_judge_preferences(self, data, batch_size=8):
66
        """
67
        批量判断偏好（提高效率）
68
        """
69
        results = []
70

71
        for i in range(0, len(data), batch_size):
72
            batch = data[i:i+batch_size]
73

74
            # 构建批量评判 prompt
75
            batch_prompt = self._build_batch_prompt(batch)
76
            batch_result = self.judge.generate(batch_prompt)
77

78
            # 解析结果
79
            parsed = self._parse_batch_result(batch_result, len(batch))
80
            results.extend(parsed)
81

82
        return results
83

84
    def _build_batch_prompt(self, batch):
85
        """
86
        构建批量评判 prompt
87
        """
88
        prompt = "Evaluate the following preference pairs:\n\n"
89

90
        for i, item in enumerate(batch):
91
            prompt += f"--- Pair {i+1} ---\n"
92
            prompt += f"Question: {item['prompt']}\n"
93
            prompt += f"Response A: {item['response_a']}\n"
94
            prompt += f"Response B: {item['response_b']}\n\n"
95

96
        prompt += "\nOutput JSON array of results:\n"
97
        prompt += '[{"pair_id": 1, "winner": "A", "reason": "..."}, ...]'
98

99
        return prompt

展开收起

4.4 偏好数据质量控制#

1
class PreferenceDataQualityControl:
2
    """
3
    偏好数据质量控制
4
    """
5

6
    def __init__(self):
7
        self.agreement_threshold = 0.6
8
        self.min_annotators = 3
9

10
    def compute_inter_annotator_agreement(self, annotations):
11
        """
12
        计算标注者一致性
13

14
        方法：Fleiss' Kappa 或 Krippendorff's Alpha
15
        """
16
        import numpy as np
17
        from scipy import stats
18

19
        # 将评分转换为矩阵
20
        # 行：样本
21
        # 列：标注者
22
        n_samples = len(set(a["sample_id"] for a in annotations))
23
        n_annotators = self.min_annotators
24

25
        ratings = np.zeros((n_samples, n_annotators))
26

27
        for ann in annotations:
28
            sid = ann["sample_id"]
29
            aid = ann["annotator_id"]
30
            score = ann["score"]
31
            ratings[sid, aid] = score
32

33
        # 计算 Fleiss' Kappa
34
        kappa = self._fleiss_kappa(ratings)
35

36
        return kappa
37

38
    def _fleiss_kappa(self, ratings):
39
        """
40
        Fleiss' Kappa 计算
41
        """
42
        n_subjects, n_raters = ratings.shape
43

44
        # 观测一致性
45
        observed = self._compute_observed_agreement(ratings)
46

47
        # 期望一致性（随机）
48
        expected = self._compute_expected_agreement(ratings)
49

50
        kappa = (observed - expected) / (1 - expected)
51

52
        return kappa
53

54
    def filter_by_consistency(self, preference_pairs):
55
        """
56
        过滤不一致的偏好
57

58
        检测方法：
59
        1. A > B, B > C, 但 A < C（循环偏好）
60
        2. 同一 prompt 多次标注结果差异大
61
        """
62
        from collections import defaultdict
63

64
        # 按 prompt 分组
65
        by_prompt = defaultdict(list)
66
        for pair in preference_pairs:
67
            by_prompt[pair["prompt"]].append(pair)
68

69
        filtered = []
70

71
        for prompt, pairs in by_prompt.items():
72
            # 检查是否有循环偏好
73
            if self._has_cycle_preference(pairs):
74
                continue  # 跳过不一致的数据
75

76
            # 检查标注一致性
77
            consistency = self._compute_pairwise_consistency(pairs)
78
            if consistency >= self.agreement_threshold:
79
                filtered.extend(pairs)
80

81
        return filtered
82

83
    def _has_cycle_preference(self, pairs):
84
        """
85
        检测循环偏好
86
        """
87
        # 如果有 A > B > C > A 的循环，返回 True
88
        # 简化为：检查是否有传递性冲突
89
        responses = {}
90

91
        for pair in pairs:
92
            chosen = pair["chosen"]
93
            rejected = pair["rejected"]
94

95
            if chosen not in responses:
96
                responses[chosen] = set()
97
            if rejected not in responses:
98
                responses[rejected] = set()
99

100
            # chosen 优于 rejected
101
            responses[chosen].add(rejected)
102

103
        # 检查传递性
104
        for a in responses:
105
            for b in responses[a]:
106
                for c in responses[b]:
107
                    if a in responses[c]:
108
                        # a > b > c > a，存在循环
109
                        return True
110

111
        return False
112

113
    def filter_by_length_ratio(self, pairs, max_ratio=5.0):
114
        """
115
        过滤长度差异过大的偏好
116

117
        理由：长度偏差可能影响人类判断
118
        """
119
        filtered = []
120

121
        for pair in pairs:
122
            len_chosen = len(pair["chosen"].split())
123
            len_rejected = len(pair["rejected"].split())
124

125
            ratio = max(len_chosen, len_rejected) / (min(len_chosen, len_rejected) + 1)
126

127
            if ratio <= max_ratio:
128
                filtered.append(pair)
129

130
        return filtered

展开收起

5. 测试集设计#

5.1 测试集的核心原则#

1
class TestSetDesign:
2
    """
3
    测试集设计的核心原则
4
    """
5

6
    PRINCIPLES = {
7
        # 1. 不泄漏训练数据
8
        "no_leakage": """
9
        测试集不应该出现在训练数据中。
10

11
        检查方法：
12
        1. 用 MinHash 检测与训练数据的重叠
13
        2. 过滤重叠度 > 5% 的样本
14
        3. 定期审计测试集
15
        """,
16

17
        # 2. 代表性
18
        "representativeness": """
19
        测试集应该代表模型将面对的真实分布。
20

21
        检查方法：
22
        1. 与训练数据分布对比
23
        2. 覆盖所有能力维度
24
        3. 考虑真实使用场景
25
        """,
26

27
        # 3. 难度分布
28
        "difficulty_distribution": """
29
        测试集应该有合理的难度分布。
30

31
        建议分布：
32
        - Easy: 30%
33
        - Medium: 50%
34
        - Hard: 20%
35

36
        理由：
37
        - 太难：无法区分模型能力
38
        - 太简单：无法区分模型能力
39
        """,
40

41
        # 4. 抗污染
42
        "contamination_resistance": """
43
        测试集应该能抵抗数据污染。
44

45
        方法：
46
        1. 使用闭卷测试（模型未见过的场景）
47
        2. 使用需要推理的问题
48
        3. 定期更新测试集
49
        """,
50
    }

展开收起

5.2 能力维度覆盖#

1
class CapabilityCoverage:
2
    """
3
    能力维度的测试覆盖
4
    """
5

6
    CAPABILITIES = {
7
        # 知识与问答
8
        "knowledge": {
9
            "description": "知识问答和事实准确性",
10
            "tests": [
11
                "Factual QA (TriviaQA, NaturalQuestions)",
12
                "World Knowledge (MMLU subjects)",
13
                "Reading Comprehension (SQuAD, HotpotQA)",
14
            ],
15
            "difficulty_levels": ["factual", "reasoning", "multi-hop"],
16
        },
17

18
        # 推理能力
19
        "reasoning": {
20
            "description": "逻辑和数学推理",
21
            "tests": [
22
                "Math (GSM8K, MATH, MATH-500)",
23
                "Logic (LogiQA, ReClor)",
24
                "Commonsense (HellaSwag, PIQA)",
25
                "Code (HumanEval, MBPP, LiveCodeBench)",
26
            ],
27
            "difficulty_levels": ["single-step", "multi-step", "adversarial"],
28
        },
29

30
        # 代码能力
31
        "coding": {
32
            "description": "代码生成和理解",
33
            "tests": [
34
                "Code Generation (HumanEval, MBPP)",
35
                "Code Completion (Fill-in-the-middle)",
36
                "Code Explanation (Code Understanding)",
37
                "Code Debugging (Bug Detection)",
38
            ],
39
            "difficulty_levels": ["simple", "medium", "complex"],
40
        },
41

42
        # 对齐质量
43
        "alignment": {
44
            "description": "指令遵循和偏好对齐",
45
            "tests": [
46
                "Instruction Following (IFEval)",
47
                "Helpfulness (MT-Bench, AlpacaEval)",
48
                "Safety (ToxiGen, HarmBench)",
49
                "Honesty (TruthfulQA)",
50
            ],
51
            "difficulty_levels": ["basic", "complex", "adversarial"],
52
        },
53

54
        # 多语言能力
55
        "multilingual": {
56
            "description": "多语言理解和生成",
57
            "tests": [
58
                "Translation (WMT, Flores)",
59
                "Multilingual QA (XQuAD, MLQA)",
60
                "Language Understanding (XGLUE)",
61
            ],
62
            "difficulty_levels": ["high-resource", "low-resource"],
63
        },
64
    }

展开收起

5.3 测试集去污染#

1
class ContaminationChecker:
2
    """
3
    测试集污染检测
4
    """
5

6
    def __init__(self, training_data):
7
        self.training_minhash = self._build_minhash_index(training_data)
8

9
    def _build_minhash_index(self, documents):
10
        """
11
        为训练数据构建 MinHash 索引
12
        """
13
        from datasketch import MinHash, MinHashLSH
14

15
        lsh = MinHashLSH(threshold=0.5, num_perm=128)
16

17
        for i, doc in enumerate(documents):
18
            m = MinHash(num_perm=128)
19
            for word in doc.split()[:1000]:  # 只用前1000词
20
                m.update(word.encode('utf8'))
21
            lsh.insert(f"doc_{i}", m)
22

23
        return lsh
24

25
    def check_contamination(self, test_sample):
26
        """
27
        检查测试样本是否与训练数据重叠
28
        """
29
        m = MinHash(num_perm=128)
30
        for word in test_sample.split()[:1000]:
31
            m.update(word.encode('utf8'))
32

33
        candidates = self.training_minhash.query(m)
34

35
        return len(candidates) > 0, len(candidates)
36

37
    def compute_overlap_score(self, test_sample):
38
        """
39
        计算与训练数据的重叠分数
40
        """
41
        test_words = set(test_sample.lower().split())
42

43
        # 采样检查（完整检查太慢）
44
        overlap_count = 0
45
        sample_size = min(len(test_words), 1000)
46
        sampled_words = random.sample(list(test_words), sample_size)
47

48
        for word in sampled_words:
49
            if self._word_in_training(word):
50
                overlap_count += 1
51

52
        return overlap_count / sample_size
53

54
    def filter_contaminated(self, test_set, overlap_threshold=0.1):
55
        """
56
        过滤被污染的测试样本
57
        """
58
        clean_set = []
59
        contaminated_count = 0
60

61
        for sample in tqdm(test_set):
62
            overlap = self.compute_overlap_score(sample["text"])
63

64
            if overlap < overlap_threshold:
65
                clean_set.append(sample)
66
            else:
67
                contaminated_count += 1
68

69
        print(f"Filtered {contaminated_count}/{len(test_set)} contaminated samples")
70

71
        return clean_set

展开收起

6. Scaling Law 与数据配比#

6.1 数据 Scaling Law#

1
class DataScalingLaw:
2
    """
3
    数据 Scaling Law
4

5
    核心发现（Hoffmann et al., 2022）：
6
    计算最优：C ≈ 20 * N
7
    其中 C 是 tokens 数量，N 是参数量
8

9
    Chinchilla 结论：
10
    模型大小和数据量应该同比 scaling
11
    """
12

13
    def compute_optimal_tokens(self, model_params, total_compute):
14
        """
15
        Chinchilla 最优：20 tokens per parameter
16

17
        LLaMA 2 (70B) 应该用 1.4T tokens
18
        实际用了 2T，已经超过 Chinchilla 最优
19
        """
20
        optimal_tokens = 20 * model_params
21

22
        return optimal_tokens
23

24
    def compute_optimal_model_size(self, total_tokens, total_compute):
25
        """
26
        给定 tokens 和 compute，估算最优模型大小
27
        """
28
        # total_compute ∝ N * T
29
        # 最优点：N ≈ C / 20
30

31
        optimal_params = total_compute / 20
32

33
        return optimal_params
34

35
    def predict_performance(self, model_params, num_tokens, quality_factor=1.0):
36
        """
37
        预测模型性能（基于 scaling law）
38

39
        损失 ~ (C / (N * T))^(alpha)
40

41
        其中：
42
        - C 是计算量
43
        - N 是参数量
44
        - T 是 tokens
45
        - alpha ≈ 0.5
46
        """
47
        alpha = 0.5
48
        quality_factor = quality_factor  # 数据质量因子
49

50
        compute = model_params * num_tokens
51

52
        # 简化的损失预测
53
        loss = (1.0 / compute) ** alpha * quality_factor
54

55
        # 转换为困惑度
56
        perplexity = np.exp(loss)
57

58
        return perplexity

展开收起

6.2 数据配比策略#

1
class DataMixingStrategy:
2
    """
3
    数据配比策略
4
    """
5

6
    def compute_optimal_mix(
7
        self,
8
        domain_importance,
9
        domain_difficulty,
10
        current_performance,
11
        target_capabilities,
12
    ):
13
        """
14
        计算最优数据配比
15

16
        考虑因素：
17
        1. 领域重要性（任务占比）
18
        2. 领域难度（需要多少数据才能掌握）
19
        3. 当前性能（短板优先）
20
        4. 目标能力（重点发展的方向）
21
        """
22
        mix = {}
23

24
        for domain, importance in domain_importance.items():
25
            difficulty = domain_difficulty.get(domain, 1.0)
26
            current = current_performance.get(domain, 0.0)
27
            target = target_capabilities.get(domain, 1.0)
28

29
            # 基础权重 = 重要性
30
            weight = importance
31

32
            # 难度调整：越难需要越多数据
33
            weight *= (1 + 0.5 * difficulty)
34

35
            # 性能调整：短板优先
36
            gap = target - current
37
            weight *= (1 + gap)
38

39
            mix[domain] = weight
40

41
        # 归一化
42
        total = sum(mix.values())
43
        mix = {k: v / total for k, v in mix.items()}
44

45
        return mix
46

47
    def expert_mixture(self, num_experts=8, expert_domain_size=100e9):
48
        """
49
        专家混合数据配比
50

51
        DeepSeek-V2 的做法：
52
        每个"专家"专注一个领域
53
        """
54
        mix = {}
55

56
        domains = [
57
            "code",          # 代码
58
            "math",          # 数学
59
            "reasoning",     # 逻辑推理
60
            "science",        # 自然科学
61
            " humanities",   # 人文
62
            "social",        # 社会科学
63
            "news",          # 新闻
64
            "web",           # 网页
65
        ]
66

67
        # 基于能力目标分配
68
        base_mix = {
69
            "code": 0.20,      # 代码是关键
70
            "math": 0.15,      # 数学推理
71
            "reasoning": 0.15, # 逻辑推理
72
            "science": 0.15,   # 自然科学
73
            "humanities": 0.10,
74
            "social": 0.10,
75
            "news": 0.05,
76
            "web": 0.10,
77
        }
78

79
        return base_mix
80

81
    def progressive_mix(self, training_steps, total_steps):
82
        """
83
        渐进式数据混合
84

85
        训练过程中动态调整数据配比
86
        """
87
        progress = training_steps / total_steps
88

89
        if progress < 0.3:
90
            # 早期：更多通用数据
91
            return {
92
                "general": 0.8,
93
                "domain": 0.2,
94
            }
95
        elif progress < 0.7:
96
            # 中期：增加领域数据
97
            return {
98
                "general": 0.5,
99
                "domain": 0.5,
100
            }
101
        else:
102
            # 后期：领域数据为主
103
            return {
104
                "general": 0.2,
105
                "domain": 0.8,
106
            }

展开收起

6.3 数据质量 vs 数量#

1
class QualityVsQuantity:
2
    """
3
    数据质量 vs 数量的权衡
4
    """
5

6
    def analyze_tradeoff(self):
7
        """
8
        分析质量与数量的权衡
9

10
        关键发现：
11
        1. Phi-3: 3.3T 高质量 tokens ≈ 15T 混合质量 tokens
12
        2. 数据质量提升 2x ≈ 数据量提升 10x
13
        3. "Textbook quality" 数据最有效
14
        """
15
        return {
16
            "phi_insight": {
17
                "description": "Phi-3 的数据策略",
18
                "high_quality_tokens": 3.3e12,
19
                "equivalent_mixed_tokens": 15e12,
20
                "quality_premium": "5x",
21
                "key": "Textbook-level quality",
22
            },
23

24
            "scaling_curve": {
25
                "x": "Data quality (normalized)",
26
                "y": "Model performance (normalized)",
27
                "shape": "log-linear",
28
                "slope": 0.5,  # 质量翻倍 → 性能提升 2^0.5 ≈ 1.4x
29
            },
30

31
            "practical_recommendation": {
32
                "high_budget": "Maximize quality, moderate quantity",
33
                "medium_budget": "Balance quality and quantity",
34
                "low_budget": "Prioritize quality in key domains",
35
            },
36
        }

展开收起

7. 数据清洗工具链#

7.1 完整清洗流水线#

1
class DataCleaningPipeline:
2
    """
3
    完整数据清洗流水线
4
    """
5

6
    def __init__(self):
7
        self.stages = [
8
            "deduplication",
9
            "language_filter",
10
            "quality_filter",
11
            "safety_filter",
12
            "format_filter",
13
        ]
14

15
    def clean(self, documents, config=None):
16
        """
17
        完整清洗流水线
18
        """
19
        if config is None:
20
            config = self.default_config()
21

22
        current = documents
23
        stats = {"input": len(documents)}
24

25
        for stage in self.stages:
26
            current = getattr(self, f"stage_{stage}")(current, config[stage])
27
            stats[stage] = len(current)
28

29
        stats["output"] = len(current)
30
        stats["retention_rate"] = len(current) / stats["input"]
31

32
        return current, stats
33

34
    def default_config(self):
35
        """
36
        默认配置
37
        """
38
        return {
39
            "deduplication": {
40
                "threshold": 0.85,  # MinHash Jaccard 阈值
41
                "method": "minhash",
42
            },
43
            "language_filter": {
44
                "target_languages": ["en", "zh", "es", "fr", "de", "ja", "ko"],
45
                "min_confidence": 0.8,
46
            },
47
            "quality_filter": {
48
                "min_length": 100,
49
                "max_length": 100000,
50
                "min_sentence_count": 3,
51
                "quality_threshold": 0.5,
52
            },
53
            "safety_filter": {
54
                "toxicity_threshold": 0.5,
55
                "pii_removal": True,
56
            },
57
            "format_filter": {
58
                "remove_html": True,
59
                "remove_urls": False,
60
                "normalize_whitespace": True,
61
            },
62
        }
63

64
    def stage_deduplication(self, documents, config):
65
        """
66
        去重阶段
67
        """
68
        return self.near_deduplication(documents, threshold=config["threshold"])
69

70
    def stage_language_filter(self, documents, config):
71
        """
72
        语言过滤阶段
73
        """
74
        from fasttext import load_model
75

76
        lid_model = load_model('lid.176.bin')
77

78
        filtered = []
79
        for doc in documents:
80
            predictions = lid_model.predict(doc, k=1)
81
            lang = predictions[0][0].replace('__label__', '')
82
            confidence = predictions[1][0]
83

84
            if lang in config["target_languages"] and confidence >= config["min_confidence"]:
85
                filtered.append(doc)
86

87
        return filtered
88

89
    def stage_quality_filter(self, documents, config):
90
        """
91
        质量过滤阶段
92
        """
93
        filtered = []
94

95
        for doc in documents:
96
            words = doc.split()
97

98
            # 长度检查
99
            if len(words) < config["min_length"]:
100
                continue
101
            if len(words) > config["max_length"]:
102
                continue
103

104
            # 句子数检查
105
            sentences = sent_tokenize(doc)
106
            if len(sentences) < config["min_sentence_count"]:
107
                continue
108

109
            # 质量评分
110
            quality = self.compute_quality_score(doc)
111
            if quality < config["quality_threshold"]:
112
                continue
113

114
            filtered.append(doc)
115

116
        return filtered
117

118
    def stage_safety_filter(self, documents, config):
119
        """
120
        安全过滤阶段
121
        """
122
        # 简化实现
123
        toxic_patterns = [
124
            "hate speech patterns",
125
            "violent content patterns",
126
        ]
127

128
        filtered = []
129
        for doc in documents:
130
            # 检查毒性
131
            toxicity_score = self.estimate_toxicity(doc)
132
            if toxicity_score > config["toxicity_threshold"]:
133
                continue
134

135
            # PII 移除
136
            if config.get("pii_removal", False):
137
                doc = self.remove_pii(doc)
138

139
            filtered.append(doc)
140

141
        return filtered
142

143
    def stage_format_filter(self, documents, config):
144
        """
145
        格式过滤阶段
146
        """
147
        cleaned = []
148

149
        for doc in documents:
150
            if config.get("remove_html", True):
151
                doc = self.remove_html_tags(doc)
152

153
            if config.get("remove_urls", True):
154
                doc = re.sub(r'http[s]?://\S+', '', doc)
155

156
            if config.get("normalize_whitespace", True):
157
                doc = re.sub(r'\s+', ' ', doc)
158

159
            cleaned.append(doc.strip())
160

161
        return cleaned

展开收起

8. 主流数据集汇总#

8.1 预训练数据集#

数据集	规模	来源	特点
CommonCrawl	~数十 TB	网页	最大、噪声多
The Pile	825 GB	22 个来源混合	多样、学术导向
RedPajama	1.2T	C4 + Reddit + Wikipedia + etc.	开源完整
RefineWeb	500B	网页	高质量
Dolma	3T	多种来源	最新、高质量
FineWeb	15T	网页	高质量过滤
Cosmopedia	25B	Synthetic	高质量教科书级

8.2 SFT 数据集#

数据集	规模	来源	特点
FLAN	1.8M	62 任务	多任务指令微调
Alpaca	52K	LLM 生成	简单、便宜
Vicuna	70K	LLM 对话	多轮对话
WizardLM	180K	LLM 进化	复杂指令
ShareGPT	~150K	用户共享	真实对话
OpenOrca	~1M	LLM 生成	CoT 推理

8.3 偏好数据集#

数据集	规模	来源	特点
HH-RLHF	160K	人类标注	Helpfulness/Harmlessness
SHP	385K	Reddit 投票	人类偏好
PKU-Alignment	340K	人类标注	安全对齐
Anthropic HH	42K	人类标注	对话偏好
UltraFeedback	64K	LLM 评判	高质量

9. 核心公式汇总#

9.1 MinHash 去重#

hmin⁡(S)=min⁡s∈Sh(s)h_{\min}(S) = \min_{s \in S} h(s)hmin(S)=s∈Sminh(s)

其中 h(s)h(s)h(s) 是元素 sss 的哈希函数，MinHash 是多个哈希函数的集合。

9.2 Jaccard 相似度#

J(A,B)=∣A∩B∣∣A∪B∣J(A, B) = \frac{|A \cap B|}{|A \cup B|}J(A,B)=∣A∪B∣∣A∩B∣

9.3 质量评分（简化）#

Q(d)=α⋅C(d)+β⋅L(d)+γ⋅R(d)Q(d) = \alpha \cdot C(d) + \beta \cdot L(d) + \gamma \cdot R(d)Q(d)=α⋅C(d)+β⋅L(d)+γ⋅R(d)

其中 CCC 是完整性、LLL 是词汇丰富度、RRR 是重复度。

9.4 Chinchilla Scaling Law#

L(N)≈(NcN)αN+(DcD)αDL(N) \approx \left(\frac{N_c}{N}\right)^{\alpha_N} + \left(\frac{D_c}{D}\right)^{\alpha_D}L(N)≈(NNc)αN+(DDc)αD

其中 NNN 是参数量，DDD 是数据量，LLL 是损失。

10. 总结#

10.1 数据构建的核心原则#

1
┌─────────────────────────────────────────────────────────────┐
2
│                 数据构建的五大核心原则                         │
3
├─────────────────────────────────────────────────────────────┤
4
│                                                              │
5
│  1. 质量第一                                                  │
6
│     → 高质量 3.3T tokens > 低质量 15T tokens                 │
7
│     → Textbooks are all you need                              │
8
│     → 投资在质量过滤上有最高的 ROI                             │
9
│                                                              │
10
│  2. 去重不可忽视                                              │
11
│     → 去重可移除 20-40% 的重复数据                            │
12
│     → MinHash 是工业级去重的标准方法                          │
13
│     → 文档级 + 句子级去重组合使用                             │
14
│                                                              │
15
│  3. 多样性很重要                                              │
16
│     → 单一来源的数据会导致能力偏差                             │
17
│     → 平衡不同领域、风格、难度的数据                           │
18
│     → 渐进式调整数据配比                                       │
19
│                                                              │
20
│  4. 测试集要独立                                              │
21
│     → 严格防止训练数据泄漏测试集                               │
22
│     → 定期审计和更新测试集                                     │
23
│     → 使用闭卷测试验证真实能力                                 │
24
│                                                              │
25
│  5. 迭代改进                                                  │
26
│     → 训练实验 → 分析失败案例 → 针对性补充数据                 │
27
│     → 数据构建是一个持续的过程                                 │
28
│     → 不要期望一次就做到完美                                   │
29
│                                                              │
30
└─────────────────────────────────────────────────────────────┘

展开收起

10.2 方法选择指南#

1
数据构建方法选择：
2

3
  预训练数据：
4
    → 质量优先：RefineWeb 流水线 > CCNet
5
    → 开源首选：RedPajama-Data > 从零爬取
6
    → 大规模：FineWeb / Dolma
7

8
  SFT 数据：
9
    → 资源充足：人类标注
10
    → 资源有限：LLM 生成 + 质量过滤
11
    → 最佳平衡：混合（50% 人工 + 50% LLM）
12

13
  偏好数据：
14
    → LLM-as-Judge 是降低成本的关键
15
    → 多标注者一致性 > 标注数量
16
    → 偏好对质量比数量更重要

展开收起

参考资料#

Wenzek, G., et al. (2020). “CCNet: Extracting High Quality Monolingual Datasets from Web Crawl Data.” LREC.
Faoorosh, G., et al. (2023). “RefineWeb: A Recipe for Efficient Data Cleaning.” arXiv.
Gao, L., et al. (2020). “The Pile: An 800GB Dataset of Diverse Text for Language Modeling.” arXiv.
Penedo, G., et al. (2023). “The RefinedWeb Dataset for Falcon LLM.” arXiv.
Soldaini, L., et al. (2024). “Dolma: An Open Corpus of 3 Trillion Tokens for Language Model Pretraining.” arXiv.
Labrak, Y., et al. (2024). “FineWeb: 15T tokens of high-quality web data.” arXiv.
Azeret, A., et al. (2024). “Cosmopedia: Synthetic Data for Education.” arXiv.
Taori, R., et al. (2023). “Alpaca: A Strong, Replicable Instruction-Following Model.” Stanford CRFM.
Peng, B., et al. (2023). “Vicuna: An Open-S chatbot Impressing GPT-4 with 90%* ChatGPT Quality.” LMSYS.
Xu, C., et al. (2023). “WizardLM: Empowering Large Language Models to Follow Complex Instructions.” ICLR.
Bai, Y., et al. (2022). “Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback.” arXiv.
Zheng, L., et al. (2023). “Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena.” NeurIPS.
Hoffmann, J., et al. (2022). “Training Compute-Optimal Large Language Models.” NeurIPS.
Touvron, H., et al. (2023). “LLaMA: Open and Efficient Foundation Language Models.” Meta AI.
Team, G., et al. (2024). “Phi-3 Technical Report.” Microsoft.

DDIM 深度剖析：从随机马尔可夫链到确定性隐式模型

Thu, 16 Jul 2026 00:00:00 GMT

1. DDPM 的致命弱点：为什么需要 DDIM？#

1.1 从 DDPM 到 DDIM 的动机#

回顾 DDPM 文章：DDPM 用一个 T=1000T = 1000T=1000 步的马尔可夫链逆向去噪，每一步都加随机噪声。生成一张 512×512 图像需要：

1
DDPM 生成时间 ≈ 1000 步 × 每步 ~100ms ≈ 100 秒 (GPU)

这对实际应用来说是不可接受的。核心矛盾：

问题	DDPM 的选择	后果
逆向过程是随机的	xt−1∼N(μ,σ2)x_{t-1} \sim \mathcal{N}(\mu, \sigma^2)xt−1∼N(μ,σ2)	每步独立随机噪声，无法利用前面的信息
马尔可夫假设	p(xt−1∥xt)p(x_{t-1}\\|x_t)p(xt−1∥xt) 只依赖 xtx_txt	无法”跳跃”——必须一步步走
生成轨迹	高度发散（随机游走）	步数少时质量急剧下降

1.2 DDIM 的核心洞察#

放弃马尔可夫假设，换取确定性和速度。

Song et al. 2021 提出的 DDIM（Denoising Diffusion Implicit Models）发现：

1
DDPM 的核心假设"逆向过程必须是马尔可夫的"是人为加的。
2
如果只要求:
3
  1. 逆向过程是"隐式概率模型"（Implicit Probabilistic Model）
4
  2. 边际分布 q(x_t) 不变（训练数据分布不变）
5
那么逆向过程可以完全自由设计——包括非马尔可夫的确定性过程！

这就是”Implicit Models”（隐式模型）名称的由来：不再显式定义 p(xt−1∣xt)p(x_{t-1}|x_t)p(xt−1∣xt) 的概率密度，而是定义一个确定性映射 xt−1=fθ(xt,t)x_{t-1} = f_\theta(x_t, t)xt−1=fθ(xt,t)，通过 ODE 积分来生成样本。

1.3 一句话概括 DDIM#

DDIM 通过放弃”逆向过程必须是马尔可夫链”的约束，把 DDPM 的随机走路径变成确定性的直线路径——从而用 10-50 步生成与 1000 步 DDPM 相当甚至更好的图像质量。

1.4 论文信息#

1
论文: "Denoising Diffusion Implicit Models" (DDIM)
2
作者: Jiaming Song, Chenlin Meng, Stefano Ermon
3
单位: Stanford University
4
发表于: ICLR 2021
5
引用: > 5,000 次 (截至 2024)
6
开源: https://github.com/ermongroup/ddim

2. 数学框架：非马尔可夫逆向过程#

2.1 隐式概率模型的定义#

DDPM 定义了一个生成过程（马尔可夫链）：

pθ(x0:T)=p(xT)∏t=1Tpθ(t)(xt−1∣xt)p_\theta(x_{0:T}) = p(x_T) \prod_{t=1}^T p_\theta^{(t)}(x_{t-1} | x_t)pθ(x0:T)=p(xT)t=1∏Tpθ(t)(xt−1∣xt)

其中 pθ(t)p_\theta^{(t)}pθ(t) 是神经网络近似的逆向分布。

DDIM 则定义了一个更一般的生成过程，称为非马尔可夫链：

pθ(x0:T)=p(xT)∏t=1Tp~θ(t)(xt−1∣xt,x0)p_\theta(x_{0:T}) = p(x_T) \prod_{t=1}^T \tilde{p}_\theta^{(t)}(x_{t-1} | x_t, x_0)pθ(x0:T)=p(xT)t=1∏Tp~θ(t)(xt−1∣xt,x0)

注意：p~θ(t)\tilde{p}_\theta^{(t)}p~θ(t) 依赖 x0x_0x0（不只是 xtx_txt）——这打破了马尔可夫假设。

关键约束：边际分布 q(xt)q(x_t)q(xt) 必须保持不变（即与前向过程定义的分布一致），否则训练目标会失效。

2.2 边际分布约束#

设 q(x1:T∣x0)q(x_{1:T}|x_0)q(x1:T∣x0) 是前向过程定义的联合分布。DDIM 要求生成模型的边际分布 pθ(xt)p_\theta(x_t)pθ(xt) 与 q(xt)q(x_t)q(xt) 匹配：

pθ(xt)=q(xt)∀tp_\theta(x_t) = q(x_t) \quad \forall tpθ(xt)=q(xt)∀t

这是一致性约束：不管用马尔可夫还是非马尔可夫方式生成，边际分布必须一致。

从 q(x1:T∣x0)q(x_{1:T}|x_0)q(x1:T∣x0) 的马尔可夫链结构出发，可以推导出 q(xt−1∣xt,x0)q(x_{t-1}|x_t, x_0)q(xt−1∣xt,x0) 的解析形式（贝叶斯定理）：

q(xt−1∣xt,x0)=N(xt−1;μ~t(xt,x0),σ~t2I)q(x_{t-1} | x_t, x_0) = \mathcal{N}\left(x_{t-1}; \tilde{\mu}_t(x_t, x_0), \tilde{\sigma}_t^2 I\right)q(xt−1∣xt,x0)=N(xt−1;μ~t(xt,x0),σ~t2I)

其中均值和方差与 DDPM 中的逆向条件分布完全相同（来自 DDPM 文章的推导）。

2.3 DDIM 逆向过程的定义#

设 σt∈[0,1]\sigma_t \in [0, 1]σt∈[0,1] 是自由参数（DDPM 中固定为 βt1/2\beta_t^{1/2}βt1/2）。DDIM 定义逆向过程为：

xt−1=αˉt−1x0+1−αˉt−1−σt2⋅ϵθ(xt,t)+σt⋅ϵx_{t-1} = \sqrt{\bar\alpha_{t-1}} x_0 + \sqrt{1 - \bar\alpha_{t-1} - \sigma_t^2} \cdot \epsilon_\theta(x_t, t) + \sigma_t \cdot \epsilonxt−1=αˉt−1x0+1−αˉt−1−σt2⋅ϵθ(xt,t)+σt⋅ϵ

其中：

ϵθ(xt,t)\epsilon_\theta(x_t, t)ϵθ(xt,t) 是神经网络预测的噪声（与 DDPM 相同）
ϵ∼N(0,I)\epsilon \sim \mathcal{N}(0, I)ϵ∼N(0,I) 是可选的随机噪声
当 σt=0\sigma_t = 0σt=0 时，过程完全确定性——这就是 DDIM 的核心

2.4 三种特殊情形#

情形	σt\sigma_tσt 取值	过程类型	采样步数影响
DDPM（完全随机）	σt=1−αˉt−1\sigma_t = \sqrt{1 - \bar\alpha_{t-1}}σt=1−αˉt−1	随机马尔可夫	必须多步，质量急剧下降
DDIM（完全确定）	σt=0\sigma_t = 0σt=0	确定性非马尔可夫	步数减少时质量下降最慢
插值（部分随机）	0<σt<1−αˉt−10 < \sigma_t < \sqrt{1 - \bar\alpha_{t-1}}0<σt<1−αˉt−1	随机非马尔可夫	介于两者之间

1
def ddim_reverse_step(xt, eps_theta, t, t_prev, sigma_t):
2
    """
3
    DDIM 单步逆向过程。
4

5
    参数:
6
        xt:        当前位置 x_t
7
        eps_theta: 神经网络预测的噪声
8
        t:         当前时间步 (整数)
9
        t_prev:    上一个时间步 (整数)
10
        sigma_t:   随机性参数 [0, 1]
11
    """
12
    # 从 xt 估计 x0
13
    alpha_bar_t = alphas_cumprod[t]
14
    sqrt_ab     = torch.sqrt(alpha_bar_t)
15
    sqrt_1m     = torch.sqrt(1 - alpha_bar_t)
16
    x0_pred     = (xt - sqrt_1m * eps_theta) / sqrt_ab
17

18
    # 计算 x_{t-1}
19
    alpha_bar_prev = alphas_cumprod[t_prev]
20
    sqrt_ab_prev   = torch.sqrt(alpha_bar_prev)
21
    sqrt_1m_prev   = torch.sqrt(1 - alpha_bar_prev)
22

23
    # 确定性方向（始终存在）
24
    pred_eps    = eps_theta
25
    x_prev_deterministic = sqrt_ab_prev * x0_pred + sqrt_1m_prev * pred_eps
26

27
    # 随机性方向（sigma_t > 0 时加入）
28
    if sigma_t > 0:
29
        noise     = torch.randn_like(xt)
30
        x_prev    = x_prev_deterministic + sigma_t * noise
31
    else:
32
        x_prev    = x_prev_deterministic
33

34
    return x_prev

展开收起

3. 轨迹视图：从 SDE 到 ODE#

3.1 把 DDPM 看作 SDE#

连续时间下，DDPM 的前向过程是一个随机微分方程（SDE）：

dx=f(x,t)dt+g(t)dwdx = f(x, t) dt + g(t) dwdx=f(x,t)dt+g(t)dw

其中 dwdwdw 是维纳过程（Wiener process），引入随机性。

DDPM 的逆向过程（在连续时间下）对应 SDE 的反向：

dxˉ=[f(xˉ,t)−g(t)2∇xlog⁡pt(xˉ)]dt+g(t)dwˉd\bar{x} = \left[ f(\bar{x}, t) - g(t)^2 \nabla_x \log p_t(\bar{x}) \right] dt + g(t) d\bar{w}dxˉ=[f(xˉ,t)−g(t)2∇xlogpt(xˉ)]dt+g(t)dwˉ

3.2 把 DDIM 看作 ODE#

当 σt=0\sigma_t = 0σt=0 时，随机项消失，DDIM 变成常微分方程（ODE）：

dxdt=[f(x,t)−12g(t)2∇xlog⁡pt(x)]\frac{dx}{dt} = \left[ f(x, t) - \frac{1}{2} g(t)^2 \nabla_x \log p_t(x) \right]dtdx=[f(x,t)−21g(t)2∇xlogpt(x)]

ODE 的好处：

1
SDE (DDPM):    随机轨迹，每步独立，无法加速
2
                ↓
3
ODE (DDIM):    确定性轨迹，可以用 ODE 求解器加速
4
                - 自适应步长
5
                - 高阶积分（Heun, RK45）
6
                - 逆向轨迹固定，步数减少时误差可预测

3.3 ODE 求解器与采样#

1
@torch.no_grad()
2
def ddim_ode_sampler(model, num_steps=50, eta=0.0):
3
    """
4
    DDIM ODE 采样器（确定性）。
5
    eta = 0   → 完全确定性 DDIM
6
    eta = 1   → 等效 DDPM（随机）
7
    """
8
    T     = len(alphas_cumprod) - 1
9
    shape = (batch_size, channels, height, width)
10

11
    # 从纯噪声开始
12
    x = torch.randn(*shape, device=device)
13

14
    # 生成时间步序列（均匀或非线性）
15
    # 非线性调度: t 越接近 0，变化越慢
16
    ts = torch.linspace(1.0, 0.0, num_steps + 1)
17

18
    for i in range(num_steps):
19
        t     = int(ts[i] * T)          # 当前时间步
20
        t_prev = int(ts[i + 1] * T)     # 上一个时间步
21

22
        eps_theta = model(x, t)
23

24
        # sigma_t: 由 eta 控制
25
        # eta = 0 → 完全确定
26
        # eta = 1 → 随机性等于 DDPM
27
        sigma_t = eta * ((1 - alphas_cumprod[t_prev] / alphas_cumprod[t]) ** 0.5) \
28
                  * ((1 - alphas_cumprod[t] / alphas_cumprod[t_prev]) ** 0.5)
29

30
        alpha_bar_t     = alphas_cumprod[t]
31
        alpha_bar_prev  = alphas_cumprod[t_prev]
32

33
        sqrt_ab         = alpha_bar_t ** 0.5
34
        sqrt_1m         = (1 - alpha_bar_t) ** 0.5
35
        sqrt_ab_prev    = alpha_bar_prev ** 0.5
36
        sqrt_1m_prev    = (1 - alpha_bar_prev) ** 0.5
37

38
        # 估计 x_0
39
        x0_pred = (x - sqrt_1m * eps_theta) / sqrt_ab
40

41
        # 预测 x_{t-1}
42
        x_pred = sqrt_ab_prev * x0_pred + sqrt_1m_prev * eps_theta
43

44
        # 加入随机性（仅当 eta > 0）
45
        if sigma_t > 0:
46
            std = sigma_t * ((1 - alpha_bar_prev) ** 0.5)
47
            x_pred = x_pred + std * torch.randn_like(x)
48

49
        x = x_pred
50

51
    return x

展开收起

3.4 自适应步长求解器#

1
from scipy.integrate import solve_ivp
2

3

4
def ddim_neural_ode_sampler(model, rtol=1e-5, atol=1e-6):
5
    """
6
    用 scipy 的自适应 ODE 求解器（精度驱动的采样）。
7

8
    与固定步数相比:
9
    - 在低噪声区域(t→0)自动用更小步长
10
    - 在高噪声区域(t→T)自动用更大步长
11
    - 总步数可能更少，精度更高
12
    """
13
    dim = channels * height * width
14

15
    def ode_function(t_float, x_flat):
16
        """
17
        t_float ∈ [0, 1]: 连续时间（0=数据，1=噪声）
18
        这里用反向时间，t=0 是噪声，t=1 是数据
19
        """
20
        t_int = int((1.0 - t_float) * (T - 1))
21
        t_int = max(0, min(T - 1, t_int))
22

23
        x = torch.tensor(x_flat.reshape(1, channels, height, width),
24
                         device=device, dtype=torch.float32)
25
        eps = model(x, t_int)
26

27
        alpha_bar     = alphas_cumprod[t_int]
28
        sqrt_ab       = alpha_bar ** 0.5
29
        sqrt_1m       = (1 - alpha_bar) ** 0.5
30
        x0_pred       = (x - sqrt_1m * eps) / sqrt_ab
31

32
        # dx/dt = -(预测方向) / T
33
        # 近似: 速度场方向
34
        dx_dt = -(eps - x0_pred).cpu().numpy().flatten() / T
35
        return dx_dt
36

37
    x0 = torch.randn(1, channels, height, width).cpu().numpy().flatten()
38
    sol = solve_ivp(
39
        ode_function,
40
        t_span=(0.0, 1.0),
41
        y0=x0,
42
        method='RK45',
43
        rtol=rtol, atol=atol,
44
        dense_output=True,
45
    )
46
    final = sol.y[:, -1].reshape(channels, height, width)
47
    return torch.from_numpy(final).to(device).unsqueeze(0)

展开收起

4. 与 DDPM 的关系：严格的数学联系#

4.1 训练目标完全相同#

最重要的事实：DDIM 的训练目标与 DDPM 完全相同！

L=Et,x0,ϵ∥ϵ−ϵθ(xt,t)∥2\mathcal{L} = \mathbb{E}_{t, x_0, \epsilon} \left\| \epsilon - \epsilon_\theta(x_t, t) \right\|^2L=Et,x0,ϵ∥ϵ−ϵθ(xt,t)∥2

这意味着：

不需要重新训练 DDPM 模型——已有的 DDPM checkpoint 可以直接用于 DDIM
DDIM 只是不同的采样策略，不是新的训练方法

4.2 采样分布的数学联系#

设 qσ(xt−1∣xt,x0)q_\sigma(x_{t-1}|x_t, x_0)qσ(xt−1∣xt,x0) 是 DDIM 定义的逆向分布，q(xt−1∣xt,x0)q(x_{t-1}|x_t, x_0)q(xt−1∣xt,x0) 是 DDPM 的真实逆向分布。

定理：对任意 σt\sigma_tσt，DDIM 的边际分布与 DDPM 的边际分布相同：

∫qσ(xt−1∣xt,x0)q(xt∣x0)dx0=q(xt−1∣x0)\int q_\sigma(x_{t-1}|x_t, x_0) q(x_t|x_0) dx_0 = q(x_{t-1}|x_0)∫qσ(xt−1∣xt,x0)q(xt∣x0)dx0=q(xt−1∣x0)

这就是 DDIM 的”隐式”特性——不显式定义概率密度函数，但边际分布自动保持一致。

4.3 采样轨迹的几何对比#

1
采样轨迹对比 (x_0 → x_T 方向为正向)
2

3
DDPM (随机马尔可夫):
4
    x_0 ──→ x_1 ──→ x_2 ──→ ... ──→ x_T
5
             ↑        ↑               ↑
6
           随机     随机             随机
7
           发散     发散             发散
8
    → 每步独立随机偏移
9
    → 步数少时轨迹高度发散
10

11
DDIM (确定性非马尔可夫):
12
    x_0 ────────→ x_{S_1}
13
        ↑                   ↑ 跳步
14
    x_{S_2} ←───────────────
15
        ↑ 跳步
16
    x_{S_3} ←───────────────
17
    ...
18
    x_T = x_{S_k}
19

20
    → 确定性轨迹，每步连续
21
    → 跳步后仍沿同一条光滑曲线
22
    → 步数减少时轨迹仍光滑

展开收起

4.4 为什么确定性更好？#

直觉解释：

1
DDPM 每步: x_{t-1} = μ + 随机噪声
2
           ↑ 每个时间步的随机偏移是独立的
3
           → 走 N 步后，总偏移 ~ √N × 单步偏移（随机游走）
4
           → 步数减半，偏移翻倍，误差爆炸
5

6
DDIM 每步: x_{t-1} = F(x_t)  (确定性函数)
7
           ↑ 跳步仍沿同一光滑 ODE 曲线
8
           → 走 N 步 = 直接求 ODE 的 N 个点
9
           → 步数减半，只是不采样中间点，不影响精度
10
           → 误差可预测（由 ODE 求解器控制）

5. 采样调度策略#

5.1 均匀时间步调度#

最简单的策略：将 [0,T][0, T][0,T] 均匀分割：

1
def uniform_timestep_schedule(num_steps, T=1000):
2
    """均匀调度。"""
3
    return torch.linspace(T - 1, 0, num_steps).long().tolist()

5.2 DDIM 论文推荐的调度#

论文发现，噪声强度越大（ttt 越接近 TTT），ODE 变化越剧烈。因此建议用非线性调度，在高噪声区域用更多步：

1
def ddim_timestep_schedule(num_steps, T=1000, eta=0.0):
2
    """
3
    DDIM 论文推荐的时间步调度。
4

5
    核心思想:
6
    - 高噪声区域 (t > 100): 变化剧烈，多放步
7
    - 低噪声区域 (t < 100): 变化平缓，少放步
8
    """
9
    # 从 [1, T] 均匀选 num_steps 个
10
    steps = torch.arange(0, T, T // num_steps).flip(0)
11

12
    # 如果步数不够 T，调整最后一步
13
    if len(steps) < num_steps:
14
        steps = torch.cat([steps, torch.tensor([0])])
15

16
    return steps[:num_steps].long().tolist()

展开收起

5.3 重要性采样：跳过 vs 插值#

DDIM 支持两种子序列策略：

策略	方法	效果
跳步 (skip)	跳过中间时间步，只采样子序列	更直接，减少步数
插值 (interpolate)	在相邻时间步之间插值	平滑过渡，更稳定

1
def get_ddim_schedule(num_steps, T=1000, method="skip"):
2
    """
3
    生成 DDIM 子序列。
4

5
    skip:       [999, 949, 899, ..., 49, 0]
6
    interpolate: [999, 997, 995, ..., 5, 3, 1]
7
    """
8
    if method == "skip":
9
        # 每 T/num_steps 步跳一次
10
        step_size = T // num_steps
11
        schedule = list(range(T - 1, -1, -step_size))
12
    elif method == "interpolate":
13
        # 均匀分布 + 线性插值
14
        uniform = torch.linspace(T - 1, 0, num_steps)
15
        schedule = uniform.round().long().tolist()
16

17
    return schedule

展开收起

5.4 噪声调度的敏感性#

1
def sensitivity_analysis():
2
    """
3
    不同 eta 和 num_steps 的 FID 表现（经验值）。
4
    来自 DDIM 论文实验。
5
    """
6
    # num_steps=100 时
7
    configs = [
8
        (100, 0.0, "DDIM (完全确定)"),   # FID ~5.8
9
        (100, 0.5, "DDIM (部分随机)"),  # FID ~6.2
10
        (100, 1.0, "DDPM (完全随机)"),  # FID ~20.0+
11
    ]
12

13
    # num_steps=20 时
14
    configs_20 = [
15
        (20, 0.0, "DDIM (完全确定)"),   # FID ~7.5
16
        (20, 0.5, "DDIM (部分随机)"),  # FID ~12.0
17
        (20, 1.0, "DDPM (完全随机)"),  # FID ~100+ (几乎无法用)
18
    ]
19

20
    return configs, configs_20

展开收起

关键发现：η=0\eta = 0η=0（完全确定）在所有步数下都优于 η=1\eta = 1η=1（完全随机）。这是 DDIM 最重要的实践结论。

6. 轨迹坍缩与多样性#

6.1 轨迹坍缩现象#

DDIM 的 η=0\eta = 0η=0（完全确定性）有一个副作用：多样性降低。

1
多样性 vs 确定性 trade-off:
2

3
高多样性 (η=1, DDPM):
4
    提示: "a cat on a sofa"
5
    生成: [猫1], [猫2], [猫3], ... [猫N]
6
    每张猫都不同（颜色、姿势、背景）
7

8
低多样性 (η=0, DDIM):
9
    提示: "a cat on a sofa"
10
    生成: [猫A], [猫A], [猫A], ... [猫A]
11
    几乎完全相同（只有微小随机性来自初始噪声）
12

13
原因: 确定性映射 f_θ: x_t → x_{t-1} 是固定的
14
     相同的初始噪声 x_T → 相同的生成结果
15
     初始噪声的多样性 = 生成的多样性

展开收起

6.2 数学解释#

随机 DDPM 的多样性来源：

每一步逆向过程都加随机噪声：

xt−1=μt+σt⋅ϵtx_{t-1} = \mu_t + \sigma_t \cdot \epsilon_txt−1=μt+σt⋅ϵt

这意味着 xt−1x_{t-1}xt−1 的分布是”展宽”的（convolution），引入额外随机性。

确定性 DDIM 的多样性来源：

无随机噪声：

xt−1=μtdetx_{t-1} = \mu_t^{\text{det}}xt−1=μtdet

只有初始噪声 xT∼N(0,I)x_T \sim \mathcal{N}(0, I)xT∼N(0,I) 引入多样性。

6.3 解决方案：调整 η\etaη#

1
def trade_off_eta():
2
    """
3
    eta 对多样性-质量 trade-off 的影响。
4

5
    η = 0.0:  质量最高，多样性最低（轨迹坍缩）
6
    η = 0.3:  质量高，多样性适中（推荐默认值）
7
    η = 0.5:  质量中等，多样性中等
8
    η = 1.0:  质量低（步数少时），多样性最高
9
    """
10
    pass

6.4 不同任务的最佳 η\etaη#

任务	推荐 η\etaη	理由
高保真生成	0.00.00.0	质量优先，不在乎多样性
多样化生成	0.3∼0.50.3 \sim 0.50.3∼0.5	平衡质量和多样性
图像编辑	0.00.00.0	确定性保证编辑一致性
Inpainting	0.00.00.0	需要精确控制
视频生成	0.0∼0.20.0 \sim 0.20.0∼0.2	时间一致性优先

7. 图像生成实验结果#

7.1 步数 vs 质量#

1
def cifar10_results():
2
    """
3
    CIFAR-10 32x32 上的 FID vs 采样步数。
4
    """
5
    results = {
6
        # (num_steps, eta): FID
7
        (1000, 1.0): 3.85,   # DDPM baseline
8
        (100,  0.0): 4.20,   # DDIM 100步确定
9
        (50,   0.0): 4.60,   # DDIM 50步确定
10
        (20,   0.0): 5.90,   # DDIM 20步确定
11
        (10,   0.0): 8.10,   # DDIM 10步确定
12
        (5,    0.0): 15.20,  # DDIM 5步确定
13
        (100,  1.0): 25.50,  # DDPM 100步随机（很差！）
14
        (50,   1.0): 55.00,  # DDPM 50步随机（几乎失败）
15
    }
16
    return results
17

18

19
def imagenet_results():
20
    """
21
    ImageNet 256x256 上的 FID vs 采样步数。
22
    """
23
    results = {
24
        (1000, 1.0, "ADM-Improved"): 10.4,
25
        (100,  0.0, "DDIM"):         11.3,
26
        (50,   0.0, "DDIM"):         12.8,
27
        (20,   0.0, "DDIM"):         17.2,
28
    }
29
    return results

展开收起

核心数据：DDIM 50 步（η=0\eta=0η=0）的 FID 与 DDPM 1000 步相当，而速度快 20 倍。

7.2 感知质量分析#

1
def perceptual_quality():
2
    """
3
    DDIM 的感知质量特点。
4
    """
5
    observations = {
6
        "优点": [
7
            "细节更锐利（确定性路径避免随机模糊）",
8
            "颜色更一致（无每步随机偏移累积）",
9
            "构图更稳定（适合编辑/控制任务）",
10
        ],
11
        "缺点": [
12
            "多样性降低（相同初始噪声产生相同结果）",
13
            "部分图像出现"超锐化"伪影（尤其 10 步以内）",
14
            "某些类别容易陷入局部最优",
15
        ]
16
    }
17
    return observations

展开收起

8. 加速技术的演进#

8.1 DDIM 的后续加速方法#

1
DDIM (2021): 非马尔可夫, 20-50步, 确定性
2

3
    ├── LDM (2022): Latent Diffusion → 压缩到隐空间，速度再次提升
4
    │
5
    ├── Consistency Models (2023): Song et al.
6
    │      一致性蒸馏 → 1-4 步
7
    │
8
    ├── LCM (Latent Consistency Models) (2023): Luo et al.
9
    │      蒸馏 + 隐空间 → 2-4 步 ★ (SD + LCM)
10
    │
11
    ├── SDXL-Turbo / SDXL-Lightning (2024)
12
    │      步蒸馏 → 1-2 步
13
    │
14
    └── Rectified Flow (2023): Liu et al.
15
           最优传输 → 4-10 步
16
           (见 [Flow Matching 文章](/posts/flow-matching/))

展开收起

8.2 为什么 DDIM 之后还需要新方法？#

维度	DDIM	新方法（LCM/RF）
采样步数	10-50 步	1-4 步
训练方式	同一训练目标	需要额外蒸馏
模型兼容性	复用 DDPM 模型	需要重新训练
生成质量	10步 FID ~8	4步 FID ~6

8.3 DDIM 的持续重要性#

尽管有更快的方法，DDIM 仍被广泛使用：

1
# Stable Diffusion 推理管线中的 DDIM
2
from diffusers import DDIMScheduler, StableDiffusionPipeline
3

4
pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
5
pipeline.scheduler = DDIMScheduler.from_config(pipeline.scheduler.config)
6

7
# 使用 DDIM 调度器
8
image = pipeline(
9
    prompt="a cat on a sofa",
10
    num_inference_steps=20,    # DDIM: 20步即可
11
    eta=0.0,                  # 完全确定性
12
).images[0]

9. 图像编辑与 DDIM#

9.1 DDIM 的编辑优势#

确定性是编辑任务的关键：

1
编辑任务: 把 [原图] 中的 [对象] 替换为 [新对象]
2

3
DDPM 方法:
4
    原图 → 加噪到某步 → 编辑条件 → 去噪
5
    问题: 加噪和去噪都是随机的
6
         同一张图多次处理会得到不同结果
7

8
DDIM 方法:
9
    原图 → 加噪到某步（确定性反向 ODE）→ 编辑条件 → 去噪
10
    优点: 加噪路径是确定性的
11
         同一张图→同一条轨迹→可逆
12
         编辑结果可复现

9.2 DDIM Inpainting#

1
def ddim_inpainting():
2
    """
3
    DDIM Inpainting 流程。
4

5
    1. 原图 x_0 → 加噪到 t=T-edit（如 T=500）
6
       (确定性反向 ODE, 可以精确恢复)
7
    2. 在隐空间中修改对应区域（inpaint mask）
8
    3. 从 x_t=500 用 DDIM 去噪到 x_0
9
    4. 合并编辑区域和原始区域
10
    """
11
    pass

9.3 引导编辑#

DDIM 与 CFG（无分类器引导）完美结合：

ϵ^=(1+w)⋅ϵθ(xt,t,c)−w⋅ϵθ(xt,t,∅)\hat{\epsilon} = (1 + w) \cdot \epsilon_\theta(x_t, t, c) - w \cdot \epsilon_\theta(x_t, t, \emptyset)ϵ^=(1+w)⋅ϵθ(xt,t,c)−w⋅ϵθ(xt,t,∅)

CFG 的引导强度 www 与 DDIM 的 η\etaη 可以独立调节：

www	效果	典型值
w=1w = 1w=1	无引导	默认
w=7w = 7w=7	中度引导	文本对齐好
w=15w = 15w=15	强度引导	高保真

10. 完整实现#

1
import torch
2
import torch.nn.functional as F
3

4

5
class DDIMScheduler:
6
    """
7
    DDIM 调度器实现。
8
    完全兼容 Diffusers 库的接口。
9
    """
10

11
    def __init__(self, num_train_timesteps=1000, beta_start=1e-4, beta_end=0.02,
12
                 beta_schedule="linear", clip_sample=False, set_alpha_to_one=False):
13
        self.num_train_timesteps = num_train_timesteps
14
        self.clip_sample = clip_sample
15
        self.final_alpha_cumprod = 1.0 if set_alpha_to_one else 0.0
16

17
        # 噪声调度
18
        if beta_schedule == "linear":
19
            betas = torch.linspace(beta_start, beta_end, num_train_timesteps)
20
        elif beta_schedule == "cosine":
21
            betas = self._cosine_beta_schedule(num_train_timesteps)
22
        else:
23
            raise ValueError(f"Unknown beta_schedule: {beta_schedule}")
24

25
        alphas = 1.0 - betas
26
        alphas_cumprod = torch.cumprod(alphas, dim=0)
27
        self.alphas_cumprod = alphas_cumprod
28

29
    def _cosine_beta_schedule(self, timesteps, s=0.008):
30
        """余弦噪声调度（更平滑）。"""
31
        steps = timesteps + 1
32
        x = torch.linspace(0, timesteps, steps)
33
        alphas_cumprod = torch.cos((x / timesteps + s) / (1 + s) * torch.pi * 0.5) ** 2
34
        alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
35
        betas = 1 - alphas_cumprod[1:] / alphas_cumprod[:-1]
36
        return torch.clamp(betas, 0.0001, 0.9999)
37

38
    def set_timesteps(self, num_inference_steps, offset=0):
39
        """
40
        设置推理时间步。
41

42
        参数:
43
            num_inference_steps: DDIM 采样步数（如 50）
44
            offset: 偏移量
45
        """
46
        self.num_inference_steps = num_inference_steps
47
        step_ratio = self.num_train_timesteps // num_inference_steps
48
        # 生成从 T-1 到 0 的时间步序列
49
        timesteps = (torch.arange(0, num_inference_steps) * step_ratio
50
                     + step_ratio // 2).flip(0)
51
        self.timesteps = timesteps.long()
52

53
    @torch.no_grad()
54
    def step(self, model_output, timestep, sample, eta=0.0):
55
        """
56
        DDIM 单步采样。
57

58
        参数:
59
            model_output:  模型预测的噪声 ε_θ(x_t, t)
60
            timestep:       当前时间步 t
61
            sample:         当前状态 x_t
62
            eta:            随机性参数 (0=确定, 1=随机)
63
        返回:
64
            prev_sample:    下一状态 x_{t-1}
65
        """
66
        t = timestep
67
        prev_t = timestep - self.num_train_timesteps // self.num_inference_steps
68

69
        # 当前和上一步的 alpha_cumprod
70
        alpha_prod_t     = self.alphas_cumprod[t]
71
        alpha_prod_t_prev = self.alphas_cumprod[prev_t] if prev_t >= 0 else self.final_alpha_cumprod
72

73
        # 预测 x_0
74
        pred_original_sample = (sample - (1 - alpha_prod_t).sqrt() * model_output) \
75
                               / alpha_prod_t.sqrt()
76
        if self.clip_sample:
77
            pred_original_sample = torch.clamp(pred_original_sample, -1, 1)
78

79
        # 计算方差
80
        # DDIM 方差 (公式来自论文)
81
        var = (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)
82
        std = (var ** 0.5) * eta
83

84
        # 预测 x_{t-1} 的方向
85
        pred_sample_direction = (1 - alpha_prod_t_prev - std ** 2).sqrt() * model_output
86

87
        # 确定性部分
88
        prev_sample = alpha_prod_t_prev.sqrt() * pred_original_sample + pred_sample_direction
89

90
        # 随机噪声（仅当 eta > 0）
91
        if eta > 0:
92
            noise = torch.randn_like(sample)
93
            variance_noise = std * noise
94
            prev_sample = prev_sample + variance_noise
95

96
        return prev_sample
97

98
    def add_noise(self, original_samples, timesteps, noise=None):
99
        """前向加噪（用于重建目标）。"""
100
        if noise is None:
101
            noise = torch.randn_like(original_samples)
102

103
        sqrt_alpha_prod     = self.alphas_cumprod[timesteps] ** 0.5
104
        sqrt_one_minus_alpha = (1 - self.alphas_cumprod[timesteps]) ** 0.5
105

106
        noisy_samples = sqrt_alpha_prod.view(-1, 1, 1, 1) * original_samples \
107
                       + sqrt_one_minus_alpha.view(-1, 1, 1, 1) * noise
108
        return noisy_samples

展开收起

11. 与 Flow Matching 的理论联系#

11.1 DDIM 的 ODE 就是 Flow Matching#

回想 Flow Matching 文章的核心：

Flow Matching 把从噪声到数据的路径建模为 ODE dx/dt=vθ(x,t)dx/dt = v_\theta(x, t)dx/dt=vθ(x,t)，训练目标是 MSE(预测速度, 真实速度)。

DDIM 的 σt=0\sigma_t = 0σt=0 情形正是 Flow Matching 的一个特例：

1
DDIM (η=0) 的 ODE:
2

3
x_{t-1} = √ᾱ_{t-1} · x̂_0 + √(1-ᾱ_{t-1}) · ε_θ(x_t, t)
4

5
将 t-1 写成 dt，并整理成 dx/dt 的形式：
6

7
dx/dt =  (ᾱ_{t-1} - ᾱ_t) / dt · x̂_0  + ...
8
       = v_θ(x_t, t)
9

10
其中 v_θ(x_t, t) = ε_θ(x_t, t) - x̂_0
11

12
而 x̂_0 = (x_t - √(1-ᾱ_t) · ε_θ) / √ᾱ_t
13

14
→ v_θ(x_t, t) = ε_θ - (x_t - √(1-ᾱ_t) · ε_θ) / √ᾱ_t
15
             = (1 - 1/√ᾱ_t) · ε_θ - x_t/√ᾱ_t
16

17
这恰好是 RF 直线路径的向量场！

展开收起

11.2 三种路径的理论统一#

1
所有扩散类生成模型 = Flow Matching 的特例
2

3
路径设计空间:
4
  ├── Rectified Flow:     直线 x_t = (1-t)x_0 + tε    → v = ε - x_0
5
  ├── DDPM (VP-SDE):      弧线 x_t = √ᾱ_t x_0 + √(1-ᾱ_t)ε
6
  └── DDIM (η=0):         与 DDPM 相同的路径，但 ODE 求解
7

8
采样策略:
9
  ├── DDPM:    SDE → 1000步随机马尔可夫
10
  ├── DDIM:    ODE → 10-50步确定性非马尔可夫
11
  └── RF:      ODE → 10-50步确定性直线

12. 总结#

12.1 核心要点#

维度	关键要点
核心洞察	逆向过程不需要是马尔可夫的，边际分布不变即可
数学框架	非马尔可夫链 + 边际分布约束
训练目标	与 DDPM 完全相同（不需要重新训练）
采样	σt=0\sigma_t = 0σt=0 时为确定性 ODE 采样
质量	50 步 DDIM ≈ 1000 步 DDPM（FID 相当）
速度	20 倍加速（1000 步 → 50 步）
权衡	η=0\eta = 0η=0 质量最高但多样性最低
编辑优势	确定性使加噪/去噪可逆，完美支持图像编辑
与 FM	DDIM η=0\eta=0η=0 的 ODE 路径是 Flow Matching 的特例

12.2 一句话总结#

DDIM 通过放弃”逆向必须是马尔可夫链”的人为约束，把 DDPM 的随机走路径变成确定性的 ODE 轨迹——不改变任何训练目标，只改变采样方式，就实现了 20 倍加速，同时保持了生成质量。这是非马尔可夫隐式概率模型在生成建模中的首次成功应用，也是 Rectified Flow 和一致性模型的理论先驱。

12.3 推荐资源#

1
论文:
2
  - DDIM (Song et al., 2021): "Denoising Diffusion Implicit Models"
3
  - Consistency Models (Song et al., 2023): "Consistency Models"
4
  - Flow Matching (Lipman et al., 2022): "Flow Matching for Causal Inference"
5

6
代码:
7
  - ermongroup/ddim (官方实现)
8
  - huggingface/diffusers (DDIMScheduler)
9
  - CompVis/stable-diffusion (集成 DDIM)
10

11
教程:
12
  - Lilian Weng 博客: "Diffusion Models as a Kind of ODE"
13
  - Hugging Face 博客: "Speed up inference with DDIM"

DDPM 深度剖析：变分推断视角下的去噪扩散概率模型

Thu, 16 Jul 2026 00:00:00 GMT

1. DDPM 的诞生：为什么是 2020 年？#

1.1 在 DDPM 之前：两段独立的历史#

扩散概率模型并非凭空出现——它由两条独立发展的理论脉络交汇而成。

历史脉络 A：变分推断

1
2002 VAE (Kingma & Welling)
2
  - 变分下界 (ELBO)
3
  - 重参数化技巧
4
  - 重构 + KL 正则化
5

6
2019 Sohl-Dickstein et al.
7
  - 扩散概率模型 (Deep Unsupervised Learning using Nonequilibrium Thermodynamics)
8
  - 马尔可夫前向 + 逆向
9
  - 但太慢，生成需 ~1000 步

历史脉络 B：Score Matching

1
2011 Score Matching (Hyvarinen)
2
  - 学 ∇_x log p(x) — 不需要归一化常数
3

4
2019 Song et al. — Noise Conditional Score Network (NCSN)
5
  - 用 Score Matching 学噪声尺度下的对数密度梯度
6
  - 朗之万采样生成
7

8
2021 Song et al. — Score SDE
9
  - 把 NCSN 统一到 SDE 框架

2020 年，Ho et al. 把两条脉络合二为一：

DDPM = Sohl-Dickstein 扩散的工程化 + Score Matching 的数学优雅 + 简单的 MSE 训练目标

1.2 一句话概括 DDPM#

DDPM 是一个分两步走的生成模型：① 预设一个将数据逐步破坏成纯噪声的前向过程（马尔可夫链）；② 训练一个神经网络去学习逆向过程——从噪声一步步恢复出原始数据。用变分推断证明，训练等价于简单的 MSE（预测噪声或 x0x_0x0）。

1.3 DDPM 论文信息#

1
论文: "Denoising Diffusion Probabilistic Models" (DDPM)
2
作者: Jonathan Ho, Ajay Jain, Pieter Abbeel
3
单位: Google Brain, UC Berkeley
4
发表于: NeurIPS 2020
5
引用: > 15,000 次 (截至 2024)
6
开源: https://github.com/hojonathanho/diffusion

2. 前向过程：预设的加噪马尔可夫链#

2.1 马尔可夫链定义#

设原始数据为 x0∼q(x0)x_0 \sim q(x_0)x0∼q(x0)（真实数据分布）。前向过程定义为一个 TTT 步马尔可夫链，每一步添加少量高斯噪声：

q(xt∣xt−1)=N(xt;1−βtxt−1,βtI)q(x_t | x_{t-1}) = \mathcal{N}\left(x_t; \sqrt{1 - \beta_t} x_{t-1}, \beta_t I\right)q(xt∣xt−1)=N(xt;1−βtxt−1,βtI)

其中 {βt}t=1T\{\beta_t\}_{t=1}^T{βt}t=1T 是噪声调度（noise schedule），满足 0<β1<β2<⋯<βT<10 < \beta_1 < \beta_2 < \cdots < \beta_T < 10<β1<β2<⋯<βT<1，βt\beta_tβt 随 ttt 递增。

1
def forward_step(xt_minus_1, beta_t):
2
    """前向过程的单步: x_{t-1} → x_t"""
3
    mean = torch.sqrt(1 - beta_t) * xt_minus_1
4
    std  = torch.sqrt(beta_t)
5
    noise = torch.randn_like(xt_minus_1)
6
    return mean + std * noise

2.2 关键性质：闭式分布#

不用逐步递推！ 通过重参数化技巧，任意时刻 ttt 的分布可以直接写出：

q(xt∣x0)=N(xt;αˉtx0, (1−αˉt)I)q(x_t | x_0) = \mathcal{N}\left(x_t; \sqrt{\bar\alpha_t} x_0, \; (1 - \bar\alpha_t) I\right)q(xt∣x0)=N(xt;αˉtx0,(1−αˉt)I)

其中：

αt=1−βt,αˉt=∏s=1tαs=∏s=1t(1−βs)\alpha_t = 1 - \beta_t, \quad \bar\alpha_t = \prod_{s=1}^t \alpha_s = \prod_{s=1}^t (1 - \beta_s)αt=1−βt,αˉt=s=1∏tαs=s=1∏t(1−βs)

推导：

设 ϵ1,ϵ2,…,ϵt∼N(0,I)\epsilon_1, \epsilon_2, \ldots, \epsilon_t \sim \mathcal{N}(0, I)ϵ1,ϵ2,…,ϵt∼N(0,I) 独立，则：

xt=αtxt−1+1−αtϵ1=αtαt−1xt−2+1−αtϵ1+αt(1−αt−1)ϵ2=⋯=αˉtx0+1−αˉtϵ\begin{aligned} x_t &= \sqrt{\alpha_t} x_{t-1} + \sqrt{1-\alpha_t} \epsilon_1 \\ &= \sqrt{\alpha_t \alpha_{t-1}} x_{t-2} + \sqrt{1-\alpha_t} \epsilon_1 + \sqrt{\alpha_t(1-\alpha_{t-1})} \epsilon_2 \\ &= \cdots \\ &= \sqrt{\bar\alpha_t} x_0 + \sqrt{1 - \bar\alpha_t} \epsilon \end{aligned}xt=αtxt−1+1−αtϵ1=αtαt−1xt−2+1−αtϵ1+αt(1−αt−1)ϵ2=⋯=αˉtx0+1−αˉtϵ

其中 ϵ∼N(0,I)\epsilon \sim \mathcal{N}(0, I)ϵ∼N(0,I) 是标准高斯噪声。

1
def closed_form_forward(x0, t, alphas_cumprod):
2
    """
3
    闭式前向分布: q(x_t | x_0)
4

5
    参数:
6
        x0: (batch_size, ...) 原始数据
7
        t:  (batch_size,) 时间步 (整数或浮点)
8
        alphas_cumprod: (T,)  cumprod 数组
9
    返回: (batch_size, ...) 在 t 时刻加噪后的数据
10
    """
11
    idx = t.long()  # 安全取整数索引
12
    sqrt_alpha_bar = torch.sqrt(alphas_cumprod[idx])
13
    sqrt_one_minus  = torch.sqrt(1 - alphas_cumprod[idx])
14
    noise = torch.randn_like(x0)
15
    B, C, H, W = x0.shape
16
    reshape_shape = (B, 1, 1, 1)
17
    return sqrt_alpha_bar.view(*reshape_shape) * x0 \
18
         + sqrt_one_minus.view(*reshape_shape) * noise

展开收起

2.3 噪声调度的直观理解#

αˉt\sqrt{\bar\alpha_t}αˉt 和 1−αˉt\sqrt{1-\bar\alpha_t}1−αˉt 随时间步变化：

t=0t = 0t=0：αˉt≈1.0,1−αˉt≈0\sqrt{\bar\alpha_t} \approx 1.0, \sqrt{1-\bar\alpha_t} \approx 0αˉt≈1.0,1−αˉt≈0，对应原始数据
t=500t = 500t=500：两者约各占一半，噪声和数据混合
t=Tt = Tt=T：αˉt≈0,1−αˉt≈1.0\sqrt{\bar\alpha_t} \approx 0, \sqrt{1-\bar\alpha_t} \approx 1.0αˉt≈0,1−αˉt≈1.0，接近纯噪声

时刻	信号占比	噪声占比	含义
t=0t = 0t=0	高	极低	原始数据
t=500t = 500t=500	中等	中等	噪声和数据各占一半
t=Tt = Tt=T	极低	高	接近纯噪声

核心洞察：当 t=Tt = Tt=T 时，αˉT≈0\bar\alpha_T \approx 0αˉT≈0，所以 xT≈N(0,I)x_T \approx \mathcal{N}(0, I)xT≈N(0,I)——前向过程把任意数据分布变成了标准高斯噪声。

2.4 常见噪声调度#

1
class NoiseSchedule:
2
    """三种常见噪声调度。"""
3

4
    @staticmethod
5
    def linear(T, beta_start=1e-4, beta_end=0.02):
6
        """线性调度 (DDPM 原始论文使用)。"""
7
        betas = torch.linspace(beta_start, beta_end, T)
8
        return betas
9

10
    @staticmethod
11
    def cosine(T, s=0.008):
12
        """余弦调度 (改进版, 噪声衰减更平滑)。"""
13
        steps = torch.arange(T + 1)
14
        # \bar\alpha_t = cos²( (t/T + s)/(1+s) * π/2 )
15
        alphas_cumprod = torch.cos((steps / T + s) / (1 + s) * (torch.pi / 2)) ** 2
16
        alphas_cumprod = alphas_cumprod / alphas_cumprod[0]  # 归一化: \bar\alpha_0 = 1
17
        betas = 1 - alphas_cumprod[1:] / alphas_cumprod[:-1]
18
        return betas
19

20
    @staticmethod
21
    def sigmoid(T, beta_start=1e-4, beta_end=0.02):
22
        """S 形调度 (介于线性和余弦之间)。"""
23
        betas = torch.linspace(-6, 6, T)
24
        betas = torch.sigmoid(betas) * (beta_end - beta_start) + beta_start
25
        return betas

展开收起

1
ᾱ_t 曲线对比:
2

3
1.0 ─────────────────────────────────────────
4
     ╲  Linear     (更快变噪声)
5
      ╲
6
       ╲ Cosine    (更平滑)
7
        ╲
8
         ╲
9
          ──────────────────────── 0.0  t

3. 逆向过程：学习的去噪分布#

3.1 逆向分布的必要性#

如果能精确知道 q(xt−1∣xt)q(x_{t-1} | x_t)q(xt−1∣xt)，就可以从纯噪声 xT∼N(0,I)x_T \sim \mathcal{N}(0, I)xT∼N(0,I) 精确恢复数据 x0x_0x0。

但 q(xt−1∣xt)q(x_{t-1} | x_t)q(xt−1∣xt) 不可直接计算（需要对所有 x0x_0x0 积分），所以用神经网络 pθ(xt−1∣xt)p_\theta(x_{t-1} | x_t)pθ(xt−1∣xt) 来近似：

pθ(xt−1∣xt)=N(xt−1;μθ(xt,t),σt2I)p_\theta(x_{t-1} | x_t) = \mathcal{N}\left(x_{t-1}; \mu_\theta(x_t, t), \sigma_t^2 I\right)pθ(xt−1∣xt)=N(xt−1;μθ(xt,t),σt2I)

3.2 逆向过程的参数化#

均值参数化（最直觉的写法）：

pθ(xt−1∣xt)=N(xt−1;μθ(xt,t),β~tI)p_\theta(x_{t-1} | x_t) = \mathcal{N}\left(x_{t-1}; \mu_\theta(x_t, t), \tilde\beta_t I\right)pθ(xt−1∣xt)=N(xt−1;μθ(xt,t),β~tI)

但更常用的是直接预测噪声或 x0x_0x0，而不是直接预测均值——因为这让目标函数成为简单的 MSE。

3.3 均值与噪声的关系#

由前向过程的闭式公式，可以反解出均值：

μt(xt,x0)=1αt(xt−1−αtϵ)\mu_t(x_t, x_0) = \frac{1}{\sqrt{\alpha_t}} \left( x_t - \sqrt{1 - \alpha_t} \epsilon \right)μt(xt,x0)=αt1(xt−1−αtϵ)

其中 ϵ\epsilonϵ 是注入的噪声。由 xt=αˉtx0+1−αˉtϵx_t = \sqrt{\bar\alpha_t} x_0 + \sqrt{1-\bar\alpha_t} \epsilonxt=αˉtx0+1−αˉtϵ，得：

x0=xt−1−αˉtϵαˉtx_0 = \frac{x_t - \sqrt{1-\bar\alpha_t} \epsilon}{\sqrt{\bar\alpha_t}}x0=αˉtxt−1−αˉtϵ

代入均值表达式：

μt(xt,x0)=1αt(xt−βt⋅xt−αˉtx01−αˉt)\mu_t(x_t, x_0) = \frac{1}{\sqrt{\alpha_t}} \left( x_t - \beta_t \cdot \frac{x_t - \sqrt{\bar\alpha_t} x_0}{\sqrt{1-\bar\alpha_t}} \right)μt(xt,x0)=αt1(xt−βt⋅1−αˉtxt−αˉtx0)

或者，直接用噪声 ϵ\epsilonϵ 和 x0x_0x0 来表达（这是神经网络实际预测的）：

μt(xt,t)=1αt(xt−βt1−αˉt⋅ϵθ(xt,t))\mu_t(x_t, t) = \frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{\beta_t}{\sqrt{1-\bar\alpha_t}} \cdot \epsilon_\theta(x_t, t) \right)μt(xt,t)=αt1(xt−1−αˉtβt⋅ϵθ(xt,t))

网络 ϵθ(xt,t)\epsilon_\theta(x_t, t)ϵθ(xt,t) 预测注入的噪声 ϵ\epsilonϵ，从而反推均值。

4. 变分推断与 ELBO#

4.1 生成模型的对数似然目标#

我们想最大化生成数据的对数似然 log⁡pθ(x0)\log p_\theta(x_0)logpθ(x0)。但 log⁡pθ(x0)\log p_\theta(x_0)logpθ(x0) 直接计算困难（需要对所有隐变量积分），所以用变分推断引入隐变量 x1:Tx_{1:T}x1:T：

log⁡pθ(x0)=log⁡∫pθ(x0:T) dx1:T\log p_\theta(x_0) = \log \int p_\theta(x_{0:T}) \, dx_{1:T}logpθ(x0)=log∫pθ(x0:T)dx1:T

4.2 ELBO 的引入#

对任意分布 q(x1:T∣x0)q(x_{1:T} | x_0)q(x1:T∣x0)（建议分布，即前向过程），用 Jensen 不等式得到变分下界 (ELBO)：

log⁡pθ(x0)≥Eq[log⁡pθ(x0:T)q(x1:T∣x0)]=L\log p_\theta(x_0) \geq \mathbb{E}_{q} \left[ \log \frac{p_\theta(x_{0:T})}{q(x_{1:T}|x_0)} \right] = \mathcal{L}logpθ(x0)≥Eq[logq(x1:T∣x0)pθ(x0:T)]=L

展开 pθ(x0:T)=p(xT)∏t=1Tpθ(xt−1∣xt)p_\theta(x_{0:T}) = p(x_T) \prod_{t=1}^T p_\theta(x_{t-1}|x_t)pθ(x0:T)=p(xT)∏t=1Tpθ(xt−1∣xt)，q(x1:T∣x0)=∏t=1Tq(xt∣xt−1)q(x_{1:T}|x_0) = \prod_{t=1}^T q(x_t|x_{t-1})q(x1:T∣x0)=∏t=1Tq(xt∣xt−1)，得到 ELBO 的分项形式：

L=−log⁡p(xT)⏟−LT+∑t=1TEq[DKL(q(xt−1∣xt,x0)∥pθ(xt−1∣xt))]⏟−Lt−1+Eq[log⁡pθ(x0∣x1)]⏟−L0\mathcal{L} = \underbrace{-\log p(x_T)}_{-\mathcal{L}_T} + \sum_{t=1}^T \underbrace{\mathbb{E}_q \left[ D_{\mathrm{KL}}(q(x_{t-1}|x_t, x_0) \| p_\theta(x_{t-1}|x_t)) \right]}_{-\mathcal{L}_{t-1}} + \underbrace{\mathbb{E}_q \left[ \log p_\theta(x_0 | x_1) \right]}_{-\mathcal{L}_0}L=−LT−logp(xT)+t=1∑T−Lt−1Eq[DKL(q(xt−1∣xt,x0)∥pθ(xt−1∣xt))]+−L0Eq[logpθ(x0∣x1)]

4.3 ELBO 三项的物理含义#

1
ELBO = -[ L_T + L_{T-1} + ... + L_1 + L_0 ]
2
        ↓       ↓         ↓       ↓
3
      终点    中间步    中间步   重构项
4
     KL项    KL项      KL项    重构损失

ELBO 包含三项，每项对应不同的训练阶段：

项	名称	含义	训练时是否需要
LT\mathcal{L}_TLT	终点先验 KL	q(xT∣x0)q(x_T \mid x_0)q(xT∣x0) 与 p(xT)p(x_T)p(xT) 的 KL，TTT 足够大时自动满足	❌
Lt−1\mathcal{L}_{t-1}Lt−1	去噪匹配项	q(xt−1∣xt,x0)q(x_{t-1} \mid x_t, x_0)q(xt−1∣xt,x0)（可解析）与 pθ(xt−1∣xt)p_\theta(x_{t-1} \mid x_t)pθ(xt−1∣xt) 的 KL	✅ 核心
L0\mathcal{L}_0L0	重构项	在 t=1t = 1t=1 时从 x1x_1x1 重构 x0x_0x0	✅（可与 L1\mathcal{L}_1L1 合并）

4.4 关键：q(xt−1∣xt,x0)q(x_{t-1} | x_t, x_0)q(xt−1∣xt,x0) 有闭式解#

定理：给定 x0x_0x0 和 xtx_txt，逆向条件分布 q(xt−1∣xt,x0)q(x_{t-1} | x_t, x_0)q(xt−1∣xt,x0) 也是高斯分布，可以直接写出均值和方差：

q(xt−1∣xt,x0)=N(xt−1;μ~t(xt,x0),β~tI)q(x_{t-1} | x_t, x_0) = \mathcal{N}\left(x_{t-1}; \tilde\mu_t(x_t, x_0), \tilde\beta_t I\right)q(xt−1∣xt,x0)=N(xt−1;μ~t(xt,x0),β~tI)

其中：

μ~t(xt,x0)=αˉt−1βt1−αˉtx0+αt(1−αˉt−1)1−αˉtxt\tilde\mu_t(x_t, x_0) = \frac{\sqrt{\bar\alpha_{t-1}} \beta_t}{1 - \bar\alpha_t} x_0 + \frac{\sqrt{\alpha_t}(1 - \bar\alpha_{t-1})}{1 - \bar\alpha_t} x_tμ~t(xt,x0)=1−αˉtαˉt−1βtx0+1−αˉtαt(1−αˉt−1)xt

方差：

β~t=1−αˉt−11−αˉt⋅βt\tilde\beta_t = \frac{1 - \bar\alpha_{t-1}}{1 - \bar\alpha_t} \cdot \beta_tβ~t=1−αˉt1−αˉt−1⋅βt

1
def reverse_posterior_mean(x0, xt, t, alphas_cumprod, betas):
2
    """
3
    计算 q(x_{t-1} | x_t, x_0) 的均值 (闭式)。
4
    这是在 ELBO 中，用于计算 L_{t-1} KL 项的 target mean。
5
    """
6
    alpha_bar_t   = alphas_cumprod[t]
7
    alpha_bar_tm1 = alphas_cumprod[t - 1]
8
    alpha_t       = 1 - betas[t]
9

10
    # 从 xt 解出注入的噪声
11
    sqrt_alpha_bar_t   = torch.sqrt(alpha_bar_t)
12
    sqrt_one_minus     = torch.sqrt(1 - alpha_bar_t)
13
    eps                = (xt - sqrt_alpha_bar_t * x0) / sqrt_one_minus
14

15
    # 计算均值 (用噪声表示)
16
    sqrt_alpha_bar_tm1 = torch.sqrt(alpha_bar_tm1)
17
    beta_t             = betas[t]
18

19
    # 方法 1: 用 x0 和 xt 表示
20
    mean = (sqrt_alpha_bar_tm1 * beta_t / (1 - alpha_bar_t)) * x0 \
21
         + (torch.sqrt(alpha_t) * (1 - alpha_bar_tm1) / (1 - alpha_bar_t)) * xt
22

23
    # 方法 2: 用噪声 ε 表示 (更常用, 神经网络预测 ε)
24
    mean_eps = (xt - torch.sqrt(1 - alpha_bar_t) * beta_t / (1 - alpha_bar_t).sqrt() * eps) \
25
             / alpha_t.sqrt()
26

27
    return mean_eps

展开收起

4.5 ELBO 的简化#

在 DDPM 论文中，Lt−1\mathcal{L}_{t-1}Lt−1 简化为：

Lt−1≈Et,x0,ϵ[βt22σt2(1−αt)(1−αˉt)∥ϵ−ϵθ(αˉtx0+1−αˉtϵ,t)∥2]\mathcal{L}_{t-1} \approx \mathbb{E}_{t, x_0, \epsilon} \left[ \frac{\beta_t^2}{2 \sigma_t^2 (1-\alpha_t)(1-\bar\alpha_t)} \left\| \epsilon - \epsilon_\theta(\sqrt{\bar\alpha_t}x_0 + \sqrt{1-\bar\alpha_t}\epsilon, t) \right\|^2 \right]Lt−1≈Et,x0,ϵ[2σt2(1−αt)(1−αˉt)βt2ϵ−ϵθ(αˉtx0+1−αˉtϵ,t)2]

论文发现，去掉复杂的系数后，简化为：

Lt−1simple=Et,x0,ϵ∥ϵ−ϵθ(xt,t)∥2\mathcal{L}_{t-1}^{\text{simple}} = \mathbb{E}_{t, x_0, \epsilon} \left\| \epsilon - \epsilon_\theta(x_t, t) \right\|^2Lt−1simple=Et,x0,ϵ∥ϵ−ϵθ(xt,t)∥2

这意味着：训练 DDPM 等价于训练一个 MSE 损失——预测注入的噪声 ϵ\epsilonϵ！

5. 三种参数化方式#

5.1 噪声预测 ϵ\epsilonϵ-prediction#

网络输出：注入的噪声 ϵ\epsilonϵ 目标：ϵθ(xt,t)≈ϵ\epsilon_\theta(x_t, t) \approx \epsilonϵθ(xt,t)≈ϵ

Lsimple=Et,x0,ϵ[∥ϵ−ϵθ(xt,t)∥2]\mathcal{L}_{\text{simple}} = \mathbb{E}_{t, x_0, \epsilon} \left[ \| \epsilon - \epsilon_\theta(x_t, t) \|^2 \right]Lsimple=Et,x0,ϵ[∥ϵ−ϵθ(xt,t)∥2]

1
def loss_eps_prediction(model, x0, t, alphas_cumprod):
2
    """DDPM 标准噪声预测损失。"""
3
    eps = torch.randn_like(x0)
4
    xt = torch.sqrt(alphas_cumprod[t]) * x0 + torch.sqrt(1 - alphas_cumprod[t]) * eps
5
    pred_eps = model(xt, t)
6
    return F.mse_loss(pred_eps, eps)

5.2 x0x_0x0 预测 (x-prediction)#

网络输出：原始数据 x0x_0x0 目标：x^0=xθ(xt,t)≈x0\hat{x}_0 = x_\theta(x_t, t) \approx x_0x^0=xθ(xt,t)≈x0

由 x0=(xt−1−αˉtϵ)/αˉtx_0 = (x_t - \sqrt{1-\bar\alpha_t} \epsilon) / \sqrt{\bar\alpha_t}x0=(xt−1−αˉtϵ)/αˉt：

x0=xt−1−αˉtϵθ(xt,t)αˉtx_0 = \frac{x_t - \sqrt{1 - \bar\alpha_t} \epsilon_\theta(x_t, t)}{\sqrt{\bar\alpha_t}}x0=αˉtxt−1−αˉtϵθ(xt,t)

1
def loss_x0_prediction(model, x0, t, alphas_cumprod):
2
    """x0 预测损失。"""
3
    eps = torch.randn_like(x0)
4
    xt = torch.sqrt(alphas_cumprod[t]) * x0 + torch.sqrt(1 - alphas_cumprod[t]) * eps
5
    pred_x0 = model(xt, t)
6
    return F.mse_loss(pred_x0, x0)

5.3 速度预测 v-prediction#

网络输出：速度向量 v=αtϵ−αˉtx0v = \alpha_t \epsilon - \sqrt{\bar\alpha_t} x_0v=αtϵ−αˉtx0 目标：vθ(xt,t)≈vv_\theta(x_t, t) \approx vvθ(xt,t)≈v

由 xt=αˉtx0+1−αˉtϵx_t = \sqrt{\bar\alpha_t} x_0 + \sqrt{1-\bar\alpha_t} \epsilonxt=αˉtx0+1−αˉtϵ，对 ttt 求导得：

v=dxtdt=−dαˉtdtx0+d1−αˉtdtϵv = \frac{dx_t}{dt} = -\frac{d\sqrt{\bar\alpha_t}}{dt} x_0 + \frac{d\sqrt{1-\bar\alpha_t}}{dt} \epsilonv=dtdxt=−dtdαˉtx0+dtd1−αˉtϵ

在实际中，DDPM++ 论文使用：

v=αˉtϵ−1−αˉtx0v = \sqrt{\bar\alpha_t} \epsilon - \sqrt{1-\bar\alpha_t} x_0v=αˉtϵ−1−αˉtx0

1
def loss_v_prediction(model, x0, t, alphas_cumprod):
2
    """速度预测损失 (DDPM++ / ADM 采用)。"""
3
    eps = torch.randn_like(x0)
4
    sqrt_ab = torch.sqrt(alphas_cumprod[t])
5
    sqrt_1m = torch.sqrt(1 - alphas_cumprod[t])
6
    xt = sqrt_ab * x0 + sqrt_1m * eps
7

8
    # 真实速度
9
    v = sqrt_ab * eps - sqrt_1m * x0
10

11
    pred_v = model(xt, t)
12
    return F.mse_loss(pred_v, v)

5.4 三种参数化对比#

维度	噪声预测 (DDPM)	x0 预测	速度预测 (DDPM++)
网络输出	ϵ\epsilonϵ	x0x_0x0	vvv
损失函数	MSE	MSE	MSE
低噪声时 (t→0t \to 0t→0)	好	最好	好
高噪声时 (t→Tt \to Tt→T)	好	差	好
训练稳定性	高	中	高
典型应用	DDPM, SD 1/2/xl	—	ADM, DiT, MMDiT

6. 训练与采样完整流程#

6.1 训练算法#

1
def ddpm_train_step(model, x0, alphas_cumprod, betas):
2
    """
3
    DDPM 训练步 (Algorithm 1, Ho et al. 2020)。
4
    """
5
    B = x0.shape[0]
6

7
    # 1) 采样时间步 t ~ Uniform{1, ..., T}
8
    t = torch.randint(0, len(alphas_cumprod), (B,), device=x0.device)
9

10
    # 2) 采样噪声 ε ~ N(0, I)
11
    eps = torch.randn_like(x0)
12

13
    # 3) 前向加噪 (闭式)
14
    sqrt_ab  = alphas_cumprod[t] ** 0.5
15
    sqrt_1m  = (1 - alphas_cumprod[t]) ** 0.5
16
    xt = sqrt_ab.view(-1, *([1] * (x0.dim() - 1))) * x0 \
17
       + sqrt_1m.view(-1, *([1] * (x0.dim() - 1))) * eps
18

19
    # 4) 计算损失 (噪声预测)
20
    eps_theta = model(xt, t)
21
    loss = F.mse_loss(eps_theta, eps)
22

23
    return loss
24

25

26
def train_ddpm(model, dataloader, T, beta_schedule_fn):
27
    """完整 DDPM 训练循环。"""
28
    betas = beta_schedule_fn(T)
29
    alphas = 1 - betas
30
    alphas_cumprod = torch.cumprod(alphas, dim=0)
31

32
    for epoch in range(num_epochs):
33
        for batch in dataloader:
34
            x0 = batch.to(device)
35
            loss = ddpm_train_step(model, x0, alphas_cumprod, betas)
36

37
            optimizer.zero_grad()
38
            loss.backward()
39
            optimizer.step()

展开收起

6.2 采样算法（DDPM 原版）#

1
@torch.no_grad()
2
def ddpm_sample(model, alphas_cumprod, betas, shape, T):
3
    """
4
    DDPM 采样 (Algorithm 2, Ho et al. 2020)。
5
    从 x_T ~ N(0, I) 开始，T 步逆向去噪。
6
    """
7
    # 1) 从纯噪声开始
8
    xt = torch.randn(*shape, device=device)
9

10
    # 2) 逐步逆向去噪 (t = T, T-1, ..., 1)
11
    for t in reversed(range(T)):
12
        t_batch = torch.full((shape[0],), t, device=device)
13

14
        # 预测噪声
15
        eps_theta = model(xt, t_batch)
16

17
        # 估计 x_0
18
        sqrt_ab  = alphas_cumprod[t] ** 0.5
19
        sqrt_1m  = (1 - alphas_cumprod[t]) ** 0.5
20
        x0_pred  = (xt - sqrt_1m * eps_theta) / sqrt_ab
21

22
        # 计算均值
23
        beta_t    = betas[t]
24
        alpha_t   = 1 - beta_t
25
        mean = (xt - (beta_t / sqrt_1m) * eps_theta) / alpha_t ** 0.5
26

27
        # 如果不是最后一步，加随机噪声
28
        if t > 0:
29
            std = (1 - alphas_cumprod[t - 1] / alphas_cumprod[t]) ** 0.5 * beta_t ** 0.5
30
            xt = mean + std * torch.randn_like(xt)
31
        else:
32
            xt = mean
33

34
    return xt

展开收起

注意：原版 DDPM 需要 1000 步才能生成高质量图像——这太慢了，所以后来有了 DDIM 等加速方法。

6.3 DDIM 加速采样#

DDIM（Denoising Diffusion Implicit Models, Song et al., 2021）放弃了马尔可夫假设，用非马尔可夫逆向过程实现确定性且更快的采样：

1
@torch.no_grad()
2
def ddim_sample(model, alphas_cumprod, betas, shape,
3
                 num_steps=50, eta=0.0):
4
    """
5
    DDIM 采样 (Song et al., 2021)。
6
    num_steps << T (比如 T=1000, num_steps=50)
7

8
    eta = 0:   完全确定性 (DDIM)
9
    eta = 1:   随机马尔可夫 (等效 DDPM)
10
    """
11
    T = len(alphas_cumprod) - 1
12
    # 选择子序列
13
    step_list = torch.linspace(0, T - 1, num_steps, dtype=torch.long).tolist()
14
    seq = step_list[::-1]  # 倒序: [T-1, ..., 0]
15

16
    # 从纯噪声开始
17
    xt = torch.randn(*shape, device=device)
18

19
    for i, t in enumerate(seq):
20
        t_batch = torch.full((shape[0],), t, device=device)
21

22
        # 预测噪声
23
        eps_theta = model(xt, t_batch)
24

25
        # 估计 x_0
26
        sqrt_ab  = alphas_cumprod[t] ** 0.5
27
        sqrt_1m  = (1 - alphas_cumprod[t]) ** 0.5
28
        x0_pred  = (xt - sqrt_1m * eps_theta) / sqrt_ab
29

30
        # 下一个时间步
31
        t_next = seq[i + 1] if i + 1 < len(seq) else 0
32
        sqrt_ab_next = alphas_cumprod[t_next] ** 0.5
33
        sqrt_1m_next = (1 - alphas_cumprod[t_next]) ** 0.5
34

35
        # DDIM 确定性方向
36
        pred_eps = eps_theta
37
        xt_next = sqrt_ab_next * x0_pred + sqrt_1m_next * pred_eps
38

39
        # 如果要随机性 (eta > 0)
40
        if eta > 0:
41
            std = eta * ((1 - alphas_cumprod[t_next]) / (1 - alphas_cumprod[t]) * betas[t]) ** 0.5
42
            xt_next = xt_next + std * torch.randn_like(xt)
43

44
        xt = xt_next
45

46
    return xt

展开收起

DDIM 的关键洞察：t→t−1t \to t-1t→t−1 的转移不需要是随机的——只要 xt−1x_{t-1}xt−1 能从 xtx_txt 推出 x0x_0x0，就是合法的逆向过程。

6.4 DDPM vs DDIM vs Flow Matching 对比#

维度	DDPM (马尔可夫)	DDIM (非马尔可夫)	Rectified Flow
逆向过程	随机（加噪声）	确定性（可选随机性）	ODE（确定性）
采样步数	1000（慢）	20-50（快）	10-50（快）
轨迹形状	弧线	取决于噪声调度	直线（OT 最优）
训练目标	MSE(ϵ\epsilonϵ)	同 DDPM	MSE(vvv)
轨迹可视化	发散	集中	最短

7. 与 Score Matching 的理论联系#

7.1 Score Matching 的核心#

Score Matching 的目标是训练一个得分网络 sθ(x,t)s_\theta(x, t)sθ(x,t) 来估计对数密度梯度：

sθ(x,t)≈∇xlog⁡pt(x)s_\theta(x, t) \approx \nabla_x \log p_t(x)sθ(x,t)≈∇xlogpt(x)

得分匹配损失（Hutchinson 估计器）：

LSM=Ex∼pdata[tr⁡(∂sθ(x)∂x)+12∥sθ(x)∥2]\mathcal{L}_{\text{SM}} = \mathbb{E}_{x \sim p_{\text{data}}} \left[ \operatorname{tr}\left( \frac{\partial s_\theta(x)}{\partial x} \right) + \frac{1}{2} \| s_\theta(x) \|^2 \right]LSM=Ex∼pdata[tr(∂x∂sθ(x))+21∥sθ(x)∥2]

7.2 DDPM 与 Score Matching 的等价性#

核心联系：DDPM 预测的噪声 ϵθ(xt,t)\epsilon_\theta(x_t, t)ϵθ(xt,t) 与得分函数成正比：

ϵθ(xt,t)=−σt∇xtlog⁡pt(xt)\epsilon_\theta(x_t, t) = -\sigma_t \nabla_{x_t} \log p_t(x_t)ϵθ(xt,t)=−σt∇xtlogpt(xt)

其中 σt=1−αˉt\sigma_t = \sqrt{1 - \bar\alpha_t}σt=1−αˉt 是噪声标准差。

推导：

pt(xt)=∫p(x0)q(xt∣x0)dx0p_t(x_t) = \int p(x_0) q(x_t | x_0) dx_0pt(xt)=∫p(x0)q(xt∣x0)dx0
q(xt∣x0)=N(αˉtx0,σt2I)q(x_t | x_0) = \mathcal{N}(\sqrt{\bar\alpha_t} x_0, \sigma_t^2 I)q(xt∣x0)=N(αˉtx0,σt2I)
∇xtlog⁡q(xt∣x0)=−xt−αˉtx0σt2=−ϵσt\nabla_{x_t} \log q(x_t | x_0) = -\frac{x_t - \sqrt{\bar\alpha_t} x_0}{\sigma_t^2} = -\frac{\epsilon}{\sigma_t}∇xtlogq(xt∣x0)=−σt2xt−αˉtx0=−σtϵ
所以 ∇xtlog⁡pt(xt)=Eq(x0∣xt)[−ϵσt]\nabla_{x_t} \log p_t(x_t) = \mathbb{E}_{q(x_0|x_t)}[-\frac{\epsilon}{\sigma_t}]∇xtlogpt(xt)=Eq(x0∣xt)[−σtϵ]

1
def score_from_eps(eps_theta, t, alphas_cumprod):
2
    """
3
    从噪声预测网络得到得分函数 ∇_x log p_t(x)。
4
    """
5
    sigma_t = torch.sqrt(1 - alphas_cumprod[t])
6
    score = -eps_theta / sigma_t  # s_θ(x_t, t) ≈ -ε / σ
7
    return score
8

9

10
def eps_from_score(score, t, alphas_cumprod):
11
    """
12
    从得分函数恢复噪声预测。
13
    """
14
    sigma_t = torch.sqrt(1 - alphas_cumprod[t])
15
    return -sigma_t * score

展开收起

7.3 为什么 DDPM 比纯 Score Matching 更好训练？#

维度	纯 Score Matching	DDPM
学习目标	∇log⁡p(x)\nabla \log p(x)∇logp(x)（得分函数）	ϵ\epsilonϵ（噪声）
额外计算	需要 Hessian 或 Hutchinson 估计	不需要
采样方式	朗之万动力学（慢，不精确）	直接从 xTx_TxT 逆向（快）
训练稳定性	差（得分函数在高维空间容易爆炸）	好（MSE 损失有界）

DDPM 通过引入预设前向过程和可闭式计算的 KL 目标，绕过了纯 Score Matching 的 Hessian 估计难题——这是它最重要的工程贡献。

8. 方差参数化：固定 vs 学习#

8.1 DDPM 的方差设计#

原版 DDPM 将逆向方差 βt\beta_tβt 固定为前向方差：

σt2=βt\sigma_t^2 = \beta_tσt2=βt

论文发现，学习方差并不会带来显著提升，反而增加训练难度。

1
# DDPM 原版: 方差固定
2
def ddpm_variance(t, betas):
3
    """逆向过程的方差（固定为 β_t）。"""
4
    return betas[t]
5

6

7
# DDPM++ (改进版): 学习方差
8
class LearnedVarianceUNet(nn.Module):
9
    """输出均值 + log 方差。"""
10
    def forward(self, xt, t):
11
        out = self.backbone(xt, t)
12
        mean, log_var = out.chunk(2, dim=1)
13
        log_var = torch.clamp(log_var, -10, 10)
14
        return mean, log_var
15

16

17
def ddpmpp_sample(model, alphas_cumprod, betas, shape, T):
18
    """DDPM++ 采样 (学习方差版本)。"""
19
    xt = torch.randn(*shape, device=device)
20

21
    for t in reversed(range(T)):
22
        t_batch = torch.full((shape[0],), t, device=device)
23
        mean, log_var = model(xt, t_batch)
24

25
        # 混合: 固定 + 学习
26
        fixed_var = betas[t]
27
        learned_var = log_var.exp()
28
        var = (1 - alphas_cumprod[t - 1] / alphas_cumprod[t]) * betas[t]
29

30
        std = var ** 0.5
31
        xt = mean + std * torch.randn_like(xt)
32

33
    return xt

展开收起

9. Classifier Guidance（分类器引导）#

9.1 为什么需要引导？#

无条件 DDPM 生成质量不错，但无法控制类别。Classifier Guidance 用外部分类器注入类别信息：

p^ϕ(y∣xt)∝p(xt∣y)γ\hat{p}_\phi(y | x_t) \propto p(x_t | y)^{\gamma}p^ϕ(y∣xt)∝p(xt∣y)γ

对应的引导得分：

∇xtlog⁡p(y∣xt)^=γ⋅∇xtlog⁡pϕ(y∣xt)\widehat{\nabla_{x_t} \log p(y | x_t)} = \gamma \cdot \nabla_{x_t} \log p_\phi(y | x_t)∇xtlogp(y∣xt)=γ⋅∇xtlogpϕ(y∣xt)

9.2 引导下的逆向分布#

ϵ^θ(xt,t∣y)=ϵθ(xt,t)−1−αˉt⋅γ⋅∇xtlog⁡pϕ(y∣xt)\hat{\epsilon}_\theta(x_t, t | y) = \epsilon_\theta(x_t, t) - \sqrt{1 - \bar\alpha_t} \cdot \gamma \cdot \nabla_{x_t} \log p_\phi(y | x_t)ϵ^θ(xt,t∣y)=ϵθ(xt,t)−1−αˉt⋅γ⋅∇xtlogpϕ(y∣xt)

1
def classifier_guidance(eps_theta, x_t, t, classifier, y, gamma=1.0):
2
    """
3
    分类器引导 (Dhariwal & Nichol, 2021)。
4

5
    在无类别噪声预测上叠加分类器梯度。
6
    """
7
    # 分类器对 x_t 求梯度
8
    x_t.requires_grad_(True)
9
    logits = classifier(x_t, t)
10
    log_probs = torch.log_softmax(logits, dim=-1)
11
    target_log_prob = log_probs.gather(-1, y.unsqueeze(-1)).squeeze(-1)
12
    grad = torch.autograd.grad(target_log_prob.sum(), x_t)[0]
13

14
    # 引导噪声预测
15
    sigma_t = ...  # 噪声标准差
16
    guided_eps = eps_theta - gamma * sigma_t * grad
17
    return guided_eps

展开收起

关键：引导强度 γ\gammaγ 控制类别保真度和多样性之间的 trade-off。γ=1\gamma = 1γ=1 是标准引导，γ\gammaγ 越大，生成越精确但越缺乏多样性。

9.3 CFG vs Classifier Guidance#

维度	Classifier Guidance	Classifier-Free Guidance
是否需要	单独训练分类器	无需额外网络
训练方式	分类器单独训，DDPM 单独训	DDPM 联合训练条件 + 无条件
质量	高（专用信号）	高（隐式引导）
多样性	低（强引导降低）	可控（由 www 调节）
典型使用	ADM, Guided Diffusion	SD 1/2/xl, SD3, FLUX

**CFG（无分类器引导）**是主流：

ϵ^θ(xt,t∣c)=(1+w)⋅ϵθ(xt,t,c)−w⋅ϵθ(xt,t,∅)\hat{\epsilon}_\theta(x_t, t | c) = (1 + w) \cdot \epsilon_\theta(x_t, t, c) - w \cdot \epsilon_\theta(x_t, t, \emptyset)ϵ^θ(xt,t∣c)=(1+w)⋅ϵθ(xt,t,c)−w⋅ϵθ(xt,t,∅)

1
def cfg_eps(eps_cond, eps_uncond, w):
2
    """无分类器引导 (CFG)。"""
3
    return (1 + w) * eps_cond - w * eps_uncond

10. DDPM 的完整 PyTorch 实现#

1
import math
2
import torch
3
import torch.nn as nn
4
import torch.nn.functional as F
5
from torch.utils.data import DataLoader
6

7

8
class SinusoidalTimeEmbedding(nn.Module):
9
    """DDPM 标准时间步嵌入。"""
10

11
    def __init__(self, dim):
12
        super().__init__()
13
        self.dim = dim
14

15
    def forward(self, t):
16
        device = t.device
17
        half = self.dim // 2
18
        embeddings = math.log(10000) / (half - 1)
19
        embeddings = torch.exp(torch.arange(half, device=device) * -embeddings)
20
        embeddings = t[:, None] * embeddings[None, :]
21
        embeddings = torch.cat([embeddings.sin(), embeddings.cos()], dim=-1)
22
        return embeddings
23

24

25
class ResidualBlock(nn.Module):
26
    """带 GroupNorm + SiLU 的残差块。"""
27

28
    def __init__(self, in_ch, out_ch, time_emb_dim, dropout=0.1):
29
        super().__init__()
30
        self.norm1 = nn.GroupNorm(32, in_ch)
31
        self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)
32
        self.time_mlp = nn.Sequential(
33
            nn.SiLU(), nn.Linear(time_emb_dim, out_ch * 2)
34
        )
35
        self.norm2 = nn.GroupNorm(32, out_ch)
36
        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)
37
        self.dropout = nn.Dropout(dropout)
38

39
        # 残差连接
40
        self.shortcut = nn.Conv2d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity()
41

42
    def forward(self, x, t_emb):
43
        h = self.norm1(x).swish()
44
        h = self.conv1(h)
45

46
        # 调制: AdaGN 风格
47
        t = self.time_mlp(t_emb)
48
        shift, scale = t.chunk(2, dim=-1)
49
        h = h * (1 + scale[:, None, :, :]) + shift[:, None, :, :]
50

51
        h = self.dropout(h)
52
        h = self.norm2(h).swish()
53
        h = self.conv2(h)
54

55
        return h + self.shortcut(x)
56

57

58
class DDPMUNet(nn.Module):
59
    """简化版 DDPM U-Net。"""
60

61
    def __init__(self, in_channels=3, base_channels=128, channel_mults=(1, 2, 4, 8)):
62
        super().__init__()
63
        self.time_embed_dim = base_channels * 4
64
        self.time_mlp = nn.Sequential(
65
            SinusoidalTimeEmbedding(base_channels),
66
            nn.Linear(base_channels, self.time_embed_dim),
67
            nn.SiLU(),
68
            nn.Linear(self.time_embed_dim, self.time_embed_dim),
69
        )
70

71
        # Encoder
72
        chs = [base_channels]
73
        for mult in channel_mults:
74
            chs.append(base_channels * mult)
75

76
        self.encoder = nn.ModuleList()
77
        for i in range(len(chs) - 1):
78
            self.encoder.append(
79
                ResidualBlock(chs[i], chs[i + 1], self.time_embed_dim)
80
            )
81

82
        # Bottleneck
83
        self.bottleneck = ResidualBlock(chs[-1], chs[-1], self.time_embed_dim)
84

85
        # Decoder
86
        self.decoder = nn.ModuleList()
87
        for i in reversed(range(len(chs) - 1)):
88
            self.decoder.append(
89
                ResidualBlock(chs[i + 1], chs[i], self.time_embed_dim)
90
            )
91

92
        # 输出层
93
        self.final = nn.Sequential(
94
            nn.GroupNorm(32, base_channels),
95
            nn.SiLU(),
96
            nn.Conv2d(base_channels, in_channels, 3, padding=1),
97
        )
98

99
    def forward(self, xt, t):
100
        """返回噪声预测 ε_θ(x_t, t)。"""
101
        t_emb = self.time_mlp(t)
102

103
        # Encoder
104
        hs = []
105
        for block in self.encoder:
106
            xt = block(xt, t_emb)
107
            hs.append(xt)
108

109
        # Bottleneck
110
        xt = self.bottleneck(xt, t_emb)
111

112
        # Decoder with skip connections
113
        for block in self.decoder:
114
            xt = torch.cat([xt, hs.pop()], dim=1)
115
            xt = block(xt, t_emb)
116

117
        return self.final(xt)

展开收起

11. DDPM 的局限与后续发展#

11.1 DDPM 的三大局限#

局限	描述	影响
推理速度	需要 1000 步，生成极慢	实际应用受限
隐空间缺失	直接在像素空间加噪，计算量大	512×512 以上极慢
无条件生成	文本条件需要额外引导机制	CFG 引入计算开销

11.2 DDPM 之后的演进时间线#

1
2020 DDPM (Ho et al.) — MSE 训练, 1000 步推理
2
  │
3
2021 DDIM (Song et al.) — 非马尔可夫, 20-50 步
4
  │
5
2021 Classifier-free Guidance (Ho & Salimans)
6
  │
7
2022 Latent Diffusion (Rombach et al.) — VAE 压缩到隐空间 ★
8
  │
9
2022 Score SDE (Song et al.) — 统一 SDE 框架
10
  │
11
2022 Consistency Models (Song et al.) — 蒸馏到 1 步
12
  │
13
2022 DiT (Peebles & Xie) — Transformer 替换 UNet
14
  │
15
2023 SDXL (Rombach et al.) — 级联扩散, 1024px
16
  │
17
2024 Rectified Flow (Liu et al.) — 最优传输, 直线路径
18
  │
19
2024 SD3 (Esser et al.) — MMDiT 双流架构 ★
20
  │
21
2024 FLUX (Black Forest Labs) — RF + MMDiT, 12B

展开收起

12. 总结#

12.1 核心公式速查#

名称	公式
前向加噪	xt=αˉtx0+1−αˉtϵx_t = \sqrt{\bar\alpha_t} x_0 + \sqrt{1 - \bar\alpha_t} \epsilonxt=αˉtx0+1−αˉtϵ
逆向均值	μt=1αt(xt−βt1−αˉtϵθ)\mu_t = \frac{1}{\sqrt{\alpha_t}} \left(x_t - \frac{\beta_t}{\sqrt{1 - \bar\alpha_t}} \epsilon_\theta\right)μt=αt1(xt−1−αˉtβtϵθ)
训练目标	Et,x0,ϵ[ ∥ϵ−ϵθ(xt,t)∥2 ]\mathbb{E}_{t, x_0, \epsilon}[\, \\| \epsilon - \epsilon_\theta(x_t, t) \\|^2\,]Et,x0,ϵ[∥ϵ−ϵθ(xt,t)∥2]
DDIM 采样	xt−1=αˉt−1x^0+1−αˉt−1ϵθx_{t-1} = \sqrt{\bar\alpha_{t-1}} \hat{x}_0 + \sqrt{1 - \bar\alpha_{t-1}} \epsilon_\thetaxt−1=αˉt−1x^0+1−αˉt−1ϵθ
CFG	ϵ^=(1+w)ϵc−wϵu\hat\epsilon = (1 + w) \epsilon_c - w \epsilon_uϵ^=(1+w)ϵc−wϵu
Score 联系	ϵθ=−σt∇xtlog⁡pt(xt)\epsilon_\theta = -\sigma_t \nabla_{x_t} \log p_t(x_t)ϵθ=−σt∇xtlogpt(xt)

12.2 一句话总结#

DDPM 用”预设前向 + 变分推断 + 简单 MSE”的三角组合，绕过了纯 Score Matching 的计算难题，把扩散模型从理论可行的方法变成了实际可训练的生成模型——这是 2020 年最重要的生成建模突破之一，也是后续所有扩散改进（Latent Diffusion、DiT、MMDiT、Rectified Flow、FLUX）的起点。

12.3 推荐资源#

1
论文:
2
  - DDPM (Ho et al., 2020): "Denoising Diffusion Probabilistic Models"
3
  - DDIM (Song et al., 2021): "Denoising Diffusion Implicit Models"
4
  - Score SDE (Song et al., 2021): "Score-Based Generative Modeling through SDEs"
5
  - CFG (Ho & Salimans, 2022): "Classifier-Free Diffusion Guidance"
6
  - Latent Diffusion (Rombach et al., 2022): "High-Resolution Image Synthesis with Latent Diffusion Models"
7

8
代码:
9
  - hojonathanho/diffusion (官方实现)
10
  - openai/guided-diffusion (ADM)
11
  - CompVis/stable-diffusion (LDM)
12
  - lucidrains/DDPM-pytorch
13

14
数学参考:
15
  - Lilian Weng 博客: "What are Diffusion Models?" (最好的技术解读)
16
  - Angus Gilmour 博客: "Deriving the DDPM Training Objective"

展开收起

DeepSeek-R1 深度解析：推理能力涌现与纯 RL 的范式突破

Thu, 16 Jul 2026 00:00:00 GMT

1. 引言：2025 年 AI 领域的里程碑事件#

1.1 一个石破天惊的发布#

2025 年 1 月 20 日，DeepSeek 团队发布了三个重量级模型：

DeepSeek-R1-Zero：完全通过纯强化学习训练，没有任何人类标注的思维链数据
DeepSeek-R1：基于 R1-Zero 的技术，加入冷启动数据和多阶段训练
DeepSeek-R1-Distill-Qwen/Llama：用 R1 蒸馏的小模型系列（1.5B ~ 32B）

DeepSeek-R1 在数学（AIME 2024: 79.8%）、代码（LiveCodeBench: 65.9%）、推理（IFEVAL: 89%）上与 OpenAI-o1-1217 持平——参数量仅为其一半，成本为其三十分之一。

1.2 为什么这次发布如此震撼？#

1
传统认知（2024 年之前）：
2
  要让模型学会"推理"（输出长思维链）：
3
  1. 大量人工标注思维链数据（成本极高）
4
  2. 或者用 OpenAI o1 的闭源 API（无法自主控制）
5

6
DeepSeek-R1 的反直觉发现：
7
  → 不需要任何人类写的思维链数据
8
  → 纯 RL 训练可以让模型自动涌现出长思维链能力
9
  → 推理能力可以通过"试错-奖励"自学而来

1.3 本系列文章关联#

本文是 LLM 对齐系列的第四篇，建议阅读顺序：

顺序	文章	定位
1	RLHF 深度解析	对齐全景：PPO/GRPO/DPO 综述
2	PPO 深度解析	信赖域优化与 LLM 对齐
3	GRPO 深度解析	DeepSeek 的高效对齐算法
4	本文	R1 系列完整技术解析

2. 背景：DeepSeek-V3 Base#

2.1 架构创新#

DeepSeek-V3 是 R1 系列的基座模型，采用了多项技术创新：

1. MoE（混合专家）架构

1
class DeepSeekV3Config:
2
    """
3
    DeepSeek-V3 关键配置
4
    """
5
    num_parameters = "236B"          # 总参数
6
    num_experts = 256               # 专家总数
7
    num_active_experts = 8          # 每个 token 激活的专家数
8
    routing_type = "auxiliary-free"  # 无辅助损失负载均衡
9

10
    # → 每个 token 实际计算: 236B / 256 × 8 ≈ 7.4B FLOPs
11
    # → 训练效率比同等dense模型高 10 倍+

2. FP8 混合精度训练

首个在超大规模 MoE 模型上验证 FP8 训练可行性的工作
训练速度提升 2 倍，显存占用减少 40%

3. Multi-head Latent Attention (MLA)

DeepSeek 特有的注意力机制，通过低秩投影减少 KV Cache 显存：

MLP(x)=RMSNorm(x)+MLP(RMSNorm(x))\text{MLP}(x) = \text{RMSNorm}(x) + \text{MLP}\big(\text{RMSNorm}(x)\big)MLP(x)=RMSNorm(x)+MLP(RMSNorm(x))

1
class MultiHeadLatentAttention(nn.Module):
2
    """
3
    MLA: 将 KV 压缩到低秩 latent space
4
    """
5
    def __init__(self, d_model=7168, n_heads=128, n_kv_heads=128, latent_dim=512):
6
        super().__init__()
7
        # 压缩后的 latent
8
        self.q_proj = nn.Linear(d_model, d_model)
9
        self.kv_latent = nn.Linear(d_model, latent_dim)  # 压缩
10
        self.kv_output = nn.Linear(latent_dim, d_model * 2)
11

12
    def forward(self, x):
13
        q = self.q_proj(x)
14
        # 压缩的 KV
15
        kv = self.kv_latent(x)              # (B, L, latent_dim)
16
        k, v = self.kv_output(kv).chunk(2, dim=-1)
17
        # ...

展开收起

2.2 基座模型能力#

DeepSeek-V3 Base 在标准 benchmark 上的表现：

Benchmark	DeepSeek-V3 Base	GPT-4o	Claude-3.5
MMLU	87.1%	88.7%	88.3%
MATH	60.6%	58.0%	61.3%
GSM8K	89.0%	92.0%	88.7%
HumanEval	70.0%	70.0%	70.0%

Tip

DeepSeek-V3 Base 本身已是 SOTA 基座模型。这为后续 RL 训练提供了强大的”知识基础”——RL 是在这个基础上学会”如何使用知识”，而非”获取知识”。

3. DeepSeek-R1-Zero：纯 RL 的突破#

3.1 核心设计哲学#

DeepSeek-R1-Zero 的设计哲学可以概括为一句话：

“给模型足够强的基座能力和足够清晰的奖励信号，它会自己学会推理。”

三条核心设计：

纯 RL，无 SFT：直接用 RL 训练基座模型，跳过指令微调
规则奖励为主：数学/代码用规则验证，不用 Reward Model
无人类思维链数据：不教模型”如何思考”，只告诉它”什么是正确的”

3.2 训练流程#

1
DeepSeek-R1-Zero 训练流程：
2

3
DeepSeek-V3 Base (无 SFT)
4
        │
5
        ↓
6
    GRPO 训练
7
    ├─ 每条 prompt 采样 G=16 个回答
8
    ├─ 用规则奖励评分（答案对错）
9
    ├─ 组内归一化优势
10
    └─ PPO-style 更新
11
        │
12
        ↓
13
  训练 数千步...
14
        │
15
        ↓
16
  涌现能力自动出现：
17
    - 长思维链（数千 tokens）
18
    - 自我反思（"等等，让我重新检查..."）
19
    - 尝试多种方法
20
    - 回溯与修正

展开收起

3.3 规则奖励设计#

DeepSeek-R1-Zero 使用纯规则奖励，不依赖 Reward Model：

1
def compute_reward_ZERO(
2
    problem: dict,
3
    response: str,
4
    answer_format_check: bool = True,
5
) -> float:
6
    """
7
    DeepSeek-R1-Zero 的奖励函数。
8
    全部基于规则，无需训练。
9
    """
10
    total_reward = 0.0
11

12
    # === 奖励 1: 答案正确性 ===
13
    if problem["type"] == "math":
14
        ground_truth = extract_ground_truth(problem)
15
        model_answer = extract_final_answer(response)
16

17
        if model_answer is not None and model_answer == ground_truth:
18
            total_reward += 1.0
19
        else:
20
            total_reward -= 0.5  # 轻微惩罚错误答案
21

22
    elif problem["type"] == "code":
23
        # 执行测试用例
24
        test_result = execute_tests(response, problem["tests"])
25
        total_reward += test_result.pass_rate
26

27
    # === 奖励 2: 格式正确性 ===
28
    # 检查回答是否包含最终答案标记
29
    has_answer_tag = (
30
        "The answer is" in response or
31
        "答案是" in response or
32
        "\\boxed{" in response or
33
        "Final Answer:" in response
34
    )
35
    if has_answer_tag:
36
        total_reward += 0.1  # 鼓励明确标记答案
37

38
    # === 奖励 3: 思考过程 ===
39
    # 鼓励模型展示推理过程（而非直接给答案）
40
    has_thinking = (
41
        "Let me" in response or
42
        "First" in response or
43
        "Step" in response or
44
        "分析" in response or
45
        len(response) > 500  # 回答足够长
46
    )
47
    if has_thinking:
48
        total_reward += 0.05
49

50
    return total_reward
51

52

53
def extract_ground_truth(problem: dict) -> str:
54
    """
55
    从问题中提取标准答案。
56
    支持：
57
      - 数学表达式的数值结果
58
      - 证明题的逻辑验证
59
      - 选择题的正确答案
60
    """
61
    gt = problem.get("answer", "").strip()
62

63
    # 归一化比较
64
    gt_normalized = normalize_answer(gt)
65
    return gt_normalized

展开收起

3.4 思维链的涌现过程#

DeepSeek 团队记录了 R1-Zero 在训练过程中的能力演化：

1
训练步数 vs 行为变化（AIME 2024 pass@1）：
2

3
Step 0 (Base):
4
  表现: 26.7%
5
  行为: 直接输出答案，无推理过程
6
  输出示例: "答案是 42。"
7

8
Step 200:
9
  表现: 40.0%
10
  行为: 开始出现短推理（1-2 句）
11
  输出示例: "让我计算... 答案是 42。"
12

13
Step 500:
14
  表现: 52.3%
15
  行为: 推理变长（5-10 句），但结构散乱
16
  输出示例: "首先... 其次... 因此... 答案是 42。"
17

18
Step 1000:
19
  表现: 68.7%
20
  行为: 出现结构化推理，开始有"步骤感"
21
  输出示例: "Step 1: ... Step 2: ... 验证: ... 最终答案是 42。"
22

23
Step 2000 (关键涌现点):
24
  表现: 73.3%
25
  行为:
26
    - 出现自我反思："等等，我需要检查第一步"
27
    - 尝试多种方法："另一种思路是..."
28
    - 回溯修正："我发现之前的计算有误，修正后..."
29
  输出示例: "我先用方法 A... 得到 ... 但检查后发现不对。
30
             让我用方法 B... 验证后正确。
31
             因此答案是 42。"
32

33
Step 5000:
34
  表现: 79.8%
35
  行为:
36
    - 长推理链（2000+ tokens）
37
    - 主动验证中间步骤
38
    - 在困难问题上"停下来思考"
39
    - 复杂的反思与修正循环

展开收起

涌现的量化证据

DeepSeek 团队还发现了一个有趣的”顿悟时刻”（Aha Moment）：在某个训练步，模型突然学会了在回答前暂停并重新审视问题——这是思维链能力的质变，而非渐变。

3.5 Group Relative Response Award (GRPO) 变体#

DeepSeek-R1-Zero 使用的 GRPO 相比标准 GRPO 有几处关键改进：

1
def deepseek_r1_grpo_loss(
2
    policy_logps: torch.Tensor,     # (B,)
3
    ref_logps: torch.Tensor,        # (B,)
4
    rewards: torch.Tensor,          # (B,) 规则奖励
5
    group_size: int = 16,
6
    beta: float = 0.04,
7
    clip_ratio: float = 0.2,
8
    entropy_coef: float = 0.02,
9
    response_masks: torch.Tensor = None,  # 思维链部分 mask
10
) -> tuple[torch.Tensor, dict]:
11
    """
12
    DeepSeek-R1-Zero 使用的 GRPO 损失。
13
    与标准 GRPO 的区别：
14
      1. 使用 self-consistency 增强（多答案采样）
15
      2. 显式鼓励长回答（长回答通常包含更多思考）
16
      3. 更高的熵正则系数
17
    """
18
    num_groups = rewards.size(0) // group_size
19

20
    # === 1. 组内归一化优势 ===
21
    rewards = rewards.view(num_groups, group_size)
22
    mean_r = rewards.mean(dim=-1, keepdim=True)
23
    std_r = rewards.std(dim=-1, keepdim=True) + 1e-8
24
    advantages = (rewards - mean_r) / std_r  # (num_groups, G)
25

26
    # === 2. PPO-style Clip Loss ===
27
    # 展平处理
28
    advantages_flat = advantages.flatten()
29

30
    ratio = torch.exp(policy_logps - ref_logps)
31
    surr1 = ratio * advantages_flat
32
    surr2 = torch.clamp(ratio, 1 - clip_ratio, 1 + clip_ratio) * advantages_flat
33
    policy_loss = -torch.min(surr1, surr2).mean()
34

35
    # === 3. KL 散度 ===
36
    kl_penalty = (policy_logps - ref_logps).mean()
37

38
    # === 4. 长度鼓励（鼓励长思维链）===
39
    # 如果最终答案正确，奖励更长的思考过程
40
    with torch.no_grad():
41
        is_correct = rewards > 0.5  # 正确答案
42
        avg_length = response_masks.sum(dim=-1).float().mean() if response_masks is not None else 0
43

44
    # === 5. 熵正则 ===
45
    # R1-Zero 使用更高的熵系数防止坍缩
46
    entropy_loss = -entropy_coef * compute_response_entropy(policy_logps)
47

48
    total_loss = policy_loss + beta * kl_penalty + entropy_loss
49

50
    # 监控指标
51
    metrics = {
52
        "policy_loss": policy_loss.item(),
53
        "kl_penalty": kl_penalty.item(),
54
        "entropy": entropy_loss.item(),
55
        "mean_reward": rewards.mean().item(),
56
        "correct_rate": is_correct.float().mean().item(),
57
        "mean_length": avg_length.item(),
58
    }
59

60
    return total_loss, metrics

展开收起

4. DeepSeek-R1：完整的四阶段训练#

4.1 为什么需要四阶段？#

DeepSeek-R1-Zero 有一个明显的问题：回答可读性差，格式混乱。

1
R1-Zero 的典型输出：
2
"Let me think... 首先... 然后...
3
等等，我发现一个问题... 让我重新...
4
最终... 所以答案是...
5
但等等，我还需要... "
6

7
问题：
8
  1. 没有明确的格式（很难提取最终答案）
9
  2. 混合语言（中英文混杂）
10
  3. 自我反思过于频繁，显得"不自信"
11
  4. 不适合直接面向用户

DeepSeek-R1 通过四阶段训练解决这个问题：

4.2 完整训练流水线#

1
┌─────────────────────────────────────────────────────────────────┐
2
│                    DeepSeek-R1 四阶段训练                          │
3
├─────────────────────────────────────────────────────────────────┤
4
│                                                                  │
5
│  阶段 0: DeepSeek-V3 Base                                       │
6
│     236B 参数，MoE 架构，无 SFT                                  │
7
│              ↓                                                   │
8
│  阶段 1: GRPO for Reasoning (冷启动数据)                          │
9
│     ├─ 加入少量高质量冷启动数据（数千条）                          │
10
│     ├─ 规则奖励：数学/代码/逻辑                                   │
11
│     ├─ 输出: S1_Checkpoint                                       │
12
│     └─ 目的: 建立初步的推理格式                                   │
13
│              ↓                                                   │
14
│  阶段 2: Reinforced Fine-Tuning (RFT)                            │
15
│     ├─ 用 S1 生成 60 万条思维链数据                                │
16
│     ├─ SFT 微调                                                  │
17
│     ├─ 输出: S2_Checkpoint                                       │
18
│     └─ 目的: 将 GRPO 知识压缩为快速推理能力                        │
19
│              ↓                                                   │
20
│  阶段 3: 拒绝采样 + Supervised Fine-Tuning                        │
21
│     ├─ 用 S2 生成 80 万条多样化数据                               │
22
│     ├─ 拒绝采样：过滤低质量回答                                   │
23
│     ├─ 混合写作/角色扮演/指令等通用数据                           │
24
│     ├─ SFT 微调                                                  │
25
│     ├─ 输出: S3_Checkpoint                                       │
26
│     └─ 目的: 扩展能力边界到所有场景                               │
27
│              ↓                                                   │
28
│  阶段 4: GRPO for All Scenarios                                  │
29
│     ├─ 混合奖励：规则奖励 + RM + 安全                            │
30
│     ├─ 通用能力对齐                                              │
31
│     └─ 输出: DeepSeek-R1                                         │
32
│                                                                  │
33
└─────────────────────────────────────────────────────────────────┘

展开收起

4.3 阶段 1：冷启动数据#

目的：让模型学会”格式”，不教”推理”。

1
def generate_cold_start_data(
2
    num_samples: int = 5000,
3
    domain: str = "math",
4
) -> list[dict]:
5
    """
6
    生成冷启动数据。
7
    注意：只提供格式引导，不提供推理内容。
8
    """
9
    cold_start_prompts = [
10
        "你是一个擅长数学推理的助手。请在回答中包含："
11
        "【思考过程】...详细的推理步骤...\n"
12
        "【最终答案】...简洁的答案..."
13
    ]
14

15
    # 收集问题（带标准答案）
16
    problems = load_verifiable_problems(domain)
17

18
    # 用少量示例 + few-shot 引导格式
19
    data = []
20
    for problem in problems[:num_samples]:
21
        data.append({
22
            "prompt": cold_start_prompts[0] + "\n\n问题：" + problem["question"],
23
            "response": "",  # 让模型自己推理
24
            "format_hint": "请使用【思考过程】和【最终答案】的格式。",
25
        })
26

27
    return data

展开收起

冷启动数据的质量要求：

来源：精选的高质量数学问题（带详细解答）
格式：使用统一的模板（｜reasoning｜ 和 ｜answer｜ 分隔）
不包含推理内容：只给出格式要求，不示范推理过程

4.4 阶段 2：Reinforced Fine-Tuning (RFT)#

RFT 是 DeepSeek 的核心创新之一：用 RL 的探索知识做 SFT。

1
def rft_stage2(
2
    grpo_model,        # 阶段 1 输出的 GRPO 模型
3
    base_model,        # DeepSeek-V3 Base
4
    num_prompts: int = 600000,
5
    sft_epochs: int = 2,
6
):
7
    """
8
    RFT: Reinforced Fine-Tuning
9

10
    核心思想：
11
      GRPO 学到了"如何推理"
12
      RFT 把这个知识压缩成"快速推理"
13
    """
14
    # === 1. 生成高质量思维链数据 ===
15
    print("Generating reasoning chains...")
16
    reasoning_data = []
17

18
    prompts = sample_prompts(num_prompts)
19
    for batch in tqdm(prompts):
20
        # 用 GRPO 模型生成（允许探索）
21
        with torch.no_grad():
22
            responses = grpo_model.generate(
23
                batch,
24
                max_tokens=8192,
25
                temperature=0.7,
26
                do_sample=True,
27
            )
28

29
        for prompt, response in zip(batch, responses):
30
            # 评分（只保留正确答案）
31
            reward = rule_based_reward(prompt, response)
32
            if reward > 0.5:  # 正确答案
33
                reasoning_data.append({
34
                    "prompt": prompt,
35
                    "response": response,
36
                    "reward": reward,
37
                })
38

39
    print(f"Generated {len(reasoning_data)} high-quality reasoning samples")
40

41
    # === 2. SFT 微调 ===
42
    print("SFT fine-tuning...")
43
    sft_model = copy.deepcopy(base_model)
44
    optimizer = torch.optim.AdamW(sft_model.parameters(), lr=1e-5)
45

46
    dataloader = DataLoader(reasoning_data, batch_size=64, shuffle=True)
47

48
    for epoch in range(sft_epochs):
49
        for batch in dataloader:
50
            loss = compute_sft_loss(
51
                sft_model,
52
                batch["prompt"],
53
                batch["response"],
54
            )
55
            loss.backward()
56
            optimizer.step()
57
            optimizer.zero_grad()
58

59
    return sft_model

展开收起

RFT 的物理含义：

1
GRPO 训练过程（慢但探索充分）：
2
  prompt → [生成多个回答] → 评分 → 优化 → 学会推理
3
              ↑
4
        需要大量计算（rollout）
5
        但探索到了高质量的推理路径
6

7
RFT 过程（快但知识压缩）：
8
  GRPO 学会的推理路径 → SFT → 快速推理能力
9
              ↑
10
        不需要 rollout
11
        直接学习输出模式

4.5 阶段 3：拒绝采样与数据扩展#

阶段 3 解决两个问题：

扩展到所有能力：不仅是数学/代码，还有写作、角色扮演、总结等
拒绝采样：从大量生成中筛选高质量数据

1
def rejection_sampling(
2
    model,
3
    prompts: list[dict],
4
    num_generations: int = 100,
5
    reward_threshold: float = 0.8,
6
) -> list[dict]:
7
    """
8
    拒绝采样：从多次生成中选择最好的回答。
9
    """
10
    selected = []
11

12
    for prompt_data in prompts:
13
        prompt = prompt_data["prompt"]
14
        domain = prompt_data.get("domain", "general")
15

16
        # 生成多次
17
        candidates = []
18
        for _ in range(num_generations):
19
            with torch.no_grad():
20
                response = model.generate(prompt, max_tokens=4096)
21

22
            if domain in ["math", "code"]:
23
                # 规则奖励
24
                reward = rule_based_reward(prompt, response)
25
            else:
26
                # RM 奖励
27
                reward = reward_model.score(prompt, response)
28

29
            candidates.append({
30
                "prompt": prompt,
31
                "response": response,
32
                "reward": reward,
33
            })
34

35
        # 拒绝采样：只保留最高质量的
36
        best = max(candidates, key=lambda x: x["reward"])
37
        if best["reward"] >= reward_threshold:
38
            selected.append(best)
39

40
    return selected
41

42

43
def build_s3_dataset():
44
    """
45
    构建阶段 3 的训练数据。
46
    """
47
    dataset = []
48

49
    # 1. 数学/代码（拒绝采样）
50
    math_data = rejection_sampling(
51
        s2_model,
52
        load_math_prompts(200000),
53
        num_generations=100,
54
        reward_threshold=0.9,
55
    )
56
    dataset.extend(math_data)
57

58
    # 2. 通用写作/总结
59
    writing_data = load_filtered_data(
60
        source="general_writing",
61
        filter_fn=lambda x: x["quality_score"] > 0.85,
62
    )
63
    dataset.extend(writing_data)
64

65
    # 3. 角色扮演/对话
66
    chat_data = load_filtered_data(
67
        source="roleplay",
68
        filter_fn=lambda x: x["engagement_score"] > 0.8,
69
    )
70
    dataset.extend(chat_data)
71

72
    # 4. 指令遵循
73
    instruction_data = load_filtered_data(
74
        source="instruction",
75
        filter_fn=lambda x: x["ifeval_score"] > 0.85,
76
    )
77
    dataset.extend(instruction_data)
78

79
    return dataset
80
    # 总计约 80 万条

展开收起

4.6 阶段 4：全场景 GRPO#

阶段 4 用混合奖励进行最终的 GRPO 对齐：

1
def multi_reward(
2
    prompt: str,
3
    response: str,
4
    domain: str,
5
    reward_model=None,
6
) -> float:
7
    """
8
    混合奖励函数。
9
    """
10
    reward = 0.0
11

12
    # 1. 准确性奖励（最重要）
13
    if domain in ["math", "code", "logic"]:
14
        reward += 0.5 * rule_accuracy_reward(prompt, response)
15

16
    # 2. 帮助性奖励
17
    if reward_model is not None:
18
        reward += 0.3 * reward_model.score(prompt, response)
19

20
    # 3. 格式奖励
21
    reward += 0.1 * format_reward(response)
22

23
    # 4. 安全奖励
24
    reward += 0.1 * safety_reward(response)
25

26
    return reward
27

28

29
def format_reward(response: str) -> float:
30
    """
31
    格式奖励：鼓励清晰的结构。
32
    """
33
    score = 0.0
34

35
    if "｜reasoning｜" in response or "【思考过程】" in response:
36
        score += 0.3
37
    if "｜answer｜" in response or "【最终答案】" in response:
38
        score += 0.3
39
    if response.count("\n") > 3:  # 有清晰的段落
40
        score += 0.2
41
    if not response.startswith(" ") and not response.startswith("\n"):
42
        score += 0.2  # 不以空白开头
43

44
    return min(score, 1.0)
45

46

47
def safety_reward(response: str) -> float:
48
    """
49
    安全奖励：惩罚有害内容。
50
    """
51
    if contains_harmful_content(response):
52
        return -1.0
53
    return 0.0

展开收起

5. 知识蒸馏：小模型的大突破#

5.1 为什么需要蒸馏？#

DeepSeek-R1 太大（236B 参数），无法在消费级 GPU 上部署。

DeepSeek 的解决方案：用 R1 生成的数据蒸馏小模型。

5.2 蒸馏数据生成#

1
def generate_distillation_data(
2
    teacher_model,      # DeepSeek-R1
3
    student_prompts: list[str],
4
    num_samples: int = 1000000,
5
) -> list[dict]:
6
    """
7
    用 R1 生成蒸馏数据。
8
    """
9
    distillation_data = []
10

11
    for prompt in tqdm(student_prompts[:num_samples]):
12
        with torch.no_grad():
13
            response = teacher_model.generate(
14
                prompt,
15
                max_tokens=8192,
16
                temperature=0.6,
17
                do_sample=True,
18
            )
19

20
        distillation_data.append({
21
            "prompt": prompt,
22
            "response": response,
23
        })
24

25
    return distillation_data

展开收起

5.3 蒸馏结果#

模型	参数量	AIME 2024	MATH-500	GPQA	LiveCodeBench
DeepSeek-R1-Distill-Qwen-1.5B	1.5B	38.7%	83.6%	-	40.1%
DeepSeek-R1-Distill-Qwen-7B	7B	55.0%	89.1%	-	59.1%
DeepSeek-R1-Distill-Qwen-14B	14B	63.6%	90.8%	-	64.1%
DeepSeek-R1-Distill-Qwen-32B	32B	72.6%	92.9%	62.1%	70.0%
GPT-4o	-	52.0%	76.6%	40.3%	51.0%
Claude-3.5	-	49.0%	78.4%	38.4%	48.0%
OpenAI-o1-mini	-	63.6%	80.3%	-	58.0%

关键发现

32B 的蒸馏 Qwen 在 AIME 2024 上超越了 OpenAI-o1-mini（72.6% vs 63.6%），证明 R1 的知识可以通过蒸馏有效传递。

5.4 蒸馏 vs 从头训练#

1
1.5B 模型不同训练方式对比（AIME 2024）：
2

3
从头训练 + GRPO (无蒸馏):
4
  32.5%  ← 需要大量计算
5

6
直接蒸馏（用 R1 数据 SFT）:
7
  38.7%  ← 仅用 1/100 计算量
8

9
蒸馏 + GRPO 精调:
10
  42.8%  ← 蒸馏 + RL 进一步提升
11

12
结论：
13
  小模型无法仅靠 RL 学会推理（计算量不够）
14
  但可以蒸馏大模型的推理能力
15
  蒸馏 + 轻量 RL = 最佳性价比

展开收起

6. 后训练技术细节#

6.1 模型合并（Model Merging）#

DeepSeek 使用 模型合并 来稳定训练：

1
def merge_with_recent_checkpoint(
2
    model,
3
    recent_ckpt_path: str,
4
    merge_alpha: float = 0.8,
5
) -> torch.nn.Module:
6
    """
7
    与最近的 checkpoint 合并，防止灾难性遗忘。
8
    """
9
    recent_ckpt = torch.load(recent_ckpt_path)
10

11
    merged_state = {}
12
    for key in model.state_dict():
13
        if key in recent_ckpt:
14
            # 线性插值
15
            merged_state[key] = (
16
                merge_alpha * model.state_dict()[key]
17
                + (1 - merge_alpha) * recent_ckpt[key]
18
            )
19
        else:
20
            merged_state[key] = model.state_dict()[key]
21

22
    model.load_state_dict(merged_state)
23
    return model

展开收起

6.2 长上下文扩展#

R1 支持 128K 上下文，关键技术：

1
class LongContextConfig:
2
    """
3
    长上下文处理：
4
    1. YaRN 位置编码外推
5
    2. 注意力稀疏化
6
    3. 滑动窗口 + 全局注意力混合
7
    """
8
    max_position = 131072    # 128K tokens
9
    rope_theta = 10000        # YaRN 参数
10

11
    # 分组注意力（GQA 已内置）
12
    num_kv_heads = 8          # 远少于 num_heads

6.3 测试时 Scaling（Inference Scaling）#

R1 展示了推理时计算Scaling的有效性：

1
def test_time_scaling(
2
    model,
3
    prompt: str,
4
    num_samples: int = 64,
5
    majority_vote: bool = True,
6
) -> str:
7
    """
8
    测试时 Scaling：
9
    生成多个回答，用投票或 RM 选择最好的。
10

11
    这与预训练的 Scaling Law 正交——
12
    推理时可以独立增加计算量。
13
    """
14
    responses = []
15
    for _ in range(num_samples):
16
        with torch.no_grad():
17
            response = model.generate(
18
                prompt,
19
                max_tokens=4096,
20
                temperature=0.8,
21
            )
22
        responses.append(response)
23

24
    if majority_vote:
25
        # 多数投票
26
        return majority_vote_answers(responses)
27
    else:
28
        # RM 选择
29
        best_response = max(
30
            responses,
31
            key=lambda r: rule_based_reward(prompt, r)
32
        )
33
        return best_response
34

35

36
# 测试时 Scaling 效果（AIME 2024）：
37
#  sample=1:   71.3%
38
#  sample=4:   75.8%
39
#  sample=16:  78.2%
40
#  sample=64:  79.8%   ← DeepSeek-R1 报告的数字

展开收起

7. 实验结果与分析#

7.1 标准 Benchmark 对比#

Benchmark	DeepSeek-R1	OpenAI-o1	OpenAI-o1-mini	GPT-4o
MATH-500	96.3%	96.8%	80.3%	76.6%
AIME 2024	79.8%	79.2%	63.6%	52.0%
GPQA Diamond	71.5%	72.6%	45.2%	40.3%
LiveCodeBench	65.9%	66.0%	58.0%	51.0%
SWE-bench Verified	49.2%	48.9%	36.0%	38.6%
IFEVAL	89.0%	84.5%	74.6%	74.5%
AlpacaEval 2.0	97.3%	92.8%	84.0%	90.5%

7.2 思维链质量分析#

DeepSeek 团队对 R1 的思维链做了系统性分析：

1. 反思行为：

1
示例（数学证明题）：
2
"让我用反证法。假设 √2 是有理数...
3
  ...
4
  这推出矛盾，所以 √2 是无理数。
5

6
  等一下，我在第一步的推导有问题。
7
  让我重新检查...
8

9
  修正后，完整的证明是..."

2. 探索多种方法：

1
示例（算法题）：
2
"方法 1：暴力枚举，时间 O(n²)...
3
  这个方法太慢了。
4

5
  方法 2：动态规划，时间 O(n)...
6
  这个更好，但空间可以优化。
7

8
  方法 3：双指针 + 贪心...
9
  这是最优解！时间 O(1) 空间..."

3. 验证与回溯：

1
示例（计算题）：
2
"首先，我用公式 A 计算：答案是 42。
3

4
  让我验证一下：
5
  用公式 B 反推：42 = ...
6
  确认无误。
7

8
  再用代入法验证：
9
  ...
10
  完全一致。最终答案是 42。"

7.3 训练 Compute Scaling#

1
训练 compute vs 推理性能（AIME 2024）：
2

3
R1-Zero (不同训练步数):
4
  32K tokens:  40.0%
5
  64K tokens:  55.0%
6
  256K tokens: 68.0%
7
  512K tokens: 75.0%
8
  1M tokens:  79.8%
9

10
对比：
11
  OpenAI-o1 (report): ~80%
12
  GPT-4 (few-shot):   ~30%

8. 局限性与未来方向#

8.1 已知的局限性#

DeepSeek-R1 的已知问题

1. 语言混合问题

模型可能在同一回答中混合中英文
尤其在非中英文问题上更明显

2. 格式一致性

思维链格式不如 o1 稳定
有时跳步，有时过于冗长

3. 软件工程任务

SWE-bench 表现虽好但未达 SOTA
长程依赖和调试能力仍有提升空间

4. 复杂多模态推理

目前只验证了文本推理
视觉/音频推理待探索

8.2 未来研究方向#

1
┌──────────────────────────────────────────────────────────────┐
2
│                    DeepSeek-R1 之后的研究方向                  │
3
├──────────────────────────────────────────────────────────────┤
4
│                                                              │
5
│  1. 更强的推理基础                                           │
6
│     → 更大规模的纯 RL 训练                                   │
7
│     → 结合搜索树（MCTS）的推理                               │
8
│                                                              │
9
│  2. 自我改进的极限                                           │
10
│     → 模型能否自己生成更难的问题？                            │
11
│     → Constitutional AI + R1 的结合                           │
12
│                                                              │
13
│  3. 多模态推理                                              │
14
│     → R1 的推理能力能否迁移到视觉/音频？                      │
15
│     → MathVista, MMMU 上的突破                               │
16
│                                                              │
17
│  4. 高效推理                                                │
18
│     → 蒸馏 + 量化 + Speculative Decoding                    │
19
│     → 让 R1 级别的推理在手机/边缘设备上运行                   │
20
│                                                              │
21
│  5. 安全与对齐                                              │
22
│     → 如何确保"长思考"不产生有害内容？                        │
23
│     → R1 的"反思"能力能否用于安全对齐？                       │
24
│                                                              │
25
└──────────────────────────────────────────────────────────────┘

展开收起

9. 核心公式汇总#

9.1 规则奖励#

r(x,y)=1[answer(y)=gt(x)]⏟准确性+1[format(y)]⋅0.1⏟格式+1[thinking(y)]⋅0.05⏟推理过程r(x, y) = \underbrace{\mathbb{1}[\text{answer}(y) = \text{gt}(x)]}_{\text{准确性}} + \underbrace{\mathbb{1}[\text{format}(y)] \cdot 0.1}_{\text{格式}} + \underbrace{\mathbb{1}[\text{thinking}(y)] \cdot 0.05}_{\text{推理过程}}r(x,y)=准确性1[answer(y)=gt(x)]+格式1[format(y)]⋅0.1+推理过程1[thinking(y)]⋅0.05

9.2 GRPO 组内归一化#

Ai=ri−μGσG,μG=1∣G∣∑j∈GrjA_i = \frac{r_i - \mu_G}{\sigma_G}, \quad \mu_G = \frac{1}{|G|}\sum_{j\in G}r_jAi=σGri−μG,μG=∣G∣1j∈G∑rj

9.3 RFT 知识压缩#

LRFT=−E(x,y)∼DRL[log⁡πSFT(y∣x)]\mathcal{L}_{\text{RFT}} = -\mathbb{E}_{(x,y)\sim D_{\text{RL}}}\Big[\log\pi_{\text{SFT}}(y|x)\Big]LRFT=−E(x,y)∼DRL[logπSFT(y∣x)]

其中 DRLD_{\text{RL}}DRL 是 GRPO 生成的高质量推理数据。

9.4 拒绝采样#

DRS={(x,y)∣reward(x,y)≥τ}⊆DrawD_{\text{RS}} = \{(x, y) \mid \text{reward}(x, y) \geq \tau\} \subseteq D_{\text{raw}}DRS={(x,y)∣reward(x,y)≥τ}⊆Draw

9.5 测试时 Scaling#

y^=Vote({yi∼π(⋅∣x,ti)}i=1N)\hat{y} = \text{Vote}\big(\{y_i \sim \pi(\cdot|x, t_i)\}_{i=1}^{N}\big)y^=Vote({yi∼π(⋅∣x,ti)}i=1N)

10. 总结#

10.1 DeepSeek-R1 的三大贡献#

1
┌─────────────────────────────────────────────────────────────┐
2
│               DeepSeek-R1 的三点核心贡献                      │
3
├─────────────────────────────────────────────────────────────┤
4
│                                                              │
5
│  1. 证明了"纯 RL 可以涌现复杂推理"                             │
6
│     → 不需要人类写的思维链数据                                │
7
│     → 推理能力可以通过"试错-奖励"自学                         │
8
│     → 对 AI 发展的路径选择有深远影响                          │
9
│                                                              │
10
│  2. 提出了高效的后训练范式                                     │
11
│     → GRPO + RFT + 拒绝采样 + 多阶段 GRPO                    │
12
│     → 降低了 SOTA 对齐的训练成本                              │
13
│                                                              │
14
│  3. 开创了知识蒸馏的新方法                                     │
15
│     → 大模型 RL 探索 → 小模型蒸馏                             │
16
│     → 让 SOTA 推理能力普惠化                                  │
17
│                                                              │
18
└─────────────────────────────────────────────────────────────┘

展开收起

10.2 对 LLM 发展的影响#

1
DeepSeek-R1 之前：
2
  推理能力 = 预训练（知识）+ SFT（模仿人类）+ RLHF（对齐）
3
             ↑                              ↑
4
          知识来自互联网               思维链来自人工标注
5

6
DeepSeek-R1 之后：
7
  推理能力 = 预训练（知识）+ 纯 RL（自我探索）
8
             ↑                    ↑
9
          知识来自互联网       思维链通过"试错-奖励"涌现
10

11
这意味着：
12
  → 未来不需要昂贵的思维链标注
13
  → 推理能力可以规模化（更多的 RL 计算 = 更强的推理）
14
  → 小模型可以通过蒸馏获得大模型的推理能力

10.3 本系列文章总结#

文章	核心算法	关键洞见
RLHF 深度解析	PPO/GRPO/DPO	对齐全景图
PPO 深度解析	TRPO → PPO	信赖域优化
GRPO 深度解析	组内归一化	去掉 Critic
本文	R1 完整流水线	纯 RL 涌现推理

参考资料#

DeepSeek-AI. (2025). “DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning.” arXiv.
DeepSeek-AI. (2025). “DeepSeek-R1-Zero: Scaling Reinforcement Learning with Zero Supervised Data.” arXiv.
DeepSeek-AI. (2025). “DeepSeek-V3 Technical Report.” arXiv.
Shao, Z., et al. (2024). “DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models.” arXiv.
Schulman, J., et al. (2017). “Proximal Policy Optimization Algorithms.” arXiv.
OpenAI. (2024). “Learning to Reason with Large Language Models.” OpenAI Blog.
Wei, J., et al. (2022). “Chain-of-Thought Prompting Elicits Reasoning in Large Language Models.” NeurIPS.
Yao, S., et al. (2023). “Tree of Thoughts: Deliberate Problem Solving with Large Language Models.” NeurIPS.
Besta, M., et al. (2024). “Graph of Thoughts: Advanced Reasoning using Large Language Models.” arXiv.
Snell, C., et al. (2024). “Scaling LLM Training with compute-optimal RL.” arXiv.

深入理解 Diffusion Transformer (DiT)：扩散模型的 Transformer 范式

Thu, 16 Jul 2026 00:00:00 GMT

1. DiT 的核心思想#

1.1 一个让人惊讶的事实#

2020-2022 年间，扩散模型在图像生成上”大杀四方”——DALL·E 2、Imagen、Stable Diffusion 接连发布，看起来已稳坐王座。

但所有这些模型都有一个共同点：

它们的核心去噪网络都是 UNet——一个专为图像设计的卷积架构。

这就引出一个自然而大胆的问题：

如果 Transformer 已经”统治”NLP、统一了视觉理解（ViT）、推翻了卷积回归——它能不能”接管”扩散模型的核心网络？

2022 年底，UC Berkeley & NYU 的 William Peebles 与 Saining Xie 给出了回答：

Diffusion Transformer (DiT)——用纯 Transformer 替代 UNet，在 ImageNet 256×256 上以 FID 1.51 击败了 LDM、ADM 等基于 UNet 的所有扩散模型，并展现出极佳的 Gflops-vs-FID 缩放曲线。

1.2 DiT 的核心思想#

一句话概括：

把扩散模型的噪声预测网络从卷积 UNet 替换为 Vision Transformer (ViT)，并通过自适应层归一化 (AdaLN-Zero) 高效注入时间步与类别条件。

1
┌────────────────────────────────────────────────────────────────┐
2
│                  DiT 整体架构示意                                 │
3
├────────────────────────────────────────────────────────────────┤
4
│                                                                │
5
│  噪声图像 x_T                                                   │
6
│      │                                                        │
7
│      ▼                                                        │
8
│  Patchify: (H/P)×(W/P) × D                                    │
9
│      │                                                        │
10
│      ▼                                                        │
11
│  ┌──────────────────────────────┐                              │
12
│  │ + Position Embedding          │  位置编码                     │
13
│  └──────────────────────────────┘                              │
14
│      │                                                        │
15
│      ▼                                                        │
16
│  ┌──────────────────────────────────────────────────┐         │
17
│  │  N × DiT Block                                   │         │
18
│  │  ┌──────────────────────────┐                    │         │
19
│  │  │ AdaLN-Zero (调制)        │ ← (t, y)           │         │
20
│  │  ├──────────────────────────┤                    │         │
21
│  │  │ Multi-Head Self-Attention │                    │         │
22
│  │  ├──────────────────────────┤                    │         │
23
│  │  │ AdaLN-Zero               │ ← (t, y)           │         │
24
│  │  ├──────────────────────────┤                    │         │
25
│  │  │ MLP                       │                    │         │
26
│  │  └──────────────────────────┘                    │         │
27
│  └──────────────────────────────────────────────────┘         │
28
│      │                                                        │
29
│      ▼                                                        │
30
│  Linear Decode: → (H/P)×(W/P) 噪声预测                          │
31
│                                                                │
32
└────────────────────────────────────────────────────────────────┘

展开收起

1.3 一个生活化的比喻#

把 DiT 与 UNet 的差别想象成两个画家：

UNet 画家：用固定大小的”画笔”（卷积核）一寸一寸画——卷积核小，效率高，但必须从局部到全局逐步建立结构
DiT 画家：先看一眼整张画（把所有 patch 全局关联起来），第一笔就能把不同区域关联在一起——更”聪明”，但需要更多计算和更多数据

1.4 DiT 在生成式 AI 历史上的位置#

1
2015-2019: GAN 主导图像生成
2
  │
3
2020: DDPM 提出 (Ho et al.) - 重新让扩散模型回到舞台
4
  │  使用 UNet + 加性噪声预测
5
  │
6
2021: LDM / Stable Diffusion
7
  │  把扩散搬到潜空间 (latent space)
8
  │  仍用 UNet
9
  │
10
2022: DiT (Peebles & Xie, 2022.12)
11
  │  **首次: 用纯 Transformer 取代 UNet**
12
  │  首次发现: 缩放律 (Gflops vs FID) 极优
13
  │
14
2023: SD3 (MM-DiT)、PixArt-α、SDXL-Turbo
15
  │  DiT 思想扩展到文本到图像
16
  │
17
2024: Sora、Wan2.1、Open Sora
18
  │  DiT 扩展到视频生成
19
  │  Transformer 彻底统一了视觉生成
20
  │
21
2025+: 视觉大模型的基础架构

展开收起

2. 为什么可以替换 UNet？#

2.1 UNet 的”假设”#

UNet 在扩散模型中之所以好用，是因为它有恰到好处的归纳偏置 (Inductive Bias)：

UNet 的偏置	对扩散的贡献
局部性	像素相邻关系强，卷积核捕获短程结构
平移不变性	同一物体出现在图像任意位置，模型处理方式一致
金字塔结构	多尺度感受野 + skip connection 保留细节
参数共享	推理不同时间步很高效

这些偏置在小数据时代非常珍贵——它们让模型”少走弯路”。

2.2 DiT 的反问#

DiT 的论文提出了一个深刻的观察：

当数据规模和模型规模足够大时，归纳偏置的重要性会下降——Transformer 可以自己”学习到”这些偏置。

而在扩散生成任务上：

已经有海量图文配对数据 (LAION 等)
算力上不封顶 (DiT-XL 用了 8 卡 A100 训练若干月)
高分辨率图像对”全局理解”要求高——正是 Transformer 的强项

2.3 UNet vs DiT 对比#

维度	UNet	DiT
核心算子	2D 卷积 (ResNet block)	Self-Attention + MLP
归纳偏置	强 (局部、平移不变)	弱 (全局，由数据学习)
信息传递	跳层连接 + 多尺度	全局自注意力
时间/类别条件	拼接或 AdaGN	AdaLN-Zero (更优)
数据效率	小数据更友好	大数据更优
缩放行为	很快饱和	强大、可预测
下限 vs 上限	上限相对受限	上限更高

3. DiT 架构详解#

3.1 整体流程图#

1
输入: 噪声图像 x_t (B, C, H, W) + 时间步 t + 类别标签 y
2
                                                 │
3
        ┌────────────────────────────────────────┘
4
        ▼                                        ▼
5
   ┌─────────┐                              ┌──────────┐
6
   │ Patchify│                              │  t Embed │ + y Embed
7
   └────┬────┘                              └────┬─────┘
8
        │ (B, S, D)                              │ (B, D_cond)
9
        ▼                                        ▼
10
   ┌──────────────────────────────┐        ┌────────────┐
11
   │ + Pos. Embedding              │        │  MLP       │
12
   └──────────────────────────────┘        └─────┬──────┘
13
        │                                       │
14
        ▼                                       │
15
   ┌─────────────────────────────────────────────┘
16
   │
17
   ▼
18
┌──────────────────────────┐
19
│  N × DiT Block           │  对每个 block:
20
│  ├ AdaLN-Zero (调制)     │  c ← (t+y) Embedding
21
│  ├ Self-Attention         │  x ← x + SelfAttn(AdaLN(c)调制后的 x)
22
│  └ MLP                    │  x ← x + MLP(AdaLN(c)调制后的 x)
23
└──────────────────────────┘
24
   │
25
   ▼
26
Linear Decoder: 还原为噪声预测 (B, C, H, W)

展开收起

3.2 Patchify (Patch 嵌入)#

DiT 的第一步是把 H×WH \times WH×W 的图像转成 token 序列（与 ViT 类似）。设 patch 大小为 ppp，则序列长度：

S=(Hp)×(Wp)S = \left(\frac{H}{p}\right) \times \left(\frac{W}{p}\right)S=(pH)×(pW)

输入尺寸	Patch	序列长度 SSS
32×32	4	64
64×64	4	256
256×256	8	1024
256×256	16	256

1
import torch
2
import torch.nn as nn
3

4

5
class PatchEmbed(nn.Module):
6
    """将图像切分为 patch (DiT 版)。"""
7

8
    def __init__(self, img_size=32, patch_size=4, in_chans=3, embed_dim=768):
9
        super().__init__()
10
        self.img_size = img_size
11
        self.patch_size = patch_size
12
        self.num_patches = (img_size // patch_size) ** 2
13

14
        # 用 Conv2d 一气呵成: 切块 + 线性投影
15
        self.proj = nn.Conv2d(
16
            in_chans, embed_dim,
17
            kernel_size=patch_size, stride=patch_size,
18
        )
19

20
    def forward(self, x):
21
        """
22
        x: (B, C, H, W)
23
        return: (B, S, D)
24
        """
25
        x = self.proj(x)            # (B, D, H/p, W/p)
26
        x = x.flatten(2)            # (B, D, S)
27
        x = x.transpose(1, 2)       # (B, S, D)
28
        return x

展开收起

3.3 时间步与类别嵌入#

1
class TimestepEmbedder(nn.Module):
2
    """时间步 t → 向量 (用正弦/余弦位置编码 + MLP)。"""
3

4
    def __init__(self, hidden_size, frequency_embedding_size=256):
5
        super().__init__()
6
        self.mlp = nn.Sequential(
7
            nn.Linear(frequency_embedding_size, hidden_size),
8
            nn.SiLU(),
9
            nn.Linear(hidden_size, hidden_size),
10
        )
11
        self.frequency_embedding_size = frequency_embedding_size
12

13
    @staticmethod
14
    def timestep_embedding(t, dim, max_period=10000):
15
        """Sinusoidal position embedding for t."""
16
        half = dim // 2
17
        freqs = torch.exp(
18
            -math.log(max_period)
19
            * torch.arange(half, dtype=torch.float32) / half
20
        ).to(t.device)
21
        args = t[:, None].float() * freqs[None]
22
        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
23
        if dim % 2:
24
            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
25
        return embedding
26

27
    def forward(self, t):
28
        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
29
        return self.mlp(t_freq)
30

31

32
class LabelEmbedder(nn.Module):
33
    """类别标签 y → 向量。DiT 默认使用 class-conditional 的 CFG。"""
34

35
    def __init__(self, num_classes, hidden_size):
36
        super().__init__()
37
        # +1 是为"无条件" (CFG 时使用)
38
        self.embedding = nn.Embedding(num_classes + 1, hidden_size)
39

40
    def forward(self, labels):
41
        return self.embedding(labels)

展开收起

把时间步和类别 embedding 相加，得到最终条件向量 ccc：

1
c = t_embedder(t) + y_embedder(y)

3.4 自注意力块 (Self-Attention)#

与 ViT 几乎完全相同——QKV 投影 + 多头 softmax-attention + 输出投影。

1
class Attention(nn.Module):
2
    """Multi-Head Self-Attention (pre-LN)。"""
3

4
    def __init__(self, dim, num_heads=8, qkv_bias=False):
5
        super().__init__()
6
        assert dim % num_heads == 0
7
        self.num_heads = num_heads
8
        self.head_dim = dim // num_heads
9
        self.scale = self.head_dim ** -0.5
10

11
        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
12
        self.proj = nn.Linear(dim, dim)
13

14
    def forward(self, x):
15
        B, N, C = x.shape
16
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim)
17
        qkv = qkv.permute(2, 0, 3, 1, 4)  # (3, B, h, N, d)
18
        q, k, v = qkv.unbind(0)
19

20
        # 使用 PyTorch 2.0 的 SDPA 加速
21
        x = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0)
22
        x = x.transpose(1, 2).reshape(B, N, C)
23
        return self.proj(x)

展开收起

3.5 MLP 块#

1
class Mlp(nn.Module):
2
    """GELU MLP, hidden = 4 × in_dim。"""
3

4
    def __init__(self, dim, mult=4):
5
        super().__init__()
6
        hidden = int(dim * mult)
7
        self.fc1 = nn.Linear(dim, hidden)
8
        self.act = nn.GELU(approximate="tanh")
9
        self.fc2 = nn.Linear(hidden, dim)
10

11
    def forward(self, x):
12
        return self.fc2(self.act(self.fc1(x)))

4. AdaLN-Zero：条件注入的关键创新#

4.1 为什么 AdaLN-Zero 如此重要？#

扩散模型需要告诉网络”现在是去噪到哪一步了 ttt”、“我要生成什么类别 yyy”。

UNet 时代最常用的做法是 adaptive group normalization (AdaGN)：把 t,yt, yt,y 编码成一个向量，对每层归一化参数做调制。

DiT 论文研究了 4 种条件注入策略，发现最优方式是 AdaLN-Zero。

4.2 四种条件注入策略#

论文的图如下（简化）：

1
策略 1: In-context              策略 2: Cross-attention
2
─────────────                  ──────────────────
3
 x → [t,y, x] → Transformer    x → Transformer ──cross-attn── c
4
                                  Q K V 从 x 来
5
条件直接拼成 token                 K, V 从 c 来
6

7
策略 3: Adaptive Layer-Norm    策略 4: adaLN-Zero ✅
8
──────────────────            ─────────────────
9
x → LN → Modulated by c       x → AdaLN(c) = γ·LN·x + α
10
  → Attention/MLP               → Attention/MLP
11
                                (γ 和残差 α 都是 c 的函数)
12
                                初始化为零 → 残差为 0
13
                                (与残差连接相同效果先初始化)

4.3 AdaLN-Zero 的核心公式#

设条件嵌入 ccc，对每个 Transformer 块计算 6 个调制参数：

γ1,γ2,β1,β2=MLP(c)∈R4D\gamma_1, \gamma_2, \beta_1, \beta_2 = \text{MLP}(c) \in \mathbb{R}^{4D}γ1,γ2,β1,β2=MLP(c)∈R4Dα1,α2=MLP(c)∈R2D\alpha_1, \alpha_2 = \text{MLP}(c) \in \mathbb{R}^{2D}α1,α2=MLP(c)∈R2D

然后在每个残差块之前做调制：

h=a1+Attention(γ1⋅LayerNorm(x)+β1)h = a_1 + \text{Attention}\left(\gamma_1 \cdot \text{LayerNorm}(x) + \beta_1\right)h=a1+Attention(γ1⋅LayerNorm(x)+β1)h=a2+MLP(γ2⋅LayerNorm(h)+β2)h = a_2 + \text{MLP}\left(\gamma_2 \cdot \text{LayerNorm}(h) + \beta_2\right)h=a2+MLP(γ2⋅LayerNorm(h)+β2)

“Zero” 的含义：

初始化 α\alphaα 为零 → 整个块输出去残差效果 (identity)
这样模型在一开始就是”什么都不做”的扩散初始化
训练过程中，α\alphaα 逐渐增长，模型逐渐学出能力

AdaLN-Zero 的三大好处：

参数高效：6 维向量调制整个块，比 cross-attention 省得多

训练稳定：zero-init 让 loss 收敛更平顺

性能领先：在所有策略中 SOTA

4.4 完整 AdaLN DiT Block 代码#

1
class DiTBlock(nn.Module):
2
    """DiT 的核心 block: AdaLN-Zero 调制 + Self-Attention + MLP。"""
3

4
    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0):
5
        super().__init__()
6
        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
7
        self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True)
8
        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
9
        self.mlp = Mlp(hidden_size, mult=mlp_ratio)
10

11
        # AdaLN 参数生成: 输入 c ∈ R^D, 输出 6 * D
12
        self.adaLN_modulation = nn.Sequential(
13
            nn.SiLU(),
14
            nn.Linear(hidden_size, 6 * hidden_size, bias=True),
15
        )
16

17
    def forward(self, x, c):
18
        """
19
        x: (B, S, D)
20
        c: (B, D)  条件嵌入 (t + y)
21
        """
22
        # 生成 6 个调制参数
23
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
24
            self.adaLN_modulation(c).chunk(6, dim=1)
25
        )
26

27
        # 1. Attention 路径 (pre-LN + AdaLN + 门控残差)
28
        x = x + gate_msa.unsqueeze(1) * self.attn(
29
            modulate(self.norm1(x), shift_msa, scale_msa)
30
        )
31

32
        # 2. MLP 路径 (pre-LN + AdaLN + 门控残差)
33
        x = x + gate_mlp.unsqueeze(1) * self.mlp(
34
            modulate(self.norm2(x), shift_mlp, scale_mlp)
35
        )
36

37
        return x
38

39

40
def modulate(x, shift, scale):
41
    """AdaLN 调制: γ·x + β (但要求 γ, β 是 token-wise 的)。"""
42
    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)

展开收起

4.5 Zero-Initialization#

1
def _zero_init(module):
2
    """把 AdaLN 的输出初始化为 0 — 这一行很重要!"""
3
    for p in module.parameters():
4
        nn.init.zeros_(p)
5

6

7
# 在构造完 DiT 后调用一次
8
for block in model.blocks:
9
    _zero_init(block.adaLN_modulation)
10
# final_layer 的 AdaLN 同样 zero init

5. 完整的 DiT 实现#

5.1 完整主类#

1
import math
2
import torch
3
import torch.nn as nn
4
import torch.nn.functional as F
5

6

7
class DiT(nn.Module):
8
    """Diffusion Transformer (class-conditional)."""
9

10
    def __init__(
11
        self,
12
        img_size=32, patch_size=4, in_chans=3,
13
        hidden_size=1152, depth=28, num_heads=16,
14
        mlp_ratio=4.0, num_classes=1000,
15
    ):
16
        super().__init__()
17
        self.img_size = img_size
18
        self.patch_size = patch_size
19
        self.in_chans = in_chans
20
        self.out_chans = in_chans  # 预测噪声, 通道数相同
21

22
        # 1) Patch 嵌入
23
        self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, hidden_size)
24
        self.num_patches = self.patch_embed.num_patches
25

26
        # 2) 位置编码 (可学习)
27
        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches, hidden_size))
28

29
        # 3) 条件嵌入
30
        self.t_embedder = TimestepEmbedder(hidden_size)
31
        self.y_embedder = LabelEmbedder(num_classes, hidden_size)
32

33
        # 4) Transformer blocks
34
        self.blocks = nn.ModuleList([
35
            DiTBlock(hidden_size, num_heads, mlp_ratio)
36
            for _ in range(depth)
37
        ])
38

39
        # 5) 最终 LayerNorm
40
        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
41
        self.adaLN_modulation_final = nn.Sequential(
42
            nn.SiLU(),
43
            nn.Linear(hidden_size, 2 * hidden_size, bias=True),
44
        )
45

46
        # 6) 线性解码器: 把 patch tokens → 噪声 patch
47
        self.final_layer = nn.Linear(hidden_size, patch_size * patch_size * self.out_chans)
48

49
        self.initialize_weights()
50

51
    def initialize_weights(self):
52
        # 1) 位置编码用 sin-cos 初始化
53
        pos_embed = get_2d_sincos_pos_embed(
54
            self.pos_embed.shape[-1], int(self.num_patches ** 0.5)
55
        )
56
        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
57

58
        # 2) 默认 Linears xavier
59
        self.apply(self._init_weights)
60

61
        # 3) AdaLN 调制层 ZERO init (关键的 AdaLN-Zero!)
62
        for block in self.blocks:
63
            nn.init.zeros_(block.adaLN_modulation[1].weight)
64
            nn.init.zeros_(block.adaLN_modulation[1].bias)
65
        nn.init.zeros_(self.adaLN_modulation_final[1].weight)
66
        nn.init.zeros_(self.adaLN_modulation_final[1].bias)
67

68
        # 4) final layer 归零
69
        nn.init.zeros_(self.final_layer.weight)
70
        nn.init.zeros_(self.final_layer.bias)
71

72
    def _init_weights(self, m):
73
        if isinstance(m, nn.Linear):
74
            nn.init.xavier_uniform_(m.weight)
75
            if m.bias is not None:
76
                nn.init.constant_(m.bias, 0)
77
        elif isinstance(m, nn.LayerNorm):
78
            if m.weight is not None:
79
                nn.init.constant_(m.weight, 1.0)
80
            if m.bias is not None:
81
                nn.init.constant_(m.bias, 0)
82

83
    def unpatchify(self, x):
84
        """
85
        x: (B, S, p*p*out_chans) → (B, C, H, W)
86
        """
87
        B = x.shape[0]
88
        p = self.patch_size
89
        h = w = int(x.shape[1] ** 0.5)
90
        x = x.reshape(B, h, w, p, p, self.out_chans)
91
        x = torch.einsum("bhwpqc->bchpwq", x)
92
        return x.reshape(B, self.out_chans, h * p, w * p)
93

94
    def forward(self, x, t, y):
95
        """
96
        x: (B, C, H, W) 噪声图像
97
        t: (B,) 时间步
98
        y: (B,) 类别标签
99
        """
100
        # 1) 切块
101
        x = self.patch_embed(x)                 # (B, S, D)
102
        x = x + self.pos_embed                   # 加位置编码
103

104
        # 2) 条件嵌入
105
        c = self.t_embedder(t) + self.y_embedder(y)
106

107
        # 3) 进入 N 个 DiT block
108
        for block in self.blocks:
109
            x = block(x, c)
110

111
        # 4) 最终 AdaLN 调制
112
        shift, scale = self.adaLN_modulation_final(c).chunk(2, dim=1)
113
        x = modulate(self.norm_final(x), shift, scale)
114

115
        # 5) 线性解码 → unpatchify → 噪声
116
        x = self.final_layer(x)                  # (B, S, p*p*C)
117
        x = self.unpatchify(x)                   # (B, C, H, W)
118
        return x

展开收起

5.2 DiT 模型规模表#

变体	隐层 DDD	头数	层数 NNN	参数量	Gflops (256×256)
DiT-S	384	6	12	33M	5.9
DiT-B	768	12	12	130M	23.0
DiT-L	1024	16	24	458M	80.8
DiT-XL	1152	16	28	676M	119.4
DiT-H	1408	16	32	985M	174.5

一个有意思的发现：DiT-S 大约等于原始 UNet 的计算量，但 FID 更低——Transformer 实际更高效。

6. 训练与采样#

6.1 训练目标#

DiT 沿用 DDPM/LDM 的标准训练目标（与 UNet 完全一致）：

1
# === 标准 DDPM 训练循环 ===
2
def training_step(model, x0):
3
    """
4
    x0: 真实图像 (B, C, H, W)
5
    y:  类别标签 (B,)
6
    """
7
    # 1. 随机时间步
8
    t = torch.randint(0, num_train_timesteps, (x0.size(0),), device=x0.device)
9

10
    # 2. 采样噪声
11
    noise = torch.randn_like(x0)
12

13
    # 3. 前向扩散 (重参数化)
14
    xt = sqrt_alpha_bars[t, None, None, None] * x0 + \
15
         sqrt_one_minus_alpha_bars[t, None, None, None] * noise
16

17
    # 4. 预测噪声 (用 DiT 而非 UNet!)
18
    pred_noise = model(xt, t, y)
19

20
    # 5. MSE loss
21
    loss = F.mse_loss(pred_noise, noise)
22
    return loss

展开收起

关键点：DiT 替代的只是 “pred_noise = …” 这一行，其他一切都不变。

6.2 采样 (Sampling)#

采样同样沿用 DDIM/DPM-Solver 流水线，无需修改（只要模型是噪声预测 ϵθ\epsilon_\thetaϵθ）：

1
@torch.no_grad()
2
def sample_ddim(model, shape, y, num_steps=50):
3
    """DDIM 采样 (50 步就足够)。"""
4
    device = next(model.parameters()).device
5

6
    # 准备时间步调度
7
    ts = torch.linspace(num_train_timesteps - 1, 0, num_steps + 1).long().to(device)
8

9
    x = torch.randn(*shape, device=device)
10

11
    for i in range(num_steps):
12
        t_cur = ts[i]
13
        t_next = ts[i + 1]
14

15
        # 预测噪声
16
        pred_noise = model(x, t_cur.expand(x.size(0)), y)
17

18
        # DDIM 更新公式
19
        alpha_cur = alpha_bars[t_cur]
20
        alpha_next = alpha_bars[t_next]
21

22
        x0_pred = (x - (1 - alpha_cur).sqrt() * pred_noise) / alpha_cur.sqrt()
23
        x = alpha_next.sqrt() * x0_pred + (1 - alpha_next).sqrt() * pred_noise
24

25
    return x

展开收起

6.3 Classifier-Free Guidance (CFG)#

DiT 与 LDMs 一样使用 CFG 来增强条件生成的视觉质量：

1
@torch.no_grad()
2
def sample_cfg(model, x, t, y, null_class, guidance_scale=4.0):
3
    """
4
    CFG = unconditional + scale * (conditional - unconditional)
5
    """
6
    eps_cond = model(x, t, y)
7
    eps_uncond = model(x, t, null_class * torch.ones_like(y))
8
    return eps_uncond + guidance_scale * (eps_cond - eps_uncond)

训练时需要有概率地把 yyy 替换为 null 类别，让网络学会无条件生成：

1
def training_step_cfg(model, x0, y, uncond_prob=0.1):
2
    t = torch.randint(0, num_train_timesteps, (x0.size(0),), device=x0.device)
3

4
    # 10% 概率用 null 类别
5
    drop_mask = torch.rand(y.shape).to(y.device) < uncond_prob
6
    y_dropped = torch.where(drop_mask, null_class, y)
7

8
    noise = torch.randn_like(x0)
9
    xt = sqrt_alpha_bars[t] * x0 + sqrt_one_minus_alpha_bars[t] * noise
10
    pred = model(xt, t, y_dropped)
11
    return F.mse_loss(pred, noise)

7. 缩放律：DiT 最惊艳的结果#

7.1 Gflops vs FID 的近线性关系#

DiT 论文的一个核心贡献是发现：

FID（Fréchet Inception Distance，衡量生成质量的关键指标）随模型 Gflops 的增加，近似呈幂律下降——而且在极高 Gflops 时仍未饱和。

模型	Gflops	FID ↓	IS ↑
DiT-S (256)	5.9	11.6	92.5
DiT-B (256)	23.0	8.7	124.0
DiT-L (256)	80.8	5.7	156.0
DiT-XL (256)	119.4	4.5	173.0
DiT-XL (512)	525.0	2.3	241.0
当时 SOTA (ADM-UNet)	-	4.0	-

这条曲线意味着：只要继续把模型做大、做更长时间、喂更多数据，DiT 还能持续变好——没有”瓶颈”。

7.2 论文的缩放图#

1
FID
2
  │
3
4 │              ● DiT-XL/512 (2.3)
4
  │          ●
5
  │       ●
6
8 │    ●
7
  │  ●
8
  │● DiT-S
9
  └─────────────────────────────▶
10
    5G   20G   80G  120G     500G
11
            Gflops

7.3 与 UNet 的对比关键数据#

在相同的 Gflops 量级下：

相同 Gflops	UNet (ADM)	DiT	差距
~80 G	FID 4.0+	FID 5.7	DiT 以更低 FID 提升空间更大
缩放空间	很快饱和	持续下降	DiT 拥有更好的缩放性

8. DiT 的扩展：从图像到视频、从文本到视觉#

8.1 Stable Diffusion 3 (MM-DiT)#

Stability AI 在 SD3 中采用了多模态 DiT (MM-DiT)：

1
传统 DiT:
2
  - 把文本 c 与图像 token 拼接 → AdaLN 调制
3
  - 文本只通过 AdaLN 影响图像
4

5
MM-DiT (SD3):
6
  - 文本 token 与图像 token 在每个 attention block 里
7
    "joint attention" 一起做注意力
8
  - 每个模态有独立的投影矩阵 (Q/K/V)
9
  - 文本与图像"双向看": 文本看图像布局，图像看文本语义

1
┌─────────────────────────────────────────────┐
2
│            MM-DiT Block (SD3)                │
3
│                                              │
4
│ img_tokens, text_tokens                      │
5
│      │                                       │
6
│      ▼                                       │
7
│ ┌──────────────────────────┐                 │
8
│ │ LayerNorm                │                 │
9
│ │ Q = img_W_q @ img, ...   │ 各自不同的 W    │
10
│ │ K, V 拼接自 img + text   │                 │
11
│ │ Attn (Q, [K;V])         │                 │
12
│ └──────────────────────────┘                 │
13
│      │                                       │
14
│      ▼                                       │
15
│ MLP (各模态独立)                             │
16
└─────────────────────────────────────────────┘

展开收起

8.2 Sora: 视频生成的基础#

2024 年 OpenAI 发布的 Sora 的核心架构正是 DiT 的视频版本：

把”图像 patch”换成”时空 patch” (spatial-temporal patches)
一个 (H, W, T) 的视频被切成小块
DiT 在时空维度上做自注意力 → 全局理解每一帧

1
视频 patch 化:
2
原始视频 (T=16 帧, 256×256×3)
3
   │
4
   ▼
5
切成 16×16×2 时空小块 (2 帧, 16×16 像素)
6
   │
7
   ▼
8
得到 (T/2) × (H/16) × (W/16) 个 token
9
对于 16×256×256: 8 × 16 × 16 = 2048 个 token
10
   │
11
   ▼
12
与图像 DiT 一样的 Transformer 主干

Sora 的核心贡献之一：

只要 Transformer 主干够大、数据够多，DiT 可以学到世界模型 (world models)——视频里的物理因果关系。

8.3 与其他工作的关系谱系#

1
2022.12  DiT (Peebles & Xie)
2
  │
3
2023     PixArt-α (轻量 DiT)
4
  │
5
2023     SD3 / SDXL-Turbo (MM-DiT)
6
  │
7
2024.2   Sora (视频 DiT)
8
  │
9
2024.5   Hunyuan (腾讯视频 DiT)
10
  │
11
2024-25  CogVideoX, Wan2.1 (开源视频 DiT)
12
  │       Wan 2.1 16s 480P 单卡可跑
13
  │
14
2025     Movie Gen, Veo 2 (电影级视频 DiT)

9. DiT 的工程实践#

9.1 加速技巧#

1
# 1) PyTorch 2.0 SDPA 加速注意力
2
torch.backends.cuda.sdp_kernel(enable_flash=True, enable_mem_efficient=True)
3
# 2) 混合精度
4
with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
5
    pred = model(x, t, y)
6
# 3) torch.compile
7
model = torch.compile(model, mode="reduce-overhead")

9.2 显存优化#

梯度检查点 (Gradient Checkpointing): 把 DiT 的中间激活不存，只在反向时重算
FlashAttention: 长序列注意力降显存
分桶并行 (Bucketed Parallel): DiT-S/B/L 用单 H100 即可训练；XL/H 需要张量并行

1
from torch.utils.checkpoint import checkpoint
2

3
class DiTBlock(nn.Module):
4
    def forward(self, x, c):
5
        # 用 checkpoint 让深度网络少占显存
6
        return checkpoint(self._forward, x, c, use_reentrant=False)

9.3 Patch Size 的权衡#

情况	小 patch (p=4)	大 patch (p=8 / 16)
序列长度	长	短
计算量	大	小
局部精细度	高	低
推荐场景	高分辨率	低分辨率或大模型

9.4 类别 / 文本条件#

1
# 多类别 (含 null 类别用于 CFG)
2
self.y_embedder = nn.Embedding(num_classes + 1, hidden_size)
3

4
# 文本条件 (SD 时代)
5
self.text_embedder = ...  # CLIP / T5 等文本编码器

10. DiT 与传统 UNet 的对比总结#

维度	UNet (LDM/ADM)	DiT
核心层	2D Conv + ResNet block	Self-Attn + MLP
归纳偏置	强 (局部/平移不变)	弱 (由数据学习)
缩放行为	饱和较快	严格单调下降
高分辨率	多尺度金字塔天然适配	需要 patch 选择 + patch 化
视频扩展	需要 3D	统一时空 patch 即可
与 LLM 融合	模型架构不同	架构一致 → 便于多模态
工程成熟度	极其成熟	仍在快速迭代
训练效率	数据效率高	数据量越大越有优势

DiT 的核心遗产：它确立了**“扩散 + Transformer”为生成模型的标准范式**——既拥抱 Transformer 的缩放性，又复用扩散模型的训练范式。

11. 局限与未来#

11.1 当前局限#

数据饥渴：小数据上 DiT 比 UNet 表现差，需要 LAION/JFT 级的预训练
高分辨率成本：2562256^22562 已是 1024 个 token；102421024^210242 就是 16384 个 token，二次方的注意力很贵
离散化训练：decoder 直接预测连续 patch 值，与 VAE decoder 配合时会有损失
CFG 仍是必需：多模态质量提升对 CFG scale 依赖大

11.2 未来方向#

1
加速方向:
2
  ├── Patch packing (更稀疏的 token)
3
  ├── Linear attention (Mamba-style) 替代 self-attention
4
  ├── 蒸馏到少步采样 (Consistency Model, SDXL-Turbo)
5
  └── 潜在 patch 化 (更紧凑表示)
6

7
能力扩展:
8
  ├── 视频 DiT (Sora, Wan, CogVideoX)
9
  ├── 3D 视觉 (从图像扩散到 3D 资产生成)
10
  ├── 音频 DiT (MusicLDM, AudioLDM 2)
11
  ├── 多模态统一 DiT (同时文本/图/音/视频)
12
  └── 具身控制 DiT (机器人动作生成)

11.3 一句话总结 DiT 的影响#

DiT 不只是替换了 UNet，它证明了生成模型的”图像 → 多模态”过渡应当建立在 Transformer 之上——这一判断现在被 Sora、Stable Diffusion 3、Gemini 视频等多模态生成系统反复印证。

12. 总结#

12.1 核心要点回顾#

维度	关键要点
核心创新	用 Transformer 替换 UNet 作为去噪网络
关键设计	Patchify + AdaLN-Zero 调制
条件注入	c=temb+yembc = t_{\text{emb}} + y_{\text{emb}}c=temb+yemb，全块 AdaLN-Zero 调制
缩放行为	FID vs Gflops 严格服从幂律下降
优势	全局感受野、缩放性好、与 LLM 同架构
劣势	数据饥渴、O(S2)O(S^2)O(S2) 注意力高分辨率昂贵
影响	Sora、SD3、Wan2.1 全部建立在 DiT 之上

12.2 推荐学习资源#

1
论文:
2
  - DiT (2022): "Scalable Diffusion Models with Transformers" (Peebles & Xie)
3
  - U-ViT (2022): "All are Worth Words"
4
  - SD3 (2024): "Scaling Rectified Flow Transformers for High-Resolution Image Synthesis"
5

6
代码:
7
  - facebookresearch/DiT (官方实现)
8
  - HuggingFace diffusers (DiTPipeline)
9
  - stabilityai/sd3 (MM-DiT 商业级)
10

11
扩展阅读:
12
  - ViT (2020): Transformer 如何进入视觉
13
  - DDPM / DDIM 原文: 扩散基础
14
  - Sora 技术报告 (2024): 视频 DiT 巅峰

一句话总结：DiT 通过”patchify + AdaLN-Zero + 纯 Transformer 主干”这一简洁设计，把扩散模型推入 Transformer 时代，确立了图像/视频生成统一范式——也是 Sora 的技术源泉。

分布式训练深度解析：Data / Tensor / Pipeline / Sequence Parallelism

Thu, 16 Jul 2026 00:00:00 GMT

1. 引言：为什么需要分布式训练#

1.1 模型规模的增长#

1
┌─────────────────────────────────────────────────────────────┐
2
│               LLM 参数规模增长趋势                             │
3
├─────────────────────────────────────────────────────────────┤
4
│                                                              │
5
│  2018: ELMo          ~ 94M 参数                             │
6
│  2018: BERT-Large    ~ 340M 参数                            │
7
│  2019: GPT-2         ~ 1.5B 参数                           │
8
│  2020: GPT-3         ~ 175B 参数                           │
9
│  2021: PaLM          ~ 540B 参数                           │
10
│  2023: GPT-4        ~ 1.8T 参数 (估计)                     │
11
│  2024: Gemini Ultra ~ 1.5T 参数                            │
12
│                                                              │
13
│  内存需求（FP16）：                                          │
14
│  ─────────────────────────────────────────────────────────  │
15
│  1B  参数  →  2 GB                                         │
16
│  7B  参数  →  14 GB                                        │
17
│  70B 参数  →  140 GB                                       │
18
│  175B 参数 →  350 GB                                       │
19
│  1T  参数  →  2 TB                                         │
20
│                                                              │
21
│  单卡 A100 80GB：仅能放下 ~35B 参数的模型                    │
22
│                                                              │
23
└─────────────────────────────────────────────────────────────┘

展开收起

1.2 计算需求#

1
class ComputeRequirements:
2
    """
3
    计算需求分析
4
    """
5

6
    def flops_estimation(self):
7
        """
8
        训练 LLM 的 FLOPs 估算
9

10
        对于 decoder-only transformer:
11
        FLOPs ≈ 2 × N × D × G
12

13
        其中:
14
        - N: 模型参数量
15
        - D: 训练 tokens 数
16
        - G: 训练轮数
17
        """
18
        return {
19
            "gpt3_175b": {
20
                "params": "175B",
21
                "tokens": "300B",
22
                "flops": "~3.1 × 10^23 FLOPs",
23
                "a100_petaflops": "~3.1 × 10^5 PF",
24
                "a100_8x_90_days": "~1.5 × 10^6 PF-days",
25
            },
26
            "llama2_70b": {
27
                "params": "70B",
28
                "tokens": "2T (预训练)",
29
                "flops": "~1.4 × 10^24 FLOPs",
30
            },
31
        }
32

33
    def time_estimation(self):
34
        """
35
        训练时间估算
36
        """
37
        return """
38
        单卡 A100 (312 TFLOPS BF16):
39

40
        LLaMA-2 70B, 2T tokens:
41

42
        Theoretical time = FLOPs / FLOPS
43
                           = 1.4 × 10^24 / 3.12 × 10^14
44
                           = 4.5 × 10^9 seconds
45
                           = ~144 years
46

47
        结论: 单卡训练不现实
48

49
        需要的 GPU 数量:
50

51
        1000 GPUs → ~53 天
52
        4000 GPUs → ~13 天
53
        """

展开收起

1.3 并行策略全景#

1
┌─────────────────────────────────────────────────────────────┐
2
│                   分布式训练并行策略                           │
3
├─────────────────────────────────────────────────────────────┤
4
│                                                              │
5
│  Data Parallelism (DP)                                       │
6
│  ─────────────────────────────────────────────────────────  │
7
│  → 复制模型，并行处理不同数据分片                             │
8
│  → 最简单，提升吞吐量                                         │
9
│  → 通信: AllReduce (梯度)                                    │
10
│                                                              │
11
│  Tensor Parallelism (TP)                                     │
12
│  ─────────────────────────────────────────────────────────  │
13
│  → 切分单个层的权重到多卡                                     │
14
│  → 用于超大型单层 (如 LLM 的 FFN/Attention)                  │
15
│  → 通信: AllReduce / AllGather                              │
16
│                                                              │
17
│  Pipeline Parallelism (PP)                                   │
18
│  ─────────────────────────────────────────────────────────  │
19
│  → 按层切分模型到多卡                                         │
20
│  → 减少单卡显存需求                                          │
21
│  → 通信: P2P (点对点)                                        │
22
│                                                              │
23
│  Sequence Parallelism (SP)                                   │
24
│  ─────────────────────────────────────────────────────────  │
25
│  → 按序列维度切分                                             │
26
│  → 配合 TP 处理超长上下文                                     │
27
│  → 通信: AllReduce                                          │
28
│                                                              │
29
│  组合: 2D / 3D Parallelism                                   │
30
│  → TP × PP × DP                                             │
31
│  → 千卡级别训练的必备                                       │
32
│                                                              │
33
└─────────────────────────────────────────────────────────────┘

展开收起

1.4 本系列文章关联#

文章	关联
PEFT/LoRA 深度解析	高效微调减少计算需求
模型量化深度解析	训练后量化减少显存
vLLM 深度解析	分布式推理技术

2. Data Parallelism#

2.1 核心原理#

1
class DataParallelism:
2
    """
3
    Data Parallelism
4
    """
5

6
    def basic_idea(self):
7
        """
8
        核心思想
9

10
        将训练数据分成多个 shard，每个 GPU 持有完整模型副本，
11
        并行处理不同数据，最后同步梯度。
12
        """
13
        return """
14
        Data Parallelism 示意:
15

16
        GPU 0: [Model Copy] → Batch 0 → Forward → Backward → Grad 0
17
        GPU 1: [Model Copy] → Batch 1 → Forward → Backward → Grad 1
18
        GPU 2: [Model Copy] → Batch 2 → Forward → Backward → Grad 2
19
        GPU 3: [Model Copy] → Batch 3 → Forward → Backward → Grad 3
20

21
        ↓ Gradient AllReduce ↓
22

23
        Avg(Grad 0, Grad 1, Grad 2, Grad 3)
24

25
        ↓ 每个 GPU 用平均梯度更新 ↓
26

27
        特点:
28
        → 每个 GPU 有完整模型
29
        → 只通信梯度
30
        → 增加 GPU 数量可线性提升吞吐量
31
        """
32

33
    def gradient_synchronization(self):
34
        """
35
        梯度同步
36
        """
37
        return '''
38
        class DataParallelGradientSync:
39
            """梯度同步实现"""
40

41
            def allreduce_sync(self, gradients, world_size):
42
                """
43
                同步 SGD (所有 GPU 同步等待)
44

45
                每个 GPU 计算梯度
46
                → AllReduce 同步梯度
47
                → 所有 GPU 都得到相同的平均梯度
48
                → 各自更新
49

50
                优点: 收敛行为与单卡相同
51
                缺点: 最慢的 GPU 拖慢整体
52
                """
53
                import torch.distributed as dist
54

55
                # 对每个参数梯度执行 AllReduce
56
                for param in model.parameters():
57
                    dist.all_reduce(param.grad, op=dist.ReduceOp.SUM)
58
                    param.grad /= world_size
59

60
            def distributed_sampler(self, dataset, num_gpus):
61
                """数据分片"""
62
                return torch.utils.data.DistributedSampler(
63
                    dataset,
64
                    num_replicas=num_gpus,
65
                    rank=self.rank,
66
                )
67
        '''
68

69
    def memory_analysis(self):
70
        """
71
        显存分析
72
        """
73
        return {
74
            "model_weights": "N GB (完整副本)",
75
            "optimizer_states": "2N GB (Adam: m + v, FP32)",
76
            "gradients": "2N GB (FP32)",
77
            "activations": "取决于 batch size",
78
            "total_per_gpu": "~6N GB + activations",
79

80
            "example_7b_fp16": {
81
                "model": "14 GB",
82
                "optimizer": "28 GB",
83
                "gradients": "28 GB",
84
                "activations_8": "16 GB (batch=8)",
85
                "total": "~86 GB per GPU",
86
                "a100_80gb": "不够，需要进一步优化",
87
            },
88
        }

展开收起

2.2 DistributedDataParallel (DDP)#

1
class DDPImplementation:
2
    """
3
    PyTorch DDP 实现
4
    """
5

6
    def basic_usage(self):
7
        """
8
        基本用法
9
        """
10
        return '''
11
        import torch
12
        import torch.distributed as dist
13
        import torch.multiprocessing as mp
14
        from torch.nn.parallel import DistributedDataParallel as DDP
15

16
        def setup(rank, world_size):
17
            os.environ["MASTER_ADDR"] = "localhost"
18
            os.environ["MASTER_PORT"] = "12355"
19
            dist.init_process_group("nccl", rank=rank, world_size=world_size)
20
            torch.cuda.set_device(rank)
21

22
        def train(rank, world_size):
23
            setup(rank, world_size)
24

25
            # 创建模型
26
            model = MyModel().cuda()
27

28
            # 包装为 DDP
29
            model = DDP(model, device_ids=[rank])
30

31
            # 数据
32
            dataset = MyDataset()
33
            sampler = DistributedSampler(dataset)
34
            dataloader = DataLoader(dataset, sampler=sampler)
35

36
            # 训练循环
37
            for epoch in range(num_epochs):
38
                for batch in dataloader:
39
                    loss = model(batch).sum()
40
                    loss.backward()
41
                    optimizer.step()
42
                    optimizer.zero_grad()
43

44
            dist.destroy_process_group()
45

46
        # 启动
47
        mp.spawn(train, args=(world_size,), nprocs=world_size)
48
        '''
49

50
    def gradient_bucket(self):
51
        """
52
        Gradient Bucketing 优化
53

54
        DDP 使用 bucket 机制减少通信开销
55
        """
56
        return """
57
        问题:
58
        → 逐参数同步梯度效率低
59
        → 每个小梯度都需要完整的 AllReduce 通信
60

61
        解决: Gradient Bucketing
62

63
        ┌─────────────────────────────────────────────────────┐
64
        │  Gradients                                          │
65
        │  [g1][g2][g3][g4][g5][g6][g7][g8]...             │
66
        └─────────────────────────────────────────────────────┘
67
                      ↓ 分桶
68
        ┌─────────────────────────────────────────────────────┐
69
        │  Bucket 0: [g1][g2][g3][g4]  (size = 4 params)    │
70
        │  Bucket 1: [g5][g6][g7][g8]  (size = 4 params)    │
71
        │  ...                                                │
72
        └─────────────────────────────────────────────────────┘
73
                      ↓ 异步 AllReduce
74
        ┌─────────────────────────────────────────────────────┐
75
        │  Bucket 0: AllReduce 完成                          │
76
        │  Bucket 1: AllReduce 进行中...                      │
77
        │  ...                                                │
78
        └─────────────────────────────────────────────────────┘
79

80
        效果:
81
        → 减少通信次数 (从 N 次到 N/bucket_size 次)
82
        → 重叠计算和通信
83
        → 带宽利用率更高
84
        """
85

86
    def broadcast_buffers(self):
87
        """
88
        Buffer 广播
89
        """
90
        return """
91
        DDP 的 Buffer 处理:
92

93
        1. 同步 Buffers
94
           - BatchNorm 的 running_mean/std
95
           - 使用 AllReduce 同步
96
           - 每个 GPU 广播自己的 buffers
97

98
        2. 非同步 Buffers
99
           - Dropout 的随机状态
100
           - 不需要同步
101

102
        3. 注册 Hook
103
           - 可以在梯度同步时执行自定义操作
104
           - 用途: gradient clipping, monitoring
105
        """

展开收起

2.3 ZeRO 优化#

1
class ZeROOptimization:
2
    """
3
    ZeRO (Zero Redundancy Optimizer)
4

5
    切分优化器状态以减少显存冗余
6
    """
7

8
    def stages(self):
9
        """
10
        ZeRO 三个阶段
11
        """
12
        return {
13
            "ZeRO-1": {
14
                "name": "分片优化器状态",
15
                "description": "将 Adam optimizer states (m, v) 分片到各 GPU",
16
                "memory_reduction": "4x (75% reduction)",
17
                "communication": "与 DDP 相同",
18
            },
19
            "ZeRO-2": {
20
                "name": "分片优化器状态 + 梯度",
21
                "description": "ZeRO-1 + 梯度也分片",
22
                "memory_reduction": "8x (~87.5% reduction)",
23
                "communication": "略增 (allreduce 变成reduce-scatter + allgather)",
24
            },
25
            "ZeRO-3": {
26
                "name": "分片所有状态",
27
                "description": "ZeRO-2 + 模型参数也分片",
28
                "memory_reduction": "线性减少 (N 倍)",
29
                "communication": "显著增加 (需要动态 gather 参数)",
30
            },
31
        }
32

33
    def zero_stage1(self):
34
        """
35
        ZeRO-1 实现
36
        """
37
        return '''
38
        # ZeRO-1: 只分片优化器状态
39

40
        # 优化器状态 (Adam):
41
        # m (first moment) - FP32
42
        # v (second moment) - FP32
43

44
        # 每个参数 p 的优化器状态:
45
        # p_fp32, m, v = 4 + 4 + 4 = 12 bytes per param
46

47
        # 原来 (DDP): 每个 GPU 持有完整的 m, v
48
        # ZeRO-1: 每个 GPU 只持有 1/N 的 m, v
49

50
        class ZeRO1Optimizer:
51
            def __init__(self, model, world_size, rank):
52
                self.world_size = world_size
53
                self.rank = rank
54
                self.param_partitions = {}  # rank → params
55

56
                # 将参数均匀分配到各 GPU
57
                all_params = list(model.parameters())
58
                for i, param in enumerate(all_params):
59
                    target_rank = i % world_size
60
                    self.param_partitions[target_rank].append(param)
61

62
            def step(self):
63
                # 只更新本地 partition 的参数
64
                for param in self.param_partitions[self.rank]:
65
                    # 更新 m, v (只存储本地部分)
66
                    self._update_adam_state(param)
67
        '''
68

69
    def zero_vs_ddp_memory(self):
70
        """
71
        显存对比
72
        """
73
        return """
74
        70B 参数模型，FP16 权重，Adam 优化器:
75

76
        DDP (完全副本):
77
        ┌──────────────────────────────────────────────────┐
78
        │  权重 (FP16):        140 GB                      │
79
        │  梯度 (FP32):        280 GB                      │
80
        │  优化器状态 (FP32):  280 GB                      │
81
        │  激活值:              ~50 GB                      │
82
        │  ─────────────────────────────────────────────  │
83
        │  总计:               ~750 GB                     │
84
        │  需要: A100 80GB × 10                          │
85
        └──────────────────────────────────────────────────┘
86

87
        ZeRO-3 + 8 GPUs:
88
        ┌──────────────────────────────────────────────────┐
89
        │  权重分片:          140 / 8 = 17.5 GB          │
90
        │  梯度分片:          280 / 8 = 35 GB            │
91
        │  优化器状态分片:    280 / 8 = 35 GB            │
92
        │  激活值:              ~50 GB                    │
93
        │  ─────────────────────────────────────────────  │
94
        │  总计:               ~137.5 GB per GPU          │
95
        │  需要: A100 80GB × 8 + Offload                  │
96
        └──────────────────────────────────────────────────┘
97
        """

展开收起

3. Tensor Parallelism#

3.1 原理与动机#

1
class TensorParallelism:
2
    """
3
    Tensor Parallelism
4

5
    将单层权重矩阵沿某个维度切分到多个 GPU
6
    """
7

8
    def why_tensor_parallel(self):
9
        """
10
        为什么需要 Tensor Parallelism
11

12
        对于 LLaMA-2 70B:
13
        """
14
        return {
15
            "embedding": "4096 × 32000 = 131M",
16
            "attention": {
17
                "q_proj": "4096 × 4096 = 16.8M × 3 = 50.4M",
18
                "k_proj": "4096 × 4096",
19
                "v_proj": "4096 × 4096",
20
                "o_proj": "4096 × 4096 = 16.8M",
21
                "total_attention": "67.2M per layer × 80 layers",
22
            },
23
            "ffn": {
24
                "gate_proj": "11008 × 4096 = 45M",
25
                "up_proj": "4096 × 11008 = 45M",
26
                "down_proj": "11008 × 4096 = 45M",
27
                "total_ffn": "135M per layer × 80 layers",
28
            },
29

30
            "largest_layer": "FFN up_proj: 45M params = 90 GB FP16",
31
            "problem": "单层 FFN 比单卡显存还大！",
32
            "solution": "切分到多卡",
33
        }
34

35
    def matrix_partitioning(self):
36
        """
37
        矩阵分块策略
38
        """
39
        return """
40
        矩阵乘法 Y = XW
41

42
        W: (M, N)
43
        X: (B, M)
44
        Y: (B, N)
45

46
        ┌─────────────────────────────────────────────────┐
47
        │              W (M×N)                             │
48
        │  ┌─────────┐ ┌─────────┐ ┌─────────┐            │
49
        │  │   W0    │ │   W1    │ │   W2    │            │
50
        │  │(M, N/3) │ │(M, N/3) │ │(M, N/3) │            │
51
        │  └─────────┘ └─────────┘ └─────────┘            │
52
        │     GPU 0      GPU 1      GPU 2                │
53
        └─────────────────────────────────────────────────┘
54

55
        X: (B, M) 复制到所有 GPU
56

57
        Y0 = X @ W0 (B, N/3) → GPU 0
58
        Y1 = X @ W1 (B, N/3) → GPU 1
59
        Y2 = X @ W2 (B, N/3) → GPU 2
60

61
        输出拼接: Y = [Y0, Y1, Y2]
62
        """

展开收起

3.2 Column Parallel 与 Row Parallel#

1
class ParallelismTypes:
2
    """
3
    列并行与行并行
4
    """
5

6
    def column_parallel(self):
7
        """
8
        Column Parallel
9

10
        切分权重矩阵的列维度
11
        """
12
        return """
13
        用于 Linear 层的 QKV 投影:
14

15
        W_qkv: (d_model, 3 * d_hidden)
16

17
        Column Parallel (切分 3*d_hidden):
18

19
        GPU 0: W_q0 (d_model, d_hidden)
20
        GPU 1: W_k0 (d_model, d_hidden)
21
        GPU 2: W_v0 (d_model, d_hidden)
22
        ...
23

24
        前向:
25
        X: (B, d_model) 复制
26
        ↓
27
        X @ W_q0 = Q0 (B, d_hidden)
28
        X @ W_k0 = K0 (B, d_hidden)
29
        X @ W_v0 = V0 (B, d_hidden)
30

31
        通信: 无 (输出已经是分片的)
32

33
        注意: 需要 AllGather 合并 Q/K/V 以进行 Attention
34
        """
35

36
    def row_parallel(self):
37
        """
38
        Row Parallel
39

40
        切分权重矩阵的行维度
41
        """
42
        return """
43
        用于 Linear 层的输出投影:
44

45
        W_o: (d_hidden, d_model)
46

47
        Row Parallel (切分 d_hidden):
48

49
        GPU 0: W_o0 (d_hidden/2, d_model)
50
        GPU 1: W_o1 (d_hidden/2, d_model)
51

52
        前向:
53
        X0: (B, d_hidden/2) ← 从 Attention 来
54
        X1: (B, d_hidden/2) ← 从 Attention 来
55
        ↓
56
        X0 @ W_o0 = Y0 (B, d_model)
57
        X1 @ W_o1 = Y1 (B, d_model)
58

59
        Y = Y0 + Y1  ← AllReduce (求和)
60

61
        通信: AllReduce
62
        """
63

64
    def attention_parallel(self):
65
        """
66
        Attention 中的 Tensor Parallelism
67
        """
68
        return """
69
        Attention 层并行化:
70

71
        ┌─────────────────────────────────────────────────────────┐
72
        │                    Attention Layer                        │
73
        ├─────────────────────────────────────────────────────────┤
74
        │                                                           │
75
        │  Input X                                                  │
76
        │      ↓                                                    │
77
        │  ┌─────────────────┐                                    │
78
        │  │  QKV Projection  │  Column Parallel (每个 GPU 3 heads) │
79
        │  └─────────────────┘                                    │
80
        │      ↓                                                    │
81
        │  ┌─────────────────────────────────────────────────┐     │
82
        │  │  Q 分布到所有 GPU                                │     │
83
        │  │  K, V 需要 AllGather 收集                      │     │
84
        │  │  各自计算部分 attention                          │     │
85
        │  │  Output 经过 AllReduce                          │     │
86
        │  └─────────────────────────────────────────────────┘     │
87
        │      ↓                                                    │
88
        │  ┌─────────────────┐                                    │
89
        │  │  Output Proj    │  Row Parallel                       │
90
        │  └─────────────────┘                                    │
91
        │                                                           │
92
        └─────────────────────────────────────────────────────────┘
93

94
        Megatron-LM 风格: 将 heads 分组，每组分配给不同 GPU
95
        """

展开收起

3.3 通信分析#

1
class TPCommunication:
2
    """
3
    Tensor Parallelism 通信分析
4
    """
5

6
    def communication_pattern(self):
7
        """
8
        通信模式
9
        """
10
        return """
11
        Transformer 层的 TP 通信:
12

13
        ┌──────────────────────────────────────────────────────────┐
14
        │  Layer N                                                   │
15
        ├──────────────────────────────────────────────────────────┤
16
        │                                                           │
17
        │  Input X                                                   │
18
        │     ↓                                                     │
19
        │  QKV: Column Parallel → 无通信                            │
20
        │     ↓                                                     │
21
        │  Q AllGather (all TP ranks) → N-1 次广播                  │
22
        │  K AllGather (all TP ranks) → N-1 次广播                  │
23
        │  V AllGather (all TP ranks) → N-1 次广播                  │
24
        │     ↓                                                     │
25
        │  Attention Score (本地)                                   │
26
        │     ↓                                                     │
27
        │  Softmax → 无通信                                         │
28
        │     ↓                                                     │
29
        │  Attention × V → AllReduce (求和输出)                     │
30
        │     ↓                                                     │
31
        │  Output Proj: Row Parallel → AllReduce                     │
32
        │     ↓                                                     │
33
        │  FFN Gate: Column Parallel → 无通信                        │
34
        │     ↓                                                     │
35
        │  FFN Up: Column Parallel → 无通信                          │
36
        │     ↓                                                     │
37
        │  FFN AllGather → AllGather                                │
38
        │     ↓                                                     │
39
        │  FFN Down: Row Parallel → AllReduce                       │
40
        │     ↓                                                     │
41
        │  Output X                                                 │
42
        │                                                           │
43
        └──────────────────────────────────────────────────────────┘
44

45
        每层通信: 4 AllGather + 2 AllReduce
46
        """
47

48
    def bandwidth_requirement(self):
49
        """
50
        带宽需求
51
        """
52
        return {
53
            "per_layer_per_gpu": "2 × seq_len × hidden_size × tp_size (bytes)",
54
            "for_70b_seq_4096": "~128 MB per layer per forward",
55
            "forward_per_epoch": "~128 MB × 80 layers × 1M steps",
56
            "backward": "2x forward",
57

58
            "bandwidth_analysis": {
59
                "A100_NVLink": "900 GB/s bidirectional",
60
                "Gen1_InfiniBand": "25 GB/s",
61
                "Gen2_InfiniBand": "50 GB/s",
62
                "Gen4_InfiniBand": "200 GB/s",
63
            },
64

65
            "recommendation": "TP 需要 NVLink 或高速互联",
66
        }

展开收起

4. Pipeline Parallelism#

4.1 基本原理#

1
class PipelineParallelism:
2
    """
3
    Pipeline Parallelism
4

5
    将模型按层切分到多个 GPU
6
    """
7

8
    def layer_distribution(self):
9
        """
10
        层分布
11
        """
12
        return """
13
        32 层模型，4-GPU Pipeline Parallel:
14

15
        GPU 0: Layers 0-7     (embedding + first 8 layers)
16
        GPU 1: Layers 8-15    (middle layers)
17
        GPU 2: Layers 16-23   (middle layers)
18
        GPU 3: Layers 24-31   (last 8 layers + output)
19

20
        Forward:
21
        ┌─────────────────────────────────────────────────────────┐
22
        │  GPU 0: [L0][L1][L2][L3][L4][L5][L6][L7] → hidden_7   │
23
        │      ↓                                                 │
24
        │  GPU 1: [L8][L9][L10][L11][L12][L13][L14][L15]→hidd_15│
25
        │      ↓                                                 │
26
        │  GPU 2: [L16]...[L23] → hidden_23                     │
27
        │      ↓                                                 │
28
        │  GPU 3: [L24]...[L31] → output                        │
29
        └─────────────────────────────────────────────────────────┘
30
        """
31

32
    def pipeline_bubble(self):
33
        """
34
        Pipeline Bubble 问题
35
        """
36
        return """
37
        Naive Pipeline (无 interleaving):
38

39
        时间 →→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→
40

41
        Microbatch 1: [GPU0][GPU1][GPU2][GPU3]
42
        Microbatch 2:      [GPU0][GPU1][GPU2][GPU3]
43
        Microbatch 3:           [GPU0][GPU1][GPU2][GPU3]
44
        Microbatch 4:                [GPU0][GPU1][GPU2][GPU3]
45

46
        Bubble (空闲):
47
        [  ][  ][  ][GPU0][  ][  ][  ][GPU1][  ][  ][  ][GPU2][  ][  ][  ][GPU3]
48

49
        问题:
50
        → Bubble 数量 ≈ (P-1) / num_microbatches
51
        → P=4, microbatches=4 → 75% bubble!
52
        → GPU 利用率极低
53
        """
54

55
    def interleaving(self):
56
        """
57
        Interleaving 减少 Bubble
58
        """
59
        return """
60
        Interleaving (degree=4):
61

62
        每个 GPU 处理多个 chunks:
63

64
        GPU 0: [c0_0][c1_0][c2_0][c3_0][c4_0][c5_0]...
65
        GPU 1: [c0_1][c1_1][c2_1][c3_1][c4_1][c5_1]...
66
        GPU 2: [c0_2][c1_2][c2_2][c3_2][c4_2][c5_2]...
67
        GPU 3: [c0_3][c1_3][c2_3][c3_3][c4_3][c5_3]...
68

69
        每个 microbatch 在不同 GPU 上轮流执行
70
        Bubble 大幅减少
71

72
        Bubble 比例 ≈ 1 / (num_microbatches × degree)
73
        """

展开收起

4.2 调度策略#

1
class ScheduleStrategies:
2
    """
3
    调度策略
4
    """
5

6
    def forward_backward(self):
7
        """
8
        Forward-Backward-Pipeline
9
        """
10
        return """
11
        GPipe 调度 (Channeh et al., 2019):
12

13
        1F1B 模式 (One Forward, One Backward):
14

15
        Forward:  [F0][F1][F2][F3]
16
        Backward: [B3][B2][B1][B0]
17

18
        完整调度 ( microbatches=8 ):
19

20
        Stage 0: [F0][F1][F2][F3][B3][B2][B1][B0]
21
        Stage 1:    [F0][F1][F2][F3][B3][B2][B1][B0]
22
        Stage 2:       [F0][F1][F2][F3][B3][B2][B1][B0]
23
        Stage 3:          [F0][F1][F2][F3][B3][B2][B1][B0]
24

25
        特点:
26
        → Forward 完成后立即开始 Backward
27
        → 稳定状态时每个 GPU 只做一个操作
28
        → Bubble 只出现在开始和结束
29
        """
30

31
    def interleaved_schedule(self):
32
        """
33
        Interleaved Schedule
34
        """
35
        return """
36
        Interleaved 1F1B:
37

38
        每个 stage 被分成多个 sub-stages:
39

40
        GPU 0: [S0_0][S0_1]...[S0_k]  (layers 0..k, k+1..2k)
41
        GPU 1: [S1_0][S1_1]...[S1_k]  (layers 2k..3k, ...)
42

43
        调度:
44
        F0_0, F0_1, F1_0, F0_2, F1_1, F1_2, ...
45

46
        优势:
47
        → 减少通信等待
48
        → 更好利用计算和通信重叠
49
        """
50

51
    def async_pipeline(self):
52
        """
53
        异步 Pipeline (实验性)
54
        """
55
        return """
56
        Chimera / PipeDream 风格:
57

58
        思想: Forward 和 Backward 可以异步
59

60
        GPU 0: [F0][F1][F2][F3]...[F0][F1][F2][F3]...  (连续 forward)
61
        GPU 1:    [F0][F1][F2][F3]...[B0][B1][B2][B3]...  (backward 滞后)
62

63
        权衡:
64
        → 更高 GPU 利用率
65
        → 收敛性可能受影响
66
        → 需要 careful weight versioning
67
        """

展开收起

4.3 PipeDream 实现#

1
class PipeDreamImplementation:
2
    """
3
    PipeDream 风格实现
4
    """
5

6
    def weight_stashing(self):
7
        """
8
        Weight Stashing
9

10
        保存每个 microbatch 对应的权重版本
11
        """
12
        return '''
13
        class PipeDreamModel:
14
            """PipeDream 实现"""
15

16
            def __init__(self, stages, num_microbatches):
17
                self.stages = stages
18
                self.num_microbatches = num_microbatches
19
                self.weights = {}  # microb_id → weights
20
                self.optimizer = Adam(self.stages.parameters())
21

22
            def forward_backward(self, input_batch, microb_ids):
23
                """
24
                1F1B 调度
25
                """
26
                # Forward
27
                for mb_id in microb_ids:
28
                    self.stages[mb_id % len(self.stages)].forward(
29
                        self.weights[mb_id]  # 使用对应版本的权重
30
                    )
31

32
                # Backward
33
                for mb_id in reversed(microb_ids):
34
                    self.stages[mb_id % len(self.stages)].backward(
35
                        self.weights[mb_id]  # 梯度累积到正确版本
36
                    )
37

38
                # Update (每隔一段时间)
39
                self.optimizer.step()
40
                self.optimizer.zero_grad()
41
        '''
42

43
    def microbatch_vs_pipeline(self):
44
        """
45
        Microbatch 与 Pipeline 对比
46
        """
47
        return {
48
            "small_microbatch": {
49
                "pros": "更好的负载均衡，减少 bubble",
50
                "cons": "更多通信开销，更多 activation 显存",
51
            },
52
            "large_microbatch": {
53
                "pros": "减少通信，简化调度",
54
                "cons": "bubble 比例增加",
55
            },
56

57
            "recommendation": {
58
                "global_batch_size": "根据 GPU 数量调整",
59
                "microbatch_size": "尽量让 pipeline 填满",
60
                "num_microbatches": "P × 4 ~ P × 10",
61
            },
62
        }

展开收起

5. Sequence Parallelism#

5.1 动机与场景#

1
class SequenceParallelism:
2
    """
3
    Sequence Parallelism
4

5
    沿序列维度切分
6
    """
7

8
    def why_sequence_parallel(self):
9
        """
10
        为什么需要 Sequence Parallelism
11

12
        当 TP 处理超长序列时:
13
        """
14
        return """
15
        场景: 100K tokens 序列, TP=8
16

17
        问题:
18
        1. Attention 计算: O(N²) 显存
19
           → 即使 TP 切分 Q/K/V，单个 head 仍需处理 N/8 tokens
20

21
        2. LayerNorm / Dropout
22
           → 需要看到完整序列
23
           → 成为瓶颈
24

25
        3. 显存分布不均
26
           → embedding / output 层: 参数量大
27
           → attention / ffn: 激活值大
28

29
        解决: Sequence Parallelism
30

31
        沿序列维度切分:
32
        seq_len = 100K
33
        SP = 8
34
        每个 GPU 处理 12.5K tokens
35
        """
36

37
    def layerwise_operations(self):
38
        """
39
        需要跨序列通信的操作
40
        """
41
        return {
42
            "attention": "需要 AllGather Q/K/V，然后 AllReduce 输出",
43
            "layernorm": "需要 AllReduce 计算 mean/variance",
44
            "dropout": "可以在本地处理",
45
            "embedding": "需要 AllGather (输入) 和 AllReduce (输出)",
46
            "loss": "需要 AllReduce",
47
        }

展开收起

5.2 环形通信#

1
class RingCommunication:
2
    """
3
    环形通信实现
4
    """
5

6
    def ring_allreduce(self):
7
        """
8
        Ring AllReduce
9
        """
10
        return """
11
        Ring AllReduce 用于 Sequence Parallelism:
12

13
        4 GPUs, 每个持有 N/4 tokens 的部分和
14

15
        Step 1: Reduce Scatter
16
        ┌─────────────────────────────────────────────────────┐
17
        │  GPU 0: [a0] → [a0+b0+c0+d0]                        │
18
        │  GPU 1: [b0] → [b0+c0+d0+a0]                        │
19
        │  GPU 2: [c0] → [c0+d0+a0+b0]                        │
20
        │  GPU 3: [d0] → [d0+a0+b0+c0]                        │
21
        └─────────────────────────────────────────────────────┘
22

23
        Step 2: AllGather
24
        ┌─────────────────────────────────────────────────────┐
25
        │  GPU 0: [sum] → [sum, sum, sum, sum]                │
26
        │  GPU 1: [sum] → [sum, sum, sum, sum]                │
27
        │  GPU 2: [sum] → [sum, sum, sum, sum]                │
28
        │  GPU 3: [sum] → [sum, sum, sum, sum]                │
29
        └─────────────────────────────────────────────────────┘
30

31
        通信量: 2 × (N-1) / num_gpus
32
        """
33

34
    def ring_attention(self):
35
        """
36
        Ring Attention
37
        """
38
        return """
39
        Ring Attention 实现:
40

41
        4 GPUs, 序列长度 N, 每个 GPU 处理 N/4 tokens
42

43
        Q[i]: 本地查询
44
        K[i], V[i]: 本地键值
45

46
        1. Q 保持在本地
47

48
        2. K, V 通过环形传递:
49
           GPU 0: Q0 + [K0,V0] → 计算 [K0,V0] 的 attention
50
           GPU 0: Q0 + [K1,V1] → 计算 [K1,V1] 的 attention (接收自 GPU 1)
51
           GPU 0: Q0 + [K2,V2] → 计算 [K2,V2] 的 attention (接收自 GPU 2)
52
           GPU 0: Q0 + [K3,V3] → 计算 [K3,V3] 的 attention (接收自 GPU 3)
53

54
        3. 累加所有 attention 结果
55

56
        通信: 每个 GPU 接收 K,V N-1 次
57
        """

展开收起

6. 混合并行策略#

6.1 2D / 3D Parallelism#

1
class HybridParallelism:
2
    """
3
    混合并行策略
4
    """
5

6
    def tensor_pipeline_2d(self):
7
        """
8
        TP × PP (2D Parallelism)
9
        """
10
        return """
11
        2D Parallelism: TP × PP
12

13
        例子: 70B 模型, 16 GPUs
14

15
        TP = 4, PP = 4:
16

17
        GPU Layout:
18
        ┌─────────┬─────────┬─────────┬─────────┐
19
        │  Stage 0 │  Stage 1 │  Stage 2 │  Stage 3 │
20
        │ (TP=4)  │ (TP=4)  │ (TP=4)  │ (TP=4)  │
21
        └─────────┴─────────┴─────────┴─────────┘
22

23
        每个 Stage 有 4 GPUs (TP group)
24
        4 个 Stages (Pipeline)
25

26
        通信:
27
        - TP 内部: AllReduce (NVLink)
28
        - PP 跨 stage: P2P (IB)
29
        """
30

31
    def dp_tp_pp_3d(self):
32
        """
33
        DP × TP × PP (3D Parallelism)
34
        """
35
        return """
36
        3D Parallelism: DP × TP × PP
37

38
        例子: 175B 模型, 64 GPUs
39

40
        DP = 8, TP = 4, PP = 2:
41
        总计: 8 × 4 × 2 = 64 GPUs
42

43
        结构:
44
        ┌─────────────────────────────────────────────────────┐
45
        │  DP Groups (8)                                       │
46
        │  ┌──────────┬──────────┬──────────┬──────────┐      │
47
        │  │  PP 0    │  PP 1    │  PP 0    │  PP 1    │ ...  │
48
        │  │ (TP=4)  │ (TP=4)  │ (TP=4)  │ (TP=4)  │      │
49
        │  └──────────┴──────────┴──────────┴──────────┘      │
50
        │       ↑                                               │
51
        │    Pipeline Communication                            │
52
        └─────────────────────────────────────────────────────┘
53

54
        内存估算 (175B FP16):
55
        - TP=4: 每个 GPU 持有 1/4 的单层权重
56
        - PP=2: 每 GPU 持有 1/2 的层数
57
        - DP=8: 每个 DP replica 独立
58

59
        每 GPU: 175B / (4 × 2) × 2 bytes = ~44 GB
60
        + activations + optimizer
61
        """
62

63
    def communication_topology(self):
64
        """
65
        通信拓扑感知
66
        """
67
        return """
68
        现代 GPU 集群的拓扑:
69

70
        Node 0: [GPU 0] [GPU 1] [GPU 2] [GPU 3]
71
        Node 1: [GPU 4] [GPU 5] [GPU 6] [GPU 7]
72
        ...
73

74
        带宽层次:
75
        ┌─────────────────────────────────────────────────────┐
76
        │  NVLink (within node): ~900 GB/s                    │
77
        │  PCIe (within node): ~64 GB/s                        │
78
        │  InfiniBand (cross node): ~50 GB/s                  │
79
        └─────────────────────────────────────────────────────┘
80

81
        最佳并行策略:
82
        ┌─────────────────────────────────────────────────────┐
83
        │  TP: within node (NVLink)                           │
84
        │  PP: across nodes (IB)                             │
85
        │  DP: across all (梯度同步相对较小)                    │
86
        └─────────────────────────────────────────────────────┘
87
        """

展开收起

6.2 Megatron-LM 架构#

1
class MegatronArchitecture:
2
    """
3
    Megatron-LM 架构
4
    """
5

6
    def parallel_layout(self):
7
        """
8
        并行布局
9
        """
10
        return '''
11
        # Megatron-LM 3D Parallel
12

13
        # 配置:
14
        # tensor_model_parallel_size = 4  (TP)
15
        # pipeline_model_parallel_size = 4  (PP)
16
        # data_parallel_size = N / (TP × PP)
17

18
        # 模型构建
19
        model = TransformerBlock(
20
            tensor_parallel=TensorParallel,
21
            pipeline_parallel=PipelineParallel,
22
        )
23

24
        # Forward 流程:
25
        # 1. TP group 内广播输入
26
        # 2. Pipeline 调度
27
        # 3. 各 GPU 计算本地层
28
        # 4. PP 跨 GPU 传递 hidden states
29
        '''
30

31
    def memory_estimation(self):
32
        """
33
        显存估算
34
        """
35
        return """
36
        LLaMA-2 70B, 4096 sequence, batch=1, BF16:
37

38
        单 GPU (无并行):
39
        - 模型: 140 GB
40
        - 激活值: ~50 GB
41
        - 优化器: 280 GB
42
        - 总计: ~470 GB
43

44
        TP=8, PP=4:
45
        - TP: 每层权重 / 8
46
        - PP: 层数 / 4
47
        - 每 GPU 权重: 140 / 32 = 4.4 GB
48
        - 每 GPU 激活: 50 / 8 = 6.25 GB (SP 配合)
49
        - 每 GPU 优化器: 280 / 32 = 8.75 GB
50
        - 总计: ~19.4 GB per GPU
51

52
        结论: 可以在 A100 80GB 上训练
53
        """

展开收起

7. 主流框架实战#

7.1 PyTorch FSDP#

1
class FSDPImplementation:
2
    """
3
    FSDP (Fully Sharded Data Parallel)
4
    """
5

6
    def basic_usage(self):
7
        """
8
        FSDP 基本用法
9
        """
10
        return '''
11
        import torch
12
        from torch.distributed.fsdp import (
13
            FullyShardedDataParallel as FSDP,
14
            ShardingStrategy,
15
            MixedPrecision,
16
            BackwardPrefetch,
17
        )
18
        from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
19

20
        # 混合精度
21
        fp16_policy = MixedPrecision(
22
            param_dtype=torch.float16,
23
            reduce_dtype=torch.float16,
24
            buffer_dtype=torch.float16,
25
        )
26

27
        # 定义模型
28
        model = Transformer()
29

30
        # FSDP 包装
31
        fsdp_model = FSDP(
32
            model,
33
            sharding_strategy=ShardingStrategy.FULL_SHARD,  # ZeRO-3
34
            mixed_precision=fp16_policy,
35
            auto_wrap_policy=transformer_auto_wrap_policy,
36
            backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
37
            device_id=torch.cuda.current_device(),
38
        )
39

40
        # 训练
41
        for batch in dataloader:
42
            loss = fsdp_model(batch)
43
            loss.backward()  # FSDP 自动分片梯度
44
            fsdp_model.optimize_step()
45
        '''
46

47
    def sharding_strategies(self):
48
        """
49
        分片策略
50
        """
51
        return {
52
            "FULL_SHARD": "参数 + 梯度 + 优化器状态全分片 (ZeRO-3)",
53
            "SHARD_GRAD_OP": "梯度 + 优化器状态分片 (ZeRO-2)",
54
            "NO_SHARD": "不分片，类似 DDP",
55
            "HYBRID_SHARD": "本地分片 + 跨节点通信",
56

57
            "recommendation": {
58
                "small_model": "SHARD_GRAD_OP (低通信开销)",
59
                "large_model": "FULL_SHARD (最低显存)",
60
                "multi_node": "HYBRID_SHARD (减少跨节点通信)",
61
            },
62
        }
63

64
    def cpu_offload(self):
65
        """
66
        CPU Offload
67
        """
68
        return '''
69
        # CPU Offload 配置
70

71
        from torch.distributed.fsdp import CPUOffload
72

73
        fsdp_model = FSDP(
74
            model,
75
            sharding_strategy=ShardingStrategy.FULL_SHARD,
76
            cpu_offload=CPUOffload(
77
                offload_params=True,  # 将参数卸载到 CPU
78
                pin_memory=True,
79
            ),
80
        )
81

82
        # 效果:
83
        # → 进一步减少 GPU 显存
84
        # → 但增加 CPU-GPU 传输开销
85
        # → 适合超大型模型
86
        '''

展开收起

7.2 DeepSpeed#

1
class DeepSpeedImplementation:
2
    """
3
    DeepSpeed 实现
4
    """
5

6
    def basic_config(self):
7
        """
8
        DeepSpeed 配置
9
        """
10
        return '''
11
        # ds_config.json
12

13
        {
14
            "train_batch_size": 32,
15
            "gradient_accumulation_steps": 4,
16
            "fp16": {
17
                "enabled": true
18
            },
19
            "zero_optimization": {
20
                "stage": 3,
21
                "offload_optimizer": {
22
                    "device": "cpu",
23
                    "pin_memory": true
24
                },
25
                "offload_param": {
26
                    "device": "cpu"
27
                },
28
                "overlap_comm": true,
29
                "contiguous_gradients": true,
30
                "sub_group_size": 1e9
31
            },
32
            "gradient_clipping": 1.0,
33
            "wall_clock_breakdown": false
34
        }
35
        '''
36

37
    def training_loop(self):
38
        """
39
        训练循环
40
        """
41
        return '''
42
        import deepspeed
43

44
        # 初始化 DeepSpeed
45
        model, optimizer, _, _ = deepspeed.initialize(
46
            model=model,
47
            optimizer=optimizer,
48
            config=ds_config,
49
            dist_init_required=True,
50
        )
51

52
        # 训练
53
        for batch in dataloader:
54
            loss = model(batch)
55
            model.backward(loss)
56
            model.step()
57
        '''
58

59
    def pipeline_parallel(self):
60
        """
61
        DeepSpeed Pipeline Parallelism
62
        """
63
        return '''
64
        # Pipeline 配置
65

66
        {
67
            "pipeline": {
68
                "pipeline_parallel_size": 4,
69
                "num_stages": 4,
70
                "pipeline_checkpoint": true
71
            },
72
            "zero_optimization": {
73
                "stage": 1,  # ZeRO-1 (配合 PP 使用)
74
                "reduce_bucket_size": 5e8,
75
                "allgather_bucket_size": 5e8
76
            }
77
        }
78

79
        # 注意: PP 和 ZeRO-1/2 配合效果最好
80
        # ZeRO-3 会增加 PP 的通信量
81
        '''

展开收起

7.3 ColossalAI#

1
class ColossalAIImplementation:
2
    """
3
    ColossalAI 实现
4
    """
5

6
    def hybrid_parallel(self):
7
        """
8
        混合并行配置
9
        """
10
        return '''
11
        from colossalai.booster import Booster
12
        from colossalai.booster.plugin import HybridParallelPlugin
13
        from colossalai.cluster import DistCoordinator
14

15
        # 插件配置
16
        plugin = HybridParallelPlugin(
17
            tp_size=4,           # Tensor Parallelism
18
            pp_size=2,           # Pipeline Parallelism
19
            zero_stage=1,        # ZeRO-1
20
            num_microbatches=8,  # Pipeline microbatches
21
        )
22

23
        # 启动 Booster
24
        booster = Booster(plugin=plugin)
25

26
        model, optimizer, dataloader, criterion = booster.enable(
27
            model=model,
28
            optimizer=optimizer,
29
            dataloader=dataloader,
30
            criterion=criterion,
31
        )
32

33
        # 训练循环
34
        for batch in dataloader:
35
            outputs = model(batch)
36
            loss = criterion(outputs, batch["labels"])
37
            booster.backward(loss, optimizer)
38
            booster.optimize_step(optimizer)
39
        '''

展开收起

8. 通信优化#

8.1 通信计算重叠#

1
class CommunicationComputationOverlap:
2
    """
3
    通信与计算重叠
4
    """
5

6
    def gradient_computation_overlap(self):
7
        """
8
        梯度计算与 AllReduce 重叠
9
        """
10
        return """
11
        异步梯度同步:
12

13
        传统 (同步):
14
        ┌─────────────────────────────────────────────────────┐
15
        │ [Forward] [AllReduce Grad] [Backward] [Optimizer]    │
16
        └─────────────────────────────────────────────────────┘
17

18
        重叠 (异步):
19
        ┌─────────────────────────────────────────────────────┐
20
        │ [F] [B1] [AR1] [B2] [AR2] [B3] [AR3] ...           │
21
        │         ↑                                           │
22
        │       Backward 和 AllReduce 并行                    │
23
        └─────────────────────────────────────────────────────┘
24

25
        实现:
26
        - 使用梯度 bucket
27
        - 前一个 bucket 的 AllReduce 和后一个 bucket 的 backward 并行
28
        - DDP 默认启用
29
        """
30

31
    def pp_communication_overlap(self):
32
        """
33
        Pipeline Parallelism 通信重叠
34
        """
35
        return """
36
        P2P 通信与计算重叠:
37

38
        ┌─────────────────────────────────────────────────────┐
39
        │  GPU 0: [F0][B0][F1][B1]...                        │
40
        │             ↕ P2P                                   │
41
        │  GPU 1:    [F0][B0][F1][B1]...                      │
42
        │                 ↕ P2P                                │
43
        │  GPU 2:       [F0][B0][F1][B1]...                  │
44
        └─────────────────────────────────────────────────────┘
45

46
        关键: 使用 CUDA 流并行执行计算和通信
47

48
        # CUDA 流
49
        compute_stream = torch.cuda.Stream()
50
        comm_stream = torch.cuda.Stream()
51

52
        with torch.cuda.stream(compute_stream):
53
            compute()
54

55
        # 异步发送
56
        dist.isend(data, comm_stream)
57
        """

展开收起

8.2 通信原语#

1
class CommunicationPrimitives:
2
    """
3
    通信原语
4
    """
5

6
    def allreduce(self):
7
        """
8
        AllReduce
9
        """
10
        return """
11
        AllReduce: 所有 GPU 获得相同的归约结果
12

13
        X = [x0, x1, x2, x3]
14

15
        AllReduce(SUM):
16
        → [x0+x1+x2+x3, x0+x1+x2+x3, x0+x1+x2+x3, x0+x1+x2+x3]
17

18
        用途:
19
        - DDP 梯度同步
20
        - TP Row Parallel 输出求和
21
        """
22

23
    def allgather(self):
24
        """
25
        AllGather
26
        """
27
        return """
28
        AllGather: 收集所有 GPU 的数据
29

30
        X = [x0, x1, x2, x3]
31

32
        AllGather:
33
        → [x0, x1, x2, x3, x0, x1, x2, x3, x0, x1, x2, x3, x0, x1, x2, x3]
34

35
        用途:
36
        - TP QKV 合并
37
        - Sequence Parallel 中收集完整张量
38
        """
39

40
    def reducescatter(self):
41
        """
42
        ReduceScatter
43
        """
44
        return """
45
        ReduceScatter: 归约后分散
46

47
        X = [x0, x0, x0, x0,
48
             x1, x1, x1, x1,
49
             x2, x2, x2, x2,
50
             x3, x3, x3, x3]
51

52
        ReduceScatter(SUM):
53
        → [sum(x0), sum(x1), sum(x2), sum(x3)]
54

55
        用途:
56
        - ZeRO 梯度分片
57
        """

展开收起

8.3 NCCL 优化#

1
class NCCLOptimization:
2
    """
3
    NCCL 优化
4
    """
5

6
    def nccl_tuning(self):
7
        """
8
        NCCL 调优
9
        """
10
        return {
11
            "NCCL_IB_CUDA": "使用 CUDA 加速 IB 传输",
12
            "NCCL_SHM_DISABLE": "禁用共享内存 (可能导致问题)",
13
            "NCCL_NVLS_ENABLE": "启用 NVLink 优化",
14
            "NCCL_MIN_NCHANNELS": "最小化 channel 数",
15

16
            "best_practices": [
17
                "使用 NCCL 2.12+",
18
                "启用 CUDA 加速",
19
                "配置合适的 PXN (NVSwitch)",
20
            ],
21
        }
22

23
    def topology_aware(self):
24
        """
25
        拓扑感知通信
26
        """
27
        return '''
28
        # PyTorch 拓扑感知
29

30
        import torch.distributed as dist
31

32
        # 检测 GPU 拓扑
33
        if torch.cuda.is_available():
34
            # 使用 NVLink/PCIe 拓扑
35
            dist.init_process_group(
36
                backend="nccl",
37
                init_method="env://",
38
                timeout=timedelta(seconds=1800),
39
            )
40

41
        # NCCL 自动检测拓扑
42
        # 优化 GPU 对等连接
43
        '''

展开收起

9. 显存优化技术#

9.1 激活值重计算#

1
class ActivationRecomputation:
2
    """
3
    激活值重计算 (Activation Checkpointing)
4
    """
5

6
    def concept(self):
7
        """
8
        概念
9
        """
10
        return """
11
        激活值重计算: 用计算换显存
12

13
        Forward: 保存部分激活值
14
        Backward: 重新计算未保存的激活值
15

16
        例子:
17
        ┌─────────────────────────────────────────────────────┐
18
        │  Forward:                                            │
19
        │  x1 = Layer1(input)                                  │
20
        │  x2 = Layer2(x1)   ← 保存                           │
21
        │  x3 = Layer3(x2)   ← 保存                           │
22
        │  x4 = Layer4(x3)   ← 保存                           │
23
        │  x5 = Layer5(x4)   ← 保存                           │
24
        └─────────────────────────────────────────────────────┘
25

26
        ┌─────────────────────────────────────────────────────┐
27
        │  Backward without recompute:                         │
28
        │  grad_x4 = softmax(x5)  ← 需要 x4, x3, x2 的激活值   │
29
        └─────────────────────────────────────────────────────┘
30

31
        ┌─────────────────────────────────────────────────────┐
32
        │  Backward with recompute:                           │
33
        │  x2 = recompute(Layer2, input)  ← 重新计算          │
34
        │  grad_x4 = softmax(x5)                               │
35
        └─────────────────────────────────────────────────────┘
36
        """
37

38
    def memory_savings(self):
39
        """
40
        显存节省
41
        """
42
        return {
43
            "full_activations": "~40% 显存用于激活值",
44
            "with_checkpointing": "节省 ~60-70% 激活值显存",
45
            "cost": "增加 ~20-30% 计算量",
46

47
            "selective_checkpointing": {
48
                "recommendation": "只对大层启用",
49
                "effect": "平衡显存和计算",
50
            },
51
        }
52

53
    def implementation(self):
54
        """
55
        实现
56
        """
57
        return '''
58
        # PyTorch 实现
59
        from torch.utils.checkpoint import checkpoint
60

61
        # 全量重计算
62
        model = Transformer()
63
        model = torch.utils.checkpoint.checkpoint_sequential(
64
            model.layers,
65
            chunks=4,
66
        )
67

68
        # 选择性重计算
69
        class MyLayer(nn.Module):
70
            def forward(self, x):
71
                return checkpoint(self.heavy_op, x)
72

73
        # FSDP 中的重计算
74
        fsdp_model = FSDP(
75
            model,
76
            activation_checkpointing=checkpoint,
77
        )
78
        '''

展开收起

9.2 混合精度训练#

1
class MixedPrecisionTraining:
2
    """
3
    混合精度训练
4
    """
5

6
    def bf16_training(self):
7
        """
8
        BF16 训练
9
        """
10
        return """
11
        BF16 vs FP16:
12

13
        FP16:
14
        - 1 sign, 5 exponent, 10 mantissa
15
        - 动态范围: ~6×10^4
16
        - 精度较高但可能溢出
17

18
        BF16:
19
        - 1 sign, 8 exponent, 7 mantissa
20
        - 动态范围: ~3×10^38
21
        - 动态范围与 FP32 相同
22
        - 训练更稳定
23

24
        推荐:
25
        → Ampere 架构 (A100, RTX 3090+) 支持 BF16
26
        → H100 原生支持 BF16
27
        → 大模型训练优先使用 BF16
28
        """
29

30
    def loss_scaling(self):
31
        """
32
        Loss Scaling
33
        """
34
        return '''
35
        # FP16 训练需要 Loss Scaling
36

37
        # PyTorch 自动管理
38
        scaler = torch.cuda.amp.GradScaler()
39

40
        for batch in dataloader:
41
            with torch.cuda.amp.autocast():
42
                outputs = model(batch)
43
                loss = criterion(outputs, labels)
44

45
            scaler.scale(loss).backward()
46
            scaler.step(optimizer)
47
            scaler.update()
48

49
        # BF16 通常不需要 loss scaling
50
        with torch.cuda.amp.autocast(dtype=torch.bfloat16):
51
            outputs = model(batch)
52
        '''

展开收起

10. 实践指南#

10.1 配置选择#

1
class ConfigurationGuide:
2
    """
3
    配置选择指南
4
    """
5

6
    def model_size_to_config(self):
7
        """
8
        根据模型大小选择配置
9
        """
10
        return {
11
            "7b_model": {
12
                "single_node_8x_a100": {
13
                    "strategy": "ZeRO-2 or ZeRO-3",
14
                    "tp": 1,
15
                    "pp": 1,
16
                    "dp": 8,
17
                    "notes": "可以直接用 FSDP",
18
                },
19
            },
20
            "13b_model": {
21
                "single_node_8x_a100": {
22
                    "strategy": "ZeRO-2 + TP=2",
23
                    "tp": 2,
24
                    "pp": 1,
25
                    "dp": 4,
26
                    "notes": "PP 可能引入 bubble",
27
                },
28
            },
29
            "70b_model": {
30
                "single_node_8x_a100_80gb": {
31
                    "strategy": "ZeRO-3 + TP=4 + PP=2",
32
                    "tp": 4,
33
                    "pp": 2,
34
                    "dp": 1,
35
                    "notes": "需要仔细调优",
36
                },
37
                "multi_node": {
38
                    "strategy": "ZeRO-3 + TP=8",
39
                    "tp": 8,
40
                    "pp": 1,
41
                    "dp": "根据节点数",
42
                    "notes": "8节点训练",
43
                },
44
            },
45
            "175b_model": {
46
                "strategy": "3D Parallel: TP × PP × DP",
47
                "tp": 8,
48
                "pp": 8,
49
                "dp": "根据 GPU 总数",
50
                "notes": "需要数百 GPU",
51
            },
52
        }
53

54
    def batch_size_recommendation(self):
55
        """
56
        Batch Size 建议
57
        """
58
        return """
59
        Global Batch Size (GBS):
60

61
        GBS = micro_batch_size × gradient_accumulation_steps × dp_size
62

63
        经验公式:
64
        - 小模型 (<10B): GBS ≈ 500K - 1M tokens
65
        - 中模型 (10-70B): GBS ≈ 1M - 4M tokens
66
        - 大模型 (>70B): GBS ≈ 4M+ tokens
67

68
        示例:
69
        模型: LLaMA-2 70B
70
        GPU: 32 × A100
71

72
        目标 GBS: 4M tokens
73
        micro_batch_size: 1 (GPU 显存限制)
74
        gradient_accumulation: 128
75
        dp_size: 32
76

77
        GBS = 1 × 128 × 32 = 4096 tokens
78
        (需要调整)
79

80
        调整:
81
        micro_batch_size: 4
82
        gradient_accumulation: 256
83
        dp_size: 8 (TP=4, PP=2)
84

85
        GBS = 4 × 256 × 8 = 8192 tokens
86
        """

展开收起

10.2 性能调优清单#

1
class TuningChecklist:
2
    """
3
    性能调优清单
4
    """
5

6
    def pre_training(self):
7
        """
8
        训练前检查
9
        """
10
        return """
11
        □ 验证模型正确性
12
          - 单 GPU 测试 small model
13
          - 对比参考实现 loss
14

15
        □ 验证并行策略
16
          - TP=2 验证正确性
17
          - PP=2 验证流水线
18
          - ZeRO 验证梯度同步
19

20
        □ 性能基准
21
          - MFU (Model FLOP Utilization)
22
          - 通信带宽利用率
23
          - GPU 利用率
24
        """
25

26
    def during_training(self):
27
        """
28
        训练中监控
29
        """
30
        return """
31
        □ 监控指标
32
          - Loss 收敛曲线
33
          - Gradient norm (检查梯度爆炸)
34
          - Learning rate schedule
35

36
        □ 资源利用
37
          - GPU 利用率 (nvidia-smi)
38
          - 通信带宽 (NCCL debug)
39
          - CPU-GPU 数据传输
40

41
        □ 异常检测
42
          - Loss NaN/Inf
43
          - GPU OOM
44
          - 通信超时
45
        """

展开收起

11. 核心公式汇总#

11.1 计算 FLOPs#

FLOPsforward=2×N×S\text{FLOPs}_{\text{forward}} = 2 \times N \times SFLOPsforward=2×N×SFLOPsbackward=4×N×S\text{FLOPs}_{\text{backward}} = 4 \times N \times SFLOPsbackward=4×N×S

其中 NNN 是参数量，SSS 是 tokens 数。

11.2 MFU 计算#

MFU=实际 FLOPs峰值 FLOPs×100%\text{MFU} = \frac{\text{实际 FLOPs}}{\text{峰值 FLOPs}} \times 100\%MFU=峰值 FLOPs实际 FLOPs×100%

11.3 通信量#

操作	通信量
DDP AllReduce	2×grad_size2 \times \text{grad\_size}2×grad_size
TP AllReduce	2×layer_output2 \times \text{layer\_output}2×layer_output
PP P2P	hidden_size×S\text{hidden\_size} \times Shidden_size×S

12. 总结#

12.1 策略对比#

1
┌─────────────────────────────────────────────────────────────┐
2
│                  并行策略对比                                  │
3
├─────────────────────────────────────────────────────────────┤
4
│                                                              │
5
│  Data Parallelism                                           │
6
│  ─────────────────────────────────────────────────────────  │
7
│  显存减少: 1x (不变)                                         │
8
│  通信量: 中 (梯度)                                           │
9
│  通信模式: AllReduce                                         │
10
│  适用: 增加吞吐                                               │
11
│                                                              │
12
│  Tensor Parallelism                                          │
13
│  ─────────────────────────────────────────────────────────  │
14
│  显存减少: TPx                                               │
15
│  通信量: 高 (每层多次)                                        │
16
│  通信模式: AllReduce/AllGather                               │
17
│  适用: 超大单层                                               │
18
│                                                              │
19
│  Pipeline Parallelism                                        │
20
│  ─────────────────────────────────────────────────────────  │
21
│  显存减少: PPx                                               │
22
│  通信量: 低 (P2P)                                            │
23
│  通信模式: 点对点                                             │
24
│  适用: 减少层数                                               │
25
│                                                              │
26
│  Sequence Parallelism                                        │
27
│  ─────────────────────────────────────────────────────────  │
28
│  显存减少: SPx                                               │
29
│  通信量: 中 (attention)                                       │
30
│  通信模式: Ring AllReduce                                     │
31
│  适用: 超长序列                                               │
32
│                                                              │
33
└─────────────────────────────────────────────────────────────┘

展开收起

12.2 最佳实践#

1
分布式训练最佳实践:
2

3
  1. 从简单开始
4
     → 先用 DDP/ZeRO 验证
5
     → 再加入 TP/PP
6

7
  2. 通信拓扑感知
8
     → TP 在节点内 (NVLink)
9
     → PP 跨节点 (InfiniBand)
10
     → DP 可跨节点
11

12
  3. 显存优化
13
     → 启用混合精度 (BF16)
14
     → 使用激活值重计算
15
     → ZeRO-3 配合 PP 时注意通信
16

17
  4. 性能监控
18
     → MFU 目标 > 50%
19
     → 监控 GPU 利用率和通信
20
     → 识别瓶颈 (计算 vs 通信)
21

22
  5. 故障恢复
23
     → 保存 checkpoint
24
     → 支持从 checkpoint 恢复
25
     → 处理 GPU 故障

展开收起

参考资料#

Shoeybi, M., et al. (2020). “Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism.” arXiv.
Huang, Y., et al. (2019). “GPipe: Efficient Training of Giant Neural Networks with Pipeline Parallelism.” NeurIPS.
Rajbhandari, S., et al. (2020). “ZeRO: Memory Optimizations Toward Training Trillion Parameter Models.” SC20.
Narayanan, D., et al. (2021). “Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM.” SC21.
Li, S., et al. (2021). “Megatron-CM3: Multi-Context Prompted Training.” arXiv.
Zhao, J., et al. (2023). “PipeDream: Balanced Pipeline Parallelism for DNN Training.” SOSP.
Narayanan, D., et al. (2019). “Pipedream: Weight-Aware Pipeline Parallelism.” MLSys.
Rasley, J., et al. (2020). “DeepSpeed: System Optimizations for Training and Inference.” Blog.
Dao, T., et al. (2022). “FlashAttention: Fast and Memory-Efficient Exact Attention.” NeurIPS.
Chen, K., et al. (2023). “ColossalAI: A Unified Training Engine for Large-Scale Models.” GitHub.