【论文精读】StableSR

摘要

       将Diffusion先验嵌入到合成模型（如Stable Diffusion）的模式在图像视频编辑领域取得了良好的结果。本文提出StableSR，将Diffusion先验嵌入到超分辨率（SR），且不对图像退化模式做明确假设。具体有：

• 为了利用Diffusion先验，所提出方法只需为SR任务微调了一个轻量级的时间感知编码器和一些特征调整层。时间感知编码器纳入了一个时间嵌入层来生成时间感知特征，因为迭代过程中冻结了原始Diffusion模型以保留Diffusion先验，故通过时间感知特征在图像恢复过程中为每个Diffusion步骤提供自适应指导，使Diffusion特征在不同的迭代轮次中自适应调整。时间感知编码器在早期迭代步骤提供更强的指导，在后期提供更弱的指导，以保持SR保真度
• 为了抑制Diffusion模型的生成随机性以及自编码器编码过程中产生的信息损失，在自编码器解码过程中应用了一个系数可调的可控特征包装模块（CFW），来细化Diffusion模型的输出
• 为了使模型适应任意分辨率图像，引入一种渐进式聚合采样策略。首先将图像划分为重叠的块，并在每个Diffusion迭代过程都使用高斯核对这些块进行融合。这个过程可以平滑图像块边界，以产生了更连贯的输出
  

       上图在实际场景下对BSRGAN、Real-ESRGAN+、FeMaSR、LDM和StableSR进行了定性比较。

框架

Guided Finetuning with Time Awareness

       为了使SR能利用Stable Diffusion的先验知识建立以下约束：

• 模型必须以低分辨率输入图像为条件，生成高分辨率图像
• 模型只对原始Stable Diffusion进行最小程度的修改，以防止破坏封装在其中的知识

Feature Modulation

       采用一个额外的编码器从退化的LR图像中提取多尺度特征 { F n } n = 1 N \{ F^n\}^N_{n=1} {Fn}n=1N​，并结合空间特征变换（SFT）调制Stable Diffusion残差块的中间特征映射：
$${\hat{F}}_{dif}^n=(1+{\alpha }^n)\odot F_{dif}^n+{\beta }^n;{\alpha }^n,{\beta }^n=M_{\theta }^n(F^n)$$

       其中，$M_{\theta }^n$表示由多个卷积层组成的小网络，${\alpha }^n$和${\beta }^n$表示SFT中的仿射参数，$n$索引Stable Diffusion的UNet空间尺度。微调过程中，会冻结Stable Diffusion的权重，只训练编码器和SFT层，该策略允许向Stable Diffusion中插入从LR图像中提取的结构信息，而不破坏Stable Diffusion捕获的生成先验。

Time-aware Guidance

       实验发现，在编码器中使用一个时间嵌入层合并时间信息，可以大大提高了生成的图像质量和保真度，这可以使Stable Diffusion在迭代过程中自适应地调整来自LR特征的条件强度。

       因为在生成过程中，随着噪声的逐渐去除，生成图像的信噪比（SNR）会逐渐增加。当SNR接近$5e^{-2}$时，图像内容会迅速填充，生成图像的SNR在$5e^{-2}$的范围内时，编码器会为Diffusion提供相对强的条件。如上图，利用SFT前后的Stable Diffusion特征之间的余弦相似性来衡量编码器提供的条件强度，并绘制不同时间步长的特征余弦相似度值，观察到，余弦相似度在$5e^{-2}$的信噪比附近达到最小值，这表明编码器施加了最强条件，验证了上述结论。上图b为从编码器中提取的特征图。

Color Correction

       Diffusion模型偶尔会出现颜色偏移，故需对生成的图像进行颜色归一化，以使其均值和方差与LR图像的均值和方差对齐。令$x$表示LR图像，$\hat{y}$表示生成的HR图像，则颜色校正后的输出$y$计算为：
$$y^c=\frac{{\hat{y}}^c-{\mu }_{\hat{y}}^c}{{\sigma }_{\hat{y}}^c}\cdot {\sigma }_x^c+{\mu }_x^c$$

       其中 c ∈ { r , g , b } c ∈ \{ r, g, b \} c∈{r,g,b}表示颜色通道，${\mu }_{\hat{y}}^c$和${\sigma }_{\hat{y}}^c$（或${\mu }_x^c$和${\sigma }_x^c$ ）分别为从$\hat{y}$（或$x$）的$c-th$通道估计的均值和方差，实验发现这种简单的校正足以弥补颜色差异。

       除了在像素域进行颜色校正外，还提出了基于小波的颜色校正方法，给定图像$I$，通过小波分解在$i-th(1\le i\le l)$尺度提取其高频分量$H^i$和低频分量$L^i$：
