《从因果关系的角度学习失真不变表示以用于图像恢复》学习笔记

paper：2303.06859

GitHub：lixinustc/Causal-IR-DIL: Distortion invariant feature learning for image restoration from a causality perspective

2023 CVPR

摘要

1、介绍

1.1 图像修复任务

1.2 失真不变表示学习

1.3 因果效应估计的挑战与方法

1.4 贡献

2、相关工作

2.1 图像恢复

2.2 因果推断

3、方法

3.1 从因果视角看图像恢复

3.2 反事实失真增强

3.3 失真不变表示学习

3.4 基于元学习的失真不变学习（DIL）实现

4、实验

4.1. 实现细节

4.2. 跨失真度实验

4.2.1 图像去噪实验结果

4.2.2 图像去噪的重建图像可视化

4.2.3 图像去模糊实验结果

4.2.4 混合失真图像恢复实验结果

4.3. 跨失真类型实验

4.3.1 真实图像超分辨率（RealSR）实验结果

4.3.2 图像去噪实验结果

4.3.3 图像去雨实验结果

4.4. 消融研究

摘要

近年来，我们见证了深度神经网络（DNNs）在图像恢复领域的巨大进步。然而，一个关键的局限性是它们无法很好地泛化到具有不同程度或类型的真实世界退化。在本文中，我们首次从因果关系的角度提出了一种新的图像恢复训练策略，以提高DNNs对未知退化的泛化能力。我们的方法称为失真不变表示学习（Distortion Invariant representation Learning, DIL），将每种失真类型和程度视为一个特定的混杂因子（confounder），并通过消除每种退化的有害混杂效应来学习失真不变表示。我们从优化的角度建模不同失真的干预，并基于因果关系的后门准则（back-door criterion）推导出DIL。特别地，我们引入了反事实失真增强（counterfactual distortion augmentation）来模拟虚拟的失真类型和程度作为混杂因子。然后，我们基于相应的失真图像，通过虚拟模型更新（virtually model updating）实例化每种失真的干预，并从元学习（meta-learning）的角度消除它们。大量实验证明了我们的DIL在未见过的失真类型和程度上的泛化能力的有效性。

1、介绍

1.1 图像修复任务

图像恢复（Image Restoration, IR）任务，包括图像超分辨率（image super-resolution）、去模糊（deblurring）、去噪（denoising）、压缩伪影去除（compression artifacts removal）等，在深度学习的推动下取得了令人瞩目的性能。一系列精心设计的骨干网络被用于提升特定退化场景下的恢复性能。卷积神经网络（CNNs）和Transformer是图像恢复任务中常用的骨干网络设计选择。然而，如图所示，当这些方法遇到未见过的退化类型时，性能会显著下降。

在训练中，恢复程度对应于标准差为20的噪声，而在测试中，退化程度不同。图像恢复中常用的训练范式，即经验风险最小化（Empirical Risk Minimization, ERM），在处理分布外退化时表现不佳。具体来说，使用ERM训练的恢复网络仅通过最小化理想重建图像 $I_o$ 与对应干净图像 $I_c$ 之间的距离，挖掘失真图像 $I_d$ 与理想重建图像 $I_o$ 之间的相关性。然而，这种方法也会捕捉到虚假相关性（spurious correlation），从而引入特定退化dd的有害混杂效应。这意味着条件概率 $P(I_o|I_d)$ 也受到失真类型或程度 $d$ 的影响（即 $d\rightarrow I_o|I_d$ ）。

一个鲁棒的、具有良好泛化能力的恢复网络应该是失真不变的（即 $d$ 不影响 $I_o|I_d$ ）。例如，给定两个失真类型相同但退化不同的图像，如果它们包含相同的内容 $I_c$ ，那么理想的恢复网络应该能够从这两个失真图像中恢复出相同的理想重建图像 $I_o$ （即 $P(I_o|I_d, d=d_1) = P(I_o|I_d, d=d_2)$ ）。以往的研究在恢复网络的鲁棒性方面可以大致分为两类：基于失真适应的方案和基于领域适应/转换的方案。基于失真适应的方案旨在估计失真类型或表示，并通过自适应调整恢复网络来处理不同的退化类型。基于领域适应/转换的方案则将不同的退化视为不同的领域，并引入领域适应或领域转换策略。然而，这些研究忽略了对恢复网络泛化能力差的根本原因进行探索。

1.2 失真不变表示学习

