Deep J-Sense: Accelerated MRI Reconstruction via Unrolled Alternating Optimization

Abstract

Accelerated multi-coil magnetic resonance imaging reconstruction has seen a substantial recent improvement combining compressed sensing with deep learning. However, most of these methods rely on estimates of the coil sensitivity profiles, or on calibration data for estimating model parameters. Prior work has shown that these methods degrade in performance when the quality of these estimators are poor or when the scan parameters differ from the training conditions. Here we introduce Deep J-Sense as a deep learning approach that builds on unrolled alternating minimization and increases robustness: our algorithm refines both the magnetization (image) kernel and the coil sensitivity maps. Experimental results on a subset of the knee fastMRI dataset show that this increases reconstruction performance and provides a significant degree of robustness to varying acceleration factors and calibration region sizes.

1. Introduction

Parallel MRI is a multi-coil acceleration technique that is standard in nearly all clinical systems [5,15,21]. The technique uses multiple receive coils to measure the signal in parallel, and thus accelerate the overall acquisition. Compressed sensing-based methods with suitably chosen priors have constituted one of the main drivers of progress in parallel MRI reconstruction for the past two decades [4,10,16,25]. While parallel MRI provides additional degrees of freedom via simultaneous measurements, it brings its own set of challenges related to estimating the spatially varying sensitivity maps of the coils, either explicitly [15,25,26] or implicitly [5,21]. These algorithms typically use a fully sampled region of k-space or a low-resolution reference scan as an auto-calibration signal (ACS), either to estimate k-space kernels [5,11], or to estimate coil sensitivity profiles [15]. Calibration-free methods have been proposed that leverage structure in the parallel MRI model; namely, that sensitivity maps smoothly vary in space [16,26] and impose low-rank structure [6,20].

Deep learning has recently enabled significant improvement to image quality for accelerated MRI when combined with ideas from compressed sensing in the form of unrolled iterative optimization [1,7,18,22,23]. Our work falls in this category, where learnable models are interleaved with optimization steps and the entire system is trained end-to-end with a supervised loss. However, there are still major open questions concerning the robustness of these models, especially when faced with distributional shifts [2], i.e., when the scan parameters at test time do not match the ones at training time or the robustness of methods across different training conditions. This is especially prudent for models that use estimated sensitivity maps, and thus require reliable estimates as input.

Our contributions are the following: i) we introduce a novel deep learning-based parallel MRI reconstruction algorithm that unrolls an alternating optimization to jointly solve for the image and sensitivity map kernels directly in k-space; ii) we train and evaluate our model on a subset of the fastMRI knee dataset and show improvements in reconstruction fidelity; and iii) we evaluate the robustness of our proposed method on distributional shifts produced by different sampling parameters and obtain state-of-the-art performance. An open-source implementation of our method is publicly available¹.

2. System Model and Related Work

In parallel MRI, the signal is measured by an array of radio-frequency receive coils distributed around the body, each with a spatially-varying sensitivity profile. In the measurement model, the image is linearly mixed with each coil sensitivity profile and sampled in the Fourier domain (k-space). Scans can be accelerated by reducing the number of acquired k-space measurements, and solving the inverse problem by leveraging redundancy across the receive channels as well as redundancy in the image representation. We consider parallel MRI acquisition with C coils. Let $\mathbf{m}\in {ℂ}^n$ and $\text{s}=\left[{\text{s}}_1,\cdots ,{\text{s}}_C\right]\in {ℂ}^{k\times C}$ be the n-dimensional image (magnetization) and set of k-dimensional sensitivity map kernels, respectively, defined directly in k-space. We assume that k_i, the k-space data of the i-th coil image, is given by the linear convolution between the two kernels as

$${\mathbf{k}}_i={\mathbf{s}}_i\ast \mathbf{m}.$$ (1)

Joint image and map reconstruction formulates (1) as a bilinear optimization problem where both variables are unknown. Given a sampling mask represented by the matrix A and letting y = Ak + n be the sampled noisy multicoil k-space data, where k is the ground-truth k-space data and n is the additive noise, the optimization problem is

$$\underset{\mathbf{s},\mathbf{m}}{\arg \min }\frac{1}{2}\Vert \mathbf{y}-\mathbf{A}(\mathbf{s}\ast \mathbf{m}){\Vert }_2^2+{\lambda }_{\mathbf{m}}{R}_{\mathbf{m}}(\mathbf{m})+{\lambda }_{\mathbf{s}}{R}_{\mathbf{s}}(\mathbf{s}).$$ (2)

R_m(m) and R_s(s) are regularization terms that enforce priors on the two variables (e.g., J-Sense [26] uses polynomial regularization for s). We let A_m and A_s denote the linear operators composed by convolution with the fixed variable and A. A solution of (2) by alternating minimization involves the steps

