[論文レビュー] Boosting Semi-supervised Image Segmentation with Global and Local Mutual Information Regularization

Jizong Peng, Marco Pedersoli+1 more|arXiv (Cornell University)|Mar 8, 2021

Advanced Neural Network Applications参考文献 59被引用数 22

ひとこと要約

本論文は、グローバルおよびローカルな特徴表現における相互情報（MI）正則化を用いて、性能を向上させる半教師あり画像セグメンテーション手法を提案する。変換された同一画像のビュー間のクラスタ割り当て間のMIを最大化することでエンコーダー内のグローバル表現不変性を強化し、ネットワーク段階間の一貫性を保証することで、医療画像ベンチマークにおいて5%のラベル付きデータでのみ、準教師ありの精度に近い結果を達成する。

ABSTRACT

The scarcity of labeled data often impedes the application of deep learning to the segmentation of medical images. Semi-supervised learning seeks to overcome this limitation by exploiting unlabeled examples in the learning process. In this paper, we present a novel semi-supervised segmentation method that leverages mutual information (MI) on categorical distributions to achieve both global representation invariance and local smoothness. In this method, we maximize the MI for intermediate feature embeddings that are taken from both the encoder and decoder of a segmentation network. We first propose a global MI loss constraining the encoder to learn an image representation that is invariant to geometric transformations. Instead of resorting to computationally-expensive techniques for estimating the MI on continuous feature embeddings, we use projection heads to map them to a discrete cluster assignment where MI can be computed efficiently. Our method also includes a local MI loss to promote spatial consistency in the feature maps of the decoder and provide a smoother segmentation. Since mutual information does not require a strict ordering of clusters in two different assignments, we incorporate a final consistency regularization loss on the output which helps align the cluster labels throughout the network. We evaluate the method on four challenging publicly-available datasets for medical image segmentation. Experimental results show our method to outperform recently-proposed approaches for semi-supervised segmentation and provide an accuracy near to full supervision while training with very few annotated images.

研究の動機と目的

• ディープラーニングベースのセマンティックセグメンテーションにおいて、ラベル付き医療画像が限られているという課題に対処すること。
• 表現不変性と空間的一致性を活用して、未ラベルデータを活用することで半教師あり学習を改善すること。
• 大規模バッチサイズを必要とするか、高次元連続特徴で困難を抱える既存の一貫性ベースおよび対照学習手法の限界を克服すること。
• 離散的クラスタリングと相互情報に基づく正則化を効果的に用いて、グローバルおよびローカルな特徴表現を正則化する手法を開発すること。

提案手法

• 同じ画像の変換されたビュー間のクラスタ割り当て間の相互情報（MI）を最大化することで、エンコーダー内でのグローバル表現不変性を強制する。
• 投影ヘッドを用いて連続的特徴埋め込みを離散的クラスタ割り当てにマッピングし、連続的密度推定を避けることで、効率的なMI推定を可能にする。
• 隣接する空間的特徴間のMIを最大化することで、デコーダー特徴マップにおける空間的一致性を促進するローカルMI損失を導入する。
• 異なるネットワーク段階間のクラスタラベルを一致させるために、最終的な一貫性正則化損失を適用し、整合的なセマンティックグループ化を保証する。
• 教師あり損失とMIに基づくグローバルおよびローカル正則化を統合最適化フレームワークで組み合わせ、二段階の事前学習・微調整設定を回避する。
• 自明なクラスタリング解を避けるために、エントロピーに基づく定式化を用いた対称的MI推定手法を採用する。

実験結果

リサーチクエスチョン

• RQ1 離散的クラスタ割り当てにおける相互情報正則化は、半教師あり医療画像セグメンテーションを改善できるか？
• RQ2 グローバルおよびローカルMI正則化を組み合わせることで、特徴表現とセグメンテーション精度がどのように向上するか？
• RQ3 教師あり損失とMIベース損失を同時に最適化することは、二段階の事前学習・微調整戦略を上回るか？
• RQ4 MIベース手法は、医療画像セグメンテーションにおいて最小限のラベル付きデータで準教師あり性能に近い結果を達成できるか？
• RQ5 標準の一貫性ベースおよび対照学習アプローチと比較して、提案手法は頑健性および効率性において優れているか？

主な発見

• 提案手法は、4つの公開医療画像セグメンテーションデータセットにおいて、Mean Teacher やエントロピー最小化といった最先端の半教師あり手法を上回る性能を達成した。
• ACDCデータセットでは、5%のラベル付きデータでのみ、平均DSCが86.7%に達し、完全教師ありの性能に近づいた。
• 二段階の事前学習・微調整ベースラインと比較して、統合最適化戦略により、セグメンテーション精度が5.10%向上した。これは、ラベル付きデータを用いたエンドツーエンドのMI正則化の利点を示している。
• 3回の独立実験において安定的かつ一貫性のある性能を示し、標準偏差が低く、ランダム初期化に対する頑健性を示した。
• アブレーションスタディにより、グローバルおよびローカルMI損失の両方が性能向上に顕著に寄与していることが確認された。グローバル損失は不変性を向上させ、ローカル損失は空間的滑らかさを強化した。
• 投影ヘッドによる離散的クラスタ割り当ての使用により、連続的MI推定の計算負荷を回避し、効率的なMI推定が可能になった。