Deep Physics-Guided Unrolling Generalization for Compressed Sensing

Abstract

By absorbing the merits of both the model- and data-driven methods, deep physics-engaged learning scheme achieves high-accuracy and interpretable image reconstruction. It has attracted growing attention and become the mainstream for inverse imaging tasks. Focusing on the image compressed sensing (CS) problem, we find the intrinsic defect of this emerging paradigm, widely implemented by deep algorithm-unrolled networks, in which more plain iterations involving real physics will bring enormous computation cost and long inference time, hindering their practical application. A novel deep Physics-guided unRolled recovery Learning (RL) framework is proposed by generalizing the traditional iterative recovery model from image domain (ID) to the high-dimensional feature domain (FD). A compact multiscale unrolling architecture is then developed to enhance the network capacity and keep real-time inference speeds. Taking two different perspectives of optimization and range-nullspace decomposition, instead of building an algorithm-specific unrolled network, we provide two implementations: PRL-PGD and PRL-RND. Experiments exhibit the significant performance and efficiency leading of PRL networks over other state-of-the-art methods with a large potential for further improvement and real application to other inverse imaging problems or optimization models.

Appendices

Appendix A: More Comparison Results

1.1 A.1 High-Level Comparisons of Deep CS Networks

Here we provide more conceptual and functional classifications, and qualitative comparisons among existing CS networks and our PRL in Table 10. We can see that the PRL not only exhibits its organic absorption of the merits of optimization algorithms and advanced network structures but also breaks through the bottlenecks including performance degradation, saturation and the surge of time cost with a more general and compact FD physical framework beyond ID unrolling. It is expected and verified to better approximate the theoretically ideal target reconstruction functions.

1.2 A.2 More Comparisons on the CBSD68 and DIV2K Benchmarks

In this subsection, we compare our PRL and the competing methods on CBSD68 (Martin et al., 2001) and DIV2K (Agustsson & Timofte, 2017) validation sets. From Table 11, we can see the consistent leading of PRL networks over the previous networks (about PSNR of 0.8–1.8dB, 1.4–2.1dB, 0.3–1.5dB and 0.7–1.6dB on Set11 (Kulkarni et al., 2016), Urban100 (Huang et al., 2015), CBSD68 and DIV2K). Figure 9 shows that PRL networks recover more details and fine edges, and can especially reconstruct the correct shapes and accurate line textures. We also find that the accuracy leading on CBSD68 is relatively less significant, which comes from the greater recovery difficulties of CBSD68 (Zhang et al., 2017a, 2021) and our default extremely unbalanced training set. More further experiments and related analyses about the training set are provided in Sect. 1.

Table 11 PSNR(dB)/SSIM comparisons of various CS methods on CBSD68 (Kulkarni et al., 2016) and DIV2K (Agustsson & Timofte, 2017). The existing physics-free, physics-engaged, and our methods are in the yellow, green and cyan backgrounds, respectively, with the corresponding highlighted best and second best results

Fig. 9

Visual comparisons on recovering two images from CBSD68 (Kulkarni et al., 2016) (top) and DIV2K (Agustsson & Timofte, 2017) (bottom) respectively under the setting of $\gamma =30\%$

Table 12 PSNR comparison and vertical increments of the trained improved network variants with seven different training configurations on Set11 (Kulkarni et al., 2016) under the setting of $\gamma =10\%$

Appendix B: More Ablation Studies and Discussions

1.1 B.1 Effect of Fully-Utilized Physics in Eliminating Degradation

Under the single-scale/path plain framework, the existing ID unrolling scheme suffers from not only the performance saturation and large inference time increase with the unfavorable noise delicacy due to its sensitive “slender” architecture, but also the training difficulty with serious degradation as our pilot experiment exhibited. It is eliminated by activating the full-channel physics engagements. Here we conduct more experiments on five variants of our PGD-unrolled improved baseline in Fig. 3d with four degrees of physics utilization to give a better understanding of the effect of physical information guidance.

