全文翻译
摘要
扩散模型和流匹配模型在文本到图像生成任务中取得了显著成功。然而，这些模型通常依赖于为所有提示词预先设定的去噪调度方案。多步反向扩散过程可以看作是一种逐步生成高质量图像的“思维链”。因此，扩散模型应该为每个实例进行推理，自适应地确定最优噪声调度方案，以实现高效采样的同时保证高生成质量。在本文中，我们为此引入了时间预测扩散模型（TPDM）。
TPDM采用了一个即插即用的时间预测模块（TPM），该模块在每个去噪步骤中根据当前的潜在特征预测下一个噪声水平。我们使用强化学习来训练TPM，以最大化一个奖励值，该奖励值鼓励生成高质量的最终图像，同时惩罚过多的去噪步骤。通过这种自适应调度器，TPDM不仅能生成与人类偏好高度一致的高质量图像，还能动态调整扩散时间和去噪步骤数量，从而同时提升性能和效率。
基于Stable Diffusion 3 Medium架构，TPDM的美学评分达到5.44，人类偏好评分（HPS）达到29.59，同时使用约50%的去噪步骤就能实现更优的性能。
引言
近年来，深度生成模型（包括扩散模型[13, 50, 52]）在多种任务中取得了非凡的性能，涵盖图像合成[18, 19, 43, 44, 47]、视频生成[4, 14, 64]以及其他任务[21, 36, 40, 49]。作为一种多步去噪框架，扩散模型通过迭代采样逐步将随机噪声细化为连贯的数据，这是其能够生成高质量、多样化输出的基础。
扩散模型的推理过程需要选择一个噪声调度器，例如从高斯噪声中去噪时噪声水平如何逐步变化。这自然形成了一种“思维链”[59]，因为它决定了模型如何逐步生成真实图像。许多研究表明，不同的噪声调度方案会极大地影响模型性能[19]。领先的流匹配模型（如Stable Diffusion 3[10]和FLUX[22]）会根据目标分辨率提供推荐的噪声调度器。Sabour[46]和Xia[61]探索了微调模型调度方案的方法，以提高效率或整体性能。此外，一些单步生成器[16, 33-35]也取得了令人印象深刻的性能。尽管这些方法性能优异，但它们大多假设存在一种普遍适用的调度方案，对所有提示词和图像都是最优的，这一点值得怀疑。
我们来看图4中的一些例子。最右侧的图像包含由更长提示词生成的更复杂视觉结构。这些例子表明，扩散模型需要更长的采样步骤链来生成更丰富的内容。相比之下，左侧的图像结构更简单（例如，单个物体和统一背景），可以通过更短的步骤链生成。此外，正如后文所讨论的，每个步骤的扩散时间t所设定的噪声水平，对于正确决定样本在各个步骤中应如何去噪以最终生成高质量图像至关重要。因此，我们提出疑问：是否有可能为每个反向扩散步骤链自适应地确定去噪步骤的数量和每个步骤的噪声水平？
在本文中，我们提出了时间预测扩散模型（TPDM），它能在推理过程中自适应地调整每个步骤的噪声水平和总步骤数。这是通过一个即插即用的时间预测模块（TPM）实现的，该模块可以根据当前步骤的潜在特征预测下一个扩散时间。通过将反向扩散过程视为每个步骤对扩散时间的一系列决策，我们采用强化学习来最大化与人类偏好一致的奖励模型[63]所计算的奖励值。该奖励不仅反映生成图像的质量，还会根据去噪步骤的数量进行惩罚。因此，TPDM直接优化以生成高质量图像，同时减少去噪步骤。
我们在多个最先进的模型（包括Stable Diffusion 3和Flux）上实现了TPDM。通过自适应噪声调度器，TPDM平均可以减少50%的步骤，同时保持相当或更优的质量（CLIP-T为0.322，美学评分为5.445，Pick评分为22.33，HPS v2.1为29.59）。我们在图1中展示了一些生成的图像以及所使用的步骤数。这些结果表明，TPDM有潜力在追求高质量图像生成的同时提高模型效率。
我们的贡献总结如下：
- 我们引入了时间预测扩散模型（TPDM）。它可以在每个去噪步骤预测下一个扩散时间，为每个推理实例确定最优噪声调度方案。
- 我们使用强化学习训练TPDM，并最大化一个同时考虑图像质量和步骤数量的奖励，直接优化最终性能和效率。
- TPDM在多个评估基准上表现出改进的性能，使用更少的推理步骤取得了更好的结果。
2. 相关工作
2.1. 扩散模型
扩散概率模型（DPMs）[13, 50]通过学习不同噪声水平下的数据分布，从纯高斯噪声中恢复原始数据。凭借对复杂数据分布的强适应性，扩散模型在图像[15, 20, 39, 40, 42]、视频[4, 14, 28, 64]等多个领域取得了显著性能，极大地推动了人工智能生成内容的发展。
2.2. 噪声调度
为了生成图像，模型必须确定每个步骤的扩散时间。这可以通过离散时间调度器[13, 51]或连续时间调度器[25, 29]实现，具体取决于模型。通常，扩散时间反映了每个步骤的噪声水平，而大多数现有方法依赖于预先确定的调度器。目前，领先的流匹配模型（Stable Diffusion 3[10]和FLUX[22]）会根据目标分辨率提供推荐的调度器。
一些方法通过优化调度器来加速采样或提升图像质量。Xia等人[61]通过预测每个步骤的新扩散时间来找到更准确的积分方向，Sabour等人[46]利用随机微积分为不同调度器和模型寻找最优采样方案。Wang等人[56]借助强化学习自动搜索扩散概率模型的最优采样调度器。此外，调整调度器在分子生成等其他领域也被证明是有效的[17]。一些单步生成器[16, 34, 35]通过扩散蒸馏[31]也实现了令人印象深刻的性能。
尽管上述方法性能优异，但它们大多假设存在一种适用于所有提示词和图像的通用最优调度器，这一假设值得怀疑。相反，TPDM能够在推理过程中自适应调整噪声调度，为每个生成任务选择具有合适步骤数的最优采样方案，从而同时提升图像质量和模型效率。
2.3. 强化学习与人类反馈学习
基于人类反馈的强化学习（RLHF）近年来在大型语言模型（LLMs）领域受到广泛关注[1, 48, 53]，并逐渐扩展到其他领域。扩散模型的研究也越来越多地融入奖励模型，以增强与人类偏好的一致性[23, 60, 63]。通过将潜在变量视为动作，扩散模型可以通过策略梯度（DDPO[3]、DPOK[11]）、actor-critic框架（DACER[58]）、直接偏好优化（Diffusion-DPO[55]）等强化学习算法进行优化。一些近期工作还研究了少步生成模型的RLHF方法[32, 37]。
然而，这些方法依赖于扩散SDE采样器的高斯反向过程，难以应用于流匹配模型。相比之下，本文将扩散时间视为动作而非潜在变量，旨在通过调整调度来以更少步骤实现更优质量，提供了一种更通用、灵活的解决方案。
3. 所提方法
本节首先简要回顾扩散模型的基本原理，随后介绍时间预测扩散模型（TPDM），最后详细阐述TPDM的训练算法。
3.1. 基础知识
扩散模型通过一个反向过程学习生成图像，该过程逐步从样本中去除噪声。目前主流的范式是流匹配[10, 25]。因此，我们在此介绍流匹配模型的工作原理以及当前最先进模型的内部详细结构。
我们考虑一个生成模型，该模型建立了从噪声分布(p_1)中抽取的样本(x_1)与数据分布(p_0)中抽取的样本(x_0)之间的映射。流匹配的目标是直接回归一个向量场(v_t)，该向量场生成一个概率流，实现从(p_1)到(p_0)的转换。
[
\mathcal{L}_{FM}(\theta) = \mathbb{E}_{t, p_t(x)} \left| v_{\theta}(x_t, t) - u(x_t, t) \right|^2
]
具有参数(\theta)的流匹配模型旨在预测噪声预测函数(v_{\theta}(x_t, t))，该函数逼近引导扩散过程从噪声分布到生成样本的干净分布的真实速度场(u(x_t, t))。由此，我们可以得到扩散常微分方程（ODE）：
[
\frac{dx_t}{dt} = v_{\theta}(x_t, t)
]
在推理过程中，假设我们用(N)个步骤生成一幅图像，每个步骤都有一个对应的时间(t_n)，该时间与噪声水平相关。那么，第(n)个生成步骤可表示为：
[
x_{t_n} = x_{t_{n-1}} + (t_n - t_{n-1}) \cdot v_{\theta}(x_{t_{n-1}}, t_{n-1})
]
也就是说，在最终生成干净图像(x_{t_N})之前，模型会形成一系列中间结果(\{x_{t_n}\}_{n=1}^{N-1})。通常，在典型的流匹配算法中，每个(t_n)都有一个预先确定的调度，而所提的TPDM将为单个样本动态确定该调度。
当前，许多最先进的扩散模型都基于DiT架构[8, 10, 22]构建。它们在去噪步骤中使用多层Transformer，使网络同时以扩散时间和文本提示为条件。在本文中，为了进行公平比较，我们将在DiT架构上构建TPDM。
3.2. 时间预测扩散模型（TPDM）
为了使模型能够动态调整噪声调度，TPDM在每个去噪步骤预测下一个扩散时间，如图2所示。这可以通过在扩散模型中添加一个轻量级的时间预测模块（TPM）来实现，如图3所示。该模块将DiT块前后的潜在特征拼接作为输入，从而同时考虑该步骤的原始含噪输入和去噪结果。随后，经过多个卷积层后，TPM对潜在特征进行池化，得到一个单一的特征向量用于预测。此外，我们在TPM中还使用了自适应归一化层[41]，使模型能够感知当前的扩散时间。
扩散时间决定了噪声水平，在整个去噪过程中应单调递减。为了避免时间倒退，在每个步骤中，TPM改为预测衰减率(r)，该衰减率量化了相邻步骤之间扩散时间(t)的减少程度。由于3.3节中的基于强化学习的训练算法需要预测时间的概率分布，TPM的最后阶段使用两个线性层处理池化后的特征，预测(r)的分布参数，而非预测一个确定性的值。
假设我们处于第(n)个去噪步骤。衰减率(r_n)的分布被建模为区间((0, 1))上的Beta分布，其两个参数(\alpha_n)和(\beta_n)由TPDM预测。需要注意的是，确保(\alpha_n > 1)和(\beta_n > 1)会得到一个单峰分布，这是理想的，因为它可以避免TPDM从多峰Beta分布中采样出可能不合理的衰减率。为了强制执行这一约束，我们对模型进行重参数化，使TPDM预测两个实值参数(a_n)和(b_n)，并通过式（4）由这两个参数确定(\alpha_n)和(\beta_n)。因此，衰减率(r_n)和下一个扩散时间(t_n)可以通过式（5）和式（6）采样得到：
[
\alpha_n = 1 + e^{a_n}, \quad \beta_n = 1 + e^{b_n} \quad (4)
]
[
r_n \sim \text{Beta}(\alpha_n, \beta_n) \quad (5)
]
[
t_n = r_n \cdot t_{n-1} \quad (6)
]
在训练过程中，我们冻结原始扩散模型的参数，只更新新引入的TPM参数。因此，模型在学习预测下一个扩散时间的同时，保留了原始的图像生成能力。
3.3. 训练算法
为了训练TPM，我们至少需要执行两个去噪步骤：第一步预测下一个扩散时间，第二步使用该时间进行去噪。一种朴素的方法是将一幅含噪图像作为第一步的输入，并使用第二步计算的重建损失来训练模型。梯度会通过预测的(t_n)反向传播，以更新第一步中的TPM。然而，我们发现经过训练的模型在推理时往往倾向于用极少的步骤完成去噪过程，导致图像质量较差。我们推测，通过监督两个步骤后的损失，模型会学习在两个步骤后生成完全去噪的图像，并为了最小化损失函数而提前停止。但问题在于，整个扩散反向过程结束后生成的最终图像才是关键，而这种方法忽略了这一点。
总结上述失败经验后，我们对TPM进行优化，以最大化整个去噪过程结束后生成的图像质量，从而实现精确的时间预测。图像质量由图像奖励模型[63]衡量。考虑到整个推理过程的计算图过深，梯度难以反向传播，因此我们使用近端策略优化（PPO）[48]训练模型，其损失函数为：
[
\mathcal{L}(\theta) = -\left[ \frac{\pi_{\theta}(y | s)}{\pi_{\text{old}}(y | s)} \hat{A}(s, y) - \lambda \text{KL} \left[ \pi_{\text{ref}}(\cdot | s), \pi_{\theta}(\cdot | s) \right] \right] \quad (7)
]
在我们的任务中，(s = (c, \epsilon))表示初始状态，包括输入提示(c)和高斯噪声(\epsilon \sim \mathcal{N}(0, 1))；(\pi_{\theta})表示策略网络（例如我们的TPDM模型）；(\pi_{\text{old}})表示用于采样轨迹的旧策略；(\pi_{\text{ref}})表示用于正则化的参考策略；(y)表示我们的策略所采取的动作（即调度时间）；(\hat{A}(s, y))表示在状态(s)下采取动作(y)的优势值。我们将在下文对这些概念进行具体说明。
将整个调度视为一个动作：通常，当模型进行一系列预测时，PPO会将每个单独的预测视为一个动作，并对这些动作进行批量优化。最近，RLOO[1]提出，当奖励只在序列末尾出现且给定序列元素后状态转移完全确定时，我们可以将整个序列视为一个动作，而不会影响性能。因此，为简单起见，我们将整个生成过程（包括调度中的所有预测时间）视为一个单一的优化动作。
因此，在计算式（7）的期望时，我们将整个轨迹作为优化中的一个训练样本。然而，TPDM只输出每个单独时间预测的分布，我们将第(n)步的分布记为(\pi_{\theta}^{(n)})。
根据链式法则，整个生成过程的概率可以计算为每个预测的乘积：
[
\pi_{\theta}(t_1, \ldots, t_N | s) = \prod_{n=1}^{N} \pi_{\theta}^{(n)}(t_n | s, t_1, \ldots, t_{n-1}) \quad (8)
]
其中(N)表示生成步骤的总数。在我们的模型中，每个因子(\pi_{\theta}^{(n)})是基于TPM在相应步骤输出的Beta分布计算得到的。
经步骤数量折扣的图像奖励
在为策略梯度训练生成样本时，我们通过使用当前的TPDM策略从高斯噪声生成图像来获取轨迹。由于这些生成的图像没有真实标签，我们选择ImageReward[63]仅根据最后一个去噪步骤的最终图像分配奖励分数。我们的目标是提升最终图像质量，同时避免无限制地增加去噪步骤，因此我们使用一个折扣因子(\gamma < 1)（一个超参数）对中间扩散步骤的奖励进行折扣，并对去噪步骤的总数(N)取平均值。该轨迹的最终奖励如下式（9）所示：
[
R(s, y) = \frac{1}{N} \sum_{n=1}^{N} \gamma^{N-n} \text{IR}(y, s)
]
其中(\text{IR}(y, s))表示根据初始状态(s)和预测动作(y)生成的图像的奖励。关于如何直接从一批(R(s, y))中计算优势值(\hat{A}(s, y))而无需价值模型进行优化，我们建议读者参考RLOO[1]。
4. 实验
4.1. 实现细节
数据集：我们收集文本提示来训练模型。这些提示要么来自COCO 2017[24]训练集，要么是通过Florence-2[62]和Llava-Next[26]对Laion-Art[54]和COYO-700M[5]数据集中的图像进行 caption 得到的。更多细节见附录A。
训练配置：我们使用AdamW优化器，beta系数为(0.9, 0.99)，学习率固定为1×10⁻⁵，最大梯度范数为1.0。我们的TPM模块通常只需200个训练步骤。每步中，我们采样256条轨迹的批次，并据此更新模型参数4次。
评估指标：我们使用FID、CLIP-T、Aesthetic v2和Pick Score作为评估指标，测试采用COCO 2017验证集的5000个提示。对于HPS v2.1，我们使用基准测试[60]提供的3200个提示。
4.2. 主要结果
不同图像的自适应调度：在图4中，我们展示了不同提示生成的图像及其对应的TPDM预测调度。当向TPDM输入更短、更简单的提示时，生成图像所需的物体和细节更少，因此扩散时间下降更快，并在相对较少的步骤内降至0。相反，当提供更复杂的提示时，模型需要更多中间步骤进行推理，以便在早期和中期生成更精细的细节，因此扩散时间下降更慢，此时TPDM在生成过程中需要更多去噪步骤。附录C提供了更多预测调度的示例供分析。
调整γ以控制步骤数量：式（9）中的γ控制奖励从最后一步传播到中间步骤时的衰减速度（如图2所示），从而影响收敛后TPDM的平均去噪步骤数。
γ对平均推理步骤数和性能的影响：如图6所示，当我们将γ从0.97降至0.85时，TPDM倾向于更快地减少扩散时间，导致去噪步骤从15.0减少到7.5。此外，与基线（黄线）相比，TPDM（紫线）在相同推理步骤数下始终获得显著更高的美学分数。
视觉对比：图5展示了TPDM-SD3-Medium与使用推荐步骤数和等效步骤数的SD3-Medium的对比，前者在细节处理能力和生成准确性上表现更优。例如，在图5（c）中，TPDM生成的图像具有比15步和28步的SD3-Medium更逼真的笔记本键盘；在图5（d）中，TPDM去除了多余的棒球棒，使图像更自然。总体而言，TPDM能够生成更准确、更逼真的图像。
4.3. 定量结果
我们将TPM应用于多个最先进的扩散模型，包括SD3-Medium、SD3.5-Large和Flux 1.0-dev，并展示了TPM如何提升其性能。除了常用指标（如FID、CLIP-T和人类偏好分数Aesthetic Score v2、HPSv2.1），我们还进行了用户研究。
定量指标：我们在表1中比较了TPDM与上述模型。在保持性能竞争力的同时，TPDM平均可在推荐步骤数的约一半内生成图像。
在所有指标中，代表人类偏好的指标提升最为显著。例如，TPDM-SD3-Medium平均仅用15.28步生成图像，HPS分数达29.59，比相同步骤的Stable Diffusion 3高1.07，也比原始28步结果高0.47。这种提升可能归功于与人类偏好高度一致的奖励模型——通过基于这种一致性优化图像生成，TPDM产生了更符合审美需求的结果。此外，由于其他调度优化方法大多在Stable Diffusion 1.5[45]等经典扩散模型上进行实验，我们在附录B中也与它们进行了严格对比。
用户研究：为更好地反映人类对这些模型的偏好，我们开展了用户研究，邀请志愿者对比不同模型生成的图像。具体而言，对于每个提示，我们提供SD3-Medium（15步和28步）和TPDM-SD3-Medium生成的图像各一张，共生成50组对比图像，并邀请15名志愿者进行选择。表2的结果表明，我们的模型生成的图像更符合人类偏好。