论文阅读：One-Step Effective Diffusion Network for Real-World Image Super-Resolution

今天介绍一篇 2024 NeurIPS 的文章，是用 Diffusion 做超分的。

Abstract

由于具有强大的生成图像先验知识，预训练的文本转图像扩散模型越来越多地被用于解决现实世界中的图像超分辨率（Real - ISR）问题。现有的大多数方法都是从随机噪声出发，在给定的低质量（LQ）图像的引导下重建高质量（HQ）图像。尽管已经取得了不错的成果，但这类 Real - ISR 方法需要多个扩散步骤来重现高质量图像，从而增加了计算成本。与此同时，随机噪声给输出结果带来了不确定性，这对图像修复任务并不友好。为了解决这些问题，我们针对 Real - ISR 问题提出了一种单步高效扩散网络，即 OSEDiff。我们认为低质量图像包含了丰富的信息来恢复其对应的高质量图像，因此给定的低质量图像可以直接作为扩散的起点，消除了随机噪声采样带来的不确定性。我们通过可训练层对预训练的扩散网络进行微调，使其适应复杂的图像退化情况。为确保单步扩散模型能够生成高质量的 Real - ISR 输出，我们在潜在空间中应用变分分数蒸馏进行 KL 散度正则化。因此，我们的 OSEDiff 模型仅需一个扩散步骤就能高效且有效地生成高质量图像。我们的实验表明，在客观指标和主观评价方面，OSEDiff 与之前需要数十步或数百步的基于扩散模型的 Real - ISR 方法相比，能取得相当甚至更好的 Real - ISR 结果。

Introduction

图像超分辨率（ISR）是一个经典且仍活跃的研究问题，其目标是从受噪声、模糊和低分辨率等退化影响的低质量（LQ）观测图像中恢复出高质量（HQ）图像。一方面，部分 ISR 研究将从高质量图像到低质量图像的退化过程简化为双三次下采样（或高斯模糊后的下采样），并专注于网络架构设计的研究，但训练出的模型很难推广到现实世界中的低质量图像，因为现实中的图像退化情况往往未知且复杂得多。因此，ISR 研究中另一个日益流行的方向是所谓的现实世界图像超分辨率（Real - ISR）问题，其旨在从现实应用中采集的低质量图像重现出感知上的高质量图像。

训练一个真实世界图像超分辨率（Real - ISR）模型存在两个主要问题。其一，如何构建低质量（LQ） - 高质量（HQ）训练图像对；其二，如何确保恢复图像的自然度，即如何确保恢复后的图像符合高质量自然图像的分布。对于第一个问题，一些研究人员提议使用长焦和短焦相机镜头来收集真实世界的低质量 - 高质量图像对。然而，这种方法成本极高，并且只能涵盖某些类型的真实世界图像退化情况。另一种更为经济的方法是通过使用复杂的图像退化流程来模拟真实世界的低质量 - 高质量图像对。具有代表性的成果包括 BSRGAN 和 Real - ESRGAN，在这两项研究中，分别使用了对基本退化算子的随机组合以及高阶退化模型来生成低质量 - 高质量图像对。

有了给定的训练数据，如何训练一个稳健的真实世界图像超分辨率（Real - ISR）模型，以输出在感知上自然且高质量的图像，就成了一个关键问题。仅仅通过基于像素损失来学习低质量（LQ） - 高质量（HQ）配对数据之间的映射网络，可能会导致结果过度平滑。将自然图像先验知识融入学习过程以重现高质量图像至关重要。为此，已经提出了一些方法。感知损失在诸如 VGG - 16 和 AlexNet 等预训练模型中探索纹理、颜色和结构先验。生成对抗网络（GANs）则通过交替训练生成器和判别器来实现，并且已被用于 Real - ISR 任务。生成器网络旨在合成高质量图像，而判别器网络旨在区分合成图像是否逼真。尽管取得了巨大成功，尤其是在特定类别图像（如人脸图像）的恢复方面，但基于 GAN 的 Real - ISR 由于对抗训练的不稳定性以及在区分不同自然场景的图像空间时存在困难，往往会生成令人不悦的细节。

