[论文解读] RaBiT: An Efficient Transformer using Bidirectional Feature Pyramid Network with Reverse Attention for Colon Polyp Segmentation

Nguyen Hoang Thuan, Thi-Oanh Nguyen+3 more|arXiv (Cornell University)|Jul 12, 2023

Radiomics and Machine Learning in Medical Imaging被引用 7

一句话总结

RaBiT 提出了一种基于轻量级 Transformer 的编码器-解码器网络，用于结肠息肉分割，采用双向特征金字塔与反向注意力（RaBiFPN）以及瓶颈优化的反向注意力模块，以增强多尺度特征融合与边界细化。该方法在高分辨率 ETIS-Larib 数据集上实现了 4.5% 的 mIoU 提升，同时相比基线模型参数量减少了 37%。

ABSTRACT

Automatic and accurate segmentation of colon polyps is essential for early diagnosis of colorectal cancer. Advanced deep learning models have shown promising results in polyp segmentation. However, they still have limitations in representing multi-scale features and generalization capability. To address these issues, this paper introduces RaBiT, an encoder-decoder model that incorporates a lightweight Transformer-based architecture in the encoder to model multiple-level global semantic relationships. The decoder consists of several bidirectional feature pyramid layers with reverse attention modules to better fuse feature maps at various levels and incrementally refine polyp boundaries. We also propose ideas to lighten the reverse attention module and make it more suitable for multi-class segmentation. Extensive experiments on several benchmark datasets show that our method outperforms existing methods across all datasets while maintaining low computational complexity. Moreover, our method demonstrates high generalization capability in cross-dataset experiments, even when the training and test sets have different characteristics.

研究动机与目标

• 解决现有模型在捕捉多尺度特征以及在多样化结肠镜数据集上泛化能力方面的局限性。
• 通过新型反向注意力机制实现的迭代细化，提升息肉边界分割性能。
• 在保持高精度语义分割性能的同时，降低计算成本。
• 通过双向特征金字塔结构，增强编码器-解码器层级之间的特征融合。
• 通过引入兼容多类别的反向注意力模块，实现有效的多类别分割。

提出的方法

• 编码器采用基于轻量级 MiT-B3 的 Transformer，以捕捉多尺度下的分层全局语义关系。
• 解码器采用堆叠的双向特征金字塔（RaBiFPN）模块，通过加权拼接与跳跃连接，融合来自不同编码器层级的特征。
• 引入一种基于瓶颈结构的轻量级反向注意力（RA）模块，以减少参数量与计算成本，同时支持对息肉边界的重复细化。
• 设计了一种多类别 RA 模块，通过强调区域边界，更好地分割非肿瘤性和肿瘤性息肉。
• RA 模块中使用 Sigmoid 或 Softmax 激活函数，实验表明在小目标分割任务中 Softmax 表现更优。
• 网络采用交叉熵损失与 Dice 损失联合端到端训练，消融实验验证了各组件的有效性。

实验结果

研究问题

• RQ1 轻量级 Transformer 编码器是否能有效建模结肠息肉分割中的全局上下文与多尺度特征？
• RQ2 与标准 FPN 或 BiFPN 相比，结合反向注意力的双向特征金字塔是否能提升特征融合与边界细化效果？
• RQ3 瓶颈优化的反向注意力模块是否能在保持或提升分割精度的同时降低计算成本？
• RQ4 该方法在图像分辨率与息肉特征各异的多样化结肠镜数据集上是否具备良好的泛化能力？
• RQ5 多类别反向注意力模块是否能提升对肿瘤性与非肿瘤性息肉的分割性能？

主要发现

• 与标准 BiFPN 相比，RaBiT 在高分辨率 ETIS-Larib 数据集上实现了 4.5% 的 mIoU 提升，表明其在大尺寸图像上具有更强的有效性。
• 在 Neo-Small 数据集上，RA 模块中使用 Softmax 激活函数使 Dice 非度量指标提升 6.43%，表明其在小目标分割方面表现更优。
• 所提出的瓶颈模块将参数量从 81.79M 降低至 51.72M，GFLOPs 从 41.37 降至 23.12，显著降低了计算成本。
• 当使用 4 个 RaBiFPN 模块时，RaBiT 在所有基准数据集上均优于 SegFormer-B3 及其他 SOTA 模型，实现最佳整体性能。
• 在跨数据集泛化测试中，即使训练集与测试集分布不同，RaBiT 仍保持优异性能，展现出高度鲁棒性。
• 消融实验表明，RaBiFPN 与瓶颈模块均能独立提升 mIoU，且在高分辨率数据集上的增益最为显著。