[论文解读] Medical Transformer: Gated Axial-Attention for Medical Image Segmentation

Jeya Maria Jose Valanarasu, Poojan Oza+2 more|arXiv (Cornell University)|Feb 21, 2021

Radiomics and Machine Learning in Medical Imaging被引用 6

一句话总结

本文提出医学变换器（MedT），一种用于医学图像分割的门控轴向注意力网络，通过将可学习门控机制融入轴向自注意力，增强长距离特征建模能力。该方法在三个医学影像数据集上实现了最先进性能，且无需预训练，通过一种新颖的局部-全局训练策略（LoGo）超越了卷积神经网络和现有的基于变换器的模型。

ABSTRACT

Over the past decade, Deep Convolutional Neural Networks have been widely adopted for medical image segmentation and shown to achieve adequate performance. However, due to the inherent inductive biases present in the convolutional architectures, they lack understanding of long-range dependencies in the image. Recently proposed Transformer-based architectures that leverage self-attention mechanism encode long-range dependencies and learn representations that are highly expressive. This motivates us to explore Transformer-based solutions and study the feasibility of using Transformer-based network architectures for medical image segmentation tasks. Majority of existing Transformer-based network architectures proposed for vision applications require large-scale datasets to train properly. However, compared to the datasets for vision applications, for medical imaging the number of data samples is relatively low, making it difficult to efficiently train transformers for medical applications. To this end, we propose a Gated Axial-Attention model which extends the existing architectures by introducing an additional control mechanism in the self-attention module. Furthermore, to train the model effectively on medical images, we propose a Local-Global training strategy (LoGo) which further improves the performance. Specifically, we operate on the whole image and patches to learn global and local features, respectively. The proposed Medical Transformer (MedT) is evaluated on three different medical image segmentation datasets and it is shown that it achieves better performance than the convolutional and other related transformer-based architectures. Code: https://github.com/jeya-maria-jose/Medical-Transformer

研究动机与目标

• 解决卷积神经网络在建模医学图像中长距离依赖关系方面的局限性。
• 克服医学影像中数据稀缺的问题，该问题阻碍了标准变换器的有效训练。
• 开发一种基于变换器的架构，通过显式建模全局上下文和局部细节来提升分割精度。
• 设计一种训练策略，有效利用整体图像和图像块级别的表征，以提升小样本医学数据集上的性能。

提出的方法

• 提出一种门控轴向注意力机制，通过在2D自注意力中引入可学习门控来控制注意力权重，通过选择性关注相关空间上下文来提升特征表征能力。
• 引入一种局部-全局（LoGo）训练策略，联合训练一个全局分支（在完整图像上）和一个局部分支（在图像块上），使模型能够同时学习全局上下文和细粒度的局部特征。
• 使用轴向注意力将2D自注意力分解为两个1D操作，降低计算成本，同时保持空间建模能力。
• 采用类似U-Net的编码器-解码器架构，用轴向注意力模块替代编码器中的标准卷积层。
• 采用两阶段训练协议：首先在10个周期内不更新门控参数以稳定特征表示，然后以基础初始学习率0.001联合训练所有组件。
• 通过仅在全局分支中保留两个编码器/解码器层，并在局部分支中采用基于图像块的处理方式，减少模型参数量，即使在多分支设计下仍保持高效。

实验结果

研究问题

• RQ1 基于变换器的架构结合轴向自注意力是否能在无需预训练的情况下，有效建模医学图像分割中的长距离依赖关系？
• RQ2 在轴向注意力机制中引入可学习门控在小样本医学影像数据集上是否能提升性能？
• RQ3 与仅在完整图像或仅在图像块上训练的标准训练策略相比，局部-全局（LoGo）训练策略在多大程度上增强了特征学习能力？
• RQ4 所提出的医学变换器（MedT）在有限医学数据上的分割精度是否优于现有的卷积神经网络和基于变换器的模型？

主要发现

• 在脑部超声（Brain US）数据集上，MedT的F1得分为88.84，优于U-Net（85.37）、Res-UNet（87.5）和Axial-UNet（88.39）。
• 消融实验表明，门控轴向注意力、LoGo训练策略以及多分支设计的每个组件均对性能提升有贡献。
• 仅使用两个全局编码器层的LoGo策略即取得优异性能，表明在结合局部块级学习时，最小程度的全局建模已足够。
• 尽管参数量仅为140万，MedT仍优于参数量显著更多的模型，如原始U-Net（1250万）和Res-UNet（532万），证明了其参数效率。
• 定性结果表明，MedT能正确分割出U-Net和Res-UNet误分类的背景区域，尤其在具有复杂背景纹理的超声扫描中表现更优。
• 该模型在无需ImageNet预训练的情况下达到最先进性能，而类似工作如TransUNet和TransFuse则依赖大规模预训练。