[论文解读] The Fully Convolutional Transformer for Medical Image Segmentation

Athanasios Tragakis, Chaitanya Kaul+2 more|arXiv (Cornell University)|Jun 1, 2022

Radiomics and Machine Learning in Medical Imaging被引用 7

一句话总结

该论文提出了一种新型、紧凑且完全卷积的Transformer架构——全卷积Transformer（FCT），用于医学图像分割。该模型结合了卷积注意力机制与多分辨率空洞卷积的广域聚焦模块，以捕捉长距离依赖关系和层次化特征。FCT在多个数据集上实现了最先进性能，Dice分数最高提升4.4%，同时参数量仅为nnUNet和TransUNet等竞争模型的五分之一。

ABSTRACT

We propose a novel transformer model, capable of segmenting medical images of varying modalities. Challenges posed by the fine grained nature of medical image analysis mean that the adaptation of the transformer for their analysis is still at nascent stages. The overwhelming success of the UNet lay in its ability to appreciate the fine-grained nature of the segmentation task, an ability which existing transformer based models do not currently posses. To address this shortcoming, we propose The Fully Convolutional Transformer (FCT), which builds on the proven ability of Convolutional Neural Networks to learn effective image representations, and combines them with the ability of Transformers to effectively capture long-term dependencies in its inputs. The FCT is the first fully convolutional Transformer model in medical imaging literature. It processes its input in two stages, where first, it learns to extract long range semantic dependencies from the input image, and then learns to capture hierarchical global attributes from the features. FCT is compact, accurate and robust. Our results show that it outperforms all existing transformer architectures by large margins across multiple medical image segmentation datasets of varying data modalities without the need for any pre-training. FCT outperforms its immediate competitor on the ACDC dataset by 1.3%, on the Synapse dataset by 4.4%, on the Spleen dataset by 1.2% and on ISIC 2017 dataset by 1.1% on the dice metric, with up to five times fewer parameters. Our code, environments and models will be available via GitHub.

研究动机与目标

• 为解决现有基于Transformer的模型在医学图像分割中缺乏UNet所具备的细粒度特征学习能力的问题。
• 开发一种完全卷积的Transformer，以有效捕捉医学图像中的长距离语义依赖关系和层次化全局上下文。
• 构建一种轻量化、高精度且鲁棒的模型，在无需预训练的情况下超越最先进性能。
• 以一种消除位置编码和分块标记化需求的方式，将卷积操作与自注意力机制相结合。

提出的方法

• FCT采用一种新颖的全卷积Transformer层，由两个主要组件构成：卷积注意力模块和广域聚焦模块。
• 卷积注意力模块使用深度可分离卷积生成重叠的、具有空间感知能力的块，替代标准的分块嵌入，并消除对位置编码的需求。
• 广域聚焦模块通过在三个分支中使用具有递增空洞率的多分辨率空洞卷积，捕捉从局部到全局的层次化特征。
• 模型处理输入分为两个阶段：首先通过卷积注意力模块学习长距离语义依赖关系，随后利用广域聚焦模块构建层次化全局表征。
• 采用编码器到解码器的跳跃连接，遵循UNet设计，以保留空间细节并改善特征传播。
• 模型采用标准数据增强进行端到端训练，推理阶段使用滑动窗口推理并结合后处理平均重叠预测结果。

实验结果

研究问题

• RQ1 全卷积Transformer架构是否能在无需预训练的情况下，超越现有的CNN和基于Transformer的模型在医学图像分割中的表现？
• RQ2 卷积注意力与多分辨率空洞卷积的结合，如何影响模型在捕捉长距离与细粒度特征方面的性能？
• RQ3 广域聚焦模块在医学图像分割中，其最优配置（分支数量与空洞率）为何种设置？
• RQ4 跳跃连接在所提出的FCT架构中对性能的贡献程度如何？
• RQ5 轻量化基于Transformer的模型是否能以显著更少的参数量实现最先进性能？

主要发现

• 在ACDC数据集上，FCT在无需预训练的情况下，Dice分数较所有现有基于Transformer和CNN的模型高出1.3%。
• 在Synapse数据集上，FCT相比其直接竞争对手Dice分数提升4.4%，且参数量仅为nnUNet的五分之一。
• 在Spleen数据集上，FCT较次优模型Dice分数提升1.2%，表明其在不同解剖结构上的鲁棒性。
• 在ISIC 2017数据集上，FCT较最接近的竞争对手Dice分数高出1.1%，证实其在多种图像模态下的泛化能力。
• 在ACDC Post-2017-MICCAI-Challenge在线测试集上，FCT创下新最先进性能，仅用3170万参数即超越大型集成模型和nnUNet。
• 消融研究证实，跳跃连接以及广域聚焦模块的最优配置（三分支、线性递增空洞率）对达到峰值性能至关重要。