EDM 推导总结

下面我们讲的 EDM 指的是 Karras 等人在论文 Elucidating the Design Space of Diffusion-Based Generative Models 中提出的 diffusion 统一框架。它不是一个全新的扩散模型，而更像是一次“把扩散模型拆开重装”的工作：把 网络参数化、训练噪声分布、损失权重、采样 ODE、采样步长调度、数值求解器 这些设计选择分离开来，然后给出一套更清晰、更稳定、更高效的推荐方案。

EDM 可以理解为：

用 $\sigma$ 作为统一噪声坐标，把扩散模型写成“从高噪声到低噪声解 ODE”的问题；网络学习在任意 $\sigma$ 下把带噪图像 $x$ 还原成干净图像 $y$；采样时根据 $x-D_\theta(x;\sigma)$ 估计噪声方向，然后沿 ODE 一步步去噪。

---

省流速通

EDM 的完整训练过程

训练时可以概括为：

1. 采样真实图像：$y \sim p_{\text{data}}$

2. 采样噪声等级：$\sigma \sim p_{\text{train}}(\sigma)$

   EDM 通常使用 log-normal 分布：$\log\sigma \sim \mathcal N(P_{\text{mean}}, P_{\text{std}}^2)$

3. 采样高斯噪声：$\epsilon \sim \mathcal N(0,I)$

4. 构造带噪样本：$x = y+\sigma\epsilon$

5. 计算网络输出：

$$D_\theta(x;\sigma) = c_{\text{skip}}x + c_{\text{out}} F_\theta(c_{\text{in}}x,c_{\text{noise}})$$

1. 最小化：

$$\mathcal L = \lambda(\sigma) |D_\theta(x;\sigma)-y|_2^2$$

也就是：

$$\boxed{ \mathcal L_{\text{EDM}} = \mathbb E_{\sigma,y,\epsilon} \left[ \frac{\sigma^2+\sigma_{\text{data}}^2} {(\sigma\sigma_{\text{data}})^2} \left| D_\theta(y+\sigma\epsilon;\sigma)-y \right|_2^2 \right] }$$

EDM 的完整采样过程

采样时：

1. 采样初始噪声：$x_0 \sim \mathcal N(0,\sigma_{\max}^2 I)$
2. 构造噪声序列：$\sigma_0 > \sigma_1 > \cdots > \sigma_N = 0$
3. 对每一步，计算：

$$d_i = \frac{x_i-D_\theta(x_i;\sigma_i)}{\sigma_i}$$

1. Euler 预测：

$$\hat x_{i+1} = x_i + (\sigma_{i+1}-\sigma_i)d_i$$

1. Heun 修正：

$$d'_{i+1} = \frac{ \hat x_{i+1} - D_\theta(\hat x_{i+1};\sigma_{i+1}) } {\sigma_{i+1}}$$

1. 更新：

