[論文レビュー] Momentum Contrastive Voxel-wise Representation Learning for Semi-supervised Volumetric Medical Image Segmentation

Chenyu You, Ruihan Zhao+2 more|arXiv (Cornell University)|May 14, 2021

Domain Adaptation and Few-Shot Learning参考文献 31被引用数 17

ひとこと要約

本稿では、ビoxel単位（バッチレベル）および次元（特徴量レベル）の対照的学習を同時に活用することで、特徴量の多様性を向上させる、対照的ビoxel単位表現学習（CVRL）を提案する。インスタンス内特徴量多様性正則化と高次元3D空間的文脈モデリングを統合することで、ベンチマークデータセットにおいて10%のラベル付きデータ条件下でDiceスコアで最大5.21%の向上を達成し、最先端の性能を実現した。

ABSTRACT

Contrastive learning (CL) aims to learn useful representation without relying on expert annotations in the context of medical image segmentation. Existing approaches mainly contrast a single positive vector (i.e., an augmentation of the same image) against a set of negatives within the entire remainder of the batch by simply mapping all input features into the same constant vector. Despite the impressive empirical performance, those methods have the following shortcomings: (1) it remains a formidable challenge to prevent the collapsing problems to trivial solutions; and (2) we argue that not all voxels within the same image are equally positive since there exist the dissimilar anatomical structures with the same image. In this work, we present a novel Contrastive Voxel-wise Representation Learning (CVRL) method to effectively learn low-level and high-level features by capturing 3D spatial context and rich anatomical information along both the feature and the batch dimensions. Specifically, we first introduce a novel CL strategy to ensure feature diversity promotion among the 3D representation dimensions. We train the framework through bi-level contrastive optimization (i.e., low-level and high-level) on 3D images. Experiments on two benchmark datasets and different labeled settings demonstrate the superiority of our proposed framework. More importantly, we also prove that our method inherits the benefit of hardness-aware property from the standard CL approaches.

研究の動機と目的

• ボリュメトリックセグメンテーションにおいて、ラベル付き3D医療スキャンの限界を克服するため、大規模なラベルなしデータを活用すること。
• 対照的学習における次元的崩壊を克服するため、インスタンス内（特徴量レベル）の制約を導入して特徴量の多様性を促進すること。
• 3D特徴空間における低レベルの空間相関と高レベルの文脈表現を同時にモデリングすることで、セグメンテーション精度を向上させること。
• 既存の半教師付き学習パイプラインと互換性がある、プラグアンドプレイ可能な対照的フレームワークを開発すること。

提案手法

• 3D部分空間における低レベルの対照的学習と、特徴空間における高レベルの対照的学習を組み合わせた、新しい二段階対照的学習戦略を導入。これにより、局所的な解剖的詳細を捉え、グローバルな文脈をモデル化できる。
• 各3Dスキャンのチャネル次元にわたる特徴量多様性を強制する次元的対照的目的関数を提案。これにより、自明な解への収束を防ぐ。
• 高次元表現を符号化する3Dプロジェクションヘッドを採用。これにより、データオーグメンテーションに対して不変な判別性の高い埋め込みを学習可能になる。
• 一貫性損失を正則化項として統合。これによりモデルのロバスト性が向上し、既存の半教師付きフレームワークとも互換性を保てる。
• トレーニング中のモーメンタムエンコーダーの安定化を図るため、モーメンタムベースの更新を採用。これにより、異なるビュー間での表現の一貫性が向上する。
• 空間的および特徴量次元の両方で強力なオーグメンテーションを適用。新規に提案された特徴量レベルの対照的損失により、オーグメント済みビューが同一ベクトルに収束することを防ぐ。

実験結果

リサーチクエスチョン

• RQ1 インスタンス内（特徴量レベル）の対照的学習は、3D医療画像表現学習における次元的崩壊を効果的に防止できるか？
• RQ2 低レベルおよび高レベルの対照的目的関数を同時に最適化することで、低ラベルスアップ条件下でもセグメンテーション性能が向上するか？
• RQ3 限られたアノテーション条件下で、提案手法は最先端の半教師付き3Dセグメンテーションモデルと比較して境界精度および一般化性能に優れているか？
• RQ4 ビoxel単位および次元的対照的学習は、モデル性能にどのように独立してかつ相乗的に寄与しているか？

主な発見

• LA-Netデータセットにおいて10%のラベル付きデータ（8枚のラベル付きスキャン）条件下で、CVRLはDiceスコア88.56%を達成。ベースラインより5.47ポイントの向上を示した。
• 20%のラベル付きデータ条件下では、Diceスコア90.45%、Jaccardインデックス83.02%を達成。境界精度と形状の一貫性に優れた性能を示した。
• アブレーションスタディの結果、高レベル・低レベル・一貫性損失を併用することで最良の性能が得られ、ベースラインと比較してASDが48.3%、95HDが58.0%低減した。
• 膵臓データセットにおいて、CVRLは最先端手法と比較してDiceスコアで3.25～5.21%の相対的改善を達成。解剖的構造の多様性にわたるロバストネスを示した。
• 本手法は、対照的学習が有するハードネスに敏感な性質を引き継ぎ、複雑な3D解剖的文脈においてより判別性の高い表現を学習可能である。
• 可視化結果から、極度のラベル不足下でも、境界領域において特に正確で一貫性のあるセグメンテーションマスクが得られることを確認した。