[论文解读] 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++ 是一种 U 形卷积神经网络-Transformer 架构，以 Bi-Level Routing Attention (BRA) 作为核心模块，并通过通道-空间跳过连接 (SCCSA) 实现高效且准确的医学图像分割，在多数据集上表现出色。

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）来保留空间细节。
• 在多样的医学影像数据集上展示相对于最先进方法的性能提升。

提出的方法

• 以 BiFormer 模块为核心构建 U 形编码器-解码器，采用 BRA 进行区域到区域和令牌到令牌的注意力。
• 采用带补丁嵌入/合并和补丁扩展的分层编码器-解码器，形成对称结构。
• 通过 SCCSA 重设计跳过连接，利用通过卷积实现的通道和空间注意力来提升跨维度交互。
• 使用深度卷积来编码位置信息，并采用数据集特定权重的混合损失（Dice + 交叉熵）。
• 利用预处理和基于补丁的标记化实现高效、查询感知的稀疏注意力，复杂度为 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 相比，BRAU-Net++ 在 Synapse 上将 DSC 提升了 4.49% 和 3.34%，HD 分别降低了 12.62 mm 和 2.48 mm。
• Bi-Level Routing Attention 使远程建模在计算成本低于普通全注意力（复杂度降低到 O((HW)^(4/3))）的情况下成为可能。
• SCCSA 跳过连接有助于恢复在下采样过程中丢失的空间信息，并增强跨维度的特征交互。
• BRAU-Net++ 在三个公开数据集（Synapse、ISIC-2018、CVC-ClinicDB）上展示出强健的泛化能力。