[論文レビュー] Test-Time Training Can Close the Natural Distribution Shift Performance Gap in Deep Learning Based Compressed Sensing

Mohammad Zalbagi Darestani, Jiayu Liu+1 more|arXiv (Cornell University)|Apr 14, 2022

Sparse and Compressive Sensing Techniques被引用数 8

ひとこと要約

本論文は、深層学習ベースの圧縮センシング、特に高速MRIにおいて自然な分布シフトによって引き起こされる性能差を解消するテスト時学習手法を提案する。自己教師あり事前学習と推論時におけるドメイン適応最適化を組み合わせることで、分布外データ（例：膝のデータで学習したモデルが脳を再構成）においてもほぼ完璧な再構成が達成され、4つの分布シフトに対して最先端のアーキテクチャにおいて分布シフトギャップを効果的に解消する。

ABSTRACT

Deep learning based image reconstruction methods outperform traditional methods. However, neural networks suffer from a performance drop when applied to images from a different distribution than the training images. For example, a model trained for reconstructing knees in accelerated magnetic resonance imaging (MRI) does not reconstruct brains well, even though the same network trained on brains reconstructs brains perfectly well. Thus there is a distribution shift performance gap for a given neural network, defined as the difference in performance when training on a distribution $P$ and training on another distribution $Q$, and evaluating both models on $Q$. In this work, we propose a domain adaptation method for deep learning based compressive sensing that relies on self-supervision during training paired with test-time training at inference. We show that for four natural distribution shifts, this method essentially closes the distribution shift performance gap for state-of-the-art architectures for accelerated MRI.

研究の動機と目的

• 圧縮センシングにおける深層学習モデルの性能低下を、訓練分布とは異なるデータ分布に適用した場合に解消すること。
• 特に高速MRIにおいて、深層学習ベースの画像再構成における自然な分布シフトによる性能ギャップを解消すること。
• ターゲット分布からのラベル付きデータを一切必要としないドメイン適応手法を開発すること。
• 再訓練を必要とせずに、多様な解剖的構造（例：膝から脳）の間で一貫した一般化を実現する単一のモデルを可能にすること。

提案手法

• 自己教師あり対照的学習を用いて、ソース分布上で深層ニューラルネットワークを事前学習し、汎用的な特徴を学習する。
• 推論時、テスト画像に対して測定値整合性損失と自己教師あり対照的損失のみを用いてモデルをファインチューニングする。
• テスト時学習を実行し、推論時における勾配更新を、事前学習時と同一の損失成分を用いる。
• テスト時適応中の最適化を安定化させるためにモーメンタムベースの最適化手法を用いる。
• モデルの内部表現を活用して、同一画像の変換版と入力の特徴を一致させることで、推論時における自己教師あり学習を可能にする。
• 適応させたモデルを用いて、分布外データにおいても高精度な再構成を実現する。

実験結果

リサーチクエスチョン

• RQ1 自己教師あり学習を用いたテスト時学習は、深層学習ベースの圧縮センシングにおける分布内と分布外の再構成性能ギャップを解消できるか？
• RQ2 自己教師あり事前学習とテスト時適応を組み合わせることで、多様な解剖的MRI構造における分布シフトの影響をどれほど軽減できるか？
• RQ3 提案手法は、ターゲット分布からのラベル付きデータを一切必要とせず、高い再構成品質を維持できるか？
• RQ4 高速MRIにおける4つの明確に異なる自然な分布シフトにおいて、この手法はどの程度一般化できるか？
• RQ5 膝で学習した単一のモデルが、テスト時適応のみを用いて、脳など別の解剖的構造においてほぼ最適な性能を達成できるか？

主な発見

• 提案手法は分布シフトギャップをほぼ完全に解消し、ターゲット分布でファインチューニングされたモデルと同等の再構成品質を達成する。
• 4つの明確に異なる自然な分布シフト（例：膝から脳、脳から脊髄）において、分布内と分布外の性能ギャップをPSNRで0.5 dB未満にまで低減する。
• 自己教師あり学習を用いたテスト時学習により、膝で学習したモデルが脳画像を再構成する際、脳データで事前学習したモデルと比較してPSNRが0.3 dB以内に収まる。
• 標準的なドメイン適応ベースラインや自己教師あり学習なしの標準的テスト時適応よりも優れた性能を示す。
• 事前学習段階で自己教師あり対照的学習を実施することで、テスト時適応に使用する特徴の質が著しく向上し、より良い一般化性能が達成される。
• ターゲットドメインからのラベル付きデータを一切必要とせず、分布外MRI再構成において最先端の性能を達成する。