相关性并不等于因果关系。学习图像恢复的失真不变表示需要获取失真图像与理想重建图像之间的因果效应，而不仅仅是它们之间的相关性。我们从因果推断的角度构建了图像恢复过程的因果图（Causal Graph）。如图所示，图像恢复中的失真 $D$ （包括失真类型 $D_t$ 或程度 $D_l$ ）是混杂因素，它会引入有害的偏差，并使得恢复过程 $I_d \rightarrow I_o$ 依赖于 $D$ ，因为通过 $I_d \rightarrow D \rightarrow I_o$ 建立了虚假的关系路径。我们希望恢复过程中失真图像Id与理想重建图像Io之间的因果关系是 $I_d \rightarrow I_o$ ，其条件概率可以表示为 $P(I_o|do(I_d))$ 。这里，采用“do”操作来切断从失真 $D$ 到 $I_d$ 的连接，从而去除失真 $D$ 对路径 $I_d \rightarrow I_o$ 的负面混杂效应，学习失真不变的特征表示（即 $D$ 不影响 $I_o|I_d$ ）。

1.3 因果效应估计的挑战与方法

因果效应估计有两个典型的调整标准：后门准则（Back-door Criterion）和前门准则（Front-door Criterion）。后门准则旨在通过遍历已知的混杂因素来去除负面的混杂效应，而前门准则则用来解决无法识别混杂因素的挑战。为了提高恢复网络的泛化能力，我们提出了失真不变表示学习（DIL）方法，并通过从优化的角度实现后门准则来解决这一问题。实现这一目标面临两个挑战：第一个挑战是如何构建混杂因素集（即失真集）。从因果推断的角度看，最好保持失真图像中其他因素不变，仅改变失真类型。然而，在现实中，收集具有相同内容但不同失真类型的图像对是不现实的。受到因果反事实和失真模拟的启发，我们提出了反事实失真增强方法，从常用数据集中选择高质量图像，并在这些图像上模拟不同的失真类型或程度，作为混杂因素。

第二个挑战是如何为后门准则找到一个稳定且合适的实例化方案。以往的研究已将因果推断应用于高层次任务中，并通过注意力干预、特征干预等方法实例化了后门准则，但在图像恢复这一低级任务中，这些方法难以实现。在本文中，我们从优化角度理论上推导出了失真不变表示学习（DIL），通过虚拟更新恢复网络并结合元学习优化策略来消除失真带来的混杂效应。通过这种方式，我们能够实现图像恢复中的因果学习，并基于后门准则实例化DIL。

1.4 贡献

本文的贡献总结如下：

我们从因果推断的角度重新审视了图像恢复任务，并指出恢复网络泛化能力差的原因是因为恢复网络并未独立于训练数据中的失真。
基于因果推断中的后门准则，我们提出了一种新的训练范式——失真不变表示学习（DIL），该方法通过反事实失真增强和元学习优化策略实例化了因果干预，并消除了退化带来的负面混杂效应。
在不同的图像恢复任务上的广泛实验表明，我们的DIL方法能够有效提高恢复网络对未见过的失真类型和程度的泛化能力。

2、相关工作

2.1 图像恢复

图像恢复（Image Restoration, IR）旨在从对应的失真图像中恢复高质量图像，这在提升人类视觉体验方面起着重要作用。随着深度学习的进展，一系列研究在多个图像恢复任务中取得了显著的进展，包括图像去噪、去模糊、超分辨率（SR）等。特别地，大多数研究都致力于根据不同的退化过程精心设计不同图像恢复任务的框架，这些方法大致可以分为两类：基于卷积神经网络（CNN）的框架和基于Transformer的框架。然而，以上的研究仅探讨了如何提高模型对特定退化的归纳偏差能力，而缺乏足够的泛化能力。为提高模型的鲁棒性，一些研究尝试将领域转换或失真自适应学习引入图像恢复任务。与此不同，本文将因果学习引入到图像恢复中，探讨了恢复网络鲁棒性差的原因，并从因果推断的角度提出了失真不变表示学习。

2.2 因果推断

因果推断（Causal Inference）旨在消除混杂因素带来的有害偏差，并发现多个变量之间的因果关系。通过调整标准（如前门准则或后门准则），实现“do”操作以估计因果效应。近年来，深度学习推动了一系列智能任务的广泛发展，例如图像分类、图像分割、目标检测、低级处理等。然而，现有的大量研究集中于拟合输入和输出之间的相关性，而忽略了因果关系。由于混杂因素的存在，网络容易捕捉到输入和输出之间的虚假相关性。例如，如果训练数据中大多数狮子出现在草丛中，模型不可避免地将草丛误认为是狮子。为了消除混杂因素带来的有害偏差，一些研究尝试将因果推断引入到深度学习中。从特征角度对混杂因素的干预进行了建模，从注意力角度将前门准则引入视觉-语言任务中。为了提高泛化能力，将因果学习引入领域适应/泛化任务。然而，上述基于因果推断的方法主要关注于高层次任务。在本文中，我们首次将因果推断应用于图像恢复任务，旨在提高恢复网络在不同失真类型和程度下的泛化能力。