最近发展起来的生成扩散模型（DM），尤其是大规模预训练的文本转图像（T2I）模型，在各种下游任务中展现出了卓越的性能。经过数十亿图像 - 文本对的训练，预训练的 T2I 模型具备强大的自然图像先验知识，这些知识可被充分利用来提升真实世界图像超分辨率（Real - ISR）输出结果的自然度和感知质量。已经有一些方法被开发出来，利用预训练的 T2I 模型来解决 Real - ISR 问题。虽然与基于生成对抗网络（GAN）的方法相比，这些方法在生成更丰富、更逼真的图像细节方面取得了令人瞩目的成果，但现有的基于扩散模型（SD）的方法仍有几个问题有待进一步解决。首先，这些方法通常将随机高斯噪声作为扩散过程的起点。尽管低质量（LQ）图像通过 ControlNet 模块被用作控制信号，但这些方法在输出的高质量（HQ）图像中引入了不必要的随机性。其次，恢复后的高质量图像通常需要通过数十甚至数百次扩散步骤才能获得，这使得 Real - ISR 过程的计算成本高昂。尽管最近已经提出了一些基于单步扩散的 Real - ISR 方法，但与多步方法相比，它们在实现高质量细节方面有所欠缺。

为解决上述问题，我们针对真实世界图像超分辨率（Real - ISR）问题，提出了一种单步高效扩散网络，简称为 OSEDiff。预训练的扩散模型（SD）中的 UNet 主干网络，具备将输入数据转换到另一领域的强大能力，而已知的低质量（LQ）图像实际上蕴含着丰富信息，可用于恢复其对应的高质量（HQ）图像。因此，我们提议直接将低质量图像输入预训练的 SD 模型，无需引入任何随机噪声。同时，我们将可训练的低秩自适应（LoRA）层集成到预训练的 UNet 中，并对 SD 模型进行微调，使其适用于 Real - ISR 任务。
另一方面，为确保单步模型能像多步模型一样生成高质量的自然图像，我们利用变分分数蒸馏（VSD）进行 KL 散度正则化。这一操作有效利用了预训练 SD 模型强大的自然图像先验知识，并使生成图像的分布与自然图像先验保持一致。如图 1 所示，我们大量的实验表明，OSEDiff 与当前基于 SD 的最先进 Real - ISR 模型相比，性能指标相当或更优，同时它将推理步骤从 N 显著减少到 1，且可训练参数最少，与近期最先进的 SeeSR 相比，推理时间加快了 10 倍以上。

• 图 1：基于 Stable - Diffusion（SD）的真实图像超分辨率重建（Real - ISR）方法的性能与效率比较。(a) 在 DrealSR 基准测试集上的性能比较 [50]。像学习感知图像块相似度（LPIPS）和自然图像质量评价（NIQE）等指标，分数越低表示图像质量越好，为便于显示，对这些指标进行了取反和归一化处理。OSEDiff 仅通过一个扩散步骤就在大多数指标上取得了领先分数。(b) 模型效率比较。推理时间是在输入图像尺寸为 512×512 的 A100 GPU 上进行测试的。OSEDiff 的可训练参数最少，并且比最近开发的 SeeSR 快 10 倍以上。

Methodology

Problem Modeling

真实世界图像超分辨率（Real - ISR）旨在从给定的低质量图像 ${\mathbf{x}}_L$ 中估计出高质量图像 ${\hat{\mathbf{x}}}_H$。该任务通常可建模为一个优化问题：${\hat{\mathbf{x}}}_H={\mathrm{argmin}}_{{\mathbf{x}}_H}({\mathcal{L}}_{data}(\Phi ({\mathbf{x}}_H),{\mathbf{x}}_L)+\lambda {\mathcal{L}}_{reg}({\mathbf{x}}_H))$，
其中 $\Phi$ 是退化函数，${\mathcal{L}}_{data}$ 是数据项，用于衡量优化输出的保真度，${\mathcal{L}}_{reg}$ 是正则化项，用于利用自然图像的先验信息，标量 $\lambda$ 是平衡参数。许多传统的图像超分辨率（ISR）方法通过假设简单且已知的退化模型，并采用手工设计的自然图像先验模型（即基于图像稀疏性的先验）来恢复所需的高质量图像。

