[論文レビュー] U-Net Transformer: Self and Cross Attention for Medical Image Segmentation

Olivier Petit, Nicolas Thome+2 more|arXiv (Cornell University)|Mar 10, 2021

Radiomics and Machine Learning in Medical Imaging被引用数 18

ひとこと要約

本論文では、多ヘッド自己注意およびクロス注意メカニズムを統合したU-Netベースのアーキテクチャ、U-Transformerを提案する。このアーキテクチャは、低コントラストおよび複雑な臓器の医療画像セグメンテーションを改善することを目的としている。長距離の文脈的依存関係をモデル化し、注意誘導型スキップ接続によって空間的詳細を精緻化することで、U-Transformerは腹部CTデータセットにおいて最先端の性能を達成し、特に膵臓のような困難な臓器において、U-Net や Attention U-Net よりも最大3.4%高いDiceスコアを達成した。

ABSTRACT

Medical image segmentation remains particularly challenging for complex and low-contrast anatomical structures. In this paper, we introduce the U-Transformer network, which combines a U-shaped architecture for image segmentation with self- and cross-attention from Transformers. U-Transformer overcomes the inability of U-Nets to model long-range contextual interactions and spatial dependencies, which are arguably crucial for accurate segmentation in challenging contexts. To this end, attention mechanisms are incorporated at two main levels: a self-attention module leverages global interactions between encoder features, while cross-attention in the skip connections allows a fine spatial recovery in the U-Net decoder by filtering out non-semantic features. Experiments on two abdominal CT-image datasets show the large performance gain brought out by U-Transformer compared to U-Net and local Attention U-Nets. We also highlight the importance of using both self- and cross-attention, and the nice interpretability features brought out by U-Transformer.

研究の動機と目的

• U-Netが低コントラストで複雑な解剖的構造における長距離の文脈的依存関係や空間的関係を捉える能力に限界があることに対処すること。
• 膵臓、胆嚢、胃のような小形で視覚的に曖昧な臓器のセグメンテーション精度を向上させること。
• 自己注意およびクロス注意メカニズムをU-Netフレームワークに統合し、グローバルな文脈をモデル化するとともに、空間的詳細を精緻化すること。
• 注意マップの可視化とアブレーションスタディーを通じて、モデルの解釈性を向上させること。
• 自己注意とクロス注意の両方のメカニズムが統合された一元化されたトランスフォーマーベースのU-Netアーキテクチャにおいて、相乗効果を実証すること。

提案手法

• エンコーダーの終端に多ヘッド自己注意（MHSA）を統合したU-Transformerネットワークを導入し、画像全体にわたるグローバルな文脈的相互作用をモデル化する。
• スキップ接続に多ヘッドクロス注意（MHCA）を採用し、高レベルの意味的特徴と高分解能特徴を融合することで、細粒度の空間的回復を可能にする。
• MHSAおよびMHCAモジュールの両方で空間的位置情報を保持するため、学習可能な位置エンコーディングを採用し、注意の局所化を向上させる。
• 異なるデコーダーステージにMHCAモジュールをスタックすることで、マルチレベルのクロス注意を実現し、予測を段階的に精緻化する。
• MHSAが特徴間の長距離依存関係を捉え、MHCAが特徴のアップサンプリング過程で非意味的ノイズをフィルタリングする二重ブランチ型の注意メカニズムを採用する。
• 残差ブロックとスキップ接続を備えた標準的なU-Netバックボーンを採用し、畳み込み層を置き換えるのではなく、注意モジュールを統合することで強化する。

実験結果

リサーチクエスチョン

• RQ1 エンコーダーにおける自己注意は、U-Netの限られた受容 field を超えて、複雑な臓器のセグメンテーションを、長距離の文脈的依存関係をモデル化することで向上させることができるか？
• RQ2 スキップ接続におけるクロス注意は、非意味的特徴をフィルタリングし、位置情報統合を可能にすることで、空間的詳細回復を向上させることができるか？
• RQ3 自己注意とクロス注意のメカニズムを個別に評価し、併用した場合のセグメンテーション性能および解釈性の観点から、両者の相違と相乗効果をどのように比較できるか？
• RQ4 位置エンコーディングとマルチレベルの注意が、デコーダパスにおけるクロス注意の性能に及ぼす影響は何か？
• RQ5 両方の注意メカニズムを統合することで、特に低コントラストで複雑な形状を示す臓器において、より良い一般化性能が得られるか？

主な発見

• IMOデータセットにおける膵臓のDice類似係数（DSC）は73.10%を達成し、U-Net（+3.4%）およびAttention U-Net（+4.45%）を上回った。
• TCIA膵臓データセットでは、U-TransformerがDSC 80.65%を達成し、U-Net（76.35%）およびAttention U-Net（77.23%）を大きく上回った。
• U-Transformerにおける自己注意とクロス注意の併用は、単独での使用よりも優れた結果をもたらし、両者の相乗効果が明確に示された。
• 位置エンコーディングは、TCIAで+1.7ポイント、IMOで+0.6ポイントのMHCA性能向上をもたらし、空間的局所化において重要であることが示された。
• マルチレベルのクロス注意は、TCIAで+1.8ポイント、IMOで+0.6ポイントの性能向上をもたらし、階層的な注意の精緻化がセグメンテーション精度を向上させることを示した。
• Hausdorff距離の結果もDiceスコアと同様の傾向を示し、U-Transformerが予測のアーチファクトを低減し、境界の正確性を向上させていることが示された。