EDM

前置

• 分数匹配

1. EDM 的扩散参数化

EDM（Elucidating the Design Space of Diffusion-Based Generative Models）
(Karras et al., 2022) 把扩散过程、网络预条件、损失权重、训练噪声分布和采样日程解耦。

考虑 $s(t)>0$、$\sigma (t)\ge 0$ 的参数化

$${\mathbf{X}}_t=s(t){\mathbf{X}}_0+s(t)\sigma (t)\mathbf{ϵ},\mathbf{ϵ}\sim \mathcal{N}(0,\mathbf{I}).$$

令

$${\mathbf{u}}_t:=\frac{{\mathbf{X}}_t}{s(t)}={\mathbf{X}}_0+\sigma (t)\mathbf{ϵ}.$$

若 $p(\mathbf{u};\sigma )$ 是 ${\mathbf{X}}_0+\sigma \mathbf{ϵ}$ 的密度，则

$$p_t(\mathbf{x})=s(t{)}^{-d}p\left(\frac{\mathbf{x}}{s(t)};\sigma (t)\right)$$

以及

$${\nabla }_{\mathbf{x}}\log p_t(\mathbf{x})=\frac{1}{s(t)}{\nabla }_{\mathbf{u}}\log p(\mathbf{u};\sigma (t)){|}_{\mathbf{u}=\mathbf{x}/s(t)}.$$

条件高斯分数为

$${\nabla }_{\mathbf{x}}\log p_{0t}(\mathbf{x}\mid {\mathbf{x}}_0)=\frac{1}{s(t)\sigma (t{)}^2}\left({\mathbf{x}}_0-\frac{\mathbf{x}}{s(t)}\right).$$

若去噪器

$$D_{\theta }(\mathbf{u},\sigma )\approx \mathbb{E}[{\mathbf{X}}_0\mid {\mathbf{X}}_0+\sigma \mathbf{ϵ}=\mathbf{u}],$$

则 Tweedie 公式给出

$${\nabla }_{\mathbf{u}}\log p(\mathbf{u};\sigma )\approx \frac{D_{\theta }(\mathbf{u},\sigma )-\mathbf{u}}{{\sigma }^2}.$$

2. EDM 的 probability-flow ODE

由线性高斯扩散的一般结论，

$$\boxed{\mathrm{d}{\mathbf{X}}_t=\left[\frac{s^{\prime }(t)}{s(t)}{\mathbf{X}}_t-s(t)\sigma (t){\sigma }^{\prime }(t){\nabla }_{\mathbf{u}}\log p(\mathbf{u};\sigma (t)){|}_{\mathbf{u}={\mathbf{X}}_t/s(t)}\right]\mathrm{d}t}.$$

代入去噪器后，

$$\boxed{\frac{\mathrm{d}{\mathbf{X}}_t}{\mathrm{d}t}=\left(\frac{s^{\prime }(t)}{s(t)}+\frac{{\sigma }^{\prime }(t)}{\sigma (t)}\right){\mathbf{X}}_t-s(t)\frac{{\sigma }^{\prime }(t)}{\sigma (t)}{D}_{\theta }\left(\frac{{\mathbf{X}}_t}{s(t)},\sigma (t)\right)}.$$

这对应 EDM 论文的式 (4)。EDM 常取

$$s(t)\equiv 1,$$

于是

$$\frac{\mathrm{d}{\mathbf{X}}_t}{\mathrm{d}t}=\frac{{\sigma }^{\prime }(t)}{\sigma (t)}\left[{\mathbf{X}}_t-D_{\theta }({\mathbf{X}}_t,\sigma (t))\right].$$

再取 $\sigma (t)=t>0$，并定义网络分数

$${\mathbf{s}}_{\theta }(\mathbf{x},t):=\frac{D_{\theta }(\mathbf{x},t)-\mathbf{x}}{t^2},$$

则

$$\boxed{\frac{\mathrm{d}{\mathbf{X}}_t}{\mathrm{d}t}=\frac{{\mathbf{X}}_t-D_{\theta }({\mathbf{X}}_t,t)}{t}=-t{\mathbf{s}}_{\theta }({\mathbf{X}}_t,t)}.$$

采样时沿噪声尺度从大到小积分，通常使用 Euler 或二阶 Heun 修正。

3. 去噪器预条件

EDM 不让主网络直接承担所有尺度变化，而定义

$$\boxed{D_{\theta }(\mathbf{x},\sigma )=c_{\mathrm{skip}}(\sigma )\mathbf{x}+c_{\mathrm{out}}(\sigma )F_{\theta }\left(c_{\mathrm{in}}(\sigma )\mathbf{x},c_{\mathrm{noise}}(\sigma )\right)}.$$

四个系数分别负责：

• $c_{\mathrm{in}}$：把网络输入缩放到近似单位方差；
• $c_{\mathrm{skip}}$：提供随噪声变化的线性去噪基线；
• $c_{\mathrm{out}}$：把网络输出恢复到合适的残差尺度；
• $c_{\mathrm{noise}}$：把噪声标准差编码为网络条件。

