IDM

PKU TPAMI 2025

架构

主干网络

噪声估计器基于 SD1.5，并对其 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.)

论文给出的结构图。对原模型进行了不少剪枝。

这部分可以结合代码理解

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:]
...

$T$ 的选取

实验部分对去噪步数分析了参数敏感性，结论是 $T$ 取 3 时综合收益最大。

可逆布线

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).

出发点是训练过程保存太多中间结果导致显存需求大。从可逆神经网络 (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}$ ，从而优化 memory，实际测试确实可以在 3090 完成训练。

Injector

在 U-Net 的 Residual Block 和 Attention Block 之后插入 Injector (见上图 (b))，把采样矩阵和测量值 $\{\mathbf{A, y}\}$ 作为特征融合进网络。提到，直接融合测量值的效果不好，因此通过卷积先进行特征提取，然后通过 pixel unshuffle 的方式把特征图和 $\hat{\mathbf{x}}$ 融合并保持形状。