【Siggraph Asia 2023】低光增强Diffusion-Low-Light-main（引入diffusion与DWT） -- part1论文精读

【Siggraph Asia 2023】本文参考论文[Siggraph Asia 2023]Low-light Image Enhancement with Wavelet-based Diffusion Models

论文地址：arxiv
源码地址：github

---

💡 DiffLL 方法概述

该论文介绍了一种名为 DiffLL 的低光图像增强框架，它独特地将扩散模型 (diffusion models) 的强大能力与小波变换 (wavelet transformation) 相结合。其核心思想是在小波域中执行计算密集型的扩散过程，特别是在图像的低频（平均）系数上。这显著减小了空间维度，从而在不牺牲关键图像信息的前提下，实现了更快的推理速度和更低的资源消耗。该框架还包括一种新颖的扩散模型训练策略和一个专门用于恢复高频细节的模块。

---

🌊 离散小波变换 (Discrete Wavelet Transformation - DWT)

该方法首先使用二维离散小波变换（2D-DWT），特别是哈尔小波 (Haar wavelets)，将输入的低光图像 ($I_{low}$) 转换到小波域。

1. 分解 (Decomposition)：单次二维离散小波变换将图像分解为四个子带：

   • $A_{low}^1$：平均系数（低频），代表了全局图像信息的降采样版本。
   • $V_{low}^1$：垂直方向的高频信息。
   • $H_{low}^1$：水平方向的高频信息。
   • $D_{low}^1$：对角线方向的高频信息。
     这些子带中的每一个都具有 $\frac{H}{2}\times \frac{W}{2}\times c$ 的维度。

2. 聚焦平均系数 (Focus on Average Coefficient)：论文强调平均系数 ($A_{low}$) 包含了最关键的全局信息。实验表明，改变这个系数会极大地改变图像内容，而对高频系数的改变对整体场景的影响较小。因此，在小波域中的主要目标是恢复这个平均系数，使其与其在正常光照下的对应部分相匹配。

3. 多级分解 (Multi-Level Decomposition)：为了进一步提高效率，该方法对前一级别的平均系数应用 $K$ 次二维离散小波变换。因此，$A_{low}^k$ 被分解为 { A l o w k + 1 , V l o w k + 1 , H l o w k + 1 , D l o w k + 1 } \{A_{low}^{k+1}, V_{low}^{k+1}, H_{low}^{k+1}, D_{low}^{k+1}\} {Alowk+1​,Vlowk+1​,Hlowk+1​,Dlowk+1​}。然后，扩散操作在最终的、最小的平均系数 $A_{low}^K$ 上执行。这导致了扩散模型处理的空间维度减少了 $4^K$ 倍。原始输入图像表示为 $A_{low}^0$。

4. 重建 (Reconstruction)：处理之后，恢复的平均系数和（单独恢复的）高频系数被用于通过二维逆离散小波变换 (2D-IDWT) 重建前一尺度的图像。重复此过程，直到获得全分辨率图像 ${\tilde{A}}_{low}^0$（最终增强图像 ${\hat{I}}_{low}$）。

---

🔄 基于小波的条件扩散模型 (Wavelet-based Conditional Diffusion Models - WCDM)

增强过程的核心在于基于小波的条件扩散模型 (WCDM)，它作用于通过离散小波变换获得的平均系数 $A_{low}^K$。

常规条件扩散模型：预备知识

扩散模型通常有两个阶段：

1. 前向扩散过程 (Forward Diffusion Process)：此过程在 $T$ 个时间步内逐渐向输入数据样本 $x_0$ 添加高斯噪声，产生一系列噪声越来越大的样本 $x_1,\dots ,x_T$。这种转变定义为：
   $q(x_t|x_{t-1})=\mathcal{N}(x_t;\sqrt{1-{\beta }_t}{x}_{t-1},{\beta }_{t}I)$，其中 ${\beta }_t$ 是预定义的方差调度。

2. 去噪（逆向）过程 (Denoising (Reverse) Process)：此过程学习逆转加噪过程。从随机高斯噪声 ${\hat{x}}_T$ 开始，它迭代地去噪以产生清晰的样本 ${\hat{x}}_0$。对于条件模型，此过程由条件输入 $\tilde{x}$ 指导：
   $p_{\theta }({\hat{x}}_{t-1}|{\hat{x}}_t,\tilde{x})=\mathcal{N}({\hat{x}}_{t-1};{\mu }_{\theta }({\hat{x}}_t,\tilde{x},t),{\sigma }_t^{2}I)$。
   模型 ${ϵ}_{\theta }({\hat{x}}_t,\tilde{x},t)$ 被训练来预测在时间步 $t$ 添加的噪声。典型的目标函数是：
   ${\mathcal{L}}_{diff}=E_{x_0,t,{ϵ}_t\sim N(0,I)}[||{ϵ}_t-{ϵ}_{\theta }(x_t,\tilde{x},t)|{|}^2]$。

