论文：https://arxiv.org/pdf/2406.08177

代码：https://github.com/cswry/OSEDiff

1 背景

1.1 问题

训练一个Real - ISR模型有两个主要问题。

• 一个是如何构建LQ - HQ训练图像对

• 另一个是如何保证复原图像的自然性，即如何保证复原图像遵循HQ自然图像的分布。

对于一个问题，通常有长短焦距镜头采集法和高阶退化模型等。

在给定训练数据的情况下，如何训练一个鲁棒的Real-ISR模型以输出高质量的自然图像成为一个关键问题。简单地学习具有逐像素损失的LQ-HQ配对数据之间的映射会导致过于平滑的结果，将自然图像先验融入到HQ图像的学习过程中是至关重要的。为此，提出了几种方法。在 VGG - 16 和 AlexNet 等预训练模型中，感知损失探索了纹理、颜色和结构先验。生成式对抗网络( GANs )交替训练生成器和判别器，并被用于Real - ISR任务。生成器网络旨在合成HQ图像，判别器网络旨在区分合成图像是否真实。虽然基于GAN的Real - ISR取得了巨大的成功，特别是在恢复特定类别的图像(如人脸图像)时，由于对抗训练不稳定和难以区分的图像空间，往往会产生不好的细节。

基于扩散模型的超分效果很好，虽然LQ图像被用作Controlnet网络的控制信号，但这些方法在输出的HQ图像中引入了不必要的随机性。其次，复原的HQ图像通常由数十甚至上百个扩散步骤获得，导致Real - ISR过程计算量大。尽管最近提出了一些基于一步扩散的Real - ISR方法，但与多步方法相比，它们无法获得高质量的细节。

1.2 提出的方法

为了解决上述问题，作者针对Real - ISR问题提出了一种一步有效扩散网络( One-step Effective Diffusion Network，简称OSEDiff )。预训练的SD模型中的UNet主干具有较强的将输入数据转移到另一个域的能力，而给定的LQ图像实际上具有丰富的信息来恢复其HQ对应物。因此，作者提出在不引入任何随机噪声的情况下，直接将LQ图像输入到预训练的SD模型中。同时，作者整合了可训练LoRA层加入到预训练的UNet中，并微调SD模型使其适应于Real - ISR任务。另一方面，为了保证一步法模型仍能产生与多步法模型相同的HQ自然图像，我们利用变分得分蒸馏( VSD ) 进行KL散度正则化。该操作有效地利用了预训练SD模型强大的自然图像先验，并将生成图像的分布与自然图像先验对齐。如图1所示，我们的大量实验表明，OSEDiff取得了与最先进的基于SD的Real - ISR模型相当或更优越的性能指标，同时它显著地将推理步数从N减少到1，并且具有最少的可训练参数，导致推理速度是之前方法（如StableSR）的100倍以上。

2 创新点

提出了针对Real-ISR问题的一步有效扩散网络。

3 方法

Real-ISR是从给定的LQ图像 $x_L$ 中估计一幅 HQ 图像 ${\hat{x}}_H$，该任务可以建模为一个优化问题：${\hat{x}}_H=argmi{n}_{x_H}(L_{data}(\Phi (x_H),x_L)+\lambda L_{reg}(x_H))$ ，其中 $\Phi$ 为衰退函数，$L_{data}$ 是衡量优化输出的保真度的数据项，$L_{reg}$ 是利用自然图像的先验信息的正则化项，$\lambda$ 是平衡参数。许多传统的ISR方法通过假设简单和已知的退化模型，并采用手工设计的自然图像先验模型来回复所需的HQ图像。

然而这种基于优化的方法的性能在很大程度上受到两个因素的阻碍。首先，退化函数 $\Phi$ 往往是未知的，在现实场景中很难建模，其次，手工构造的正则化项 $L_{reg}$ 很难有效的对自然的复杂图像先验进行建模。随着深度学习技术的发展，从$(x_L,x_H)$ 对的训练数据集 $S$ 中学习一个 $\theta$ 参数化的神经网络 $G_{\theta }$ 来将LQ图像映射到HQ图像已经变得非常普遍，网络训练可以表述为如下学习问题：

