PUERT: Probabilistic Under-sampling and Explicable Reconstruction Network for CS-MRI

Nested optimization methods formulate the sampling matrix learning problem as two nested optimization issues. The outer issue is to find an optimal sub-sampling mask or trajectory that behaves best under the current reconstruction algorithm, which is then taken as input for the inner problem to further optimize the reconstruction method. Several heuristic or greedy algorithms [43, 44, 45, 46, 47, 48, 49, 50, 51] have been proposed to tackle the above nested pair of problems, and has been demonstrated to achieve better reconstructions than conventional methods [50, 52] in an empirical study [53]. Most recently, in [54], continuous optimization methods are applied to a formulated bilevel optimization problem, with the aim of learning sparse sampling patterns for MRI within a supervised learning framework for a given variational reconstruction method. However, the above algorithms suffer from high computational complexity when dealing with large-scale training data, and often lack the flexibility to explore different types of sampling patterns.

Fueled by the rise of deep learning, some end-to-end data-driven sampling optimization methods [55, 56, 57, 58, 59] are proposed to achieve computational efficiency and improve reconstruction. PILOT [57] introduces hardware constraints into the learning pipeline and explores an optimal physically viable $k$-space trajectory that is enabled by interpolation on the discrete $k$-space grid. However, the learned trajectory significantly depends on the initialization and lacks exploration. J-MoDL [59] proposes a multichannel sampling model consisting of a non-uniform Fourier operator with continuously defined sampling locations to promote the optimization process without approximations. However, it constrains the sampling mask as the tensor product of two 1D sampling patterns and restricts its performance. Note that, due to learning a deterministic mask, the above methods inhibit flexibility and randomness, and are against the stochastic strategy of CS theory [60, 61].

LOUPE [56] assumes that each binary sampling location is an independent Bernoulli random variable and learns a probabilistic sampling pattern instead of a deterministic mask. In this aspect, both LOUPE and our PUERT have similar inspirations, but in fact, our PUERT enjoys three superior advantages compared to LOUPE. First, due to the relaxation of the binarization operation, LOUPE not only pays a performance penalty, but also needs to retrain the reconstruction model with the learned binary mask, whereas our PUERT directly uses the binarization function in the forward pass and introduces an efficient gradient estimation strategy for the backward pass, which overcomes the above two shortcomings and is shown to promote the recovery quality. Second, LOUPE mainly adopts a U-Net network and does not point out the superiority of DUNs against non-DUNs on further exploitation of the sampling subnet. However, we emphasize that, thanks to DUN’s intrinsic structural feature of fully utilizing knowledge on the mask in each stage, it’s of great significance to use DUNs to facilitate exploration and training of the sampling subnet. Third, when testing, LOUPE directly samples from the learned probabilistic pattern and gets a mask with an inexact sampling ratio, whereas in PUERT, the vanilla binarization function is replaced by a greedy version, thus achieving precise control of sampling ratio and promoting fair comparisons.

In CS-MRI, the acquisition process can be formulated as:

where $\mathbf{x}\in\mathbb{C}^{N_{x}\times N_{y}}$ and $\mathbf{y}\in\mathbb{C}^{N_{x}\times N_{y}}$ denote the original fully-sampled image to be reconstructed and the under-sampled $k$-space observation, respectively. $\mathbf{\varepsilon}\in\mathbb{C}^{N_{x}\times N_{y}}$ is the noise generated during acquisition. $\mathbf{F}$ and $\odot$ denote the Fourier Transform (FT) and the element-wise multiplication operation, respectively. $\mathbf{M}\in\mathbb{R}^{N_{x}\times N_{y}}$ is a binary sampling mask matrix composed of 1 and 0, which separately represent whether to sample or not at the corresponding $k$-space position. Note that the CS sampling ratio, denoted by $\alpha$, is defined as the ratio of value 1 in $\mathbf{M}$, i.e., $\alpha=\frac{\left\|\mathbf{M}\right\|_{0}}{{N_{x}\times N_{y}}}$.