然而，这类基于优化的方法的性能在很大程度上受到两个因素的制约。其一，在现实场景中，退化函数 $\Phi$ 往往是未知的，且难以建模。其二，手工设计的正则化项 ${\mathcal{L}}_{reg}$ 很难有效地对复杂的自然图像先验进行建模。

随着深度学习技术的发展，通过从包含（${\mathbf{x}}_L$，${\mathbf{x}}_H$）图像对的训练数据集 $S$ 中学习由参数 $\theta$ 表征的神经网络 $G_{\theta }$，从而将低质量图像映射为高质量图像，已变得十分普遍。该网络训练可描述为如下学习问题：

$${\theta }^{\ast }=\underset{\theta }{\mathrm{argmin}}{\mathbb{E}}_{(x_L,x_H)\sim S}\left[{\mathcal{L}}_{\text{data}}(G_{\theta }(x_L),x_H)+\lambda {\mathcal{L}}_{\text{reg}}(G_{\theta }(x_L))\right]\text{\hspace{1em}(1)}$$

其中 ${\mathcal{L}}_{data}$ 和 ${\mathcal{L}}_{reg}$ 是损失函数。${\mathcal{L}}_{data}$ 要求网络输出 $ \hat{\mathbf{x}}H = G{\theta}(x_L)$ 尽可能接近真实值 $x_H$，这可以通过诸如 $L_1$ 范数、$L_2$ 范数和 LPIPS 等指标来量化。仅使用 ${\mathcal{L}}_{data}$ 损失从零开始训练网络可能会导致对训练数据集过拟合。在这项工作中，我们提议对预训练的生成网络，更具体地说是 Stable Diffusion（SD）网络进行微调，以提高 $G_{\theta }$ 的泛化能力。此外，正则化损失 ${\mathcal{L}}_{reg}$ 对于提高 $G_{\theta }$ 的泛化能力以及增强输出的高分辨率（HQ）图像 ${\hat{\mathbf{x}}}_H$ 的自然度至关重要。假设我们有真实世界高分辨率图像的分布，记为 $p({\mathbf{x}}_H)$，Kullback-Leibler（KL）散度是作为 $L_{reg}$ 损失函数的理想选择；也就是说，恢复的高分辨率图像的分布，记为 $q({\hat{\mathbf{x}}}_H)$，应尽可能与 $p({\mathbf{x}}_H)$ 相同。正则化损失可以定义为：

$${\mathcal{L}}_{reg}={\mathcal{D}}_{KL}(q({\hat{\mathbf{x}}}_H)||p({\mathbf{x}}_H))\text{\hspace{1em}(2)}$$

现有研究大多通过对抗训练来实现上述目标，这需要学习一个判别器，以区分生成的高分辨率（HQ）图像 ${\hat{\mathbf{x}}}_H$ 和真实的高分辨率图像 ${\mathbf{x}}_H)$，并更新生成器 $G_{\theta }$，使 ${\hat{\mathbf{x}}}_H$ 和 ${\mathbf{x}}_H$ 难以区分。然而，判别器通常与生成器一起从零开始训练。它们可能无法掌握高分辨率图像的完整分布，并且缺乏足够的判别能力，从而导致真实图像超分辨率重建（Real - ISR）的性能不理想。

