[Paper Review] Robustness of Deep Learning for Accelerated MRI: Benefits of Diverse Training Data

Kang Lin, Reinhard Heckel|arXiv (Cornell University)|Dec 16, 2023

Advanced X-ray and CT Imaging4 citations

TL;DR

This paper demonstrates that training deep learning models for accelerated MRI on a diverse combination of data distributions—such as different MRI scanners, anatomies, contrasts, and field strengths—significantly improves robustness to distribution shifts without sacrificing in-distribution performance. The key result is that a single model trained on a large, diverse dataset outperforms models trained on individual, narrower datasets, even closing much of the performance gap on out-of-distribution data.

ABSTRACT

Deep learning based methods for image reconstruction are state-of-the-art for a variety of imaging tasks. However, neural networks often perform worse if the training data differs significantly from the data they are applied to. For example, a model trained for accelerated magnetic resonance imaging (MRI) on one scanner performs worse on another scanner. In this work, we investigate the impact of the training data on a model's performance and robustness for accelerated MRI. We find that models trained on the combination of various data distributions, such as those obtained from different MRI scanners and anatomies, exhibit robustness equal or superior to models trained on the best single distribution for a specific target distribution. Thus training on such diverse data tends to improve robustness. Furthermore, training on such a diverse dataset does not compromise in-distribution performance, i.e., a model trained on diverse data yields in-distribution performance at least as good as models trained on the more narrow individual distributions. Our results suggest that training a model for imaging on a variety of distributions tends to yield a more effective and robust model than maintaining separate models for individual distributions.

Motivation & Objective

• To investigate whether training deep learning models on multiple MRI data distributions improves robustness to distribution shifts.
• To evaluate whether diverse training compromises in-distribution performance on individual distributions.
• To assess whether models trained on diverse data can generalize to unseen pathologies or non-healthy subjects.
• To examine the impact of training duration on robustness, particularly the risk of distributional overfitting.
• To validate that a single model trained on a large, diverse dataset outperforms models trained on isolated datasets like fastMRI.

Proposed method

• Training a U-Net, Vision Transformer (ViT), and VarNet on a combined dataset of 13 diverse MRI datasets, including fastMRI brain and knee, CC-359, Stanford 2D, M4Raw GRE, and NYU data.
• Preprocessing all datasets to follow the fastMRI convention, excluding slices with pure noise, resulting in 413k training slices.
• Using early stopping based on performance on the fastMRI knee validation set to prevent distributional overfitting.
• Evaluating in-distribution performance on fastMRI knee and brain, and out-of-distribution performance on four additional datasets with diverse scanner vendors, anatomical views, contrasts, and sequences.
• Comparing the diverse-dataset model against baselines trained only on fastMRI brain or fastMRI knee data, and against models fine-tuned on each out-of-distribution dataset.
• Leveraging existing code from the fastMRI repository for implementation and training.

Experimental results

Research questions

• RQ1 Does training a single deep learning model on multiple MRI data distributions improve its robustness to distribution shifts compared to training on individual distributions?
• RQ2 Does training on diverse data compromise in-distribution performance on any individual data source?
• RQ3 Can a model trained on diverse healthy subjects reconstruct images with pathologies accurately, even without seeing such data during training?
• RQ4 Does prolonged training lead to distributional overfitting, where out-of-distribution performance degrades despite continued in-distribution improvement?
• RQ5 Can a single model trained on a large, diverse dataset close the performance gap on out-of-distribution data compared to models trained specifically on that data?

Key findings

• A model trained on a diverse dataset spanning multiple MRI scanners, anatomies, contrasts, and sequences achieved superior out-of-distribution performance compared to models trained on individual datasets, including fastMRI.
• The model trained on the diverse dataset $π_{P}$ matched or exceeded the in-distribution performance of models trained on individual datasets like fastMRI brain or knee.
• On the CC-359 sagittal view dataset, the VarNet trained on $π_{P}$ nearly closed the distribution shift performance gap, achieving reconstruction quality close to that of a model fine-tuned on the same data.
• Out-of-distribution reconstructions from the diverse-dataset model showed higher image quality and fewer artifacts than those from fastMRI-trained models.
• Distributional overfitting was observed: out-of-distribution performance sharply declined with longer training, even as in-distribution performance continued to improve slightly.
• Early stopping was effective in preserving robustness by achieving near-optimal in-distribution performance without sacrificing out-of-distribution generalization.