Specifically, we retrain two fully-activated networks with $r\equiv 1$, $K\in \left\{ 20, 30\right\}$, $C=D\equiv 64$ and the maximized physical feature dimensionality $q=64$, and three partially-activated ones with $K=30$ and $q\in \left\{ 1,4,16\right\}$ by disabling the latter $(64-q)$ Conv kernels and activations in $\mathcal {H}^{(k)}_\text {grad}$ for comparison. Fig. 10 shows the training processes and exhibits that the 30-stage (212-layer) recovery networks are inevitably difficult to converge as expected when $q\in \left\{ 1, 4\right\}$, and the degradation is alleviated by $q=16$ and eliminated by $q=64$. Table 12 gives the final evaluations of all these trained networks on Set11 (Kulkarni et al., 2016) and demonstrates the effectiveness of the physics-feature fusion enhancement with a 0.74dB PSNR improvement in total brought by the sufficiently large q values. Note that our neural physical guidance is implemented by multiscale unrolling in FD, and can not be achieved by only the introduction of extra parameters. We also find that this plain architecture quickly goes saturated when $K\ge 20$, and obtain similar conclusions from extra trainings with $C=D\equiv 32$ and $K\in \{80,100,120\}$. These results reveal the importance of adequate physics injection in stabilizing and enhancing the learning processes of deep physics-engaged networks for improving performance.

Fig. 10

Training loss (left) and test PSNR (right) curves on Set11 (Kulkarni et al., 2016) of our five retrained improved variants with $\gamma =10\%$. The learning of 30-stage (212-layer) network is stabilized and enhanced by more feature channels of physics guidance from $\{\textbf{A}, \textbf{y}\}$

Fig. 11

Visualization of the trajectories with relative data positions based on the principal component analysis (PCA) technique about recovering an image from CBSD68 (Martin et al., 2001) by the plain OPINE-Net$^+$ (Zhang et al., 2020a), MADUN (Song et al., 2021) and our PRL networks with/without noise ${\epsilon }$ of a large standard deviation (noise level) $\sigma =50$ in observation $\textbf{y}=\textbf{Ax}+{\epsilon }$ (top/middle) when $\gamma =10\%$. The recovered results (bottom) are marked by $\sigma$ and PSNR (dB) with corresponding drops. Compared with the existing plain unrolling networks, PRL can stabilize the recovery process with heavy noise interference

Fig. 12

Visualization of the [0, 255]-normalized intermediate images/features (left) of recovering three Set11 (Kulkarni et al., 2016) images (including “Lena”, “Monarch” and “Peppers”) from the 5-/15-th stage of OPINE-Net$^+$ (top) and our default PRL-PGD (bottom) with $\sigma \in \{0,50\}$ and $\gamma =10\%$. Features from the third level of PRL-PGD are upscaled to be in 8-dimensional space and visualized channel-by-channel. Their distribution curves are on the right side. All visualizations are marked by PSNR/SSIM of the final recoveries (with drop margins) in their upper left corner. PRL-PGD recovers better under noise disturbance by preserving the image structure and feature distribution shapes based on its effective U-shaped FD unrolling

1.2 B.2 Effect of U-Shaped Residual Recovery in Improving the Noise Robustness

To better understand the reconstruction characteristics of different unrolled architectures, based on Fig. 7, we conduct further image- and feature-level visualizations on the refinement processings of various networks from the perspective of noise robustness. Figure 11 provides the restoration trajectories of four PGD-unrolled networks including the plain 9-stage OPINE-Net$^+$ (Zhang et al., 2020a) and 25-stage MADUN (Song et al., 2021) in ID, and our 30-stage PRL-PGD and PRL-RND in FD about the recovery of an image from CBSD68 (Martin et al., 2001) (with/without observation noise). Here we use the principal component analysis (PCA) instead of the more advanced t-SNE (Van der Maaten & Hinton, 2008) to reduce the image data dimensionality for all stages, since the PCA technique is simpler and more stable without hyperparameters or randomness in our evaluation. We map the high-dimensional features in PRL networks from FD to ID by $\times r$-upscaling and transforming them to be with the shape of $H\times W\times 1$ through PixelShuffle (Shi et al., 2016) and $\mathcal {H}_\text {rec}$. From the visualization in Fig. 11, we observe that compared with OPINE-Net$^+$, MADUN reconstructs better with 25-step milder content refinements in the noise-free case. Both of them are severely affected by observation noise which brings large error accumulation and a “domino” effect of amplification in such a long and slender architecture, especially the deeper ones (see the blocking/ripple/honeycomb artifacts in OPINE-Net$^+$ recovery with a 6.20dB PSNR drop and more totally distorted blocks from MADUN with a 10.56dB PSNR drop). Similar phenomena are observed in other plain networks (Zhang et al., 2020c; You et al., 2021b) and our baseline variants. PRL preserves the trajectory structures and well keeps the recovery stable under the interference of noise (with <3.50dB PSNR drops). Note that the intra-stage connections and inter-stage short-/long-term memories are introduced in OPINE-Net$^+$ and MADUN for transmission enhancements. These results demonstrate the positive effect of improving noise robustness from our compact multiscale unrolling with additive long-range skip connections, which conducts three levels of residual FD recoveries and plays a major role in stabilizing the 30-stage reconstructions since our PGD-/RND-unrolling networks share the similar recovery trajectories and are with minor differences in the path detail.