设干净数据 $\mathbf{y}$ 零均值、每个分量方差
${\sigma }_{\mathrm{data}}^2$，噪声
$\mathbf{n}\sim \mathcal{N}(0,{\sigma }^2\mathbf{I})$ 与数据不相关，并令

$$\mathbf{x}=\mathbf{y}+\mathbf{n}.$$

EDM 的系数为

$$\boxed{c_{\mathrm{skip}}(\sigma )=\frac{{\sigma }_{\mathrm{data}}^2}{{\sigma }^2+{\sigma }_{\mathrm{data}}^2}},$$ $$\boxed{c_{\mathrm{out}}(\sigma )=\frac{\sigma {\sigma }_{\mathrm{data}}}{\sqrt{{\sigma }^2+{\sigma }_{\mathrm{data}}^2}}},$$ $$\boxed{c_{\mathrm{in}}(\sigma )=\frac{1}{\sqrt{{\sigma }^2+{\sigma }_{\mathrm{data}}^2}},c_{\mathrm{noise}}(\sigma )=\frac{1}{4}\log \sigma }.$$

3.1 输入和输出尺度

因为

$$\mathrm{Var}(x_j)={\sigma }_{\mathrm{data}}^2+{\sigma }^2,$$

$c_{\mathrm{in}}$ 使输入每个分量具有近似单位方差。

去掉 skip 后，网络要预测的目标是

$$\mathbf{r}=\frac{\mathbf{y}-c_{\mathrm{skip}}\mathbf{x}}{c_{\mathrm{out}}}.$$

用下节得到的 $c_{\mathrm{skip}}$ 计算，

$$\mathrm{Var}(\mathbf{y}-c_{\mathrm{skip}}\mathbf{x})=\frac{{\sigma }^2{\sigma }_{\mathrm{data}}^2}{{\sigma }^2+{\sigma }_{\mathrm{data}}^2},$$

其标准差正是 $c_{\mathrm{out}}$，所以网络目标在各噪声级保持近似单位尺度。

3.2 $c_{\mathrm{skip}}$ 的 LMMSE 推导

在标量线性估计族

$$\hat{\mathbf{y}}=a\mathbf{x}$$

中最小化

$$J(a)=\mathbb{E}\Vert \mathbf{y}-a\mathbf{x}{\Vert }_2^2.$$

求导并令零：

$$\begin{array}{rl}{J}^{\prime }(a)& =-2\mathbb{E}[{\mathbf{y}}^{\top }\mathbf{x}]+2a\mathbb{E}[{\mathbf{x}}^{\top }\mathbf{x}]=0,\\ a^{\star }& =\frac{\mathbb{E}[{\mathbf{y}}^{\top }\mathbf{x}]}{\mathbb{E}[{\mathbf{x}}^{\top }\mathbf{x}]}.\end{array}$$

由 $\mathbf{x}=\mathbf{y}+\mathbf{n}$ 及不相关性，

$$\mathbb{E}[{\mathbf{y}}^{\top }\mathbf{x}]=d{\sigma }_{\mathrm{data}}^2,$$ $$\mathbb{E}[{\mathbf{x}}^{\top }\mathbf{x}]=d({\sigma }_{\mathrm{data}}^2+{\sigma }^2).$$

因此

$$a^{\star }=\frac{{\sigma }_{\mathrm{data}}^2}{{\sigma }_{\mathrm{data}}^2+{\sigma }^2}=c_{\mathrm{skip}}(\sigma ).$$

所以当主网络输出为零时，skip 分支已经是该各向同性标量线性类中的 LMMSE 估计。它不声称对任意非高斯数据都是完整 MMSE 后验均值；非线性残差由网络学习。

4. 训练目标与权重

EDM 的去噪损失为

$$\mathcal{L}={\mathbb{E}}_{\sigma ,\mathbf{y},\mathbf{n}}\left[\lambda (\sigma ){\Vert D_{\theta }(\mathbf{y}+\mathbf{n},\sigma )-\mathbf{y}\Vert }_2^2\right],\mathbf{n}\sim \mathcal{N}(0,{\sigma }^2\mathbf{I}).$$

选择

$$\boxed{\lambda (\sigma )=\frac{1}{c_{\mathrm{out}}(\sigma {)}^2}=\frac{{\sigma }^2+{\sigma }_{\mathrm{data}}^2}{{\sigma }^2{\sigma }_{\mathrm{data}}^2}}$$

后，改写成网络输出空间：

$$\begin{array}{rl}\mathcal{L}=\mathbb{E}[& \lambda (\sigma )c_{\mathrm{out}}(\sigma {)}^2\\ & \cdot {\Vert F_{\theta }(c_{\mathrm{in}}\mathbf{x},c_{\mathrm{noise}})-\frac{\mathbf{y}-c_{\mathrm{skip}}\mathbf{x}}{c_{\mathrm{out}}}\Vert }_2^2].\end{array}$$

