[论文解读] 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 Imaging被引用 4

一句话总结

本论文表明，在多种不同的数据分布组合上训练深度学习模型用于加速磁共振成像（MRI），例如不同MRI扫描仪、解剖结构、对比度和场强，能显著提升对分布偏移的鲁棒性，同时不损害分布内性能。关键结果是，一个在大规模、多样化数据集上训练的单一模型，其性能优于在单一、狭窄数据集上训练的模型，甚至在分布外数据上大幅缩小了性能差距。

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.

研究动机与目标

• 调查在多个MRI数据分布上训练深度学习模型是否能提升对分布偏移的鲁棒性。
• 评估多样化训练是否会影响在任一单一分布上的分布内性能。
• 评估在多样化数据上训练的模型是否能泛化到未见过的病理或非健康受试者。
• 检查训练时长对鲁棒性的影响，特别是分布过拟合的风险。
• 验证在大规模、多样化数据集上训练的单一模型是否优于仅在孤立数据集（如fastMRI）上训练的模型。

提出的方法

• 在包含fastMRI脑部和膝部、CC-359、Stanford 2D、M4Raw GRE和NYU数据在内的13个多样化MRI数据集的组合数据集上，训练U-Net、视觉Transformer（ViT）和VarNet。
• 将所有数据集预处理为符合fastMRI规范，排除纯噪声的切片，最终得到41.3万张训练切片。
• 基于fastMRI膝部验证集的表现使用早停法，以防止分布过拟合。
• 在fastMRI膝部和脑部分布内数据集上评估分布内性能，在另外四个具有不同扫描仪厂商、解剖视角、对比度和序列的数据集上评估分布外性能。
• 将多样化数据集模型与仅在fastMRI脑部或膝部数据上训练的基线模型，以及在每个分布外数据集上微调的模型进行比较。
• 复用fastMRI代码库中的现有代码进行实现和训练。

实验结果

研究问题

• RQ1 在多个MRI数据分布上训练单一深度学习模型，是否相比在单一分布上训练，能提升其对分布偏移的鲁棒性？
• RQ2 在多样化数据上训练是否会影响任一单一数据源上的分布内性能？
• RQ3 在多样化健康受试者上训练的模型，是否能在未在训练中见过病理数据的情况下，准确重建病理图像？
• RQ4 长时间训练是否会导致分布过拟合，即尽管分布内性能持续提升，但分布外性能却下降？
• RQ5 在大规模、多样化数据集上训练的单一模型，是否能缩小与在特定数据上专门训练的模型在分布外数据上的性能差距？

主要发现

• 在涵盖多种MRI扫描仪、解剖结构、对比度和序列的多样化数据集上训练的模型，其分布外性能优于在单一数据集上训练的模型，包括fastMRI。
• 在多样化数据集 $π_{P}$ 上训练的模型，在分布内性能上与在fastMRI脑部或膝部等单一数据集上训练的模型相当或更优。
• 在CC-359矢状位数据集上，基于 $π_{P}$ 训练的VarNet几乎完全弥合了分布偏移的性能差距，其重建质量接近于在相同数据上微调的模型。
• 与fastMRI训练的模型相比，多样化数据集模型生成的分布外重建图像质量更高，伪影更少。
• 观察到分布过拟合现象：随着训练时间延长，分布外性能急剧下降，尽管分布内性能仍在轻微提升。
• 早停法有效保留了鲁棒性，实现了接近最优的分布内性能，同时未牺牲分布外泛化能力。