UGoDIT: Unsupervised Group Deep Image Prior Via Transferable Weights

TL;DR

This paper proposes UGoDIT, an Unsupervised Group DIP via Transferable weights, designed for the low-data regime where only a very small number of sub-sampled measurement vectors are available during training, and achieves performance competitive with SOTA DM-based and supervised approaches, despite not requiring large amounts of clean training data.

Abstract

Recent advances in data-centric deep generative models have led to significant progress in solving inverse imaging problems. However, these models (e.g., diffusion models (DMs)) typically require large amounts of fully sampled (clean) training data, which is often impractical in medical and scientific settings such as dynamic imaging. On the other hand, training-data-free approaches like the Deep Image Prior (DIP) do not require clean ground-truth images but suffer from noise overfitting and can be computationally expensive as the network parameters need to be optimized for each measurement set independently. Moreover, DIP-based methods often overlook the potential of learning a prior using a small number of sub-sampled measurements (or degraded images) available during training. In this paper, we propose UGoDIT, an Unsupervised Group DIP via Transferable weights, designed for the low-data regime where only a very small number, M, of sub-sampled measurement vectors are available during training. Our method learns a set of transferable weights by optimizing a shared encoder and M disentangled decoders. At test time, we reconstruct the unseen degraded image using a DIP network, where part of the parameters are fixed to the learned weights, while the remaining are optimized to enforce measurement consistency. We evaluate UGoDIT on both medical (multi-coil MRI) and natural (super resolution and non-linear deblurring) image recovery tasks under various settings. Compared to recent standalone DIP methods, UGoDIT provides accelerated convergence and notable improvement in reconstruction quality. Furthermore, our method achieves performance competitive with SOTA DM-based and supervised approaches, despite not requiring large amounts of clean training data.

Figures (19)

• Figure 1: Illustrative block diagram UGoDIT. During training (left), we learn the shared encoder by optimizing over $\phi, \psi_1,\dots, \psi_M$. At inference (right), given some unseen measurement vector $\mathbf{y}$ and its forward operator $\mathbf{A}$, we reconstruct the image $\hat{\mathbf{x}}$ by optimizing over $\psi$.
• Figure 2: SR averaged PSNR curves of SR ($6$ for training and $20$ for testing) of UGoDIT vs. the shared-decoder case (i.e., \ref{['eqn: shared decoder training updates']}).
• Figure 3: PSNR curves of UGoDIT-$M$ and baselines for the tasks of MRI (left), SR (middle), and NDB (right) averaged over 20 test images. Iterations (x-axis) correspond to $NK$.
• Figure 4: Reconstructed/recovered images using UGoDIT (last column) and baselines (columns 3 to 5). As observed, UGoDIT return comparable PSNR results with data-intensive methods all without the need of large amount of clean images (or fully sampled measurement vectors). See Appendix \ref{['sec: appen more visual']} for more visualizations.
• Figure 5: Average PSNR of $20$ MRI brain test scans using a knee-trained (with $M=6$) UGoDIT-OOD (blue) vs. the case where at test-time, we use the parameters of the pre-trained encoder to only initialize the test-time encoder (red). Then, we optimize over both the encoder and decoder. Running aSeqDIP independently is also included (black). Iterations in the x-axis correspond to $NK=30000$.