---

Restoring Vision in Adverse Weather Conditions with Patch-Based Denoising Diffusion Models

ÖZDENIZCI O, LEGGENSTEIN R. Restoring vision in adverse weather conditions with patch-based denoising diffusion models[J]. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2023, 45(7): 8202-8215.

https://github.com/IGITUGraz/WeatherDiffusion

---

DDPM

前向过程

$$\begin{array}{}\text{(1)}& q({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1})=\mathcal{N}({\mathbf{x}}_t;\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1},{\beta }_t\mathbf{I})\end{array}$$$$\begin{array}{}\text{(2)}& q({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0)=\prod _{t=1}^{T}q({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1})\end{array}$$

反向过程

$$\begin{array}{}\text{(3)}& p_{\theta }\left({\mathbf{x}}_{0:T}\right)=p\left({\mathbf{x}}_T\right)\prod _{t=1}^T{p}_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t)\end{array}$$$$\begin{array}{}\text{(4)}& p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t)=\mathcal{N}({\mathbf{x}}_{t-1};{\mathit{\mu }}_{\theta }({\mathbf{x}}_t,t),{\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_t,t))\end{array}$$

损失函数

$$\begin{array}{}\text{(5)}& \begin{array}{rl}{L}_{\theta }={\mathbb{E}}_q& [\underset{L_T}{\underbrace{D_{\mathrm{KL}}(q({\mathbf{x}}_T\mid {\mathbf{x}}_0)\parallel p\left({\mathbf{x}}_T\right))}}\underset{L_0}{\underbrace{-\log p_{\theta }({\mathbf{x}}_0\mid {\mathbf{x}}_1)}}\\ & +\sum _{t>1}\underset{L_{t-1}}{\underbrace{D_{\mathrm{KL}}(q({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0)\parallel p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t))}}]\end{array}\end{array}$$

该损失函数可通过随机梯度下降对随机采样的 $L_{t-1}$ 项进行高效优化，这得益于我们可以通过以下方式对高斯扩散过程进行边缘化，从而直接从干净数据 $x_0$ 采样得到中间状态 $x_t$：

$$\begin{array}{}\text{(6)}& q(x_t\mid x_0)=\mathcal{N}(x_t;\sqrt{{\overline{\alpha }}_t}{x}_0,(1-{\overline{\alpha }}_t)I)\end{array}$$

该式也可写成如下闭式形式：

$$\begin{array}{}\text{(7)}& x_t=\sqrt{{\overline{\alpha }}_t}{x}_0+\sqrt{1-{\overline{\alpha }}_t}{ϵ}_t\end{array}$$

其中 ${\alpha }_t=1-{\beta }_t$，${\overline{\alpha }}_t=\prod _{i=1}^t{\alpha }_i$，且 ${ϵ}_t\sim \mathcal{N}(0,I)$ 与数据 $x_0$ 及潜变量 $x_t$ 维度相同。

此处，式 $(5)$ 中的 $L_{t-1}$ 项比较了两个高斯分布之间的 KL散度，即式 $(4)$ 中的 $p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t)$ 与 $q(x_{t-1}\mid x_t,x_0)$，后者是以 $x_0$ 为条件的真实未知生成过程后验分布，表示为：

$$\begin{array}{}\text{(8)}& q(x_{x-1}\mid x_t,x_0)=\mathcal{N}(x_{x-1};{\tilde{\mu }}_t(x_t,x_0),{\tilde{\beta }}_{t}I)\end{array}$$

其分布参数可写作：

$$\begin{array}{}\text{(9)}& {\tilde{\mu }}_t=\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{{\beta }_t}{\sqrt{1-{\overline{\alpha }}_t}}{ϵ}_t),{\tilde{\beta }}_t=\frac{(1-{\overline{\alpha }}_{t-1})}{(1-{\overline{\alpha }}_t)}{\beta }_t\end{array}$$

为了方便训练目标，我们可以考虑固定的反向过程方差 ${\mathrm{\Sigma }}_{\theta }(x_t,t)={\sigma }_t^{2}I$ （例如 ${\sigma }_t^2={\tilde{\beta }}_t$）,或者通过混合学习目标优化 ${\mathrm{\Sigma }}_{\theta }(x_t,t)$。

