Invertible Diffusion Models Accelerate Image Reconstruction – TPAMI 2025

Paper Information

Title: Invertible Diffusion Models for Compressed Sensing

Authors: Bin Chen, Zhenyu Zhang, Weiqi Li, Chen Zhao, Jiwen Yu, Shijie Zhao, Jie Chen, Jian Zhang

Affiliations: Peking University, KAUST, ByteDance

Venue: IEEE Transactions on Pattern Analysis and Machine Intelligence (TPAMI), 2025

Links: https://ieeexplore.ieee.org/document/10874182 | https://arxiv.org/abs/2403.17006 | https://github.com/Guaishou74851/IDM

Task Background

Diffusion models are trained to estimate noise from a noisy image using a UNet‑style network, but image reconstruction requires recovering the original image from degraded measurements. This creates a mismatch between the training objective (noise estimation) and the deployment goal (reconstruction). Additionally, diffusion inference typically needs hundreds of iterative denoising steps, leading to slow runtime and high computational cost.

Challenge 1: Gap between “noise estimation” (training) and “image reconstruction” (deployment).

Challenge 2: Inference latency caused by many sampling steps.

To close the gap and accelerate inference, the authors propose Invertible Diffusion Models (IDM), which combine an end‑to‑end reconstruction‑oriented training framework with a reversible network architecture.

Main Contributions

1. End‑to‑end diffusion sampling as a reconstruction network

Instead of training the diffusion model solely for noise estimation, the entire sampling process is re‑defined as a single reconstruction network. All parameters are optimized directly for the reconstruction loss (e.g., L2 between reconstructed and ground‑truth images). Using Stable Diffusion pretrained weights, IDM achieves:

≈2 dB higher PSNR on compressed‑sensing benchmarks.

Sampling steps reduced from 100 to 3.

Inference speedup of roughly 15×.

2. Dual‑layer invertible network for memory‑efficient training

Training large diffusion models end‑to‑end requires storing intermediate feature maps for every diffusion step, which quickly exceeds GPU memory. IDM introduces a dual‑layer reversible architecture where each layer’s output can be recomputed from its input, eliminating the need to cache full activations. The reversible design is applied to both the outer diffusion steps and the inner noise‑estimation UNet. As a result, IDM can be trained on a single 1080Ti (≈11 GB VRAM) without sacrificing reconstruction quality.

Experimental Results

1. Compressed‑sensing image reconstruction

IDM is compared with state‑of‑the‑art end‑to‑end networks and diffusion‑based methods (e.g., DDNM). Across four metrics—PSNR, SSIM, FID, and LPIPS—IDM consistently outperforms baselines. The quantitative table shows PSNR gains of 1.8–2.2 dB and SSIM improvements of 0.02–0.04.

2. Image inpainting and medical imaging

On a 90 % mask inpainting benchmark, IDM restores complex structures such as windows that DDNM fails to reconstruct. The method also transfers to MRI and CT reconstruction, delivering comparable or superior visual fidelity and quantitative scores (e.g., PSNR ≈ 33 dB on the MRI test set).

3. Computational cost and inference time

Traditional diffusion‑based reconstruction of a 256×256 image requires ≈9 s (≈100 steps). IDM reduces inference to 0.63 s (3 steps) while also lowering peak GPU memory usage from >12 GB to <8 GB. Compared with DDNM, IDM shows a >14× speedup in inference and a ≈2× reduction in training memory footprint, without degrading reconstruction quality.