DiffIR: Efficient Diffusion Model for Image Restoration

发布于:2024-11-29 ⋅ 阅读:(17) ⋅ 点赞:(0)

摘要:

扩散模型(DM)通过将图像合成过程建模为去噪网络的顺序应用,达到了最先进的性能。然而,与图像合成不同,图像恢复(IR)有一个强约束,要求生成的结果必须符合真实图像。因此,对于图像恢复任务,传统的扩散模型在大规模模型上运行大量迭代来估计整个图像或特征图是低效的。为了解决这个问题,我们提出了一种高效的图像恢复扩散模型(DiffIR),该模型包括紧凑的图像恢复先验提取网络(CPEN)、动态图像恢复变换器(DIRformer)和去噪网络。具体来说,DiffIR有两个训练阶段:预训练和扩散模型训练。在预训练阶段,我们将真实图像输入到CPENS1中,捕捉一个紧凑的图像恢复先验表示(IPR),以指导DIRformer。在第二阶段,我们训练扩散模型,仅使用低质量图像(LQ)来直接估计与预训练的CPENS1相同的IRP。我们观察到,由于IPR只是一个紧凑的向量,DiffIR比传统的扩散模型使用更少的迭代次数就能获得准确的估计,并生成更加稳定和真实的结果。由于迭代次数较少,我们的DiffIR可以采用CPENS2、DIRformer和去噪网络的联合优化,从而进一步减少估计误差的影响。我们在多个图像恢复任务中进行了广泛的实验,并在减少计算成本的同时实现了最先进的性能。代码可在 https://github.com/Zj-BinXia/DiffIR 上获取。

解释:

这段话介绍了一种新提出的图像恢复方法 DiffIR,它基于扩散模型(Diffusion Models,DM)。扩散模型通常用于图像合成,效果非常好,它通过逐步去除噪声,生成最终图像。然而,在图像恢复任务中,要求恢复出的图像必须尽量与真实图像一致,这比图像合成要求更高。因此,传统的扩散模型需要大量的计算和迭代步骤来恢复整个图像,效率较低。

DiffIR的创新之处在于通过两种主要方法来提高效率:

  1. 紧凑的图像恢复先验提取网络(CPEN):它从真实图像中提取出一个简化的“图像恢复先验”表示,这个先验可以指导恢复过程,减少不必要的计算。
  2. 动态图像恢复变换器(DIRformer):它利用变换器模型,能够处理图像中各个像素之间的长距离关系,帮助模型更好地恢复图像细节。

DiffIR的训练分为两个阶段:首先用真实图像训练一个先验提取网络(CPEN),然后用低质量图像训练扩散模型,从而生成更为准确的恢复图像。由于图像恢复先验只是一个简单的向量,DiffIR可以用更少的迭代来得到更精确的结果,从而减少了计算资源的消耗,并提高了恢复结果的稳定性和真实感。此外,DiffIR还采用了联合优化策略,进一步提高了图像恢复的精度,并减少了误差。

通过这些创新,DiffIR在多个图像恢复任务中表现出色,并且在减少计算成本的同时达到了最先进的性能。


1. 引言

图像恢复(IR)是一个长期存在的问题,因为它具有广泛的应用价值且本质上是一个病态问题。图像恢复的目标是从低质量(LQ)图像中恢复出高质量(HQ)图像,低质量图像通常会受到各种降质因素(如模糊、遮罩、降采样等)的影响。目前,基于深度学习的图像恢复方法已经取得了显著的成功,因为它们能够从大规模数据集中学习到强大的先验知识。最近,扩散模型(DMs)[54],一种基于去噪自编码器层级构建的模型,在图像合成[23, 55, 12, 24]和图像恢复任务(如修复[40, 50]和超分辨率[52])中取得了优异的结果。具体来说,扩散模型通过逆转扩散过程,迭代地去噪图像。扩散模型表明,基于概率的扩散建模能够实现从随机采样的高斯噪声到复杂目标分布(如逼真的图像或潜在分布[50])的高质量映射,而不会像生成对抗网络(GANs)那样遭遇模式崩溃和训练不稳定问题。

作为一种基于似然的模型,扩散模型需要在大规模去噪模型上执行大量的迭代步骤(大约50到1000步)来建模数据的精确细节,这需要大量的计算资源。与图像合成任务从头生成每个像素不同,图像恢复任务只需要在给定的低质量图像上添加准确的细节。因此,如果扩散模型采用图像合成的方式进行图像恢复,不仅会浪费大量计算资源,还可能生成一些与给定低质量图像不匹配的细节。