对于前者，当 $p_{\theta }(x_{t-1}\mid x_t)=\mathcal{N}(x_{t-1};{\mu }_{\theta }(x_t,t),{\sigma }_t^{2}I)$ 时，整体的训练目标对应于训练一个网络 ${\mu }_{\theta }(x_t,t)$ 来预测 ${\tilde{\mu }}_t$。通过以下方式对反向过程进行重新参数化：

$$\begin{array}{}\text{(10)}& {\mu }_{\theta }(x_t,t)=\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{{\beta }_t}{\sqrt{1-{\overline{\alpha }}_t}}{ϵ}_{\theta }(x_t,t))\end{array}$$

该模型可以通过优化重新加权的简化目标来训练以预测噪声向量 ${ϵ}_{\theta }(x_t,t)$：

$$\begin{array}{}\text{(11)}& {\mathbb{E}}_{x_0,t,{ϵ}_t\sim \mathcal{N}(0,I)}\left[{\Vert {ϵ}_t-{ϵ}_{\theta }(\sqrt{{\alpha }_t}{x}_0+\sqrt{1-{\alpha }_t}{ϵ}_t,t)\Vert }^2\right]\end{array}$$

在这种设置下，我们优化一个网络，使其能够从 $x_t$ 预测时刻 $t$ 的噪声 ${ϵ}_t$。使用学习到的参数化高斯转移分布 $p_{\theta }(x_{t-1}\mid x_t)$ 进行采样时，可以从 $x_T\sim \mathcal{N}(0,I)$ 开始，通过下式进行：

$$\begin{array}{}\text{(12)}& x_{x-1}=\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{{\beta }_t}{\sqrt{1-{\alpha }_t}}{ϵ}_{\theta }(x_t,t))+{\sigma }_{t}z\end{array}$$

其中 $z\sim \mathcal{N}(0,I)$，这类似于通过朗之万动力学进行进一步采样。

高斯分布采样

从 $\mathcal{N}(\mu ,{\sigma }^{2}I)$ 采样的标准方法为 $x=\mu +\sigma z$，其中 $z\sim \mathcal{N}(0,I)$

较大的 $T$ 和前向过程中较小的 ${\beta }_t$ 允许我们假设反向过程接近高斯分布，但这会导致采样成本高昂，例如当 $T=1000$ 时。方差调度通常选择为 ${\beta }_1<{\beta }_2<\dots <{\beta }_T$，从而对噪声更大的样本执行更大的更新。我们专注于使用固定的、线性增加的方差调度，如最初在 DDPM 中发现的那样已经足够，当然也可以基于例如信噪比估计来学习这个调度。

方差调度

• 早期：β小，加噪少 → 主要学习"细微去噪"
• 后期：β大，加噪多 → 主要学习"大幅去噪"

朗之万方程

对于一个粒子在势能场 $U(x)$ 中的运动，郎之万方程为：

$$\frac{dx}{dt}=\underset{\text{确定性漂移项}}{\underbrace{-\mathrm{\nabla }U(x)}}+\underset{\text{随机扩散项}}{\underbrace{\sqrt{2D}\cdot \eta (t)}}$$

将连续时间方程离散化：

$$x_{k+1}=x_k-\gamma \mathrm{\nabla }U\left(x_k\right)+\sqrt{2\gamma D}\cdot z_k$$

其中 $z_k\sim \mathcal{N}(0,I)$

DDIM

DDIM 提出了一种针对预训练扩散模型的的加速确定性采样方法，经验证能够产生一致性更好、质量更高的图像样本。隐式采样利用了广义的非马尔可夫前向过程公式：

$$\begin{array}{}\text{(13)}& q_{\lambda }(x_{1:T}\mid x_0)=q_{\lambda }(x_T\mid x_0)\prod _{t=2}^T{q}_{\lambda }(x_{t-1}\mid x_t,x_0)\end{array}$$

其中，将式 $(8)$ 中的分布按其标准差 ${\lambda }_t$ 的特定选择重新表述为：