Table 13 Evaluation of the recovery accuracy in PSNR, inference times and parameter numbers of four prior networks (left) and twenty PRL-PGD variants (right) on Set11 (Kulkarni et al., 2016) when CS ratio $\gamma =50\%$

We further extract the intermediate refined images/features of three Set11 (Kulkarni et al., 2016) instances from the recoveries of the 5-/15-th stage in OPINE-Net$^+$/PRL-PGD, and visualize them in Fig. 12 with $\sigma \in \{0,50\}$. Specifically, we provide their original distribution density curves (right) and visualize stage outputs channel-by-channel (left) with [0, 255]-normalization and upscaling $(H/4)\times (W/4)\times 128$ features from PRL-PGD to the shape of $H\times W\times 8$. From Fig. 12 we observe that the ID recoveries of OPINE-Net$^+$ are with unstable signal distributions and loss of accurate structures of both the image content and distribution shapes under the noise disturbance (with an average PSNR/SSIM drop of 9.25dB/0.4448). Our FD recoveries (with a drop of 5.95dB/0.2458) decompose the images into multiple physics channels and preserve their spatial structures with smooth and stable distribution curves. We also find that with only an end-to-end $\ell _2$ loss function, OPINE-Net$^+$ can be regarded as a special FD network reconstructing a single-channel feature instead of totally working in ID since its non-linearly transformed intermediate results tend to be distributed in $[-0.3,1.2]$ (not the original [0, 1] for ID). PRL-PGD recovers better in 8-dimensional space with wider and sparser zero-mean distributions which may be helpful to enhance the edges and textures. These results validate the necessity of our FD restoration for supporting the high-throughput transmissions and sufficient physics guidance with leaving large freedom degrees for information distributions that makes PRL robust.

Fig. 13

Visualization of the LAM (Gu & Dong, 2021) attribution results of PRL$^*$-PGD and PRL-PGD with $K=1$ and $C=D=8$ about the recovery of two images named “Boats” (top) and “Barbara” (bottom) from Set11 (Kulkarni et al., 2016) when $\gamma =10\%$. The LAM maps represent the importance of different pixels for the selected central $16\times 16$ local patches and the more informative regions are marked with darker red color. a is the original image with the selected region marked in a red box; b is the LAM result; c is the LAM result with input image; d is the informative area with input image; e is the CS reconstructed result. The upper/lower rows of each instance correspond to the single-/multi-scale PRL$^*$/PRL, respectively. Their recovery PSNR (dB) with accuracy distances are exhibited on the right side. One can see that PRL produces much more accurate recovered outputs by significantly increasing the range of pivotal information utilization

Table 14 PSNR(dB)/SSIM comparisons among our default PGD-unrolled PRL networks with three different training set configurations (a), (b) and (c) on four benchmarks in the case of $\gamma =50\%$

1.3 B.3 Study of FD Dimensionality and the Inter-Stage Weight Sharing Strategy

We evaluate 20 PRL-PGD variants with the FD dimensionalities $d\in \{2,4,8,16\}$, the stage numbers of each group $K\in \{1,3,5,7\}$ and the weight sharing across all the stages in each group with $K=5$, and report their recovery accuracy, speeds, and parameter numbers when CS ratio $\gamma =50\%$ in Table 13 with four competing unrolled methods (Zhang et al., 2020a, c; You et al., 2021b; Song et al., 2021). We observe that the PRL performance gives full play when $d\in \{8,16\}$ and is restricted by $d\in \{2,4\}$ with poor feature capacity, so a larger FD dimensionality can sufficiently stimulate their high performance. The inter-stage weight sharing strategy brings little impact on recovery accuracy (only 0.04dB PSNR drop on average) and speed (0.506ms time saving) with about $\times (1/K)$ parameter reduction. It benefits from the compact and efficient PRL structure and is quite effective and considerable in deployments. The PRL-PGD variants most comparable to other methods (in parameter number, depth and speed) bring the average improvements including PSNR leading of 0.66dB, 16% time reduction and 26% parameter number saving under such a high sampling rate. Note that these PRL variants are limited by insufficient configurations but the competing networks are their own optimal and saturated versions. These results verify the higher structural efficiency and appealing compression potential of PRL with its flexible and competent optional settings for real-world applications.