To tackle the ill-posed inversion of reconstructing $\mathbf{x}$ from sub-sampled data $\mathbf{y}$ under a given mask $\mathbf{M}$, traditional CS-MRI reconstruction methods usually reconstruct the original image $\mathbf{x}$ by solving the following optimization problem:

where $\psi(\mathbf{x})$ denotes an image prior-regularized term and $\lambda$ is a weight to balance between fidelity and regularization. Despite the broad task defined above, note that we only consider partial Fourier reconstruction here, excluding parallel imaging [62, 63, 6] and structured low-rank matrix methods [64, 65, 66].

In pursuit of superior recovery performance, we propose to concurrently explore an optimal sampling pattern while learning reconstruction network parameters. To this end, we take $\mathbf{M}$ as learnable parameters and construct a novel end-to-end Probabilistic Under-sampling and Explicable Reconstruction neTwork (PUERT), which is composed of a sampling subnet and a reconstruction subnet, as illustrated in Fig. 1.

In the Sampling Subnet (SampNet), we first explore a Probabilistic Under-sampling (PU) scheme to preserve robustness and stochastics, and then adopt a Dynamic Gradient Estimation (DGE) strategy to enable efficient training and promote network performance.

Probabilistic Under-sampling (PU): Building on stochastic strategies and the robustness requirement of compressed sensing, stochastic sub-sampling patterns are able to create noise-like artifacts that are relatively easier to remove [2]. Therefore, we propose to directly optimize a probabilistic sampling pattern rather than a fixed binary mask. Furthermore, considering that adopting a classic sampling pattern (e.g., Gaussian [4] or Uniform [2] distribution) exhibits great limitation and lacks adaptability, we propose to learn independent Bernoulli random variables for each $k$-space point, thus learning a sampling pattern highly customized to specific training data and recovery methods.

Concretely, a learnable probabilistic sampling pattern $\mathbf{P}\in\mathbb{R}^{N_{x}\times N_{y}}$ is introduced and each value $\mathbf{P}_{i,j}\in[0,1]$ stands for the probability of $\mathbf{M}_{i,j}$ taking value 1, i.e., $\mathbf{M}_{i,j}\sim\mathcal{B}(\mathbf{P}_{i,j})$, where $\mathcal{B}(z)$ denotes a Bernoulli random variable with parameter $z$. In order to draw a sample from $\mathbf{P}$ and generate a binary mask $\mathbf{M}$, we introduce a matrix $\mathbf{U}\in\mathbb{R}^{N_{x}\times N_{y}}$ with each element independently drawn from a Uniform distribution on $[0,1]$, i.e., $\mathbf{U}_{i,j}\sim\mathcal{U}(0,1)$, and then binarize the difference between $\mathbf{P}_{i,j}$ and $\mathbf{U}_{i,j}$ by a function $Bina$. So the generation of $\mathbf{M}$ is formulated as:

where $Bina$ is implemented by a vanilla binarization function $Bina_{v}$ in the training process:

Note that $\mathbf{U}$ is not fixed, but randomly generated every time $\mathbf{P}$ generates $\mathbf{M}$. In this way, the sampling subnet is able to learn a probabilistic sampling pattern $\mathbf{P}$ that expresses belief or importance across all $k$-space locations, thus gaining robustness and stochastics for a more reliable CS-MRI reconstruction. Note that the use of $\mathbf{U}$ is similar to the re-parameterization trick used in VAE[67], which generates non-uniform random numbers by transforming some base distribution, so as to recast the statistical expression and tackle the stochastic node.

Another critical issue is how to control the sparsity of the generated mask $\mathbf{M}$ at a given target sampling ratio $\alpha$. To this end, we first introduce a rescale operator to adjust the average value of $\mathbf{P}$, and then adopt a greedy binarization operator specifically for testing to implement a totally accurate control.

To be concrete, before sampling out $\mathbf{M}$ as Eq. (3), the probabilistic sampling pattern $\mathbf{P}$ is rescaled as follows:

Here, $\bar{p}$ stands for the average value of $\mathbf{P}$, i.e., $\bar{p}=\frac{\sum_{i,j}{\mathbf{P}_{i,j}}}{{N_{x}\times N_{y}}}$. It can be proven that Eq. (5) yields $\mathbf{P}$ with its average value rescaled to the given CS ratio $\alpha$. With this rule, the ratio of the generated binary mask $\mathbf{M}$ would be close to the target $\alpha$.

During the testing process, in order to achieve an accurate control of the ratio of $\mathbf{M}$, a greedy binarization function $Bina_{g}$ is further utilized to replace the vanilla one $Bina_{v}$ in Eq. (3). Specifically, the greedy binarization function $Bina_{g}$ used during testing is defined as follows:

where $\mathbf{\Omega}$ represents the set $\left\{\mathbf{P}_{i,j}-\mathbf{U}_{i,j}\right\}$ and $b(\mathbf{\Omega},\alpha)$ stands for the $(\alpha\times\left|\mathbf{\Omega}\right|)$-$th$ largest number in the set $\mathbf{\Omega}$. In this way, our method precisely controls the ratio of $\mathbf{M}$ and implements a fair comparison with other methods. Note that since the learned probabilistic pattern $\mathbf{P}$ represents importance across all $k$-space locations, our value-maximization design as Eq. (6) for the testing process is actually cooperative with the training period and exhibits nice rationality. Besides, thanks to the adoption of $\mathbf{U}$ in Eq. (6), our method outputs a sampling mask with randomness rather than a fixed mask.

Dynamic Gradient Estimation (DGE): In order to make the above sampling subnet trainable, inspired by recent works in Binarized Neural Networks (BNNs) [68, 69], a dynamic gradient estimation strategy is further introduced for the vanilla binarization function $Bina_{v}$ in Eq. (4). To be specific, we adopt the following dynamic function $g$ as a progressive approximation of $Bina_{v}$ during backward propagation:

whose derivative is used as the backward gradient of $Bina_{v}$. Here $t$ and $k$ are control variables and change as follows during the training period:

where $i$ is the current epoch and $N_{e}$ is the number of epochs, $T_{\min}=10^{-1}$ and $T_{\max}=10^{1}$. Note that $t$ and $k$ controls the slope and output range of function $g$, respectively.

Essentially, this is a progressive two-stage strategy to approximate the vanilla binarization function $Bina_{v}$. As shown in Fig. 2, during Stage 1, i.e., the early stage of training, we focus more on the updating ability and speed of the backward propagation algorithm. Therefore, the gradient estimation function’s derivative value is kept close to one, and then the output range is reduced progressively from a large scope to $[0,1]$. As for Stage 2, we lay more importance on retaining accurate gradients for parameters around zero. Consequently, we keep the function output range as $[0,1]$ and gradually increase the slope so as to push the estimation curse to the shape of the binarization function $Bina_{v}$, thus achieving the consistency of forward and backward propagation during the late stage of training. Based on the above two-stage scheme, our proposed DGE gradually approximates the binarization function in backward propagation and reasonably updates all parameters, which is shown to efficiently retain the gradient information and promote network performance.

In the Reconstruction Subnet (RecNet), without losing generality, we construct a deep unfolding network by casting the traditional ISTA method into a deep network form, so as to achieve simplicity, effectiveness and interpretability.

The classic CS-MRI optimization problem Eq. (2) can be efficiently solved with ISTA by iterating the following two update steps:

Here, $k$ denotes the iteration index, $\rho$ is the step size, $\mathbf{F}^{H}$ denotes the Inverse Fourier Transform (IFT) and the proximal mapping operator of regularizer $\psi$ is defined as $\mathtt{prox}_{\lambda\psi}(\mathbf{r})=\arg\min_{\mathbf{x}}\frac{1}{2}||\mathbf{x}-\mathbf{r}||^{2}_{2}+\lambda\psi(\mathbf{x})$. Eq. (9) and Eq. (10) are usually called gradient descent step (GDS) and proximal mapping step (PMS), respectively.

