2025-CVPR-Adaptive Non-Uniform Timestep Sampling for Diffusion Model Training全文翻译

全文翻译

摘要

作为一种具有高表达能力的生成模型，扩散模型已在包括图像生成、自然语言处理和组合优化等多个领域展现出卓越的成功。然而，随着数据分布变得愈发复杂，将这些模型训练至收敛所需的计算资源也日益增加。尽管扩散模型通常采用均匀时间步长采样进行训练，但我们的研究表明，随机梯度的方差在不同时间步长间存在显著差异，高方差的时间步长成为阻碍更快收敛的瓶颈。为解决这一问题，我们引入了一种非均匀时间步长采样方法，该方法优先处理这些更为关键的时间步长。我们的方法通过跟踪每个时间步长的梯度更新对目标函数的影响，自适应地选择最有可能有效最小化目标函数的时间步长。实验结果表明，这种方法不仅加速了训练过程，还在收敛时提升了性能。此外，我们的方法在各种数据集、调度策略和扩散架构上均表现出稳健的性能，优于此前提出的缺乏这种稳健性的时间步长采样和加权启发式方法。

1. 引言

近年来，扩散模型展现出了强大的表达能力，并在众多领域取得了显著成功，包括图像生成[9,23,24]、视频生成[18,35]、文本处理[4,12,22]、音频生成[26]、组合优化[25,30]、强化学习[1,15,32]等。然而，随着目标数据分布的复杂性不断增加，实现扩散模型训练的收敛在计算上变得越来越具有挑战性。例如，使用Stable-Diffusion-2.0[24]训练图像需要24,000个A100 GPU小时，而Open-Sora[35]进行视频训练则需要48,000个H800 GPU小时。这些高昂的计算成本不仅对推进生成式人工智能应用构成了重大障碍，还对环境产生了负面影响。

尽管已有许多研究探索了加速扩散模型训练的方法，但我们的研究专门针对扩散模型在扩散时间步长上的均匀训练过程进行改进。众所周知，跨时间步长的非均匀训练方案可以提高训练速度。例如，已经提出了几种启发式方法[7,11,17,31]。

然而，这些方法成功的根本原因在很大程度上仍未被探索，它们通常被视为启发式方法，可能在不同实验条件下的稳健性方面存在局限性。

通过一系列实验，我们观察到在不同的扩散训练时间步长中，随机梯度的方差存在显著差异。

我们假设这种可变性在均匀采样中观察到的不平衡训练中起着关键作用，并且可以解释先前提出的启发式方法在改善扩散模型训练方面的有效性。然而，由于不同时间步长的梯度之间存在强烈的相互依赖性，我们发现仅这一见解并不能直接导致扩散训练的最佳加速效果。

为此，我们提出了一种更直接的方法，开发了一种高效的算法，用于采样预计能最大程度降低目标函数的时间步长。该算法近似每个时间步长的梯度更新对变分下界的影响，并增加需要进一步优化的特定时间步长的采样频率。

这使得扩散训练过程在各种设置下的收敛速度显著快于以前的启发式方法。总而言之，我们的三个主要贡献如下：

1. 通过对扩散模型训练中梯度方差的分析，我们解释了为什么跨时间步长的非均匀训练过程可以导致更快的收敛。
2. 与以往基于启发式的扩散加速方法不同，我们提出了一种基于学习的方法，通过近似最小化负对数似然的变分边界来自适应地采样时间步长。
3. 我们在各种图像数据集、噪声调度策略和扩散架构上进行了实验，以证明我们方法的稳健性。

2. 相关工作

扩散模型的非均匀训练

多项研究探索了扩散模型的时间步长非均匀训练以加速训练过程。这些方法可主要分为两类：

• 1) 加权方法：这类方法对每个时间步长的损失应用手工设计的权重。例如，P2[7]为包含感知丰富内容的时间步长分配更大权重，而对涉及不可感知细节的时间步长分配较小权重。类似地，Min-SNR[11]基于钳制信噪比调整损失权重。