本文的目标是设计一种基于扩散模型的图像恢复网络,能够充分且高效地利用扩散模型强大的分布映射能力来恢复图像。为此,我们提出了DiffIR。由于变换器(Transformer)能够建模长距离像素依赖关系,我们将变换器块作为DiffIR的基本单元。我们将变换器块堆叠成Unet形状,形成动态图像恢复变换器(DIRformer),用于提取和聚合多层特征。我们将DiffIR的训练分为两个阶段:(1)在第一阶段(图2(a)),我们开发了一个紧凑的图像恢复先验提取网络(CPEN),从真实图像中提取一个紧凑的图像恢复先验表示(IPR),以指导DIRformer。此外,我们还为DIRformer开发了动态门控前馈网络(DGFN)和动态多头反向注意力(DMTA),充分利用IPR。需要注意的是,CPEN和DIRformer是一起优化的。(2)在第二阶段(图2(b)),我们训练扩散模型,直接从低质量图像(LQ)中估计准确的IPR。由于IPR是轻量级的,只需为恢复过程添加细节,我们的扩散模型能够通过几次迭代准确地估计IPR,并获得稳定的视觉效果。

除了上述的方案和架构创新外,我们还展示了联合优化的有效性。在第二阶段,我们观察到估计的IPR可能仍然存在微小的误差,这会影响DIRformer的性能。然而,传统的扩散模型需要许多迭代步骤,这使得无法将解码器与去噪网络一起优化。由于我们的DiffIR所需的迭代次数较少,我们可以运行所有迭代并得到估计的IPR,进而与DIRformer联合优化。如图1所示,我们的DiffIR在消耗比其他基于扩散模型的方法(如RePaint [40]和LDM [50])少得多的计算资源的情况下,达到了SOTA性能。特别地,DiffIR比RePaint效率高1000倍。我们的主要贡献有三点:

  • 我们提出了DiffIR,一种强大、简单且高效的基于扩散模型的图像恢复基准方法。与图像合成不同,图像恢复任务中的大多数像素已经给定。因此,我们利用扩散模型强大的映射能力,估计一个紧凑的图像恢复先验(IPR)来指导图像恢复,从而提高了图像恢复的效率和稳定性。
  • 我们为动态图像恢复变换器(DIRformer)提出了动态多头反向注意力(DGTA)和动态门控前馈网络(DGFN),充分利用IPR。与之前的潜在扩散模型仅优化去噪网络不同,我们提出了去噪网络与解码器(即DIRformer)联合优化的方法,进一步提高了估计误差的鲁棒性。
  • 大量实验表明,所提出的DiffIR在图像恢复任务中能够达到SOTA性能,同时与其他基于扩散模型的方法相比,消耗的计算资源大大减少。

解释:

这段文字讨论了一个图像恢复新方法,称为 DiffIR,该方法是基于扩散模型的,并且在效率和性能上进行了创新。图像恢复的任务是将低质量的图像(例如,模糊、降采样或遮挡的图像)恢复成高质量的图像。现有的扩散模型在图像合成和恢复任务中取得了不错的效果,但它们通常需要很长时间来迭代处理图像,这导致计算资源的浪费。

DiffIR的创新在于它通过两阶段训练来提高效率并减少计算成本:

  1. 第一阶段:从真实的高质量图像中提取图像恢复的先验信息(IPR),这个信息将指导恢复过程。这个先验信息通过一个叫做CPEN的网络提取,并通过DIRformer网络使用。这一阶段的特点是用较少的计算资源就能获得有效的先验信息。

  2. 第二阶段:使用低质量图像来训练扩散模型,直接从这些图像中估计图像恢复先验(IPR),并通过少量的迭代步骤获得稳定的恢复结果。

与其他方法相比,DiffIR不需要大量的迭代步骤,就能在恢复过程中生成更加准确和稳定的结果,计算效率大大提高。它通过联合优化去噪网络和解码器(DIRformer),增强了对误差的容忍度,进一步提高了恢复的质量和效率。


2. 相关工作

图像恢复(IR)