$$\begin{array}{}\text{(14)}& q_{\lambda }(x_{t-1}\mid x_t,x_0)=\mathcal{N}(x_{t-1};{\tilde{\mu }}_t(x_t,x_0),{\lambda }_t^{2}I)\end{array}$$

且均值用方差表示为：

$$\begin{array}{}\text{(15)}& {\tilde{\mu }}_t=\sqrt{{\overline{\alpha }}_{t-1}}{x}_0+\sqrt{1-{\overline{\alpha }}_{t-1}-{\lambda }_t^2}\cdot {ϵ}_t\end{array}$$

这是将 $(7)$ 代入 ${\tilde{\mu }}_t(x_t,x_0)$ 得到的。此处，若令 ${\lambda }_t^2={\beta }_t$，则前向过程变为马尔可夫过程，并恢复为先前描述的原始扩散模型公式。重要的是，训练目标 $(11)$ 保持不变，仅在推理阶段利用了嵌入的非马尔可夫前向过程。

确定性隐式采样通过设定 ${\lambda }_t^2=0$ 实现，因此从边缘噪声分布生成初始 $x_T$ 后，采样过程变为确定性过程。我们将类似地通过设置 ${\lambda }_t^2=0$ 来使用我们的模型。使用噪声估计网络进行隐式采样的公式如下：

$$\begin{array}{}\text{(16)}& x_{t-1}=\sqrt{{\overline{\alpha }}_{t-1}}\left(\frac{x_t-\sqrt{1-{\overline{\alpha }}_t}\cdot {ϵ}_{\theta }(x_t,t)}{\sqrt{{\overline{\alpha }}_t}}\right)+\sqrt{1-{\overline{\alpha }}_{t-1}}\cdot {ϵ}_{\theta }(x_t,t)\end{array}$$

在加速采样过程中，我们只需要完整时间步索引 $\{1,\dots ,T\}$ 的一个子序列 ${\tau }_1,{\tau }_2,\dots ,{\tau }_S$。这可以将采样所需要的时间步数减少多达两个数量级。我们通过从 $\{1,\dots ,T\}$ 中均匀间隔选取来确定该子序列：

$$\begin{array}{}\text{(17)}& {\tau }_i=(i-1)\cdot T/S+1\end{array}$$

这使反向采样的最后一步对应 ${\tau }_1=1$。

条件扩散模型

条件扩散模型的核心思想是学习一个条件反射过程 $p_{\theta }(x_0:T\mid \tilde{x})$，从而无需修改针对 $x$ 的扩散过程 $q(x_1:T\mid x_0)$，使得采样的 $x$ 在给定条件 $\tilde{x}$ 下，能够高度保真地反映对应的数据分布。

在训练过程中，我们从配对数据分布（如一张清晰图像 $x_0$ 和一张天气降质的图像 $\tilde{x}$）中采样 $(x_0,\tilde{x})\sim q(x_0,\tilde{x})$，并学习一个条件扩散模型。在此模型中，我们将 $\tilde{x}$ 作为输入提供给反向过程：

$$\begin{array}{}\text{(18)}& p_{\theta }(x_0:T\mid \tilde{x})=p\left(x_T\right)\prod _{t=1}^T{p}_{\theta }(x_{t-1}\mid x_t,\tilde{x})\end{array}$$

我们先前通过公式 $(11)$ 优化噪声估计网络的表达，现在则使用 ${ϵ}_{\theta }(x_t,\tilde{x},t)$。对于基于图像的条件输入，输入 $x$ 和 $\tilde{x}$ 会在通道维度上进行拼接，从而得到一个六通道的输入图像。

需要注意的是，让反向过程以 $\tilde{x}$ 为条件，并不影响其与隐式采样的兼容性。在此框架下，从 ${\mathbf{x}}_{t-1}\sim p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,\tilde{\mathbf{x}})$ 采样的公式如下：

$$\begin{array}{}\text{(19)}& \begin{array}{rl}{\mathbf{x}}_{t-1}=& \sqrt{{\overline{\alpha }}_{t-1}}\left(\frac{{\mathbf{x}}_t-\sqrt{1-{\overline{\alpha }}_t}\cdot {\mathit{ϵ}}_{\theta }({\mathbf{x}}_t,\tilde{\mathbf{x}},t)}{\sqrt{{\overline{\alpha }}_t}}\right)\\ & +\sqrt{1-{\overline{\alpha }}_{t-1}}\cdot {\mathit{ϵ}}_{\theta }({\mathbf{x}}_t,\tilde{\mathbf{x}},t)\end{array}\end{array}$$