• 2) 采样方法：这些方法使用手工设计的采样器以非均匀方式选择时间步长。例如，Log-Normal[17]从对数正态分布中采样时间步长；SpeeD[31]修改采样概率，降低低损失时间步长的权重，并优先考虑过程增量快速变化的时间步长。

尽管这些启发式方法已证明能有效加速扩散模型训练，但其与底层目标函数的直接关系仍不明确，这引发了对其在未来可能出现的新场景（尤其是更复杂场景）中适用性的担忧。

学习优化（L2O）

L2O是机器学习中专注于学习优化算法的领域，通常涉及由多个训练和验证数据集组成的元训练集。利用该元训练集，L2O通过监督学习[3,33]或强化学习[6,20]学习基于梯度更新参数，以最小化验证集误差。尽管我们的目标同样聚焦于学习优化，但我们解决的是更具挑战性的任务——不假设用于迁移学习的元训练集，而是专注于在线学习优化，无需先验训练数据。

3. 背景：扩散模型

扩散模型[14,27]包含两个关键过程：预定义的前向加噪过程和反向去噪过程，后者通过变分推理学习以匹配前向过程。

前向过程

从前向过程从数据 $x_0 \sim q(x_0)$ 开始，作为马尔可夫高斯过程逐步引入噪声，生成一系列噪声潜变量 $\{x_1, x_2, \dots, x_T\}$：

$$q(x_t | x_{t-1}) := \mathcal{N}(x_t ; \sqrt{1 - \beta_t}x_{t-1}, \beta_t I), \quad (1)$$

其中 $\beta_t$ 是噪声调度。潜变量可通过 $x_t := \sqrt{\bar{\alpha}_t}x_0 + \sqrt{1 - \bar{\alpha}_t}\epsilon$ 等价采样，$\epsilon \sim \mathcal{N}(0, I)$，其中 $\bar{\alpha}_t := \prod_{s=1}^t (1 - \alpha_s)$ 且 $\alpha_t := 1 - \beta_t$。

反向过程

扩散模型通过学习的反向过程 $p_\theta(x_{t-1} | x_t)$ 从 $x_T$ 生成数据 $x_0$，逐步从高斯噪声中对潜变量去噪：

$$p_\theta(x_{t-1} | x_t) := \mathcal{N}(x_{t-1} ; \mu_\theta(x_t, t), \sigma_t^2 I), \quad (2)$$ $$\mu_\theta(x_t, t) := \frac{1}{\sqrt{\alpha_t}}\left(x_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}}\epsilon_\theta(x_t, t)\right),$$

其中 $\sigma_t^2$ 是去噪过程的方差，$\mu_\theta$ 按Ho等人[14]的建议参数化以简化计算。$\epsilon_\theta$ 是参数为 $\theta$ 的深度神经网络，基于 $x_t$ 和 $t$ 预测噪声向量 $\epsilon$。

训练

扩散模型的训练通过优化对数似然的变分下界实现，该下界表示为高斯分布对比的KL散度项之和，并可通过归一化噪声重新表述：

$$\begin{aligned} \mathcal{L}_{\text{VLB}}(\theta) &:= \mathbb{E}_{x_0}\left[\sum_t D_{\text{KL}}\left(q(x_{t-1} | x_t, x_0) \| p_\theta(x_{t-1} | x_t)\right)\right] \\ &= \mathbb{E}_{x_0, \epsilon, t}\left[\frac{\beta_t^2}{2\sigma_t^2\alpha_t(1 - \bar{\alpha}_t)}\|\epsilon - \epsilon_\theta(x_t, t)\|^2\right], \end{aligned}$$

其中 $t \sim U\{1, T\}$。训练扩散模型的标准方法是忽略常数项[14]，得到目标函数：

$$\mathcal{L}_{\text{DDPM}}(\theta) := \mathbb{E}_{x_0, \epsilon, t}\left[\|\epsilon - \epsilon_\theta\left(\sqrt{\bar{\alpha}_t}x_0 + \sqrt{1 - \bar{\alpha}_t}\epsilon, t\right)\|^2\right].$$

4. 揭秘跨时间步长的非均匀扩散模型训练