作为开创性的工作,SRCNN [15]、DnCNN [84] 和 ARCNN [14] 采用紧凑的卷积神经网络(CNN)在图像恢复任务上取得了显著的成绩。之后,基于CNN的方法比传统的图像恢复方法更受欢迎。直到现在,研究人员从不同的角度开展了对CNN的研究,设计了更精细的网络架构和学习方案,如残差块 [29, 81, 6]、生成对抗网络(GAN) [21, 65, 48]、注意力机制 [86, 66, 11, 72, 71, 68, 73]、知识蒸馏 [67] 等等。最近,变换器(Transformer)作为一种自然语言处理模型,在计算机视觉领域获得了广泛的关注。与CNN相比,变换器可以建模不同区域之间的全局交互,并取得了SOTA(最先进的)性能。目前,变换器已被广泛应用于多个视觉任务,如图像识别 [17, 60]、分割 [62, 69, 87, 49]、目标检测 [5, 89] 和图像恢复 [7, 38, 74, 36, 8]。

扩散模型(DM)

扩散模型(DM)[23]在密度估计 [31] 和样本质量 [12] 方面取得了最先进的结果。扩散模型采用参数化的马尔可夫链来优化似然函数的下界,使其能够生成比其他生成模型(如GAN)更准确的目标分布。近年来,扩散模型在图像恢复任务中变得越来越重要,例如超分辨率 [28, 52] 和修复 [40, 50, 10]。SR3 [52] 和 SRdiff [35] 将扩散模型引入图像超分辨率,并取得了比基于GAN的SOTA方法更好的性能。此外,Palette [51] 受到条件生成模型 [44] 的启发,提出了一种用于图像恢复的条件扩散模型。LDM [50] 提出了在潜在空间进行扩散模型,从而提高恢复效率。此外,RePaint [40] 设计了一种通过重新采样扩散模型的迭代来改进去噪策略,用于图像修复。然而,这些基于扩散模型的图像恢复方法直接采用了图像合成中的扩散模型范式。然而,在图像恢复中,大多数像素已经给定,不需要对整个图像或特征图执行扩散模型。我们提出的DiffIR对紧凑的图像恢复先验(IPR)执行扩散模型,这使得扩散模型在图像恢复任务中更加高效和稳定。

解释:

这段文字主要讨论了图像恢复(IR)和扩散模型(DM)在计算机视觉中的应用,尤其是在图像恢复任务中的表现。

  1. 图像恢复(IR)

    • 图像恢复的研究起初主要依赖于卷积神经网络(CNN),如SRCNN、DnCNN等,它们取得了不错的效果。随着研究的深入,许多基于CNN的改进方法相继出现,例如使用残差块、生成对抗网络(GAN)、注意力机制等。近期,变换器(Transformer)也被引入图像恢复领域,尤其是在处理图像的全局依赖关系时,比传统的CNN更有优势,能够取得更好的性能。
  2. 扩散模型(DM)

    • 扩散模型是一类通过反向扩散过程来生成高质量图像的生成模型。它们通过逐步去噪从高斯噪声生成逼真的图像,并在许多任务中取得了SOTA的表现,包括图像恢复任务中的超分辨率和修复。与生成对抗网络(GAN)相比,扩散模型通过优化似然函数的下界来生成更准确的目标分布,因此在许多任务中表现更好。
    • 尽管扩散模型在图像恢复中很有潜力,传统的扩散模型通常是对整个图像或特征图执行去噪,但在图像恢复任务中,许多像素已经是已知的,不需要从头生成。DiffIR方法的创新之处在于,它对一个紧凑的图像恢复先验(IPR)进行扩散建模,从而提高了计算效率,同时也使图像恢复过程更加稳定。

3. 前提知识:扩散模型

在本文中,我们采用扩散模型(DMs)[23]来生成准确的图像恢复先验表示(IPR)。在训练阶段,扩散模型方法定义了一个扩散过程,该过程通过 T T T 次迭代将输入图像 x 0 x_0 x0 转换为高斯噪声 x T ∼ N ( 0 , 1 ) x_T \sim N(0, 1) xTN(0,1)。扩散过程的每次迭代可以描述如下:

q ( x t ∣ x t − 1 ) = N ( x t ; 1 − β t   x t − 1 , β t I ) q(x_t | x_{t-1}) = N \left(x_t; \sqrt{1 - \beta_t} \, x_{t-1}, \beta_t I \right) q(xtxt1)=N(xt;1βt xt1,βtI)

其中, x t x_t xt 是时间步 t t t 时的噪声图像, β t \beta_t βt 是预定义的缩放因子, N N N 表示高斯分布。公式(1)可以进一步简化为:

q ( x t ∣ x 0 ) = N ( x t ; α t ˉ   x 0 , ( 1 − α t ˉ ) I ) q(x_t | x_0) = N \left(x_t; \sqrt{\bar{\alpha_t}} \, x_0, (1 - \bar{\alpha_t}) I \right) q(xtx0)=N(xt;αtˉ x0,(1αtˉ)I)

其中, α t = 1 − β t \alpha_t = 1 - \beta_t αt=1βt α t ˉ = ∏ i = 0 t α i \bar{\alpha_t} = \prod_{i=0}^t \alpha_i αtˉ=i=0tαi

在推理阶段(反向过程),扩散模型方法从高斯随机噪声图像 x T x_T xT 开始,然后逐步去噪,直到得到高质量输出 x 0 x_0 x0

p ( x t − 1 ∣ x t , x 0 ) = N ( x t − 1 ; μ t ( x t , x 0 ) , σ t 2 I ) p(x_{t-1} | x_t, x_0) = N \left(x_{t-1}; \mu_t(x_t, x_0), \sigma_t^2 I \right) p(xt1xt,x0)=N(xt1;μt(xt,x0),σt2I)

其中,均值 μ t ( x t , x 0 ) \mu_t(x_t, x_0) μt(xt,x0) 为:

μ t ( x t , x 0 ) = 1 − α t α t ˉ x t − ϵ 1 − α t \mu_t(x_t, x_0) = \frac{\sqrt{1 - \alpha_t}}{\sqrt{\bar{\alpha_t}}} x_t - \frac{\epsilon}{\sqrt{1 - \alpha_t}} μt(xt,x0)=αtˉ 1αt xt1αt ϵ

方差 σ t 2 \sigma_t^2 σt2 为:

σ t 2 = 1 − α t − 1 ˉ 1 − α t ˉ β t \sigma_t^2 = \frac{1 - \bar{\alpha_{t-1}}}{1 - \bar{\alpha_t}} \beta_t σt2=1αtˉ1αt1ˉβt

其中, ϵ \epsilon ϵ 表示 x t x_t xt 中的噪声,它是反向过程中唯一的不确定变量。扩散模型使用去噪网络 ϵ θ ( x t , t ) \epsilon_\theta(x_t, t) ϵθ(xt,t) 来估计 ϵ \epsilon ϵ。为了训练 ϵ θ ( x t , t ) \epsilon_\theta(x_t, t) ϵθ(xt,t),给定一个干净的图像 x 0 x_0 x0,扩散模型随机采样一个时间步 t t t 和噪声 ϵ ∼ N ( 0 , I ) \epsilon \sim N(0, I) ϵN(0,I),根据公式(2)生成噪声图像 x t x_t xt。然后,扩散模型通过以下公式优化网络参数 θ \theta θ

∇ θ ∥ ϵ − ϵ θ ( α t ˉ   x 0 + ϵ 1 − α t , t ) ∥ 2 2 \nabla_\theta \left\| \epsilon - \epsilon_\theta \left( \sqrt{\bar{\alpha_t}} \, x_0 + \epsilon \sqrt{1 - \alpha_t}, t \right) \right\|^2_2 θ ϵϵθ(αtˉ x0+ϵ1αt ,t) 22

解释:

这段话详细介绍了扩散模型(DM)如何应用于图像恢复任务,特别是如何通过反向过程从噪声生成清晰的图像。

  1. 扩散过程:扩散模型的工作原理是将一个清晰图像通过多个步骤逐渐加上噪声,直到变成纯噪声。每一步的噪声加法是根据一个高斯分布进行的,公式中通过参数 β t \beta_t βt 控制每步加噪的强度。最终,图像变成了一个随机的高斯噪声。

  2. 反向过程:在推理阶段,扩散模型的目标是从这个噪声图像( x T x_T xT)开始,逐步去除噪声,最终恢复出原始的清晰图像( x 0 x_0 x0)。每一步去噪是通过一个去噪网络 ϵ θ ( x t , t ) \epsilon_\theta(x_t, t) ϵθ(xt,t) 来估计噪声 ϵ \epsilon ϵ,并且反向过程会利用这个去噪网络来逐渐改进图像的质量。

  3. 优化训练:为了训练去噪网络,扩散模型会从原始图像随机选择一个时间步,然后加入噪声,生成一个噪声图像 x t x_t xt。接着,模型通过最小化预测噪声与真实噪声之间的误差来优化去噪网络的参数。

简而言之,扩散模型通过模拟图像从清晰到噪声的过程,并在反向过程中恢复清晰图像。这种方式使得图像恢复任务能够在多个迭代步骤中逐步去噪,从而实现高质量的图像恢复。