其中 $L_{data}$ 和 $L_{reg}$ 是损失函数。$L_{data}$ 让网络输出 ${\hat{x}}_H=G_{\theta }(x_L)$ 可尽可能接近真值 $x_H$，可以通过 $L_1$ 范数，$L_2$ 范数和 $LPIPS$ 等指标进行量化。只用 $L_{data}$ 损失从头训练网络可能会导致对训练数据的过拟合。在这项工作中，作者提出微调一个预训练的生成网络，更具体的说是SD网络，以提高 $G_{\theta }$ 的泛化能力。此外，正则化损失 $L_{reg}$ 对于提高 $G_{\theta }$ 的泛化能力以及增强输出HQ图像 ${\hat{x}}_H$ 至关重要。假设我们有真实世界HQ的分布 $p(x_H)$， KL散度是作为 $L_{reg}$ 损失函数的理想选择，也就是说，恢复的HQ图像的分布 $q_{\theta }({\hat{x}}_H)$ 应尽可能与 $p(x_H)$ 相同。正则化损失可以定义为：

现有的工作大多通过对抗训练来实例化上述目标。然而判别器通常是和生成器一起从头训练的，他们无法获取全分布的HQ图像，且缺乏足够的判别能力，从而导致次优的Real-ISR性能。最近发展起来的 T2I 扩散模型，如SD为 $L_{reg}$ 提供了新的选择。这些模型在数十亿的图像-文本对上训练，能够有效刻画图像在潜在空间中的自然分布。一些得分蒸馏方法被报道使用SD以KL散度为目标来优化。特别的，变分得分蒸馏VSD从基于例子的变分优化中诱导出这样一个基于KL散度的目标来对其两个扩散模型所表示的分布。基于以上讨论，作者提出将式（1）中的学习目标实例化，设计了一个高效且有效的一步扩散网络。具体来说，使用LoRA作为Real-ISR主干网络来微调预训练的SD，并使用VSD作为正则化器来对其网络输出与自然HQ图像的分布。

4 模块

4.1 整体框架

现有的基于SD的Real-ISR方法以随机噪声为起点，LQ图像为控制信号，执行多个时间步来估计HQ图像，这些方法是资源密集型的，并且会固定的引入随机性。本文提出了一个用于Real-ISR的一步有效扩散网络OSEDiff，其训练网络如图2所示。要训练的生成器 $G_{\theta }$ 由一个可训练的编码器 $E_{\theta }$，一个微调的扩散网络 ${ϵ}_{\theta }$ 和一个冻结的解码器 $D_{\theta }$ 组成。为保证 $G_{\theta }$ 的泛化能力，将扩散网络 ${ϵ}_{\theta }$ 的输出送入两个正则化网络，其中VSD损失在隐空间执行，正则化损失反向传播以更新 $E_{\theta }$ 和 ${ϵ}_{\theta }$。一旦训练完成，只需要生成器 $G_{\theta }$ 进行推理。

4.2 网络结构设计

用 $E_{\theta }$，${ϵ}_{\theta }$ 和 $D_{\theta }$ 分别表示预训练SD模型的VAE编码器，隐扩散网络和VAE解码器，其中 $\Phi$ 表示模型参数。作者采用LoRA对预训练的SD进行微调，以获得期望的生成器 $G_{\theta }$。

如图2左图所示，为了保持SD原有的生成能力，作者在编码器 $E_{\theta }$ 和扩散网络 ${ϵ}_{\theta }$ 中引入可训练的LoRA层，并利用训练数据对齐微调。对于解码器，冻结其参数，直接设置 $D_{\theta }=D_{\Phi }$。这是为了保证扩散网络的输出空间与正则化器一致。

回归一下，扩散模型通过 $z_t={\alpha }_{t}z+{\beta }_tϵ$ 对输入的潜在特征 $z$ 进行扩散，其中 ${\alpha }_t,{\beta }_t$ 是依赖于扩散时间步长 t ∈ { 1 , … , T } t∈\{1,…,T\} t∈{ 1,…,T}的标量。通过一个可以预测 $z_t$ 中噪声的神经网络，记为 $\widehat{ϵ}$，可以得到去噪后的隐层为 ${\hat{z}}_0=\frac{z_t-{\beta }_t\widehat{ϵ}}{{\alpha }_t}$，比 $z_t$ 更干净逼真。此外，SD是一种基于文本条件的生成模型，通过从给定的文本描述 $y$ 中提取文本嵌入 $c_y$，预测噪声可以表示为 $\widehat{ϵ}={ϵ}_{\theta }(z_t;t,c_y)$。

作者将上述文本到图像的去噪过程改写成Real-ISR任务，并将LQ到HQ的应变换 $F_{\theta }$ 表示为文本条件下的图像到图像去噪过程：

其中，只对LQ隐变量 $z_L$ 进行一步去噪，而不引入任何噪声。期望去噪输出 ${\hat{z}}_H$ 比 $z_L$ 更逼真。对于文本嵌入，作者将一些文本提示提取器，如DAPE，应用到LQ输入 $x_L$ 中，得到 $c_y=Y(x_L)$。最后，整个LQ-to-HQ过程可以写为：

为了提高Real-ISR模型的性能，需要同时使用数据项 $L_{data}$ 和正则化项 $L_{reg}$ 来监督生成器的训练，如图2右半部分所示。作者采用VSD作为正则项。除了使用SD模型作为预训练的正则化器 ${ϵ}_{\varphi }$ 外，VSD还引入了一个微调的正则化器，即在LoRA生成图像的分布 $q_{\theta }({\hat{x}}_H)$ 上微调潜在扩散模块 ${ϵ}_{{\varphi }^{\prime }}$ 。

${ϵ}_{\theta }$ 去噪完再加噪让 ${ϵ}_{\varphi }$ 去噪。

4.3 损失函数

作者使用数据损失 $L_{data}$ 和正则化损失 $L_{reg}$ 来训练生成器 $G_{\theta }$。将 $L_{data}$ 设置为MSE损失和LPIPS损失的加权和：

其中 ${\lambda }_1$ 是一个加权标量。对于 $L_{reg}$，使用VSD损失：

给定任意可训练的图像形状特征 $x$，其潜在编码 $z=E_{\varphi }(x)$ 和编码后的文本提示条件 $c_y$，VSD对 $x$ 进行优化，使其与文本提示 $y$ 一致：

其中梯度的期望是在所有的扩散时间 t ∈ { 1 , … , T } t∈\{1,…,T\} t∈{ 1,…,T} 上进行的，且 $ϵN(0,1)$。因此生成器 $G_{\theta }$ 的总体训练目标是：

式中 ${\lambda }_2$ 为加权标量。此外，根据VSD的要求，微调的正则化器 ${ϵ}_{{\varphi }^{\prime }}$ 也是可训练的，其训练目标为：

原始的VSD计算图像空间中的梯度。当它用于训练一个基于SD的生成器网络时，在计算 $L_{reg}$ 时会出现重复的潜在解码/编码。这是很昂贵的，并且使得正则化的效果较差。考虑到一个训练好的VAE应该满足 $E_{\varphi }(x)=E_{\varphi }(D_{\varphi }(z))\approx z$，我们可以近似地令 $E_{\varphi }({\hat{x}}_H)={\hat{z}}_H$ .在这种情况下，我们可以在计算正则化损失时消除冗余的潜在编码/解码，正如我们遵循DMD [ 58 ]来优化潜在状态空间而不是噪声空间中的分布损失。正则化损失网络参数θ在隐空间中的梯度为：

5 效果

5.1 和SOTA方法对比

和SOTA方法的指标对比。

和SOTA方法的耗时对比。

和SOTA方法的可视化对比。从这张图奥巴马头部右边的文字可以看出，对字母会有修改。

5.2 消融实验

不同损失的消融实验。

不同文本提示提取器的消融实验。

VAE中LoRA不同秩的消融实验。

Unet中LoRA不同秩的消融实验。

是否微调encoder和decoder的消融实验。