Source-linked AI summary
Learned D-AMP: Principled Neural Network based Compressive Image Recovery
Christopher A. Metzler, Ali Mousavi, Richard G. Baraniuk
TL;DR
Compressive image recovery must reconstruct images from underdetermined measurements while balancing accuracy, speed, and principled behavior. The paper unrolls D-AMP into the Learned D-AMP network, combining an iterative recovery structure with learned denoisers. LDAMP is reported to outperform BM3D-AMP and NLR-CS in accuracy and run time, while retaining state-evolution analysis and broad measurement-matrix applicability.
Problem
Compressive image recovery seeks to reconstruct an image from low-dimensional measurements using prior information, while existing accurate-prior methods can be too slow and data-driven methods can be unprincipled black boxes.
Method
LDAMP unrolls the D-AMP algorithm into a neural network with AMP layers and learned denoisers, combining hand-designed recovery structure with data-driven priors.
Results
LDAMP outperforms state-of-the-art BM3D-AMP and NLR-CS in both accuracy and run time, with state evolution accurately predicting its performance.
Takeaways & Limitations
Unrolling a conventional iterative algorithm produces a neural network with well-understood behavior, performance guarantees, and predictable shortcomings.
Takeaways & Limitations
D-IT’s effective noise deviates from additive white Gaussian noise as iterations proceed, reducing denoising effectiveness and slowing convergence.
Abstract
from arXiv · showhide
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\times$ faster than BM3D-AMP and hundreds of times faster than NLR-CS.
1 Introduction
Compressive sensing recovers signals from underdetermined measurements by using prior knowledge, with recovery methods ranging from hand-designed models to data-driven and hybrid approaches.
- Problem: When m < n, recovering x from y = Ax + ϵ is underdetermined, so compressive sensing uses prior knowledge about the signal.The framework applies across imaging domains including synthetic aperture radar, functional MRI, and CT.
- Hand-designed methods: Hand-designed methods use interpretable priors and often have understood behavior or convergence guarantees, but accurate-prior methods can be too slow for real-time applications.They also do not exploit potentially available training data.
- Data-driven methods: Data-driven neural networks learn signal structure from large training datasets and can compete in accuracy while running thousands of times faster than state-of-the-art methods.Examples include stacked denoising autoencoders, convolutional networks, and mixed convolutional–fully connected architectures.
- Hybrid methods: Hybrid methods combine expert-designed recovery algorithms with learned priors, aiming to retain data-driven accuracy while preserving interpretability and theoretical guarantees.They include black-box neural networks embedded within recovery algorithms and unrolled iterative algorithms whose parameters are learned.
- Unrolled algorithms: Unrolling converts well-understood iterative algorithms such as ISTA and AMP into neural networks whose algorithmic parameters can be optimized using training data.LISTA and LAMP exemplify this approach.
2 Learned D-AMP
LDAMP unrolls D-IT and D-AMP into a trainable network that combines measurement-based updates with learned image denoisers. Its design uses Onsager correction to preserve an AWGN-like effective-noise model while replacing non-differentiable denoisers with gradient-friendly convolutional networks.
- Core framework: LDAMP is a mixed hand-designed/data-driven compressive recovery framework built by unrolling D-IT and D-AMP.The network retains the structure of the recovery algorithms while learning denoiser parameters from data.
- D-IT and D-AMP: D-IT alternates a gradient step toward the measurement-consistent affine subspace with denoising toward the natural-image set C.Starting from x0 = 0, these steps repeat until convergence.
- D-IT and D-AMP: D-AMP adds an Onsager correction that removes intermediate-solution bias, keeping effective noise closer to the AWGN model expected by image denoisers.The correction addresses the deviation from AWGN that makes D-IT denoising less effective and slows convergence.
- Network architecture: LDAMP contains 10 AMP layers, each with two tied-weight denoisers: one updates x_l and the other estimates divergence.The divergence uses a Monte-Carlo approximation.
- Network architecture: Because BM3D cannot readily propagate gradients, LDAMP uses trainable denoisers such as DnCNN, which has 16 to 20 convolutional layers.The forward and backward measurement operators are represented by A and A^H, though function handles can also be used.
3 Training the LDIT and LDAMP networks
The paper compares three training strategies for unrolled LDIT and LDAMP networks, finding that LDAMP’s structured training remains optimal while denoiser-by-denoiser training generalizes across sampling rates.
- Training methods: End-to-end training jointly optimizes all network weights, whereas layer-by-layer training progressively adds and trains AMP or D-IT layers.Layer-by-layer training continues until the network reaches 10 layers.
- Training methods: Denoiser-by-denoiser training independently trains denoisers on AWGN problems across noise levels and selects weights using the estimated noise standard deviation during inference.The expected noise range is discretized when selecting denoiser weights.
- Training comparisons: For LDAMP, layer-by-layer and denoiser-by-denoiser training are theoretically MMSE optimal, and experiments support these predictions.The evaluated networks contain 10 unrolled DAMP or D-IT layers with 16-layer DnCNN denoisers.
- Training comparisons: Additional end-to-end training does not improve layer-by-layer LDAMP, whereas it dramatically improves layer-by-layer LDIT performance.This contrast distinguishes LDAMP’s behavior from that of a typical neural network represented by LDIT.
- Training comparisons: Denoiser-by-denoiser LDAMP performs slightly worse than the other LDAMP variants, likely because discretized noise levels are absent from the theory.Finer noise-level discretization or deeper denoisers could reduce this gap.
- Sampling-rate generalization: Denoiser-by-denoiser trained networks generalize well to different sampling rates because they are not trained at one specific sampling rate.The supplied passage contrasts this with networks trained end-to-end or layer-by-layer at a fixed rate.
4 Theoretical analysis of LDAMP
The theoretical analysis extends state evolution to LDAMP and uses it to establish optimality results for layer-by-layer and denoiser-by-denoiser training under stated conditions.
- 4.1 State-evolution: State evolution predicts LDAMP’s intermediate MSE across network layers, and empirical results show that these predictions are accurate.The comparison includes LDAMP and LDIT intermediate reconstructions and their predicted state evolution.
- 4.1 State-evolution: For D-AMP and LDAMP, state evolution remains observational rather than rigorously established.Rigorous asymptotic accuracy is stated only for AMP with i.i.d. Gaussian measurements.
- 4.1 State-evolution: LDAMP is well predicted by state evolution at the reported sampling rate, whereas LDIT is not.The supplied passage reports this contrast for n = 0.1.
- 4.1 State-evolution: For large m and n, state evolution predicts each LDAMP layer’s mean square error when the network starts from x0 = 0.The result assumes suitable measurement, noise, and denoiser conditions.
- 4.2 Layer-by-layer and denoiser-by-denoiser training is optimal: The state-evolution framework proves that layer-by-layer and denoiser-by-denoiser training of LDAMP are MMSE optimal.The proof proceeds through a lemma concerning monotone denoisers and sequential minimization of expected layer errors.
5 Experiments
The experiments evaluate LDAMP on Gaussian and coded-diffraction measurements, using PSNR and runtime for 128 × 128 reconstructions and visual comparisons at 512 × 512.
- Implementation: The implementation uses the DnCNN denoiser and provides LDAMP and LDIT versions in TensorFlow and MatConvnet.Public implementations are available in the D-AMP Toolbox repository.
- Testing setup: The evaluation tests i.i.d. Gaussian and randomly sampled coded-diffraction measurements, reporting PSNR and runtime.Except for Figure 7, experiments use 128 × 128 images; Figure 7 compares 512 × 512 reconstructions.
- Gaussian measurements: With noise-free Gaussian measurements, LDAMP produces the best reconstructions at every sampling rate and on every image except Fingerprints.Fingerprints differs substantially from the natural images used for training.
- Gaussian measurements: LDIT and LDAMP reconstruct significantly faster than competing methods under noise-free Gaussian measurements.At sampling rate m/n = 0.25, LDAMP is only about 25% slower than LDIT, with matrix multiplications dominating computation.
- Coded diffraction measurements: With noise-free coded-diffraction measurements, LDAMP gives the best reconstruction on every image except Fingerprints and is significantly faster than competing methods.Because the forward and backward operators run in O(n log n), denoising becomes the dominant computation; LDAMP takes roughly twice as long as LDIT.
- Coded diffraction measurements: At high resolutions, LDAMP reconstructions are incrementally better than BM3D-AMP while computing over 60× faster.The comparison uses 512 × 512 reconstructions of the Boat test image.
6 Conclusions
The paper develops and validates LDAMP, a neural network that mimics D-AMP. It reports broad applicability and improved accuracy and runtime over BM3D-AMP and NLR-CS.
- Conclusions: LDAMP is a novel neural network architecture that mimics the behavior of the D-AMP signal recovery algorithm.The paper develops, analyzes, and validates the architecture.
- Conclusions: LDAMP is easy to train, applies to varied measurement matrices, and has a state-evolution heuristic that accurately predicts performance.The paper characterizes these properties as part of the method’s validated behavior.
- Conclusions: LDAMP outperforms BM3D-AMP and NLR-CS in both accuracy and runtime.The conclusion identifies this as the paper’s most important result.
- Conclusions: Unrolling an iterative algorithm produces a neural network with well-understood behavior, performance guarantees, and predictable shortcomings.The paper presents LDAMP as an example of using training data to improve iterative algorithms.
Supplementary Results
The supplementary results include PSNR tables for Gaussian and coded-diffraction measurements, plus a noisy Boat-image reconstruction table.
- Gaussian measurements: Table 3 reports PSNR for 128 × 128 reconstructions using i.i.d. Gaussian measurements without measurement noise.
- Coded diffraction measurements: Table 4 reports PSNR for 128 × 128 reconstructions using coded-diffraction measurements without measurement noise.
- Noisy measurements: Table 5 reports PSNR for a 128 × 128 Boat reconstruction with additive white Gaussian measurement noise at various standard deviations.