[論文レビュー] A Volumetric Transformer for Accurate 3D Tumor Segmentation

Himashi Peiris, Munawar Hayat+3 more|arXiv (Cornell University)|Nov 26, 2021

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

ひとこと要約

本稿では、3次元ボリュームデータの局所的およびグローバルな空間的特徴を同時に捉えるためにエンコーダーに二重自己注意機構を採用し、デコーダーでは並列なシフトされたウィンドウベースの自己注意およびクロス注意機構にフーリエ位置エンコーディングを組み合わせた、計算効率に優れた3次元ボリュームトランスフォーマーであるVT-UNetを提案する。このモデルは、BraTS 2021、Pancreas、Liver データセットにおいて最先端の性能を達成し、データの劣化に対しても強力な汎化性と耐性を示した。

ABSTRACT

This paper presents a Transformer architecture for volumetric medical image segmentation. Designing a computationally efficient Transformer architecture for volumetric segmentation is a challenging task. It requires keeping a complex balance in encoding local and global spatial cues, and preserving information along all axes of the volumetric data. The proposed volumetric Transformer has a U-shaped encoder-decoder design that processes the input voxels in their entirety. Our encoder has two consecutive self-attention layers to simultaneously encode local and global cues, and our decoder has novel parallel shifted window based self and cross attention blocks to capture fine details for boundary refinement by subsuming Fourier position encoding. Our proposed design choices result in a computationally efficient architecture, which demonstrates promising results on Brain Tumor Segmentation (BraTS) 2021, and Medical Segmentation Decathlon (Pancreas and Liver) datasets for tumor segmentation. We further show that the representations learned by our model transfer better across-datasets and are robust against data corruptions. \href{https://github.com/himashi92/VT-UNet}{Our code implementation is publicly available}.

研究の動機と目的

• 3次元ボリューム医療画像セグメンテーションのための計算的に効率的なトランスフォーマーを設計する課題に対処すること。
• ボリュームデータの3方向すべてにおいて、局所的およびグローバルな空間的依存関係をバランスよく表現すること。
• 新しい注意メカニズムを用いて腫瘍セグメンテーションにおける境界の精緻化を向上させること。
• 異なる医療画像データセット間での表現の転送可能性と、データの劣化に対する耐性を強化すること。

提案手法

• モデルは、パッチ分割を行わず、完全な3次元ボリューム入力を処理するU字型のエンコーダー・デコーダー構造を採用する。
• エンコーダーでは、連続する2つの自己注意層を用いて、局所的およびグローバルな空間的コンテキストを同時に符号化する。
• デコーダーでは、細粒度な詳細学習を強化するため、並列なシフトされたウィンドウベースの自己注意およびクロス注意ブロックを導入する。
• 計算コストの増加を伴わずに位置情報モデリングを改善するため、フーリエ位置エンコーディングを統合する。
• 空間的情報を3軸すべてにわたって保持しつつ、効率性を維持するようにアーキテクチャを設計する。
• 標準のセグメンテーション損失関数を用いて、BraTS 2021、Pancreas、Liver データセット上でエンドツーエンドにモデルを学習する。

実験結果

リサーチクエスチョン

• RQ1 トランスフォーマーに基づくアーキテクチャは、計算効率を保ちながらも、3次元腫瘍セグメンテーションで高い精度を達成できるか？
• RQ2 このモデルは、3次元医療画像において局所的およびグローバルな空間的依存関係をどの程度効果的に捉えられるか？
• RQ3 提案された注意メカニズムは、腫瘍セグメンテーションにおける境界の明確化をどの程度向上できるか？
• RQ4 学習された特徴表現は、異なる医療画像データセット間でどの程度汎化可能であり、データの劣化に対しても耐性を示すか？
• RQ5 フーリエ位置エンコーディングの統合は、計算コストの増加を伴わずに性能を向上させるか？

主な発見

• 提案されたVT-UNetは、脳腫瘍セグメンテーションのBraTS 2021データセットにおいて最先端の性能を達成した。
• モデルは強力な汎化能力を示し、膵臓および肝臓腫瘍セグメンテーションのMedical Segmentation Decathlonデータセットでも性能が向上した。
• ベースラインモデルと比較して、学習済み表現の異なるデータセット間への転送性が優れていた。
• さまざまなデータ劣化に対して耐性を示し、信頼性の高い特徴学習が可能であることが示された。
• 並列なシフトされたウィンドウ注意とフーリエ位置エンコーディングの組み合わせにより、計算コストの増加を抑えながら境界の精緻化が向上した。
• エンドツーエンドで完全な3次元ボリュームを処理しながらも、計算効率を維持した。