回复
去噪步数减少50%，图像生成质量反而更好！西湖大学等提出TPDM：自适应噪声调度
文章链接：https://arxiv.org/pdf/2412.01243
亮点直击
- 提出了时间预测扩散模型（TPDM），该模型可以在推理过程中自适应地调整噪声调度，实现图像质量和模型效率之间的平衡。
- 为了训练TPDM，通过强化学习最大化图像质量，并根据去噪步骤数折扣，直接优化最终的性能和效率。
- 模型在多个评估基准上表现优越，在减少推理步数的同时取得了更好的结果。
总览全文
扩散模型和流模型在文本到图像生成等多种应用中取得了显著成功。然而，这些模型通常在推理过程中依赖于相同的预定去噪调度策略，这可能限制了推理效率以及在处理不同提示时的灵活性。本文认为，最优的噪声调度应该适应每个推理实例，并提出了时间预测扩散模型（TPDM）来实现这一目标。
TPDM采用了一个即插即用的时间预测模块（TPM），该模块在每个去噪步骤中根据当前的隐空间特征预测下一个噪声水平。使用强化学习来训练TPM，目标是最大化一个奖励，该奖励通过去噪步骤的数量来折扣最终图像质量。通过这种自适应调度器，TPDM不仅生成与人类偏好高度一致的高质量图像，还能动态调整去噪步骤的数量和时间，从而提升性能和效率。
在多个扩散模型基准上训练了TPDM。在Stable Diffusion 3 Medium架构下，TPDM实现了5.44的美学评分和29.59的人类偏好评分（HPS），同时使用大约50%更少的去噪步骤，取得了更好的性能。
动机
以下图4中的几张图像为例，右侧的图像内容丰富，需要更多的去噪步骤来捕捉更细致的细节。相比之下，左侧的图像相对简单，可以使用较少的步骤生成，而不影响质量。此外，Karras也证明了不同的噪声调度对生成质量有很大影响。
那么，不禁要问：是否可以在推理过程中自适应地调整去噪步骤的数量和每个步骤的噪声水平，而无需用户的手动干预？
基于此，本文提出了时间预测扩散模型（TPDM），该模型能够在推理过程中自适应地调整去噪步骤的数量和去噪强度。具体而言，TPDM通过实现一个即插即用的时间预测模块（TPM），该模块可以根据当前步骤的隐空间特征预测下一个扩散时间，从而使噪声调度能够动态调整。
TPM通过强化学习进行训练。将多步骤的去噪过程视为一个完整的轨迹，并将图像质量（经过去噪步骤数的折扣）作为奖励进行优化。图像质量通过与人类偏好对齐的奖励模型进行衡量。
TPM可以轻松集成到任何扩散模型中，几乎不增加额外的计算负担，并使模型能够自动调整超参数，例如样本步骤和每个步骤的噪声水平，从而在不需要人工干预的情况下实现图像质量和效率之间的平衡。此外，在训练过程中，模型的扩散过程与推理过程中保持一致，直接优化推理性能并减少去噪步骤。
在多个最先进的模型上实现了TPDM，包括Stable Diffusion和Flux。通过自适应噪声调度，模型在生成图像时平均使用了50%更少的步骤，且图像质量与Stable Diffusion 3持平或略有提高（0.322 CLIP-T，5.445美学评分，22.33选择评分，29.59 HPSv2.1）。这些结果表明，TPDM具有在追求高质量图像生成和提高模型效率之间找到平衡的潜力。
方法
在本节中，我们首先简要回顾扩散模型的基本原理，然后介绍时间预测模块（TPM），最后详细说明TPM的训练算法。
基础知识
扩散模型通过一个反向过程学习生成图像，该过程逐渐去除样本中的噪声。实现这一反向过程的主要范式是流匹配（Flow Matching）。因此，在此介绍流匹配模型的工作原理以及当前最先进模型的详细结构。
目前，许多最先进的扩散模型都基于DiT架构，它们只关注条件图像生成，并采用调制机制来条件化网络，既考虑扩散去噪步骤中的扩散时间，也考虑文本提示。
使用这样的模型，可以根据隐特征和当前噪声水平执行单个去噪步骤。
时间预测扩散模型（TPDM）
如前所述，需要一系列去噪步骤才能使用训练好的扩散模型生成图像。通常，在这个过程中会对所有提示应用固定的噪声调度，为每个步骤分配一个预定的噪声水平。
TPDM仅需要对原始扩散模型进行最小的修改：添加一个轻量级的时间预测模块（TPM），如下图3所示。该模块将transformer块前后的隐特征作为输入，从而考虑到当前噪声输入和本步骤的预测结果。经过几层卷积后，TPM将隐空间特征池化成一个单一的特征向量，并通过两个线性层预测a和b。我们还在TPM中使用了自适应归一化层，使模型能够感知当前的时间嵌入。
在训练过程中，冻结原始的扩散模型，只更新新引入的TPM。因此，模型在保持原有图像生成能力的同时，学习预测下一个扩散时间。
训练算法
因此，我们优化TPM，以最大化整个去噪过程后的图像质量，从而实现精确的时间预测。图像质量通过图像奖励模型进行衡量。考虑到整个推理计算图过于深度，无法进行梯度反向传播，我们使用邻近策略优化（PPO）进行训练，其损失函数公式为：
将在后文中具体说明PPO中使用的动作和优势。
通常，当模型做出一系列预测时，PPO将每个单独的预测视为一个动作并以批量优化。然而，最近的RLOO研究指出，当奖励信号只出现在序列的末尾，并且环境动态完全确定时，可以将整个序列视为一个动作，而不会影响性能。因此，为简化起见，将整个生成过程，包括调度中的所有预测时间，视为一个单独的动作进行优化。
其中，(N) 表示生成步骤的总数。
图像奖励与步骤数的折扣
实验
实现细节
数据集收集了用于训练模型的提示词。这些提示词是通过 Florence-2和 Llava-Next生成的，用于生成 Laion-Art和 COYO-700M数据集的图像描述，并利用这些提示词构成我们的训练集。
训练配置
主要结果
不同图像的动态调度在前面图 4 中，展示了使用不同提示词生成的图像及其由 TPDM 预测的对应调度。当输入较短且简单的提示时，生成的图像包含较少的物体和细节，因此扩散时间会更快地下降，并在相对较少的步骤中达到 0。相反，当提供较长且更复杂的提示时，模型需要生成更多的视觉细节，因此扩散时间下降得较慢，以便生成精细的细节。在这种情况下，TPDM 在生成过程中需要更多的去噪步骤。
调整 γ 对不同步骤数的影响公式（9）中的 γ 控制图像奖励在更多生成步骤中的折扣方式，从而影响扩散时间在去噪过程中的衰减速度，进而影响我们模型的平均去噪步骤数。
如下图 6 所示，当将 γ 从 0.97 降低到 0.85 时，TPDM 趋向于更快速地减少扩散时间，从而减少了采样步骤，从 15.0 步减少到 7.5 步。此外，与基准模型（黄色线）相比，TPDM（紫色线）在相同的推理步骤数下，始终能获得显著更高的美学分数，达到了模型效率和生成性能之间的良好平衡。
视觉对比我们的方法在生成细粒度细节方面表现出色。TPDM 生成的图像相比 SD3-Medium（图 5C 中的图像）和结果中的图像，展示了更为真实的笔记本键盘。
定量结果
在多个最先进的扩散模型上应用了 TPM，包括 Stable Diffusion 3 Medium、Stable Diffusion 3.5 Large 和 Flux 1.0 dev，展示了 TPM 如何提升它们的性能。主要评估了两类指标：第一类是客观指标，包括 FID、与给定提示的对齐度（CLIP-T）以及人类偏好分数（美学分数 v2 和 HPSv2.1）。第二类是通过用户研究直接比较不同模型生成的图像。
定量指标在表 1 中比较了 TPDM 和上述模型。尽管保持竞争力的性能，所有这些模型平均可以在推荐的步骤数的一半左右生成图像。
此外，代表人类偏好的指标提高得最多。例如，通过仅使用平均 15.28 步生成的图像，TPDM-SD3-Medium 获得了 29.59 的 HPS 分数，比使用相似步骤的 Stable Diffusion 3 高出 +1.07，比原始的 28 步结果高出 +0.47。这可能归因于我们在优化中使用的奖励模型，从而生成符合人类偏好的美学图像。
用户研究为了更好地反映人类对这些模型的态度，通过邀请志愿者对比不同模型生成的图像，并选择他们偏好的图像，进行了用户研究。
具体来说，对于每个提示，我们提供了从 SD3-Medium 生成的 15 步和 28 步图像，以及从 TPDM-SD3-Medium 生成的图像。我们邀请了 15 名志愿者评估 50 个提示生成的图像。结果如表 2 所示，表明我们的模型生成的图像更符合人类的偏好。
模块架构的消融实验
在本节中，我们对 TPM 输入的选择进行了消融实验。如表 3 所示，将第一层和最后一层的特征都输入到 TPM 中，比仅使用其中任意一层的特征表现更好。
结论与局限性
本文提出了时间预测扩散模型（TPDM），这是一个具有灵活去噪调度器的文本到图像扩散模型，可以针对不同的提示自动调整去噪调度。通过引入时间预测模块，通过强化学习和奖励模型有效地训练了 TPDM。基于当前领先的扩散模型架构（Stable Diffusion 3 Medium），我们训练了一个强大的 MM-DiT 基础 TPDM，并在多个文本到图像生成基准测试中表现出了竞争力的定量性能。
尽管 TPDM 展现了良好的性能，但仍有一些局限性。例如，在本文中，我们仅为 TPM 设计了一个相对简单的架构，如何改进该模块以获得更好的性能仍然是一个未解之题。其次，我们冻结了原始模型的参数，并采用我们的训练方法更新了扩散模型的参数，从而取得了更好的结果，这一方法还需要进一步探索。
本文转自AI生成未来 ，作者：AI生成未来
赞
收藏
回复
相关推荐
