[論文レビュー] UNETR: Transformers for 3D Medical Image Segmentation

Ali Hatamizadeh, Tang, Yucheng+6 more|arXiv (Cornell University)|Mar 18, 2021

Radiomics and Machine Learning in Medical Imaging参考文献 52被引用数 210

ひとこと要約

UNETR は、3D 医用ボリュームをパッチのシーケンスとして処理するトランスフォーマーエンコーダを使用し、CNNベースのデコーダとスキップ接続で接続して正確な3Dセグメンテーションを実現し、BTCVとMSDデータセットで最先端の結果を達成します。

ABSTRACT

Fully Convolutional Neural Networks (FCNNs) with contracting and expanding paths have shown prominence for the majority of medical image segmentation applications since the past decade. In FCNNs, the encoder plays an integral role by learning both global and local features and contextual representations which can be utilized for semantic output prediction by the decoder. Despite their success, the locality of convolutional layers in FCNNs, limits the capability of learning long-range spatial dependencies. Inspired by the recent success of transformers for Natural Language Processing (NLP) in long-range sequence learning, we reformulate the task of volumetric (3D) medical image segmentation as a sequence-to-sequence prediction problem. We introduce a novel architecture, dubbed as UNEt TRansformers (UNETR), that utilizes a transformer as the encoder to learn sequence representations of the input volume and effectively capture the global multi-scale information, while also following the successful "U-shaped" network design for the encoder and decoder. The transformer encoder is directly connected to a decoder via skip connections at different resolutions to compute the final semantic segmentation output. We have validated the performance of our method on the Multi Atlas Labeling Beyond The Cranial Vault (BTCV) dataset for multi-organ segmentation and the Medical Segmentation Decathlon (MSD) dataset for brain tumor and spleen segmentation tasks. Our benchmarks demonstrate new state-of-the-art performance on the BTCV leaderboard. Code: https://monai.io/research/unetr

研究の動機と目的

• 医療画像セグメンテーションにおいて長距離の3D コンテキストを捉えるためにトランスフォーマーの活用を動機づける。
• トランスフォーマーエンコーダを直接CNNデコーダへスキップ接続で接続するUNETR アーキテクチャを提案する。
• BTCV の多臓器セグメンテーションと MSD の脳腫瘍および脾臓セグメンテーションデータセットでの有効性を示す。

提案手法

• 3D ボリュームを非重畳のパッチとして表現し、K次元埋め込みへ投影する。
• パッチ列を ViT-B16 スタイルのトランスフォーマーエンコーダ（L=12, K=768, パッチサイズ 16^3）で処理する。
• タスクがセマンティックセグメンテーションであるため、1D の位置埋め込みを追加しクラストークンは省略する。
• 中間のトランスフォーマー表現（z3, z6, z9, z12）を取り出し、空間テンソルへリシェイプし、スキップ接続によってCNNベースのデコーダと融合する。
• 3x3x3 の畳み込みを用いてトランスフォーマー特徴を複数の解像度でデコーダへ射影する。デコーダへアップサンプルするためのデコンボリューションを適用し、最終的には 1x1x1 の畳み込みとソフトマックスでボクセル単位の予測を行う。
• 連結されたソフト Dice とクロスエントロピー損失の組み合わせで訓練する。パッチベースのスライディングウィンドウ推論を0.5のオーバーラップで適用する。

実験結果

リサーチクエスチョン

• RQ1 3Dパッチで訓練されたトランスフォーマーエンコーダは、3D 医用画像のセグメンテーションのために長距離依存性を捉えられるのか。
• RQ2 トランスフォーマー由来の特徴を多解像度のスキップ接続を介してCNNベースのデコーダへ接続することは、CNN-のみまたはトランスフォーマーのみのベースラインよりセグメンテーション精度を向上させるのか。
• RQ3 デコーダ設計、パッチ解像度、およびモデルサイズが3D 医用画像のセグメンテーション性能に与える影響は何か。

主な発見

• UNETR は BTCV の Standard および Free Competitions の両方で最先端の成績を達成。
• MSD の脳腫瘍および脾臓セグメンテーションでは、UNETR は競合手法を上回り、特に胆嚢や副腎などの小さな構造での性能向上が見られた。
• BTCV では平均 Dice スコアがベースラインより明確に改善され、小さな臓器で顕著な向上を示す。
• MSD では脳腫瘍のサブ領域および脾臓セグメンテーション全体で、最も強力なベースラインより高い Dice スコアを示す。
• モデルは約 92.58M パラメータと 41.19G FLOPs を持ち、他のトランスフォーマーベース手法と比較して推論時間（平均 ~12.08s）も競争力がある。