INFusion: Diffusion Regularized Implicit Neural Representations for 2D and 3D accelerated MRI reconstruction

TL;DR

INFusion introduces diffusion-regularized Implicit Neural Representations to accelerate MRI reconstruction under undersampling. By coupling a learned diffusion prior with INR optimization through data-consistency and diffusion-based regularization on random slices, it enables improved 2D reconstructions and feasible 3D applications using 2D priors. The approach demonstrates lower reconstruction error in 2D brain data and shows feasibility for large-scale 3D knee data, with training times on the order of minutes per slice on contemporary GPUs. This work broadens INR-based MRI acceleration by decoupling priors from the measurement model and enabling 3D applicability via 2D diffusion models, potentially reducing scan time while maintaining image fidelity.

Abstract

Implicit Neural Representations (INRs) are a learning-based approach to accelerate Magnetic Resonance Imaging (MRI) acquisitions, particularly in scan-specific settings when only data from the under-sampled scan itself are available. Previous work demonstrates that INRs improve rapid MRI through inherent regularization imposed by neural network architectures. Typically parameterized by fully-connected neural networks, INRs support continuous image representations by taking a physical coordinate location as input and outputting the intensity at that coordinate. Previous work has applied unlearned regularization priors during INR training and have been limited to 2D or low-resolution 3D acquisitions. Meanwhile, diffusion based generative models have received recent attention as they learn powerful image priors decoupled from the measurement model. This work proposes INFusion, a technique that regularizes the optimization of INRs from under-sampled MR measurements with pre-trained diffusion models for improved image reconstruction. In addition, we propose a hybrid 3D approach with our diffusion regularization that enables INR application on large-scale 3D MR datasets. 2D experiments demonstrate improved INR training with our proposed diffusion regularization, and 3D experiments demonstrate feasibility of INR training with diffusion regularization on 3D matrix sizes of 256 by 256 by 80.

Figures (4)

• Figure 1: Our proposed approach, INFusion, first computes standard data consistency loss $L_{data}$ by evaluating the INR at discrete voxel locations, applying the MRI forward model to the resultant image, and then comparing to the acquired k-space data. A second loss $L_{diffusion}$ is computed by adding gaussian noise to random slices produced by the INR, running prior sampling with the diffusion model, and comparing the resultant prior sample to the INR produced random slices with perceptual loss. The diffusion loss guides the INR to produce images that better match the learned prior.
• Figure 2: Quantitative comparison of standard L1-Wavelet CS and INRs trained with Wavelet and our proposed regularization at $R = 8$ and $R = 9$ across 96 samples. Our proposed INFusion approach improves NRMSE across the test dataset.
• Figure 3: Reconstructions of discrete images on the (A.) 2D single-coil brain data at $R=4$ and (B.) 2D multi-coil brain data at $R=6$ with standard L1-wavelet Compressed Sensing (CS) and INRs with none, L1-wavelet, and our proposed diffusion regularization. The proposed INFusion approach yields images with lowest NRMSE and best qualitative image quality.
• Figure 4: (A.) INFusion outperforms no regularization on the modest, single-coil 3D dataset, but encoding all three spatial dimensions in the INR required prohibitively large GPU memory usage and compute. (B.) Applying diffusion regularization to random slices at each iteration and encoding only the x-y coordinates enables training of INRs with diffusion regularization on a realistically sized $256 \times 256 \times 80$ 3D k-space volume.