[論文レビュー] BRAU-Net++: U-Shaped Hybrid CNN-Transformer Network for Medical Image Segmentation

Libin Lan, Pengzhou Cai+4 more|arXiv (Cornell University)|Jan 1, 2024

Radiomics and Machine Learning in Medical Imaging被引用数 24

ひとこと要約

BRAU-Net++ は、Bi-Level Routing Attention (BRA) をコアブロックとして用い、チャンネル-空間のスキップ接続 (SCCSA) を組み込んだ U 字型の CNN-Transformer アーキテクチャで、複数データセットに対して効率的で高精度な医用画像分割を実現します。

ABSTRACT

Accurate medical image segmentation is essential for clinical quantification, disease diagnosis, treatment planning and many other applications. Both convolution-based and transformer-based u-shaped architectures have made significant success in various medical image segmentation tasks. The former can efficiently learn local information of images while requiring much more image-specific inductive biases inherent to convolution operation. The latter can effectively capture long-range dependency at different feature scales using self-attention, whereas it typically encounters the challenges of quadratic compute and memory requirements with sequence length increasing. To address this problem, through integrating the merits of these two paradigms in a well-designed u-shaped architecture, we propose a hybrid yet effective CNN-Transformer network, named BRAU-Net++, for an accurate medical image segmentation task. Specifically, BRAU-Net++ uses bi-level routing attention as the core building block to design our u-shaped encoder-decoder structure, in which both encoder and decoder are hierarchically constructed, so as to learn global semantic information while reducing computational complexity. Furthermore, this network restructures skip connection by incorporating channel-spatial attention which adopts convolution operations, aiming to minimize local spatial information loss and amplify global dimension-interaction of multi-scale features. Extensive experiments on three public benchmark datasets demonstrate that our proposed approach surpasses other state-of-the-art methods including its baseline: BRAU-Net under almost all evaluation metrics. We achieve the average Dice-Similarity Coefficient (DSC) of 82.47, 90.10, and 92.94 on Synapse multi-organ segmentation, ISIC-2018 Challenge, and CVC-ClinicDB, as well as the mIoU of 84.01 and 88.17 on ISIC-2018 Challenge and CVC-ClinicDB, respectively.

研究の動機と目的

• CNN の局所モデリングと Transformer ベースの全体依存性を組み合わせて、正確な医用画像分割を動機づける。
• Bi-Level Routing Attention (BRA) によって長距離アテンションの計算量を削減する。
• チャネル-空間のアテンション駆動のスキップ接続 (SCCSA) により空間的なディテールを保持する。
• SOTA 手法と比較して多様な医用画像データセットで性能向上を実証する。

提案手法

• BiFormer ブロックをコアとするU字型のエンコーダ-デコーダを構築し、BRA を region-to-region および token-to-token アテンションに適用する。
• パッチ埋め込み/統合とパッチ展開を用いて対称的な階層型エンコーダ-デコーダを採用する。
• SCCSA を用いてスキップ接続を再設計し、チャネルと空間のアテンションを畳み込みで実装して次元間相互作用を高める。
• 位置情報を符号化するための深層畳み込みと、データセット固有の重み付けを持つハイブリッド損失（Dice + Cross-Entropy）。
• 前処理とパッチベースのトークン化を活用して、O((HW)^(4/3)) の計算量で効率的なクエリアウェアなスパースアテンションを実現する。
• Synapse、ISIC-2018、CVC-ClinicDB でデータ拡張と標準的な最適化スケジュールを用いて訓練する。

実験結果

リサーチクエスチョン

• RQ1 BRAU-Net++ は多様なモダリティに対して最先端の CNN、Transformer、ハイブリッド型の医用画像分割法を上回ることができるか。
• RQ2 Bi-Level Routing Attention は計算量とメモリを削減しつつ長距離依存性の効果的なモデリングを可能にするか。
• RQ3 SCCSA スキップ接続戦略はデコード時の空間ディテールの保持とスケール横断的な特徴統合を改善するか。

主な発見

• BRAU-Net++ はベンチマーク全体で高い分割性能を達成し、DSC および mIoU の平均指標を報告（例: ISIC-2018 で DSC 82.47%、mIoU 84.01%、CVC-ClinicDB でも類似）。
• TransUNet および Swin-Unet と比較して、Synapse で DSC を 4.49%、3.34% 向上させ、HD をそれぞれ 12.62 mm、2.48 mm 減少させた。
• Bi-Level Routing Attention は従来の全アテンションより低い計算コストで長距離モデリングを実現可能。
• SCCSA スキップ接続はダウンサンプリング中に失われた空間情報の回復と次元横断的特徴相互作用の強化に寄与。
• BRAU-Net++ は三つの公開データセット（Synapse、ISIC-2018、CVC-ClinicDB）で堅牢な一般化を示す。