摘要：扩散模型被广泛应用于图像和视频生成，但其迭代生成过程缓慢且资源消耗大。尽管现有的蒸馏方法已显示出在图像领域实现一步生成的潜力，但它们仍存在显著的质量退化问题。在本研究中，我们提出了一种在扩散预训练后针对真实数据进行的对抗后训练（APT）方法，用于一步视频生成。为了提高训练稳定性和质量，我们对模型架构和训练过程进行了多项改进，并引入了一种近似R1正则化目标。实验表明，我们的对抗后训练模型Seaweed-APT能够使用单个前向评估步骤实时生成2秒、1280x720分辨率、24帧每秒的视频。此外，我们的模型还能够在一步内生成1024px分辨率的图像，质量可与最先进的方法相媲美。Huggingface链接：Paper page 论文链接：2501.08316
一、引言
随着生成对抗网络（GANs）和扩散模型（Diffusion Models）的兴起，图像和视频生成领域取得了显著进展。然而，传统的扩散模型由于其迭代生成过程，导致生成成本高且耗时。尽管现有的蒸馏方法已经展示了在图像领域实现一步生成的可能性，但这些方法往往伴随着显著的质量下降。为了解决这个问题，本文提出了一种新的方法，即对抗后训练（Adversarial Post-Training, APT），通过对抗训练来加速扩散模型，实现图像和视频的一步生成。
二、背景与相关工作
1. 扩散模型加速
扩散模型已成为学习大规模图像和视频生成的标准方法。为了减少生成成本，研究者们提出了多种方法，其中扩散步长蒸馏（Diffusion Step Distillation）是一种有效的方法。这种方法通过预训练的扩散模型作为教师模型，生成目标，然后使用知识蒸馏技术训练学生模型，使学生模型能够使用更少的扩散步长复制教师模型的输出。尽管这些方法在减少推理步长方面取得了进展，但一步生成仍然面临巨大挑战，特别是在实现细粒度细节、最小化伪影和保持结构完整性方面。
2. 一步视频生成
一步视频生成的工作可以追溯到使用GANs的方法，如DVD-GAN、MoCoGAN-HD和StyleGAN-V等。然而，这些方法通常只能在受限的数据域上生成低质量的视频。最近，一些研究尝试将图像扩散模型蒸馏为一步视频生成模型，但这些方法生成的视频分辨率较低，且存在显著的质量下降。
3. 稳定对抗训练
在GANs中，R1正则化已被证明对于促进收敛是有效的。然而，在大规模对抗工作中，由于高阶梯度计算的计算成本高昂，且不被现代深度学习软件栈支持，因此R1正则化的使用受到限制。本文提出了一种近似R1正则化方法，以解决这一问题。
三、方法
1. 概述
本文的目标是将一个文本到视频的扩散模型转换为一步生成器。为此，我们采用对抗后训练（APT）方法，通过对抗优化来实现这一目标。整个训练过程遵循一个最小-最大游戏，其中判别器D尝试区分真实样本和生成样本，而生成器G则试图生成能够欺骗判别器的样本。
2. 生成器
直接对抗训练扩散模型会导致训练崩溃。为了解决这个问题，我们首先使用确定性蒸馏方法初始化生成器。具体来说，我们采用离散时间一致性蒸馏方法，通过均方误差损失进行简单蒸馏。蒸馏后的模型虽然生成的样本很模糊，但为后续的对抗训练提供了一个有效的初始化。
在后续的训练中，我们主要关注一步生成能力，并始终将最终时间步长T传递给基础模型。生成器的损失函数定义为：
其中，gG(x)=log(1−σ(x))，σ(x)是sigmoid函数。
3. 判别器
判别器被训练为产生一个标量对数，以有效区分真实样本x和生成样本x^。为了实现稳定训练和提高质量，我们对判别器进行了多项改进：
- 初始化：判别器骨干网络使用预训练的扩散模型进行初始化，并直接在潜在空间中操作。
- 架构修改：在判别器的第16、26和36层引入新的仅交叉注意力变换器块，每个块使用一个可学习令牌作为查询，以交叉注意力到所有视觉令牌。
- 输入处理：直接为判别器提供原始样本x和x^，避免引入伪影。由于判别器骨干网络是从扩散模型初始化的，且扩散预训练目标在t=0时没有意义，因此我们使用不同时间步长的集合作为输入。
判别器的损失函数定义为：
其中，
，λ是正则化项的权重，σ是高斯噪声的方差。
4. 近似R1正则化
由于高阶梯度计算的计算成本高昂，且不被现代深度学习软件栈支持，因此我们提出了一种近似R1正则化方法。具体来说，我们通过对真实数据添加高斯噪声来扰动真实数据，并鼓励判别器对真实数据和其扰动之间的预测保持接近。这种方法减少了判别器在真实数据上的梯度，从而实现了与原始R1正则化一致的目标。
四、实验结果
1. 定性评估
我们首先在图像生成方面比较了我们的APT模型与原始扩散模型。实验结果显示，APT模型在保持细节和真实感方面表现更好。此外，我们还与其他一步图像生成方法进行了比较，结果显示我们的方法在细节保持和结构完整性方面表现优异。
在视频生成方面，APT模型在视觉细节和真实感方面也有所提升，但在结构完整性和文本对齐方面仍然存在一定的降解。然而，即使如此，APT模型生成的视频在1280×720分辨率下仍然保持了相当的质量。
2. 用户研究
我们进行了一系列用户研究，以评估生成样本的视觉保真度、结构完整性和文本对齐性。实验结果显示，我们的APT模型在视觉保真度方面表现优于其他一步生成方法，但在结构完整性和文本对齐性方面仍存在一定的不足。尽管如此，我们的模型仍然达到了与当前最先进方法相当的性能。
3. 消融研究
我们进行了一系列消融研究，以评估不同组件对模型性能的影响。实验结果显示，近似R1正则化对于维持稳定训练至关重要，没有它训练会迅速崩溃。此外，使用更深的判别器网络和多层特征可以显著提高图像质量。我们还发现，指数移动平均（EMA）和较大的批量大小对于稳定训练和提高模型性能也是有益的。
五、结论与局限
本文提出了一种新的对抗后训练方法（APT），用于加速扩散模型的一步生成。通过引入近似R1正则化和对判别器进行多项改进，我们成功训练了一个能够实时生成高分辨率视频和图像的模型。尽管我们的模型在视觉保真度方面表现优异，但在结构完整性和文本对齐性方面仍存在一定的不足。未来的工作将致力于进一步提高模型的结构完整性和文本对齐性，并探索生成更长时间视频的可能性。
六、局限性与未来工作
尽管我们的方法在一步生成高分辨率视频方面取得了显著进展，但仍存在一些局限性。首先，由于计算资源的限制，我们目前只能训练模型生成最多两秒的视频。未来的工作将探索使用更多的计算资源来生成更长时间的视频。其次，我们观察到APT可能会对文本对齐产生负面影响，这将是未来工作中的一个重要研究方向。最后，我们将继续探索如何进一步提高模型的结构完整性和生成质量。
本文通过引入对抗后训练方法（APT），成功实现了扩散模型的一步生成，显著提高了生成速度和效率。尽管仍存在一些局限性，但本文的研究为未来的工作提供了有价值的参考和启示。
