[論文レビュー] 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スキャナーや臓器、コントラスト、磁場強度を含む多様なデータ分布の組み合わせを用いた深層学習モデルのトレーニングが、分布シフトに対する耐性を著しく向上させることを示している。性能が低下することなく、インディストリビューション性能を維持する。主な結果として、1つのモデルを大規模で多様なデータセットでトレーニングすることで、個々の限定的データセットでトレーニングされたモデルよりも優れた性能を達成しており、特に分布外データに対する性能ギャップを大幅に縮小している。

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データ分布で深層学習モデルをトレーニングすることで、分布シフトに対する耐性が向上するかどうかを調査すること。
• 多様なデータでトレーニングすることで、個々のデータ分布におけるインディストリビューション性能が損なわれるかどうかを評価すること。
• 多様なデータでトレーニングされたモデルが、トレーニング中に観測されていない疾患や健康でない被験者に対しても一般化できるかどうかを評価すること。
• トレーニング期間が耐性に与える影響、特に分布的過適合のリスクを調査すること。
• 1つのモデルを大規模で多様なデータセットでトレーニングした場合、fastMRIのような個別のデータセットでトレーニングされたモデルよりも優れた性能を示すかどうかを検証すること。

提案手法

• fastMRIの脳・膝、CC-359、Stanford 2D、M4Raw GRE、NYUデータを含む13の多様なMRIデータセットを統合したデータセット上で、U-Net、ビジョントランスフォーマー（ViT）、VarNetをトレーニングする。
• すべてのデータセットをfastMRIの規格に従って前処理し、完全にノイズだけのスライスを除外した結果、413,000枚のトレーニングスライスが得られた。
• 分布的過適合を防ぐために、fastMRI膝のバリデーションセットでの性能に基づいて早期停止を実施する。
• fastMRI膝および脳のインディストリビューション性能を評価し、スキャナーメーカー、解剖的視点、コントラスト、シーケンスが多様な4つの追加データセットで分布外性能を評価する。
• fastMRI脳または膝データにのみトレーニングされたベースラインモデルと比較し、また、各分布外データセットでファインチューニングされたモデルとも比較する。
• 実装およびトレーニングに、fastMRIリポジトリの既存コードを活用する。

実験結果

リサーチクエスチョン

• RQ1 複数のMRIデータ分布で1つの深層学習モデルをトレーニングすることで、個々の分布でトレーニングした場合と比較して、分布シフトに対する耐性が向上するか？
• RQ2 多様なデータでトレーニングすることで、特定の1つのデータソースにおけるインディストリビューション性能が損なわれるか？
• RQ3 健康な被験者でトレーニングしたモデルが、トレーニング中に観測されていない疾患を有する画像を正確に再構築できるか？
• RQ4 長期間のトレーニングにより、インディストリビューション性能が改善し続ける一方で、分布外性能が低下する分布的過適合が生じるか？
• RQ5 大規模で多様なデータセットでトレーニングされた1つのモデルが、特定のデータセットでトレーニングされたモデルと比較して、分布外データにおける性能ギャップを埋めることができるか？

主な発見

• 複数のMRIスキャナーや臓器、コントラスト、シーケンスをカバーする多様なデータセットでトレーニングしたモデルは、fastMRIのような個々のデータセットでトレーニングしたモデルと比較して、より優れた分布外性能を達成した。
• 多様なデータセット $π_{P}$ でトレーニングされたモデルは、fastMRI脳や膝でトレーニングされたモデルのインディストリビューション性能と同等またはそれを上回った。
• CC-359の矢状断面データセットにおいて、$π_{P}$ でトレーニングされたVarNetは、分布シフトの性能ギャップをほぼ埋め、同じデータセットでファインチューニングされたモデルと同等の再構築品質に近づいた。
• 多様なデータセットでトレーニングされたモデルの分布外再構築画像は、fastMRIでトレーニングされたモデルのものよりも高品質で、アーチファクトが少なく、より優れた画像品質を示した。
• 分布的過適合が観察された：長期間のトレーニングにより、インディストリビューション性能がわずかに向上し続ける一方で、分布外性能は急激に低下した。
• 早期停止は、インディストリビューション性能をほぼ最適化しつつ、分布外一般化性能を損なわずに耐性を維持するのに有効であった。