[Paper Review] Neural Proximal Gradient Descent for Compressive Imaging

Morteza Mardani, Qingyun Sun+5 more|arXiv (Cornell University)|Jun 1, 2018

Sparse and Compressive Sensing Techniques117 citations

TL;DR

The paper unrolls proximal gradient iterations and learns the proximal map with a recurrent ResNet, achieving faster, data-feasible MRI reconstruction and improved super-resolution compared to traditional CS methods and non-recurrent nets.

ABSTRACT

Recovering high-resolution images from limited sensory data typically leads to a serious ill-posed inverse problem, demanding inversion algorithms that effectively capture the prior information. Learning a good inverse mapping from training data faces severe challenges, including: (i) scarcity of training data; (ii) need for plausible reconstructions that are physically feasible; (iii) need for fast reconstruction, especially in real-time applications. We develop a successful system solving all these challenges, using as basic architecture the recurrent application of proximal gradient algorithm. We learn a proximal map that works well with real images based on residual networks. Contraction of the resulting map is analyzed, and incoherence conditions are investigated that drive the convergence of the iterates. Extensive experiments are carried out under different settings: (a) reconstructing abdominal MRI of pediatric patients from highly undersampled Fourier-space data and (b) superresolving natural face images. Our key findings include: 1. a recurrent ResNet with a single residual block unrolled from an iterative algorithm yields an effective proximal which accurately reveals MR image details. 2. Our architecture significantly outperforms conventional non-recurrent deep ResNets by 2dB SNR; it is also trained much more rapidly. 3. It outperforms state-of-the-art compressed-sensing Wavelet-based methods by 4dB SNR, with 100x speedups in reconstruction time.

Motivation & Objective

• Motivate rapid, plausible reconstruction in ill-posed linear inverse problems.
• Incorporate physical data consistency with learned priors.
• Develop a neural proximal gradient framework by unrolling proximal gradient iterations.
• Analyze contraction properties and empirical convergence of the learned proximal.
• Evaluate on pediatric MRI reconstruction and natural image super-resolution.

Proposed method

• Cast proximal gradient iterations as a recurrent network to model data-consistent updates.
• Model the proximal operator with a multi-layer neural network (ResNet-based).
• Train end-to-end by minimizing empirical risk over a training set with a mixed loss (reconstruction and measurement-consistency).
• Use a state-space form with s_{t+1}=g(x_t;y) and x_{t+1}=P_ψ(s_{t+1}) to couple data fidelity with learned proximal.
• Explore different activation mask schemes (D(z)) in the proximal network and address vanishing gradients with ResNet/DiracNet variants.

Experimental results

Research questions

• RQ1 Can a recurrent neural network implementing a learned proximal operator improve reconstruction quality for undersampled inverse problems?
• RQ2 Does repeating a small proximal block outperform a single deep network in MRI reconstruction?
• RQ3 Do contraction properties hold empirically for the learned proximal, ensuring convergence of the unrolled iterations?
• RQ4 How does the neural proximal approach compare to traditional CS-based methods and non-recurrent deep nets in reconstruction quality and speed?

Key findings

• A recurrent ResNet with a single residual block (RB) unrolled over iterations can effectively model the proximal and reveal MR details.
• The recurrent architecture outperforms non-recurrent deep ResNets by about 2 dB in SNR and trains more rapidly.
• The method outperforms state-of-the-art CS-WV methods by about 4 dB SNR with up to 100x faster reconstruction time.
• Increasing the number of iterations (T) improves SNR and SSIM, but very deep/prolific networks with many RBs may not yield proportional gains for a single iteration.
• For MRI, 10 iterations with 1 RB offers a good trade-off between quality and computation, enabling practical clinical use.
• Contraction analysis suggests that the dominant contraction comes from the fixed linear part, with the learned proximal contributing a smaller perturbation, and later iterations show decaying perturbation effects.