最近开发的文生图（T2I）扩散模型，如 Stable Diffusion（SD），为我们构建损失函数 ${\mathcal{L}}_{reg}$ 提供了新的选择。这些在数十亿图像 - 文本对数据上训练的模型，能够有效地描绘潜在空间中的自然图像分布。已有一些分数蒸馏方法被报道使用 SD 模型，以 KL 散度作为目标来优化图像。特别地，变分分数蒸馏（VSD）从基于粒子的变分优化中推导出了这样一个基于 KL 散度的目标，用于对齐两个扩散模型所表示的分布。基于上述讨论，我们提议通过设计一个高效且有效的单步扩散网络来实现公式（1）中的学习目标。具体而言，我们使用低秩自适应（LoRA）对预训练的 SD 模型进行微调，将其作为我们真实图像超分辨率重建（Real - ISR）的骨干网络，并采用 VSD 作为正则化器，使网络输出的分布与自然高分辨率图像的分布对齐。具体细节将在下一节中给出。

One-Step Effective Diffusion Network

• 图 2：OSEDiff 的训练框架。低质量（LQ）图像通过一个可训练的编码器 $E_{\theta }$、一个经低秩自适应（LoRA）微调的扩散网络 ${ϵ}_{\theta }$ 以及一个固定的解码器 $D_{\theta }$，以获得期望的高质量（HQ）图像。此外，从低质量图像中提取文本提示，并将其输入到扩散网络中，以激发其生成能力。同时，扩散网络 ${ϵ}_{\theta }$ 的输出将被发送到两个正则化器网络（一个是固定的预训练网络，另一个是微调网络），在潜在空间中进行变分分数蒸馏，以确保 ${ϵ}_{\theta }$ 的输出遵循高质量自然图像的分布。正则化损失将通过反向传播来更新和。一旦训练完成，推理过程中仅使用 $E_{\theta }$、${ϵ}_{\theta }$ 和 $D_{\theta }$。

如前面所述，现有的基于 Stable Diffusion（SD）的真实图像超分辨率重建（Real - ISR）方法以随机噪声为起点、低质量（LQ）图像为控制信号，通过多个时间步来估计高质量（HQ）图像。这些方法耗费大量资源，并且本质上会引入随机性。基于第 3.1 节的公式推导，我们提出了一种用于真实图像超分辨率重建的单步高效扩散（OSEDiff）网络，其训练框架如图 2 所示。

我们待训练的生成器由 $G_{\theta }$ 可训练的编码器 $E_{\theta }$、微调后的扩散网络 ${ϵ}_{\theta }$ 和固定的解码器 $D_{\theta }$ 组成。为确保生成器 $G_{\theta }$ 的泛化能力，扩散网络 ${ϵ}_{\theta }$ 的输出将被输入到两个正则化网络，在潜在空间中计算变分分数蒸馏（VSD）损失。正则化损失通过反向传播来更新编码器 $E_{\theta }$ 和扩散网络 ${ϵ}_{\theta }$。训练完成后，推理阶段仅使用生成器 $G_{\theta }$。接下来，我们将详细探讨 OSEDiff 的架构设计及其相关的训练损失。

我们用 $E_{\varphi }$、${ϵ}_{\varphi }$ 和 $D_{\varphi }$ 分别表示预训练的 Stable Diffusion（SD）模型的变分自编码器（VAE）编码器、潜在扩散网络和 VAE 解码器，其中 $\varphi$ 表示模型参数。受低秩自适应（LoRA）近期在将 SD 模型微调至下游任务中取得成功的启发，我们采用 LoRA 方法在真实图像超分辨率重建（Real - ISR）任务中对预训练的 SD 模型进行微调，以获得所需的生成器 $G_{\theta }$。

如图 2 左半部分所示，为了保留 SD 模型原有的生成能力，我们在编码器 $E_{\varphi }$ 和扩散网络 ${ϵ}_{\varphi }$ 中引入可训练的 LoRA 层，利用我们的训练数据将它们微调为 $E_{\theta }$ 和 ${ϵ}_{\theta }$。对于解码器，我们固定其参数，直接令 $D_{\theta }=D_{\varphi }$。这样做是为了确保扩散网络的输出空间与正则化器保持一致。

