[Paper Review] ISTA-Net++: Flexible Deep Unfolding Network for Compressive Sensing

Di You, Jingfen Xie+1 more|arXiv (Cornell University)|Mar 22, 2021

Sparse and Compressive Sensing Techniques26 references8 citations

TL;DR

ISTA-Net++ proposes a flexible deep unfolding network for compressive sensing that dynamically adapts to multiple sampling ratios using a condition module, while employing a cross-block strategy and a balanced training dataset (Train400) to reduce blocking artifacts and improve robustness across diverse image scenes. It achieves state-of-the-art performance on four datasets with consistent gains across multi-ratio and multi-scene tasks.

ABSTRACT

While deep neural networks have achieved impressive success in image compressive sensing (CS), most of them lack flexibility when dealing with multi-ratio tasks and multi-scene images in practical applications. To tackle these challenges, we propose a novel end-to-end flexible ISTA-unfolding deep network, dubbed ISTA-Net++, with superior performance and strong flexibility. Specifically, by developing a dynamic unfolding strategy, our model enjoys the adaptability of handling CS problems with different ratios, i.e., multi-ratio tasks, through a single model. A cross-block strategy is further utilized to reduce blocking artifacts and enhance the CS recovery quality. Furthermore, we adopt a balanced dataset for training, which brings more robustness when reconstructing images of multiple scenes. Extensive experiments on four datasets show that ISTA-Net++ achieves state-of-the-art results in terms of both quantitative metrics and visual quality. Considering its flexibility, effectiveness and practicability, our model is expected to serve as a suitable baseline in future CS research. The source code is available on https://github.com/jianzhangcs/ISTA-Netpp.

Motivation & Objective

• Address the lack of flexibility in existing deep learning-based compressive sensing methods when handling multiple sampling ratios in practical applications.
• Improve reconstruction robustness across diverse image scenes, especially for 'rarely seen' scenarios, by mitigating dataset imbalance in training data.
• Reduce blocking artifacts in reconstructed images through a novel cross-block strategy.
• Develop a unified model capable of multi-ratio CS reconstruction without retraining, enhancing practicality and efficiency.

Proposed method

• Introduces a dynamic unfolding strategy (DUS) that uses a condition module to inject sampling ratio information into each network stage, enabling single-model adaptation to multiple CS ratios.
• Employs a dynamic gradient descent module (DGDM) and dynamic proximal mapping module (DPMM) that are conditioned on the sampling ratio, allowing adaptive optimization during reconstruction.
• Applies a cross-block strategy (CBS) that enables feature exchange across adjacent image blocks to suppress blocking artifacts and improve texture fidelity.
• Trains the network on a balanced dataset (Train400), which includes diverse scenes beyond the dominant 'plant' category found in standard datasets like Train91.
• Uses a three-part framework: sampling (modeled as a convolutional layer), initialization (to resolve dimension mismatch), and reconstruction (via deep unfolding with DUS and CBS).
• Leverages end-to-end training with a loss function minimizing reconstruction error, enabling joint optimization of all components.

Experimental results

Research questions

• RQ1 Can a single deep learning model effectively handle compressive sensing reconstruction across multiple sampling ratios without retraining?
• RQ2 How does dataset imbalance affect the generalization of CS models on unseen image scenes, and can this be mitigated?
• RQ3 To what extent does a cross-block strategy reduce blocking artifacts in deep unfolding-based CS reconstruction?
• RQ4 Does using a more balanced training dataset like Train400 lead to measurable improvements in reconstruction robustness across diverse scenes?

Key findings

• ISTA-Net++ achieves an average PSNR of 38.73 dB at 50% sampling ratio on Set11, outperforming the baseline without the condition module by 0.26 dB on average across 10%, 30%, and 50% ratios.
• The cross-block strategy (CBS) reduces blocking artifacts and improves PSNR by up to 0.64 dB compared to the baseline without CBS.
• The model trained on the balanced Train400 dataset achieves 34.81 dB PSNR at 30% ratio, outperforming the model trained on unbalanced Train91 (34.25 dB) with the same number of epochs.
• On the DIV2K validation set at 30% ratio, the ISTA-Net++ model trained on Train400 outperforms the one trained on Train91 by 0.97 dB in the 'house' scene, demonstrating improved robustness to rare scenes.
• The combination of dynamic unfolding and cross-block strategy leads to consistent performance gains across all tested ratios and datasets, confirming the effectiveness of both components.
• Experiments show that dataset balance has a greater impact on performance than dataset size, as adding more data to Train400 (e.g., Train400+Train91) yields negligible further gains.