Source-linked AI summary
AMP-Net: Denoising based Deep Unfolding for Compressive Image Sensing
Zhonghao Zhang, Yipeng Liu, Jiani Liu, Fei Wen, Ce Zhu
TL;DR
Visual image compressive sensing needs reconstruction methods that balance the interpretability of model-based optimization with the speed of deep networks. AMP-Net addresses this by unfolding AMP’s iterative denoising process, adding trainable deblocking and joint sampling-matrix training; experiments report better reconstruction performance than other state-of-the-art methods, including ten methods in one comparison.
Problem
Classical deep networks quickly reconstruct compressive measurements but have limited interpretation and theoretical guarantees, while iterative methods are slower.
Method
AMP-Net unfolds AMP’s denoising perspective into a network and combines CNN-based denoising, trainable deblocking, and jointly trained sampling.
Results
AMP-Net-2-BM has better performance than ten other state-of-the-art methods.
Takeaways & Limitations
AMP-Net combines reconstruction accuracy with high speed and a small number of network parameters when the number of reconstruction modules is small.
Abstract
from arXiv · showhide
Most compressive sensing (CS) reconstruction methods can be divided into two categories, i.e. model-based methods and classical deep network methods. By unfolding the iterative optimization algorithm for model-based methods onto networks, deep unfolding methods have the good interpretation of model-based methods and the high speed of classical deep network methods. In this paper, to solve the visual image CS problem, we propose a deep unfolding model dubbed AMP-Net. Rather than learning regularization terms, it is established by unfolding the iterative denoising process of the well-known approximate message passing algorithm. Furthermore, AMP-Net integrates deblocking modules in order to eliminate the blocking artifacts that usually appear in CS of visual images. In addition, the sampling matrix is jointly trained with other network parameters to enhance the reconstruction performance. Experimental results show that the proposed AMP-Net has better reconstruction accuracy than other state-of-the-art methods with high reconstruction speed and a small number of network parameters.
I. INTRODUCTION
AMP-Net is a deep unfolding model for visual image compressive sensing that combines AMP-inspired denoising, trainable deblocking, and jointly optimized sampling. It targets the trade-off between interpretable iterative methods and fast classical deep networks.
- Compressive sensing reconstructs signals from fewer measurements than Nyquist sampling and supports applications including single-pixel cameras, MRI, and SCI.
- Classical deep networks reconstruct images quickly but lack the interpretation and theoretical guarantees associated with iterative optimization methods.
- AMP-Net unfolds AMP’s iterative denoising perspective into a reconstruction network rather than learning conventional regularization terms.The model interprets AMP’s nonlinear operations as denoising processes dependent on image priors.
- AMP-Net uses a sampling model and a K-iteration reconstruction model containing initialization, denoising, and whole-image deblocking modules.Images are measured block-by-block with the same sampling matrix.
- AMP-Net combines CNNs with the sampling matrix to estimate AMP noise terms and introduces a trainable deblocking module for block-based reconstruction artifacts.The deblocking module follows the denoising module and can also be plugged into other deep learning methods.
- The sampling matrix is jointly trained with other network parameters because it affects both sampling and reconstruction.The authors state that the resulting data-driven matrix can improve performance for AMP-Net and other reconstruction methods.
II. BACKGROUND
The AMP algorithm’s denoising perspective rewrites its iterative linear operation as the original signal plus noise, allowing the nonlinear step to be learned as denoising rather than specified through a regularizer.
- AMP rewrites the linear operation A^T(y−Ax′)+x′ as the original data plus a noise term.The noise term is e = (A^TA − I)(x̄ − x^k−1).
- Under the i.i.d. Gaussian sampling-matrix assumption, the residual component is Gaussian distributed with variance determined by the current estimation error.The matrix entries satisfy A_ij ∼ N(0, 1/M).
- The AMP algorithm interprets each nonlinear iteration as denoising that depends on the assumed signal prior.For sparse signals without value limitation, the denoiser can be a soft-thresholding function.
- The noise term is independent of the model-based regularizer and can therefore be fitted by supervised learning while learning image priors flexibly.This avoids requiring a specific regularization-term form.
- As the estimation-error norm decreases across iterations, the Euclidean distance between reconstruction and original data decreases.
B. Related Works
Prior work spans regularization-based recovery, direct deep networks, and deep unfolding; AMP-Net differs by learning noise through the sampling matrix rather than directly estimating the entire noise term.
- Model-based methods: Model-based image recovery uses structure-inducing regularizers such as sparsity, low rank, and transformation-domain priors.Fixed transformation domains may result in poor performance, motivating more elaborate priors such as denoising priors.
- Classical deep network methods: Classical deep networks directly map compressive samples to image estimates using architectures including autoencoders, CNNs, and GANs.
- Deep unfolding methods: Deep unfolding maps iterative restoration algorithms onto neural networks to balance reconstruction speed with model-based interpretability.Examples include unfolded ISTA, AMP, and IHT algorithms.
- Denoising-prior unfolding methods: DPDNN learns regularizer information with two control parameters, whereas AMP-Net learns noise approximation with one control parameter.
- Denoising-prior unfolding methods: AMP-Net estimates (A^TA − I)(x̄ − x^k−1) by fitting only (x̄ − x^k−1) with a CNN before multiplying by the sampling matrix.LDIT instead estimates the entire noise term directly; AMP-Net also introduces trainable α_k.
III. AMP-NET
AMP-Net combines block-based sampling with a deep-unfolded denoising reconstruction process, trainable deblocking modules, and jointly optimized sampling parameters.
- Sampling Model: AMP-Net uses a block-based sampling model with a shared sampling matrix A, producing measurements for each image block.The sampling matrix is trained jointly with the other model parameters because it contributes to both sampling and reconstruction.
- Reconstruction Model: The reconstruction model unfolds K iterative denoising steps inspired by the denoising perspective of approximate message passing.It contains an initialization module followed by stacked reconstruction modules, each representing one iteration.
- Denoising Module: Each reconstruction module applies a CNN denoiser to image blocks and uses a trainable control parameter αk to regulate noise and reconstruction behavior.The control parameter acts similarly to a step size and reduces to the corresponding unscaled process when αk equals 1.
- Joint Sampling Optimization: The sampling matrix A directly participates in the reconstruction process, so an appropriate learned matrix can improve recovery without extra floating-point operations from block reshaping functions.The paper identifies A as important to reconstruction performance and jointly trains it with the network parameters.
- Deblocking Module: A trainable deblocking module processes the concatenated image to fit and remove additive blocking artifacts introduced by non-overlapping reconstruction.The module uses a ResNet-like structure and operates on the whole image rather than individual blocks.
C. Loss Function
AMP-Net variants combine denoising modules, deblocking modules, and sampling-matrix training, and are optimized using an MSE loss over reconstructed images.
- Model Variants: AMP-Net-K denotes a model with K denoising modules, while the B, M, and BM suffixes indicate deblocking, matrix training, or both.AMP-Net-K-BM includes both deblocking modules and the sampling-matrix training strategy.
- Trainable Parameters: AMP-Net-K-BM trains the measurement matrix A, initialization matrix B, control parameters Sα, denoiser parameters SΘ, and deblocking parameters SΩ.The parameter sets include trainable parameters for every reconstruction module.
- Forward Propagation: The forward pass measures the input image block-by-block, initializes X0 from BY, and iteratively updates the reconstruction for K modules.Algorithm 1 explicitly includes measurement, initialization, residual calculation, and iterative reconstruction steps.
- Loss Function: Mean square error measures the difference between original and recovered images and defines the training loss for AMP-Net-K-BM.Losses for the other variants are obtained by removing the corresponding model elements from the forward process and loss formulation.
IV. EXPERIMENTAL RESULTS
The experiments evaluate AMP-Net components individually and compare the full model with state-of-the-art reconstruction methods for visual image compressive sensing.
- Experimental Design: The experiments separately validate the unfolding strategy, deblocking module, and sampling-matrix learning strategy.They then compare AMP-Net with deblocking and sampling-matrix optimization against state-of-the-art reconstruction methods.
- BSDS500 Evaluation: The comparison includes AMP-Net variants and eight comparison methods on the BSDS500 test set without deblocking or sampling-matrix strategies.Table III reports results for AMP-Net-6 and AMP-Net-9 alongside the comparison methods.
A. Experimental settings
Experiments use BSDS500 for training, validation, and testing, Set11 for testing, and block-based models trained with specified image-block datasets and shared initialization conditions.
- Datasets: BSDS500 contains 500 color images split into 200 training, 100 validation, and 200 test images, while Set11 contains 11 grayscale test images.Only the luminance components of BSDS500 are used for training, validation, and testing.
- Training Data: Image blocks use size n = 33, with separate training sets for models with and without trainable deblocking operations.The two sets contain randomly extracted 99 × 99 sub-images and 33 × 33 image patches, respectively.
- Training Protocol: The sampling matrix is initialized as an orthogonalized Gaussian matrix, and models at the same compressive-sensing ratio use the same initialization.Training uses Adam, with reconstruction-module counts of 2, 4, 6, or 9.
- Baselines: The methods compared with AMP-Net include model-based, classical deep-network, and deep-unfolding approaches adapted to block-based image compressive sensing.The listed baselines include TVAL3, D-AMP, DCS, ReconNet, CSNet+, LDIT, LDAMP, DPDNN, NN, and ISTA-Net+.
B. Validating the unfolding strategy of AMP-Net
The experiments evaluate AMP-Net’s unfolded denoising strategy against model-based, classical deep, and other deep unfolding methods across CS ratios and reconstruction depths. AMP-Net generally performs strongly, particularly at lower CS ratios and on BSDS500.
- Experimental setup: The unfolding-strategy evaluation trains the compared trainable models on training set 2 and excludes deblocking and sampling-matrix training.The setup isolates the reconstruction strategy before the additional modules and sampling optimization are evaluated.
- Cross-method comparison: AMP-Net-6 outperforms DPDNN and NN, while its AMP-derived design achieves higher PSNR and SSIM than other AMP-based methods at most CS ratios.The comparison includes TVAL3, D-AMP, ReconNet, LDIT, LDAMP, DPDNN, NN, and ISTA-Net+.
- Generalization: AMP-Net-6 and AMP-Net-9 generalize favorably to BSDS500, with AMP-Net-6 outperforming five comparison methods and AMP-Net-9 surpassing ISTA-Net+ at 10%, 4%, and 1%.The evaluated comparison methods are ReconNet, LDIT, LDAMP, DPDNN, NN, and ISTA-Net+.
- Reconstruction depth: Across reconstruction-module counts, AMP-Net-K performs best at a CS ratio of 10% and remains competitive at 30% when using fewer than nine modules.The analysis uses average PSNR on BSDS500 and also evaluates a variant with α_k fixed to 1.
C. Validating the Capability of the Deblocking Module
This section tests whether the trainable deblocking module improves block-based CS reconstruction and transfers to other reconstruction models. Results compare no deblocking, BM3D, and the learned module on Set11.
- Effectiveness: Models with the trainable deblocking module achieve higher PSNR and SSIM than versions without it or using BM3D.The improvement is attributed to joint training of the deblocking module and the other model parameters.
- Artifact removal: Block-by-block reconstruction produces visible blocking artifacts without deblocking, while BM3D still leaves visible artifacts because it targets Gaussian noise.The learned module removes the blocking artifacts more effectively in the reported visual comparison.
- Universality: The deblocking module improves models beyond AMP-Net, indicating applicability to ReconNet, LDIT, LDAMP, ISTA-Net+, DPDNN, and NN.The comparison uses models with and without the module under the same evaluation framework.
E. Validating the performance of AMP-Net-K-BM
AMP-Net-K-BM combines deep unfolding, deblocking, and trained sampling matrices for visual image CS. Comprehensive comparisons report stronger reconstruction performance than competing methods, with performance improving as reconstruction depth increases.
- Experimental setup: The comparison includes AMP-Net-K-BM with K values of 2, 4, 6, and 9 alongside ten state-of-the-art methods.Compared models follow the settings reported in their original papers, including trained sampling matrices for selected methods.
- Evaluation scope: Tables VII and VIII report reconstruction quality across Set11 and BSDS500, while also providing reconstruction time, parameter counts, and visual results for Monarch and Parrots.The visual comparisons use CS ratios of 30% and 10%.
- Overall comparison: With deblocking and trained sampling, AMP-Net outperforms competing methods, including ReconNet and CSNet+, which also use both enhancements.The reported advantage is attributed to AMP-Net’s deep unfolding strategy.
- Component combination: Adding both deblocking and a trained sampling matrix gives AMP-Net its best performance among the evaluated variants.The deblocking module increases parameter count while improving performance.
- Overall comparison: AMP-Net-2-BM outperforms ten state-of-the-art methods, and increasing the number of reconstruction modules further improves results.The conclusion reports this result across the comprehensive comparison of AMP-Net-K-BM variants.
APPENDIX A DERIVATION OF (4)
The appendix derives the gradients needed to train the sampling matrix jointly through AMP-Net’s sampling and reconstruction pathways. It begins from the AMP-related residual and linear operation.
- AMP algebra: The derivation starts with y = Ax and z_k = y − Ax_k, then rewrites A^Tz_0 + x_0 as the AMP linear operation.The resulting expression is A^T(y − Ax_0) + x_0.
- Gradient decomposition: The gradient of the sampling matrix is decomposed into sampling and reconstruction contributions, written as ∇_A L = A_Sam + A_Rec.The derivation considers one image block because the matrix processes blocks during sampling and reconstruction.
- Gradient decomposition: A_Sam denotes the gradient contribution from the sampling model, while A_Rec aggregates contributions from the reconstruction modules.The reconstruction contribution is indexed by module k and summed over K modules.
- Gradient derivation: The appendix introduces assumptions on gradients of intermediate estimates x_0 and x_k to derive the sampling and reconstruction terms.These assumptions support the subsequent expressions for A_Sam and A_Rec.
APPENDIX C DERIVATION OF (11)
The appendix derives equation (11) by introducing a new parameter, applying equation (23) to the image block process, and extending it iteratively.
- A new parameter α is introduced in distinction from equation (17).
- Equation (23) is applied to the image block process and developed into an iterative version.
- The preceding process constitutes the derivation of equation (11).