该公式定义了一条从 $x_T\sim \mathcal{N}(0,I)$ 开始的、朝向 $x_0$ 的确定性路径，并保证了生成结果对条件 $\tilde{x}$ 的保真度。

基于块的扩散式图像恢复

图像修复的基准测试集以及现实世界中的图片，都包含各种尺寸的图像。相反，现有的生成模型架构大多是为固定尺寸的图像处理而定制的。最近有一项扩散模型研究工作探索了尺寸无关的模糊图像恢复。他们的模型使用固定尺寸的图像块进行优化，然后在测试时通过简单地向模型提供任意尺寸的输入来进行去模糊，因此严格依赖于一种修改过的全卷积网络架构。这也导致了高昂的测试时计算需求，以确保整个图像能在内存中处理。与之不同，我们在测试时也将图像分解为重叠的固定尺寸块，并在采样过程中将它们融合。

为什么分块

扩散模型通常要求固定输入尺寸（如256x256）。处理任意大图时，直接缩放会丢失细节，而分块处理可以保持局部细节，并允许在有限显存下处理大图。

基于块的图像修复的总体思路是：对从图像中提取的局部块进行操作，并以最优方式合并结果。这种方法迄今为止的一个重要缺点是，生成的图像可能包含来自独立修复的中间结果的合并伪影，这在传统修复方法中已被广泛研究。我们将通过引导反向采样过程，使相邻块之间趋于平滑，从而解决这个问题，避免出现边缘伪影。

块间不一致

如果每个块独立进行扩散逆过程（从噪声生成清晰块），相邻块的边界处会出现不连续、撕裂或模糊，因为模型不知道块之间的关系。

我们将任意大小的未知真实图像定义为 $X_0$，将天气降质的观测图像定义为 $\tilde{X}$，并将 $P_i$ 定义为一个与 $X_0$ 和 $\tilde{X}$ 维度相同的二值掩码矩阵，指示图像中第 $i$ 个 $p\times p$ 块的位置。在训练方法中我们学习条件反向过程：

$$\begin{array}{}\text{(20)}& p_{\theta }({\mathbf{x}}_{0:T}^{(i)}\mid {\tilde{\mathbf{x}}}^{(i)})=p\left({\mathbf{x}}_T^{(i)}\right)\prod _{t=1}^T{p}_{\theta }({\mathbf{x}}_{t-1}^{(i)}\mid {\mathbf{x}}_t^{(i)},{\tilde{\mathbf{x}}}^{(i)})\end{array}$$

其中 $X_0^{(i)}=\mathrm{Crop}(P_i\circ X_0)$ 和 ${\tilde{X}}^{(i)}=\mathrm{Crop}(P_i\circ \tilde{X})$ 表示来自训练集图像对 $(X_0,\tilde{X})$ 的 $p\times p$ 图像块，$\mathrm{Crop}(\cdot )$ 操作从 $P_i$ 指示的位置提取图像块。训练期间，我们在图像尺寸的完整范围内随机采样 $p\times p$ 图像块的位置。

我们通过从网格状排列的解析方案中提取所有重叠的 $p\times p$ 图像块，来分解任意大小的图像 $X$。我们在整个图像上考虑一个网格状排列，其中每个网格单元包含 $r\times r$ 像素，并通过在水平和垂直维度上以步长 $r$ 移动此网格来提取所有 $p\times p$ 图像块。我们将 $D$ 定义为提取的图像块总数，定义了一个重叠图像块位置的字典。

当基于相邻的重叠图像块执行条件反向采样时，对于重叠的网格单元将获得不同的修复估计。我们通过在任意给定的去噪时间步 $t$，对重叠图像块区域中的每个像素基于其平均估计噪声来执行反向采样，从而缓解此问题。我们的方法有效地引导了反向采样过程，以确保在所有相关的相邻图像块之间具有更高的保真度。更具体地说，在采样的每个时间步 $t$：

