扩散小波

论文阅读

摘要

扩散模型作为一种强大的高保真图像生成解决方案正在崛起,在许多情况下,其质量超过了gan。然而,它们缓慢的训练和推理速度是一个巨大的瓶颈,阻碍了它们在实时应用中的应用。最近的一种扩散GAN方法通过将采样步骤从数千个减少到几个,显著地减少了模型的运行时间,但它们的速度仍然大大落后于GAN的同类。本文提出了一种新的基于小波的扩散方案来减小速度差距。我们通过小波分解从图像和特征层提取低频和高频成分,并自适应处理这些成分,以提高处理速度,同时保持良好的生成质量。此外,我们提出使用重构项,有效地提高了模型训练的收敛性。在CelebA-HQ, CIFAR-10, LSUN-Church和STL-10数据集上的实验结果表明,我们的解决方案是提供实时和高保真扩散模型的基石。

介绍

这些模型将扩散过程还原为从随机噪声输入生成干净、高质量的输出。这些技术应用于各种数据领域和应用程序,但在图像生成任务中表现出最显著的成功。在各种数据集上,扩散模型在生成质量上优于最先进的生成对抗网络(GANs)[4,38]。更值得注意的是,扩散模型提供了更好的模式覆盖[14,22,41],并提供了一种灵活的方式来处理不同类型的条件输入,例如语义地图、文本、表示和图像[36]。。最近基于扩散的文本到图像生成模型[1,34,38]允许用户仅通过文本输入生成令人难以置信的逼真图像,开启了基于人工智能的数字艺术的新时代,并有望应用于其他各种领域。

尽管扩散模型显示出巨大的潜力,但它的运行速度非常慢,这是一个关键的弱点,阻碍了它们像gan一样被广泛采用。基础工作去噪扩散概率模型(Denoising Diffusion Probabilistic Models, ddpm)[13]需要1000个采样步骤来产生所需的输出质量,需要几分钟才能生成一张图像。已经提出了许多减少推理时间的技术[25,40],主要是通过减少采样步骤。然而,在DiffusionGAN之前最快的算法仍然需要几秒钟才能生成32×32图像,这比GAN慢了大约100倍。DiffusionGAN[50]通过将扩散和gan结合在一个系统中,在提高推理速度方面取得了突破性进展,最终将采样步骤减少到4步,生成32×32图像的推理时间缩短到几分之一秒。

尽管如此,它至少比StyleGAN慢4倍,并且随着输出分辨率的增加,速度差距不断扩大。此外,DiffusionGAN仍然需要较长的训练时间和较慢的收敛速度,这证实了扩散模型还没有为大规模或实时应用做好准备。

本文旨在通过引入一种新的基于小波的扩散方案来弥补速度差距。我们的解决方案依赖于离散小波变换,它将每个输入分解为低(LL)和高频(LH, HL, HH)分量的四个子带。我们在图像和特征级别上应用该变换。这使我们能够在保持输出质量相对不变的情况下显著减少训练和推理时间。在图像级别上,我们通过将空间分辨率降低四倍来获得高速提升。在特征层面上,我们强调了小波信息在发生器不同块上的重要性。通过这样的设计,我们可以在只产生边际计算开销的情况下获得相当大的性能改进。

我们提出的小波扩散提供了最先进的训练和推理速度,同时保持了高生成质量,通过包括CIFAR-10, STL-10, CelebA-HQ和LSUN-Church在内的标准基准实验得到了彻底的证实。我们的模型显著减少了扩散模型和gan之间的速度差距,针对大规模和实时系统。

  • 我们提出了一种新的小波扩散框架,该框架利用小波子带的降维来加速扩散模型,同时通过高频分量保持生成结果的良好视觉质量。
  • 我们在图像和特征空间中使用小波分解来提高生成模型的鲁棒性和执行速度。
  • 我们提出的小波扩散提供了最先进的训练和推理速度,这是实现实时和高保真扩散模型的基石。

相关知识

小波方法

真实世界的数据或者信号经常表现出缓慢变化的趋势或因瞬态而出现的震荡,另一方面,图像具有被边缘中断或者对比度突然变化的平滑区域,傅里叶变换不能有效代表突然的变化,这是因为傅里叶变换将数据表示为未在时间或空间上定位的正弦波之和,这些正弦波永远震荡。

为了很好准确分析突然变化的信号和图像,我们需要使用在时间和频率上都很好定位的一类新功能,就是小波变换。

img

方法

基于小波的扩散方案

首先,我们描述了如何在扩散过程中引入小波变换。我们将输入图像分解为四个小波子带,并将它们作为单个目标连接起来进行去噪处理(如图2所示)。这种模型不是在原始图像空间上执行,而是在小波谱上执行。因此,我们的模型可以利用高频信息来进一步增加生成图像的细节。同时,小波子带的空间面积比原始图像小4倍,大大降低了采样过程的计算复杂度。

