[论文解读] ScaleFormer: Revisiting the Transformer-based Backbones from a Scale-wise Perspective for Medical Image Segmentation

Huimin Huang, Shiao Xie+5 more|arXiv (Cornell University)|Jul 29, 2022

Advanced Neural Network Applications被引用 13

一句话总结

ScaleFormer 提出了一种用于医学图像分割的新型基于 Transformer 的主干网络，通过双分支设计解决尺度相关挑战：在尺度内采用轻量级双轴多头自注意力（DA-MSA）模块，实现每个尺度内局部与全局线索的融合；在尺度间通过空间感知 Transformer 模块增强跨尺度特征交互。该方法在多个医学图像分割基准上取得了最先进性能。

ABSTRACT

Recently, a variety of vision transformers have been developed as their capability of modeling long-range dependency. In current transformer-based backbones for medical image segmentation, convolutional layers were replaced with pure transformers, or transformers were added to the deepest encoder to learn global context. However, there are mainly two challenges in a scale-wise perspective: (1) intra-scale problem: the existing methods lacked in extracting local-global cues in each scale, which may impact the signal propagation of small objects; (2) inter-scale problem: the existing methods failed to explore distinctive information from multiple scales, which may hinder the representation learning from objects with widely variable size, shape and location. To address these limitations, we propose a novel backbone, namely ScaleFormer, with two appealing designs: (1) A scale-wise intra-scale transformer is designed to couple the CNN-based local features with the transformer-based global cues in each scale, where the row-wise and column-wise global dependencies can be extracted by a lightweight Dual-Axis MSA. (2) A simple and effective spatial-aware inter-scale transformer is designed to interact among consensual regions in multiple scales, which can highlight the cross-scale dependency and resolve the complex scale variations. Experimental results on different benchmarks demonstrate that our Scale-Former outperforms the current state-of-the-art methods. The code is publicly available at: https://github.com/ZJUGiveLab/ScaleFormer.

研究动机与目标

• 为解决现有视觉 Transformer 中尺度内局部与全局特征融合不充分的问题，提升小目标的表征学习能力。
• 解决跨尺度特征交互不足的问题，尤其针对尺寸和形状多变的医学图像对象。
• 通过显式建模特征层级中的尺度依赖关系，提升医学图像分割中的表征学习能力。
• 开发一种轻量化、高效的主干网络，提升小目标的信号传播能力，并增强对尺度变化的鲁棒性。

提出的方法

• 提出一种基于尺度的尺度内 Transformer 模块，利用轻量级双轴多头自注意力（DA-MSA）机制，将基于 CNN 的局部特征与基于 Transformer 的全局线索进行融合。
• 在 DA-MSA 中分别应用行方向和列方向的注意力机制，独立捕获两个方向的长程依赖，提升局部与全局特征的融合效果。
• 设计一种空间感知的跨尺度 Transformer 模块，通过关注不同尺度间对应区域，显式建模跨尺度交互。
• 将编码器多个阶段的特征图作为跨尺度模块的输入，实现层级化特征优化。
• 采用残差连接和层归一化以稳定训练过程并改善梯度流动。
• 在医学图像数据集上使用标准分割损失（如 Dice 损失）端到端训练完整模型。

实验结果

研究问题

• RQ1 如何在每个尺度内有效融合局部与全局特征，以提升小目标的表征学习能力？
• RQ2 跨尺度特征交互在多大程度上能增强模型处理医学图像中尺寸和形状各异对象的能力？
• RQ3 轻量化、尺度感知的 Transformer 模块是否能在不显著增加计算成本的前提下提升性能？
• RQ4 所提出的尺度内与跨尺度设计相较于标准视觉 Transformer 主干网络，在医学图像分割中表现如何？

主要发现

• ScaleFormer 在多个医学图像分割基准（包括 nnU-Net 和其他标准数据集）上达到最先进性能。
• 消融实验表明，尺度内与跨尺度模块均对性能提升有显著贡献，尤其在小目标上表现突出。
• 由于更优的局部-全局特征融合，模型在小尺寸和不规则形状结构上的分割精度显著提高。
• 空间感知的跨尺度模块有效捕捉了跨尺度依赖关系，提升了尺度层级上的特征表示鲁棒性。
• 与现有基于 Transformer 的主干网络相比，该方法在推理速度和参数效率方面保持了竞争力。