4. 方法论

传统的扩散模型(DM)[54, 50, 40]需要大量的迭代次数、计算资源和模型参数,才能生成准确且逼真的图像或潜在特征图。尽管扩散模型在从头生成图像(图像合成)方面取得了令人印象深刻的成绩,但直接将图像合成的扩散模型范式应用于图像恢复(IR)是低效的。因为在图像恢复任务中,大部分像素和信息已经给定,直接对整个图像或特征图执行扩散模型不仅会浪费大量的迭代次数和计算资源,而且容易生成更多伪影。总体来说,扩散模型具有强大的数据估计能力,但将现有的图像合成扩散模型范式应用于图像恢复是不高效的。为了解决这个问题,我们提出了一种高效的扩散模型用于图像恢复(即DiffIR),该模型采用扩散模型估计一个紧凑的图像恢复先验表示(IPR)来指导网络恢复图像。由于IPR非常轻量,DiffIR的模型大小和迭代次数可以大大减少,与传统的扩散模型相比,能够生成更准确的估计。在本节中,我们将介绍DiffIR。如图2所示,DiffIR主要由一个紧凑的图像恢复先验提取网络(CPEN)、动态IRformer(DIRformer)和去噪网络组成。我们分两个阶段训练DiffIR,包括预训练DiffIR和训练扩散模型。在接下来的章节中,我们将首先介绍第4.1节中的DiffIR的预训练过程。然后,在第4.2节中,我们将提供DiffIR高效扩散模型训练的细节。

解释:

这段话介绍了DiffIR模型的设计思想和方法。扩散模型(DM)是一个强大的生成模型,它可以通过反向过程从噪声中恢复清晰图像,但传统的扩散模型需要大量计算资源和迭代步骤,尤其是在图像恢复(IR)任务中,这种计算是浪费的,因为在IR任务中大部分信息已经给定。

为了解决这个问题,DiffIR提出了一种更加高效的方法,它通过扩散模型估计一个紧凑的图像恢复先验表示(IPR),从而指导网络进行图像恢复。相比于传统的扩散模型,DiffIR大大减少了计算和迭代次数,因为它只需要对轻量级的IPR进行估计,而不是直接处理整个图像或特征图。

DiffIR的核心包括三个部分:

  1. CPEN(紧凑图像恢复先验提取网络):用来从真实图像中提取IPR。
  2. DIRformer(动态IRformer):负责根据提取的IPR进行图像恢复。
  3. 去噪网络:用来去除噪声,进一步提高恢复图像的质量。

DiffIR的训练分为两个阶段:第一个阶段是预训练,第二个阶段是训练扩散模型。在接下来的章节中,文章将详细介绍这两个阶段的具体方法和技术细节。


4.1. 预训练DiffIR

在介绍预训练DiffIR之前,我们首先介绍第一阶段的两个网络,包括紧凑的图像恢复先验提取网络(CPEN)和动态IRformer(DIRformer)。CPEN的结构如图2黄色框所示,主要由残差块和线性层堆叠而成,用于提取紧凑的图像恢复先验表示(IPR)。之后,DIRformer可以利用提取的IPR来恢复低质量(LQ)图像。DIRformer的结构如图2粉色框所示,采用了以Unet形状堆叠的动态transformer块。动态transformer块包括动态多头反向注意力(DMTA,图2绿色框)和动态门控前馈网络(DGFN,图2深蓝框),这些组件可以将IPR作为动态调节参数,加入恢复细节到特征图中。