在本节中，我们提出一个假设来解释为何非均匀训练策略可能比Ho等人[14]使用的均匀训练方法更具优势。为简化起见，我们首先假设跨时间步长的参数化是独立的，表示为：

$$\epsilon_\theta(x_t, t) = \epsilon_{\theta_t}(x_t), \forall t, \quad (6)$$

其中每个时间步长有独立的参数集 $\theta = \{\theta_t\}_{t=1}^T$。这种参数化使每个时间步长具有独立的优化目标，即 $L_{\text{VLB}}(\theta) = \sum_{t=1}^T \mathbb{E}_{x_0, \epsilon}[L_t(\theta)]$，其中

$$\mathcal{L}_t(\theta) := c_t \|\epsilon - \epsilon_{\theta_t}(x_t)\|^2, \quad c_t := \frac{\beta_t^2}{2\sigma_t^2\alpha_t(1 - \bar{\alpha}_t)}. \quad (7)$$

我们现在将问题分解为T个独立的子问题，其中均匀时间步长采样的训练对应于均匀选择这些子问题 $L_t$ 之一，并对其对应参数 $\theta_t$ 执行随机梯度步骤。通俗地说，如果某些子问题特别具有挑战性，需要更多迭代才能优化，那么非均匀训练策略就有改进的潜力。

从这个意义上讲，我们通过测量随机梯度的方差 $Var_{x_0, \epsilon}[\nabla_{\theta_t} L_t(\theta)]$ 来评估这些子问题的难度。众所周知，随机梯度下降的收敛性仅在学习率与梯度方差成比例时得到保证：

定理1（Garrigos和Gower[10]，非正式）在几个合理假设下，对于每个 $\varepsilon > 0$，只要

$$\gamma = \mathcal{O}\left(\frac{1}{\sqrt{K}}\left(Var_{x_0, \epsilon}\left[\nabla_{\theta_t} \mathcal{L}_t(\theta_*)\right]\right)^{-\frac{1}{2}}\right),$$ $$K \geq \mathcal{O}\left(\frac{1}{\varepsilon^2} Var_{x_0, \epsilon}\left[\nabla_{\theta_t} \mathcal{L}_t(\theta_*)\right]\right),$$

就可以保证 $\mathbb{E}[L_t(\theta_K) - \inf L_t] \leq \varepsilon$。

这里，$\gamma$ 是随机梯度下降（SGD）的学习率，$\theta_K$ 是经过K次梯度更新后的参数，$\theta_*$ 是最优参数。K是达到精度 $\varepsilon$ 所需的总迭代次数。

简而言之，定理1表明学习率应与随机梯度的标准差成反比缩放。通过适当调整学习率，SGD收敛所需的迭代次数与梯度的方差成比例增加。由于现代优化器（如Adam[19]）的设计通过将学习率与指数移动标准差成反比缩放来确保收敛，我们可以推断子问题 $L_t$ 收敛所需的迭代次数大约与最优参数下的随机梯度方差成比例增长。

扩散模型的梯度方差

考虑到每个时间步长的梯度方差反映了该时间步长优化问题的难度，我们在图2的顶行展示了扩散模型训练期间测量的梯度方差。结果清楚地表明，时间步长之间的梯度方差存在不平衡。值得注意的是，对于所有三种调度策略，早期时间步长的梯度方差显著高于后期时间步长，且这种趋势持续存在，不会快速衰减。由于最优参数下的方差是收敛所需梯度迭代次数的下界（定理1），我们可以预期当均匀采样时间步长时，早期时间步长可能成为训练效率的瓶颈。

图2的底行显示了实际的时间步长扩散损失，正如预期的那样，后期时间步长附近的损失比早期时间步长附近的损失下降得快得多。这些结果表明，对早期时间步长进行过采样，对后期时间步长进行欠采样，可以通过有效解决早期时间步长子问题带来的瓶颈来提高整体训练效率。

$c_t$ 的相互依赖性

我们通过展示跨时间步长的梯度方差变化来推动时间步长的非均匀采样，并由此得出非均匀采样器应与梯度方差成比例的结论。然而，如果没有基于特定参数化的子问题 $L_t$ 独立性假设，这种采样器的最优性无法保证。

