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

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

Medical Image Segmentation Techniques被引用数 9

ひとこと要約

本稿では、3D MRIボリュームをスライスせず、階層的自己注意をエンコーダーで、並列の自己およびクロス注意をデコーダーで使用する、完全なボリュームトランスフォーマー・アーキテクチャであるVT-UNetを提案する。Fourier位置エンコーディングを用いることで、BraTSデータセットで最先端の性能を達成し、MRIアーティファクトに対して高い耐性を示しながら、計算量も低く抑えられる。

ABSTRACT

We propose a Transformer architecture for volumetric segmentation, a challenging task that requires keeping a complex balance in encoding local and global spatial cues, and preserving information along all axes of the volume. Encoder of the proposed design benefits from self-attention mechanism to simultaneously encode local and global cues, while the decoder employs a parallel self and cross attention formulation to capture fine details for boundary refinement. Empirically, we show that the proposed design choices result in a computationally efficient model, with competitive and promising results on the Medical Segmentation Decathlon (MSD) brain tumor segmentation (BraTS) Task. We further show that the representations learned by our model are robust against data corruptions. \href{https://github.com/himashi92/VT-UNet}{Our code implementation is publicly available}.

研究の動機と目的

• 2Dスライスベースの処理を避けることで、3D医療画像セグメンテーションにおける完全なボリューム空間的コンテキストの保持に挑戦すること。
• 計算効率を維持しつつ、3D全軸にわたる長距離依存関係を捉えることができる、純粋なトランスフォーマーに基づくUNetアーキテクチャの設計。
• デコーダーにおける並列注意メカニズムを活用して特徴の精錬を向上させることで、腫瘍セグメンテーションの境界明確化を実現すること。
• モーション、ゴースト、スパイクアーティファクトなどの一般的なMRIアーティファクトに対するモデルの耐性を高めること。
• 3Dにおける自己注意メカニズムが、畳み込みニューラルネットワーク（CNN）ベースやハイブリッドモデルを上回る精度と耐性を示すかどうかを実証すること。

提案手法

• エンコーダーでは、階層的かつシフトされたウィンドウ形式の2つの連続する自己注意層を用い、3Dボリュームにおける局所的およびグローバルなコンテキスト手がかりを同時に捉える。
• デコーダーでは、共有クエリプロジェクションを用いた並列の自己およびクロス注意メカニズムを採用し、アップサンプリング中にグローバルコンテキストを保持する。
• 空間の誘導的バイアスを強化し、特徴表現を改善するために、デコーディングプロセスにFourier位置エンコーディングを組み込む。
• 並列な注意ヘッドからの特徴を凸結合戦略で融合することで、表現品質を向上させる。
• 学習率1e-4でAdamWを用い、データオーグメンテーション（回転、ノイズ、ぼかし、ガンマ補正）を実施し、ImageNet-22Kで事前学習されたSwin-T重みを初期化に使用する。
• 本アーキテクチャは、2Dスライス分解を一切行わず、MSD BraTSデータセットの完全な3D MRIボリューム上で評価される。

実験結果

リサーチクエスチョン

• RQ1 純粋なボリュームトランスフォーマー・アーキテクチャは、計算コストを低く抑えつつ、CNNベースやハイブリッドモデルを上回る性能を発揮できるか？
• RQ2 デコーダーにおける並列の自己およびクロス注意は、3Dセグメンテーションにおける境界精度と特徴精錬をどのように向上させるか？
• RQ3 Fourier位置エンコーディングは、3Dトランスフォーマーに基づくセグメンテーションネットワークにおける空間モデリングをどの程度向上させるか？
• RQ4 提案アーキテクチャは、モーション、ゴースト、スパイクアーティファクトなどの一般的なMRIアーティファクトに対して、改善された耐性を示すか？
• RQ5 完全に注意ベースの3D UNet設計は、畳み込みの誘導的バイアスに依存せずに、最先端の性能を達成できるか？

主な発見

• VT-UNet-Bは、全腫瘍（WT）クラスで91.9のDice Similarity Coefficient（DSC）を達成し、BraTSデータセットにおけるすべてのSOTA手法を上回った。
• WTクラスにおけるハウスドルフ距離（HD95）は3.43にまで低下し、比較対象すべての手法の中で最低であり、優れた境界精度を示した。
• VT-UNetはMRIアーティファクトに対して優れた耐性を示した：スパイクアーティファクトではDSCが86.6、モーションアーティファクトでは85.8を記録し、nnFormerを含むすべての腐敗タイプで優位であった。
• アブレーションスタディの結果、Fourier位置エンコーディング、凸結合、並列注意メカニズムの組み合わせが性能向上に顕著に寄与しており、これらの要素をすべて除去するとDSCが1.78ポイント低下した。
• VT-UNet-SとVT-UNet-Bは、nnFormer（4.05 GFLOPs）を含む大多数のSOTAモデルよりも低いFLOPs（3.84および3.43 GFLOPs）を達成しながら、より高い精度を維持した。
• 定性的な結果から、VT-UNetは鋭い腫瘍境界と、周腫瘍浮腫や強化腫瘍のような複雑な構造のより正確なセグメンテーションを実現していることが示された。