在预训练阶段(图2(a)),我们将CPENS1和DIRformer一起训练。具体来说,我们首先将真实图像(IGT)和低质量图像(ILQ)拼接起来,并使用PixelUnshuffle操作对它们进行下采样,以获得CPENS1的输入。然后,CPENS1提取IPR Z ∈ R 4 C ′ Z \in \mathbb{R}^{4C'} ZR4C,计算公式为:

Z = CPENS1 ( PixelUnshuffle ( Concat ( I GT , I LQ ) ) ) . Z = \text{CPENS1}(\text{PixelUnshuffle}(\text{Concat}(I_{\text{GT}}, I_{\text{LQ}}))). Z=CPENS1(PixelUnshuffle(Concat(IGT,ILQ))).

然后,将提取的IPR Z Z Z作为动态调节参数输入到DIRformer的DGFN和DMTA中,指导恢复过程:

F ′ = W l 1 Z ⊙ Norm ( F ) + W l 2 Z , F' = W^1_l Z \odot \text{Norm}(F) + W^2_l Z, F=Wl1ZNorm(F)+Wl2Z,

其中, ⊙ \odot 表示逐元素乘法,Norm表示层归一化[2], W l W_l Wl表示线性层, F F F F ′ ∈ R H ^ × W ^ × C ′ F' \in \mathbb{R}^{\hat{H} \times \hat{W} \times C'} FRH^×W^×C分别为输入和输出特征图。接着,在DMTA中,我们聚合全局空间信息。具体地, F ′ F' F被投影为查询 Q = W d Q W Q c F ′ Q = W^Q_d W^c_Q F' Q=WdQWQcF, 键 K = W d K W K c F ′ K = W^K_d W^c_K F' K=WdKWKcF,和值 V = W d V W V c F ′ V = W^V_d W^c_V F' V=WdVWVcF,其中 W c W^c Wc 1 × 1 1 \times 1 1×1点卷积, W d W^d Wd 3 × 3 3 \times 3 3×3深度卷积。接着,我们重塑查询 Q ′ ∈ R H ^ W ^ × C Q' \in \mathbb{R}^{\hat{H} \hat{W} \times C} QRH^W^×C,键 K ′ ∈ R C × H ^ W ^ K' \in \mathbb{R}^{C \times \hat{H} \hat{W}} KRC×H^W^和值 V ′ ∈ R H ^ W ^ × C V' \in \mathbb{R}^{\hat{H} \hat{W} \times C} VRH^W^×C。之后,进行查询 Q ′ Q' Q与键 K ′ K' K的点积,生成一个转置注意力图 A A A,大小为 R C × C \mathbb{R}^{C \times C} RC×C,这种方式比常规的注意力图 R H ^ W ^ × H ^ W ^ \mathbb{R}^{\hat{H} \hat{W} \times \hat{H} \hat{W}} RH^W^×H^W^更加高效。DMTA的总体过程如下:

F ^ = W c V ′ ⋅ Softmax ( K ^ ⋅ Q ′ / γ ) + F , \hat{F} = W^c V' \cdot \text{Softmax}(\hat{K} \cdot Q' / \gamma) + F, F^=WcVSoftmax(K^Q/γ)+F,

其中, γ \gamma γ是一个可学习的缩放参数。与常规的多头自注意力[17, 7]类似,我们将通道分离为多个头,并计算注意力图。

接着,在DGFN中,我们聚合局部特征。我们使用 1 × 1 1 \times 1 1×1卷积聚合来自不同通道的信息,并采用 3 × 3 3 \times 3 3×3深度卷积聚合来自空间邻近像素的信息。此外,我们采用门控机制来增强信息编码。DGFN的总体过程定义为:

F ^ = GELU ( W d 1 W c 1 F ′ ⊙ W d 2 W c 2 F ′ ) + F . \hat{F} = \text{GELU}(W^1_d W^1_c F' \odot W^2_d W^2_c F') + F. F^=GELU(Wd1Wc1FWd2Wc2F)+F.

我们将CPENS1和DIRformer一起训练,这可以使DIRformer充分利用CPENS1提取的IPR进行恢复。训练损失定义为:

L rec = ∥ I GT − I ^ HQ ∥ 1 , L_{\text{rec}} = \| I_{\text{GT}} - \hat{I}_{\text{HQ}} \|_1, Lrec=IGTI^HQ1,

其中, I GT I_{\text{GT}} IGT I ^ HQ \hat{I}_{\text{HQ}} I^HQ分别是真实图像和恢复的高质量图像。 ∥ ⋅ ∥ 1 \|\cdot\|_1 1表示L1范数。如果某些任务强调视觉质量,例如图像修复和单图像超分辨率(SISR),我们可以进一步添加感知损失和对抗损失。更多细节请参见补充材料。

解释:

这段话描述了DiffIR的预训练过程。首先,DiffIR使用两个关键组件来进行图像恢复:一个是紧凑的图像恢复先验提取网络(CPEN),另一个是动态IRformer(DIRformer)

  • CPEN:这个网络的任务是从真实图像和低质量图像中提取一个紧凑的图像恢复先验表示(IPR),它包含了有助于图像恢复的重要信息。

  • DIRformer:这个网络是DiffIR的核心,它使用**动态多头反向注意力(DMTA)动态门控前馈网络(DGFN)**来恢复图像。它能够利用CPEN提取的IPR作为调节参数,细致地恢复图像中的缺失部分。