$$\mathbf{m}=\underset{\mathbf{m}}{\arg \min }\frac{1}{2}{\Vert \mathbf{y}-{\mathbf{A}}_{\mathbf{m}}\mathbf{m}\Vert }_2^2+{\lambda }_{\mathbf{m}}{R}_{\mathbf{m}}(\mathbf{m}),$$ (3a)

$$\text{s}=\underset{\text{s}}{\arg \min }\frac{1}{2}{\Vert \mathbf{y}-{\mathbf{A}}_{\text{s}}\mathbf{s}\Vert }_2^2+{\lambda }_{\text{s}}{R}_{\text{s}}(\text{s}).$$ (3b)

If ℓ₂-regularization is used for the R terms, each sub-problem is a linear least squares minimization, and the Conjugate Gradient (CG) algorithm [19] with a fixed number of steps can be applied to obtain an approximate solution.

2.1. Deep Learning for MRI Reconstruction

Model-based deep learning architectures for accelerated MRI reconstruction have recently demonstrated state-of-the-art performance [1,7,22]. The MoDL algorithm [1] in particular is used to solve (1) when only the image kernel variable is unknown and a deep neural network $\mathcal{D}$ is used in R_m(m) as

$$\underset{\mathbf{m}}{\arg \min }\frac{1}{2}{\Vert \mathbf{y}-{\mathbf{A}}_{\mathbf{m}}\mathbf{m}\Vert }_2^2+\lambda \Vert \mathcal{D}(\mathbf{m})-\mathbf{m}{\Vert }_2^2.$$ (4)

To unroll the optimization in (4), the authors split each step in two different sub-problems. The first sub-problem treats $\mathcal{D}(\mathbf{m})$ as a constant and uses the CG algorithm to update m. The second sub-problem treats $\mathcal{D}$ as a proximal operator and is solved by direct assignment, i.e., ${\mathbf{m}}^{+}=\mathcal{D}(\mathbf{m})$. In our work, we use the same approach for unrolling the optimization, but we use a pair of deep neural networks, one for each variable in (2). Unlike [1], our work does not rely on a pre-computed estimate of the sensitivity maps, but instead treats them as an optimization variable.

The idea of learning a sensitivity map estimator using neural networks was first described in [3]. Recently, the work in [22] introduced the E2E-VarNet architecture that addresses the issue of estimating the sensitivity maps by training a sensitivity map estimation module in the form of a deep neural network. Like the ESPiRiT algorithm [25], E2E-VarNet additionally enforces that the sensitivity maps are normalized per-pixel and uses the same sensitivity maps across all unrolls. This architecture – which uses gradient descent instead of CG – is then trained end-to-end, using the estimated sensitivity maps and the forward operator A_m. The major difference between our work and [22] is that we iteratively update the maps instead of using a single-shot data-based approach [22] or the ESPiRiT algorithm [1,7]. As our results show in the sequel, this has a significant impact on the out-of-distribution robustness of the approach on scans whose parameters differ from the training set.

Concurrent work also extends E2E-VarNet and proposes to jointly optimize the image and sensitivity maps via the Joint ICNet architecture [9]. Our main difference here is the usage of a forward model directly in k-space, and the fact that we do not impose any normalization constraints on the sensitivity maps. Finally, the work in [12] proposes a supervised approach for sensitivity map estimation, but this still requires an external algorithm and a target for map regression, which may not be easily obtained.

3. Deep J-Sense: Unrolled Alternating Optimization

We unroll (2) by alternating between optimizing the two variables as

$$\text{s}=\underset{\text{s}}{\arg \min }\frac{1}{2}{\Vert \mathbf{y}-{\mathbf{A}}_{\text{s}}\mathbf{s}\Vert }_2^2+{\lambda }_{\text{s}}{R}_{\text{s}}(\mathbf{s})$$ (5a)

$${\text{s}}^{+}={\mathcal{D}}_{\text{s}}(\text{s}),$$ (5b)

$$\mathbf{m}=\underset{\mathbf{m}}{\arg \min }\frac{1}{2}{\Vert \mathbf{y}-{\mathbf{A}}_{\mathbf{m}}\mathbf{m}\Vert }_2^2+{\lambda }_{\mathbf{m}}{R}_{\mathbf{m}}(\mathbf{m}),$$ (5c)

$${\mathbf{m}}^{+}={\mathcal{D}}_{\mathbf{m}}(\mathbf{m}),$$ (5d)

where R is defined as

