Inference Stage Denoising for Undersampled MRI Reconstruction

TL;DR

This paper tackles the challenge of generalising MRI reconstruction to distribution shifts and noise without resorting to data augmentation. It introduces a conditional hypernetwork that modulates a reconstruction model via AdaIN conditioning on a hyperparameter $oldsymbol{lambda}$, enabling robust inference across varying noise levels. A cosine-based scheduler accelerates training by progressively adjusting the DC influence, yielding faster convergence and better performance than baselines. Experiments on FastMRI knee data show that the approach achieves superior image quality, especially under low-SNR conditions, and that augmentation offers limited gains only in specific high-noise scenarios, underscoring the practicality of augmentation-free robustness in clinical settings.

Abstract

Reconstruction of magnetic resonance imaging (MRI) data has been positively affected by deep learning. A key challenge remains: to improve generalisation to distribution shifts between the training and testing data. Most approaches aim to address this via inductive design or data augmentation. However, they can be affected by misleading data, e.g. random noise, and cases where the inference stage data do not match assumptions in the modelled shifts. In this work, by employing a conditional hyperparameter network, we eliminate the need of augmentation, yet maintain robust performance under various levels of Gaussian noise. We demonstrate that our model withstands various input noise levels while producing high-definition reconstructions during the test stage. Moreover, we present a hyperparameter sampling strategy that accelerates the convergence of training. Our proposed method achieves the highest accuracy and image quality in all settings compared to baseline methods.

Figures (4)

• Figure 1: Overview of the training process. The orange arrows indicate the data flow. The scheduler samples various $\lambda$ during training, and the thick blue arrows show $\lambda$'s involvement.
• Figure 2: Detailed illustration of how $\lambda$ is integrated. During training, $\lambda$ is sampled and the MLP generates $\gamma_\lambda$ and $\beta_\lambda$, which corrupt the training data representation $\mathbf{z}_c$ using AdaIN huang2017arbitrary; During inference, we select the optimal $\lambda_{opt}$ to help the noisy representation $\mathbf{z}_n$ decode to the optimised generation.
• Figure 3: A comparison of training schemes. The fixed value (blue, $\lambda=0$) is only optimised for a certain case. The uniform (green) performs the worst. The proposed scheduler (orange) shows the fastest convergence, and the advantage of using the proposed scheduler is to include a variety of choices in a single model.
• Figure 4: A comparison of the reconstruction performance among models in a $10^{-4}$ level of noise input reveals that Cond has superior results compared to Unet and DIDN. Reconstruction generated by Cond exhibits clear structure and preserve more details compared to the other models.