$$x_{i+1} = x_i + (\sigma_{i+1}-\sigma_i) \frac{d_i+d'_{i+1}}{2}$$

当 $\sigma_{i+1}=0$ 时，最后一步通常只做 Euler，避免除以 0。

---

后面将会比较详细的说明 EDM 的具体过程推导

EDM 的核心思想

传统 DDPM 通常写成：

$$q(x_t \mid x_0) = \mathcal N \left( x_t; \sqrt{\bar\alpha_t}x_0, (1-\bar\alpha_t)I \right)$$

也就是：

$$x_t = \sqrt{\bar\alpha_t}x_0 + \sqrt{1-\bar\alpha_t}\epsilon, \qquad \epsilon \sim \mathcal N(0,I)$$

EDM 换了一个更直观的写法：直接用噪声标准差 $\sigma$ 表示噪声强度。

令真实数据为 $y$，加噪后的样本为：

$$x = y + n, \qquad n \sim \mathcal N(0, \sigma^2 I)$$

也就是：

$$x = y + \sigma \epsilon, \qquad \epsilon \sim \mathcal N(0,I)$$

这里的 $\sigma$ 就是噪声等级。 $\sigma$ 大，图像几乎全是噪声；$\sigma$ 小，图像接近真实数据。

所以 EDM 的基本建模对象是：

$$p(x;\sigma) = p_{\text{data}}(y) * \mathcal N(0,\sigma^2 I)$$

也就是“真实数据分布经过高斯模糊后的分布”。

---

从 score 到 denoiser

扩散模型里有一个非常核心的量，叫 score：

$$\nabla_x \log p(x;\sigma)$$

它表示：在噪声等级 $\sigma$ 下，当前点 $x$ 应该往哪个方向移动，才能更接近高概率区域。

但 EDM 不直接让网络预测 score，而是让网络预测 denoised image，即：

$$D_\theta(x;\sigma) \approx y$$

也就是：给定带噪图像 $x = y + \sigma\epsilon$，网络输出干净图像 $y$ 的估计。

为什么 denoiser 和 score 是等价的？这里用到 Tweedie 公式：

$$\nabla_x \log p(x;\sigma) = \frac{\mathbb E[y \mid x] - x}{\sigma^2}$$

如果我们用神经网络 $D_\theta(x;\sigma)$ 近似条件期望：

$$D_\theta(x;\sigma) \approx \mathbb E[y \mid x]$$

那么 score 可以写成：

$$\nabla_x \log p(x;\sigma) \approx \frac{D_\theta(x;\sigma)-x}{\sigma^2}$$

反过来也可以写成：

$$D_\theta(x;\sigma) \approx x + \sigma^2 \nabla_x \log p(x;\sigma)$$

所以：

预测 denoised image，本质上等价于预测 score。

这点非常重要。EDM 之所以喜欢用 denoiser，是因为它在训练和采样时更稳定、更直观。

---

概率流 ODE 推导

EDM 把采样过程看成一个连续时间 ODE。我们从高噪声开始：

$$x(\sigma_{\max}) \sim \mathcal N(0,\sigma_{\max}^2 I)$$

然后逐渐降低噪声：

$$\sigma_{\max} \to \sigma_{\min} \to 0$$

最终得到干净图像。

对于 variance exploding 形式的扩散过程，有概率流 ODE：

$$\frac{dx}{dt} = -\dot\sigma(t)\sigma(t) \nabla_x \log p(x;\sigma(t))$$

该方程具体推导过程见文末

如果令时间变量直接等于噪声强度：

$$t = \sigma$$

那么：

$$\dot\sigma(t)=1$$

于是：

$$\frac{dx}{d\sigma} = -\sigma \nabla_x \log p(x;\sigma)$$

代入刚才的 score-denoiser 关系：

$$\nabla_x \log p(x;\sigma) \approx \frac{D_\theta(x;\sigma)-x}{\sigma^2}$$

得到：

$$\frac{dx}{d\sigma} = -\sigma \cdot \frac{D_\theta(x;\sigma)-x}{\sigma^2}$$

化简：

$$\frac{dx}{d\sigma} = \frac{x - D_\theta(x;\sigma)}{\sigma}$$

这就是 EDM 采样 ODE 的核心形式：

$$\boxed{ \frac{dx}{d\sigma} = \frac{x - D_\theta(x;\sigma)}{\sigma} }$$

它的含义很漂亮：

• $x$ 是当前带噪图像；
• $D_\theta(x;\sigma)$ 是网络估计的干净图像；
• $x-D_\theta(x;\sigma)$ 是“当前噪声成分”；
• 除以 $\sigma$ 后得到沿噪声尺度变化的速度。

所以采样过程就是沿着这个 ODE，从大 $\sigma$ 积分到小 $\sigma$。

---

EDM 的采样公式

假设当前噪声等级是 $\sigma_i$，下一个噪声等级是 $\sigma_{i+1}$，其中：

$$\sigma_i > \sigma_{i+1}$$

当前状态为 $x_i$。根据 ODE：

$$\frac{dx}{d\sigma} = d_i = \frac{x_i - D_\theta(x_i;\sigma_i)}{\sigma_i}$$

用一阶 Euler 方法：

$$x_{i+1} = x_i + (\sigma_{i+1}-\sigma_i)d_i$$

即：

$$\boxed{ x_{i+1} = x_i + (\sigma_{i+1}-\sigma_i) \frac{x_i-D_\theta(x_i;\sigma_i)}{\sigma_i} }$$

因为 $\sigma_{i+1}-\sigma_i < 0$，所以这一步实际上是在“减去噪声方向”。

• DDPM / DDIM 的习惯（倒计时）：使用的是物理时间步 $t$。定义 $t = 0$ 是干净图像，$t = T$（如 1000）是纯噪声。因此，在采样生成时，时间是倒着走的，从 $t = T \rightarrow T - 1 \rightarrow \cdots \rightarrow 0$。
• EDM 的习惯（正向计步）：使用的是算法循环的步数计数器 $i$。定义 $i = 0$ 是采样的开始（纯噪声），$i = N$（如 30）是采样结束（干净图像）。因此，在采样生成时，计数器是正着走的，从 $i = 0 \rightarrow 1 \rightarrow \cdots \rightarrow N$。

EDM 更推荐用二阶 Heun 方法。先做 Euler 预测：

$$\hat x_{i+1} = x_i + (\sigma_{i+1}-\sigma_i)d_i$$

然后在新位置计算新的斜率：

$$d_{i+1} = \frac{ \hat x_{i+1} - D_\theta(\hat x_{i+1};\sigma_{i+1}) } {\sigma_{i+1}}$$

最后取两个斜率的平均：

$$x_{i+1} = x_i + (\sigma_{i+1}-\sigma_i) \frac{d_i+d_{i+1}}{2}$$

也就是：

$$\boxed{ x_{i+1} = x_i + (\sigma_{i+1}-\sigma_i) \frac{1}{2} \left[ \frac{x_i-D_\theta(x_i;\sigma_i)}{\sigma_i} + \frac{\hat x_{i+1}-D_\theta(\hat x_{i+1};\sigma_{i+1})}{\sigma_{i+1}} \right] }$$

这就是 EDM 采样算法的核心。论文和官方实现也提供了对应代码与采样器。

---

EDM 的噪声调度

EDM 认为采样时 $\sigma$ 的选择非常关键。不是简单线性下降，而是使用 Karras noise schedule：

$$\sigma_i = \left( \sigma_{\max}^{1/\rho} + \frac{i}{N-1} \left( \sigma_{\min}^{1/\rho} - \sigma_{\max}^{1/\rho} \right) \right)^\rho$$

其中：

$$i = 0,1,\dots,N-1$$

通常 $\rho = 7$。

直觉上，这个 schedule 会在小噪声区域分配更多采样步数。 为什么？因为图像细节主要在低噪声阶段形成。高噪声阶段只决定大致结构，小噪声阶段负责纹理、边缘、颜色、局部细节。

所以 EDM 的采样噪声序列是：

$$\sigma_0 = \sigma_{\max} > \sigma_1 > \dots > \sigma_{N-1} = \sigma_{\min} > \sigma_N = 0$$

最后一步到 $\sigma=0$，得到最终图像。

---

EDM 的训练目标

EDM 的训练非常直接。

从数据集中采样干净图像：

$$y \sim p_{\text{data}}$$

采样噪声等级：

$$\sigma \sim p_{\text{train}}(\sigma)$$

采样标准高斯噪声：

$$\epsilon \sim \mathcal N(0,I)$$

构造带噪输入：

$$x = y + \sigma\epsilon$$

然后训练网络：

$$D_\theta(x;\sigma) \approx y$$

最基本的 denoising loss 是：

$$\mathcal L = \mathbb E_{y,\epsilon,\sigma} \left[ |D_\theta(y+\sigma\epsilon;\sigma)-y|_2^2 \right]$$

但 EDM 加了一个噪声相关的权重：

$$\mathcal L = \mathbb E_{y,\epsilon,\sigma} \left[ \lambda(\sigma) |D_\theta(y+\sigma\epsilon;\sigma)-y|_2^2 \right]$$

即：

$$\boxed{ \mathcal L_{\text{EDM}} = \mathbb E_{\sigma,y,\epsilon} \left[ \lambda(\sigma) \left| D_\theta(y+\sigma\epsilon;\sigma) - y \right|_2^2 \right] }$$

这里 $\lambda(\sigma)$ 用来平衡不同噪声等级的训练难度。

EDM 的预条件化设计

它不直接让网络输出 $D_\theta(x;\sigma)$，而是写成：

$$D_\theta(x;\sigma) = c_{\text{skip}}(\sigma)x + c_{\text{out}}(\sigma) F_\theta \left( c_{\text{in}}(\sigma)x, c_{\text{noise}}(\sigma) \right)$$

其中 $F_\theta$ 是真正的神经网络。

具体推导细节见文末

这个式子的意思是：

• $c_{\text{skip}}(\sigma)x$：保留一部分输入；
• $F_\theta$：预测需要补充或修正的部分；
• $c_{\text{in}}(\sigma)$：把输入标准化；
• $c_{\text{out}}(\sigma)$：控制网络输出尺度；
• $c_{\text{noise}}(\sigma)$：把噪声等级编码后送入网络。

EDM 推荐的系数是：

$$c_{\text{skip}}(\sigma) = \frac{\sigma_{\text{data}}^2} {\sigma^2+\sigma_{\text{data}}^2}$$ $$c_{\text{out}}(\sigma) = \frac{\sigma\sigma_{\text{data}}} {\sqrt{\sigma^2+\sigma_{\text{data}}^2}}$$ $$c_{\text{in}}(\sigma) = \frac{1} {\sqrt{\sigma^2+\sigma_{\text{data}}^2}}$$ $$c_{\text{noise}}(\sigma) = \frac{1}{4}\log \sigma$$

其中 $\sigma_{\text{data}}$ 表示数据本身的标准差。对于归一化到大约 $[-1,1]$ 的图像，论文中常用：

$$\sigma_{\text{data}} = 0.5$$

---

一些细节推导与理解

1. $\frac{dx}{dt} = -\dot{\sigma}(t)\sigma(t)\nabla_x \log p(x; \sigma(t))$

EDM 这里的推导，其实是直接借用了 SDE 框架中关于“方差爆炸（Variance Exploding, VE）”模型的现成结论。

• 第一步：定义方差爆炸（VE）的前向过程

在连续时间 $t$ 下，假设一个前向加噪过程。对于 Variance Exploding (VE) 模型，它的特点是不衰减原图信号，只不断往里面加方差越来越大的噪声。

写成随机微分方程（SDE）的标准形式：

$$dx = f(x, t)dt + g(t)dw$$

• $f(x, t)$ 是漂移项（Drift coefficient），控制信号的衰减趋势。
• $g(t)$ 是扩散项（Diffusion coefficient），控制加入噪声的强度。
• $w$ 是标准布朗运动（维纳过程）。

因为 VE 模型不缩放原图信号，所以漂移项为 0，即 $f(x, t) = 0$。于是它的前向 SDE 形式非常简单：

$$dx = g(t)dw$$

• 第二步：确定扩散项 $g(t)$ 到底是什么

我们知道，VE 模型在时刻 $t$ 的边缘分布，相当于给原图加了方差为 $\sigma^2(t)$ 的噪声。也就是说，在时间从 0 积分到 $t$ 的过程中，累积的方差必须等于 $\sigma^2(t)$。

根据随机微积分的性质，SDE $dx = g(t)dw$ 在 $dt$ 时间内产生的方差增量是 $g^2(t)dt$。因此，总方差对时间的导数，就是 $g^2(t)$：

$$\frac{d}{dt}[\sigma^2(t)] = g^2(t)$$

利用微积分的链式法则展开左边：

$$2\sigma(t)\dot{\sigma}(t) = g^2(t)$$

这就是 VE 模型中极其关键的关系式。此时，我们可以把前向 SDE 具体写为：

$$dx = \sqrt{2\sigma(t)\dot{\sigma}(t)}dw$$

• 第三步：引入“概率流常微分方程（Probability Flow ODE）”定理

这里需要用到 Yang Song 论文中最核心的定理之一：对于任意一个 SDE，都存在一个对应的常微分方程（ODE），它们的边缘概率密度函数 $p_t(x)$ 随着时间的演化是完全一样的。

对于一般形式的 SDE：$dx = f(x, t)dt + g(t)dw$ 它所对应的概率流 ODE (Probability Flow ODE) 的通用公式被证明为：

$$dx = \left[ f(x, t) - \frac{1}{2}g^2(t)\nabla_x \log p_t(x) \right] dt$$

这个公式的物理意义是：我们可以把原本带有随机性（$dw$）的布朗运动，等效转换成一个完全确定性的流体运动（ODE），只要我们沿着对数概率密度的梯度（也就是 Score，$\nabla_x \log p_t(x)$）进行引导。

• 第四步：代入 VE 模型的参数完成推导

现在，已知 VE 模型的参数为：

1. $f(x, t) = 0$
2. $g^2(t) = 2\sigma(t)\dot{\sigma}(t)$

代入概率流 ODE 的通用公式：

$$dx = \left[ 0 - \frac{1}{2} \left( 2\sigma(t)\dot{\sigma}(t) \right) \nabla_x \log p_t(x) \right] dt$$

化简括号里的式子，$\frac{1}{2}$ 和 2 抵消：

$$dx = \left[ -\sigma(t)\dot{\sigma}(t) \nabla_x \log p_t(x) \right] dt$$

两边同除以 $dt$，就得到了你在图中看到的那行概率流 ODE 核心公式：

$$\frac{dx}{dt} = -\dot{\sigma}(t)\sigma(t)\nabla_x \log p(x; \sigma(t))$$

（注：这里的 $p_t(x)$ 和图中的 $p(x; \sigma(t))$ 是同一种东西的不同写法，都代表在时间 $t$ 或噪声级别 $\sigma(t)$ 下的边缘概率分布）。

EDM 巧妙的地方在于，它拿到这个现成的结论后，做了一个极其简单却绝妙的变量替换——直接让 $t = \sigma$，从而消灭了对抽象时间 $t$ 的依赖，把整个微分方程变成了完全由“噪声强度 $\sigma$”主导的简洁系统。

1. $D_\theta(x;\sigma) = c_{\text{skip}}(\sigma)x + c_{\text{out}}(\sigma) F_\theta \left( c_{\text{in}}(\sigma)x, c_{\text{noise}}(\sigma) \right)$

我们来推一下它的直觉来源。

假设数据满足：

$$y \sim \mathcal N(0,\sigma_{\text{data}}^2 I)$$

加噪后：

$$x = y + \sigma\epsilon$$

所以：

$$\operatorname{Var}(x) = \sigma_{\text{data}}^2 + \sigma^2$$

为了让网络输入在不同噪声等级下都有相近方差，希望：

$$\operatorname{Var}(c_{\text{in}}x) \approx 1$$

因此：

$$c_{\text{in}}^2 (\sigma_{\text{data}}^2+\sigma^2) = 1$$

所以：

$$\boxed{ c_{\text{in}} = \frac{1} {\sqrt{\sigma_{\text{data}}^2+\sigma^2}} }$$

接着考虑最优线性估计。若 $y$ 和噪声都是高斯，则：

$$\mathbb E[y\mid x] = \frac{\sigma_{\text{data}}^2} {\sigma_{\text{data}}^2+\sigma^2} x$$

我们已知：

1. 真实数据（先验）：$y \sim \mathcal{N}(0, \sigma_{\text{data}}^2 I)$
2. 观测数据：$x = y + \sigma \epsilon$，其中噪声 $\epsilon \sim \mathcal{N}(0, I)$，且 $y$ 与 $\epsilon$ 相互独立。

因为 $y$ 和 $\epsilon$ 都是高斯分布，且它们是线性组合的，所以 $x$ 和 $y$ 服从联合高斯分布（Joint Gaussian Distribution）。

在多变量高斯分布中，如果已知变量 $X$ 的值，要估计另一个变量 $Y$ 的条件期望 $\mathbb{E}[Y \mid X]$，有一个标准的通用公式：

$$\mathbb{E}[Y \mid X] = \mu_Y + \frac{\text{Cov}(Y, X)}{\text{Var}(X)} (X - \mu_X)$$

现在，我们把具体的参数代入这个公式中进行计算：

• $\mu_Y = \mathbb{E}[y] = 0$
• $\mu_X = \mathbb{E}[x] = \mathbb{E}[y + \sigma \epsilon] = \mathbb{E}[y] + \sigma \mathbb{E}[\epsilon] = 0 + 0 = 0$

因为 $y$ 和 $\epsilon$ 是独立的，所以方差可以直接相加：

• $\text{Var}(x) = \text{Var}(y + \sigma \epsilon) = \text{Var}(y) + \text{Var}(\sigma \epsilon) = \sigma_{\text{data}}^2 + \sigma^2$

计算 $x$ 和 $y$ 的协方差（$\text{Cov}(Y, X)$）：

• $\text{Cov}(y, x) = \mathbb{E}[(y - \mu_Y)(x - \mu_X)^T] = \mathbb{E}[y x^T]$
• 将 $x = y + \sigma \epsilon$ 代入： $\mathbb{E}[y (y + \sigma \epsilon)^T] = \mathbb{E}[y y^T] + \mathbb{E}[y \sigma \epsilon^T]$
• 因为 $y$ 和 $\epsilon$ 独立，$\mathbb{E}[y \epsilon^T] = 0$。而 $\mathbb{E}[y y^T]$ 就是 $y$ 的方差： $\text{Cov}(y, x) = \sigma_{\text{data}}^2$

将上述结果代入标准条件期望公式：

$$\mathbb{E}[y \mid x] = 0 + \frac{\sigma_{\text{data}}^2}{\sigma_{\text{data}}^2 + \sigma^2} (x - 0)$$ $$\mathbb{E}[y \mid x] = \frac{\sigma_{\text{data}}^2}{\sigma_{\text{data}}^2 + \sigma^2} x$$

所以自然得到 skip connection：

$$\boxed{ c_{\text{skip}} = \frac{\sigma_{\text{data}}^2} {\sigma_{\text{data}}^2+\sigma^2} }$$

这很好理解：

当 $\sigma$ 很小时：

$$c_{\text{skip}} \approx 1$$

此时输入 $x$ 已经很干净，所以直接保留输入。

当 $\sigma$ 很大时：

$$c_{\text{skip}} \approx 0$$

此时输入几乎全是噪声，不能相信输入本身。

然后 $c_{\text{out}}$ 的作用是让网络 $F_\theta$ 的输出目标具有大致稳定的尺度。EDM 选择：

$$\boxed{ c_{\text{out}} = \frac{\sigma\sigma_{\text{data}}} {\sqrt{\sigma^2+\sigma_{\text{data}}^2}} }$$

这样可以让网络在所有 $\sigma$ 下看到的输入和输出尺度都比较稳定。

最后，EDM 的 loss weight 选择为：

$$\lambda(\sigma) = \frac{1} {c_{\text{out}}(\sigma)^2}$$

代入 $c_{\text{out}}$：

$$\lambda(\sigma) = \frac{\sigma^2+\sigma_{\text{data}}^2} {\sigma^2\sigma_{\text{data}}^2}$$

也就是：

$$\boxed{ \lambda(\sigma) = \frac{\sigma^2+\sigma_{\text{data}}^2} {(\sigma\sigma_{\text{data}})^2} }$$

这样做的目的，是让不同噪声等级对训练的贡献更平衡。

1. EDM 与 DDPM / DDIM 的关系

可以这样理解：

DDPM 写法：

$$x_t = \sqrt{\bar\alpha_t}x_0 + \sqrt{1-\bar\alpha_t}\epsilon$$

把两边除以 $\sqrt{\bar\alpha_t}$：

$$\frac{x_t}{\sqrt{\bar\alpha_t}} = x_0 + \sqrt{\frac{1-\bar\alpha_t}{\bar\alpha_t}}\epsilon$$

令：

$$\sigma_t = \sqrt{ \frac{1-\bar\alpha_t} {\bar\alpha_t} }$$

以及：

$$\tilde x_t = \frac{x_t}{\sqrt{\bar\alpha_t}}$$

那么：

$$\tilde x_t = x_0 + \sigma_t \epsilon$$

这就变成了 EDM 的形式：

$$x = y + \sigma\epsilon$$

所以 EDM 不是和 DDPM 完全不同的东西，而是把 DDPM 的加噪过程重新参数化为 $\sigma$ 形式。

DDIM 也可以放进 EDM 的 ODE 框架里。DDIM 的确定性采样本质上也是在解一个 probability flow ODE，只是它的时间参数、噪声调度、网络预测形式和 EDM 不一样。