[论文解读] Multi-Compound Transformer for Accurate Biomedical Image Segmentation

Yuanfeng Ji, Ruimao Zhang+5 more|arXiv (Cornell University)|Jun 28, 2021

Advanced Neural Network Applications参考文献 23被引用 15

一句话总结

本文提出了一种新型视觉Transformer架构——多化合物Transformer（MCTrans），通过自注意力与交叉注意力机制建模跨尺度特征依赖关系和语义关系，从而提升生物医学图像分割性能。MCTrans与UNet类网络结合，在六个基准测试中实现最先进性能，Dice分数相比UNet最高提升4.34%。

ABSTRACT

The recent vision transformer(i.e.for image classification) learns non-local attentive interaction of different patch tokens. However, prior arts miss learning the cross-scale dependencies of different pixels, the semantic correspondence of different labels, and the consistency of the feature representations and semantic embeddings, which are critical for biomedical segmentation. In this paper, we tackle the above issues by proposing a unified transformer network, termed Multi-Compound Transformer (MCTrans), which incorporates rich feature learning and semantic structure mining into a unified framework. Specifically, MCTrans embeds the multi-scale convolutional features as a sequence of tokens and performs intra- and inter-scale self-attention, rather than single-scale attention in previous works. In addition, a learnable proxy embedding is also introduced to model semantic relationship and feature enhancement by using self-attention and cross-attention, respectively. MCTrans can be easily plugged into a UNet-like network and attains a significant improvement over the state-of-the-art methods in biomedical image segmentation in six standard benchmarks. For example, MCTrans outperforms UNet by 3.64%, 3.71%, 4.34%, 2.8%, 1.88%, 1.57% in Pannuke, CVC-Clinic, CVC-Colon, Etis, Kavirs, ISIC2018 dataset, respectively. Code is available at https://github.com/JiYuanFeng/MCTrans.

研究动机与目标

• 解决现有视觉Transformer在生物医学分割中的局限性，特别是其无法建模标签间的跨尺度依赖关系与语义关系。
• 提升医学图像中不同解剖结构的特征表示一致性与可分性。
• 通过超越单尺度注意力的长距离上下文建模，更好地捕捉大小多变的病灶。
• 构建统一框架，结合多尺度特征学习与语义结构挖掘，实现精确分割。
• 在多样化生物医学图像分割基准上验证所提方法的有效性，且计算开销极低。

提出的方法

• MCTrans采用Transformer-自注意力（TSA）模块，对多尺度特征图中的像素级关系进行建模，实现全面的跨尺度上下文编码。
• 引入可学习代理嵌入以表示语义类别，并通过自注意力机制优化，捕捉类别特异性依赖关系。
• Transformer-交叉注意力（TCA）模块实现特征表示与代理嵌入之间的交互，增强特征的可分性与一致性。
• 对更新后的代理嵌入应用辅助损失，显式提升同类类别内的特征相关性与不同类别间的可分性。
• MCTrans模块无缝集成至UNet类编码器-解码器架构中，支持端到端训练与推理。
• 多尺度卷积特征被嵌入为token，经TSA与TCA模块处理后，上采样并融合生成分割输出。

实验结果

研究问题

• RQ1 与单尺度自注意力相比，跨尺度注意力机制是否能提升生物医学图像分割中的特征表示？
• RQ2 学习不同标签之间的语义关系对特征可分性与分割精度有何影响？
• RQ3 在小样本与复杂数据集上，引入带辅助监督的可学习代理嵌入能在多大程度上提升模型性能？
• RQ4 与标准UNet相比，所提出的MCTrans框架是否仅带来轻微计算成本增加，同时保持高性能？
• RQ5 MCTrans是否能在多样化的生物医学分割任务中泛化，包括病灶与解剖结构分割？

主要发现

• 在Pannuke数据集上，MCTrans相比UNet的Dice分数提升3.64%；在CVC-Clinic、CVC-Colon、ETIS、Kavirs和ISIC2018数据集上，分别提升3.71%、4.34%、2.8%、1.88%和1.57%。
• 在Pannuke数据集上，使用ResNet-34主干网络时，MCTrans的Dice分数达到68.40%，优于UNet++（66.08%）与CENet（66.50%），尽管参数量相近。
• 消融实验表明，若移除代理嵌入的辅助损失，性能从68.40%降至67.87%，证实其在增强特征相关性与可分性方面的重要作用。
• TSA与TCA模块的最佳数量为4个，超过此数量性能下降，表明模型容量需与数据集规模相平衡。
• MCTrans仅带来适度的FLOPs增加（18.065G vs. 14.037G for UNet），同时实现性能提升，展现出高效性与准确性的良好平衡。
• 图4的定性结果表明，MCTrans生成的分割掩码更精确且一致，尤其在小尺寸或不规则形状病灶上表现更优。