Source-linked AI summary
COAST: COntrollable Arbitrary-Sampling NeTwork for Compressive Sensing
Di You, Jian Zhang, Jingfen Xie, Bin Chen, Siwei Ma
TL;DR
Existing CS networks often require separate models for different sampling matrices, limiting arbitrary-sampling reconstruction. COAST uses one optimization-inspired model with sampling augmentation and controllable reconstruction modules, and experiments report strong arbitrary-matrix performance, including unseen matrices.
Problem
Most existing CS networks treat different sampling matrices as separate tasks and require retraining for each matrix.
Method
COAST uses random projection augmentation, controllable proximal mapping, and plug-and-play deblocking to process arbitrary sampling matrices with one model.
Results
COAST matches separately trained models on arbitrary seen matrices, performs similarly on unseen matrices, and improves state-of-the-art CS results on FRGM and DALM.
Takeaways & Limitations
COAST supports arbitrary sampling in a single CS model while retaining strong reconstruction performance and computational speed.
Abstract
from arXiv · showhide
Recent deep network-based compressive sensing (CS) methods have achieved great success. However, most of them regard different sampling matrices as different independent tasks and need to train a specific model for each target sampling matrix. Such practices give rise to inefficiency in computing and suffer from poor generalization ability. In this paper, we propose a novel COntrollable Arbitrary-Sampling neTwork, dubbed COAST, to solve CS problems of arbitrary-sampling matrices (including unseen sampling matrices) with one single model. Under the optimization-inspired deep unfolding framework, our COAST exhibits good interpretability. In COAST, a random projection augmentation (RPA) strategy is proposed to promote the training diversity in the sampling space to enable arbitrary sampling, and a controllable proximal mapping module (CPMM) and a plug-and-play deblocking (PnP-D) strategy are further developed to dynamically modulate the network features and effectively eliminate the blocking artifacts, respectively. Extensive experiments on widely used benchmark datasets demonstrate that our proposed COAST is not only able to handle arbitrary sampling matrices with one single model but also to achieve state-of-the-art performance with fast speed. The source code is available on https://github.com/jianzhangcs/COAST.
I. INTRODUCTION
Compressive sensing reconstructs natural signals from few linear measurements, but existing deep CS models usually require separate training for each sampling matrix. COAST instead targets arbitrary, including unseen, sampling matrices with one model while maintaining strong reconstruction performance.
- For γ = 10%, COAST significantly surpasses ISTA-Net+, especially on unseen sampling matrices.
- Existing deep CS methods commonly treat each sampling matrix as an independent task and require a separately trained model.
- COAST handles arbitrary sampling matrices, including unseen matrices, with one single model trained once.
- COAST combines random projection augmentation, controllable proximal mapping, and plug-and-play deblocking to improve arbitrary-sampling generalization and reconstruction quality.
- COAST achieves state-of-the-art performance for arbitrary sampling matrices with one single model, supporting real-world CS applications.
II. BACKGROUND
CS reconstruction methods span optimization-based approaches and network-based approaches, with deep unfolding combining their interpretability and learning advantages. A central unresolved issue is handling arbitrary sampling matrices without repeated search or retraining.
- Traditional CS methods solve regularized reconstruction problems iteratively using hand-crafted image priors, while network methods learn reconstructions with lower time complexity.
- Deep unfolding networks combine optimization-based structure with deep-network learning for interpretable and effective CS reconstruction.
- Existing methods typically treat different sampling matrices as separate tasks, requiring iterative re-search or model retraining for each matrix.
- ISTA alternates a gradient-style update with proximal mapping, whose efficient implementation is important for reconstruction.
III. APPROACH
COAST is formulated as an optimization-inspired network for natural-image CS and is organized into sampling, initialization, and recovery subnets. Its learnable prior term uses a residual network to implement proximal mapping.
- COAST partitions its architecture into a sampling subnet, initialization subnet, and recovery subnet.
- The method models natural-image CS measurements as linear sampling of image patches corrupted by additive Gaussian noise.
- COAST replaces a hand-crafted image prior with learnable parameters and uses a residual neural network to solve its proximal mapping operator.
B. Architecture Design of COAST
COAST combines sampling, initialization, and recovery subnets within an optimization-inspired design to process arbitrary sampling matrices using one model. RPA broadens sampling during training, while CPMM and PnP-D address condition-dependent reconstruction and blocking artifacts.
- Overall architecture: COAST comprises sampling, initialization, and recovery subnets, with the recovery subnet implemented as an interpretable deep unfolding model.The sampling subnet incorporates RPA, the initialization subnet avoids matrix-specific parameterization, and the recovery subnet maps ISTA into a network.
- Sampling Subnet: COAST trains with augmented sampling matrices but uses the supplied sampling matrix directly at testing, enabling one model to handle arbitrary matrices.Previous methods use a fixed matrix and produce Φ-specific models, whereas COAST’s RPA creates a Φ-agnostic training setup.
- Sampling Subnet: RPA augments each given M-by-N sampling matrix with additional same-dimension matrices and can extend across multiple arbitrary dimensions.The augmented set supports varying CS ratios and patch sizes; training randomly selects among augmented matrices.
- Initialization Subnet: The initialization subnet sets x̂(0)=0 and ρ(0)=1, an initialization equivalent to x̂(0)=Φ^T y that accommodates arbitrary matrix dimensions.This avoids the large parameterization and fixed input dimensions associated with fully connected initialization layers.
- Recovery Subnet: CPMM conditions proximal mapping on z=[γ,σ] to dynamically modulate features for different CS ratios and noise levels.The supported CS ratios are 10%, 20%, 30%, 40%, and 50%, with σ∈[0,1].
- Recovery Subnet: PnP-D jointly reconstructs independently sampled image patches at testing without extra parameters, allowing plug-in use while targeting blocking artifacts.It combines folding and unfolding operations and can be embedded into existing CS systems without changing their original model structures.
C. Network Parameters and Loss Function
COAST uses RPA-augmented sampling matrices and takes the measurement, sampling matrix, and control variables as inputs to learn reconstruction parameters. Its loss minimizes reconstruction discrepancy, while testing supports arbitrary sampling matrices.
- C. Network Parameters and Loss Function: RPA generates an augmented sampling-matrix set from sampling matrices with different dimensions for COAST training.The method defines the augmented set before forming the training objective.
- C. Network Parameters and Loss Function: COAST minimizes the discrepancy between each training signal x_i and its reconstruction H_COAST(y_i,t, Φ_t, z_i).The control input is z_i = [γ_i, σ_i].
- C. Network Parameters and Loss Function: The learnable parameter set Θ includes the parameters of the GDM and the controllable proximal mapping modules across phases.
- C. Network Parameters and Loss Function: At testing, users may provide either trained-on or unseen arbitrary sampling matrices Φ to obtain measurements and reconstructions from H_COAST(y, Φ, z).The input control vector is z = [γ, σ].
A. Implementation Details
COAST is trained on luminance patches extracted from 91 images using PyTorch and Adam, with 20 phases requiring roughly four days of training.
- A. Implementation Details: COAST training uses 88,912 luminance patches of size 33 × 33 extracted from 91 images, with batch size 64.The implementation uses an Intel Core i7-6820 CPU and GTX1080Ti GPU workstation.
- A. Implementation Details: Training COAST with 20 phases takes roughly four days using Adam with momentum 0.9 and weight decay 0.999.
B. Comparison with State-of-the-Art Methods
COAST is evaluated against representative model-driven, deep, unfolding, and unsupervised CS methods across fixed, learned, and unseen sampling matrices. It achieves the strongest reported reconstruction performance while retaining fast speed and generalization across unseen matrices.
- B. Comparison with State-of-the-Art Methods: Visual comparisons cover FRGM at γ = 10%, DALM at γ = 10% and 50%, and USM at γ = 30% and 50% across Set11 and BSD68.
- B. Comparison with State-of-the-Art Methods: COAST uses RPA, CPMM, and PnP-D to improve reconstruction quality while maintaining fast speed.The method is compared with ten representative state-of-the-art CS methods.
- B. Comparison with State-of-the-Art Methods: COAST achieves the highest PSNR/SSIM results on FRGM, with average PSNR gains of 1.46/1.54 dB over NLR-CSNet and 1.24/0.88 dB over ISTA-Net+ on Set11/BSD68.The comparison also reports computationally attractive real-time speed.
- B. Comparison with State-of-the-Art Methods: On DALM, COAST achieves the best performance among competing methods and handles multiple learned sampling matrices with one model.It improves average PSNR over OPINE-Net by roughly 0.32 dB on Set11 and 0.21 dB on BSD68 across three CS ratios.
- B. Comparison with State-of-the-Art Methods: COAST generalizes to unseen sampling matrices without additional training, outperforming computationally expensive TVAL3 and BM3D-AMP while remaining comparable to trained-matrix performance.Φ-specific ISTA-Net+, ReconNet, and DPDNN perform poorly on arbitrary unseen matrices.
C. Single Model for Arbitrary Sampling Matrices
COAST-A handles arbitrary CS ratios and patch sizes with one model, matching separately trained models closely, while COAST-N remains similarly robust across noise levels.
- C. Single Model for Arbitrary Sampling Matrices: COAST is evaluated for arbitrary sampling using noiseless and noisy measurements across varied CS ratios, patch sizes, and noise levels.
- C. Single Model for Arbitrary Sampling Matrices: The USM experiments report COAST performance on unseen CS ratios outside the training set and compare it with separately trained settings.
- C. Single Model for Arbitrary Sampling Matrices: COAST-A handles arbitrary CS ratios and patch sizes with one model, matching COAST-S within 0.2 dB PSNR across the tested tasks.COAST-A even outperforms COAST-S at γ = 10% with patch size N = 33.
- C. Single Model for Arbitrary Sampling Matrices: COAST-N remains robust to noise, staying within 0.1 dB PSNR of separately trained COAST-S across all tested noise levels and CS ratios.At γ = 10% and σ = 5, COAST-N surpasses COAST-S.
D. Ablation Studies and Discussions
Ablations show that RPA improves reconstruction across sampling matrices, CPMM's CU improves robustness, shared parameters reduce model size without performance loss, and PnP-D removes blocking artifacts. Increasing phase number improves PSNR until the gains flatten, motivating 20 phases.
- Effect of RPA: RPA improves reconstruction performance on both FRGM and DALM, with gains of 0.8 dB and 0.29 dB respectively over no augmentation.The selected defaults are NS = 25 for FRGM and NS = 5 for DALM, balancing memory use and recovery performance.
- Effect of CU in CPMM: CPMM with CU consistently outperforms CPMM without CU across five CS ratios, achieving average gains of 0.07 dB and higher noise robustness.The CU dynamically modulates network features within CPMM.
- Shared vs. Unshared: Sharing CU parameters reduces the parameter count from 5760 to 96 without performance loss compared with unshared parameters.COAST contains 60 CUs, each with 96 parameters, so sharing preserves performance while substantially reducing storage.
- With vs. Without PnP-D: PnP-D boosts performance across all sampling ratios, with improvements up to 0.93 dB, while eliminating blocking artifacts in reconstructed images.The visual comparison supports the artifact-removal effect, while the ablation text supplies the quantitative comparison.
- Performance of Phase Number: PSNR increases with phase number but becomes nearly flat at NP ≥ 20, leading COAST to use 20 phases as a complexity-performance trade-off.Each COAST phase corresponds to one iteration in traditional ISTA.
- Discussion: COAST is compared with methods including SCSNet, RACS, and DRNN because arbitrary sampling ratios and sampling matrices remain central discussion points.COAST handles given sampling matrices, while SCSNet and RACS learn the sampling matrix; its RPA also operates at the matrix level rather than only the ratio level.
E. Analysis on USM for Deep Unfolding Methods
The analysis attributes deep unfolding methods' poor performance on unseen sampling matrices to sampling-matrix-dependent initialization and learned reconstruction parameters, compounded by short, large-step iterations. COAST shows only slight performance changes on unseen matrices, indicating stronger generality.
- Initialization: ISTA-Net+ cannot handle unseen sampling ratios because its initialization mapping Qinit has a fixed dimension tied to the training sampling matrix.It also performs very poorly on unseen matrices with seen CS ratios.
- Reconstruction: COAST's performance gain or drop on unseen sampling matrices is slight, verifying superior generality across sampling ratios.The comparison uses zero-initialized ISTA-Net+, BM3D-AMP, and COAST on Set11.
- Reconstruction: Deep unfolding parameters learned from training pairs contain information about the training sampling matrix, reducing generality on unseen matrices.This dependence includes learned image transforms, step sizes, and shrinkage thresholds.
- Reconstruction: ISTA-Net+'s nine forced iterations produce larger steps and stronger proximal mappings that can become suboptimal when the unseen target problem changes.The analysis contrasts these short learned iterations with traditional optimization-based methods.
- Conclusion of analysis: The two mechanisms together cause deep unfolding models to overfit the specific sampling matrix used during training.The paper identifies initialization dependence and learned-parameter dependence as the underlying reasons.
V. CONCLUSION
COAST is an optimization-inspired deep unfolding network that handles arbitrary sampling matrices with one model while offering interpretability, robustness, and fast computation. Experiments show performance on par with separately trained models and improved state-of-the-art results on FRGM and DALM.
- Conclusion: COAST handles arbitrary sampling matrices with one single model under an optimization-inspired deep unfolding framework.The conclusion highlights interpretability, robustness, and attractive computational speed.
- Conclusion: COAST matches separately trained single-matrix models and improves state-of-the-art CS results on both FRGM and DALM.The conclusion frames these results as outcomes of extensive experiments.