Following [24], our ISTA-unfolding reconstruction subnet consists of $K$ stages and each stage corresponds to one iteration in ISTA. Concretely, each stage is composed of a gradient descent module (GDM) and a proximal mapping module (PMM), which correspond to the above two update steps Eq. (9) and Eq. (10), respectively.

For GDM, to preserve the ISTA structure while increasing network flexibility, the step size $\rho$ is allowed to vary across iterations, and GDS, i.e., Eq. (9) is casted as follows:

As for PMM in each stage, we propose a simple yet effective module as shown in Fig. 1, which can be formulated as:

Here, PMM consists of 2 residual blocks (RBs) $\mathcal{H}_{k}^{RB,1}$ and $\mathcal{H}_{k}^{RB,2}$, two convolution layers $\mathcal{H}^{ext}_{k}$, $\mathcal{H}^{rec}_{k}$, which devote to extracting the image features and reconstruction, respectively, and a long residual connection with input $\mathbf{r}_{k}$.

Given the training dataset with full-resolution images $\{\mathbf{x}^{(i)}\}_{i=1}^{N_{D}}$ , the sampling subnet first uses a learnable sampling mask $\mathbf{M}$ to get simulation measurements $\mathbf{y}^{(i)}=\mathbf{M}\odot\mathbf{F}\mathbf{x}^{(i)}$, and then, with initialization $\mathbf{x}^{(i)}_{0}=\mathbf{F}^{H}\mathbf{y}^{(i)}$ as input, the reconstruction subnet outputs the recovery results, aiming to reduce the discrepancy between $\mathbf{x}^{(i)}$ and $\mathbf{\hat{x}}^{(i)}_{K}$. Therefore, we use the following end-to-end loss function to train our PUERT:

where $N_{D}$ is the number of training images and $N$ is the size of each image $\mathbf{x}^{(i)}$, i.e., $N=N_{x}\times N_{y}$. $\mathbf{\Theta}$ denotes the learnable parameter set in the two subnets of PUERT. Note that, in the sampling subnet, the learnable parameters are not directly the probabilistic sampling pattern $\mathbf{P}$, but an unconstrained image $\mathbf{O}$ with the same size, such that $\mathbf{P}_{i,j}=\sigma(\mathbf{O}_{i,j})$ to realize $\mathbf{P}_{i,j}\in(0,1)$. Here, $\sigma$ denotes the sigmoid fuction with the slope set be to 5. It is also worth emphasizing that, thanks to the rescale operator as Eq. (5), a loss term to constrain the sparsity ratio of the learned sampling mask is unnecessary.

To validate the effectiveness of our proposed sampling pattern optimization scheme, i.e., our SampNet, we compare our learnable masks with four types of classic fixed masks under the conditions of four reconstruction methods.

• •

  For classic fixed masks, we consider four categories that are widely used in the literature: Cartesian [3] with skipped lines (dubbed VD-1D), pseudo Radial [25], Random Uniform [2] and Variable Density under 2D (dubbed VD-2D) [4] masks, where the first category is 1D sub-sampling mask and the last three are 2D masks. Note that a fixed calibration region of size $32\times 32$ is adopted in the center of the $k$-space for Uniform and VD-2D masks in order to yield better recovery performance.

  Ratios	Methods	1D Settings		2D Settings
  		VD-1D	SampNet	Radial	Uniform	VD-2D	SampNet
  5%	PANO [75]	17.51	28.53	27.70	30.26	33.04	34.14
  	BM3D-MRI [76]	25.43	28.72	27.78	30.77	34.26	35.27
  	U-Net [72]	30.03	32.09	30.49	32.91	34.66	35.65
  	RecNet	30.64	32.26	31.02	33.00	35.79	36.74
  10%	PANO [75]	31.66	32.79	31.93	32.69	36.26	36.75
  	BM3D-MRI [76]	32.06	33.65	33.19	33.83	37.58	38.09
  	U-Net [72]	33.32	33.68	33.45	33.82	37.31	38.03
  	RecNet	34.22	35.51	34.88	36.03	38.45	39.13

  	Example Slices	Optimized under 5%	Optimized under 10%
  Knee
  Brain

• •

  For reconstruction methods, we consider two traditional model-based methods, i.e., PANO [75] and BM3D-MRI [76], and one widely used deep learning model called U-Net [72]. Our reconstruction subnet (RecNet) is also extracted as a recovery model for fair comparisons.

In this experiment, the aforementioned four classic masks as well as our proposed SampNet, including 1D and 2D versions, are experimented under the above four reconstruction methods, thus leading to 24 possible combinations. Note that when combining the two traditional methods PANO and BM3D-MRI with our SampNet, we use the fixed sampling masks learned by our PUERT to conduct the reconstruction. While for U-Net, it is combined with our proposed SampNet to jointly optimize the parameters for under-sampling and reconstruction. Note that for the 1D setting, our SampNet adopts a Bernoulli probabilistic pattern $\mathbf{P}\in\mathbb{R}^{N_{x}}$ and generates a binary vector $\mathbf{V}\in\mathbb{R}^{N_{x}}$, i.e., $\mathbf{V}_{i}\sim\mathcal{B}(\mathbf{P}_{i})$, which is then expanded on the column dimension to a 1D sampling mask $\mathbf{M}\in\mathbb{R}^{N_{x}\times N_{y}}$.

Fig. 3 shows the visual comparisons among the classic masks and our optimized sub-sampling matrices trained on Brain dataset for two different ratios $\alpha=5\%$ and $\alpha=10\%$. One can see that the learned 2D masks are similar to VD-2D masks, both showing a strong preference for lower frequencies and exhibiting a denser sampling pattern closer to the origin of $k$-space. While for high frequency values, compared with VD-2D masks, our optimized masks show a relatively smaller density when $\alpha=5\%$ but larger when $\alpha=10\%$, which demonstrates that our PUERT enjoys the ability to adaptively optimize sub-sampling patterns for specific ratios.

Furthermore, the optimized masks by our PUERT for the knee and brain anatomies are compared in Fig. 5. We observed that, for the knee MR images, more attention has been paid to lateral frequencies rather than ventral frequencies, owing to the unique asymmetric feature of the knee anatomy, where there is dramatically more tissue contrast in the lateral direction. The masks learned on Brain dataset, on the other hand, exhibit a more radially symmetric feature. This comparison highlights the importance of adapting the learnable under-sampling pattern to the anatomy and data at hand.

Fig. 4 shows box plots for PSNR values of reconstructions obtained with four reconstruction methods, six different sampling schemes, and two sub-sampling ratios on Brain dataset. One can intuitively observe that overall our proposed SampNet yields attractively better results than competing masks when combined with any of the four reconstruction methods and under any of the two CS ratios, which demonstrates the superiority of our proposed sub-sampling learning scheme. It is worth emphasizing that U-Net with our proposed SampNet achieves higher PSNR than competing classic masks, thus validating that the proposed PUERT enjoys the generality of being extended to other reconstruction networks.

Table I lists the concrete PSNR values of Fig. 4. We can observe that, benefiting from our efficient probabilistic under-sampling scheme and dynamic gradient estimation strategy, the proposed SampNet achieves the highest PSNR results among all competing masks. Taking RecNet under ratio 10% for 2D setting as an example, our SampNet achieves remarkable 0.68 dB PSNR gains over the state-of-the-art classic VD-2D masks. It is also noteworthy that, when equipped with our SampNet, U-Net still performs worse than our RecNet, thus demonstrating that, compared with the classic data-driven reconstruction model, a deep unfolding network is more suitable to facilitate exploration and efficient training of our SampNet. We attribute such superiority to its intrinsic structural feature of fully utilizing knowledge on the sampling mask in each stage of the deep unfolding network.

Fig. 6 further shows the visual reconstruction comparisons of different masks under RecNet at ratio 5% for 1D and 2D settings on Brain dataset. As can be intuitively appreciated from the two test images, the PUERT is able to consistently recover more details and much sharper edges than other competing masks, even under the aggressive 5% sampling ratio, thus validating the effectiveness, efficiency and practicability of our proposed PUERT.

We compare our proposed PUERT (including 1D and 2D versions) with six representative state-of-the-art methods, namely UNet-DC [16], ADMMNet [25], ISTA-Net [24], RDN [20], CDDN [19] and LOUPE [56]. The first five methods are reconstruction networks trained under some fixed sampling mask, while the last method jointly optimizes the learnable sampling pattern and the reconstruction network parameters. Following [25] and [72], for those trained under some fixed mask, we adopt Radial masks (2D) and Cartesian masks with skipped lines (1D) on Brain and FastMRI datasets respectively. Note that LOUPE is also able to implement both 1D and 2D sub-sampling optimization schemes.

The average PSNR/SSIM performance reconstructions of various methods on two datasets with respect to three CS ratios are summarized in Table  II. One can observe that methods with learnable sampling patterns, especially our PUERT, overall yield higher PSNR than those with fixed masks, which verifies the superiority of sampling pattern optimization schemes. Concretely, for Brain dataset, our PUERT-1D and PUERT-2D achieve on average 0.55 dB and 4.38 dB PSNR gain over the state-of-the-art method CDDN respectively. It is noteworthy that the Radial masks used for Brain dataset belong to a 2D setting, whereas our PUERT under the 1D setting still enjoys superiority by a large margin. As for FastMRI dataset, our PUERT-1D and PUERT-2D achieve on average 2.32 dB and 3.98 dB PSNR gain over the state-of-the-art method CDDN respectively.

When compared with the sampling pattern optimization scheme LOUPE, one can observe that, on FastMRI dataset, LOUPE-2D has a slightly better SSIM while PUERT-2D obtains a remarkably higher PSNR on average. As for Brain dataset, the proposed PUERT-2D consistently produces higher PSNR and SSIM results than LOUPE-2D across three CS ratios. The performance advantage of PUERT under 2D is also shown under 1D. The above results verify the superiority of the proposed PUERT against LOUPE. Note that, LOUPE mainly adopts the U-Net architecture as the recovery network, it also shows an example of reconstructing with a DUN called CascadeNet [18, 56], with an observation that the performance metrics are very close between different reconstruction networks. However, we highlight the importance of adopting DUNs to facilitate exploration and efficient training of the sampling subnet, owing to the intrinsic structural feature of fully utilizing knowledge on the sampling mask. More discussions on the importance of DUN are shown in Section IV-E.

As shown in Table II, six representative state-of-the-art methods are based on neural networks, and thus achieve a real-time reconstruction speed, with the inference GPU time less then 0.1s. Among the five methods trained under some fixed mask, CDDN needs the largest amount of inference time, owing to the dense blocks and dilated convolutions adopted in the network. Besides, compared with LOUPE, our PUERT achieves a comparable inference speed on average, but with higher PSNR in reconstruction accuracy, thus demonstrating the efficiency and superiority of the proposed PUERT.

Furthermore, visual comparisons of all the competing methods on FastMRI dataset are shown in Fig. 7 with CS ratio $\alpha=10\%$. Obviously, the proposed PUERT is able to produce more faithful and clearer results than the other competitive methods. Overall, the above experiments on two widely used MRI datasets demonstrate that our proposed PUERT performs favorably against state-of-the-arts in terms of both quantitative metrics and visual quality.

In this section, we present an ablation study on Probabilistic Under-Sampling (PU) in our sampling subnet, so as to emphasize the importance of learning a probabilistic sampling pattern rather than a deterministic mask.

First, we consider two possible mask learning schemes uncorrelated with probability (without PU), named case (a) and (b), so as to demonstrate the superiority of PU during training. Then, we consider testing without the Uniform distribution $\mathbf{U}$, called case (c), to investigate the contribution of adopting $\mathbf{U}$ during testing. Table III shows the PSNR comparisons for the above three cases and our PUERT when ratio is 5% and 10% under 1D and 2D settings on Brain dataset. Fig. 8 further illustrates their progression curves of test PSNR when ratio is 10% under the 1D setting. We detail the above three cases and their results in the following three paragraphs.

In case (a), the sampling subnet directly binarizes the unconstrained parameter matrix $\mathbf{O}$ into the sampling mask $\mathbf{M}$ via $\mathbf{M}_{i,j}=Bina(\mathbf{O}_{i,j})$ as a replacement of Eq. (3). Here, a L2 loss term to constrain the sparsity ratio of the learned sampling mask is necessary. However, such loss term still can not stabilize the ratio to the target setup, and therefore, in pursuit of fair comparisons, we have to test case (a) via greedy binarization Eq. (6) but with $\mathbf{\Omega}$ being the set $\{\mathbf{O}_{i,j}\}$. Table III shows that the final PSNR result of case (a) is inferior against PUERT under all conditions. This is mainly due to: 1) the limited optimization space without PU, 2) the mismatch between training and testing, and 3) the fluctuation caused by the poor ratio control, as shown in Fig. 8.

In order to eliminate the influence of the latter two limitations in case (a), we design case (b) to adopt greedy binarization during both training and testing, via Eq. (6) with $\mathbf{\Omega}$ being the set $\{\mathbf{O}_{i,j}\}$. Table III reports that case (b) achieves PSNR scores comparable with case (a), but still lower than PUERT. Fig. 8 further shows that, during the later period of training, the curve stability of case (b) is better than case (a). However, in the earlier training period, the PNSR upgrade of case (b) is obviously slow. We attribute such slow update speed to the greedy binarization used for training, which adopts a movable and relative boundary, rather than a fixed boundary adopted in vanilla binarization (i.e., the constant 1).

With the aim of investigating the contribution of adopting $\mathbf{U}$ during testing, we consider case (c) which tests PUERT without the Uniform distribution $\mathbf{U}$, i.e., testing via Eq. (6) but with $\mathbf{\Omega}$ being the set $\{\mathbf{P}_{i,j}\}$. Without the randomness and exploration introduced by $\mathbf{U}$, case (c) completely depends on the values in $\mathbf{P}$, and thus, with an earlier $\mathbf{P}$ which has not yet fully learned the probability value, case (c) achieves obviously lower PSNR than PUERT in the earlier period, as shown in Fig. 8. Table III also reports that case (c) achieves slightly lower PSNR than PUERT. Note that, in actual scanning, we might use a fixed mask as case (c), but this experiment still verifies the contribution of adopting randomness.

Overall, the above results corroborate the importance of learning a probabilistic sampling pattern rather than a deterministic mask, thus demonstrating the superiority of our design on Probabilistic Under-Sampling (PU) in the sampling subnet.

In this section, we present an ablation study on our proposed Dynamic Gradient Estimation (DGE) in our sampling subnet, so as to underscore the great significance of two-step gradient estimation. Note that, in LOUPE, the binarization operator is directly relaxed to the sigmoid function to enable optimizing the network by back-propagation. However, this relaxation not only causes a network performance penalty, but also leads to a non-binary output mask, which means that, after one has obtained the optimized non-binary mask, the reconstruction network needs to be retrained with the learned binary mask. To overcome the above two limitations, our PUERT chooses to still adopt the binarization in the forward pass, but develops an efficient gradient estimator for the backward pass.

Such gradient estimators have been widely investigated in Binarized Neural Networks (BNNs) [68, 77]. One typical and simple estimator is Straight-Through Estimator (STE), which was first proposed by Hinton [78] to train networks with binary activations (i.e., binary neuron). In STE, the values are passed through a binarization layer that evaluates the sign in the forward pass and performs the identity function during the backward pass. Adopting STE as an ablation study, we replace the dynamic gradient estimation function $g(x)$ in Eq. (7) simply by the identity function $h(x)=x$ and conduct an experiment. Besides, we also consider an STE variant proposed in [79], which uses the sigmoid function $\sigma(x)=\frac{1}{1+e^{-x}}$ as a replacement of the identity function.

