[論文レビュー] Multi-Task Attention-Based Semi-Supervised Learning for Medical Image Segmentation

Shuai Chen, Gerda Bortsova+3 more|arXiv (Cornell University)|Jul 29, 2019

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

ひとこと要約

本論文では、前景と背景の再構築を分離するためにアテンション機構を用いたアテンションベースのマルチタスク半教師あり学習（MASSL）を提案する。この手法は、限定的なラベル付きデータで学習された場合、完全に教師ありの手法や事前学習済みのCNNを上回り、交互に学習する戦略とアテンション機構により、脳腫瘍および白色髄質巣状病変のデータセットにおける特徴の識別性とセグメンテーション精度が顕著に向上する。

ABSTRACT

We propose a novel semi-supervised image segmentation method that simultaneously optimizes a supervised segmentation and an unsupervised reconstruction objectives. The reconstruction objective uses an attention mechanism that separates the reconstruction of image areas corresponding to different classes. The proposed approach was evaluated on two applications: brain tumor and white matter hyperintensities segmentation. Our method, trained on unlabeled and a small number of labeled images, outperformed supervised CNNs trained with the same number of images and CNNs pre-trained on unlabeled data. In ablation experiments, we observed that the proposed attention mechanism substantially improves segmentation performance. We explore two multi-task training strategies: joint training and alternating training. Alternating training requires fewer hyperparameters and achieves a better, more stable performance than joint training. Finally, we analyze the features learned by different methods and find that the attention mechanism helps to learn more discriminative features in the deeper layers of encoders.

研究の動機と目的

• ラベル付き医療画像が限られている課題に対処するため、半教師あり学習を用いて大量の未ラベルデータを活用する。
• U-Netアーキテクチャにおける標準的なオートエンコーダー基盤の再構築手法は、スキップ接続のため、クラス固有の情報を保持できないという限界を克服する。
• アテンション機構を介したセグメンテーションに配慮した再構築を統合することで、深層エンコーダーにおける特徴学習を改善する。
• 2つのマルチタスク学習戦略（統合学習と交互学習）がセグメンテーション性能に与える影響を評価する。
• アテンションガイドド再構築が、標準的な再構築と比較して、エンコーダーのより深い層におけるより識別性の高い特徴を学習することを示す。

提案手法

• 共有エンコーダーがセグメンテーションヘッド（スキップ接続あり）と再構築ヘッド（スキップ接続なし）に接続されるマルチタスクフレームワークを提案し、U-Netに類似したアーキテクチャを形成する。
• ソフトセグメンテーション予測から前景および背景マスクを生成するアテンション機構を導入し、クラスに配慮した再構築を可能にする。
• アテンションマスクを用いてクラス固有の領域を分離することで、前景および背景領域を個別に重み付きMSEで再構築するための再構築損失（L2）を定義する。
• 2つの学習戦略を実装する：統合学習（両タスクの同時最適化）と交互学習（タスク固有の学習率を用いた逐次最適化）。
• 両タスクに共有エンコーダーを用い、再構築目的がエンコーダーにおけるより識別性の高い特徴の学習をガイドする。特に、より深い層で顕著に効果を発揮する。
• 本手法を2つの医療画像処理タスクに適用する：脳腫瘍（BraTS18）および白色髄質巣状病変（WMH17）のセグメンテーション。ラベル付き画像は少数、未ラベルデータは多数を用いる。

実験結果

リサーチクエスチョン

• RQ1 アテンションベースの再構築目的は、標準的なオートエンコーダー基盤の再構築と比較して、半教師あり医療画像セグメンテーションの性能を向上させるか？
• RQ2 アテンションを介したセグメンテーションと再構築の統合は、標準的なマルチタスク学習と比較して、エンコーダーにおけるより識別性の高い特徴を学習するか？
• RQ3 統合学習と交互学習の2つの戦略は、マルチタスク半教師ありセグメンテーションにおける性能と安定性にどのように影響するか？
• RQ4 限定的なラベル付き画像しか利用できない状況下で、本手法は完全に教師ありのCNNや事前学習ベースラインをどの程度上回るか？
• RQ5 アテンション機構は、線形識別力で測定した場合、エンコーダーのより深い層で学習される特徴の質にどのように影響を与えるか？

主な発見

• 交互学習を用いたMASSLは、BraTS18（50ラベル付き、70未ラベル付き）でDiceスコア0.7111 ± 0.02を達成し、完全に教師ありのCNNベースライン（0.6670 ± 0.03）およびMSSLバージョン（0.6646 ± 0.03）を顕著に上回った。
• WMH17では、交互学習を用いたMASSLがDiceスコア0.7204 ± 0.02を達成し、完全に教師ありのCNNベースライン（0.6806 ± 0.03）およびMSSLベースライン（0.6880 ± 0.02）を上回った。
• アテンション機構は性能向上に顕著な寄与を示した：交互学習を用いたMASSLはBraTS18で0.7111 ± 0.02を達成したが、非アテンションのMSSLベースラインは0.6670 ± 0.03にとどまった。
• 交互学習は安定性と性能の両面で統合学習を上回り、損失重み係数γのハイパーパramータチューニングを一切行わずに、全データセットでより高いDiceスコアを達成した。
• 特徴分析の結果、アテンション機構により、エンコーダーのより深い層での識別性が向上した：BraTS18では、MASSLはレベル5（最も深い層）でR² = 0.535 ± 0.03を達成したのに対し、標準的なCNNでは0.486 ± 0.04であった。
• アブレーションスタディにより、アテンション機構が性能向上の主因であることが確認された。MSSL（アテンションなし）は、同じ学習戦略と損失重みを用いても、MASSLを下回った。