[論文レビュー] DAE-Former: Dual Attention-guided Efficient Transformer for Medical Image Segmentation

Reza Azad, René Arimond+3 more|arXiv (Cornell University)|Dec 27, 2022

Advanced Neural Network Applications被引用数 5

ひとこと要約

DAE-Formerは、空間的およびチャネル次元の両方における長距離依存関係をモデル化するための二重アテンション機構を導入することで、計算効率が良く、畳み込みを含まないU-Netに類似したTransformerアーキテクチャを提案する。この手法は、事前学習を必要とせず、心臓および皮膚病変のセグメンテーションベンチマークで最先端の性能を達成する。特徴の融合と局在化を向上させるために、新規のスカイプ接続クロスアテンション（SCCA）モジュールを採用している。

ABSTRACT

Transformers have recently gained attention in the computer vision domain due to their ability to model long-range dependencies. However, the self-attention mechanism, which is the core part of the Transformer model, usually suffers from quadratic computational complexity with respect to the number of tokens. Many architectures attempt to reduce model complexity by limiting the self-attention mechanism to local regions or by redesigning the tokenization process. In this paper, we propose DAE-Former, a novel method that seeks to provide an alternative perspective by efficiently designing the self-attention mechanism. More specifically, we reformulate the self-attention mechanism to capture both spatial and channel relations across the whole feature dimension while staying computationally efficient. Furthermore, we redesign the skip connection path by including the cross-attention module to ensure the feature reusability and enhance the localization power. Our method outperforms state-of-the-art methods on multi-organ cardiac and skin lesion segmentation datasets without requiring pre-training weights. The code is publicly available at https://github.com/mindflow-institue/DAEFormer.

研究の動機と目的

• 高解像度の医療画像におけるTransformerの標準自己アテンションの2次的計算複雑性を軽減すること。
• 医療画像セグメンテーションにおけるCNNの長距離空間的・構造的依存関係を捉える能力の制限を克服すること。
• 畳み込みニューラルネットワーク（CNN）バックボーンに依存せずに、高い効率性と強力な局在化能力を維持する純粋なTransformerベースのU-Netアーキテクチャを設計すること。
• エンコーダーとデコーダー間の特徴の融合を向上させるために、スカイプ接続クロスアテンション（SCCA）モジュールを導入し、表現の再利用性とセグメンテーション精度を向上させること。
• 事前学習重みを必要とせず、多臓器および皮膚病変セグメンテーションタスクで最先端の性能を達成すること。

提案手法

• 入力画像を重複する4×4のパッチでトークン化するためのパッチエンベッディングを採用し、畳み込み演算を排除する。
• 空間次元をダウンサンプリングおよびアップサンプリングしながらチャネル数を倍増させるパッチマージングおよび拡張レイヤーを用いた、階層的なエンコーダー・デコーダー構造を実装し、3段階のスタックされたブロックを構成する。
• 二重アテンションブロックを導入し、2段階の処理を実施する：(1) 効率的アテンションにより、正規化されたクエリおよびキーを用いて空間的重要性を計算し、O(N²)の複雑性をO(d²N)に低減する。 (2) トランスポーズアテンションにより、キーと値の行列を転置することでチャネルごとの依存関係をモデル化する。
• ソフトマックス正規化されたクエリおよびキーを用いた効率的アテンション機構を適用し、値との行列乗算を実行してグローバルコンテキストベクトルを生成する。これにより、計算量を削減しつつ高い表現力が得られる。
• エンコーダーとデコーダーの特徴を、空間的およびチャネル表現の両方に注目することで統合するスカイプ接続クロスアテンション（SCCA）モジュールを設計し、特徴の再利用性と局在化精度を向上させる。
• 最終段階で線形プロジェクション層を用いて、精練された特徴からセグメンテーションマップを生成する。

実験結果

リサーチクエスチョン

• RQ1 純粋なTransformerベースのU-Netアーキテクチャは、計算効率を維持したまま、医療画像セグメンテーションで最先端の性能を達成できるか？
• RQ2 自己アテンション機構をどのように再設計すれば、2次的複雑性を負担せずに空間的およびチャネルワイドな依存関係を効率的にモデル化できるか？
• RQ3 スカイプ接続クロスアテンション（SCCA）モジュールを導入することで、Transformerベースのセグメンテーションネットワークにおける特徴の統合と局在化はどの程度向上するか？
• RQ4 二重アテンション機構（効率的空間アテンション＋トランスポーズチャネルアテンション）に依存するモデルが、CNN-Transformerハイブリッドや既存の純粋なTransformerを凌駕できるか、事前学習を必要とせずに？
• RQ5 提案されたアーキテクチャは、心臓および皮膚病変セグメンテーションなど、多様な医療画像タスクに一般化できるか？

主な発見

• DAE-Formerは、事前学習重みを必要とせず、多臓器心臓セグメンテーションおよび皮膚病変セグメンテーションデータセットで最先端の性能を達成した。
• モデルは強力な一般化能力を示し、心臓および皮膚病変セグメンテーションベンチマークの両方で、既存の最先端手法を上回った。
• 二重アテンション機構（効率的空間アテンションに続くトランスポーズチャネルアテンション）により、計算コストを低減しつつ、空間的およびチャネル次元の両方における長距離依存関係を効果的にモデル化できた。
• スカイプ接続クロスアテンション（SCCA）モジュールは、エンコーダーとデコーダーの特徴間のクロスモダリティアテンションを可能にすることで、特徴の再利用性を著しく向上させ、局在化精度を改善した。
• 畳み込みバックボーンを含まない純粋なTransformer設計に加え、効率的アテンションを採用することで、高い効率性を維持した。このため、高解像度の医療画像に適している。
• アブレーションスタディの結果、二重アテンション機構およびSCCAモジュールの両方が、性能向上に顕著な貢献をしていることが確認され、設計選択の妥当性が裏付けられた。