$$L^i=C_i(L^{i-1},k),H^i=L^{i-1}-L^i$$

       其中$L_0=I$, $C_i$表示dilation rate为$2^i$的空洞卷积算子，$k$为卷积核，定义为:
k = [ 1 16 1 8 1 16 1 8 1 4 1 8 1 16 1 8 1 16 ] k=

$$\begin{bmatrix} \frac 1 {16} & \frac 1 {8} & \frac 1 {16} \\ \frac 1 {8} & \frac 1 {4} & \frac 1 {8} \\ \frac 1 {16} & \frac 1 {8} & \frac 1 {16} \end{bmatrix}$$ k= ​161​81​161​​81​41​81​​161​81​161​​ ​

       将$x$（或$\hat{y}$）的$l-th$的低频和高频成分表示为$L_x^l$和$H_x^l$（或$L_{\hat{y}}^l$和$H_{\hat{y}}^l$），则输出$y$表示如下:
$$y=H_{\hat{y}}^l+L_x^l$$

       通过将$\hat{y}$的低频分量$L_{\hat{y}}^l$替换为$L_x^l$来纠正颜色偏差。本文默认采用像素域的颜色校正。

Fidelity-Realism Trade-off

       由于Diffusion模型具有随机性，其生成图像可能偏离真实情况，故引入一个可控特征包装（CFW）模块，以平衡生成图像的真实感和保真度。

       因为Stable Diffusion利用了编码器的潜空间进行生成，因此可以通过自编码器的编码器特征来调制相应的解码器特征，以进一步提高保真度。令$F_e$和$F_d$分别表示编码器和解码器特征，引入可调系数$w\in [0,1]$来控制调制权重：
$$F_m=F_d+C(F_e,F_d;\theta )\times w$$

       其中$C(\cdot ;\theta )$表示具有可训练参数$\theta$的卷积层，模型整体结果如上图。

       这种设计下，采用较小的$w$可以更大权重利用Stable Diffusion的生成能力，可以增强输出的真实感；而较大的$w$允许从LR图像中获得更强的结构指导，可以增强输出的保真度。实验观察到$w=0.5$时可以在真实感和保真度之间取得良好的平衡。（CFW只在这个特定阶段训练）

Aggregation Sampling

       Stable Diffusion的注意力层对图像分辨率有高度敏感性，对于与其训练设置不同分辨率的图像，其会产生劣质输出，对该问题，常见的解决方案是将较大的图像分割为几个重叠的较小的patch，并分别进行处理。但patchs之间的差异会在Diffusion迭代过程中被复合和放大，如上图左中对比。

       本文采用一种渐进式块聚合采样算法来处理任意分辨率的图像。具体为，首先将低分辨率图像通过编码器编码为潜特征图$F\in R^{h\times w}$，然后将其划分为$M$个重叠的patch { F Ω n } n = 1 M \{F_{Ω_n} \}^M_{ n=1} {FΩn​​}n=1M​，每个patch的分辨率为$64\times 64$，${\Omega }_n$是特征图$F$的第$n$个patch的坐标集。reverse sampling过程的每个step中，会通过StableSR处理每个patch，然后对处理后的patch进行合并。为了合并处理后的patch，对每个patch $F_{{\Omega }_n}$生成一个权重图$w_{{\Omega }_n}\in R^{h\times w}$，该步骤会将patch从$64\times 64$扩展为$h\times w$的分辨率。$w_{{\Omega }_n}$中，${\Omega }_n$区域为特征图经过高斯滤波器得到的值，${\Omega }_n$之外的区域由函数$f(\cdot )$填充0。随后，patchs根据各自的高斯权重图进行加权相加得到最终的合并输出。

       每个patch的输出定义为${ϵ}_{\theta }(Z_{{\Omega }_n}^{(t)},F_{{\Omega }_n},t)$，其中$Z_{{\Omega }_n}^{(t)}$为噪声$Z^{(t)}$的第$n$个patch，$\theta$是Diffusion模型的参数，所有patch合并的结果为:
$${ϵ}_{\theta }(Z^{(t)},F,t)=\sum _{n=1}^M\frac{w{\Omega }_n}{\hat{w}}\odot f({ϵ}_{\theta }(Z_{{\Omega }_n}^{(t)},F_{{\Omega }_n},t))$$

       其中$\hat{w}={\sum }_n{w}_{{\Omega }_n}$。采用$Sampler(Z^{(t)},{ϵ}_{\theta }(Z^{(t)},F,t))$采样 $Z^{(t-1)}$，随后将 $Z^{(t-1)}$重新分割为重叠的patch，并重复上述步骤，直到$t=1$。整体流程如下：

       实验表明，这种渐进式聚合方法可以大大提高输出质量，如图4右。

实验

