浅析扩散模型与图像生成【应用篇】(十二)——DiffI2I

12. DiffI2I: Efficient Diffusion Model for Image-to-Image Translation

  该文提出一种基于扩散模型的图像到图像的转换算法（DiffI2I）,可用于图像修复、超分辨率提升、图像去模糊、语义分割等任务。作者指出一般的扩散模型，虽然在图像生成任务中表现很好，但是在图像转换任务表现不佳，因为许多图像转换任务都包含一个约束条件，要求生成结果要能与GroundTruth保持一致，而图像生成任务中并没有此类要求。为此，作者提出一个图像转换先验信息表征（I2I Prior Representation， IPR）的概念，用于引导一个动态的I2IFormer模型（DI2Iformer）来完成图像转换。此外，由于很多扩散模型是直接对图像进行扩散和去噪的，这就导致模型的维度很大，计算量很高，这也导致很多扩散模型不能处理高分辨率的图像，且生成速度很慢。作者采用类似LDM的思想，不对图像直接进行扩散和生成，而是对压缩过的特征信息IPR进行扩散和生成，这就极大的提升了图像转换的效率。在多个任务中的转换效果和效率与其他算法的对比如下图所示

  那么作者是怎么实现这样一个又好又快的方法的呢？在DiffI2I中包含三个关键的部分：1. 紧凑的图像转换先验提取网络（CPEN）用于根据输入图像和GT生成对应的IPR；2.动态图像转换Former（DI2Iformer）在IPR的引导下根据输入图像生成输出图像；3. 扩散和去噪网络，用于根据输入的IPR $Z$生成估计的$\hat{Z}$。整个算法的流程如下图所示

  整个DiffI2I的训练过程分成两个阶段，第一阶段是对CPEN和DI2Iformer进行预训练，第二阶段是对去噪模型以及DI2Iformer进行联合训练。首先看第一阶段，将输入图像和对应的GT级联起来，经过PixelUnshuffle算子进行下采样后，输入到CPEN网络中提取IPR 。CPEN网络结构如图中黄色方框内所示，主要由卷积层、残差块、平均池化层、线性层和激活层等结构组成，输出的是一个经过紧凑的带有先验信息的特征图 $Z$。然后，在 $Z$的引导下，DI2Iformer根据输入图像$I_{input}$输出目标图像$I_{output}$。DI2Iformer是一个由多个动态Transformer块组成的UNet结构网络，如图中灰色的方框中所示。每个动态Transformer块由动态注意力模块（DA）和动态前向传播网络（DFFN）构成，DA和DFFN模块的结构分别如图中绿色和蓝色方框内所示。对于DA而言，输入特征$F$经过规则化层，$1\ast 1$卷积和$3\ast 3$的深度维度卷积得到$F^{\prime }$$${\mathbf{F}}^{\prime }=W_d{W}_c\mathrm{Norm}(\mathbf{F})$$然后经过一个SimpleGate （SG）激活操作得到$F_{SG}^{\prime }$ $${\mathbf{F}}_{SG}^{\prime }=\mathrm{SG}\left({\mathbf{F}}_1^{\prime },{\mathbf{F}}_2^{\prime }\right)={\mathbf{F}}_1^{\prime }\odot {\mathbf{F}}_2^{\prime }$$采用通道注意力机制对$F_{SG}^{\prime }$的各个通道进行加权求和得到$F_{CA}^{\prime }$ $${\mathbf{F}}_{CA}^{\prime }={\mathbf{F}}_{SG}^{\prime }\odot \varphi \left({\mathbf{F}}_{SG}^{\prime }\right)$$其中$\varphi$表示全局平均池化操作，最后$F_{CA}^{\prime }$再次经过$3\ast 3$的深度维度卷积，并与经过线性映射层的先验信息$Z$和最初输入的特征$F$求和，得到DA模块的输出$\hat{F}$ $$\hat{\mathbf{F}}=W_d{\mathbf{F}}_{CA}^{\prime }+W_l\mathbf{Z}+\mathbf{F}$$DFFN模块的结构与DA模块非常类似，只是缺少计算全局平均池化和进行通道注意力加权的部分，计算方式如下$$\hat{\mathbf{F}}=W_d^2\mathrm{SG}\left(W_d^1{W}_c\mathrm{Norm}(\mathbf{F})\right)+W_l\mathbf{Z}+\mathbf{F}$$关于DA模块和DFFN模块的输出是怎么进行融合的，原文中我没看到，但是图中可以看出来动态Transformer块的输入和输出维度不变，那么大概率是做了一个相加的操作。经过多个级联的动态Transformer块处理后，输出第一阶段的生成结果$I_{output}$，用于计算损失，对CPEN和DI2Iforme进行预训练。
  完成预训练后，我们进入第二阶段对去噪模型以及DI2Iformer进行联合训练。首先，在扩散阶段输入图像和GT级联起来，并经过PixelUnshuffle算子进行下采样后，使用预训练得到的CPEN模块进行处理得到IPR $Z$，然后经过一系列的扩散过程得到噪声$Z_T$，接着以$Z_T$作为起点进行反向去噪，去噪网络${ϵ}_{\theta }\left(\mathrm{Concat}\left({\hat{\mathbf{Z}}}_t,t,\mathbf{D}\right)\right)$包含三个输入，${\hat{\mathbf{Z}}}_t$是前一时刻$t$的输出结果，初始时刻${\hat{\mathbf{Z}}}_T={\mathbf{Z}}_T$，$t$表示前一时刻，$\mathbf{D}$表示IPR。要注意的是，这里的IPR也就是$\mathbf{D}$是用一个全新的CPEN模型仅对输入图像进行处理得到的，如下$$\mathbf{D}={\mathrm{CPEN}}_{\mathrm{S}2}\left(\text{ PixelUnshuffle }\left(I_{\text{input }}\right)\right)$$最后，经过生成过程得到的$\hat{\mathbf{Z}}$将于输入图像$I_{input}$一起输入到DI2Iformer中得到输出结果。${\mathrm{CPEN}}_{\mathrm{S}2}$，去噪模型${ϵ}_{\theta }$和DI2Iformer一起进行训练，损失函数如下$${\mathcal{L}}_{\text{diff }}=\frac{1}{4C^{\prime }}\sum _{i=1}^{4C^{\prime }}|\hat{\mathbf{Z}}(i)-\mathbf{Z}(i)|,{\mathcal{L}}_{\text{all }}={\mathcal{L}}_{\text{task }}+{\mathcal{L}}_{\text{diff }}$$其中${\mathcal{L}}_{\text{task }}$与具体任务有关，比如去噪任务就是用GT和输出结果之间的L1距离，语义分割任务就用交叉熵损失。
  在预测时，输入图像经过${\mathrm{CPEN}}_{\mathrm{S}2}$处理得到 $\mathbf{D}$，然后与采样得到的高斯噪声${\hat{\mathbf{Z}}}_T$一起经过反向去噪过程得到IPR $\hat{\mathbf{Z}}$，将其与输入图像一起输入到DI2Iformer模块中得到预测结果。训练和预测过程如下


  由于扩散和生成阶段是对紧凑的先验信息IPR进行处理的，因此扩散和生成阶段计算复杂度较小，只需要迭代4次($T=4$)就可以生成较好的结果。在图像修复任务中与其他方法的结果对比如下

运动去模糊任务

超分辨率提升任务