$$R(\text{x})={\Vert \mathcal{F}\left\{{\mathcal{D}}_{\text{x}}\left({\mathcal{F}}^{-1}\left\{\text{x}\right\}\right)\right\}-\text{x}\Vert }_2^2,$$ (6)

for both m and s, $\mathcal{F}$ is the Fourier transform, and $\mathcal{D}$ is the deep neural network corresponding to each variable. Similar to MoDL, we set ${\mathcal{D}}_{\mathbf{m}}^{(j)}={\mathcal{D}}_{\mathbf{m}}$ and ${\mathcal{D}}_{\text{s}}^{(j)}={\mathcal{D}}_{\text{s}}$ across all unrolls, leading to the efficient use of learnable weights. The coefficients λ_s and λ_m are also learnable. The optimization is initialized with m⁽⁰⁾and s⁽⁰⁾, obtained using a simple root sum-of-squares (RSS) estimate. Steps (5a) and (5c) are approximately solved with n₁ and n₂ steps of the CG algorithm, respectively, while steps (5b) and (5d) represent direct assignments. The two neural networks serve as generalized denoisers applied in the image domain and are trained in an end-to-end fashion after unrolling the alternating optimization for a number of N outer steps. A block diagram of one unroll is shown in Fig.1.

Fig. 1.

A single unroll of the proposed scheme. The CG algorithm is executed on the loss given by the undersampled data y and the measurement matrices A_s and A_m. Each block is matched to its corresponding equation.

Using the estimated image and map kernels after N outer steps, we train the end-to-end network using the estimated RSS image as

$$\widehat{\mathbf{x}}=\sqrt{\sum _{i=1}^C{\left|{\mathcal{F}}^{-1}\left\{{\mathbf{s}}_i^{(N)}\ast {\mathbf{m}}^{(N)}\right\}\right|}^2},$$ (7)

where the supervised loss is the structural similarity index (SSIM) between the estimated image and ground truth image as $L=-\mathrm{SSIM}(\widehat{x},x)$.

Our model can be seen as a unification of MoDL and J-Sense. For example, by setting n₁ = 0 and ${\mathcal{D}}_{\text{s}}(\text{s})=\text{s}$, the sensitivity maps are never updated and the proposed approach becomes MoDL. At the same time, removing the deep neural networks by setting $\mathcal{D}(\cdot )=0$ and removing steps (5b) and (5d) leads to L2-regularized J-Sense. One important difference is that, unlike the ESPiRiT algorithm or E2E-VarNet, our model does not perform pixel-wise normalization of the sensitivity maps, and thus does not impart a spatially-varying weighting across the image. We use a forward model in the k-space domain based on a linear convolution instead of prior work that uses the image domain. This allows us to use a small-sized kernel for the sensitivity maps as an implicit smoothness regularizer and reduction in memory.

4. Experimental Results

We compare the performance of the proposed approach against MoDL [1] and E2E-VarNet [22]. We train and evaluate all methods on a subset of the fastMRI knee dataset [27] to achieve reasonable computation times. For training, we use the five central slices from each scan in the training set, for a total of 4580 training slices. For evaluation, we use the five central slices from each scan in the validation set, for a total of 950 validation slices. All algorithms are implemented in PyTorch [14] and SigPy [13]. Detailed architectural and hyper-parameter choices are given in the material.

To evaluate the impact of optimizing the sensitivity map kernel, we compare the performance of the proposed approach with MoDL trained on the same data, that uses the same number of unrolls (both inner and outer) and the same architecture for the image denoising network ${\mathcal{D}}_{\text{m}}$ We compare our robust performance with E2E-VarNet trained on the same data, and having four times more parameters, to compensate for the run-time cost of updating the sensitivity maps.

4.1. Performance on Matching Test-Time Conditions

We compare the performance of our method with that of MoDL and E2E-VarNet when the test-time conditions match those at training time on knee data accelerated by a factor of R = 4. For MoDL, we use a denoising network with the same number of parameters as our image denoising network and the same number of outer and inner CG steps. We use sensitivity maps estimated by the ESPiRiT algorithm via the BART toolbox [24], where a SURE-calibrated version [8] is used to select the first threshold, and we set the eigenvalue threshold to zero so as to not unfairly penalize MoDL, since both evaluation metrics are influenced by background noise. For E2E-VarNet, we use the same number of N = 6 unrolls (called cascades in [22]) and U-Nets for all refinement modules.

Table 1 shows statistical performance on the validation data. The comparison with MoDL allows us to evaluate the benefit of iteratively updating the sensitivity maps, which leads to a significant gain in both metrics. Furthermore, our method obtains a superior performance to E2E-VarNet while using four times fewer trainable weights and the same number of outer unrolls. This demonstrates the benefit of trading off the number of parameters for computational complexity, since our model executes more CG iterations than both baselines. Importantly, Deep J-Sense shows a much lower variance of the reconstruction performance across the validation dataset, with nearly one order of magnitude gain against MoDL and three times lower than E2E-VarNet, delivering a more consistent reconstruction performance across a heterogeneity of patient scans.

