Learning to Sample MRI via Variational Information Maximization
Abstract
Accelerating MRI scans requires optimal sampling of k-space data.
This is however a daunting task due to the discrete and non-convex nature of sampling optimization.
To cope with this challenge, we put forth a novel deep learning framework that leverages uncertainty autoencoders to enable joint optimization of sampling pattern and reconstruction of MRI scans.
We represent the encoder as a non-uniform
Fast Fourier Transform that allows continuous optimization of k-space samples on a non-Cartesian plane, while the decoder is a deep reconstruction network.
Our approach is universal in a sense that it can be used with any reconstruction network.
Experiments with knee MRI shows improved reconstruction quality of our data-driven sampling over the prevailing variable-density sampling.
Introduction
Deep learning methods has shown great promise at solving inverse problems.
Of particular interest is
MR reconstruction problem, where the goal is to reconstruct high ﬁdelity images from a subset of spatial Fourier domain (k-space) representations.
Even though deep learning methods have shown improved image quality when solving MR reconstruction problems [1–9], undersampling patterns that specify the selected subset of Fourier domain representation are typically chosen heuristically.
The reconstruction models are optimized for a pre-determined acquisition (encoding) model without taking advantage of possible gains that can be obtained from learning undersampling patterns.
Recently, end-to-end deep learning methods have been proposed for learning undersampling patterns for MRI reconstruction problem.
Active acquisition strategies [10–12] attempt to predict the next k-space samples to be acquired using information from existing samples.
The non-active strategies can be grouped into two categories.
[13], [14] focus on the Cartesian sampling case and model binary sampling masks probabilistically.
These methods employ certain relaxations to maintain differentiability with respect to sampling probabilities.
The second type of methods [15, 16] directly optimize for the k-space coordinates instead of estimating sampling probabilities.
Authors of [15] restrict the optimization to a set of separable variables such as horizontal and vertical directions in 2-D plane to reduce the search space.
As a result, their optimized sampling masks are restricted to a uniform grid pattern and do not cover the more general sampling patterns.
In this work, we present a variational information maximization method that enables joint optimization of acquisition and reconstruction of MRI scans in a data-driven manner.
Our method enables learning an undersampling pattern tuned speﬁcally to the reconstruction network, and vice versa, to obtain improved reconstruction performance.
We represent the acquisition (encoder) model with the nonuniform Fast Fourier Transform (nuFFT) operator [17] that is parameterized by the sampling locations in k-space.
This allows interpolation of non-cartesian coordinates and enables continuous optimization of the sampling pattern.
On the reconstruction (decoder) side, we use an unrolled reconstruction network which mimics the proximal gradient based solutions to compressed sensing problems.
Methods
We consider the MR signal model under the additive white complex Gaussian noise as
[Equation]
where x ∈ C N is the image, z ∈ C M is the measured data in k-space domain, f φ ( · ) is the forward model that describes the imaging system parameterized by k-space sample coordinates φ ∈ [ − 0 .
5] M , and (cid:15) ∼ N c
0 , σ 2 I ) is the measurement noise.
The imaging model f φ includes non-uniform Fast Fourier Transform (nuFFT) operation F nu , and for the multi-coil scenario, contains signal modulation by coil sensitivity maps S .
Given an acceleration factor R = N/M , the goal is to ﬁnd the optimal set of samples φ along with a reconstruction function g ( z ) that maintains the full k-space data image quality.
2.1 Variational Information Maximization for Acquisition and Reconstruction
Under the setting in 1, the distribution of measurements Z given image X has a complex valued
Gaussian distribution q φ ( Z | X )
= N c
( f φ ( X ) , σ 2 I ) .
We adapt the uncertainty autoencoder framework deﬁned in [18] and we make use of the InfoMax principle [19, 20] that maximizes the mutual information between the measurements and noisy latent representations
[Equation]
Since H ( X ) does not depend on k-space coordinates φ , this is equivalent to
[Equation]
where in the last line we introduce a variational approximation to the model posterior q φ ( X | Z ) to obtain a lower bound.
The variational parameters θ correspond to the weights of the reconstruction network.
Using the bound on 4 we express our objective as
[Equation]
where we estimate the expectation with respect to q φ ( X ) using Monte-Carlo sampling from the dataset D .
Thus, we propose optimizing shared φ and θ for the entire dataset rather than speciﬁc acquisition and reconstruction parameters for each example in the dataset.
In addition, for a different anatomy (hence for a different dataset), the optimization must be repeated.
Similarly, we estimate the expectation with respect to q φ ( Y | x ) using Monte-Carlo methods, and estimate the gradients with respect to φ using reparameterization trick.
Depending on the observation model used for p θ ( · ) , we end up with a different loss function.
For example, we can have (cid:96) 2 or (cid:96) 1 distances between reconstruction and ground truth as loss functions for Gaussian and Laplacian observation models, respectively.
The main difference between [18] and our formulation is that we enforce the latent space to correspond to the Fourier domain coefﬁcients by using a nuFFT operator rather than a fully parameterized neural network.
2.2 Multi-Channel Acquisition and Reconstruction Models
Our overall network architecture is illustrated in Figure 1.
We represent the acquisition (encoding) model with the nuFFT operator that is parameterized by the sampling locations φ in k-space.
This enables us to represent k-space sampling locations as continuous variables and directly optimize for the sampling pattern by backpropagating through φ .
In the multi-channel MRI setting, the acquisition model admits
[Equation]
(6) where C is the number of channels,
S i ∈ C N × N is a diagonal matrix containing coil sensitivity proﬁles for coil i , and F nu ( φ ) :
C N → C M/C is the nuFFT operator at sampling locations φ .
Notice that 5 allows any deep neural network for reconstruction as the decoder.
In this work, we used an unrolled reconstruction network similar to [7, 8] that mimic the proximal gradient based solutions for non-smooth compressed sensing.
The unrolled architecture also incorporates MR physics, where the reconstruction is performed by alternating between data consistency (Fig. 1b) and proximal steps (Fig. 1c) for a ﬁxed number of iterations.
The data-consistency renders the reconstruction a function of both φ and θ , hence we share the parameters of the nuFFT encoder with the decoder network.
The unrolled nature of the network also allows for multiple gradient paths to φ from the decoder.
Experiments 3.1 Data
We used the "Stanford Fully Sampled 3D FSE Knees" dataset available in mridata.org [21] which contains 3D knee scans from 20 subjects.
Each 3D volume consists of 320 slices with a matrix size of 256 × 320 and was acquired with 8 channels.
The sensitivity maps of each coil were estimated using
ESPIRiT [22].
Each of the slices were treated as separate examples during training and validation.
The datasets were divided according to subjects where 15 subjects (4800 slices) were used for training, 2 subjects (640 slices) were used for validation, and 3 subjects (960 slices) were used for testing.
3.2 Network Implementation and Training Details
All of our models were implemented on TensorFlow and trained on a NVIDIA Titan Xp graphics card with 12GB of memory.
We used the TensorFlow implementation of nuFFT available in [23].
As the sampling locations change during the training, we did not use density compensation while calculating the adjoint nuFFT.
Instead we relied on the proximal block to compensate for the density.
The proximal block consists of 3 Residual Blocks each having 64 channels and we unrolled the network for 4 iterations.
We also used instance normalization at the output of convolutional layers.
We considered the Gaussian observation model for p θ ( · ) in Equation 5 which corresponds to the (cid:96)reconstruction loss.
Due to GPU memory constraints, we used a batch size of 2.
We also used the
Adam optimizer with a learning rate of 0 .
01 , β 1 = 0 .
9 , β 2 = 0 .
999 for network optimization.
To assess the image quality, we adopt peak signal-to-noise ratio (pSNR) and structural similarity index (SSIM) between the reconstruction and fully-sampled ground truth.
pSNR was evaluated directly on complex-valued images, while for SSIM complex-valued images are treated as two channels.
3.3 Experiment Details
We considered four acceleration factors in our experiments: R = 5 , 10 , 15 , 20 .
For each acceleration factor, we initialized the k-space sampling pattern ( φ ) by sampling realizations from a 2D Gaussian [Figure] [Figure] distribution centered at the origin (Variable Density Sampling [24]) and picked a variance σ = 0 .for
the measurement noise (cid:15) .
We then ran our model and compared the result with the case where we only optimized the reconstruction network.
Figure 2a,b shows pSNR and SSIM metrics, respectively.
We observe that joint optimization using our framework improves reconstruction quality.
Figure 3 shows the initial and optimized sampling patterns along with the reconstruction results on a sample test slice for R = 5 and R = 10 .
Sampling optimization accumulates more samples near k-space center and reduces the gaps between them which enhances reconstruction quality.
In addition, zero-ﬁlling reconstruction using the optimized samples have less severe aliasing artifacts, making it easier for the decoder network to remove these artifacts.
Figure 2c shows the point spread function (PSF) of the sampling pattern before and after optimization for R = 5 .
The sidelobes of the PSF are suppressed which reduces the amount of aliasing due to undersampling.
Conclusion and Discussion
This work explores the use of variational information maximization for learning the sampling pattern in MRI jointly with image reconstruction.
Using the nuFFT operator in the encoder enables continuous optimization of sampling variables.
On a knee MR dataset, we have shown that the optimized sampling pattern improves reconstruction quality highlighting possible beneﬁts that can be obtained by learning undersampling patterns.
One observation is that the peripheral samples do not change as much as the samples near origin during training.
Incorporating additional regularizers that emphasize high frequency information can improve the gradient ﬂow to these samples.
Our framework allows incorporating other MR physics related artifacts beyond the Gaussian measurement noise such as motion and off-resonance effects.
These artifacts can be simulated as a part of the forward model to ﬁnd the optimal sampling locations in the presence of such effects.
Acknowledgments and Disclosure of Funding
This work is supported by NIH R01EB009690 and R01EB026136, as well as GE Healthcare.