图4说明了扩散模型训练期间子问题之间的相互依赖程度。该图显示，未训练时间步长的损失增加显著大于已训练时间步长的损失减少，突显了子问题之间的强相互依赖性。鉴于梯度方差可能相差高达百倍（图2），显然仅按梯度方差成比例采样可能不足以有效降低后期时间步长的损失。通过实验我们证实，按梯度方差成比例采样时间步长会导致性能不佳（表4）。

5. 用于扩散模型训练的自适应非均匀时间步长采样

基于前一节讨论的非均匀时间步长采样方案的需求，本节介绍一种自适应时间步长采样器$\pi_{\phi}$，旨在提升扩散模型的训练效率。我们提出了一种将该采样器与扩散模型以在线方式联合训练的方法。借助$\pi_{\phi}$，这种方法可加速扩散训练过程，且无需为其优化准备先验训练数据。

5.1 非均匀时间步长采样器

我们提出一种用于扩散模型训练的非均匀时间步长采样策略，利用时间步长采样器$\pi$来提升训练效率：

$$\theta \leftarrow \theta - \gamma \nabla_{\theta} \mathcal{L}_t, \ t \sim \pi(\cdot | x_0), \ \mathcal{L}_t = c_t \|\epsilon - \epsilon_{\theta}(x_t, t)\|^2, \quad (10)$$

其中$c_t$是如公式(7)所定义的依赖于调度的常数。请注意，我们不再使用跨时间步长的独立参数化。

为利用扩散训练损失在紧密相关的时间步长间具有相似性这一特点，我们对生成两个正实标量的神经网络进行参数化，随后利用这两个标量从Beta分布中采样$t$：

$$a, b \sim \pi_{\phi}(\cdot | x_0), \ a > 0, \ b > 0, \ t \sim \text{Beta}(a, b). \quad (11)$$

这确保我们的时间步长采样器能够考虑到相似时间步长之间的高相关性。为简化表述，我们将继续用$t \sim \pi_{\phi}(\cdot | x_0)$表示这一过程。

我们的目标是训练时间步长采样器$\pi_{\phi}$，以在有限数量的训练步骤内尽可能高效地最小化原始目标函数（即变分下界$L_{\text{VLB}}$）。为实现这一目标，我们寻求优化采样器$\pi_{\phi_k}$，使其在每次迭代$k$中采样能使$L_{\text{VLB}}$降幅最大的时间步长$t$。换句话说，采样器$\pi_{\phi_k}$经训练后，可最大化在梯度步骤$k$使用采样时间步长$t$更新扩散模型前后的$L_{\text{VLB}}$差值：

$$\phi_k = \underset{\phi_k}{\text{arg max}} \ \mathbb{E}_{x_0, \epsilon, t \sim \pi_{\phi_k}(\cdot | x_0)}[\Delta_k^t], \ \text{其中} \quad (12)$$ $$\Delta_k^t = \frac{1}{T} \sum_{\tau=1}^T \delta_{k, \tau}^t = \frac{1}{T} \sum_{\tau=1}^T \mathcal{L}_{\tau}(\theta_k) - \mathcal{L}_{\tau}(\theta_k - \gamma \nabla_{\theta_k} \mathcal{L}_t).$$

通过这样优化后的$\phi_k$，使用采样器$\pi_{\phi_k}$的扩散模型能够在整个训练过程中自适应选择最佳时间步长，从而在更新$\theta_k$后最小化$L_{\text{VLB}}$。在实际应用中，我们还在$\Delta_k^t$中为采样器$\pi_{\phi}$加入熵正则化，以避免过早收敛。

前瞻挑战

评估$\mathcal{L}_{\tau}(\theta_k - \gamma \nabla_{\theta_k} \mathcal{L}_t)$需要计算更新后$\theta$的变化，这本质上涉及跨不同$t$的前瞻过程，导致计算上不可行。为开发实用算法，我们选择对$\phi_k$进行单梯度更新，以$\theta$的最新更新作为参考。通过应用似然比梯度，我们对$\phi_k$使用以下更新：