Implementation Details

       StableSR基于Stable Diffusion 2.1-base构建，所提出的时间感知编码器类似于Stable Diffusion中的denoising U-Net的收缩路径，但更轻量（∼105M，包括SFT层）。

       首先对Diffusion模型微调117个epoch，batch size为192，prompt固定为null，并采用学习率为$5\times 10^{-5}$的Adam优化器优化模型。训练过程中使用了8个NVIDIA Tesla 32G-V100 gpu。

       对于CFW的训练，首先按照Real-ESRGAN中的degradation pipeline生成100k个分辨率为$512\times 512$的合成LR-HR对，并以LR为条件，采用微调后的Diffusion模型生成相应的潜代码$Z_0$来训练CFW。训练损失与LDM中使用的自编码器损失相同，且使用固定的对抗性损失权重0.025。

       推理采用200个step的DDPM采样。为了处理任意大小的图像，对分辨率大于$512\times 512$的图像采用了所提出的聚合采样策略。对分辨率小于$512\times 512$的图像，首先放大LR图像，使较短的边具有512的长度，并在生成后将结果重新缩放回目标分辨率（推理前，输入首先被上采样到与输出相同的大小）。

Experimental Settings

Training Datasets

       训练数据采用Real-ESRGAN的degradation pipeline在DIV2K、DIV8K、Flickr2K、OutdoorSceneTraining和FFHQ数据集的5000张人脸图像数据集上合成LR/HR对。

Testing Datasets

       本文在合成数据集和真实世界数据集上评估了所提出方法。合成数据集遵循Real-ESRGAN的degradation pipeline，从DIV2K Valid中生成了3k的 LR-HR对。真实世界数据集采用了RealSR、DRealSR、DPED-iPhone及从互联网上收集的40张图像。

Compared Methods

       为了验证所提方法的有效性，将StableSR与RealSR 、BSRGAN、Real ESRGAN+、DASR、FeMaSR、LDM、SwinIR-GAN 、DeepFloyd IF_III进行了比较。由于LDM是在分辨率为$256\times 256$的图像上训练的，故按照StableSR的相同训练设置对其进行微调，其他方法直接使用官方代码和模型进行测试。

Evaluation Metrics

       对DIV2K Valid、RealSR和DRealSR，采用LPIPS6、FID、CLIP-IQA、MUSIQ、PSNR和SSIM评分等感知指标，评估生成图像的感知质量。

Comparison with Existing Methods

Quantitative Comparisons

       定量对比结果如上表。观察到，stableSR在多个感知指标方面优于最先进的SR方法，包括FID、CLIP-IQA和MUSIQ。在合成数据集基准DIV2K Valid上，StableSR（$w=0.5$）实现了24.44 FID分数，比LDM低7.7%，比其他基于GAN的方法至少低32.9%。 StableSR（$w=0.5$）在两个常用的现实世界数据集基准上也取得了最高的CLIP-IQA分数，表明StableSR的优越性。

Qualitative Comparisons

       上图展示了stableSR重建的真实世界数据集基准和互联网的真实世界图像的视觉结果。实验结果表明，StableSR在伪影去除和细节生成方面都优于之前的方法。

       上图进一步证明了StableSR在分辨率大于$512\times 512$的图像上的优越性。

User Study

       为了进一步检验StableSR的有效性，本实验从互联网上收集了16幅真实的LR图像，采用StableSR与BSRGAN, Real-ESRGAN+和LDM分别进行图像重建，并挑选了25个人类受试者对$3\times 16$的生成图像对进行评估，总共投票$3\times 400$次。

       结果如上图，StableSR以很大的优势超过了所有三种竞争方法，始终获得超过75％的投票。

Comparison with Concurrent Diffusion Applications

       本文在真实世界数据集的图像上进一步与ControlNet-tile、IF_III upscaler 进行了比较，所有方法都使用eta 1.0和timestep 200的DDIM采样，随机数种子固定为42。 在StableSR中进一步设置了$w=0.0$，以避免由于CFW带来的额外改进。对于ControlNet-tile，使用stable-Diffusion-WebUI生成额外的prompt，以获得更好的性能。对于IF_III upscaler ，按照官方示例将噪音等级设置为100 w/o。

       结果如上图，由于没有针对SR的具体设计，ControlNet-tile保真度较差。与IF_III upscaler相比，StableSR能够生成边缘更清晰的细节。实验进一步验证了StableSR的优越性。

Ablation Study

Importance of Time-aware Guidance and Color Correction

       本节研究了时间感知引导和颜色校正在真实世界数据集基准上的有效性。

       消融结果如上表，删除时间感知引导或颜色校正都会导致较差的SSIM和LPIPS。

       上图可视化了在没有时间感知引导或颜色校正情况下的生成结果，整体显示出较差的生成性能，这表明时间感知引导和颜色校正的有效性。

       上图可视化了采用小波颜色校正的模型生成质量，显示出较好的生成质量，这可能进一步促进实际使用。