回顾一下，扩散模型通过 $z_t={\alpha }_{t}z+{\beta }_tϵ$ 对输入的潜在特征 $z$ 进行扩散操作，其中 ${\alpha }_t$ 和 ${\beta }_t$ 是依赖于扩散时间步 t ∈ { 1 , 2 , 3 , . . . T } t \in\{1, 2, 3, ... T \} t∈{1,2,3,...T} 的标量。借助一个能够预测 $z_t$ 中噪声（记为 $\widehat{ϵ}$）的神经网络，去噪后的潜在特征可以通过 ${\hat{z}}_0=z_t-\frac{{\alpha }_t}{{\beta }_t}\widehat{ϵ}$ 得到，我们期望 ${\hat{z}}_0$ 比 $z_t$ 更清晰、更接近真实照片效果。此外，SD 是一个文本条件生成模型。通过从给定的文本描述 $y$ 中提取文本嵌入（记为 $c_y$），可以进行噪声预测，即 $\widehat{ϵ}={ϵ}_{\theta }(z_t;t,c_y)$。

我们将上述从文本到图像的去噪过程应用于真实图像超分辨率重建（Real - ISR）任务，并将从低质量（LQ）到高质量（HQ）的潜在变换 $F_{\theta }$ 表述为一个基于文本条件的图像到图像去噪过程，具体如下：

$${\hat{z}}_H=F_{\theta }(z_L;c_y)\triangleq \frac{z_L-{\beta }_T{ϵ}_{\theta }(z_L;T,c_y)}{{\alpha }_T}\text{\hspace{1em}(3)}$$

其中，我们在第 $T$ 个扩散时间步仅对低质量（LQ）潜在特征 $z_L$ 进行单步去噪，且不引入任何噪声。去噪输出 ${\hat{z}}_H$ 有望比 $z_L$ 更具照片真实感。至于文本嵌入，我们将一些文本提示提取器（如 DAPE）应用于低质量输入 $x_L$，并得到 $c_y=Y(x_L)$。最后，整个从低质量（LQ）到高质量（HQ）的图像合成过程可写为：

$${\hat{x}}_H=G_{\theta }(x_L)\triangleq D_{\theta }(F_{\theta }(E_{\theta }(x_L);Y(x_L)))\text{\hspace{1em}(4)}$$

为了提高真实图像超分辨率重建（Real - ISR）模型的性能，需要同时使用数据项 $L_{data}$ 和正则化 $L_{reg}$ 项来监督生成器的训练。
如图 2 的右半部分所示，我们提议采用变分分数蒸馏（VSD）作为正则化项。除了将 Stable Diffusion（SD）模型用作预训练的正则化器 ${ϵ}_{\varphi }$ 之外，VSD 还引入了一个经过微调的正则化器，即在生成图像的分布 $q_{\theta }({\hat{x}}_H)$ 上通过低秩自适应（LoRA）进行微调的潜在扩散模块。我们将这个微调后的扩散模块记为 ${ϵ}_{{\varphi }^{\prime }}$。

我们使用数据损失 $L_{data}$ 和正则化损失 $L_{reg}$ 来训练生成器 $G_{\theta }$。我们将 $L_{data}$ 设置为均方误差（MSE）损失和学习感知图像块相似度（LPIPS）损失的加权和：

$${\mathcal{L}}_{data}(G_{\theta }(x_L),x_H)={\mathcal{L}}_{MSE}(G_{\theta }(x_L),x_H)+{\lambda }_1{\mathcal{L}}_{LPIPS}(G_{\theta }(x_L),x_H)\text{\hspace{1em}(5)}$$

其中 ${\lambda }_1$ 是一个加权标量。至于 $L_{reg}$，我们通过以下方式采用变分分数蒸馏（VSD）损失：

$${\mathcal{L}}_{reg}(G_{\theta }(x_L))={\mathcal{L}}_{VSD}(G_{\theta }(x_L),c_y)={\mathcal{L}}_{VSD}(G_{\theta }(x_L),Y(x_L))\text{\hspace{1em}(6)}$$