（1）我们使用 ${ϵ}_{\theta }(x_t^{(d)},{\tilde{x}}^{(d)},t)$ 估计所有重叠图像块位置 $d\in \{1,\dots ,D\}$ 的加性噪声；

（2）将这些重叠的噪声估计累积到与整个图像尺寸相同的矩阵 ${\hat{\mathrm{\Omega }}}_t$ 中它们各自的位置；

（3）根据每个像素接收到的估计数量对 ${\hat{\mathrm{\Omega }}}_t$ 进行归一化；

（4）使用平滑后的全图噪声估计 ${\hat{\mathrm{\Omega }}}_t$ 执行一次隐式采样更新。

算法2 基于块的扩散式图像恢复

输入: 天气退化的图像$\tilde{X}$,条件扩散模型${ϵ}_{\theta }({\mathbf{x}}_t,\tilde{\mathbf{x}},t)$,隐式采样步骤$S$, 包含$D$个重叠图像块位置的字典.

1. ${\mathbf{X}}_t\sim \mathcal{N}(\mathbf{0},\mathbf{I})$
2. for $i=S,\dots ,1$ do
3.     $t=(i-1)\cdot T/S+1$
4.     $t_{\text{next}}=(i-2)\cdot T/S+1$ if $i>1$ else $0$
5.     ${\hat{\mathbf{\Omega }}}_t=\mathbf{0}$ and $\mathbf{M}=\mathbf{0}$
6.     for $d=1,\dots ,D$ do
7.         ${\mathbf{x}}_t^{(d)}=\text{Crop}({\mathbf{P}}_d\circ {\mathbf{X}}_t)$ and ${\tilde{\mathbf{x}}}^{(d)}=\text{Crop}({\mathbf{P}}_d\circ \tilde{\mathbf{X}})$
8.         ${\hat{\mathbf{\Omega }}}_t={\hat{\mathbf{\Omega }}}_t+{\mathbf{P}}_d\cdot {ϵ}_{\theta }({\mathbf{x}}_t^{(d)},{\tilde{\mathbf{x}}}^{(d)},t)$
9.         $\mathbf{M}=\mathbf{M}+{\mathbf{P}}_d$
10.     end for
11.     ${\hat{\mathrm{\Omega }}}_t={\hat{\mathrm{\Omega }}}_t\oslash \mathbf{M}$     // $\oslash$ 表示逐元素除法
12.     ${\mathbf{X}}_t\leftarrow \sqrt{{\overline{\alpha }}_{t_{\text{next }}}}\left(\frac{{\mathbf{X}}_t-\sqrt{1-{\overline{\alpha }}_t}\cdot {\hat{\mathrm{\Omega }}}_t}{\sqrt{{\overline{\alpha }}_t}}\right)+\sqrt{1-{\overline{\alpha }}_{t_{\text{next }}}}\cdot {\hat{\mathrm{\Omega }}}_t$
13. end for
14. return ${\mathbf{X}}_t$

我们的方法不同于在采样后平均重叠最终重建结果的简单基线方法。如果在采样后应用，这种方法会破坏局部图像块分布与学习到的后验分布的保真度。然而，与我们的基于重叠图像块的引导采样原理不同但精神相似，最近也有一些成功的图像编辑方法基于在潜在空间中引导反向过程，以从学习到的密度的一个压缩子空间中采样。

需要注意的是，较小的 $r$ 会增加图像块之间的重叠，从而提高平滑度，但同时也会增加计算负担。我们使用 $p=64$ 或 $128$ 像素用于 $P_i$，以及 $r=16$ 像素。在处理之前，我们将整个图像的尺寸调整为 $16$ 的倍数。

我们提出的基于块的条件扩散建模方法是与任务无关的，并且当在训练时观察到来自多种天气降质混合的示例图像对时，可以进一步扩展到同时处理多种天气破坏

意义

• 处理任意大小的高分辨率图像。
• 保持全局一致性和局部细节。
• 适用于多种图像修复任务（去雨、去雪、去雾等），且无需在测试时指定任务类型（模型从数据中自动学习混合分布）。