Invertible Diffusion Models for Compressed Sensing

要約

ディープニューラルネットワーク(NN)は再構成品質を向上させることで画像圧縮センシング(CS)を大幅に進歩させるが、現在のCS NNをゼロから訓練する必要があるため、その有効性が制約され、迅速な展開が妨げられる。最近の手法では、画像再構成のために事前に訓練された拡散モデルを利用しているが、推論に時間がかかり、CSへの適応性が制限されるという問題を抱えている。これらの課題に取り組むため、本論文では、新しい効率的でエンドツーエンドの拡散ベースのCS手法であるInvertible Diffusion Models（IDM）を提案する。IDMは、大規模拡散サンプリング過程を再構成モデルとして再利用し、エンドツーエンドで微調整することで、CS測定値から直接原画像を復元し、従来のワンステップノイズ推定学習のパラダイムを超える。このようなメモリ集約的なエンドツーエンドの微調整を可能にするために、我々は、(1)マルチステップサンプリングプロセスと、(2)各ステップにおけるノイズ推定U-Netの両方を可逆ネットワークに変換する、新しい2レベル可逆設計を提案する。その結果、ほとんどの中間特徴量は学習中にクリアされ、最大93.8%のGPUメモリを削減できる。さらに、我々は、ノイズ推定器に測定値を注入するための一連の軽量モジュールを開発し、再構成をさらに容易にする。実験により、IDMは既存の最先端CSネットワークをPSNRで最大2.64dB上回ることが実証された。最近の拡散ベースのアプローチであるDDNMと比較して、我々のIDMは最大10.09dBのPSNR向上と14.54倍の高速推論を達成している。コードはhttps://github.com/Guaishou74851/IDM。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL