生成扩散模型漫谈(二十二):信噪比与大图生成(上)
By 苏剑林 | 2024-04-08 | 41111位读者 |盘点主流的图像扩散模型作品,我们会发现一个特点:当前多数做高分辨率图像生成(下面简称“大图生成”)的工作,都是先通过Encoder变换到Latent空间进行的(即LDM,Latent Diffusion Model),直接在原始Pixel空间训练的扩散模型,大多数分辨率都不超过64*64,而恰好,LDM通过AutoEncoder变换后的Latent,大小通常也不超过64*64。这就自然引出了一系列问题:扩散模型是不是对于高分辨率生成存在固有困难?能否在Pixel空间直接生成高分辨率图像?
论文《Simple diffusion: End-to-end diffusion for high resolution images》尝试回答了这个问题,它通过“信噪比”分析了大图生成的困难,并以此来优化noise schdule,同时提出只需在最低分辨率feature上对架构进行scale up、多尺度Loss等技巧来保证训练效率和效果,这些改动使得原论文成功在Pixel空间上训练了分辨率高达1024*1024的图像扩散模型。
LDM回顾 #
在进入正题之前,我们不妨先反过来想一想:为什么LDM能成功成为主流的扩散模型做法?笔者认为,主要原始是两方面:
1、不管是应用还是学术,用LDM的主要原因想必是效率:当前主流的工作都直接重用了LDM论文所开源的训练好的AutoEncoder,它的Encoder部分会将512*512的图像变成了64*64的Latent,相当于说只用到64*64分辨率这个级别的算力和时间,就可以生成512*512的图像,这个效率显然是非常吸引人的;
2、LDM契合了FID这个指标,这让它看起来是效果无损的:FID全称是“Fréchet Inception Distance”,其中Inception是指用ImageNet预训练的InceptionV3模型作为Encoder编码图片,然后假设编码特征服从高斯分布来算$\mathcal{W}$距离,而LDM也是先Encoder编码,两个Encoder虽然不完全相同,但也有一定共性,因此在FID上表现为几乎无损。
我们还可以稍微展开一下。LDM的AutoEncoder在训练阶段组合了很多内容——它的重构Loss并不只有常规的MAE或者MSE,还包括对抗Loss和Perceptual Loss,对抗Loss用来保证重构结果的清晰度,而Perceptual Loss用来保证重构结果的语义和风格的相似性。Perceptual Loss跟FID很相似,都是用ImageNet模型的特征计算的相似性指标,只不过用的不是InceptionV3而是VGG-16,由于训练任务的相似性,可以猜测两者特征有很多共性,因此Perceptual Loss的加入变相地保证了FID的损失尽可能少。
此外,LDM的Encoder对原始图像来说是降维的,比如原始图像大小为512*512*3,直接patchify的话结果是64*64*192,但LDM的Encoder出来的特征是64*64*4,降低到了1/48,同时为了进一步降低编码特征的方差,避免模型“死记硬背”,LDM还对Encoder出来的特征加了相应的正则项,可选的有VAE的KL散度项或VQ-VAE的VQ正则化。降维和正则的设计,都会压缩特征的多样性,提高特征的泛化能力,但也会导致重构难度增加,最终导致了有损的重构结果。
到这里,LDM能成功的原因其实就“豁然开朗”了:“降维 + 正则”的组合,降低了Latent的信息量,从而降低了在Latent空间学习扩散模型的难度,同时Perceptual Loss的存在,保证了重构虽然有损但FID几乎无损(Perceptual Loss的Encoder跟FID一样都用InceptionV3理论上更好)。这样一来,对于FID这个指标来说,LDM几乎就是免费午餐了,因此不管是学术和工程都乐意沿用它。
信噪之比 #
尽管LDM简单高效,但它毕竟是有损的,其Latent只能保持宏观上的语义,局部细节可能会严重缺失。而在之前的文章《“闭门造车”之多模态思路浅谈(一):无损输入》中,笔者表达过一个观点:当作为输入时,图像最好的表示方式就是原始Pixel数组。基于这个观点,笔者最近都比较关注直接在Pixel空间上训练的扩散模型。
然而,将低分辨率(比如64*64)图像的扩散模型配置直接应用于高分辨率(比如512*512)的大图生成时,会存在算力消耗过大、收敛速度太慢等问题,而且效果上也比不上LDM(至少FID指标如此),Simple diffusion逐一分析了这些问题并提出了相应的解决方案。其中,笔者认为利用“信噪比(Signal-to-Noise Ratio,SNR)”的概念来分析高分辨率扩散模型的学习效率低问题最为精彩。
具体来说,Simple diffusion观察到,如果我们给高分辨率图像加上某个方差的noise,那么相对于加上同样方差的noise的低分辨率图像来说,它的信噪比其实更高,原论文的Figure 3非常直观地演示了这一点,如下图所示。第一行图片,是由512*512的图片加了特定方差的noise后再降采样(平均Pooling)到64*64的,而第二行则是直接在64*64的图片加上同样方差的noise,很明显第一行的图片更加清晰,也就是相对信噪比更高了。
所谓“信噪比”,顾名思义即“信号与噪声的强度之比”,信噪比更高(即噪声的占比更低)意味着去噪更容易,换言之训练阶段Denoiser面对的更多是简单样本,但实际上大图生成的难度显然更高,也就是说我们的目标是一个更难的模型,但却给了更简单的样本,因此导致了学习效率的低下。
向低看齐 #
我们也可以从数学上描述这一点。沿用本系列的记号,通过加噪来构造$\boldsymbol{x}_t$的运算可以表示为
\begin{equation}\boldsymbol{x}_t = \bar{\alpha}_t \boldsymbol{x}_0 + \bar{\beta}_t \boldsymbol{\varepsilon},\quad \boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0},\boldsymbol{I})\end{equation}
其中$\bar{\alpha}_t,\bar{\beta}_t$就称为noise schedule,满足$\bar{\alpha}_0=\bar{\beta}_T=1, \bar{\alpha}_T=\bar{\beta}_0=0$,此外一般来说它们还有额外的约束,比如在DDPM中通常是$\bar{\alpha}_t^2 + \bar{\beta}_t^2=1$,本文将沿用这个约束。
对于一个随机变量来说,信噪比是均值平方与方差之比。给定$\boldsymbol{x}_0$,$\boldsymbol{x}_t$的均值显然是$\bar{\alpha}_t \mathbb{E}[\boldsymbol{x}_0]$,方差则是$\bar{\beta}_t^2$,于是信噪比为$\frac{\bar{\alpha}_t^2}{\bar{\beta}_t^2}\mathbb{E}[\boldsymbol{x}_0]^2$,由于我们总是在给定$\boldsymbol{x}_0$下讨论,因此我们也可以简单地说信噪比就是$SNR(t) = \frac{\bar{\alpha}_t^2}{\bar{\beta}_t^2}$
当我们对$\boldsymbol{x}_t$使用$s\times s$大小的平均Pooling时,每个$s\times s$的patch通过取平均变成了一个标量,即
\begin{equation}\frac{1}{s^2}\sum_{i=1}^s \sum_{j=1}^s\boldsymbol{x}_t^{(i,j)} = \bar{\alpha}_t\left(\frac{1}{s^2}\sum_{i=1}^s \sum_{j=1}^s \boldsymbol{x}_0^{(i,j)}\right) + \bar{\beta}_t\left(\frac{1}{s^2}\sum_{i=1}^s \sum_{j=1}^s \boldsymbol{\varepsilon}^{(i,j)}\right) ,\quad \boldsymbol{\varepsilon}^{(i,j)}\sim\mathcal{N}(0, 1)\end{equation}
平均Pooling不改变均值,但会降低方差,从而提高信噪比,这是因为正态分布的可加性得出
\begin{equation}\frac{1}{s^2}\sum_{i=1}^s \sum_{j=1}^s \boldsymbol{\varepsilon}^{(i,j)}\sim\mathcal{N}(0, 1/s^2)\end{equation}
所以在同一noise schedule下,如果我们将高分辨率图像通过平均Pooling来对齐低分辨率,那么就会发现信噪比更高,是原来的$s^2$倍:
\begin{equation}SNR^{w\times h\to w/s\times h/s}(t) = SNR^{w/s\times h/s}(t) \times s^2 \end{equation}
反过来想,如果我们已经有一个在低分辨率图像上调好了的noise schedule $\bar{\alpha}_t^{w/s\times h/s},\bar{\beta}_t^{w/s\times h/s}$,那么当我们想要scale up到更高分辨率时,应该要调整noise schedule为$\bar{\alpha}_t^{w\times h},\bar{\beta}_t^{w\times h}$,使得它降采样到低分辨率后,其信噪比能够跟已经调好的低分辨率的noise schedule对齐,这样才能最大程度上“传承”已经低分辨率扩散模型的学习效率,即
\begin{equation} \frac{(\bar{\alpha}_t^{w\times h})^2}{(\bar{\beta}_t^{w\times h})^2} \times s^2 = \frac{(\bar{\alpha}_t^{w/s\times h/s})^2}{(\bar{\beta}_t^{w/s\times h/s})^2} \end{equation}
如果加上约束$\bar{\alpha}_t^2 + \bar{\beta}_t^2=1$,那么就可以从$\bar{\alpha}_t^{w/s\times h/s},\bar{\beta}_t^{w/s\times h/s}$中唯一地解出$\bar{\alpha}_t^{w\times h},\bar{\beta}_t^{w\times h}$。这就解决了高分辨率扩散的noise schedule设置问题。
架构拓展 #
为了做好大图的扩散生成,除了要调整noise schedule之外,我们还需要把架构也scale up上去,因为前面我们也已经说了,大图生成是一个更难的问题,因此理应需要更加重量级的架构。
扩散模型常用的就是U-Net或者U-Vit,两者都是先逐渐降采样然后逐渐上采样,比如512*512的输入,一般先进行一个block的运算,然后降采样到256*256,接着进行新一个block的运算,在降采样到128*128,依此类推,降采样到一个最低的分辨率16*16,接下来再次重复这个过程,但将降采样改为上采样,直到分辨率恢复512*512。默认设置下,我们会将参数平均分到每一个block中,但这样一来靠近输入和输出的block由于输入尺寸都很大,因此计算量会急剧增加,导致模型训练效率低下甚至不可行。
Simple diffusion提出了两个应对方案。第一,它提出可以直接在第一层(而不是第一个block,每个block有多个层)之后就降采样,并且考虑一步到位低降到128*128甚至64*64,最后输出的时候,也是在最后一层之前才从64*64或者128*128直接上采样到512*512,这样模型的大部分block所处理的分辨率都降低了,从而降低了整体计算量;第二,它提出将模型所scale up的层都放到最低分辨率(即16*16)之后,而不是平摊到每一个分辨率的block,即新增的层处理的都是16*16的输入,包括Dropout也都只加入到低分辨率的层中,这样一来分辨率增加带来的计算压力就明显减少了。
此外,为了进一步稳定训练,论文提出了“多尺度Loss”的训练目标。默认情况下,扩散模型的Loss等价于MSE损失
\begin{equation}\mathcal{L}=\frac{1}{wh}\Vert \boldsymbol{\varepsilon} - \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\bar{\alpha}_t \boldsymbol{x}_0 + \bar{\beta}_t \boldsymbol{\varepsilon}, t)\Vert^2\end{equation}
Simple diffusion将它泛化为
\begin{equation}\mathcal{L}_{s\times s} = \frac{1}{(w/s)(h/s)}\big\Vert \mathcal{D}_{w/s\times h/s}[\boldsymbol{\varepsilon}] - \mathcal{D}_{w/s\times h/s}[\boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\bar{\alpha}_t \boldsymbol{x}_0 + \bar{\beta}_t \boldsymbol{\varepsilon}, t)]\big\Vert^2\end{equation}
其中$\mathcal{D}_{w/s\times h/s}[\cdot]$是通过平均Pooling将输入变换到$w/s\times h/s$的下采样算子,原论文取了多个$s$对应的Loss进行平均,作为最终的训练目标。这个多尺度Loss的目标也很明显,跟通过信噪比对齐来调整noise schedule一样,都是保证训练出来的高分辨率扩散模型至少不差于直接训练的低分辨率模型。
至于实验部分,大家自行看原论文就好。Simple diffusion实验的最大分辨率是1024*1024(在附录中提到),效果都尚可,并且对比实验表明上述提出的一些技巧都是有提升的,最终直接在Pixel空间中训练出来的扩撒模型,相比LDM也取得了有竞争力的效果。
文章小结 #
在这篇文章中,我们介绍了Simple diffusion,这是一篇探索如何直接在Pixel空间中端到端地训练图像扩散模型的工作,利用了信噪比的概念介绍了高分辨率扩散模型的训练效率低问题,并由此来指标调整新的noise schedule,以及探索了如何尽可能节约算力成本地scale up模型架构。
转载到请包括本文地址:https://www.kexue.fm/archives/10047
更详细的转载事宜请参考:《科学空间FAQ》
如果您还有什么疑惑或建议,欢迎在下方评论区继续讨论。
如果您觉得本文还不错,欢迎分享/打赏本文。打赏并非要从中获得收益,而是希望知道科学空间获得了多少读者的真心关注。当然,如果你无视它,也不会影响你的阅读。再次表示欢迎和感谢!
如果您需要引用本文,请参考:
苏剑林. (Apr. 08, 2024). 《生成扩散模型漫谈(二十二):信噪比与大图生成(上) 》[Blog post]. Retrieved from https://www.kexue.fm/archives/10047
@online{kexuefm-10047,
title={生成扩散模型漫谈(二十二):信噪比与大图生成(上)},
author={苏剑林},
year={2024},
month={Apr},
url={\url{https://www.kexue.fm/archives/10047}},
}
April 8th, 2024
最近看到一个改进LDM的工作,好像实现了无条件的2048*1024的图片生成.效果看着还不错,苏神怎么看。 https://arxiv.org/abs/2403.12915
看上去是通过架构上的优化来实现的,金字塔特征的做法也算比较常规了,个人暂时不大关注架构细节上的调优。
另外跟本篇方法并不矛盾。因为LDM往上做的话,即便encoder之后的特征图,也可能会有512*512这个级别。
April 8th, 2024
苏神,对于flow模型来说,是不是就不存在scheduler的问题,理论上各种分辨率都能统一
flow模型本身就没有noise schedule,所以不存在这个问题,但normalizing flow模型目前属于效果最差的生成模型之一。。。
continuous normalizing flow效果很好的,SD3就用了rectified flow
rectified flow不算normalizing flow吧?要不然扩散模型也算是normalizing flow了。
我理解的normalizing flow,是构建一个可逆模型,算出它的雅可比行列式,然后利用最大似然作为目标函数去训练的,并不是只要可逆就算normalizing flow了。
那么rectified flow会不会有这个问题呢?
推了一下公式,好像降采样到低分辨率后信噪比依旧是变成原本的$s^2$倍
本文的设定是$\boldsymbol{x}_t = \bar{\alpha}_t \boldsymbol{x}_0 + \bar{\beta}_t \boldsymbol{\varepsilon},\quad \boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0},\boldsymbol{I})$,reflow也是这个形式,自然也有这个问题
苏神,请问高分辨率SNR对齐低分辨率,输入给网络的$t$是不变的,只改变$\alpha_t$和$\beta_t$吗?
是的
如果把模型架构和训练方式分开考虑的话,rectified flow是一种特定的训练continuous normalizing flow的方式。它抛弃了最大似然估计,而利用(从扩散模型和得分匹配那里学来的)flow matching的方式训练flow。
所以,在我心目中只要可逆就算flow,扩散模型也算flow的一种。当然,它们更广义的定义也许是基于输运(transportation)的生成模型。这样一来,当下好用的生成模型就两类:基于输运的和基于自回归的。
如果GAN的生成器取输入输出维度一致呢?虽然没有理论保证用这样用GAN训练出来的生成器可逆,但从数值计算的概率来讲,数学上严格不可逆的概率几乎为0(比如对于线性变换,不可逆意味着行列式为0,而行列式精确等于0的概率显然是非常小的)
April 11th, 2024
苏神,之前有论文提出了HDiT架构,能直接在1024×1024的Pixel Space训练。https://arxiv.org/abs/2401.11605
感谢分享,我读了之后再来交流。
确实,本文所介绍的论文已经是去年初的工作了,对于瞬息万变的AI领域来说算是很老了,只不过感觉里边的分析还是很有借鉴意义的,所以拿来学习一下。
April 12th, 2024
苏神你好,我本人也很关注pixel space的diffusion。结合你上篇文章我觉得你说的很对 latent这种方式局限很多,需要一个task specific的metric来指导压缩。个人也觉得长期以来不可取。而且本身这种架构对于一些应用就很局限。
关于你提到的高分辨率生成,还有一篇on the importance of noise scheduling for diffusion models里面也提到了noise schedule的问题,表达的意思是一样的,但他的处理稍微有点不一样。希望听听你的想法。另外 他提到了RIN这个结构,感觉没多少人follow 但我也觉得挺有意思的,可以在forward过程中使用一个固定长度token的数量来减少计算量。因为作者认为分辨率虽然大了 但信息量可以仍然是一定的。
谢谢认可和推荐。
on the importance of noise scheduling for diffusion models 感觉偏向于实验了,不过结论应该差不多;RIN我稍微扫了一下,感觉它能节省计算量的本质原因是训练阶段强行截断了递归的梯度(没细看,如果错误请指正),这样设计痕迹太严重,并且让人怀疑它是否scaleable~
April 13th, 2024
苏神好,感謝您的分享。
想請問LDM關於KL-reg和VQ-reg,原論文在附錄G.中提到:
(ii) For a VQ-regularized latent space, we extract z before the quantization layer
and absorb the quantization operation into the decoder, i.e. it can be interpreted as the first layer of D.
VQ-reg使用Encode後量化前的特徵,相比於KL-reg的特征分佈,VQ-reg的特徵分佈似乎不是"接近正态分布",這樣LDM的訓練還行得通嗎? 是否能多分享些有關LDM的細節,感謝。
KL-reg和VQ-reg都是对特征的正则项,目的都是让特征的信息有所损失(信息瓶颈),从而导致有损的重构。在信息瓶颈下迫使模型完成重构,则会让模型学习到更有泛化能力的特征。
关于VQ-reg对特征的有损性,我们可以参考 https://kexue.fm/archives/9984 一文对VQ的熵分析。简单来说,对序列严重缩短的特征进行VQ之后,信息会有非常明显的损失,而使用VQ前的特征则好一些,但由于学习过程中VQ前的特征会尽量向VQ后的特征靠齐,所以依然有正则化作用。
April 18th, 2024
sd3(https://arxiv.org/pdf/2403.03206v1.pdf)也给出了类似的观点,即对于高分辨率图需要添加更强的noise。不过他们针对的是多尺度训练的问题
实际上sd3似乎是从另一个角度解释了低分辨率训练模型直接生成高分辨率图像存在的gap。根据sd3的观点,是不是《Upsample Guidance: Scale Up Diffusion Models without Training》中高分辨率细节项ϵθ(xt,t)在使用时,也可以对t做一个重新映射
看了看,sd3对$t$所做的调整,貌似都是训练阶段就必须引入的,而不是推理阶段才做的免训练改动?
April 28th, 2024
苏神,信噪比章节的对比方式感觉不严谨。
高斯分布样本的均值的分布方差是原来的 1/n, 均值 pooling 之后 实际上等效于对于下采样的图片加入 $\frac{1}{n}\sigma^2$ n 是pooling 窗口的pixel 个数 的噪声,感觉文章将二者认为同样的噪声水平不妥
没看懂你这里想表达什么。你想表达Pooling之后等效于对低分辨率图加入$\frac{1}{n}\sigma^2$方差的噪声?然后呢?哪里不妥?
原文论点是:
> 如果我们给高分辨率图像加上某个方差的noise,那么相对于加上同样方差的noise的低分辨率图像来说,它的信噪比其实更高
而对比的是高分辨率加噪再下采样的结果, 不能说明 高分辨率加噪 和 低分辨率 加噪的信噪比的差异。 毕竟下采样带低通效果,也就是附带了去噪过程。 等效噪声是 1/n 。
你的个别字眼我始终没理解想要表达什么(比如最后一句突然出现的“等效噪声是 1/n”)...
总的来说,博客中想要表达的意思是【对于不同分辨率的同一幅图像,如果加上同样强度的高斯噪声,那么高分辨率图像远远看起来会更加清晰】。我不是很理解,你想要强调的是原文的表达不大严谨,还是想要否定这个结论...
如果是前者,可能确实有点不严谨,因为高分辨率图像加上同样噪声模糊的是高分辨率细节,不论一概而论说信噪比更高了,但就本文的主题来说,这样表达也似乎不会引起歧义。
June 30th, 2024
苏神你好,我想问一下,这个预训练好了的autoencoder可以直接拿来用于其他领域任务吗,比如医学领域,还是说得用医学数据再重新训练比较好?
一般垂直领域都要微调一下比较好
August 6th, 2024
假如使用DDPM T=1000步来训练的话,t=1000时的$\bar\alpha_t$并非0,也就导致信噪比仍然存在。那这时是不是就会导致inference和training的不一致呢?
另外信噪比是不是还可以和预测原图loss及预测噪声loss来建立关系哈
从数值精度的视角来看,其实$\bar{\alpha}_T$跟0没有实质差别了。
不过Stable Diffusion里设置的数值还是有差别的,snr在T步并不小,感觉这个还是会有些影响大
其实我们也可以认为存在一个绝对等于零的$\bar{\alpha}_t$,只不过既然绝对等于零,也就是没有任何有效输入,那么训练起来就没有意义了,干脆就不训了。这不会影响后面关于采样过程包括相关加速采样方法的推导。
* inference与training是否一致的关键是: P(x_T)的分布是否是高斯
* x_T|x_0信噪比存在 => x_T/x_0不是纯噪音 => P(x_T|x_0)不是高斯 ≠> P(x_T)不是高斯
* LDM中,P(x_0)本身就高斯分布,P(x_T) = ∫P(x_T|x_0)P(x_0)d(x_0)基本上就是纯高斯了。