给定任何可训练的图像形状特征 $x$ ，其潜在编码 $z=E_{\varphi }(x)$ 以及编码后的文本提示条件 $c_y$，变分分数蒸馏（VSD）通过以下方式优化 $x$，使其与文本提示 $y$ 一致：

$${\nabla }_{\mathbf{x}}{\mathcal{L}}_{\mathrm{VSD}}(\mathbf{x},c_y)={\mathbb{E}}_{t,ϵ}\left[\omega (t)\left({ϵ}_{\varphi }(z_t;t,c_y)-{ϵ}_{{\varphi }^{\prime }}(z_t;t,c_y)\right)\frac{\partial z}{\partial \mathbf{x}}\right]\text{\hspace{1em}(7)}$$

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

$$\mathcal{L}(G_{\theta }(x_L),x_H)={\mathcal{L}}_{data}(G_{\theta }(x_L),x_H)+{\lambda }_2{\mathcal{L}}_{reg}(G_{\theta }(x_L))\text{\hspace{1em}(8)}$$

其中 ${\lambda }_2$ 是一个加权标量。此外，根据变分分数蒸馏（VSD）的要求，微调后的正则化器 ${ϵ}_{{\varphi }^{\prime }}$ 也应该是可训练的，其训练目标为：

$${\mathcal{L}}_{\text{diff}}={\mathbb{E}}_{t,ϵ,c_y=Y(x_L),{\hat{z}}_H=F_{\theta }(E_{\theta }(x_L);Y(x_L))}{\mathcal{L}}_{\text{MSE}}({ϵ}_{{\varphi }^{\prime }}({\alpha }_t{\hat{z}}_H+{\beta }_tϵ;t,c_y),ϵ)\text{\hspace{1em}(9)}$$

• Algorithm
  

原始的变分分数蒸馏在图像空间中计算梯度。当它用于训练基于 Stable - Diffusion（SD）的生成器网络时，在计算正则化损失 ${\mathcal{L}}_{reg}$ 的过程中会出现重复的潜在解码/编码操作。这不仅计算成本高昂，还会降低正则化的效果。考虑到一个训练良好的变分自编码器（VAE）应满足 $E_{\varphi }(x)=E_{\varphi }(D_{\varphi }(z))\approx z$ 这一事实，我们可以近似地令 $E_{\varphi }({\hat{x}}_H)={\hat{z}}_H$。在这种情况下，我们可以在计算正则化损失时消除冗余的潜在编码 / 解码操作。潜在空间中关于网络参数 $\theta$ 的正则化损失梯度为：

$$\begin{array}{rl}{\nabla }_{\theta }{\mathcal{L}}_{\mathrm{VSD}}(G_{\theta }(x_L),c_y)& ={\nabla }_{{\hat{x}}_H}{\mathcal{L}}_{\mathrm{VSD}}({\hat{x}}_H,c_y)\frac{\partial {\hat{x}}_H}{\partial \theta }\\ & ={\mathbb{E}}_{t,ϵ,{\hat{z}}_t={\alpha }_t{E}_{\varphi }({\hat{x}}_H)+{\beta }_tϵ}\left[\omega (t)\left({ϵ}_{\varphi }({\hat{z}}_t;t,c_y)-{ϵ}_{{\varphi }^{\prime }}({\hat{z}}_t;t,c_y)\right)\frac{\partial {\hat{z}}_H}{\partial {\hat{x}}_H}\frac{\partial {\hat{x}}_H}{\partial \theta }\right]\\ & ={\mathbb{E}}_{t,ϵ,{\hat{z}}_t={\alpha }_t{\hat{z}}_H+{\beta }_tϵ}\left[\omega (t)\left({ϵ}_{\varphi }({\hat{z}}_t;t,c_y)-{ϵ}_{{\varphi }^{\prime }}({\hat{z}}_t;t,c_y)\right)\frac{\partial {\hat{z}}_H}{\partial \theta }\right]\end{array}\text{\hspace{1em}(10)}$$