【超分辨率专题】OSEDiff:针对Real-World ISR的单步Diffusion
OSEDiff:针对Real-World ISR的单步Diffusion(2024,NeurIPS)
- 专题介绍
- 一、研究背景
- 二、方法细节
- 2.1 何为单步Diff
- 2.2 损失函数是关键
- 2.3 其他巧思
- 三、实验论证
- 3.1 定量与定性
- 3.2 消融实验
- 四、总结与思考
本文将对《One-Step Effective Diffusion Network for Real-World Image Super-Resolution》这篇文章进行解读,这是一篇关于SR领域内的单步Diff技术方案,简称OSEDiff,目标是进一步提升Real-World SR的效果和Diffusion-based方案的推理速度及稳定性。参考资料如下:
[1]. OSEDiff论文地址
[2]. OSEDiff代码地址
专题介绍
现在是数字化时代,图像与视频早已成为信息传递的关键载体。超分辨率(super resolution,SR)技术能够突破数据源的信息瓶颈,挖掘并增强低分辨率图像与视频的潜能,重塑更高品质的视觉内容,是底层视觉的核心研究方向之一。并且SR技术已有几十年的发展历程,方案也从最早的邻域插值迭代至现今的深度学习SR,但无论是经典算法还是AI算法,都在视觉应用领域内发挥着重要作用。
本专题旨在跟进和解读超分辨率技术的发展趋势,为读者分享有价值的超分辨率方法,欢迎一起探讨交流。
系列文章如下:
【1】SR+Codec Benchmark
一、研究背景
Real-World ISR(Real-ISR)是目前超分领域重点关注的方向之一,毕竟现实数据退化要复杂的多,导致很多训练效果不错的模型,面对Real-World数据时表现欠佳。因此,很多方案都在想办法提升模型在Real-World问题上的效果。real-esrgan就是其中比较出色且具有代表性的方法之一。
训练Real-ISR模型其实就聚焦于两个主要难题。
-
如何构建符合真实退化的LQ-HQ训练图像对;
-
如何保证恢复图像的自然度,也就是恢复图像遵循HQ自然图像的分布。
那么针对问题一有两类方式:
-
采集Real-World成对数据,例如长短焦数据组合采集,但成本太高,覆盖面太少;
-
构建以BSRGAN、Real-ESRGAN为代表的复杂合成退化流程,基础退化算子随机混合+高阶退化模型生成LQ-HQ对。(目前的主流方案)
针对问题二有三类方式:
-
传统映射式网络+pixel损失级别的方案。但容易导致结果过于平滑。
-
基于GAN的生成对抗网络,配合感知类损失。在SR效果上有突破,但对抗训练不稳定且易出瑕疵。
-
基于生成扩散模型的方案,特别是text-to-image(T2I)模型,具备丰富的自然图像先验,可以很好地利用这些信息来提升任务的自然度和感知质量。但这类方案的输出存在随机性且推理成本高(现有方案推理要数十步甚至数百步)。
为了进一步解决上述问题,由此衍生出one-step diffusion的SR方案,但在生成高质量细节方面仍差强人意。为此,作者提出了OSEDiff,取名非常直白,“有效的单步扩散模型”。
其核心亮点在于:
- 非噪声输入:直接以低质量(LQ)图像作为扩散过程起点,而非随机噪声,消除输出不确定性;
- 轻量又高效:通过 LoRA 微调预训练 Stable Diffusion(SD)模型,并采用变分分数蒸馏(variational score distillation,VSD)正则化策略,使得生成分布更接近真实分布,且仅需一步前向传播即可生成高质量图像,推理速度提升显著,训练参数量仅8.5M;
二、方法细节
我们可将Real-ISR任务建模成一个优化问题,核心公式如下:
x^H=argminxH(Ldata(Φ(xH),xL)+λLreg(xH))\hat{x}_H = \arg\min_{x_H} \left( \mathcal{L}_{\text{data}}(\Phi(x_H), x_L) + \lambda \mathcal{L}_{\text{reg}}(x_H) \right)x^H=argxHmin(Ldata(Φ(xH),xL)+λLreg(xH))
其中xL{x}_LxL为低质量图像,x^H\hat{x}_Hx^H为预估的高质量图像,xH{x}_HxH为真实高质量信息,Φ\PhiΦ为退化函数,Ldata{L}_{\text{data}}Ldata为衡量优化保真度的数据项结果,Lreg{L}_{\text{reg}}Lreg为利用自然图像先验信息的正则化项,λ\lambdaλ为平衡系数。
若用神经网络训练的方式解此问题,可将公式转化为:
θ∗=argminθE(xL,xH)∼S[Ldata(Gθ(xL),xH)+λLreg(Gθ(xL))]\theta^* = \arg\min_{\theta} \mathbb{E}_{(x_L, x_H) \sim S} \left[ \mathcal{L}_{\text{data}} \left( G_{\theta}(x_L), x_H \right) + \lambda \mathcal{L}_{\text{reg}} \left( G_{\theta}(x_L) \right) \right]θ∗=argθminE(xL,xH)∼S[Ldata(Gθ(xL),xH)+λLreg(Gθ(xL))]
其中θ\thetaθ和θ∗\theta^*θ∗为模型GθG_{\theta}Gθ的参数。Ldata{L}_{\text{data}}Ldata和Lreg{L}_{\text{reg}}Lreg分别代表不同的损失函数。当前Diffusion-based的主流方式是拿预训练的SD模型做微调,以LQ图像作为控制信号,来实现SR任务。例如PASD,SeeSR,CCSR,SUPIR,SinSR(也是单步DIff SR)。该论文自然也直接用SD作为主要模块了。同时,以KL散度作为正则项,为了使输出分布与真实世界的HQ分布更加一致。
关于正则项的选择,作者认为现有工作用对抗训练来实例化上述目标,但判别器通常也是从头训练,不一定能学到完整的HQ分布,容易出现次优解。所以作者选用了预训练T2I模型(也是SD)作为约束分布的新选择,有丰富的自然图像分布先验,并利用分数蒸馏(score distillation)方法来构建KL散度,论文中采用了变分分数蒸馏(VSD),更具有鲁棒性和多样性,可探索更广泛的分布空间。
简单理解,VSD的作用就是让单步Diff学习到多步Diff的能力,同时使得潜空间分布与真实分布更加靠近。
让我们来看一下具体网络架构。
其实就分为两部分,生成网络和正则化网络。
- 生成网络GθG_{\theta}Gθ由一个可训练VAE编码器、可微调SD生成模型和固定参数的VAE解码器构成,分别记为EθE_{\theta}Eθ,ϵθ\epsilon_{\theta}ϵθ和DθD_{\theta}Dθ。会计算输出x^H\hat{x}_Hx^H与GTxH{x}_HxH的MSE损失和LPIPS损失,以更新生成网络的编码器和生成器的参数。
- 正则化网络包括两个预训练SD模型,一个可微调,另一个固定,分别记为ϵϕ′\epsilon_{\phi'}ϵϕ′和ϵϕ\epsilon_{\phi}ϵϕ。会计算潜空间的VSD损失,同样用于更新生成网络的编码器和生成器,同时会计算Difff损失,微调模型。
- 除此之外,由于SD是文本控制生成模型,因此也加入了文本embeddings,更好地去指导模型生成。文中采用了SeeSR中的DAPE模块从LQ中提取文本embeddings。
训练时,直接以LQ数据为输入,而非随机高斯噪声。采用了LoRA对EθE_{\theta}Eθ、ϵθ\epsilon_{\theta}ϵθ和ϵϕ′\epsilon_{\phi'}ϵϕ′做微调。训练完成后,仅需推理GθG_{\theta}Gθ即可得到SR结果,且只做一步Diff。
整体框架和流程还是比较清晰,不过我们来进一步探究和分析其中的关键内容。不熟悉的小伙伴肯定有很多疑问。
2.1 何为单步Diff
回想一下扩散模型的基本过程,即通过zt=αtz+βtϵz_t = \alpha_t z + \beta_t \epsilonzt=αtz+βtϵ的加噪方式来扩散输入的潜特征zzz。α\alphaα和β\betaβ是依赖于扩散时间步长ttt的标量。我们可以通过神经网络来预测每一步ztz_tzt的噪声ϵ^\hat{\epsilon}ϵ^(这就是逆扩散过程),经过多步逆扩散后,即可得到去噪后的干净潜特征,表征为z^0=zt−βtϵ^αt\hat{z}_0 = \frac{z_t - \beta_t \hat{\epsilon}}{\alpha_t}z^0=αtzt−βtϵ^。如果把文本embeddings(记作cyc_ycy)加进来,则噪声预测可表示为ϵ^=ϵθ(zt;t,cy)\hat{\epsilon} = \epsilon_{\theta}(z_t; t, c_y)ϵ^=ϵθ(zt;t,cy)。
那么我们直接将上述去噪过程应用到SR任务中,也就是将LQ到HQ的潜特征变换表征为基于本文控制的图到图(image to image)去噪过程,
zH^=Fθ(zL;cy)≜zL−βTϵθ(zL;T,cy)αT\hat{z_H} = F_{\theta}(z_L; c_y) \triangleq \frac{z_L - \beta_T \epsilon_{\theta}(z_L; T, c_y)}{\alpha_T}zH^=Fθ(zL;cy)≜αTzL−βTϵθ(zL;T,cy)
在扩散时间步长T时刻,用LQ的潜特征zLz_LzL来代替ztz_tzt,且只对zLz_LzL做一步去噪(虽描述为去噪,其实可理解为复原),不引入任何噪声。以期望去噪后输出潜特征比输入潜特征更真实。cyc_ycy可由文本提示提取器获取,记为cy=Y(xL)c_y = Y(x_L)cy=Y(xL)。最后整合LQ-to-HQ的过程可写成:
x^H=Gθ(xL)≜Dθ(Fθ(Eθ(xL);Y(xL)))\hat{x}_H = G_{\theta}(x_L) \triangleq D_{\theta}(F_{\theta}(E_{\theta}(x_L); Y(x_L)))x^H=Gθ(xL)≜Dθ(Fθ(Eθ(xL);Y(xL)))
这就是该论文的单步Diff过程。
简单来说,该论文的单步Diff,跳过了完整的扩散过程,直接以LQ作为输入,学习LQ到HQ数据的单步映射关系。然后特殊的训练方式把这一步映射做好,以达到多步Diff的效果。(当然也有用纯noise一步生成结果的方式)
本来从噪声(起点)到结果(终点),可能要走100步。现在空降到中间的某个节点,然后迈一大步,直接到了终点或终点附近。
2.2 损失函数是关键
那怎么才能把“一步映射”训练出类似多步的效果呢?得靠损失函数,该论文同时使用了数据项损失和正则化损失,前者就是MSE损失和LPIPS损失,后者则采用了VSD,也是文章的重点改进之一。数据项损失如下,比较好理解,像素级约束+LPIPS感知型约束。
Ldata(Gθ(xL),xH)=LMSE(Gθ(xL),xH)+λ1LLPIPS(Gθ(xL),xH)\mathcal{L}_{\text{data}} \left( G_{\theta}(x_L), x_H \right) = \mathcal{L}_{\text{MSE}} \left( G_{\theta}(x_L), x_H \right) + \lambda_1 \mathcal{L}_{\text{LPIPS}} \left( G_{\theta}(x_L), x_H \right)Ldata(Gθ(xL),xH)=LMSE(Gθ(xL),xH)+λ1LLPIPS(Gθ(xL),xH)
怎么理解VSD?
论文里对VSD的描述较为简洁,未过多解释VSD的原理,可以看本文末尾的几篇引文加以了解。VSD其实是针对分数蒸馏采样(Score Distillation Sampling,SDS)技术的改进,这里基于个人理解简述一下,SDS的过程主要包括两步:
- 对生成模型(可理解为学生模型)的输出做正向扩散过程,也就是随机选择某个时刻t,加噪声。
- 然后用一个预训练T2I模型(可理解为教师模型)去预测扩散噪声,与所加噪声计算一个分数,作为正则损失去优化学生模型
但是SDS存在过度饱和,过度平滑和多样性差的问题。原因在于学生模型会过度依赖单个教师模型,仅学到了局部知识, 出现模式坍塌现象(数据分布较窄)。于是,“Prolificdreamer”这篇文章中提出了VSD,即再增加一个动态教师模型(可训练),相当于辅助或补充,旨在让学生模型学到更全面的知识。 实则通过计算动态教师模型与固定教师模型的评分差异,来提高正则化约束能力,扩大学生模型的学习范围。
直接看一下VSD的梯度公式
Lreg(Gθ(xL))=LVSD(Gθ(xL),cy)=LVSD(Gθ(xL),Y(xL))\mathcal{L}_{\text{reg}} \left( G_{\theta}(x_L) \right) = \mathcal{L}_{\text{VSD}} \left( G_{\theta}(x_L), c_y \right) = \mathcal{L}_{\text{VSD}} \left( G_{\theta}(x_L), Y(x_L) \right)Lreg(Gθ(xL))=LVSD(Gθ(xL),cy)=LVSD(Gθ(xL),Y(xL))
∇xLVSD(x,cy)=Et,ϵ[ω(t)(ϵϕ(zt;t,cy)−ϵϕ′(zt,t;cy))∂z∂x]\nabla_{\boldsymbol{x}} \mathcal{L}_{\text{VSD}} \left( \boldsymbol{x}, c_y \right) = \mathbb{E}_{t, \epsilon} \left[ \omega(t) \left( \boldsymbol{\epsilon}_{\phi}(\boldsymbol{z}_t; t, c_y) - \boldsymbol{\epsilon}_{\phi'}(\boldsymbol{z}_t, t; c_y) \right) \frac{\partial \boldsymbol{z}}{\partial \boldsymbol{x}} \right]∇xLVSD(x,cy)=Et,ϵ[ω(t)(ϵϕ(zt;t,cy)−ϵϕ′(zt,t;cy))∂x∂z]
其中梯度的期望计算包括了所有的扩散时间步t∈{1,⋯,T}t \in \{1, \cdots, T\}t∈{1,⋯,T}和ϵ∼N(0,I)\epsilon \sim \mathcal{N}(0, \mathbf{I})ϵ∼N(0,I)。
整体的损失约束
针对生成模型的总损失如下:
L(Gθ(xL),xH)=Ldata(Gθ(xL),xH)+λ2Lreg(Gθ(xL))\mathcal{L} \left( G_{\theta}(x_L), x_H \right) = \mathcal{L}_{\text{data}} \left( G_{\theta}(x_L), x_H \right) + \lambda_2 \mathcal{L}_{\text{reg}} \left( G_{\theta}(x_L) \right)L(Gθ(xL),xH)=Ldata(Gθ(xL),xH)+λ2Lreg(Gθ(xL))
可微调教师模型ϵϕ′\epsilon_{\phi'}ϵϕ′的损失如下:
Ldiff=Et,ϵ,cy=Y(xL),z^H=Fθ(Eθ(xL);Y(xL))LMSE(ϵϕ′(αtz^H+βtϵ;t,cy),ϵ)\mathcal{L}_{\text{diff}} = \mathbb{E}_{t, \epsilon, c_y=Y(x_L), \hat{z}_H=F_{\theta}(E_{\theta}(x_L); Y(x_L))} \mathcal{L}_{\text{MSE}} \left( \epsilon_{\phi'} \left( \alpha_t \hat{z}_H + \beta_t \epsilon; t, c_y \right), \epsilon \right)Ldiff=Et,ϵ,cy=Y(xL),z^H=Fθ(Eθ(xL);Y(xL))LMSE(ϵϕ′(αtz^H+βtϵ;t,cy),ϵ)
从算法流程上看,两者交替更新,先更新生成模型GθG_{\theta}Gθ,后更新可微调教师模型ϵϕ′\epsilon_{\phi'}ϵϕ′。
2.3 其他巧思
为何冻结VAE解码器 && 为何VSD要在潜空间里计算
主要是作者为了简化过程和减少训练计算成本。在潜空间中去应用VSD,能够避免反复解码/编码。因此,需要将解码器固定,这样在反向传播时,不会因为解码器的参数变化而影响到编码器和生成模型,从而保证了潜空间特征一致性。当然,这波操作做了消融实验,效果确实有提升。
三、实验论证
3.1 定量与定性
使用完整的LSDIR和FFHQ中前10K人脸图像作为训练数据集,从开源代码中可以看到作者做了两个预训练模型,场景和人脸。
测试集包括合成数据和real-world数据,前者取自DIV2K-Val(基于Real-ESRGAN做退化),后者来自RealSR和DRealSR的LQ-HQ对。
分别跟DM-based类的方法和GAN-based类的方法做了不同的对比。评价方式采用了全参考和无参考度量,以及用户真实评价。
- PSNR,SSIM在Y通道上评价保真度
- LPIPS,DISTS衡量感知质量
- FID评估GT和SR结果的分布距离
- NIQE,MANIQA-pipal,MUSIQ,CLIPIQA均为无参考指标。
table1 展示了DM-based类对比结果,OSEDiff在real-world数据集上有明显优势,特别是在全参考感知质量度量指标LPIPS和DISTS,分布对齐度量指标FID和语义质量度量指标CLIPIQA上较为突出。跟同为单步Diff的SinSR相比在保真度和感知质量方面均有明显提升。
虽然多步Diff方案SeeSR和PASD在无参考指标上表现更佳,会生成更丰富的细节,但从视觉效果上看,细节并不一定真实和自然。如下图所示。
因此OSEDiff在综合效果方面更能打,细节表征还不错又比较真实。
不过仔细看,也能发现OSEDiff在部分局部特征存在不自然的问题。
table8展示了GAN-based类对比结果,GAN-based方法在保真度指标上更优。而OSEDiff在无参考指标上优势明显。
视觉效果上,OSEDiff的纹理表现也更丰富更自然,如图。
用户评价,则是挑选了20张real-wrold数据进行处理,通过15名志愿者的主观选择,结果如图,OSEDiff的选择占比排第二,稍落后于SeeSR,但是推理效率比SeeSR要快数十倍,因此OSEDiff的性价比更高。
但看下来感觉OSEDiff最大的区别之一在于对比度提升明显,其他模型基本不改变原图对比度,这一点可能会影响主观评价结果。
table2展示训练参数量和推理效率的对比,优势显著,要训练的参数仅8.5M,512x512@A100上的推理速度仅需0.11s,比StableSR快100倍。
3.2 消融实验
VSD Loss的价值——
从实验结果上看,VSD在无参考指标上有优势,也就是有助于提升感知质量。且放在潜空间约束比在像素空间约束会更有效。在像素空间下,VSD和GAN的效果就比较类似了。
不同Text Prompt Extractors的影响——
不用文本提升,全参考指标更好。加入文本提示,无参考指标有改善。分析是文本提示能更好地激发T2I模型的生成能力,产生更丰富的合成细节,但会降低全参考索引。同时用更复杂的文本提示提取器并没有带来较大效果提升,因此选用DAPE即可,性价比更高。
LoRA Rank大小的影响——
低秩会导致训练不稳定,无法收敛;高秩会容易过拟合,使得输出的细节有所损失。测试发现,rank为4能够达到较好的平衡。
微调VAE Encoder和Decoder的影响——
OSEDiff的策略是仅调VAE编码器,固定VAE解码器。
- 与都不微调相比,MUSIQ大幅提升,说明微调编码器,有助于消除退化;
- 与都微调相比,CLIPIQ大幅提升,说明固定解码器,有助于保持潜空间的稳定性,从而获得更好的感知质量。
四、总结与思考
总结一下,OSEDiff是一种面向真实世界图像超分辨率(Real-ISR)的单步扩散网络。其核心创新包括:
- 无需随机噪声:以LQ图像的潜码为起点,避免传统扩散模型依赖随机噪声导致的输出不确定性。
- 轻量化适配:仅微调预训练SD模型的LoRA层(8.5M参数),快速适配复杂退化模式。
- 潜空间正则化:通过VSD损失,促使输出分布与自然图像先验做对齐,使单步推理达到多步扩散效果。
结论部分也提到了局限性,1)细节生成能力仍有待进一步提升;2)对文字、建筑线条等精细几何特征(小尺度结构)的重建能力有限。这也是单步Diff技术在SR领域中需要研究和改进的问题。
据Github上描述,以落地应用于OPPO FindX8上,猜测应该用在长焦的数字变倍中吧。感兴趣的可以去测测。
最后,附上VSD相关的参考文献,当然不仅限于这几篇。
- SwiftBrush One-Step Text-to-Image Diffusion Model with Variational Score Distillation
- Prolificdreamer: High-fidelity and diverse text-to-3d generation with variational score distillation
- One-step Diffusion with Distribution Matching Distillation
感谢阅读,欢迎留言或私信,一起探讨和交流。
如果对你有帮助的话,也希望可以给博主点一个关注,感谢。