3、方法

3.1 从因果视角看图像恢复

图像恢复旨在恢复失真的图像，图像的退化过程可以表示为一个函数 $I_d = g(I_c,D)$ 。其中， $I_c$ 和 $I_d$ 分别表示干净图像和失真图像， $D$ 表示退化。恢复网络 $f$ 是通过损失函数进行训练，以最小化其理想重建图像 $I_o$ 和原始干净图像 $I_c$ 之间的差异。我们用因果结构图来建模这个过程。在图中， $D$ 表示退化过程， $I_d = g(I_c,D)$ 。 $I_c$ 和 $I_o$ 表示通过最大化 $P(I_c | I_o)$ 的概率， $I_o$ 在 $I_c$ 的监督下进行学习。

此外， $D \to I_o$ 表示从 $D$ 到 $I_o$ 的知识学习过程。 $I_d \to I_o$ 是恢复过程，经过恢复网络 $f$ 得到。

从因果推断的角度来看，图像恢复的因果表示要求恢复网络 $f$ 获取 $I_d \to I_o$ 之间的因果关系（即， $P(I_o \to I_d)$ ）。然而，有两个额外的路径， $I_d \to D \to I_o$ 和 $I_d \to I_c \to I_o$ ，它们引入了虚假相关性，其中 $I_c$ 和 $D$ 是因果推断中的混杂因素。值得注意的是， $I_c$ 在数据集中的多样性和为重建图像 $I_o$ 带来更多生动纹理的特性，使它成为一种有利的混杂因素。我们在本文中不考虑混杂因素 $I_c$ 。

我们旨在提高恢复网络对未知失真的鲁棒性，而这些失真常受到混杂因素 $D$ 的不良影响。如图所示， $I_d \leftarrow D \to I_o$ 的存在使得恢复网络 $f$ 学习的条件概率 $P(I_o| I_d)$ 也会受到失真 $D$ 的影响，即， $f$ 拟合的条件概率实际上是 $P(I_o| I_d,D)$ 。因此，恢复网络 $f$ 对不同失真的鲁棒性较差，因为它并不独立于不同的失真 $D$ 。

一个鲁棒的恢复网络 $f$ 应该独立于不同的失真。为了实现这一目标，我们采用了因果推断中的“后门准则”来实现失真不变学习（DIL）。我们将图像恢复中的后门准则公式：

$P(I_o| do(I_d)) = \sum_{d_i\in D} P(I_o|I_d, d_i)P(d_i) P(d_i) = \frac{1}{n}$

其中，因果条件概率 $P(I_o| do(I_d))$ 是恢复网络 $f$ 在失真不变学习上的优化方向。为了简化优化，我们将每种失真 $d_i$ 的概率设为 $\frac{1}{n}$ ，其中 n 是混杂因素中存在的失真类型和程度的数量。

实现这一目标面临两个关键挑战：1）在现实世界中收集具有相同内容的不同真实失真图像并不简单。2）如何将不同失真对重建过程的干预（即 $P(I_o|I_d, d_i)$ ）实例化到图像恢复中？我们通过反事实失真增强和失真不变表示学习来解决这些问题。

3.2 反事实失真增强

为了让恢复网络学习到失真不变的表示，构建失真集合 $D$ （即混杂因素）是至关重要的。例如，如果我们希望恢复网络具备针对不同失真度数的泛化能力，就需要构建一个包含不同失真级别的失真集合 $D$ 。类似地，通过构建包含不同但相关失真类型的 $D$ ，我们可以提高恢复网络对未知失真类型的泛化能力。此外，为了避免图像内容的影响，每张干净图像最好对应有不同失真类型或度数的失真图像在 $D$ 中。然而，在现实世界中，收集符合上述原则的数据集并不容易，这既费时又费力。

在本文中，我们通过合成失真来构建失真集合 $D$ ，这些失真可以被称为因果推断中的虚拟混杂因素。具体来说，我们收集了一系列高质量的图像 $I_c$ ，并通过修改退化过程生成失真图像，表示为 $I_d = g(I_c, d_i)$ ，其中 $d_i \in D$ 。

3.3 失真不变表示学习