$$\phi_{k+1} = \phi_k + \gamma \cdot \Delta_k^t \cdot \nabla_{\phi_k} \log \pi_{\phi_k}(\cdot | x_0), \quad (13)$$

其中我们使用了恒等式$\nabla_{\phi} \pi_{\phi} = \pi_{\phi} \cdot \nabla_{\phi} \log \pi_{\phi}$。这是一种使用采样$\epsilon$和Beta分布参数$a$、$b$的随机近似。算法1概述了基于此$\phi_k$梯度更新的训练过程。请注意，$c_t$未在算法中显式包含，因为若需要，它可有效归入采样器$\pi_{\phi}$。

5.2 $\Delta_k^t$的近似

即使采用上述替代算法降低了计算量，评估$\Delta_k^t$仍需计算$\theta_k$更新前后的精确目标函数$L_{\text{VLB}}$，即$\Delta_k^t = L_{\text{VLB}}(\theta_k) - L_{\text{VLB}}(\theta_{k+1})$。幸运的是，借助附近时间步长损失的相关性，我们可使用一小部分时间步长$S \subset \{t\}_{t=1}^T$精确近似$\Delta_k^t$，即：

$$\Delta_k^t \approx \overline{\Delta}_k^t = \frac{1}{|S|} \sum_{\tau \in S} \delta_{k, \tau}^t = \frac{1}{|S|} \sum_{\tau \in S} \mathcal{L}_{\tau}(\theta_k) - \mathcal{L}_{\tau}(\theta_{k+1}).$$

为确定$S$，对于每个$k$，我们为单个$x_0$计算$\{\delta_{k, \tau}^t\}_{\tau=1}^T$，并将这些值存储在队列中。然后，使用任何现成的特征选择算法，基于队列中的数据识别出最能解释$\Delta_k^t$的前$S$个时间步长，其中$|S|$为超参数。我们主要使用极小的子集$|S| = 3$来近似$\Delta_k^t$。后续实验表明，该小子集足以在墙钟时间方面实现出色性能（图7）。

使用此特征选择方法近似$\Delta_k^t$的过程如算法2所述。实际应用中，我们通过仅在扩散模型每$f_S$次梯度迭代时更新$\phi$，进一步降低计算负担。我们算法的整体结构如图3所示。所提出的时间步长采样机制引入的计算开销见附录第8节。

6. 实验

在所有实验中，我们使用官方实现来训练每个扩散架构。同时，时间步长采样器$\pi_{\phi}$使用$\Delta_{k}^{t}$进行训练，该值近似于更新扩散模型前后$L_{VLB}$的变化。我们首先在6.1节中评估我们的学习方法与启发式方法的有效性，然后在6.2节中评估我们的模型与启发式方法的兼容性，最后在6.3节中提供额外分析以进一步支持我们的方法。

实验设置

使用弗雷歇初始距离（FID）[13]来评估生成图像的保真度和覆盖范围。我们主要使用配备四个NVIDIA RTX 3090 GPU的机器来训练模型。详细的超参数值和实现细节可在附录第9节和第10节中找到。

6.1 与启发式加速方法的比较

不同噪声调度的比较

为了比较我们的模型和其他加速技术在不同噪声调度上的表现，我们使用CIFAR-10 32×32数据集[2]，并以DDPM[14]作为骨干架构。
表1表明，尽管我们的方法基于学习，但与启发式方法[7,11,17,31]相比，它通常在每次迭代中实现更快的收敛和更好的性能。此外，我们的方法在不同的噪声调度下表现出一致的鲁棒性。相比之下，虽然Min-SNR在使用线性和余弦噪声调度时是启发式方法中表现最好的，但其在二次噪声调度下的性能显著下降。如图5所示，我们的方法通过在整个训练过程中自适应地采样时间步长来实现鲁棒性能，而启发式方法依赖于固定的采样策略。此外，图1显示，与启发式方法相比，我们的方法在墙钟时间方面也实现了更快的收敛。

高分辨率图像的比较

