[論文レビュー] TransDeepLab: Convolution-Free Transformer-based DeepLab v3+ for Medical Image Segmentation

Reza Azad, Moein Heidari+5 more|arXiv (Cornell University)|Jan 1, 2022

Advanced Neural Network Applications被引用数 5

ひとこと要約

TransDeepLabは、階層的スウィン・トランスフォーマーとシフトドウインドウを用い、新規のマルチスケール特徴マッピングモジュール（SSPP）を統合することで、長距離依存関係をモデル化し境界局所化を向上させる、畳み込みを排除した純トランスフォーマー型のDeepLabv3+アーキテクチャを提案する。このモデルは、顕著に低いモデル複雑度（21.14Mパラメータ）で複数の医療画像ベンチマークでSOTA性能を達成し、CNNベースおよびハイブリッド型トランスフォーマー-CNNモデルを上回る。

ABSTRACT

Convolutional neural networks (CNNs) have been the de facto standard in a diverse set of computer vision tasks for many years. Especially, deep neural networks based on seminal architectures such as U-shaped models with skip-connections or atrous convolution with pyramid pooling have been tailored to a wide range of medical image analysis tasks. The main advantage of such architectures is that they are prone to detaining versatile local features. However, as a general consensus, CNNs fail to capture long-range dependencies and spatial correlations due to the intrinsic property of confined receptive field size of convolution operations. Alternatively, Transformer, profiting from global information modelling that stems from the self-attention mechanism, has recently attained remarkable performance in natural language processing and computer vision. Nevertheless, previous studies prove that both local and global features are critical for a deep model in dense prediction, such as segmenting complicated structures with disparate shapes and configurations. To this end, this paper proposes TransDeepLab, a novel DeepLab-like pure Transformer for medical image segmentation. Specifically, we exploit hierarchical Swin-Transformer with shifted windows to extend the DeepLabv3 and model the Atrous Spatial Pyramid Pooling (ASPP) module. A thorough search of the relevant literature yielded that we are the first to model the seminal DeepLab model with a pure Transformer-based model. Extensive experiments on various medical image segmentation tasks verify that our approach performs superior or on par with most contemporary works on an amalgamation of Vision Transformer and CNN-based methods, along with a significant reduction of model complexity. The codes and trained models are publicly available at https://github.com/rezazad68/transdeeplab

研究の動機と目的

• 医療画像セグメンテーションにおけるCNNの長距離空間的依存関係および文脈的相関の捉えにくさを是正すること。
• 密度予測タスクにおける標準的ビジョン・トランスフォーマーの高い計算コストと低レベル特徴の学習が不十分である問題を克服すること。
• 従来のDeepLabv3+アーキテクチャの純トランスフォーマー型代替案を構築し、モデル複雑度を低減しつつ高い性能を維持すること。
• トランスフォーマーのエンコーダ-デコーダフレームワーク内に、新規のSSPPモジュールを用いて階層的特徴表現とマルチスケール文脈モデリングを統合すること。

提案手法

• DeepLabv3+のCNNベースのエンコーダを、シフトドウインドウ自己注意を用いた階層的スウィン・トランスフォーマー・エンコーダに置き換え、局所的およびグローバルな特徴学習を可能にする。
• スウィン・トランスフォーマーの異なる段階間でのクロスアテンションを介したマルチスケール特徴マッピング機構を導入し、文脈表現を強化する。
• スウィン・トランスフォーマーの複数段階からの特徴を異なるレートで集約する、新規のアトロス空間的ピラミッドプーリング（ASPP）の変種（SSPP）を採用し、マルチスケール文脈を捉える。
• スキップ接続を用いた軽量デコーダヘッドを採用し、セグメンテーションマップの精細な詳細を保持する。
• スウィン・トランスフォーマーの階層的かつ局所接続構造を活用することで、二次的アテンション計算量を低減し、学習効率を最適化する。
• 多様な医療画像データセット上で、クロスエントロピー損失とDice損失を用いてエンドツーエンドで学習し、アブレーションスタディにより各モジュールの寄与度を検証する。

実験結果

リサーチクエスチョン

• RQ1 純トランスフォーマー型アーキテクチャは、モデル複雑度を低減しつつ、CNNベースのモデルを上回る性能を達成できるか？
• RQ2 階層的スウィン・トランスフォーマーと変更されたASPPモジュール（SSPP）の統合は、医療画像セグメンテーションにおけるマルチスケール文脈モデリングをどのように向上させるか？
• RQ3 トランスフォーマー段階間でのクロスアテンションベースの特徴マッピングは、単純な連結や要素演算と比較して、長距離依存関係学習をどのように改善するか？
• RQ4 DeepLabv3+のCNNエンコーダをトランスフォーマーに置き換えることで、複雑な解剖的構造におけるセグメンテーション精度はどの程度向上するか？
• RQ5 本手法は、ハイブリッド型トランスフォーマー-CNNモデルおよび確立されたCNNベースのベースラインと比較して、効率性と性能の両面で優れているか？

主な発見

• TransDeepLabは、Synapseデータセットで85.68の平均Diceスコアを達成し、ベースラインのCNNベースエンコーダ（75.89）を上回り、複数のベンチマークでSOTA手法と同等またはそれを上回る性能を示した。
• パラメータ数を21.14Mにまで低減し、Swin-UNet（27.17M）、DeepLabv3（54.70M）、Trans-UNet（105M）と比較して顕著なモデル効率性を示した。
• アブレーションスタディにより、クロスアテンションベースの特徴マッピング戦略が、基本的な連結マッピングと比較してDiceスコアを4.27%向上させることを確認した。
• 2段階のSSPPモジュールを採用した場合に最良のDiceスコア（80.16）が得られ、マルチスケール文脈の捉え方と計算コストの最適なバランスが実現された。
• 可視化結果から、強化された長距離依存関係モデリングのおかげで、特に複雑な解剖的領域においてより正確な境界予測が得られていることが示された。
• 2段階SSPPを用いた場合、SynapseデータセットでHausdorff距離が21.25にまで低下し、他の設定と比較して局所化精度が向上した。