在构建了虚拟混杂因素/失真集合 $D = \{d_i|1\leq i\leq n \}$ 后，我们可以通过实现因果推断中的后门准则来实现失真不变的表示学习。首先，我们介绍恢复网络中的条件概率 $P(I_o | I_d)$ 和常用训练范式——经验风险最小化（ERM）之间的关系。

在图像恢复中，期望通过最大化条件概率 $P(I_o | I_d)$ 来学习理想的重建图像 $I_o$ ，其损失函数表示为 $L(f_\theta (I_d), I_c)$ 。ERM 被用于通过最小化训练数据集 $D = \{I_d, I_c\}$ 上的损失函数来优化网络 $f$ :

$\theta = \arg \min_{\theta} \mathbb{E}_{(I_d, I_c) \sim D}[L(f_\theta (I_d), I_c)]$

实现这一目标的关键挑战在于如何建模条件概率 $P(I_o|I_d,d_i)$ （即，如何建模失真 $d_i \in D$ 对恢复过程 $P(I_o | I_d)$ 的干预）。

本文提出通过优化网络参数来建模失真 $d_i \in D$ 对恢复过程的干预（即， $P(I_o| I_d,d_i)$ 。从上述分析中我们知道，在使用经验风险最小化（ERM）训练配对数据（ $I_d, I_c$ ）时，恢复网络 $f$ 实际上是以失真 $d_i$ 为条件的。这启发我们通过基于 ERM 更新模型参数来实例化不同失真类型或程度 $d_i \in D$ 的干预：

$\phi_{d_{i}} =\theta -\alpha \bigtriangledown _\theta L(f_\theta(I_{d_i}),I_c )$

其中 $\phi_{d_{i}}$ 表示恢复网络经过一步更新后，条件于混杂因素 $d_i$ 的参数。因此，通过最小化损失 $L(f\phi _{d_i}( I_d), I_c)$ ，可以获得条件概率 $P(I_o|I_d,d_i)$ 的最大值。最大化因果条件概率 $P(I_o| do(I_d))$ 的优化方向为：

$\theta = \arg \min_{\theta} \mathbb{E}_{(I_d, I_c) \sim D} \left[\sum_{d_i \in D}^{}L(f\phi _{d_i}(I_d), I_c) \right]$

其中 $D$ 表示包含 $n$ 种失真类型或度数的混杂因素集合。基于以上优化目标，我们从因果推断的角度学习失真不变的表示。

3.4 基于元学习的失真不变学习（DIL）实现

一个有趣的发现是，从因果推断角度得出的失真不变学习（DIL）优化方向与一种典型的元学习策略——MAML（Model-Agnostic Meta-Learning）——是一致的，尽管它们有不同的目的。MAML旨在使网络能够快速适应少样本任务，而我们的目标是提高恢复网络的泛化能力。我们基于这种元学习策略，将DIL应用到图像恢复中。

然而，将优化方向直接融入到实际训练过程中是非常困难的，这在计算上是不可行的。原因在于它需要多次计算和更新梯度，这对于像素级图像恢复任务尤其昂贵。为了简化这一过程，我们利用泰勒展开和逆向展开将公式转化为以下形式：

$\theta = \arg \min_{\theta} \mathbb{E}_{(I_d, I_c)\sim D} \left[ L(f_{\phi _{\bar{d}}}(I_d), I_c) \right]$

其中 $\phi _{\bar{d}}$ 表示恢复网络 $f$ 的参数，在所有失真 $D = \{d_i\},1\leqslant i\leqslant n$ 上进行虚拟更新。我们将这种策略定义为 平行采样，它将DIL的复杂训练过程简化为两步。在本文中，我们称原始采样策略为 串行采样。串行采样和平行采样之间的对比如图所示。

我们还探讨了两种不同的梯度更新策略。根据公式 8 和公式 5，我们可以看到它们都需要二阶梯度，因为梯度是通过两步前向传播计算得到的，如图所示。为了简化计算，Reptile提出了一种替代策略，即通过使用一阶梯度对二阶梯度进行近似。这种优化方向（即梯度）通过计算初始参数和最后一步参数之间的偏差来获得。我们将其集成到我们的DIL中，并称之为 一阶优化。与之对比，公式中的原始优化方法被称为 二阶优化。总之，我们为DIL提出了四种变体，基于上述两种策略：

DIL_sf 使用串行采样和一阶梯度优化。
DIL_pf 使用平行采样和一阶优化。
DIL_ss / DIL_ps 使用二阶优化和串行/平行采样。

4、实验

