Invertible Diffusion Models for Compressed Sensing

2024-03-25

arXiv.org

AI Summary

To address the poor reconstruction quality, slow inference, and weak generalization of existing deep learning-based compressive sensing (CS) methods for images, this paper proposes the Invertible Diffusion Model (IDM). IDM reformulates the entire pre-trained diffusion sampling process as an end-to-end differentiable reconstruction network, enabling direct image recovery from linear measurements. Its core innovations are: (1) a two-level fully invertible architecture—ensuring invertibility both across multiple sampling steps and within each step’s noise-estimation U-Net; and (2) a lightweight measurement injection module that breaks the conventional single-step denoising paradigm. By integrating invertible neural networks, end-to-end fine-tuning, gradient checkpointing, and customized measurement embedding, IDM achieves a 2.64 dB PSNR gain over prior state-of-the-art methods, a 10.09 dB improvement over DDNM, 14.54× faster inference, and up to 93.8% reduction in GPU memory consumption.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Learning on the Edge & Model Compression
• Data Mining & Knowledge Management: Data Compression

Application Category

• Security and Privacy: Large-scale security measurements
• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation

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.

Problem

• Deep Learning
• Image Compression
• Efficiency

Innovation

• Reversible Diffusion Model
• Image Restoration
• Efficiency Optimization