Invertible Diffusion Model

PKU TPAMI-2025

网络架构

主干网络

IDM 采用了 SD1.5 的噪声估计器，复用其在 LAION 数据集上预训练的网络权重，并对 U-Net 结构进行剪枝 (we further simplify the U-Net by removing its time embeddings, cross-attention layers, and the final three scales to balance efficiency and performance.) 实际上，这个操作将原始参数规模 (~860M) 降至 116.6M，这使得在 3090 上训练模型成为可能。

结构图如下。

这部分可以结合代码理解

del unet.time_embedding, unet.mid_block
unet.down_blocks = unet.down_blocks[:-2]
unet.down_blocks[-1].downsamplers = None
unet.up_blocks = unet.up_blocks[2:]
...

其代码源于 diffusers 库，其中 Stable Diffusion v1.5 是潜扩散模型 (LDM) 的开源实现。为了进行剪枝并加入自定义模块，需采用动态注册和替换，工程实现上存在一定难度。另一类方法 (如 DMP-DUN) 选择从开源扩散模型中自行提取并封装扩散模型代码段，提供了较高的灵活性和可扩展性。两类方案各有优势，实践中灵活选用。

时间步 $T$ 的选取

DDIM 采样过程一般应说明起始步数和步长，如 $T=8$ 可以是从预训练的噪声表的第 96 步开始采样，每隔 12 步采样。(p.s. 一般从噪声表较低步数开始采样，因为噪声分量更低)。比较奇怪的是，IDM 淡化了这一概念，且直接移除了时间步嵌入，这样 $T$ 就变成了一个纯纯的超参数。

实验部分对去噪步数分析了参数敏感性，如下图所示。

可逆布线

IDM 引入可逆神经网络思想，例如：

Notably, with our two level invertible models, the peak memory use per GPU is suppressed to 14.64 GB, making it manageable not only for A100 but also for other GPUs like 3090 and 4090 (24GB).

其出发点是，训练过程保存了太多中间结果，这导致显存需求大，尤其对于扩散模型，显著提升硬件需求。IDM 从可逆神经网络 (INN) 获得灵感，对于模型中的所有级联的保形网络模块，将上一级输入短路到下一级输出，参考上图 (a)，其过程可以描述为：

\begin{aligned} \text{Forward:} \quad & \begin{cases} \hat{\mathbf{x}}_{t-1} = u_t \mathcal{F}_t(\hat{\mathbf{x}}_t; \mathbf{y}, \mathbf{A}) + v_t \mathbf{h}_t \\ \mathbf{h}_{t-1} = \hat{\mathbf{x}}_t \end{cases} \\[10pt] \text{Inverse:} \quad & \begin{cases} \hat{\mathbf{x}}_t = \mathbf{h}_{t-1} \\ \mathbf{h}_t = \frac{\hat{\mathbf{x}}_{t-1} - u_t \mathcal{F}_t(\hat{\mathbf{x}}_t; \mathbf{y}, \mathbf{A})}{v_t} \end{cases} \end{aligned}

其中 $u_t, v_t$ 是可学习的参数， $u_t + v_t = 1$ 。

这样之后在反向传播过程中可以使用 $\hat{\mathbf{x}_t}$ 恢复 $\hat{\mathbf{x}}_{t-1}$ ，从而优化显存开销。

Injector (物理信息嵌入)

在 U-Net 的 Residual Block 和 Attention Block 之后，IDM 插入了自设计的 Injector (见上图 (b))，将采样矩阵和测量值 $\{\mathbf{A, y}\}$ 作为条件信息嵌入网络，以引导生成方向，同时提高重建质量。文中提到，直接融合测量先验 $\{\mathbf{A}, \mathbf{y}\}$ 的效果并不好，因此通过卷积先进行浅层特征提取，后通过 concat 拼接 $\hat{\mathbf{x}}$ 以实施嵌入。

这里实际上形成了一个特征域 (Res, Attn) 与图像域 (Injector) 交替的结构，这使得重建信号可以充分融合测量先验与生成式先验。另一些工作指出，这种双域跨域重建可能潜在地滤除了噪声信号，进一步提升重建质量。这部分不展开介绍。

实验复现

原文工作量较大，在消融实验部分，对比了两个预训练模型 SD1.5 和 UIDM，且针对 Y 通道重建和 RGB 三通道重建分别训练模型。受限于实验条件，复现时只训练了灰度模型。

Setup

单卡 3090 训练了 11 天 (中途 screen 挂了一次)，实验配置未 $\gamma=0.1$ ，batch_size 按照默认 8，显存消耗约 20G 左右，每个 epoch 耗时 5h 左右。学习率阶梯衰减，共训练 50 epoch，Adam 优化器采用默认配置 $\beta_!=0.9, \beta_2=0.999$ ，训练数据集为 WED，验证集为 Set11。