[論文レビュー] Dual Multi-scale Mean Teacher Network for Semi-supervised Infection Segmentation in Chest CT Volume for COVID-19

Liansheng Wang, Jiacheng Wang+7 more|arXiv (Cornell University)|Nov 10, 2022

COVID-19 diagnosis using AI被引用数 5

ひとこと要約

本稿では、胸部CTボリュームにおける自己教師付き3次元COVID-19肺感染症セグメンテーションのためのデュアルマルチスケールミアンテイチャーネットワーク（DM²T-Net）を提案する。複数次元アテンションCNN（MDA-CNN）を用いて階層的マルチスケール特徴を捉え、複数の予測ヘッドを通じて学生・教師ネットワーク間の一貫性を強制する。本手法は、COVID-19-P20データセットで72.59%のDiceスコア、MosMedDataで60.19%のDiceスコアを達成し、低データ環境下でも優れた精度と一般化性能を示した。

ABSTRACT

Automated detecting lung infections from computed tomography (CT) data plays an important role for combating COVID-19. However, there are still some challenges for developing AI system. 1) Most current COVID-19 infection segmentation methods mainly relied on 2D CT images, which lack 3D sequential constraint. 2) Existing 3D CT segmentation methods focus on single-scale representations, which do not achieve the multiple level receptive field sizes on 3D volume. 3) The emergent breaking out of COVID-19 makes it hard to annotate sufficient CT volumes for training deep model. To address these issues, we first build a multiple dimensional-attention convolutional neural network (MDA-CNN) to aggregate multi-scale information along different dimension of input feature maps and impose supervision on multiple predictions from different CNN layers. Second, we assign this MDA-CNN as a basic network into a novel dual multi-scale mean teacher network (DM${^2}$T-Net) for semi-supervised COVID-19 lung infection segmentation on CT volumes by leveraging unlabeled data and exploring the multi-scale information. Our DM${^2}$T-Net encourages multiple predictions at different CNN layers from the student and teacher networks to be consistent for computing a multi-scale consistency loss on unlabeled data, which is then added to the supervised loss on the labeled data from multiple predictions of MDA-CNN. Third, we collect two COVID-19 segmentation datasets to evaluate our method. The experimental results show that our network consistently outperforms the compared state-of-the-art methods.

研究の動機と目的

• 胸部CTにおける感染症セグメンテーションの正確な実現を図るため、2次元ベースの手法が3次元空間的文脈やスライス間関係を捉えきれていないという限界を解消すること。
• 既存の3次元セグメンテーションモデルに見られるマルチスケール受容野の不足を解消するため、階層的特徴学習を可能にする多次元アテンション機構を設計すること。
• 自己教師付き学習フレームワークを構築することで、ラベルなしデータを効果的に活用し、限られた3次元CTデータのラベル付きデータの不足という課題に対処すること。
• マルチスケール一貫性損失とマルチ予測監督を統合することで、多様なデータセット間でのモデル一般化性能を向上させること。

提案手法

• 3次元CTボリュームの異なる空間次元に沿ってマルチスケール特徴を統合するため、複数次元アテンションCNN（MDA-CNN）を設計し、階層的表現学習を強化する。
• MDA-CNNをデュアルマルチスケールミアンテイチャーネットワーク（DM²T-Net）に統合し、学生および教師ネットワークが異なるネットワーク層からマルチスケール予測を生成する。
• 複数スケールおよびネットワーク深さにわたるラベルなしデータ上の学生・教師予測の一貫性を強制するため、マルチスケール一貫性損失を導入する。
• ラベルありデータに対する教師付き損失と、ラベルなしデータに対するマルチスケール一貫性損失を組み合わせた学習目的を採用し、一般化性能を向上させる自己教師付き学習を実現する。
• 教師ネットワークのパラメータは、学生の重みの指数的移動平均により更新され、学習の安定化と予測の一貫性向上を図る。
• 本手法の包括的評価に用いるために、2つの新しい3次元胸部CTデータセット—COVID-19-P20とMosMedData—を収集した。

実験結果

リサーチクエスチョン

• RQ1 次元特化アテンションを備えたマルチスケール3次元CNNは、単一スケールまたは2次元アプローチと比較して、感染症セグメンテーションの正確性を向上させることができるか？
• RQ2 提案されたマルチスケール一貫性損失は、自己教師付き3次元セグメンテーションにおけるラベルなしデータの活用にどの程度効果的か？
• RQ3 既存のSOTA手法と比較して、デュアルマルチスケールミアンテイチャーフレームワークは、多様なCT画像プロトコルや感染パターンにわたる一般化性能が優れているか？
• RQ4 複数のネットワーク層からのマルチ予測監督は、モデルのロバストネスおよびセグメンテーション性能をどの程度向上させるか？

主な発見

• 提案されたDM²T-Netは、COVID-19-P20データセットで72.59%のDiceスコアを達成し、SOTA手法を顕著に上回った。
• MosMedDataデータセットでは60.19%のDiceスコアを達成し、より困難で多様なデータセットでも強力な性能を示した。
• マルチスケール一貫性損失は60エポック後に安定して収束し、初期学習段階で段階的な増加を示しており、ラベルなしデータの効果的活用が確認された。
• 可視化比較により、教師ネットワークの予測が学生の予測よりも一貫して正確であることが示され、マルチスケール監督による知識蒸留の有効性が裏付けられた。
• 一般化分析の結果、DM²T-Netは1つのデータセットで学習し、別のデータセットでテストした場合でも優れた性能（Dice: 48.81±23.51、HD95: 35.34±30.57）を維持し、2D U-Net++および3D UA-MTを上回った。
• クロスデータセット評価において、本モデルは画像のばらつきや感染パターンの多様性に対してより高いロバストネスを示し、ベースラインモデルと比較してHD95値が低く、Diceスコアが高かった。