有效前因子
$\lambda c_{\mathrm{out}}^2=1$，同时目标方差也被归一化，从而减弱噪声尺度造成的训练失衡。

5. 采样噪声日程

取 $N$ 个正噪声级，按采样方向从
${\sigma }_{\max }$ 降到 ${\sigma }_{\min }$：

$$\boxed{{\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,\dots ,N-1}.$$

再附加

$${\sigma }_N=0.$$

于是

$${\sigma }_0={\sigma }_{\max }>\cdots >{\sigma }_{N-1}={\sigma }_{\min }>{\sigma }_N=0.$$

$\rho$ 控制步数在噪声尺度上的分配；EDM 图像实验常取
$\rho =7$，RFD3 (Butcher et al., 2025) 实现也取 $7$。

原图 caption：EDM 的时间步设计。
图意是沿
${\sigma }_{\max }\to {\sigma }_{\min }\to 0$
反向积分；上式明确了数组方向，避免把 ${\sigma }_0=0$ 与采样起点混在一起。

6. ${\sigma }_{\mathrm{data}}$ 的含义

${\sigma }_{\mathrm{data}}$ 是用于预条件的典型数据标准差尺度。在各维近似同尺度时，可取训练集按分量计算的经验 RMS/标准差并汇总为标量；它不是由数据取值区间唯一决定的常数。

若像素被归一化到 $[-1,1]$：

• 若真是均匀分布，标准差为 $$\sqrt{\frac{(1-(-1){)}^2}{12}}=\frac{1}{\sqrt{3}}\approx 0.577;$$
• “约 $95\%$ 数据落在 $\pm 2\sigma$”的正态经验会给出
  ${\sigma }_{\mathrm{data}}\approx 0.5$，但这只是尺度启发，不是由区间推出的定理。

EDM 图像设置常用

$${\sigma }_{\mathrm{data}}=\mathrm{0.5.}$$

对蛋白质坐标等不在 $[-1,1]$ 的数据，应根据实际数据尺度估计，而不是照搬 $0.5$。

7. 训练时的 lognormal 噪声采样

EDM 令

$$\log \sigma \sim \mathcal{N}(P_{\mathrm{mean}},P_{\mathrm{std}}^2),$$

等价于

$$\sigma =\exp (P_{\mathrm{mean}}+P_{\mathrm{std}}ϵ),ϵ\sim \mathcal{N}(0,1).$$

这里

$$\mathrm{median}(\sigma )=e^{P_{\mathrm{mean}}},$$

而均值为

$$\mathbb{E}[\sigma ]=\exp \left(P_{\mathrm{mean}}+\frac{1}{2}{P}_{\mathrm{std}}^2\right).$$

因此 $e^{P_{\mathrm{mean}}}$ 应称为中位噪声尺度，不是均值或“精修尺度”。

EDM 图像参数

$$P_{\mathrm{mean}}=-1.2,P_{\mathrm{std}}=1.2$$

对应：

• $\log \sigma$ 落在均值 $\pm 1$ 个标准差时，
  $\sigma \in [e^{-2.4},e^0]\approx [0.0907,1]$；
• 落在 $\pm 2$ 个标准差时，
  $\sigma \in [e^{-3.6},e^{1.2}]\approx [0.0273,3.32]$。

来源中的 RFD3 参数

$$P_{\mathrm{mean}}=-1.2,P_{\mathrm{std}}=1.5$$

对应：

• $\pm 1$ 个标准差：
  $\sigma \in [e^{-2.7},e^{0.3}]\approx [0.0672,1.35]$；
• $\pm 2$ 个标准差：
  $\sigma \in [e^{-4.2},e^{1.8}]\approx [0.0150,6.05]$。

这些区间不是 lognormal 分布的支持集或硬截断“范围”：分布支持是
$(0,\infty )$。它们分别对应正态变量约 $68.27\%$ 与 $95.45\%$ 的中心概率区间。

若实现采用

$$\sigma ={\sigma }_{\mathrm{data}}\exp (P_{\mathrm{mean}}+P_{\mathrm{std}}ϵ),$$

则所有分位数整体乘以 ${\sigma }_{\mathrm{data}}$，适合把无量纲日程映射到非图像数据的物理尺度；这与原始 EDM 不带该乘子时的参数命名应分开记录。

Butcher, J., Krishna, R., Mitra, R., Brent, R. I., Li, Y., Corley, N., Kim, P. T., Funk, J., Mathis, S., Salike, S., & others. (2025). De Novo Design of All-Atom Biomolecular Interactions with RFdiffusion3. bioRxiv.

Karras, T., Aittala, M., Aila, T., & Laine, S. (2022). Elucidating the Design Space of Diffusion-Based Generative Models. Advances in Neural Information Processing Systems, 35, 26565–26577.