WCDM：应对挑战

传统的扩散模型在图像恢复方面面临两个主要挑战：

• 高计算成本和延迟：由于需要大量的时间步长 ($T$)。
• 内容不一致：初始噪声 ${\hat{x}}_T$ 的随机性可能导致输出多样化，即使有条件约束，这对于需要确定性和忠实重建的恢复任务来说也是不可取的。

WCDM 通过以下方式解决这些问题：

1. 效率 (Efficiency)：通过在小得多的平均系数 $A_{low}^K$ （在训练中表示为条件 $\tilde{x}$，而 $x_0$ 是目标 $A_{high}^K$）上执行扩散，WCDM 显著减少了计算负载和推理时间。

2. 新颖的训练策略以保证一致性 (Novel Training Strategy for Consistency)：这是一项关键创新。

   • WCDM 的训练阶段同时包括前向扩散过程（用于训练噪声估计器 ${ϵ}_{\theta }$，如公式 8）和去噪过程（采样）。
   • 在训练的去噪部分，模型从随机噪声 ${\hat{x}}_T\sim \mathcal{N}(0,I)$ 开始，并在低光平均系数 $A_{low}^K$ （表示为 $\tilde{x}$）的条件下，生成恢复的系数 ${\hat{A}}_{low}^K$。
   • 关键在于，通过在损失函数中添加这个生成的 ${\hat{A}}_{low}^K$ 与真实正常光照平均系数 $A_{high}^K$ （表示为 $x_0$）之间的 L2 距离项来强制实现内容一致性。这仅在训练期间进行，因为 $A_{high}^K$ 在推理时不可用。
   • 修改后的目标函数变为：
     ${\mathcal{L}}_{diff}=E_{x_0,t,{ϵ}_t\sim N(0,I)}[||{ϵ}_t-{ϵ}_{\theta }(x_t,\tilde{x},t)|{|}^2]+||{\hat{A}}_{low}^K-A_{high}^K|{|}^2$
     这使得模型能够在推理过程中学习执行稳定的采样，减少随机性并避免多样化或混乱的输出。训练过程在算法 1 中有详细描述。

3. 推理 (Inference)：在推理过程中（算法 2），仅使用学习到的去噪过程。它从随机噪声 ${\hat{x}}_T$ 开始，并使用训练好的 ${ϵ}_{\theta }$ 网络，在输入的低光平均系数 $A_{low}^K$ 的条件下，迭代地产生增强的 ${\hat{A}}_{low}^K$。

4. 相对于潜在扩散（基于VAE）的优势 (Advantages over Latent Diffusion (VAE-based))：

   • 小波变换是无损的，不像 VAE 编码器可能会丢失信息。
   • 小波变换是一种线性操作，不会增加模型参数，这与 VAE 不同。
   • 它避免了如果将在正常光照图像上训练的 VAE 用于低光数据而不进行重新训练时可能出现的域偏移问题。

---

✨ 高频恢复模块 (High-Frequency Restoration Module - HFRM)

虽然 WCDM 通过平均系数 $A_{low}^K$ 恢复全局光照，但来自每个小波分解级别 $k\in [1,K]$ 的高频系数（$V_{low}^k,H_{low}^k,D_{low}^k$）包含稀疏的局部细节，这些细节也需要恢复。高频恢复模块 (HFRM) 就是为此目的而设计的。

$HFR{M}_k$（针对每个级别 $k$）的架构如下（图 5）：

1. 特征提取 (Feature Extraction)：三个初始的深度可分离卷积 (depth-wise separable convolutions) 高效地处理输入的高频系数（$V_{low}^k,H_{low}^k,D_{low}^k$）。
2. 信息融合（交叉注意力）(Information Fusion (Cross-Attention))：使用两个交叉注意力层来利用垂直 ($V_{low}^k$) 和水平 ($H_{low}^k$) 系数中的信息来补充和增强对角线 ($D_{low}^k$) 系数中的细节。这是一个关键方面，因为对角线细节通常对精细纹理至关重要。
3. 渐进式空洞残差块 (Progressive Dilation Resblock)：受 [Hai 等人, 2022] 的启发，该块使用一系列具有不同空洞率（例如 d = { 1 , 2 , 3 , 2 , 1 } d=\{1,2,3,2,1\} d={1,2,3,2,1}）的 $3\times 3$ 卷积。
   • 初始和最终的卷积提取局部信息。
   • 中间的空洞卷积扩大了感受野以整合远程信息。
   • 逐渐增加然后减少空洞率有助于避免网格伪影。
   • 存在跳跃连接。