Table III reports the PSNR comparisons for 1) STE, 2) Sigmoid, and 3) our DGE, when ratio is 5% and 10% under 1D and 2D settings on Brain dataset, which shows that DGE outperforms the other two methods across all situations. Fig. 9 further provides the progression curves of test PSNR results for the above three methods, when ratio is 10% under the 2D setting. One can intuitively observe that, STE converges faster than Sigmoid during the earlier training process, whereas lacks the ability to stably and consistently improve the reconstruction accuracy (resembling Sigmoid) during the later training period. Such different performances are the result of their different focuses. Concretely, STE directly passes the gradients through the binarization operator so as to guarantee the updating speed, while Sigmoid estimates the gradients with high similarity to the binarization operator to ensure estimation accuracy. As a combination of the above two focuses, our proposed two-step DGE is able to dynamically estimate the gradients and emphasize different priorities at the right time, thus implementing both fast convergence speed and high reconstruction accuracy, as can be appreciated from Fig. 9.

In the previous statements, we have emphasized to adopt Deep Unfolding Networks (DUNs) for our PUERT in order to fully utilize the knowledge on the learned sampling mask. In this section, two experiments are further conducted to investigate the contribution of adopting DUN in our PUERT.

Firstly, considering a DUN and a non-DUN with similar performances, we respectively equip them with our sampling subnet (SampNet), so as to compare the different performance gains brought by SampNet. Concretely, we select U-Net [72] as an instance of non-DUNs, and choose our RecNet with 4 stages, named RecNet_s4, to represent DUNs. Table V provides the PSNR results of U-Net with and without SampNet, as well as RecNet_s4 with and without SampNet, when ratio is 5% and 10% under 1D and 2D settings on Brain dataset. It can be observed that, without SampNet, U-Net and RecNet_s4 achieve similar PSNR results on average. However, when equipped with SampNet, RecNet_s4 realizes remarkably higher PSNR increase than U-Net (1.46 dB v.s. 1.03 dB on average). This result confirms the superiority of DUNs against non-DUNs on fully exploiting the knowledge on the learned mask.

Secondly, in order to validate the promising superiority of DUNs on further exploiting the information in SampNet and improving the performance, we propose an advanced version of PUERT, dubbed PUERT⁺, which integrates the probabilistic sampling pattern $\mathbf{P}$ in each stage of RecNet. Specifically, for each stage, we add a module called GDM^P, which is the same as GDM in Eq. (11) except for replacing $\mathbf{M}$ with $\mathbf{P}$, formulated as:

The input of each PMM is correspondingly modified to the concatenation of $\mathbf{r}_{k}^{\mathbf{P}}$ and $\mathbf{r}_{k}$, formulated as:

where $[\cdot]$ is the concatenation operator. In this way, each stage can fully use the information of not only the sampling mask $\mathbf{M}$ but also the probabilistic sampling pattern $\mathbf{P}$ from SampNet. Compared to the hard-sampling GDM with the binary mask $\mathbf{M}$, the newly added module GDM^P with the non-binary $\mathbf{P}$ can be regarded as a soft-sampling version. Table V provides the PSNR comparisons between PUERT and PUERT⁺, when ratio is 5% and 10% under 1D and 2D on Brain dataset. One can clearly see that PUERT⁺ obtains consistently higher scores than PUERT across all situations, with average PSNR increased from 35.91 dB to 36.05 dB. The 2D version achieves higher PSNR improvements than 1D, due to the utilization of more information. Note that, taking the 2D version as an example, PUERT⁺ only adds 2K learnable parameters, compared to PUERT with 404K parameters. Above results verify the promising superiority of DUNs on further exploiting the information in SampNet and improving the performance. More elaborate and efficient designs on further exploitation of SampNet are considered as our important future direction.