[論文レビュー] Enhancing Medical Image Segmentation with TransCeption: A Multi-Scale Feature Fusion Approach

Reza Azad, Yiwei Jia+3 more|arXiv (Cornell University)|Jan 25, 2023

AI in cancer detection被引用数 41

ひとこと要約

TransCeption は ResInception Patch Merging と Multi-Branch Transformer を用い、Intra-stage Feature Fusion と Dual Transformer Bridge を備えた純粋なトランスフォーマー U-Net 派生で、多スケール特徴融合を実現し、医用画像分割における最先端性能を達成します。

ABSTRACT

While CNN-based methods have been the cornerstone of medical image segmentation due to their promising performance and robustness, they suffer from limitations in capturing long-range dependencies. Transformer-based approaches are currently prevailing since they enlarge the reception field to model global contextual correlation. To further extract rich representations, some extensions of the U-Net employ multi-scale feature extraction and fusion modules and obtain improved performance. Inspired by this idea, we propose TransCeption for medical image segmentation, a pure transformer-based U-shape network featured by incorporating the inception-like module into the encoder and adopting a contextual bridge for better feature fusion. The design proposed in this work is based on three core principles: (1) The patch merging module in the encoder is redesigned with ResInception Patch Merging (RIPM). Multi-branch transformer (MB transformer) adopts the same number of branches as the outputs of RIPM. Combining the two modules enables the model to capture a multi-scale representation within a single stage. (2) We construct an Intra-stage Feature Fusion (IFF) module following the MB transformer to enhance the aggregation of feature maps from all the branches and particularly focus on the interaction between the different channels of all the scales. (3) In contrast to a bridge that only contains token-wise self-attention, we propose a Dual Transformer Bridge that also includes channel-wise self-attention to exploit correlations between scales at different stages from a dual perspective. Extensive experiments on multi-organ and skin lesion segmentation tasks present the superior performance of TransCeption compared to previous work. The code is publicly available at \url{https://github.com/mindflow-institue/TransCeption}.

研究の動機と目的

• 標準的な CNN および単一スケールのトランスフォーマーを超えた医用画像分割におけるグローバルコンテキストモデリングの改善を動機づける。
• エンコーダー段内および段間で多スケール特徴を融合する純粋なトランスフォーマー U-Net 派生（TransCeption）を提案する。
• RIPM、MB Transformer、IFF、Dual Transformer Bridge のアーキテクチャブロックを導入し、スケール間の特徴融合を強化する。

提案手法

• ResInception Patch Merging (RIPM) によるエンコーダのパッチマージを再設計し、1 ステージ内で多スケール表現を捉える。
• RIPM から得られる 3 本の並列特徴マップ（3x3、5x5、7x7）と追加のローカルディテール用 3x3 ブランチを処理する Multi-Branch (MB) Transformer ブロックを用いる。
• Intra-stage Feature Fusion (IFF) を導入し、チャネル指向の位置を preserve する注意機構でマルチブランチ出力を融合する。
• エンコーダ-デコーダブリッジを通じて、トークン認識型とチャネル認識型の注意を組み合わせた Dual Transformer Bridge を用い、マルチステージ・マルチスケール特徴を融合する。
• CoaT スタイルモジュールに基づく純粋なトランスフォーマー枠組みで、4 ステージのエンコーダとオーバラップした Patch Embedding (OPE) および Patch Expanding をデコーダと共に採用する。
• Bridge におけるスケール縮小を伴うトークン認識型トランスフォーマと、スケール間通信のためのチャネル認識型トランスフォーマを採用し、計算量を低減する。

Figure 1 : The overview of the proposed TransCeption. The network is composed of the 4-stage encoder, the decoder, and a Dual Transformer Bridge. In the encoder, the first stage utilizes an overlapped Patch Embedding module and an efficient transformer block to capture the low-level semantic feature

実験結果

リサーチクエスチョン

• RQ1 エンコーダー段内および段間での多スケール特徴融合は、既存のトランスフォーマーに基づくアプローチを超えて医用画像のセグメンテーション性能を向上させるか。
• RQ2 RIPM、MB Transformer、IFF、Dual Transformer Bridge のようなアーキテクチャを、スケール間およびチャネル間の依存性を効率的にモデルするように設計するにはどうすればよいか。
• RQ3 これらの多スケールブロックを備えた純粋なトランスフォーマー U-Net が、マルチオーガンCT分割および皮膚病変分割で最先端の結果を達成するか。
• RQ4 内部ステージ間および外部ステージ間の融合が境界の精度とノイズ耐性に与える影響はどの程度か。

主な発見

手法	DSC ↑	大動脈	胆嚢	左腎	右腎	肝臓	膵臓	脾臓	胃	HD ↓
V-Net	68.81	75.34	51.87	77.10	80.75	87.84	40.05	80.56	56.98	-
DARR	69.77	74.74	53.77	72.31	73.24	94.08	54.18	89.90	45.96	-
R50 U-Net	74.68	87.47	66.36	80.60	78.19	93.74	56.90	85.87	74.16	36.87
U-Net	76.85	89.07	69.72	77.77	68.60	93.43	53.98	86.67	75.58	39.70
R50 Att-UNet	75.57	55.92	63.91	79.20	72.71	93.56	49.37	87.19	74.95	36.97
Att-UNet	77.77	89.55	68.88	77.98	71.11	93.57	58.04	87.30	75.75	36.02
R50 ViT	71.29	73.73	55.13	75.80	72.20	91.51	45.99	81.99	73.95	32.87
TransUNet	77.48	87.23	63.13	81.87	77.02	94.08	55.86	85.08	75.62	31.69
TransNorm	78.40	86.23	65.10	82.18	78.63	94.22	55.34	89.50	76.01	30.25
Swin-Unet	79.13	85.47	66.53	83.28	79.61	94.29	56.58	90.66	76.60	21.55
TransDeepLab	80.16	86.04	69.16	84.08	79.88	93.53	61.19	89.00	78.40	21.25
HiFormer	80.39	86.21	65.69	85.23	79.77	94.61	59.52	90.99	81.08	14.70
MISSFormer	81.96	86.99	68.65	85.21	82.00	94.41	65.67	91.92	80.81	18.20
TransCeption	82.24	87.60	71.82	86.23	80.29	95.01	65.27	91.68	80.02	20.89

• TransCeption は Synapse の多臓器分割で DSC 82.24% および HD 20.89% を達成し、純粋なトランスフォーマーを含む従来手法を上回る。
• ISIC 2018 の皮膚病変分割では DSC 0.9124、ACC 0.9628、SE 0.9192、SP 0.9744 を達成し、いくつかの CNN およびトランスフォーマーベースのベースラインを上回る。
• 多スケール RIPM および MB Transformer の設計と IFF により、単一ステージ内のスケール間およびチャネル間特徴統合が改善される。
• Dual Transformer Bridge はトークン認識型とチャネル認識型の注意を組み合わせてマルチステージ・マルチスケール特徴を効果的に融合し、エンコーダー段のグローバルコンテキストモデリングを改善する。
• TransCeption は事前学習なしでゼロショットで学習され、評価データセット上でいくつかの事前学習済みトランスフォーマーベースのベースラインを上回る。

Figure 2 : Illustration of Dual Transformer Bridge. The four input feature maps $Y_{1}$ to $Y_{4}$ from stages in the encoder are flattened to the token sequences of dimension $N_{i}\times C$ . The token sequences are then concatenated into a long sequence $Y_{C}$ containing multi-stage information.