1.4 B.4 Effect of Multiscale Architecture and Large Training Set in Improving Recovery Accuracy

One important characteristic of PRL different from existing unrolled networks is the multiscale design that not only speeds up inference but also provides inter-scale FD communications for capturing the image context and structure. With a parameter number close to PRL, the single-scale U-shaped PRL$^*$ has a receptive field size of $R_{\text {PRL}^*}=(84K+13)^2=(2L+1)^2$ in its longest fully-convolutional path, where K is group stage number and L is the $3\times 3$ Conv layer number or network depth. But PRL has a much larger size of $R_{\text {PRL}}=(140K+8)^2=(10L/3-12)^2$ based on the introduction of $2\times 2$ strided/transposed (S-/T-) Conv layers, which leads to $R_{\text {PRL}}=(5/3)^2 R_{\text {PRL}^*}$ when $L\rightarrow \infty$. It is quite important to reconstruct a specific region by utilizing a larger informative image part (Zhang et al., 2017b). Here we analyze the effect of the receptive field differences between the single-/multi-scale PRL variants by using the local attribution map (LAM) (Gu & Dong, 2021) to compare the input image parts that have a strong influence on the CS recovered outputs when recovering two images from Set11 (Kulkarni et al., 2016) with $\gamma =10\%$. The two LAM visualization results in Fig. 13 indicate that both PGD-unrolled PRL and PRL$^*$ perform the sampling and recovery for the selected local patch with some edges and textures by intensively utilizing a region with rich and relevant information. PRL significantly takes notice of a wider range of informative clean areas to guide its more accurate recoveries and achieves an average PSNR gain of 0.36dB over the PRL$^*$, thus revealing the stronger information capture capability and higher structural efficiency of our multiscale design. We also find that due to the block-based CS scheme (with a globally unified block size of $32\times 32$) and the iterative stage-by-stage physical knowledge injections of $\{\textbf{A}, \textbf{y}\}$, the actual network receptive field is larger than their fully-convolutional paths as we calculated above. The most heavily utilized informative regions are with complicated and content-correlated specific spatial shapes and distributions.

To enhance the training sample diversity, we follow (Zhang et al., 2021; You et al., 2021a; Song et al., 2021) to use four datasets including T91 (Yang et al., 2008; Dong et al., 2014a), Train400 (Chen & Pock, 2016; Zhang et al., 2017a), DIV2K (Agustsson & Timofte, 2017) training set and the Waterloo exploration database (WED) (Ma et al., 2016). They cover a large image space and can enrich the network prior to many classic image restoration tasks (Zhang et al., 2021). Here we conduct more PRL-PGD trainings with three different dataset settings, report their evaluation results in Table 14, and observe from which that our network trainings on only the Train400 dataset can bring the consistent recovery accuracy average leading of PRL (PGD-unrolled versions) over the unrolled network MADUN (Song et al., 2021) in PSNR/SSIM of about 0.77dB/0.0008, 1.56dB/0.0036, 0.47dB/0.0009 and 0.71dB/0.0011 on Set11 (Kulkarni et al., 2016), Urban100 (Huang et al., 2015), CBSD68 (Martin et al., 2001) and DIV2K validation set. The introductions of T91/DIV2K and WED generally bring negative and positive effects in most benchmarks, respectively. The combination of them provides the average accuracy changes of +0.13dB/+0.0006, +0.10dB/+0.0007, −0.14dB/-0.0005, and +0.08dB/+0.0003. Especially, WED overturns the negative impacts of T91/DIV2K in Set11 and Urban100 with appreciable improvements. We also find that the recovery performance on CBSD68 can not directly benefit from the larger training set due to its unbalanced sample distribution (only $6.6\%$ training images are $180\times 180$ grayscale BSD patches), thus verifying the data bias limitation that restricts network learning and causes the overfitting on PRL variants with large capacities.