在这一部分，我们首先描述了实验的实现细节。接着，我们通过两个典型的跨分布场景验证了我们提出的DIL（分布式图像学习）方法的有效性：分别为跨失真度（Cross Distortion Degrees）和跨失真类型（Cross Distortion Types）。特别地，对于跨失真度，我们使用了已见失真度进行训练，并在未见失真度的测试集上进行验证；对于跨失真类型，我们使用合成的失真进行训练，并在真实世界或其他类型的失真上进行验证。

4.1. 实现细节

我们采用了典型的RRDB（残差残差块）作为图像恢复的主干网络，它在各种低层次图像任务中已经展示出了显著的性能。所有实验都在四张NVIDIA 2080Ti GPU上进行。Adam优化器被用来优化ERM和DIL训练框架中的网络参数。

4.2. 跨失真度实验

4.2.1 图像去噪实验结果

在图像去噪任务中，训练数据由带有不同噪声级别的失真图像及其对应的干净图像组成。在训练完恢复网络后，我们在包含未见噪声级别的测试集上进行验证。我们比较了经验风险最小化（ERM）与我们提出的四个DIL变种：DILsf, DILpf, DILss 和 DILps。

实验结果如表1所示。我们可以观察到，所有四个DIL变种在多个未见噪声级别上都表现出了比传统ERM更好的泛化能力。在多个典型场景中，包括自然图像（CBSD68、Kodak24、McMaster）、建筑图像（Urban100）和卡通图像（Manga109），我们的DIL方法甚至比ERM有了高达8.74dB的显著提升。特别地，随着训练和测试数据之间的分布差距增大，DIL能为ERM带来更大的性能提升。进一步地，DILsf在跨失真度任务中表现出了最好的泛化能力，其通过串行采样和一阶优化方法取得了最佳效果。

4.2.2 图像去噪的重建图像可视化

如图4所示，DILsf在未见噪声级别（=30）下，能够恢复更加生动、清晰的图像，而ERM则未能去除噪声，且重建图像明显存在噪声失真。

4.2.3 图像去模糊实验结果

在图像去模糊实验中，我们训练恢复网络，使用不同的高斯模糊级别（[1.0, 2.0, 3.0, 4.0]），并在更加严重的模糊级别（4.2、4.4、4.6、4.8、5.0）上验证其泛化能力。结果如表2所示，随着模糊级别的增加，使用ERM训练的网络性能严重下降，因为未见模糊级别与训练时的模糊级别相差较大。而我们的DIL方法则能在每个未见模糊级别上都取得比ERM更好的恢复效果，特别是在卡通场景Manga109上，我们在模糊级别5.0下取得了2.09dB的增益。

4.2.4 混合失真图像恢复实验结果

除了单一的失真，我们还探讨了DIL方法在混合失真图像恢复中的泛化能力。根据设置，混合失真图像分别在轻微、中等和严重失真级别上进行训练和验证。结果如表3所示，DIL在轻度和中度混合失真上的平均增益分别为1.05dB和0.66dB，相较于ERM方法，DIL能够在未见的混合失真级别上保持更好的性能。

4.3. 跨失真类型实验

4.3.1 真实图像超分辨率（RealSR）实验结果

真实图像超分辨率（RealSR）是一个广受关注的任务，特别是在图像包含复杂的混合失真时（如模糊、低分辨率、噪声等）。然而，真实超分辨率图像的失真/干净图像对非常难以收集。因此，我们采用了Real-ESRGAN的真实失真模拟方法来生成训练数据集，并在未见的RealSRV3和DRealSR数据集上进行验证。实验结果如表5所示，DILsf在RealSRV3和DRealSR数据集上分别超越了ERM 0.29dB和0.26dB。特别地，DILps在跨失真类型任务上比DILsf表现更好，分别在RealSRV3和DRealSR数据集上超过ERM 0.47dB和0.85dB。

4.3.2 图像去噪实验结果

我们还研究了DIL在真实图像去噪任务中的泛化能力。具体来说，我们基于相机ISP过程中的四种颜色空间模拟了四种失真，并生成了训练图像对。然后，我们在常用的真实去噪数据集SIDD和DND上进行了验证。实验结果表明，DILps在SIDD数据集上达到了39.92 dB的PSNR，比ERM高出1.02dB。

4.3.3 图像去雨实验结果

作为扩展实验，我们将DIL应用于图像去雨任务。在不同数据集之间，去雨的类型和级别差异很大。我们使用了三个去雨数据集（DID-MDN、Rain14000、Heavy Rain Dataset）来训练恢复网络，并在未见的Rain100L、Rain12和Rain800数据集上进行了验证。实验结果显示，DILps在Rain100L和Rain12数据集上分别取得了0.76dB和1.63dB的增益。