[Paper Review] Compressed Sensing with Deep Image Prior and Learned Regularization

Dave Van Veen, Ajil Jalal+4 more|arXiv (Cornell University)|Jun 17, 2018

Sparse and Compressive Sensing Techniques83 references167 citations

TL;DR

The paper introduces CS-DIP, a compressed sensing method using an untrained deep image prior with a learned regularization term to improve reconstruction without large pretraining datasets.

ABSTRACT

We propose a novel method for compressed sensing recovery using untrained deep generative models. Our method is based on the recently proposed Deep Image Prior (DIP), wherein the convolutional weights of the network are optimized to match the observed measurements. We show that this approach can be applied to solve any differentiable linear inverse problem, outperforming previous unlearned methods. Unlike various learned approaches based on generative models, our method does not require pre-training over large datasets. We further introduce a novel learned regularization technique, which incorporates prior information on the network weights. This reduces reconstruction error, especially for noisy measurements. Finally, we prove that, using the DIP optimization approach, moderately overparameterized single-layer networks can perfectly fit any signal despite the non-convex nature of the fitting problem. This theoretical result provides justification for early stopping.

Motivation & Objective

• Explore the use of untrained deep generative networks (DIP) for solving differentiable linear inverse problems like compressed sensing.
• Introduce a learned regularization term that encodes prior information about network weights to reduce reconstruction error.
• Prove theoretical convergence/fit guarantees for overparameterized single-layer networks under DIP optimization.
• Demonstrate empirical improvements over state-of-the-art unlearned methods on medical and natural images.

Proposed method

• Optimize the weights w of an untrained DCGAN generator G(z;w) to minimize ||y - A G(z;w)||^2 subject to regularization.
• Incorporate total variation regularization TV(G(z;w)) with weight λ_T into the objective.
• Introduce a learned regularization term LR(w) = (w - μ)^T Σ^{-1} (w - μ) with prior parameters μ, Σ.
• Learn μ, Σ layer-wise from a small set of similar measurements by solving a MAP problem with a Gaussian prior on w.
• Use early stopping to mitigate overfitting and justify it with theoretical results for a wide single-layer network.
• Compare CS-DIP to BM3D-AMP, TVAL3, and Lasso on Gaussian and Fourier measurement models.

Experimental results

Research questions

• RQ1 Can an untrained deep image prior solve compressed sensing problems with differentiable forward operators?
• RQ2 Does adding a learned weight regularization improve reconstruction quality under noise and strong undersampling?
• RQ3 How does CS-DIP performance compare to learned pre-trained generative priors and unlearned baselines across datasets and measurement models?
• RQ4 What theoretical guarantees explain the need for early stopping in DIP-based reconstructions?

Key findings

• CS-DIP outperforms state-of-the-art unlearned methods on MNIST and chest X-ray datasets under various measurement regimes.
• Learned regularization (LR) provides notable gains, especially with noise and few measurements, while vanilla L2 regularization does not.
• Pre-trained DCGAN-based methods outperform CS-DIP at very low measurements, but CS-DIP scales better as measurements increase.
• A theoretical result shows gradient descent can fit any signal with a sufficiently wide single-layer network, justifying early stopping.