Table 1.

Validation performance on a subset of the fastMRI knee dataset. Higher average/median SSIM (lower NMSE) indicates better performance. Lower standard deviations are an additional desired quality.

	Avg. SSIM	Med. SSIM	σ SSIM	Avg. NMSE	Med. NMSE	σ NMSE
MoDL	0.814	0.840	0.115	0.0164	0.0087	0.0724
E2E	0.824	0.851	0.107	0.0111	0.0068	0.0299
Ours	0.832	0.857	0.104	0.0091	0.0064	0.0095

Randomly chosen reconstructions of scans and the estimated sensitivity maps are shown in Fig. 2. We notice that our maps capture higher frequency components than those of MoDL (estimated via ESPiRiT), but do not contain spurious noise patterns outside the region spanned by the physical knee. In contrast, the maps from E2E-VarNet exhibit such patterns and, in the case of the second row, produce spurious patterns even inside the region of interest, suggesting that the knee anatomy “leaks” into the sensitivity maps. Figure 3 shows the projection of the fully-sampled k-space data on the null space of the pixel-wise normalized estimated sensitivity maps. Note that since we do not normalize the sensitivity maps this represents a different set of maps that those learned during training. As Fig. 3 shows, the residual obtained is similar to that of E2E-VarNet.

Fig. 2.

Example reconstructions at R = 4 and matching train-test conditions. The first and third rows represent RSS images from two different scans. The second and fourth rows represents the magnitude of the estimated sensitivity map (no ground truth available, instead the final column shows the coil image) for a specific coil from each scan, respectively. The maps under the MoDL column are estimated with the ESPiRiT algorithm.

Fig. 3.

Magnitude of the projection of the ground truth k-space data onto the null space of the normal operator given by the estimated sensitivity maps for various algorithms (residuals amplified 5×). The sample is chosen at random; both Deep J-Sense and E2E-VarNet produce map estimates that deviate from the linear model, but improve end-to-end reconstruction of the RSS image.

4.2. Robustness to Test-Time Varying Acceleration Factors

Figure 4 shows the performance obtained at acceleration factors between 2 and 6, with models trained only at R = 4. The modest performance gain for E2E-VarNet at R < 4 confirms the findings in [2]: certain models cannot efficiently use additional measurements if there is a train-test mismatch. At the same time, MoDL and the proposed method are able to overcome this effect, with our method significantly outperforming the baselines across all accelerations. Importantly, there is a significant decrease of the performance loss slope against MoDL, rather than an additive gain, showing the benefit of estimating sensitivity maps using all the acquired measurements.

Fig. 4.

Average SSIM on the fastMRI knee validation dataset evaluated at acceleration factors R between 2 and 6 (with granularity 0.1) using models trained at R = 4. The vertical lines are proportional to the SSIM standard deviation in each case, from which no noticeable difference can be seen.

4.3. Robustness to Train-Time Varying ACS Size

We investigate the performance of the proposed method and E2E-VarNet as a function of the ACS region size (expressed as number of acquired lines in the phase encode direction). All models are trained at R = 4 and ACS sizes {1, 6, 12, 26, 56} and tested in the same conditions, giving a total of ten models. While there is no train-test mismatch in this experiment, Fig. 5 shows a significant performance gain of the proposed approach when the calibration region is small (below six lines) and shows that overall, end-to-end performance is robust to the ACS size, with the drops at 1 and 12 lines not being statistically significant. This is in contrast to E2E-VarNet, which explicitly uses the ACS in its map estimation module and suffers a loss when this region is small.

Fig. 5.

Average SSIM on the fastMRI knee validation dataset evaluated at different sizes of the fully sampled auto-calibration region, at acceleration factor R = 4. The vertical lines are proportional to the SSIM standard deviation. Each model is trained and tested on the ACS size indicated by the x-axis.

5. Discussion and Conclusions

In this paper, we have introduced an end-to-end unrolled alternating optimization approach for accelerated parallel MRI reconstruction. Deep J-Sense jointly solves for the image and sensitivity map kernels directly in the k-space domain and generalizes several prior CS and deep learning methods. Results show that Deep J-Sense has superior reconstruction performance on a subset of the fastMRI knee dataset and is robust to distributional shifts induced by varying acceleration factors and ACS sizes. A possible extension of our work could include unrolling the estimates of multiple sets of sensitivity maps to account for scenarios with motion or a reduced field of view [17,25].