Invertible Diffusion Models for Compressed Sensing

TL;DR

This paper introduces Invertible Diffusion Models (IDM), a novel end-to-end diffusion-based approach for image compressed sensing that improves reconstruction quality, reduces memory usage, and accelerates inference compared to existing methods.

Contribution

The paper proposes a new invertible diffusion model framework for compressed sensing that enables efficient end-to-end training and significantly enhances reconstruction performance.

Findings

IDM outperforms state-of-the-art CS networks by up to 2.64dB PSNR.

IDM achieves up to 10.09dB PSNR gain over DDNM.

IDM reduces GPU memory usage by up to 93.8%.

Abstract

While deep neural networks (NN) significantly advance image compressed sensing (CS) by improving reconstruction quality, the necessity of training current CS NNs from scratch constrains their effectiveness and hampers rapid deployment. Although recent methods utilize pre-trained diffusion models for image reconstruction, they struggle with slow inference and restricted adaptability to CS. To tackle these challenges, this paper proposes Invertible Diffusion Models (IDM), a novel efficient, end-to-end diffusion-based CS method. IDM repurposes a large-scale diffusion sampling process as a reconstruction model, and fine-tunes it end-to-end to recover original images directly from CS measurements, moving beyond the traditional paradigm of one-step noise estimation learning. To enable such memory-intensive end-to-end fine-tuning, we propose a novel two-level invertible design to transform both (1) multi-step sampling process and (2) noise estimation U-Net in each step into invertible networks. As a result, most intermediate features are cleared during training to reduce up to 93.8% GPU memory. In addition, we develop a set of lightweight modules to inject measurements into noise estimator to further facilitate reconstruction. Experiments demonstrate that IDM outperforms existing state-of-the-art CS networks by up to 2.64dB in PSNR. Compared to the recent diffusion-based approach DDNM, our IDM achieves up to 10.09dB PSNR gain and 14.54 times faster inference. Code is available at https://github.com/Guaishou74851/IDM.