[Paper Review] Learned D-AMP: Principled Neural Network based Compressive Image Recovery

Christopher A. Metzler, Ali Mousavi+1 more|arXiv (Cornell University)|Apr 21, 2017

Sparse and Compressive Sensing TechniquesEngineering192 citations

TL;DR

LDAMP unrolls the D-AMP algorithm into a trainable neural network using a denoising CNN denoiser; it achieves higher accuracy and much faster runtimes than BM3D-AMP and NLR-CS across tested matrices.

ABSTRACT

Compressive image recovery is a challenging problem that requires fast and accurate algorithms. Recently, neural networks have been applied to this problem with promising results. By exploiting massively parallel GPU processing architectures and oodles of training data, they can run orders of magnitude faster than existing techniques. However, these methods are largely unprincipled black boxes that are difficult to train and often-times specific to a single measurement matrix. It was recently demonstrated that iterative sparse-signal-recovery algorithms can be "unrolled" to form interpretable deep networks. Taking inspiration from this work, we develop a novel neural network architecture that mimics the behavior of the denoising-based approximate message passing (D-AMP) algorithm. We call this new network Learned D-AMP (LDAMP). The LDAMP network is easy to train, can be applied to a variety of different measurement matrices, and comes with a state-evolution heuristic that accurately predicts its performance. Most importantly, it outperforms the state-of-the-art BM3D-AMP and NLR-CS algorithms in terms of both accuracy and run time. At high resolutions, and when used with sensing matrices that have fast implementations, LDAMP runs over $50 imes$ faster than BM3D-AMP and hundreds of times faster than NLR-CS.

Motivation & Objective

• Motivate improved compressive image recovery by combining principled iterative ideas with data-driven learning.
• Develop a neural network architecture that mimics D-AMP/D-IT and uses a learnable denoiser.
• Provide training strategies and theoretical support (state evolution) for LDAMP.
• Demonstrate LDAMP’s performance advantages over existing methods across matrix types and resolutions.

Proposed method

• Unroll the D-AMP/D-IT iterations into a multi-layer neural network (LDAMP).
• Incorporate a DnCNN-based denoiser within each layer to enforce natural-image priors.
• Use Onsager correction to maintain AWGN-like effective noise per iteration.
• Train only the denoiser weights across layers, with variants including end-to-end, layer-by-layer, and denoiser-by-denoiser training.
• Employ state evolution as a heuristic to predict LDAMP performance.
• Compare with hand-designed (TVAL3) and other learning-based methods (BM3D-AMP, NLR-CS).

Experimental results

Research questions

• RQ1 How can we map D-AMP/D-IT iterations into trainable neural network layers while preserving principled behavior?
• RQ2 Does LDAMP generalize across different measurement matrices and noise levels?
• RQ3 What training strategy (end-to-end vs layer-by-layer vs denoiser-by-denoiser) is theoretically MMSE-optimal for LDAMP?
• RQ4 How does LDAMP compare to state-of-the-art CS recovery algorithms in accuracy and speed?
• RQ5 Can state-evolution predict LDAMP performance in practice?

Key findings

• LDAMP outperforms BM3D-AMP and NLR-CS in PSNR and run time across tested setups.
• With Gaussian and coded diffraction measurements, LDAMP achieves higher PSNRs and substantially faster runtimes than competitors (e.g., up to 50x faster than BM3D-AMP at high resolutions).
• Layer-by-layer and denoiser-by-denoiser training are MMSE-optimal for LDAMP, matching or exceeding end-to-end training in performance.
• State evolution accurately predicts LDAMP’s intermediate MSE across layers under i.i.d. Gaussian A and noise assumptions.
• LDAMP trained at one sampling rate generalizes well to other sampling rates when using denoiser-by-denoiser training.
• LDAMP remains effective across different measurement matrices and scales (e.g., 128x128 and higher resolutions) with fast forward/backward operators.