在DiffIR的预训练阶段,我们将这两个网络一起训练。训练的步骤如下:

  1. 图像预处理:我们首先将真实图像和低质量图像拼接,然后通过下采样来生成CPEN的输入。
  2. IPR提取:CPEN提取出图像的IPR。
  3. 图像恢复:DIRformer使用这些IPR来恢复图像,并通过动态调节参数引导网络恢复细节。
  4. 训练目标:我们通过计算恢复图像与真实图像之间的差异来优化网络,确保恢复的图像尽可能接近真实图像。

整个预训练过程使得DiffIR能够高效地利用IPR来进行图像恢复,并准备好进行后续的扩散模型训练。


在这里插入图片描述
在这里插入图片描述
在这里插入图片描述
在这里插入图片描述
在这里插入图片描述

图2. 提出的DiffIR概述,包括DIRformer、CPEN和去噪网络。DiffIR有两个训练阶段:(a) 在第一阶段,CPENS1将真实图像作为输入,并输出一个IPR Z Z Z,以引导DIRformer恢复图像。我们将CPENS1与DiffIRS1一起优化,使得DiffIRS1能够充分利用提取的IPR。(b) 在第二阶段,我们利用扩散模型(DM)强大的数据估计能力来估计由预训练的CPENS1提取的IPR。值得注意的是,我们不会将真实图像输入到CPENS2和去噪网络中。在推理阶段,我们只使用扩散模型的反向过程。

解释:

这段话讲的是DiffIR模型的整体结构和训练过程。DiffIR有两个主要部分:DIRformer、CPEN和去噪网络。

  1. 第一阶段:在训练的第一阶段,CPEN(图像恢复先验提取网络)从真实图像中提取一个“先验表示”(IPR),这个IPR指导DIRformer恢复图像。为了确保DIRformer能充分利用CPEN提取的IPR,我们同时训练CPENS1和DiffIRS1(即DiffIR的初始模型)。

  2. 第二阶段:在第二阶段,我们利用扩散模型的强大数据估计能力来估计由CPENS1提取的IPR。这里有一个关键点:我们不再将真实图像输入到CPENS2和去噪网络中,依赖的是通过训练得到的IPR。推理阶段则完全使用扩散模型的反向过程进行图像恢复。

简而言之,DiffIR的训练过程分为两个阶段,第一阶段使用真实图像提取IPR来帮助恢复,第二阶段则依赖扩散模型强大的数据处理能力来进一步恢复图像,而不直接依赖真实图像。


4.2. 用于图像恢复的扩散模型

在第二阶段(图2 (b)),我们利用扩散模型(DM)的强大数据估计能力来估计IPR。具体来说,我们使用预训练的CPENS1来捕获IPR Z ∈ R 4 C ′ Z \in R^{4C'} ZR4C。之后,我们对 Z Z Z应用扩散过程,以采样 Z T ∈ R 4 C ′ Z_T \in R^{4C'} ZTR4C,这个过程可以描述为:

q ( Z T ∣ Z ) = N ( Z T ; α ˉ T Z , ( 1 − α ˉ T ) I ) (10) q(Z_T | Z) = N \left(Z_T; \sqrt{\bar{\alpha}_T} Z, (1 - \bar{\alpha}_T) I \right) \quad \text{(10)} q(ZTZ)=N(ZT;αˉT Z,(1αˉT)I)(10)

其中 T T T是总的迭代次数, α \alpha α α ˉ \bar{\alpha} αˉ在公式(1)和(2)中定义(即 α ˉ T = ∏ i = 0 T α i \bar{\alpha}_T = \prod_{i=0}^T \alpha_i αˉT=i=0Tαi)。

在反向过程,由于IPR是紧凑的,DiffIRS2可以使用更少的迭代和更小的模型尺寸来获得比传统的扩散模型[50, 40]更好的估计。传统的扩散模型在迭代过程中具有巨大的计算成本,它们必须随机采样一个时间步 t ∈ [ 1 , T ] t \in [1, T] t[1,T],并且仅在该时间步优化去噪网络(公式(1)、(2)、(3)和(4))。由于去噪网络和解码器(即DIRformer)没有联合训练,去噪网络产生的微小估计误差会使DIRformer无法达到其潜力。相比之下,DiffIR从 T T T-th时间步(公式(10))开始,并运行所有去噪迭代(公式(11))来获得$ \hat{Z}$,然后将其送入DIRformer进行联合优化。

