Learning Diffusion Texture Priors for Image Restoration（DTPM）-CVPR2024

Paper：论文地址

Code：Coming Soon

Baseline：DDPM

文章目录

[3.3. Stage I: Learning Diffusion Texture Priors](#3.3. Stage I: Learning Diffusion Texture Priors)

[3.4. Stage II: Finetuning Diffusion Texture Priors Model with Adapters](#3.4. Stage II: Finetuning Diffusion Texture Priors Model with Adapters)

前言

问题：如何在随机扩散过程中保持高水平的图像保真度，生成准确的背景结构和逼真的纹理细节是其中的关键挑战。

**动机：**挑战和限制需要重新思考在图像恢复中使用扩散范式:

（1）为了保持恢复的高保真度，我们建议使用扩散模型仅恢复纹理层。

（2）认识到扩散范式是一种强大的表示学习工具，需要大量的数据来收敛，我们在高质量图像的大规模数据集上预训练了我们的模型，命名为扩散纹理先验模型(DTPM)。

该模型训练了大约55K个样本，用于重建纹理层，在扩散过程中嵌入纹理先验。

**贡献：**DTPM通过扩散过程显式地建模高质量的纹理细节，而不是全局上下文内容

在训练阶段的第一阶段，我们在大约55K的高质量图像样本上预训练DTPM，之后我们冻结了它的大部分参数。

在第二阶段，我们将条件引导适配器插入 DTPM，并为其配备初始预测器，从而促进其快速适应下游图像恢复任务。

我们的DTPM框架的特点是：(a)一种新的扩散建模方法， 通过残差学习专注于纹理层**，提高了下游恢复任务的效率。(b)** 引入语义代码作为约束**，引导扩散模型产生语义连贯的纹理。(c)为了保持扩散纹理先验的完整性，预训练后DTPM的大部分参数保持固定。我们引入了一个** 条件退化适配器**，以促进快速适应各种恢复任务，而不会有灾难性遗忘的风险。**

一、整体框架

3.3. Stage I: Learning Diffusion Texture Priors

（ 1）Motivation

然而，传统的方法在精确恢复细粒度纹理细节方面遇到了很大的困难，这对于真实的图像重建至关重要。虽然扩散模型在生成逼真图像方面表现出色，但它们在图像恢复任务中的直接应用并不理想，产生的结果一般。我们的方法试图通过将纹理特定的先验整合到扩散模型中来弥补这一点。

这种集成使扩散模型能够专门专注于重建纹理，而不是仅仅专注于降解消除。这种策略利用了扩散模型的固有优势，并减轻了它的弱点，使我们能够利用大量高质量的自然图像来获取先验知识。该方法不仅显著提高了模型的泛化能力，而且显著提高了采样过程的效率。

（ 2）Basic Denoising Block

我们设计了一个基于NAFNet[8]块的去噪块，旨在增强模型对广泛的自然图像内容的适应性。该模块旨在平衡峰值模型性能与减少计算需求和易于处理的不同输入分辨率。如图4所示。

为了结合时间条件，我们利用 MLP 将时间嵌入转换为特定信道的尺度和移位参数。

为了保持和 之间的语义一致性，我们集成了一个用于语义条件的MLP层，该层将语义代码合并为关键条件因素。

（3 ）Semantic code

为了使扩散模型能够产生逼真的纹理细节，我们 使用预训练的ResNet-18[21]来提取语义嵌入作为全局语义代码**，然后将其集成到每个去噪块中，如图4所示。**

这种使用语义代码的方法施加了鲁棒的语义约束，而不是仅仅依赖于相邻的上下文信息，从而允许我们的模型生成与上下文语义一致的纹理。因此，这保证了我们的扩散模型的高保真结果。

（4 ）Training objective

为了优先考虑纹理重建而不是图像的内容，我们的扩散模型被设计为捕获 的残差分布。因此，阶段1的训练目标定义为:

这里， 表示噪声向量，而 是我们的扩散模型。项 表示与噪声表相关的分布。在这种情况下， 表示高质量的输入图像，而 表示其相应的平滑图像。 是通过使用边缘保持平滑方法获得的[63- 论文地址]，在保留上下文内容的同时消除大部分纹理。

3.4. Stage II: Finetuning Diffusion Texture Priors Model with Adapters

（ 1）Initial Predictor

为了在第二阶段增强扩散模型对下游任务的适应性(见图3)，我们从DvSR中提取，集成了一个紧凑和轻量级的U-Net来生成初始输出。

这一初始步骤有效地捕获了最终恢复图像的确定性成分，这对于捕获基本结构信息至关重要。

此外，这个初始输出作为扩散模型的关键指导元素，与我们的扩散纹理先验模型形成了有效的补充。

（ 2）Conditional Guidance Adapter

我们希望DPTM能够充分利用这些知识来适应下游的图像恢复任务，而不是忘记它。因此，我们冻结了DPTM中的大部分参数，并通过图3 (b)和(c)所示的条件引导适配器引入了灵活且参数高效的自适应机制。

条件引导适配器的设计带来了 两个好处：****(i)参数高效学习。 由于DPTM的大部分参数都得到了充分的训练，因此在引入适配器后，我们的模型将更快更好地收敛于下游任务。 (ii)引入附加条件**。引入初始预测器输出的粗结构信息，为模型提供更灵活的条件信息，从而产生更符合预期的纹理细节。**

其中 M 表示将输入特征 的维数 缩放到维数 的MLP层， 表示初始预测器的输出，因此 。

（ 3）Training objective

阶段II的训练目标定义为：

其中 表示初始预测器，如图3所示。重要的是要注意并不需要额外的监督损失或预训练。这是因为上述损失的梯度通过 传播并随后影响 。

二、数据集

我们的方法在 五个具体的图像恢复任务 上进行了评估： 图像运动去模糊、单图像离焦去模糊、图像去噪和去雾、图像去雪和图像雨滴去除 。在初始阶段，我们开发了一个全面的扩散纹理先验模型，在随后的阶段，我们对每个任务量身定制的单个完整模型进行微调。

三、实验与结果

实验设置

我们使用Pytorch框架实现了DTPM，利用了 四个NVIDIA RTX 4090 GPU 。训练包括 两个阶段： 阶段1有800K次迭代，阶段2有600K次迭代。我们使用Adam优化器，将动量值设置为0.9和0.999。初始学习率为 ，采用 余弦退火策略 逐步降低学习率。

扩散过程的结构超过1000个步骤(表示为T)，结合了噪声调度 ，在整个训练过程中从0.0001线性升级到0.02。输入数据始终使用256×256裁剪补丁。数据增强技术，包括水平翻转和随机图像旋转45°和90°，在训练期间应用。我们使用去噪扩散隐式模型(DDIM)[55]来实现我们的扩散模型，从而显著提高了我们的采样速度。