[論文レビュー] D-Former: A U-shaped Dilated Transformer for 3D Medical Image Segmentation

Yixuan Wu, Kuanlun Liao+5 more|arXiv (Cornell University)|Jan 3, 2022

AI in cancer detection被引用数 9

ひとこと要約

D-Formerは、トークン数を増加させずに、拡張されたマルチヘッド自己注意を用いてパッチ間の長距離依存関係を捉えることで、計算コストを低減するU字型の拡張型Transformerを3D医療画像セグメンテーションに提案する。動的位置符号化機構を導入することで空間モデリングを向上させ、Swin-Unet や nnFormer などの最先端モデルを上回る性能を達成し、FLOPs やパラメータ数が競合するCNNおよびTransformerモデルよりも少ない。

ABSTRACT

Computer-aided medical image segmentation has been applied widely in diagnosis and treatment to obtain clinically useful information of shapes and volumes of target organs and tissues. In the past several years, convolutional neural network (CNN) based methods (e.g., U-Net) have dominated this area, but still suffered from inadequate long-range information capturing. Hence, recent work presented computer vision Transformer variants for medical image segmentation tasks and obtained promising performances. Such Transformers model long-range dependency by computing pair-wise patch relations. However, they incur prohibitive computational costs, especially on 3D medical images (e.g., CT and MRI). In this paper, we propose a new method called Dilated Transformer, which conducts self-attention for pair-wise patch relations captured alternately in local and global scopes. Inspired by dilated convolution kernels, we conduct the global self-attention in a dilated manner, enlarging receptive fields without increasing the patches involved and thus reducing computational costs. Based on this design of Dilated Transformer, we construct a U-shaped encoder-decoder hierarchical architecture called D-Former for 3D medical image segmentation. Experiments on the Synapse and ACDC datasets show that our D-Former model, trained from scratch, outperforms various competitive CNN-based or Transformer-based segmentation models at a low computational cost without time-consuming per-training process.

研究の動機と目的

• 3D医療画像セグメンテーションにおけるTransformerのグローバル自己注意の高コスト問題に対処すること。
• 空間分解能を維持しながら、3D医療画像における長距離特徴モデリングを向上させ、パラメータ数を削減すること。
• 限られた計算リソース下でも効率的かつスケーラブルな3D医療画像向けTransformerアーキテクチャを設計すること。
• 入力に依存する位置関係に適応する動的位置符号化を統合し、より良い空間理解を実現すること。

提案手法

• 局所的およびグローバルな自己注意を交互に拡張する形で計算する、拡張自己注意機構を提案し、パッチ数を増加させずに受容 field を拡大する。
• 局所的スコープモジュール（LSM）とグローバルスコープモジュール（GSM）を交互に配置した、D-Formerと呼ばれるU字型エンコーダデコーダアーキテクチャを導入する。
• 最初の局所的自己注意モジュールの直前に配置された動的位置符号化（DPE）ブロックを採用し、柔軟に位置に依存する表現を学習する。
• 複数の段階を含む階層的設計を採用し、各段階でLSMとGSMモジュールを交互に配置することで、マルチスケールの文脈的特徴を段階的に捉える。
• 入力に依存する空間的関係に適応する学習済みの位置埋め込みを適用し、3D空間構造のモデリングを向上させる。
• 事前学習を用いず、スクラッチから訓練することで、低い計算コストで高い性能を達成する。

実験結果

リサーチクエスチョン

• RQ1 拡張自己注意機構は、3D医療画像セグメンテーションにおける計算コストを低減しつつ、長距離依存関係のモデリングを維持できるか？
• RQ2 標準TransformerおよびCNNと比較して、D-Formerアーキテクチャは3D医療画像認識タスクにおいて、セグメンテーション精度とFLOPsの両面でどのように差をつけるか？
• RQ3 動的位置符号化の最適な配置と設計は、3D空間表現の向上にどのように寄与するか？
• RQ4 グローバル自己注意を拡張されたグローバルスコープモジュールに置き換えることで、トークン数を増加させずに性能が向上するか？
• RQ5 Synapse や ACDC といった3D医療画像データセットをスクラッチから訓練した場合、D-Formerモデルはどの程度の効果を示すか？

主な発見

• D-Formerは、Synapseデータセットで平均DSC 88.83%を達成し、Swin-Unet や nnFormer などの最先端モデルを上回る。
• 標準Transformerと比較してFLOPsとパラメータ数を削減しながらも、優れたセグメンテーション精度を維持している。
• 提案された動的位置符号化（DPE）は、絶対位置符号化と比較して平均DSCを4.05%向上させ、相対位置符号化と比較しても2.48%向上させた。
• DPEブロックを最初の局所的自己注意モジュールの直前に配置した場合が最も高い性能を示し、空間モデリングにおけるその重要性を確認した。
• 拡張された自己注意を用いたグローバルスコープモジュール（GSM）は、局所的スコープモジュールのみを用いた場合に比べて精度を0.66%向上させ、長距離モデリングにおける有効性を示した。
• アブレーションスタディの結果、拡張自己注意機構とDPEの両方が、低い計算コストで最先端の性能を達成するために不可欠であることが確認された。