跳转到内容
Jason's Study

EDM — 把扩散模型的训练配方一次拆清楚

是什么

EDM（Elucidating the Design Space of Diffusion-Based Generative Models，阐明扩散模型的设计空间）是 NVIDIA 的 Tero Karras 团队在 2022 年发的一篇方法论文。它解决一件具体的事：ddpm、ddim-2020、score-based-2020 三条流派把训练、采样、网络结构搅在一起，调参全靠玄学。EDM 把它们拆成正交模块，一项项消融，最后给出一套可落地的最优配方。

日常类比：

之前的扩散论文像”奶奶的家传菜谱”——盐少许、火候自己把握、按祖训来。EDM 是把它重写成实验室手册：盐 3.2g，180℃ 烤 12 分钟，原因是糖在 165℃ 焦糖化、蛋白在 70℃ 凝固。每一步独立测、独立调，最后写下确定的工序。

EDM 不是新模型，它是一份扩散训练的”标尺”。后来的 Stable Diffusion 3、SD-XL Turbo、stable-diffusion 改进版、Consistency Models、Flow Matching 全部直接继承它的预处理和采样器配方。

为什么重要

不理解 EDM，下面这些事都解释不通：

  • 为什么 2022 年之后所有扩散论文符号几乎都长一个样（sigma、c_skip、c_out）——EDM 把符号系统统一了
  • 为什么 SD3 的采样器叫”Heun” 而不是”DDIM”——Heun 是 EDM 推荐的二阶 ODE 求解器
  • 为什么训练时噪声 sigma 用 lognormal 分布而不是均匀分布——EDM 消融出来的最优分布
  • 为什么扩散模型 ImageNet-64 的 FID 从 2.07 一夜降到 1.36——不是新架构，是把配方调对了

核心要点

EDM 把扩散模型拆成 四个正交模块，每个独立消融、独立调：

  1. 采样器（ODE 离散方式）：从噪声走回干净图的过程是一个 ODE。用 Euler 一阶还是 Heun 二阶？时间步怎么取？EDM 推荐 Heun + rho=7 的非均匀网格，35 步顶过去 250 步。

  2. 训练时 sigma 噪声分布：训练时给图像加哪些 sigma 的噪声？均匀采太分散，要按”网络在哪个 sigma 最难学”加权。EDM 用 lognormal（mean=-1.2, std=1.2），让训练集中在中等噪声段。

  3. 网络预处理（preconditioning）：网络不直接预测噪声 epsilon，也不预测 x_0，而是输出一个被 c_skip / c_out / c_in / c_noise 四个系数包裹的中间量。这四个系数从”让网络输入输出方差归一”这一条原理推出来——纯数学，不调参。

  4. 损失加权：不同 sigma 的样本对损失的贡献权重。EDM 给出的权重保证训练梯度量级不随 sigma 漂移。

加一句关键事实：这四件事之前都被混在一起调。EDM 的贡献不是发明新模块，而是把它们解耦——这样后人可以替换其中任一个而不动其他。

实践案例

案例 1：preconditioning 系数怎么算

EDM 网络的输出形式是：

# x: 加噪图像，sigma: 当前噪声标准差
# F: U-Net，sigma_data: 数据集 std（CIFAR-10 约 0.5）
c_skip = sigma_data**2 / (sigma**2 + sigma_data**2)
c_out  = sigma * sigma_data / sqrt(sigma**2 + sigma_data**2)
c_in   = 1 / sqrt(sigma**2 + sigma_data**2)
c_noise = 0.25 * log(sigma)


D = c_skip * x + c_out * F(c_in * x, c_noise)

D 就是去噪后的 x_0 估计。这四个系数的来源是一条原则：让 F 的输入和输出方差恒为 1，无论 sigma 多大。这样网络在所有噪声尺度上看到的都是”标准化”输入，训练稳定、可迁移。

案例 2：Heun 二阶采样器为什么省步

Euler 一阶每步用当前点斜率走一步：

x = x + (sigma_next - sigma) * d_cur  # 一阶

Heun 二阶多算一次终点斜率，取平均：

x_euler = x + (sigma_next - sigma) * d_cur
d_next = denoise(x_euler, sigma_next)
x = x + (sigma_next - sigma) * 0.5 * (d_cur + d_next)  # 二阶

代价：每步多一次网络前向。收益：步数可以从 100+ 降到 35。总 NFE（神经网络评估次数）反而少——这就是 EDM 推荐 Heun 的原因。

案例 3：sigma 时间表 rho=7

采样时从 sigma_max=80 走到 sigma_min=0.002，但中间不要均匀分。EDM 的网格：

# i 从 0 到 N-1，N=35
sigma_i = (sigma_max**(1/rho) + i/(N-1) * (sigma_min**(1/rho) - sigma_max**(1/rho)))**rho

rho=7 让前期（高噪声）走得疏、后期（低噪声）走得密——因为低 sigma 时图像细节快速形成，需要更密的步长。rho=1 退化为线性，rho=7 是消融出来的最优。

案例 4：FID 数字背后的拆解