Z ^ t − 1 = 1 α t Z ^ t − ε 1 − α t 1 − α ˉ t (11) \hat{Z}_{t-1} = \sqrt{\frac{1}{\alpha_t}} \hat{Z}_t - \varepsilon \sqrt{\frac{1 - \alpha_t}{1 - \bar{\alpha}_t}} \quad \text{(11)} Z^t1=αt1 Z^tε1αˉt1αt (11)

其中 ε \varepsilon ε表示相同的噪声,我们使用CPENS2和去噪网络来预测噪声,如公式(3)所示。值得注意的是,与传统的扩散模型(公式(3))不同,我们的DiffIRS2删除了方差估计,并发现这种方法有助于准确估计IPR并获得更好的性能(第6节)。

在扩散模型的反向过程中,我们首先使用CPENS2从低质量(LQ)图像中获取条件向量 D ∈ R 4 C ′ D \in R^{4C'} DR4C

D = C P E N S 2 ( PixelUnshuffle ( I L Q ) ) (12) D = CPENS2(\text{PixelUnshuffle}(I_{LQ})) \quad \text{(12)} D=CPENS2(PixelUnshuffle(ILQ))(12)

其中CPENS2的结构与CPENS1相同,唯一不同的是第一层卷积的输入维度。然后,我们使用去噪网络 ε θ \varepsilon_\theta εθ在每个时间步 t t t估计噪声,公式为 ε θ ( Concat ( Z ^ t , t , D ) ) \varepsilon_\theta(\text{Concat}(\hat{Z}_t, t, D)) εθ(Concat(Z^t,t,D))。估计的噪声被代入公式(11)以获得 Z ^ t − 1 \hat{Z}_{t-1} Z^t1,开始下一次迭代。经过 T T T次迭代后,我们得到最终估计的IPR Z ^ ∈ R 4 C ′ \hat{Z} \in R^{4C'} Z^R4C

我们联合训练CPENS2、去噪网络和DIRformer,使用损失函数 L a l l L_{all} Lall

L d i f f = 1 4 C ′ ∑ i = 1 4 C ′ ( Z ^ ( i ) − Z ( i ) ) , L a l l = L r e c + L d i f f (13) L_{diff} = \frac{1}{4C'} \sum_{i=1}^{4C'} \left(\hat{Z}^{(i)} - Z^{(i)}\right), \quad L_{all} = L_{rec} + L_{diff} \quad \text{(13)} Ldiff=4C1i=14C(Z^(i)Z(i)),Lall=Lrec+Ldiff(13)

其中我们可以进一步在 L a l l L_{all} Lall中加入感知损失和对抗损失,以获得更好的视觉质量,如公式(9)所示。

在推理阶段,我们只使用扩散模型的反向过程(图2 (b)的下部分)。CPENS2从低质量图像中提取条件向量 D D D,然后我们随机采样一个高斯噪声 Z T Z_T ZT。去噪网络利用 Z ^ T \hat{Z}_T Z^T D D D来估计IPR Z ^ \hat{Z} Z^,经过 T T T次迭代后得到。然后,DIRformer利用估计的IPR恢复低质量图像。

解释:

这一段讲的是如何在第二阶段使用扩散模型(DM)来估计图像恢复中的先验表示(IPR)。具体来说:

  1. 使用扩散过程:首先,使用预训练的CPENS1从真实图像中提取一个紧凑的IPR( Z Z Z)。然后,应用扩散过程(公式10),对 Z Z Z进行迭代采样,得到 Z T Z_T ZT,并通过反向过程(公式11)来逐步恢复IPR。

  2. 去噪网络与DIRformer的联合优化:与传统的扩散模型不同,DiffIR通过将去噪网络和DIRformer联合训练,从而在更少的迭代和更小的模型尺寸下获得更加准确的IPR估计。通过这种方法,DiffIR能够减少计算量并避免生成更多的伪影。

  3. 条件向量的提取:在反向扩散过程中,我们通过CPENS2从低质量(LQ)图像中提取条件向量 D D D。然后,去噪网络使用这个条件向量和当前的IPR估计进行迭代,最终得到准确的IPR。

  4. 损失函数:为了训练模型,DiffIR使用了一个联合损失函数 L a l l L_{all} Lall,该函数包含了恢复损失( L r e c L_{rec} Lrec)和扩散过程损失( L d i f f L_{diff} Ldiff)。还可以加入感知损失和对抗损失来提高图像的视觉质量。

最后,在推理阶段,只需通过扩散模型的反向过程估计IPR,然后用DIRformer恢复图像。