image-20240422172619071

我们的方法建立在DDGAN模型上,其中输入是小波变换的4个小波子带。给定输入图像x∈R3×H×W,我们将其分解为一组低、高子带,并将它们进一步连接形成矩阵y∈R12×H 2 ×W 2。

然后,该输入通过第一个线性层投射到基本通道D,与DDGAN相比,保持网络宽度不变。因此,大多数网络都受益于空间维度减少4倍,大大减少了计算量。

设y0为时间步长为t的干净样本,而yt为从q(yt|y0)采样的损坏样本。在去噪过程中,生成器接收变量yt的元组,潜函数z ~ N(0, I)和时间步长t,以生成原始信号y0的近似值:y ‘ 0 = G(yt, z, t)。然后从可处理的后验分布q(yt - 1|yt, y ‘ 0)中提取预测的噪声样本y ‘ t - 1。鉴别器的作用是区分实对(yt - 1, yt)和假对(y ‘ t - 1, yt)。

根据[50],我们通过对抗性损失来优化生成器和鉴别器

在Eq.(4)中的对抗目标基础上,我们增加了一个重构项,既防止了频率信息的丢失,又保持了小波子带的一致性。它被表示为生成的图像与其基真值之间的L1损失:

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

经过定义的几个采样步骤后,我们获得了估计的去噪子带y ‘ 0。最后的图像可以通过小波逆变换x ‘ 0 = IWT(y ‘ 0)恢复。我们在算法1中描述了采样过程。

小波嵌入网络

接下来,我们通过生成器将小波信息进一步纳入特征空间,增强对高频成分的感知。这有利于最终图像的清晰度和质量。

图3说明了我们提出的小波嵌入发生器的结构。它遵循[44]的UNet结构,M个下采样和M个上采样块,加上相同分辨率的块之间的跳过连接,M是预定义的。然而,我们没有使用正常的下采样和上采样算子,而是用频率感知块代替它们。在最低分辨率下,我们采用频率瓶颈块来更好地关注低频和高频组件。最后,为了将原始信号Y合并到编码器的不同特征金字塔中,我们使用小波下采样层引入了频率残差连接。设Y为输入图像,Fi为Y的第i个中间特征映射。我们将在下面讨论新引入的分量:.

频率感知下采样和上采样块。传统方法依赖于模糊核的下采样和上采样过程,以减轻混叠伪影。相反,我们利用小波变换的固有特性来进行更好的上采样和下采样(如图4所示)。实际上,这加强了对这些操作的高频信息的认识。特别是,下采样块接收输入特征Fi、潜伏z和时间嵌入t的元组,然后通过一系列层处理以返回下采样特征和高频子带。这些返回的子带作为基于上采样块中的频率线索的上采样特征的额外输入。

image-20240422173330672

嵌入式小波发生器的说明。为了简化,忽略时间步嵌入t和潜嵌入z,而是将它们注入到去噪过程的各个块中。输入是时间步长为t的形状为[12 × H × W]的噪声小波子带,这些子带由我们提出的一系列组件处理,包括频率感知的上采样和下采样块、频率残差连接和一个全新的频率瓶颈块。模型的输出是无扰动输入的近似值。

image-20240422173429792

频率瓶颈块位于中间阶段,中间包括两个频率瓶颈块和一个注意块。每个频率瓶颈块首先将特征映射Fi划分为低频子带Fi、ll和高频子带Fi、H的拼接。然后将Fi,ll作为输入传递给重新分配块以进行更深入的处理。处理后的低频特征映射和原高频子带Fi、H通过IWT变换回原空间。有了这样的瓶颈,该模型可以专注于学习低频子带的中间特征表示,同时保留高频细节。

[44]中网络的原始设计通过跨行卷积下采样层将原始信号Y合并到编码器的不同特征金字塔中。相反,我们使用小波下采样层将输入Y的剩余快捷方式映射到相应的特征维度,然后将其添加到每个特征金字塔中。具体来说,Y的残差捷径被分解成四个子带,然后将这些子带连接并馈送到卷积层进行特征投影。这个捷径的目的是丰富对特征嵌入的频率源的感知。

实验

结论

本文介绍了一种新的基于小波的扩散方案,该方案在图像保真度和采样速度方面都有优异的性能。通过对图像和特征空间进行小波变换,我们的方法可以达到扩散模型最先进的运行速度,缩小了与StyleGAN模型[19,20,57]的差距,同时获得与StyleGAN2和其他扩散模型相当的图像生成质量。此外,我们的方法提供了比基线DDGAN更快的收敛速度[50],证实了我们提出的框架的效率。有了这些初步结果,我们希望我们的方法可以促进未来对实时和高保真扩散模型的研究。

代码解读

-------------已经到底啦!-------------