为了评估高分辨率图像的加速方法，我们使用CelebA-HQ 256×256数据集[16]，并以LDM[24]作为骨干架构。表2表明，尽管启发式方法最初表现更好，但我们的基于学习的方法随着时间的推移会超越它们。我们还发现，在这种情况下，P2的表现与表1中的结果相当，这表明启发式方法的性能可能因不同的设置而异。我们的方法和基线生成的CelebA-HQ图像的可视化结果见附录第11节。

6.2 与启发式方法的兼容性

在6.1节中，我们证明了虽然我们的基于学习的方法由于其基于学习的性质，最初比启发式方法加速更慢，但它最终会在后续迭代中收敛到更好的最优解。然而，我们发现启发式加权方案可以纳入我们的模型，以从一开始就实现更快的收敛。实际上，启发式权重可以被视为依赖于调度的常数$c_{t}$的一部分，并且可以有效地纳入我们的采样器$\pi_{\phi}$（5.1节）。因此，我们的采样器可以有效地考虑通过启发式方法纳入扩散损失的额外权重。
为了评估我们的模型与启发式方法的兼容性，我们在ImageNet 256×256基准[8]上使用Min-SNR加权训练我们的模型。我们使用LDM[24]的编码器将尺寸为256×256×3的图像转换为尺寸为32×32×4的潜在表示。我们使用U-ViT[5]作为我们的扩散骨干。表3显示，我们的带有启发式加权的方法也可以从一开始就实现更快的收敛，最终收敛到比基线更低的FID分数。

6.3 分析

非均匀时间步长采样分析

图6（右列）展示了我们的方法应用于DDPM时采样的时间步长分布。图6（左、中列）展示了在不使用我们提出的方法的情况下，使用$L_{VLB}$而不是$L_{DDPM}$训练DDPM时观察到的梯度方差和扩散损失的变化。使用$L_{VLB}$进行训练与使用$L_{DDPM}$进行传统扩散训练（图2）的不同之处在于，$L_{VLB}$使用$c_{t}$对跨时间步长的损失进行不同的加权。由于$c_{t}$在后期时间步长具有较大的值，这会导致这些时间步长的损失和梯度方差增加，尤其是在余弦调度中，如图6所示。
在第4节中，我们认为假设跨时间步长的独立性，收敛所需的迭代次数将与收敛时的梯度方差成比例。图中显示的时间步长采样分布支持了这一观点，线性调度大量采样低时间步长，而余弦调度主要从低时间步长和高时间步长中采样。这表明，最有效地减少原始目标$L_{VLB}$的方向在某种程度上与采样具有较高梯度方差的时间步长一致。另一方面，表4表明，按梯度方差成比例采样时间步长会导致性能不佳，这可能是由于子问题$L_{t}$之间的相互依赖性。

关于$|S|$的消融研究

图7展示了改变所选时间步长数量$|S|$对我们提出的模型性能的影响。所有$|S|$的选择都显示出比基线的加速，这表明我们可以仅使用时间步长的一小部分来近似$\Delta_{k}^{t}$。我们观察到$|S|=3$在墙钟时间方面实现了最佳性能。因此，在我们的所有实验中，我们主要使用$|S|=3$来近似$\Delta_{k}^{t}$。

7. 结论

在本文中，我们发现了不同扩散训练时间步长中随机梯度方差的变化，并指出这种变化可能导致收敛缓慢。然而，不同时间步长的梯度之间存在强相互依赖性，这表明仅靠这一见解并不能直接实现扩散训练的最优加速。为此，我们提出了一种更直接的方法，该方法估计每个时间步长的梯度更新对变分下界的影响，并调整采样频率以优先处理需要进一步优化的时间步长。实验结果表明，与以往的启发式方法相比，我们的方法显著加速了扩散训练过程的收敛，同时在各种数据集、调度策略和扩散架构上表现出稳健的性能。

局限性与未来工作

我们的基于学习的方法与启发式方法相比，计算成本更高，在较大的问题领域中资源需求显著增加。此外，我们尚未探索该方法在基于分数的扩散框架中的应用，该框架将基于分数的生成建模与扩散概率建模相结合，在图像生成任务中已展现出显著效果[17,29,34]。解决这些计算效率方面的局限性，并将该方法的适用性扩展到基于分数的扩散模型中，仍是未来研究的一个有前途的方向。