4. 输出卷积 (Output Convolutions)：三个最终的深度可分离卷积减少通道维度，以产生恢复的高频系数 ${\hat{V}}_{low}^k,{\hat{H}}_{low}^k,{\hat{D}}_{low}^k$。

恢复的平均系数 ${\hat{A}}_{low}^k$ （如果 $k=K$，则来自 WCDM；如果 $k<K$，则来自级别 $k+1$ 的 IDWT）和来自 $HFR{M}_k$ 的恢复的高频系数 { V ^ l o w k , H ^ l o w k , D ^ l o w k } \{\hat{V}_{low}^{k}, \hat{H}_{low}^{k}, \hat{D}_{low}^{k}\} {V^lowk​,H^lowk​,D^lowk​} 然后使用二维逆离散小波变换组合，以获得前一个更大尺度的恢复平均系数 ${\hat{A}}_{low}^{k-1}$：
A ^ l o w k − 1 = 2 D − I D W T ( { A ^ l o w k , V ^ l o w k , H ^ l o w k , D ^ l o w k } ) \hat{A}_{low}^{k-1} = 2D-IDWT(\{\hat{A}_{low}^{k}, \hat{V}_{low}^{k}, \hat{H}_{low}^{k}, \hat{D}_{low}^{k}\}) A^lowk−1​=2D−IDWT({A^lowk​,V^lowk​,H^lowk​,D^lowk​})。

---

📉 网络训练和损失函数

整个网络使用多种损失函数进行训练：

1. 扩散损失 (${\mathcal{L}}_{diff}$)：如公式 9 所定义，该损失优化 WCDM 以实现平均系数的噪声预测和内容一致性。
2. 细节保持损失 (${\mathcal{L}}_{detail}$)：此损失应用于 HFRM 的输出，以确保高频细节的准确重建。它结合了均方误差 (MSE) 损失和全变分 (TV) 损失：
   L d e t a i l = λ 1 ∑ k = 1 K ∣ ∣ { V ^ l o w k , H ^ l o w k , D ^ l o w k } − { V h i g h k , H h i g h k , D h i g h k } ∣ ∣ 2 + λ 2 ∑ k = 1 K T V ( { V ^ l o w k , H ^ l o w k , D ^ l o w k } ) \mathcal{L}_{detail} = \lambda_1 \sum_{k=1}^{K} ||\{\hat{V}_{low}^{k}, \hat{H}_{low}^{k}, \hat{D}_{low}^{k}\} - \{V_{high}^{k}, H_{high}^{k}, D_{high}^{k}\}||^2 + \lambda_2 \sum_{k=1}^{K} TV(\{\hat{V}_{low}^{k}, \hat{H}_{low}^{k}, \hat{D}_{low}^{k}\}) Ldetail​=λ1​∑k=1K​∣∣{V^lowk​,H^lowk​,D^lowk​}−{Vhighk​,Hhighk​,Dhighk​}∣∣2+λ2​∑k=1K​TV({V^lowk​,H^lowk​,D^lowk​})。
   权重 ${\lambda }_1$ 和 ${\lambda }_2$ 分别设置为 0.1 和 0.01。
3. 内容损失 (${\mathcal{L}}_{content}$)：应用于最终恢复的图像 ${\hat{I}}_{low}$（即 ${\tilde{A}}_{low}^0$）和参考正常光照图像 $I_{high}$。它结合了 L1 损失和 SSIM 损失：
   ${\mathcal{L}}_{content}=||{\hat{I}}_{low}-I_{high}|{|}_1+(1-SSIM({\hat{I}}_{low},I_{high}))$。

总损失 (Total loss) 是这些组成部分的总和：
${\mathcal{L}}_{total}={\mathcal{L}}_{diff}+{\mathcal{L}}_{detail}+{\mathcal{L}}_{content}$。

这种全面的方法旨在确保 DiffLL 模型产生的增强图像不仅光照良好、细节丰富，而且在感知上令人愉悦且与输入内容保持一致。