[Paper Review] 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 Imaging24 citations

TL;DR

BRAU-Net++ is a U-shaped CNN-Transformer architecture that uses Bi-Level Routing Attention (BRA) as the core block and a channel-spatial skip connection (SCCSA) to achieve efficient and accurate medical image segmentation across multiple datasets.

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.

Motivation & Objective

• Motivate accurate medical image segmentation by combining CNN local modeling with Transformer-based global dependencies.
• Reduce computational complexity of long-range attention via Bi-Level Routing Attention (BRA).
• Preserve spatial detail through a channel-spatial attention–driven skip connection (SCCSA).
• Demonstrate performance gains on diverse medical imaging datasets compared with SOTA methods.

Proposed method

• Build a U-shaped encoder-decoder using BiFormer blocks as the core, with BRA for region-to-region and token-to-token attention.
• Employ a hierarchical encoder-decoder with patch embedding/merging and patch expanding to form a symmetric architecture.
• Redesign skip connections with SCCSA to boost cross-dimension interactions using channel and spatial attention implemented via convolutions.
• Use depth-wise convolutions to encode positional information and a hybrid loss (Dice + Cross-Entropy) with dataset-specific weighting.
• Leverage preprocessing and patch-based tokenization to enable efficient, query-aware sparse attention with O((HW)^(4/3)) complexity.
• Train on Synapse, ISIC-2018, and CVC-ClinicDB with data augmentation and standard optimization schedules.

Experimental results

Research questions

• RQ1 Can BRAU-Net++ outperform state-of-the-art CNN-, Transformer-, and hybrid-based medical image segmentation methods across diverse modalities?
• RQ2 Does Bi-Level Routing Attention enable effective long-range dependency modeling with reduced computation and memory?
• RQ3 Does the SCCSA skip-connection strategy improve preservation of spatial detail and cross-scale feature integration during decoding?

Key findings

• BRAU-Net++ achieves high segmentation performance across benchmarks, with reported average metrics including DSC and mIoU (e.g., DSC 82.47%, mIoU 84.01% on ISIC-2018 and similar on CVC-ClinicDB).
• Compared with TransUNet and Swin-Unet, BRAU-Net++ improves DSC by 4.49% and 3.34% and reduces HD by 12.62 mm and 2.48 mm on Synapse, respectively.
• Bi-Level Routing Attention enables long-range modeling with lower computational cost than vanilla full attention (complexity reduced to O((HW)^(4/3))).
• SCCSA skip connections help recover spatial information lost during down-sampling and enhance cross-dimension feature interactions.
• BRAU-Net++ demonstrates robust generalization across three public datasets (Synapse, ISIC-2018, CVC-ClinicDB).