受到经典图像表示方法拉普拉斯金字塔(Laplacian Pyramid)的启发,英伟达最近公布了一种叫做 Laplacian Diffusion Model (拉普拉斯扩散模型,后文简称 LaDM)的新型像素空间扩散模型,并用这种架构实现了文生图、超分辨率、ControlNet 等多种任务。在这篇博文里,我们来着重学习一下这种新型扩散模型的设计思想。
以往工作
扩散模型奠基之作 DDPM 及其升级版 ADM (Diffusion Models Beat GANs on Image Synthesis) 都是像素空间里的扩散模型。相比 LDM (隐扩散模型,即 Stable Diffusion),这类扩散模型不需要额外的自编码器来压缩图像,避免了编码解码带来的精度损失。
将图像从分辨率的维度拆解是一种很常见的思想。比如 Cascaded Diffusion Models 就是一种先生成低分辨率图像,再不断超分的扩散模型。今年比较有名的 VAR(Visual Autoregressive Modeling: Scalable Image Generation via Next-Scale Prediction)也是一种按分辨率自回归的生成模型。
和这篇工作非常相关的早期工作是苹果在 2022 发表的f-DM: A Multi-stage Diffusion Model via Progressive Signal Transformation。f-DM 将扩散模型的加噪推广到了降采样、模糊等其他退化策略上。降采样版的 f-DM 有非常多的设计和本工作很像。苹果该团队次年发表的Matryoshka Diffusion Models 也用到了按分辨率逐次生成的设计。
将拉普拉斯金字塔融入扩散模型
拉普拉斯金字塔是一种图像表示方法,它把图像按频率成分拆成几张分辨率不同的图像,分辨率越低的图像表示频率越低的图像成分。我们直接通过下面的例子学习它的原理。假如x
是原图,那么x(3)=down(down(x))
,x(2)=down(x)-up(x(3))
,x(1)=x-up(down(x))
。对x(1), x(2), x(3)
求加权和就可以还原输入图像。
受此启发,LaDM 将扩散模型的训练过程也用类似的方法分解:设 为训练图片集合, 分别是拉普拉斯金字塔不同成分构成的集合,那么我们在,, 上分别训练三个去噪模型。也就是说,不同分辨率的模型生成不同层级的拉普拉斯金字塔复原结果。
根据经验,扩散模型早期(加噪后期)生成低频内容,后期(加噪前期)生成高频内容,所以训练时我们让不同分辨率的输入图像随噪声的衰退速度也不同。图像所代表的频率越低,衰减速度越慢,越需要从早期开始去噪。这样,在生成时,我们能生成到中途后再逐渐加上高频细节。
在采样过程中,我们按照下图所示的路线从低频到高频生成图像。有了该分辨率的初始图像后,按正常 DDPM 采样的步骤就可以生成当前分辨率的图像了。问题在于某分辨率的初始图像怎么从上一个分辨率过渡而来。
在切换当前带噪图像的分辨率时,我们既要放大其中的清晰图像(信号),也要放大其中的噪声。观察上一张图和下面的图,在分辨率切换时,新的高频成分(上图中的在时刻 3 及 在时刻 2)是一张纯黑图,新信号为零,所以对于信号的部分我们可以直接放大。而放大噪声时,我们要做一些噪声强度上的修改,保证放大后信噪比不变。这部分的细节详见论文。
1K 图像生成
为了生成 分辨率的图像,LaDM 采用了两阶段 Cascaded Diffusion Model 的设计,让生成高分辨率的图像约束于低分辨率图像。另外,由于注意力操作的时间复杂度很高,一般的像素扩散模型只能做到 大小。为了解决此问题,LaDM 依然用一个 的去噪模型来生成 1K 图片,但输入前后用小波变换来压缩/复原图像。
批判性分析与总结
这篇文章是一篇由公司发表的技术报告,展示了很多可视化结果,却没有任何定量结果,代码也没有开源,不知道它的生成能力和其他模型比起来如何。
这篇文章提出的模型虽然是像素空间扩散模型,但是其拉普拉斯金字塔的设计与模型是像素空间模型还是隐空间模型无关。我们完全可以把这套设计搬到隐空间上。VAR 已经向我们证明了对隐空间图像做拉普拉斯分解是可行的。另外,这篇文章的主干网络是 U-Net 而不是 DiT。想对这个工作做一点简单的改进的话,可以弄一个 LDM + DiT 版本的。
LaDM 设计最巧妙的点是其加噪过程,频率越高的成分越早变成纯噪声。这样的话我们可以在图像生成到一半的时候再直接把高频成分加上。如果高频成分一直在的话,我们还需要额外的设计在切换分辨率时把缺少的高频加上。
有工作证明神经网络不擅长拟合高频信息。因此,在图像任务中,手动将输入图像拆成不同频率成分可能有助于网络的学习。我们可能可以沿着这个思路去改进之前多种图像任务的输入。