Compressed Sensing MRI with ℓ₁-Wavelet Reconstruction Revisited Using Modern Data Science Tools

Abstract

Deep learning (DL) has emerged as a powerful tool for improving the reconstruction quality of accelerated MRI. These methods usually show enhanced performance compared to conventional methods, such as compressed sensing (CS) and parallel imaging. However, in most scenarios, CS is implemented with two or three empirically-tuned hyperparameters, while a plethora of advanced data science tools are used in DL. In this work, we revisit ℓ₁-wavelet CS for accelerated MRI using modern data science tools. By using tools like algorithm unrolling and end-to-end training with stochastic gradient descent over large databases that DL algorithms utilize, and combining these with conventional concepts like wavelet subband processing and reweighted ℓ₁ minimization, we show that ℓ₁-wavelet CS can be fine-tuned to a level comparable to DL methods. While DL uses hundreds of thousands of parameters, the proposed optimized ℓ₁-wavelet CS with sub-band training and reweighting uses only 128 parameters, and employs a fully-explainable convex reconstruction model.

I. INTRODUCTION

Slow data acquisition remains a challenge for MRI, requiring accelerated imaging strategies. Conventional methods, such as parallel imaging [1], [2] and compressed sensing (CS) [3] are used clinically, but typically their acceleration rates are limited by noise amplification and residual aliasing artifacts in reconstructed images. Recently, deep learning (DL) methods for accelerated MRI [4]–[8] have emerged as a powerful tool for MRI reconstruction, with improved performance over conventional methods in many studies. Among DL methods, physics-guided DL (PG-DL) methods that unroll conventional optimization algorithms that incorporate the encoding operator have received attention [6]–[9]. While CS uses a linear transform-based representation of images for regularization, PG-DL methods utilize a non-linear representation for regularization, which is implicitly learned through neural networks.

DL reconstruction methods are trained using large databases, include a large number (usually more than hundreds of thousands or millions [6], [10], [11]) of parameters, incorporate sophisticated optimization algorithms for training [12], and utilize state-of-the-art loss functions [13], [14]. On the other hand, when CS reconstruction methods are implemented for comparison, they typically use two or three parameters, which are frequently hand-tuned using a simple grid search. Although some automatic tuning methods have been proposed [15], [16], these have not leveraged the widely-available and popular modern data science tools from the DL era.

In this work, we use these data science tools to revisit ℓ₁-wavelet CS for accelerated MRI. Similar to PG-DL methods, we unroll an ADMM algorithm and train it end-to-end, while using only 4 orthogonal wavelet bases for the regularizing transforms for a total of 12 tunable parameters. Building on this naive model, we further incorporate processing of each individual wavelet subband [17], and reweighted ℓ₁ minimization [18], leading to 64 and 128 tunable parameters respectively. Results show that even the naive model closes the gap in reconstruction performance to advanced PG-DL methods, while the incorporation of subband and reweighting further improves the quality of reconstruction to a level comparable to PG-DL methods. All the models proposed here enable a linear representation for interpretable and convex sparse image reconstruction at inference time.

II. MATERIALS AND METHODS

A. Inverse Problem for Accelerated MRI

The forward model for accelerated MRI is given as

$$y=Ex+n$$ (1)

where $x\in {ℂ}^N$ is the image to be reconstructed, $y\in {ℂ}^M$ is the undersampled k-space data from all coils, $E:{ℂ}^N\to {ℂ}^M$ is the linear forward encoding operator containing coil sensitivity maps and partial Fourier matrix for undersampling in k-space [19], and n is the measurement noise. The inverse problem involves solving the objective function:

$$\widehat{x}=\text{arg\hspace{0.17em}}\underset{x}{\text{min}}\frac{1}{2}{\Vert y-Ex\Vert }_2^2+\mathcal{R}(x)$$ (2)

where ${\Vert y-Ex\Vert }_2^2$ enforces data consistency (DC) and $\mathcal{R}(x)$ is a regularizer.

In conventional CS MRI reconstruction, the form of $\mathcal{R}(x)$ is often a weighted ℓ₁-norm of transform coefficients, i.e. $\mathcal{R}(x)={\sum }_{l=1}^L{\lambda }_l{\Vert W_{l}x\Vert }_1$, where W_l is a pre-specified linear (often orthogonal) transform, such as a discrete wavelet transform (DWT) [3], and L is the number of linear transforms used for regularization. The resulting convex objective function is solved via an iterative optimization algorithm [20]. These algorithms are conventionally run until a stopping criterion is met, making hyperparameter tuning difficult.

On the other hand, in PG-DL reconstruction, the inverse problem is usually solved by unrolling an iterative optimization algorithm for a fixed number of iterations [21], [22]. Typically, the solutions are decoupled to a series of regularizer and DC units. The regularizer in PG-DL is implemented implicitly via convolutional neural networks (CNNs), while the DC unit is solved by linear methods, such as gradient descent or conjugate gradient [7]. The network is trained end-to-end as:

$$\underset{\theta }{\text{min}}\frac{1}{N}\sum _{n=1}^N\mathcal{L}\left(y_{\text{ref}}^n,E_{\text{full }}^n\left(f\left(y^n,E^n;\theta \right)\right)\right),$$ (3)

where $y_{\text{ref}}^n$ denotes the fully-sampled reference k-space of the n^th subject, f(yⁿ, Eⁿ; θ) denotes network output of the unrolled network with parameters θ of the n^th subject, $E_{\text{full }}^n$ is the fully sampled multi-coil encoding operator of the n^th subject, N is the number of datasets in the training database, and $\mathcal{L}(\cdot ,\cdot )$ is a loss function between the network output and the reference. Common choices for $\mathcal{L}(\cdot ,\cdot )$ include ℓ₂ norm, ℓ₁ norm, mixed norms and perception-based loss [8], [14].

B. Proposed Learning of ℓ₁-Wavelet CS Reconstruction

We optimize conventional ℓ₁-Wavelet CS reconstruction by using the data science tools utilized in PG-DL techniques. First, we utilize ADMM to set

$$x^{(t+1)}={\left(E^{H}E+\sum _{l=1}^L{\rho }_{l}I\right)}^{-1}\left(E^{H}y+\sum _{l=1}^L{\rho }_l{W}_l^H\left(z_l^{(t)}-{\mathit{\beta }}_l^{(t)}\right)\right)$$ (4a)

$$z_l^{(t+1)}=\text{soft}\left(W_l{x}^{(t+1)}+{\beta }_l^{(t)};{\lambda }_l/{\rho }_l\right)$$ (4b)

$${\beta }_l^{(t+1)}={\beta }_l^{(t)}+{\eta }_l^{(t+1)}\left(W_l{x}^{(t+1)}-z_l^{(t+1)}\right)$$ (4c)

where z_l are auxiliary variables in wavelet domain, β_l are dual variables, soft(·; λ_l/ρ_l) is the ℓ₁ soft-thresholding operator parameterized by λ_l/ρ_l, and t denotes the iteration count. The algorithm is unrolled for T iterations, as depicted in Figure 1.

Fig. 1:

Schematic of the unrolled ADMM for ℓ₁-wavelet compressed sensing (CS). One unrolled iteration of ADMM with ℓ₁-wavelet regularizers consists of regularizer (R), data consistency (DC) and dual update (DWT: Discrete wavelet transform). In accordance with the ADMM framework, learnable parameters are shared across different unrolled iterations. In its simplest form, this leads to 3 trainable parameters per wavelet. Further enhancements, such as separate thresholds for wavelet subbands and reweighted ℓ₁ minimization increase this number to 16 trainable parameters per wavelet.

The learnable parameters in this algorithm are ρ_l, λ_l/ρ_l and η_l, which correspond to parameters for augmented Lagrangian relaxation, ℓ₁ soft-thresholding and the dual update per each wavelet transform. We note that these are shared across all the unrolled iterations to ensure that the objective function in (2) remains unchanged throughout the iterations, and interpretability of the algorithm can be maintained. Thus, there are 3 · L learnable parameters for the whole algorithm when using L orthogonal DWTs.

This approach serves as the foundation for all our proposed models, and is subsequently referred to as the learned naive ℓ₁-Wavelet reconstruction. In all our models, the input to the network is the zerofilled image, x⁽⁰⁾ = E^Hy. Furthermore, since the regularizer in (2) scales with ∥x∥_∞, while the DC term in (2) scales with ${\Vert x\Vert }_{\infty }^2$, λ_l/ρ_l is parametrized as ${\gamma }_l{\Vert W_l{x}^{(0)}\Vert }_{\infty }$, and the scaling-invariant parameter γ_l is learned. Overall, the learned parameters are ${\left\{{\rho }_l,{\gamma }_l,{\eta }_l\right\}}_{l=1}^L$.

C. Further Enhancements for Optimized ℓ₁-Wavelet CS Reconstruction

The naive optimized ℓ₁-Wavelet approach can further be enhanced using our understanding of wavelet representations [17] and of ℓ₁ minimization problems [18]. In particular, we use the fact that signal scaling changes severely between different wavelet subbands for the former, and that reweighted ℓ₁ minimization helps recover finer details for the latter.

Learning ℓ₁-wavelet reconstruction with subband processing:

For different subbands of a wavelet transform, the soft-thresholding parameters may be different. To this end, let $D_l^s$ be an operator that select the s^th subband of the l^th wavelet transform. We propose to use the regularizer

$$\mathcal{R}(x)=\sum _{l=1}^L\sum _{s=1}^S{\lambda }_{l,s}{\Vert D_l^s{W}_{l}x\Vert }_1,$$ (5)

which also lead to the following modified update in (4b).

$$D_l^s{z}_l^{(t+1)}=\text{soft}\left(D_l^s\left(W_l{x}^{(t+1)}+{\beta }_l^{(t)}\right);\frac{{\lambda }_{l,s}}{{\rho }_l}\right)$$ (6)

for all s ∈ {1, … , S}. Thus, the learnable soft-thresholding parameters are λ_l,s/ρ_l for the s^th subband of the l^th wavelet transform. During end-to-end training, this parameter is again implemented in a scaling-invariant manner by defining ${\gamma }_{l,s}=\left({\lambda }_{l,s}/{\rho }_l\right)/{\Vert D_l^s{W}_l{x}^{(0)}\Vert }_{\infty }$ and learning {γ_l,s} for l ∈ {1, … , L} and s ∈ {1, … , S}. Thus this approach leads to a total of L · (S + 2) learnable parameters when using S subbands and L orthogonal DWTs.

Learning reweighted ℓ₁-wavelet reconstruction with subband processing:

A further improvement in performance, especially in the lower SNR regimes, may be achieved using reweighted ℓ₁ minimization, which has been shown to improve recovery of small coefficients [18]. To this end, let ${\widehat{x}}_{\text{sb}}$ denote the output of the learned subband ℓ₁-wavelet reconstruction. We define a diagonal weight matrix U_l whose (k, k)^th entry is given as

$${\left(U_l\right)}_{(k,k)}=\frac{1}{\left|{\left(W_l{\widehat{x}}_{\text{sb}}\right)}_k\right|+ϵ},$$ (7)

where (·)_k denotes the k^th coefficient of the vector (·), and ϵ is a small constant to avoid numerical issues when dividing by zero. This weight matrix is used to define the reweighted ℓ₁ regularizer with subband processing as:

$$\mathcal{R}(x)=\sum _{l=1}^L\sum _{s=1}^S{\lambda }_{l,s}{\Vert D_l^s{U}_l{W}_{l}x\Vert }_1.$$ (8)

This leads to the following modified update in (4b).

$$D_l^s{z}_l^{(t+1)}=\text{soft}\left(D_l^s\left(W_l{x}^{(t+1)}+{\beta }_l^{(t)}\right);\frac{{\lambda }_{l,s}}{{\rho }_l}{D}_l^s\text{\hspace{0.17em}diag}\left(U_l\right)\right).$$ (9)

for all s ∈ {1, … , S}. We note that the regularizer in (8) does not change if x is scaled by a constant α, while the DC term in (2) still scales with ${\Vert x\Vert }_{\infty }^2$. Thus, we define a scaling-invariant thresholding factor ${\gamma }_{l,s}^r=\left({\lambda }_{l,s}/{\rho }_l\right)/{\Vert D_l^s{W}_l{x}^{(0)}\Vert }_{\infty }^2$. During end-to-end training, $\left\{{\gamma }_{l,s}^r\right\}$ is learned for l ∈ {1, …, L} and s ∈ {1, …, S} in addition to ${\left\{{\rho }_l,{\eta }_l\right\}}_{l=1}^L$.

This approach still has L · (S + 2) learnable parameters during the reweighting stage, even though signal-dependent weights are incorporated via (7). Including the learned subband reconstruction, which is used to determine the weights in (7) leads to a total of 2L · (S + 2) learnable parameters for the whole reconstruction pipeline. We also note that once these parameters are learned, they can be applied for multiple reweightings, k_rew, during testing, since the scaling of $\left\{{\gamma }_{l,s}^r\right\}$ remains on the same order.

D. Imaging Data

Fully-sampled coronal proton density (PD), and PD with fat-suppression (PD-FS) knee data obtained from the NYU-fastMRI database [23] were used throughout the experiments. Relevant imaging parameters were: matrix size = 320×368, in-plane resolution = 0.49 × 0.44 mm², slice thickness = 3 mm. The datasets were retrospectively under-sampled with a random mask (R = 4 with 24 ACS lines). Training was performed on 300 slices from 10 different subjects. Testing was performed on all slices from 10 different subjects. Coil sensitivity maps were generated using ESPIRiT [24].

E. Implementation Details

For all models, L = 4 wavelets were used, corresponding to Daubechies1–4 orthogonal wavelets with 14 subbands for each. Thus, the total number of learnable parameters were 12, 64 and 128 for the learned naive ℓ₁-wavelet, ℓ₁-wavelet with subbands, and reweighted ℓ₁-wavelet with subbands reconstructions, respectively. For the last method, k_rew = 2 was used in testing, similar to [18].

ADMM algorithm was unrolled for T = 10 for all models. ϵ was set to 10⁻⁹ in (7). DC subproblem was solved using conjugate gradient [7] with 5 iterations and warm-start. All tunable parameters were randomly initialized. Adam optimizer with learning rate 5 × 10⁻³ was used for training over 100 epochs, with a batch size of 1. Supervised training was performed with a normalized ℓ₁−ℓ₂ loss in k-space [8], [10], using TensorFlow in Python.

For comparison, a PG-DL approach was implemented using the same ADMM unrolling except for using a ResNet-based regularizer unit. The ResNet was originally adapted from the winner of a super-resolution challenge [25], and has been used in multiple recent MRI studies successfully [9], [10]. The PG-DL approach has a total of 592,130 learnable parameters. Note this constitutes a head-to-head comparison, with the only difference being in the $\mathcal{R}(x)$ term, where our approaches employ ℓ₁-norm of wavelets for solving a convex problem, while PG-DL uses a CNN for implicit regularization. All results were quantitatively compared using SSIM and NMSE.

III. RESULTS

Figure 2 shows a representative slice from coronal PD knee MRI, reconstructed using PG-DL and the three ℓ_l-wavelet CS approaches optimized with modern data science tools, respectively. For this high SNR acquisition, even the learned naive ℓ₁-wavelet results in a high-quality reconstruction, while the learned ℓ₁-wavelet with subbands leads to a slightly sharper reconstruction. Learned reweighted ℓ₁-wavelet with subbands performs the best among the three optimized ℓ_l-wavelet CS approaches, resulting in a sharp reconstruction, showing comparable visual quality and quantitative metrics to PG-DL.

Fig. 2:

A representative slice from coronal PD knee MRI, reconstructed using PG-DL, learned naive ℓ₁-wavelet, learned ℓ₁-wavelet with subbands, and learned reweighted ℓ₁-wavelet with subbands. The proposed optimized ℓ₁-wavelet reconstructions perform closely to PG-DL. Learned reweighted ℓ₁-wavelet with subbands performs the best among these ℓ₁-wavelet CS variants, resulting in sharp images with similar quantitative metrics to PG-DL.

Figure 3 depicts a representative coronal PD-FS knee MRI slice, reconstructed using PG-DL and the three optimized ℓ_l-wavelet CS approaches. The PD-FS dataset has an inherently lower SNR compared to PD. As such, the learned naive ℓ₁-wavelet results in a blurry image due to the difficulty of fine-tuning a single thresholding parameter. This is improved with subband processing, but sharpness is fully recovered only with the learned reweighted ℓ₁-wavelet with subbands approach, which results in a visibly similar reconstruction to the PG-DL method.

Fig. 3:

A representative slice from coronal PD–FS knee MRI, reconstructed using PG-DL, learned naive ℓ₁-wavelet, learned ℓ₁-wavelet with subbands, and learned reweighted ℓ₁-wavelet with subbands. For this lower SNR acquisition, the learned naive ℓ₁-wavelet suffers from visible blurring artifacts. These are improved using subband processing, while the sharpness is only fully recovered using the learned reweighted ℓ₁-wavelet with subbands method, which leads to a visibly similar reconstruction to PG-DL.

Table I summarizes quantitative results from knee MRI. While PG-DL has the best metrics, the gap between proposed optimized ℓ_l-wavelet CS and PG-DL is small, with < 0.01 for SSIM and < 0.0008 for NMSE.

TABLE I:

The median and the interquartile range [25^th, 75^th percentile] of the NMSE and SSIM metrics on test slices from 10 subjects for coronal PD and PD-FS datasets. While PG-DL has the best metrics as expected, the gap in SSIM between PG-DL and the learned reweighted ℓ₁-wavelet with subbands is < 0.01.

	Coronal PD		Coronal PD-FS
	NMSE	SSIM	NMSE	SSIM
PG-DL	.0013 [.0010, .0017]	0.9673 [0.9547, 0.9775]	.0074 [.0049, 0.0107]	0.8795 [0.8322, 0.9146]
Learned Naive ℓ1-wavelet	.0019 [.0014, .0024]	0.9574 [0.9404, 0.9692]	.0090 [.0062, 0.0125]	0.8598 [0.8121, 0.9023]
Learned ℓ1-wavelet with Subbands	.0019 [.0014, .0024]	0.9588 [0.9423, 0.9699]	.0081 [.0055, .0121]	0.8662 [0.8202, 0.9057]
Learned reweighted ℓ1-wavelet with Subbands	.0017 [.0013, .0022]	0.9602 [0.9444, 0.9709]	.0082 [.0053, .0116]	0.8694 [0.8211, 0.9068]

IV. DISCUSSION AND CONCLUSION

In this study, we revisited ℓ_l-wavelet CS for accelerated MRI using modern data science tools for fine tuning. As expected, PG-DL outperformed our three optimized ℓ_l-wavelet CS approaches, but the performance gap was smaller than previously published literature. This is interesting for a number of reasons. First, PG-DL used a sophisticated non-linear representation for the underlying images during regularization with a large number of learnable parameters. On the other hand, the wavelet-based representations we used were linear, involved only a small number of parameters, and allowed for convex optimization. Interestingly, there was <0.01 difference in SSIM between our proposed learned reweighted ℓ₁-wavelet with subbands that used 128 parameters and the PG-DL approach that used >500,000 parameters. Second, while PG-DL can be further improved with more advanced neural networks and training strategies [26], our CS approach used one of the simplest linear models described by fixed orthogonal wavelets, and did not involve learning of the representation. Our results also showed that the performance gap decreased as we proceeded from learned naive ℓ₁-wavelet to learned reweighted ℓ₁-wavelet with subbands, demonstrating subband training and reweighting help improve reconstruction sharpness and quantitative metrics to a level comparable to PG-DL. Further gains for CS may be possible via learning linear representations/frames [27], [28], which warrants investigation.