[論文レビュー] Rethinking Data Augmentation for Single-source Domain Generalization in Medical Image Segmentation

Zixian Su, Kai Yao+4 more|arXiv (Cornell University)|Nov 27, 2022

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

ひとこと要約

本論文は、医療画像セグメンテーションにおける単一ソースドメイン一般化のための新しいデータオーグメンテーションフレームワーク、Saliency-balancing Location-scale Augmentation (SLAug) を提案する。クラスレベルのロケーション・スケールオーグメンテーションとサリエンシー Guided フュージョンを組み合わせることで、オーグメンテーションデータの多様性と情報量を向上させ、腹部CT-MRIタスクでERM比4.3%のDiceスコア向上、心臓bSSFP-LGEタスクで2.1%の向上を達成し、最良の性能を発揮するとともに、理論的に一般化リスクの上限を導出する。

ABSTRACT

Single-source domain generalization (SDG) in medical image segmentation is a challenging yet essential task as domain shifts are quite common among clinical image datasets. Previous attempts most conduct global-only/random augmentation. Their augmented samples are usually insufficient in diversity and informativeness, thus failing to cover the possible target domain distribution. In this paper, we rethink the data augmentation strategy for SDG in medical image segmentation. Motivated by the class-level representation invariance and style mutability of medical images, we hypothesize that unseen target data can be sampled from a linear combination of $C$ (the class number) random variables, where each variable follows a location-scale distribution at the class level. Accordingly, data augmented can be readily made by sampling the random variables through a general form. On the empirical front, we implement such strategy with constrained B$\acute{ m e}$zier transformation on both global and local (i.e. class-level) regions, which can largely increase the augmentation diversity. A Saliency-balancing Fusion mechanism is further proposed to enrich the informativeness by engaging the gradient information, guiding augmentation with proper orientation and magnitude. As an important contribution, we prove theoretically that our proposed augmentation can lead to an upper bound of the generalization risk on the unseen target domain, thus confirming our hypothesis. Combining the two strategies, our Saliency-balancing Location-scale Augmentation (SLAug) exceeds the state-of-the-art works by a large margin in two challenging SDG tasks. Code is available at https://github.com/Kaiseem/SLAug .

研究の動機と目的

• 単一のソースドメインしか利用できない状況における医療画像セグメンテーションにおけるドメインシフトの課題に対処すること。
• 多様性と情報量に欠けるグローバル・オンリーやランダムオーグメンテーション戦略の限界を克服すること。
• 医療画像におけるクラスレベルの表現不変性とスタイルの可変性を活用し、より強力なオーグメンテーション戦略を設計すること。
• 提案されたオーグメンテーションフレームワーク下での一般化リスクの理論的上限を導出し、その有効性を検証すること。
• ターゲットドメインデータを必要とせず、モデルのロバスト性を向上させるプラグアンドプレイモジュールを開発すること。

提案手法

• 各クラスごとにロケーション・スケール分布からサンプリングされるオーグメント済み画像を線形結合で定義するモデルを提案：$ x_{U}^{c} = \alpha^{c}x_{S}^{c} + \beta^{c} $、ここで $ \alpha^{c} $ と $ \beta^{c} $ は学習可能なスケールおよびロケーションパラメータである。
• 画像全体のグローバルな強度シフトを模倣するため、グローバル・ロケーション・スケールオーグメンテーション（GLA）を実装する。
• 正例マスクを用いてクラス固有の変換を適用することで、クラスレベルでのドメインシフトを模倣するローカル・ロケーション・スケールオーグメンテーション（LLA）を導入する。
• 勾配ベースのサリエンシーマップを用いて融合の方向性と大きさをガイドする、サリエンシー・バランスの統合（SBF）メカニズムを設計する。
• GLAおよびLLAの両方で滑らかで現実的な画像変換を保証するため、制約付きBézier変換を適用する。
• 理論的分析により、提案されたオーグメンテーション戦略がSDG設定下で一般化リスクの上界を導くことを証明する。

実験結果

リサーチクエスチョン

• RQ1 医療画像におけるクラスレベルの表現不変性とスタイルの可変性を活用することで、ドメイン一般化のためのより効果的なデータオーグメンテーション戦略を設計できるか？
• RQ2 ロケーション・スケール分布から抽出されたクラスレベルの確率的変数の線形結合は、グローバル・オンリーまたはランダムオーグメンテーションに比べて、未観測ドメインの分布シフトをよりよく捉えられるか？
• RQ3 オーグメント済みサンプルのサリエンシー・ガイドド統合により、高勾配で意思決定境界に関連する領域に注目することで、モデルの一般化性能が向上するか？
• RQ4 提案されたオーグメンテーション戦略は、単一ソースドメイン一般化において理論的に一般化リスクが上限を持つか？
• RQ5 個々のモジュール（GLA, LLA, SBF）は全体の性能にどのように寄与しており、互いに相乗効果をもたらすか？

主な発見

• 腹部CT-MRIタスクにおいて、SLAugはERM比4.3%のDiceスコア向上を達成し、臓器全体で88.63、心臓で86.69のDiceスコアを記録した。
• 心臓bSSFP-LGEタスクでは、SLAugはERM比2.1%の性能向上を達成し、平均87.66%のDiceスコアを達成した。
• アブレーションスタディにより、GLA, LLA, SBFのすべてのモジュールが不可欠であることが確認された。いずれかを除去すると顕著な性能低下が生じた。
• バリエーション4（GLA + SBF）はバリエーション1（GLAのみ）を上回り、GLAで生成された画像のサリエンシーマップが効果的に分布拡大を導くことが示された。
• ランダム統合はSFBより2.9%低い性能を示し、サリエンシーに基づく統合がより情報量が多く、解釈可能であることを確認した。
• SBFはERMやランダム統合よりもソースドメインの知識をよりよく保持しており、それぞれ0.8および4.0のDiceポイントの向上を達成した。