[論文レビュー] AFTer-UNet: Axial Fusion Transformer UNet for Medical Image Segmentation

Xiangyi Yan, Hao Tang+4 more|arXiv (Cornell University)|Oct 20, 2021

Advanced Neural Network Applications被引用数 10

ひとこと要約

AFTer-UNet は、U-Net アーキテクチャに軸方向融合トランスフォーマーを統合することで、インtra-slice および inter-slice の長距離依存関係を共同でモデル化するエンドツーエンドの 3D 医療画像セグメンテーションフレームワークを提案する。軸方向および空間方向に別々に自己注意を計算し、それらを統合することで、計算コストと GPU メモリ使用量を削減しながら、4.34% の平均 Dice スコア向上（U-Net に対して）および 2.56% の向上（CoTr に対して）を達成し、最先端の性能を実現した。

ABSTRACT

Recent advances in transformer-based models have drawn attention to exploring these techniques in medical image segmentation, especially in conjunction with the U-Net model (or its variants), which has shown great success in medical image segmentation, under both 2D and 3D settings. Current 2D based methods either directly replace convolutional layers with pure transformers or consider a transformer as an additional intermediate encoder between the encoder and decoder of U-Net. However, these approaches only consider the attention encoding within one single slice and do not utilize the axial-axis information naturally provided by a 3D volume. In the 3D setting, convolution on volumetric data and transformers both consume large GPU memory. One has to either downsample the image or use cropped local patches to reduce GPU memory usage, which limits its performance. In this paper, we propose Axial Fusion Transformer UNet (AFTer-UNet), which takes both advantages of convolutional layers' capability of extracting detailed features and transformers' strength on long sequence modeling. It considers both intra-slice and inter-slice long-range cues to guide the segmentation. Meanwhile, it has fewer parameters and takes less GPU memory to train than the previous transformer-based models. Extensive experiments on three multi-organ segmentation datasets demonstrate that our method outperforms current state-of-the-art methods.

研究の動機と目的

• 3D バリュームにおける軸方向（inter-slice）コンテキスト情報を効果的に活用できない既存の 2D や 3D トランスフォーマーベースのモデルの限界を是正すること。
• ボリュメトリックデータにおける 3D 自己注意の高い GPU メモリ消費量を低減しつつ、長距離モデリング能力を維持すること。
• 特に細長いおよび解剖学的に変動の大きい臓器のセグメンテーション精度を、複数スライスにわたるコンテキスト手がかりを統合することで向上させること。
• 従来のトランスフォーマーベースのモデルよりもパrameter 数が少なく、メモリ使用量も低い効率的なモデルを維持すること。
• 臨床現場での導入に適した、コンsumer GPU と互換性があるスケーラブルでエンドツーエンドのフレームワークを開発すること。

提案手法

• 軸方向と空間方向に別々に自己注意を計算し、それらの出力を統合する軸方向融合機構を導入する。
• U-Net のエンコーダーとデコーダーの間に軸方向融合トランスフォーマーを埋め込み、U-Net の U 字型スキップ接続を保持する。
• 特徴抽出には 2D CNN のエンコーダーとデコーダーを用い、軸方向融合トランスフォーマーはスライス間の長距離依存関係を捉える。
• 軸方向および各 2D スライス内でのマルチヘッド自己注意を、相対的位置エンコーディングを用いて実装し、長距離コンテキストをモデル化する。
• スライス間の注意の密度を制御するためのサンプリング周波数 $N_f$ を導入し、近接する隣接スライスに優先的に注目することで、効率性と性能を向上させる。
• 軸方向と空間方向の注意を分離することで、計算複雑度を $O(HW D^2)$ から $O(HW D + H^2W^2)$ に低減し、1枚の 11GB GPU での学習を可能にした。

実験結果

リサーチクエスチョン

• RQ1 インtra-slice および inter-slice のアテンションを軸方向に統合することで、純粋な 2D や 3D トランスフォーマーベースのモデルを上回る 3D 医療画像セグメンテーションが達成可能か？
• RQ2 軸方向と空間方向の注意を分離することで、GPU メモリ使用量を低減しつつ、長距離モデリング能力を維持できるか？
• RQ3 軸方向の隣接スライス数が、特に細長い臓器のセグメンテーション性能にどのように影響するか？
• RQ4 軸方向融合トランスフォーマー層の数を増やすことで、セグメンテーション精度がどの程度向上するか？
• RQ5 軸方向に沿ったサンプリング周波数が、モデルの性能と効率にどのように影響するか？

主な発見

• Thorax-85 データセットでは、U-Net に対して 2.13% の平均 Dice スコア向上、CoTr に対して 0.69% の向上を達成した。
• 挑戦的な細長い臓器である食道では、U-Net に対して 2.91% の Dice スコア向上、CoTr に対して 1.45% の向上を達成した。
• マルチ臓器データセットでは、AFTer-UNet は肝臓で U-Net より 3.17%、左腎臓で 7.30%、右腎臓で 7.11%、膵臓で 6.64% の向上を達成した。
• CoTr と比較して、肝臓で 3.85%、左腎臓で 5.26%、右腎臓で 6.61%、膵臓で 5.68% の高い Dice スコアを達成した。
• AFTer-UNet は GPU メモリ使用量を顕著に削減し、1枚の RTX-2080Ti（11GB）で学習が可能となった。これに対して、従来の 3D トランスフォーマーベースのモデルは複数 GPU 必要であった。
• 41.5M パラメータで、TransUNet（43.5M）や CoTr（41.9M）と同等のモデルサイズを維持しており、パrameter の増加なしに効率性を実証した。