DiffusionGAN ——最快的小波扩散模型应用研究

介绍

扩散模型最近出现并迅速发展，吸引了许多研究人员的兴趣。这些模型能从随机的噪声输入生成高质量的图像。在图像生成任务中，它们的表现尤其优于最先进的生成模型（GANs）。扩散模型可以灵活地处理各种条件输入，从而实现广泛的应用，如文本到图像的生成、图像到图像的转换和图像复原。这在基于人工智能的数字艺术和其他领域有着潜在的应用前景。

扩散模型具有巨大的潜力，但其非常缓慢的估算速度使其无法像 GAN 那样被广泛采用。基本的扩散模型需要几分钟才能获得所需的输出质量。为了缩短估算时间，人们进行了许多研究，但即使是最快的算法也需要数秒才能生成一幅 32 x 32 的图像；通过将扩散模型与 GAN 相结合，DiffusionGAN 极大地改善了估算时间，但它仍然不适合大型或复杂的应用。不适合实时应用。

为了用于实时应用，本评论文章提出了一种名为小波扩散（WaveletDiffusion）的新扩散方法。它使用离散小波变换将输入转换为低频和高频分量，从而将输入压缩了四倍，大大缩短了推理时间。此外，还提出了一种小波专用生成器，以更有效地利用小波特征并保持输出质量。实验结果证实，小波扩散在扩散模型中速度最快，同时还能保持较高的图像质量。
论文地址：https://arxiv.org/abs/2211.16152
源码地址：https://github.com/vinairesearch/wavediff

建议方法

基于小波的扩散方案

在这里插入图片描述

图 1.小波扩散方案。

在本文中，输入图像被分解成四个小波子带，并将其作为单一对象串联到扩散过程中（如图 1 所示）。这种模型在小波频谱而非原始图像空间上运行。因此，该模型可以利用高频信息为生成的图像添加更多细节。另一方面，小波子带比原始图像小四倍，这大大降低了采样过程的计算复杂度。

本文的方法基于 DDGAN 模型，输入是小波变换的四个小波子带。给定一个输入图像 x∈R 3 × H × W，将其分解为低子带和高子带，然后进一步串联形成矩阵 y∈R 12 × H 2 × W 2。该输入通过第一线性层投射到基础通道 D 上，与 DDGAN 相比，不会改变网络的宽度。因此，大多数网络的空间维度减少了四个，从而大大减少了计算量。

学习的损失函数

敌对损失

与 DDGAN 类似，它也是通过对抗性损失来优化生成器和判别器：

・重建损失和总体损失函数

除了上述敌对损失外，还增加了一个重建项，以防止频率信息的损失，并保持小波子带的一致性。这就是生成的图像与其地面实况之间的 L1 损失。

生成器的总体目标是对抗损失和重建损失的下一个线性组合：

其中，λ 是加权超参数。经过一定数量的采样步骤后，就能得到估计的去噪子带 y’0。最终图像可通过小波逆变换恢复：x’0 = IWT(y’0)。

包含小波的发电机

图 2：包含小波的发生器概览。

图 2 显示了拟议的小波嵌入生成器的结构。拟议的生成器采用 UNet 结构，包含 M 个下采样块和 M 个上采样块。相同分辨率的块之间也有跳转连接。不过，使用的是频率敏感块，而不是通常的下采样和上采样运算符。最低分辨率采用频率瓶颈块，以更好地关注低频和高频成分。

最后，利用小波下采样层引入频率残差连接，将原始信号 Y 纳入编码器的不同特征金字塔中。这里，Y 代表输入图像，Fi 代表 Y 的第 i 个中间特征图。

频率感知下采样和上采样模块

传统方法在下采样和上采样过程中使用模糊核来减少混叠伪影。本文则利用小波变换的固有特性，更好地进行上采样和下采样（如图 3 所示）。

这增强了这些操作对高频信息的识别能力。具体来说，下采样块接收输入特征 Fi、电位 z 和时间嵌入 t 的元组，并通过一系列层处理，返回下采样特征和高频子带。这些返回的子带可作为额外的输入，在上采样块中根据频率线索对特征进行上采样。

图 3. 频率感知下采样和上采样模块概览。

试验

数据集

实验在 32 x 32 CIFAR-10、64 x 64 STL-10 和 256 x 256 CelebA-HQ 及 LSUN-Church 数据集上进行。此外，还在 CelebA-HQ (512 和 1024) 高分辨率图像上进行了实验，以验证天安方法在高分辨率下的有效性。

估值指数

图像质量以弗雷谢特起始距离（FID）衡量，样本多样性以召回率（Recall）衡量；与 DDGAN 一样，FID 和召回率以 50 000 个生成样本计算。估算速度通过 300 次试验的平均推理时间来衡量，批量大小为 100。高分辨率图像（如 CelebA-HQ 512 x 512）的推理时间也是根据 25 个样本的批次计算得出的。

实验结果

图 4：Celeba 总部的发电实例

图 5：LSUN 中的生成示例。

在这里插入图片描述

表 1.CIFAR 的比较10

表 2. Celeba 总部的比较

表 3.LSUN 的比较

表 1、表 2 和表 3 列出了针对每个数据集与典型生成模型（如 VAE、GAN 和扩散模型）的比较结果。与VAESOTA或最强模型相比，所提出的方法在所有评价指标上都明显优于VAESOTA。特别是，图像质量的 FID 比VAE高出四倍多。

与GAN 的SOTA 相比，估计速度基本持平，图像质量也更高。在多样性方面，它在所有情况下都比GAN高 10%。

与 Diffusion模型和 DDGAN 相比，所提出的方法是Diffusion模型中估计速度最快的。特别是，它比 Diffusion的 SOTA快 500 多倍。图像质量和多样性也是最高的，在某些情况下比扩散模型的 SOTA高出1~2 个百分点。在所有评估指标上，它也都优于早期的DDGAN研究。

包含小波的发电机的有效性

我们在 CelebA-HQ 256×256 上测试了所提议的生成器每个单独组件的有效性。在这里，完整模型包括残差连接、上采样、下采样块和瓶颈块。如表 4 所示，每个部分都对模型的性能产生了积极影响。通过应用所有三个建议的组件，达到了 5.94 的最佳性能。然而，性能的提高在估算速度方面付出了较小的代价。

表 4.包含小波的发电机的有效性

生成一张图像的运行时间

此外，正如在实际应用中所预期的那样，所提出的方法在单幅图像中显示出极佳的速度。表 5 显示了时间和关键参数。所提出的方法能够在短短 0.1 秒内生成最大 1024 x 1024 的图像，这是第一个达到接近实时性能的扩散模型。

图 5.使用我们的完整模型对每个基准集生成的单幅图像进行估计所需的时间

结论

本研究引入了一种名为小波扩散的新扩散模型，该模型在图像质量和采样率方面都表现出色。通过将小波变换纳入图像和特征空间，所提出的方法达到了扩散模型中最先进的执行速度，缩小了与 GAN 中 SOTA 的差距，并获得了与 StyleGAN2 和其他扩散模型几乎相当的图像生成质量。此外，与基线 DDGAN 相比，所提方法的收敛速度更快，从而证实了所提框架的高效性。