Flexibility of Fidelity-realism Trade-off

       CFW模块受CodeFormer的启发，允许灵活的现实-保真权衡。给定一个可控系数$w$，范围为$[0,1]$，小的w的CFW倾向于生成逼真的结果，而大的$w$的CFW提高了保真度。详细实验结果如上表，与StableSR（$w=0.0$）相比，具有较大的$w$值（例如，0.75）的StableSR在所有三个成对基准上取得了 更高的PSNR和SSIM，表明更好的保真度。相比之下，StableSR （$w=0.0$）通过更高的CLIP-IQA分数和MUSIQ分数实现了更好的感知质量。

       上图为相应的可视化结果，观察到适当的$w$可以促使保真度和感知质量的提高。StableSR（$w=0.5$）显示了与StableSR（$w=1.0$）相当的PSNR和SSIM，但同时具有更好的感知度量分数。因此，StableSR将系数$w$默认设置为0.5，以便在感知质量和保真度之间进行权衡。

Complexity Comparison

       StableSR是一种基于Diffusion的方法，需要多步采样来生成图像。如上表，当采样步数设置为200时，StableSR在一个NVIDIA Tesla 32G V100 GPU上需要15.16秒才能生成$512\times 512$的图像，这与IF_III upscaler相当，但比基于GAN的SR方法（如Real-ESRGAN+和SwinIR-GAN）慢。若需提高模型效率，可采用快速采样策略、模型蒸馏或缩短Diffusion过程链等方法。

Inference Strategies

       StableSR在合成和真实世界的数据集基准上展示了卓越的性能，本节讨论了采样过程中的几种有效策略，可以进一步提高生成结果的视觉质量，而无需额外的微调。

Classifier-free Guidance with Negative Prompts

       StableSR的prompt默认为null，然而通过实验观察到StableSR可以对prompt做出反应，特别是negative prompt，故本节研究prompt的使用，以进一步提高采样期间的生成图像质量。

       给定两个分别以null prompt${ϵ}_{\theta }(Z^{(t)},F,[],t)$和negative prompt${ϵ}_{\theta }(Z^{(t)},F,c,t)$为条件的StableSR模型，则新的抽样过程可以使用negative prompt的模型估计与guidance scale为s的线性组合：
$${\widetilde{ϵ}}_{\theta }={ϵ}_{\theta }(Z^{(t)},F,c,t)+s({ϵ}_{\theta }(Z^{(t)},F,[],t)-{ϵ}_{\theta }(Z^{(t)},F,c,t))$$

       其中$c$为negative prompt，$s=0$时相当于直接使用negative prompt，$s=1$相当使用null prompt。

       上表比较了StableSR在各种positive prompts和negative prompt下的性能，positive prompts有“(masterpiece:2), (best quality:2), (realistic:2), (very clear:2)”、“Good photo.“，negative prompt有“3d, cartoon, anime, sketches,(worst quality:2), (low quality:2)”。

       从上表观察到，negative prompt对性能指标有显著影响，更高的$s$会导致更高的CLIP-IQA和MUSIQ分数，表现出更清晰的结果。

       对应的可视化结果如上图，结果进一步验证了上述结论。另外观察到，过于强烈的引导（$s=7.5$），会导致生成图像过度锐化。

Pre-cleaning for Severe Degradations

       当LR图像严重降质、伴有明显的模糊或噪声时，StableSR可能产生次优的结果，如上图第一列。故在StableSR之前加入了一个辅助的预清洗阶段，以解决严重退化下的场景。

       首先采用通用SR的Real-ESRGAN+和用于人脸SR的CodeFormer初步生成，为了抑制预清洗阶段产生的伪影放大，预清洗后再采用2×双三次下采样操作，随后使用StableSR生成最终输出。如上图右，这种预清洗阶段大大提高了StableSR的鲁棒性。

Appendix

Details of Time-aware Encoder

       上表为时间感知编码器的详细设置。

User Study Settings

       用户研究采用谷歌表单作为研究平台，如上图。

Additional Visual Results

Visual Results on Fixed Resolution

       上图为真实世界数据集图像的$128\times 128$：LR图像 to $512\times 512$：HR图像的定性比较，观察到StableSR成功地产生了更精细细节和更清晰边缘的输出，显著优于之前最先进的方法。

Visual Results on Arbitrary Resolution

       上图显示了StableSR能够生成超过4倍分辨率的高质量SR图像，表明其在现实应用中的实际应用。

       上图显示，StableSR可以在建筑和文本等不同和复杂的真实场景下生成真实的纹理。

reference

Jianyi, W. , Zongsheng, Y. , Shangchen, Z. , Kelvin C.K. C. , & Chen C. L. . (2023). Exploiting Diffusion Prior for Real-World Image Super-Resolution.