[論文レビュー] UTNet: A Hybrid Transformer Architecture for Medical Image Segmentation

Yunhe Gao, Mu Zhou+1 more|arXiv (Cornell University)|Jul 2, 2021

Radiomics and Machine Learning in Medical Imaging参考文献 26被引用数 49

ひとこと要約

UTNet は自己注意を CNN ベースの U-Net に組み込み、効率的な注意機構と相対位置エンコーディングでグローバルな文脈を複数スケールで捉え、事前学習なしで心臓 MRI のセグメンテーションを上回り、ベンダー横断の頑健性を実現します。

ABSTRACT

Transformer architecture has emerged to be successful in a number of natural language processing tasks. However, its applications to medical vision remain largely unexplored. In this study, we present UTNet, a simple yet powerful hybrid Transformer architecture that integrates self-attention into a convolutional neural network for enhancing medical image segmentation. UTNet applies self-attention modules in both encoder and decoder for capturing long-range dependency at different scales with minimal overhead. To this end, we propose an efficient self-attention mechanism along with relative position encoding that reduces the complexity of self-attention operation significantly from $O(n^2)$ to approximate $O(n)$. A new self-attention decoder is also proposed to recover fine-grained details from the skipped connections in the encoder. Our approach addresses the dilemma that Transformer requires huge amounts of data to learn vision inductive bias. Our hybrid layer design allows the initialization of Transformer into convolutional networks without a need of pre-training. We have evaluated UTNet on the multi-label, multi-vendor cardiac magnetic resonance imaging cohort. UTNet demonstrates superior segmentation performance and robustness against the state-of-the-art approaches, holding the promise to generalize well on other medical image segmentations.

研究の動機と目的

• 従来のCNNを超える長距離の文脈が医用画像セグメンテーションに必要であることを動機づける。
• エンコーダ/デコーダの複数レベルに効率的な自己注意を注入する U字型ハイブリッド Transformer ネットワーク（UTNet）を提案する。
• 畳み込みの帰納バイアスを通じて事前学習なしで Transformer の統合を可能にする。
• 高解像度の医用画像で境界中心の正確なセグメンテーションを実現しつつ、計算効率を維持する。

提案手法

• KeyとValueを低次元空間へ射影することにより計算量を O(n^2) からおおよそ O(n) に削減する効率的な自己注意機構を導入する。
• U-Net のようなアーキテクチャで複数レベル（エンコーダとデコーダ）に自己注意を適用し、多スケールのグローバルコンテキストを捉える。
• 医用画像の内容位置関係をモデル化するために2D相対位置エンコーディングを組み込む。
• UTNet 内のビルディングブロックとして事前活性化残差ブロックと Transformer ブロックを用い、スキップ接続にはアイデンティティマッピングを採用する。
• Dice と cross-entropy 損失の組み合わせを用いて、事前学習なしにスクラッチから訓練する。
• マルチラベル・マルチベンダ心臓MRIデータで、UTNet を UNet、ResUNet、CBAM、Dual-Attention ネットワークと比較する。

実験結果

リサーチクエスチョン

• RQ1 大規模な事前学習なしで、高解像度の医用画像における境界重視のセグメンテーションを向上させるハイブリッド CNN-Transformer アーキテクチャは実現できるか。
• RQ2 相対位置エンコーディングを備えた多レベル自己注意は、異なるベンダー間でのセグメンテーションの頑健性を向上させるか。
• RQ3 効率的な自己注意とネットワーク内の配置が、セグメンテーション性能と計算効率に与える影響は何か。
• RQ4 マルチベンダ心臓MRIデータセットにおいて、UTNet は最先端のCNNベースのセグメンテーションモデルと比較してどのように性能を示すか。

主な発見

• UTNet は vendor A データで LV, MYO, RV の Dice スコアで最高を達成（LV 93.1, MYO 83.5, RV 88.2; Average Dice 88.3）。
• UTNet は、いくつかの注意機構ベースのベースラインと比較して、パラメータ数が競合的/劣ることがあり、推論時間も同様（Params 9.53M; Inference Time 0.145 s）。
• アブレーション研究により、自己注意をエンコーダ/デコーダの高いレベルに配置し、8次元の縮小サイズ射影を使用すると最良の性能を示し、相対的位置エンコーディングは不可欠である。
• UTNet はベンダー横断評価で優れた頑健性を示し、未知のベンダー C および D で他モデルが劣化する中、競争力のあるセグメンテーション性能を維持する。
• Dual-Attention（quadratic complexity）と比較して、UTNet はメモリを低く抑え、推論時間が速く、より高いセグメンテーション精度を達成している。