[論文レビュー] On-the-Fly Test-time Adaptation for Medical Image Segmentation

Jeya Maria Jose Valanarasu, Guo, Pengfei+2 more|arXiv (Cornell University)|Mar 10, 2022

Domain Adaptation and Few-Shot Learning被引用数 4

ひとこと要約

本稿では、バックプロパゲーションを伴わずに1枚のテスト画像にのみ適応するゼロショットでエピソード的なフレームワーク、On-the-Fly Test-Time Adaptationを提案する。Adaptive UNetとドメインコード駆動のアダプティブバッチ正規化を用いることで、2次元および3次元のドメインシフトにおいて最先端の性能を達成し、より厳しい臨床的制約のもとでも、既存のテスト時適応手法を上回る。

ABSTRACT

One major problem in deep learning-based solutions for medical imaging is the drop in performance when a model is tested on a data distribution different from the one that it is trained on. Adapting the source model to target data distribution at test-time is an efficient solution for the data-shift problem. Previous methods solve this by adapting the model to target distribution by using techniques like entropy minimization or regularization. In these methods, the models are still updated by back-propagation using an unsupervised loss on complete test data distribution. In real-world clinical settings, it makes more sense to adapt a model to a new test image on-the-fly and avoid model update during inference due to privacy concerns and lack of computing resource at deployment. To this end, we propose a new setting - On-the-Fly Adaptation which is zero-shot and episodic (i.e., the model is adapted to a single image at a time and also does not perform any back-propagation during test-time). To achieve this, we propose a new framework called Adaptive UNet where each convolutional block is equipped with an adaptive batch normalization layer to adapt the features with respect to a domain code. The domain code is generated using a pre-trained encoder trained on a large corpus of medical images. During test-time, the model takes in just the new test image and generates a domain code to adapt the features of source model according to the test data. We validate the performance on both 2D and 3D data distribution shifts where we get a better performance compared to previous test-time adaptation methods. Code is available at https://github.com/jeya-maria-jose/On-The-Fly-Adaptation

研究の動機と目的

• トレーニングデータとテストデータの間でドメインシフトが生じた場合の医療画像セグメンテーション性能の低下を解消すること。
• バックプロパゲーションを回避し、テスト時に全ターゲット分布へのアクセスを必要としない、臨床現場で実用可能な適応手法を開発すること。
• 個々の画像ごとにゼロショットで即時適応を可能にし、プライバシー保護や限られた計算リソースといった現実の臨床的制約に適合させること。
• 事前に学習されたエンコーダーから得られるドメインコードを用いて、モデルの特徴を動的に適応させるフレームワークを設計すること。
• 顕著なドメインシフトが生じる状況下で、多様な2次元および3次元医療画像セグメンテーションベンチマークを用いて手法を検証すること。

提案手法

• 推論時にモデル更新を行わず、エピソード的かつゼロショットな適応が可能な、新たな適応設定「On-the-Fly Test-Time Adaptation」を提案する。
• ドメインコードを用いて特徴を適応させるアダプティブバッチ正規化層を備えた、標準UNetを拡張した「Adaptive UNet」を導入する。
• 多数の医療画像を用いて事前学習されたエンコーダーを用いて、各テスト画像のドメインコードを生成し、モalityやスキャナ固有の特徴を捉える。
• ドメインコードを用いてエンコーダーおよびデコーダーのバッチ正規化層の累積統計を調整し、インスタンスごとの特徴適応を可能にする。
• ソースデータやテスト時の学習を必要とせず、事前学習モデルから得られるドメインプライアを用いて適応をガイドする。
• 各推論で1枚の画像にのみ適応を適用し、推論後に適応状態をリセットすることで、プライバシー保護と低リソース消費を実現する。

実験結果

リサーチクエスチョン

• RQ1 推論時にいかなるバックプロパゲーションも行わないテスト時適応手法が、高い性能を達成できるか？
• RQ2 全ターゲット分布が入手不可な状況下でも、1枚ごとのゼロショットで即時適応がどれほど有効か？
• RQ3 学習されたドメインコードを用いたアダプティブバッチ正規化は、2次元および3次元の医療画像分野におけるドメインシフトを効果的に処理できるか？
• RQ4 より厳しい臨床現場での展開制約のもとで、提案手法は既存のテスト時適応ベースラインを上回るか？
• RQ5 MRIシーケンス間や眼底画像デバイス間の異なるモダリティシフトにおいて、手法の一般化性能はいかがなものか？

主な発見

• Adaptive UNetを用いたOn-the-Fly Adaptationは、2次元眼底画像のドメインシフトにおいて63.14のDiceスコアを達成し、TENT（57.10）およびRN+CR（59.15）を上回った。
• 3次元MRIセグメンテーションにおいては、T1→T1ceシフトで60.66/58.73/39.30のDiceスコアを達成し、TENT（55.47/52.89/39.17）およびRN+CR（57.30/54.68/40.01）を上回った。
• 2次元および3次元の両設定で、9つのドメインシフトにおいて直接テストやすべてのベースラインと比較して一貫した性能向上を示した。
• 定性的な結果では、Adaptive UNetが正解とオラクル予測に近いセグメンテーションマスクを生成し、過剰セグメンテーションや誤分類を低減していることが確認された。
• 2次元シフトでは3次元シフトよりも性能が優れており、これはモダリティ間MRIシーケンスシフトの複雑さが影響していると考えられるが、依然として強力な性能を示した。
• アブレーション実験により、ドメインコードに基づく適応が極めて重要であることが確認され、その除去により非適応モデルと同程度の性能低下が生じた。