PE-CSNet: An equivariant network architecture with learnable patch-based sparse representation
Authors: Kai Li, Haitao Long, Bo Zhang, Haiwen Zhang, Zhi Zhou
Organizations: Department of Applied Mathematics, The Hong Kong Polytechnic University, Hung Hom, Hong Kong. · SKLMS and Academy of Mathematics and Systems Science, Chinese Academy of Sciences, Beijing 100190, China and School of Mathematical Sciences, University of Chinese Academy of Sciences, Beijing 100049, China · SKLMS and Academy of Mathematics and Systems Science, Chinese Academy of Sciences, Beijing 100190, China
Compressive sensing (CS) enables accurate signal reconstruction from sparse measurements and is widely applied in medical imaging, remote sensing, and image compression. However, designing an effective, task-specific sparse transform and the corresponding optimization procedure for high-quality CS remains challenging. This process typically requires expert domain knowledge and laborious parameter tuning. To address this issue, we present a Patch-based Equivariant deep unrolling architecture, termed PE-CSNet, for accurate CS recovery. While traditional CS methods generally use predefined patch-based transform sparsity, we generalize this idea by incorporating learnable transform sparsity that adapts to the specific CS task through an optimization-driven process. Specifically, we first establish a generalized patch-based CS model, which we solve via a block coordinate descent (BCD) algorithm. The BCD solver is then unrolled into a deep neural network, where all parameters of both the CS model and solver are learned through end-to-end training. To improve data efficiency, we introduce a stochastic equivariant training strategy that exploits the patch-wise structure of the network, enabling PE-CSNet to learn effectively even from limited data. We further provide a simpler, parameter-shared version of PE-CSNet and briefly discuss its convergence as an iterative solver. For practical applications, the network uses stage-specific (non-shared) parameters to enhance its expressive power and thereby improve its performance. On the tasks of CS magnetic resonance imaging (CS-MRI) and CS coded diffraction patterns (CS-CDP), PE-CSNet achieves state-of-the-art accuracy with fast computational speed, outperforming traditional methods and existing deep unrolling methods.
Figures & tables
Figure 1: Visualization of the patch-extraction operator Dij for a (4×4) patch grid ( Np=4 ).
Figure 2: The architecture of PE-CSNet. The top part shows the overall unrolled architecture, which alternates between the reconstruction module (X(n)) and the auxiliary patch update module (Zij(n)) . The bottom part provides a detailed view of the residual U-Net that implements (Zij(n)) .
Figure 3: Visualization of the undersampling masks used in CS-MRI experiments. White and black pixels represent sampled and unsampled k -space locations, respectively. (a) 1D Cartesian mask; (b) 2D random mask.
Figure 4: Reconstruction results of PE-CSNet on CS-MRI. (a) Results for 1D Cartesian masks at 20%, 30%, and 40% sampling rates. (b) Results for 2D random masks at 5%, 10%, and 20% sampling rates. For each case, the magnitude image of the reconstruction, the corresponding error map (magnitude of the complex difference), a detailed crop, and the magnitude image of the ground truth are displayed.
Figure 5: Empirical convergence analysis of PE-CSNet on CS-MRI for 1D Cartesian masks. (a) Stage-wise reconstruction outputs for the 1D Cartesian mask at 40% sampling rate. Magnitude images of the intermediate outputs x(n)(n=0,1,…,4) are shown, followed by the ground truth. Corresponding PSNR values are annotated below each output. (b) NRMSE convergence curves for a representative sample under 1D Cartesian masks at 20%, 30%, and 40% sampling rates. (c) Mean NRMSE ±1 standard deviation (shaded region) across the entire test set for the 1D Cartesian mask at the sampling rate of 30%.
Figure 6: Empirical convergence analysis of PE-CSNet on CS-MRI for 2D random masks. (a) Stage-wise reconstruction outputs for the 2D random mask at 20% sampling rate. Magnitude images of the intermediate outputs x(n)(n=0,1,…,4) are shown, followed by the ground truth. Corresponding PSNR values are annotated below each output. (b) NRMSE convergence curves for a representative sample under 2D random masks at 5%, 10%, and 20% sampling rates. (c) Mean NRMSE ±1 standard deviation (shaded region) across the entire test set for the 2D random mask at the sampling rate of 10%.
Method
20%
30%
40%
Time
NRMSE
PSNR
SSIM
NRMSE
PSNR
SSIM
NRMSE
PSNR
SSIM
CPU/GPU
Zero-filling [ 5 ]
27.43
25.93
61.93
21.68
28.10
70.38
20.01
28.62
70.96
2ms/-
PANO [ 47 ]
20.24
28.55
73.78
16.17
30.48
80.25
12.12
33.02
83.48
-/-
ISTA-Net [ 61 ]
16.64
31.70
80.59
11.21
34.39
86.71
8.16
37.10
90.65
-/8ms
ADMM-CSNet [ 57 ]
15.98
31.38
81.15
11.19
34.40
87.50
8.08
37.27
90.76
-/458ms
PE-CSNet
14.69
31.87
81.68
10.33
34.94
88.02
7.82
37.53
91.44
87ms/7ms
Table 1: Comparison of average reconstruction quality (NRMSE ( ×10−2 ), PSNR in dB, SSIM ( ×10−2 )) and average inference times per sample on CS-MRI for 1D Cartesian masks with sampling rates of 20%, 30%, and 40%. Best performance in each column is highlighted in bold.
Method
5%
10%
20%
Time
NRMSE
PSNR
SSIM
NRMSE
PSNR
SSIM
NRMSE
PSNR
SSIM
CPU/GPU
Zero-filling [ 5 ]
34.34
23.95
51.11
29.83
25.37
56.29
25.73
26.73
60.45
2ms/-
PANO [ 47 ]
24.44
27.21
66.68
19.79
29.12
72.34
15.11
31.22
76.91
-/-
ISTA-Net [ 61 ]
19.81
29.70
70.20
15.61
31.69
76.45
10.54
34.52
82.77
-/8ms
ADMM-CSNet [ 57 ]
20.07
29.52
73.40
15.04
31.90
78.16
10.53
34.73
83.47
-/458ms
PE-CSNet
18.24
30.19
75.54
13.91
32.33
79.65
10.02
35.04
84.86
82ms/7ms
Table 2: Comparison of average reconstruction quality (NRMSE ( ×10−2 ), PSNR (dB), SSIM ( ×10−2 )) and average inference times per sample on CS-MRI for 2D random masks with sampling rates of 5%, 10%, and 20%. Best performance in each column is highlighted in bold.
Figure 7: Visual comparison of reconstructed magnitude images (and corresponding zoomed-in views) from different CS-MRI methods under a 1D Cartesian mask (30% sampling rate). Red boxes indicate the regions of interest for detailed visual comparison.
Figure 8: Reconstruction results of PE-CSNet on CS-CDP at 5%, 10%, and 20% sampling rates. For each reconstruction, the corresponding error map, a detailed crop, and the corresponding ground truth are displayed.
Figure 9: Empirical convergence analysis of PE-CSNet for CS-CDP. (a) Stage-wise reconstruction outputs for a 20% sampling rate. The intermediate outputs x(n)(n=0,1,…,6) are shown, followed by the ground truth. Corresponding PSNR values are annotated below each output. (b) NRMSE convergence curves for a representative sample under 5%, 10%, and 20% sampling rates. (c) Mean NRMSE ±1 standard deviation (shaded region) across the entire test set for the sampling rate of 10%.
Method
5%
10%
20%
Time
NRMSE
PSNR
SSIM
NRMSE
PSNR
SSIM
NRMSE
PSNR
SSIM
GPU
ISTA-Net [ 61 ]
15.41
21.42
55.02
7.72
27.43
80.34
4.38
32.43
90.54
7ms
ADMM-CSNet [ 57 ]
11.70
23.78
63.07
6.79
28.55
79.06
3.76
33.87
90.43
301ms
PE-CSNet
8.34
26.74
78.15
5.10
31.10
89.08
2.80
36.50
95.49
11ms
Table 3: Comparison of average reconstruction quality (NRMSE ( ×10−2 ), PSNR (dB), and SSIM ( ×10−2 )) and average inference times per sample on CS-CDP for sampling rates of 5%, 10%, and 20%. Best performance in each column is highlighted in bold.
Figure 10: Visual comparison of reconstructions (and corresponding zoomed-in views) from different CS-CDP methods under sampling rates of 5%, 10%, and 20%. Red boxes indicate the regions of interest for detailed visual comparison.
Method
20%
30%
40%
NRMSE
PSNR
SSIM
NRMSE
PSNR
SSIM
NRMSE
PSNR
SSIM
PE-CSNet (shared)
17.01
30.68
78.97
12.00
33.47
84.85
8.84
36.35
89.02
PE-CSNet (global)
16.35
31.14
80.04
11.40
34.07
85.72
8.31
36.94
89.89
PE-CSNet
14.69
31.87
81.68
10.33
34.94
88.02
7.82
37.53
91.44
Table 4: Comparisons of average reconstruction quality (NRMSE ( ×10−2 ), PSNR (dB), and SSIM ( ×10−2 )) and average inference times per sample on CS-MRI for 1D Cartesian masks at 20%, 30%, and 40% sampling rates. Best performance in each column is highlighted in bold.
Deep unfolding turns iterative algorithms into trainable networks, combining model-based structure with learned components for efficient inference and optimization.