NO-MRI: A Unified Model for Compressed Sensing MRI Across Undersampling Patterns

Armeet Singh Jatyani*, Jiayun Wang*, Aditi Chandrashekar, Zihui Wu, Miguel Liu-Schiaffini, Bahareh Tolooshams, Animashree Anandkumar

California Institute of Technology

Compressed Sensing MRI reconstructs images of the body's internal anatomy from undersampled measurements, thereby reducing the scan time - the time subjects need to remain still. Recently, deep neural networks have shown great potential for reconstructing high-fidelity images from highly undersampled measurements in the frequency space. However, one needs to train multiple models for different undersampling patterns and desired output image resolutions, since most networks operate on a fixed discretization. Such approaches are highly impractical in clinical settings, where undersampling patterns and image resolutions are frequently changed to accommodate different real-time imaging and diagnostic requirements. We propose a unified model robust to different measurement undersampling patterns and image resolutions in compressed sensing MRI. Our model is based on neural operators, a discretization-agnostic architecture. Neural operators are employed in both image and measurement space, which capture local and global image features for MRI reconstruction. Empirically, we achieve consistent performance across different undersampling rates and patterns, with an average 11 percent SSIM and 4dB PSNR improvement over a state-of-the-art CNN, End-to-End VarNet. For efficiency, our inference speed is also 1,400x faster than diffusion methods. The resolution-agnostic design also enhances zero-shot super-resolution and extended field of view in reconstructed images. Our unified model offers a versatile solution for MRI, adapting seamlessly to various measurement undersampling and imaging resolutions, making it highly effective for flexible and reliable clinical imaging.

Introduction

Magnetic Resonance Imaging (MRI) is a widely-used non-invasive imaging technology, essential in medical and scientific fields such as neurosurgery, clinical oncology, diagnostic testing, neuroscience, and pharmaceutical research. However, its slow data acquisition process often requires patients to remain still for extended periods, limiting its efficiency and practicality.

Compressed Sensing (CS) techniques offer a solution by enabling MRI at sub-Nyquist rates, significantly reducing scan times. Traditional CS approaches rely on sparse priors in transform domains, while recent deep learning models take this a step further, learning deeper priors to improve reconstruction performance. Despite these advances, existing models often struggle to generalize across varying resolutions. Existing approaches need to be retrained for different undersampling patterns and rates, quickly becoming intractable.

We propose a unified neural operator model for accelerated MRI reconstruction from undersampled measurements across various patterns and rates.

$4\times$ undersampling patterns used for a variety of different diagnostic purposes. Top (rectilinear): Equispaced, Random, Magic. Bottom (irregular): Gaussian, Radial, Poisson. To reconstruct across 6 patterns and 3 subsampling rates, we would have to train 18 separate models. Our model can be trained once for multiple patterns and rates.

Neural Operators (NOs) are a cutting-edge deep learning framework that address this challenge by learning mappings between infinite-dimensional function spaces. Unlike traditional methods, NOs are resolution-agnostic, making them ideal for tasks involving data at multiple resolutions. By integrating discrete-continuous (DISCO) convolutions, a realization of the neural operator framework, we propose a unified model robust to various undersampling patterns and resolutions in compressed sensing MRI. This approach captures both local and global features of images, ensuring consistent performance and enabling advanced capabilities like zero-shot super-resolution and extended field-of-view reconstruction!

Join us as we explore how this novel framework outperforms state-of-the-art methods and sets a new benchmark in efficient, flexible MRI imaging.

Neural Operator Architecture

We follow an unrolled network design. First, the $NO_k$ inpaints missing measurements in measurement (frequency/$k$ space). Then, the inpainted $k$-space is brought to image space through an inverse Fourier transform 𝓕$^{-1}$ and finally iteratively processed by $n$ cascades of $NO_i$ and data-consistency modules.

The $NO_k$ neural operator makes our framework agnostic to different measurement undersampling patterns, and the $NO_i$ neural operator makes the framework agnostic to different image resolutions. Accurate MRI reconstruction relies on capturing both local and global image features, and our proposed architecture incorporates inductive biases to handle both effectively. The learning in both frequency and image space allows the model to capture both local and global features of images due to the duality of the Fourier transform that connects the frequency and image space. The resolution-agnostic design also enables superresolution in both frequency and image space, allowing the extended field of view (FOV) and super-resolution of the reconstructed MR images.

The $NO_i$ is implemented as a UDNO that operates in image space and the $NO_k$ is implemented as a UDNO that operates in measurement (frequency) space. We describe our in-place resolution invariant replacement for a U-shaped architecture, the UDNO, and the DISCO convolution kernels that enable resolution invariance.

U-Shaped DISCO Neural Operator (UDNO)

We introduce the U-Shaped DISCO Neural Operator, which is an in-place resolution invariant replacement for the U-Net. We follow the encoder/decoder architecture of the U-Net closely, replacing conventional convolutions with corresponding DISCO layers. The UDNO model is available for use in the neuraloperator library.

Discrete-Continuous Spherical Convolutions (DISCO)

When input resolution increases, standard convolutional kernels shrink in size relative to the feature map. While kernels can be interpolated to larger sizes to mitigate some of these issues, this often introduces artifacts at new resolutions.

In contrast, DISCO kernels are parametrized in function space, ensuring that their size remains consistent across resolutions. This inherent scaling makes DISCO kernels agnostic to resolution changes.

Zero-Shot Super-resolution and Super-FOV

Left: Super-FOV $NO_k$. Right: Super-Resolution $NO_i$.

Super-FOV $NO_k$

Zero-shot extended FOV reconstructions: Our NO model shows fewer artifacts and higher PSNR in the reconstructed brain slices compared to the CNN-based E2E-VN on $4\times$ Gaussian, despite neither model seeing data outside the initial $160 \times 160$ FOV during training.

Increasing resolution in $k$-space (frequency) corresponds to denser discretization between fixed minimum and maximum values, while the overall domain range remains constant. Due to the duality between frequency and image space in MRI, this corresponds to higher FOV. We demonstrate that our NO model which is trained for $160 \times 160$ low FOV reconstruction (top) can reconstruct larger FOV $320 \times 320$ images in a zero-shot setting.

Super-Resolution $NO_i$

Zero-shot super-resolution reconstructions in image space on $2\times$ radial: With input resolution increased to $640 \times 640$ through bilinear interpolation, our NO model preserves reconstruction quality, while E2E-VN produces visible artifacts.

Our NO can be trained on lower resolution (image-space) $320 \times 320$ images; then our $NO_i$ can operate on higher-resolution $640 \times 640$ images for super-resolution outputs!

Reconstruction on Multiple Patterns

MRI reconstructions with different undersampling patterns for different methods: NO (ours), E2E-VN++, E2E-VN, L1-Wavelet (compressed sensing), and CSGM (diffusion). Our NO reconstructs high-fidelity images across various undersampling patterns.

Performance across different undersampling patterns and rates of NO (ours) and baseline methods: End-to-End VarNet, diffusion and learning-free (classical). Our NO remains relatively consistent in performance when evaluated at different undersampling patterns and rates. Note that a high undersampling rate makes the task more difficult and thus a worse score is expected.

Citation

@article{jatyani2024unifying,
  title   = {Unifying Subsampling Pattern Variations for Compressed Sensing MRI with Neural Operators},
  author  = {Jatyani, Armeet Singh and Wang, Jiayun and Wu, Zihui and Liu-Schiaffini, Miguel and Tolooshams, Bahareh and Anandkumar, Anima},
  journal = {arXiv preprint arXiv:2410.16290},
  year    = {2024}
}