EDM 论文里有一张关键表，逐项打开每个改动带来的收益：

改动CIFAR-10 FID
baseline (DDPM++)2.51
+ 采样器换 Heun2.21
+ sigma 时间表 rho=72.04
+ 训练 sigma 用 lognormal1.94
+ preconditioning 重写1.79

每一项独立可见效，不是全调才有用。这是 EDM 最大的方法论贡献——它教会扩散圈子做正经的消融。

踩过的坑

  1. preconditioning 依赖 sigma_data。系数推导假定数据 std 已知。把 EDM 移植到新数据集时要先估 sigma_data，不然 c_skip / c_out / c_in 全错位，训练直接发散。

  2. Heun 在低步数时反而拉胯。每步多一次前向，步数 < 20 时一阶 Euler 反而总 NFE 更少。EDM 默认 35 步是甜点，别盲目套用到 10 步场景——那是 consistency-models 的地盘。

  3. 训练 sigma 用 lognormal 不是 uniform。不少复现版图省事用均匀采样，FID 直接掉 1-2 个点。这个细节最容易忘。

  4. sigma 视角和 alpha_bar 视角符号不互通。EDM 用 sigma-only（前向 x_t = x_0 + sigma * eps），DDPM 用 alpha_bar（前向 x_t = sqrt(a) * x_0 + sqrt(1-a) * eps）。两者数学等价但映射要小心，许多 bug 来自这里。

  5. rho=7 是 CIFAR/ImageNet 经验值。换分辨率（512×512 / 1024×1024）或换数据（人脸 / 文档）需要重调；rho 在 5-10 之间扫一遍是标准操作。

  6. EDM 的损失权重和 v-prediction 不一样。后人把 EDM 改成 v-parametrization（SD3 用），权重要重新推。直接套 EDM 默认权重训 v-pred 网络会让大 sigma 段欠拟合。

适用 vs 不适用场景

适用：

  • 想训新扩散模型，找一份”已被验证”的默认配方——直接抄 EDM
  • 已有 ddpm / ddim-2020 风格代码想升级——按四个模块逐项替换，可单独消融
  • 需要可解释的扩散数学——EDM 的符号系统是教材级，公式推导自洽

不适用：

  • 极低步数采样（1-4 步）——用 consistency-models / Rectified Flow / DMD 蒸馏，EDM 默认 35 步太多
  • 需要严格似然评估——EDM 重在 FID/采样质量，似然不是它的目标
  • 训练目标已经是 v-prediction 或 flow matching——EDM 的损失权重要重推，不能直接套

历史小故事（可跳过）

  • 2020 年 6 月：Ho 等发 ddpm，扩散登场。
  • 2020 年 10 月：Song 等发 ddim-2020 和 score-based-2020，三条流派并立。
  • 2021-2022 年：扩散论文井喷，但每篇都自带一套 hyperparameter 玄学，复现率低。
  • 2022 年 6 月：NVIDIA 的 Karras 团队发 EDM。Karras 是 StyleGAN / StyleGAN2 / StyleGAN3 的作者，他做研究出了名的”做完整消融、写干净代码”。EDM 是他第一次进扩散，一出手就把整个领域的方法论提了一档。
  • 2022 年底-2023 年：所有新扩散论文要么直接用 EDM 配方，要么以 EDM 为基线对比。Stability AI 训 SD3 时几乎全盘采纳。
  • 2024 年：Karras 团队跟进发 EDM2，把这套方法推到 ImageNet-512，FID 1.81 再创 SOTA。

学到什么

  1. 解耦 = 进步。把混在一起的设计拆成正交模块，每个单独消融，是科研基本功。EDM 的贡献不是任何一个模块，是让每个模块可以单独被改进。
  2. preconditioning 是数学，不是调参。从”输入输出方差归一”一条原理推出 c_skip / c_out / c_in，零超参。这种”由原理导系数”的做法值得学。
  3. 符号统一推动领域加速。EDM 之前每篇扩散论文符号都不一样，读起来痛苦；EDM 之后大家收敛到 sigma-only。一个清晰的写法等于给整个领域装了普通话。
  4. 消融表是论文的灵魂。EDM 的逐项 FID 表（baseline → +Heun → +rho → +lognormal → +precond）一目了然告诉读者每改动值多少分。这是研究论文该有的样子。
  5. 方法论论文 > 单点突破。EDM 没造新模型，但它把整个领域的研究和工程都往前推了一大步。“把已有东西拆清楚”和”造新东西”同等重要。

延伸阅读

关联

  • ddpm —— EDM 的起点；EDM 把 DDPM 的训练目标重写成 sigma 视角
  • ddim-2020 —— EDM 的采样器是 DDIM 的 ODE 视角的高阶推广（Heun 二阶）
  • score-based-2020 —— EDM 的数学框架直接继承 score-based 的 ODE/SDE 视角
  • stable-diffusion —— SD3 / SD-XL Turbo 大量采用 EDM 配方
  • classifier-free-guidance-2022 —— EDM 与 CFG 正交可叠加
  • consistency-models —— 在 EDM 框架上做蒸馏，把 35 步压到 1-4 步