[论文解读] AFTer-UNet: Axial Fusion Transformer UNet for Medical Image Segmentation

Xiangyi Yan, Hao Tang+4 more|arXiv (Cornell University)|Oct 20, 2021

Advanced Neural Network Applications被引用 10

一句话总结

AFTer-UNet 提出了一种端到端的3D医学图像分割框架，将轴向融合变换器整合到U-Net架构中，以联合建模切片内和切片间长距离依赖关系。通过在融合前分别沿轴向和空间轴计算自注意力机制，该方法在保持SOTA性能的同时，降低了计算成本和GPU显存使用量，在多器官分割基准上相比U-Net实现4.34%的平均Dice分数提升，相比CoTr提升2.56%。

ABSTRACT

Recent advances in transformer-based models have drawn attention to exploring these techniques in medical image segmentation, especially in conjunction with the U-Net model (or its variants), which has shown great success in medical image segmentation, under both 2D and 3D settings. Current 2D based methods either directly replace convolutional layers with pure transformers or consider a transformer as an additional intermediate encoder between the encoder and decoder of U-Net. However, these approaches only consider the attention encoding within one single slice and do not utilize the axial-axis information naturally provided by a 3D volume. In the 3D setting, convolution on volumetric data and transformers both consume large GPU memory. One has to either downsample the image or use cropped local patches to reduce GPU memory usage, which limits its performance. In this paper, we propose Axial Fusion Transformer UNet (AFTer-UNet), which takes both advantages of convolutional layers' capability of extracting detailed features and transformers' strength on long sequence modeling. It considers both intra-slice and inter-slice long-range cues to guide the segmentation. Meanwhile, it has fewer parameters and takes less GPU memory to train than the previous transformer-based models. Extensive experiments on three multi-organ segmentation datasets demonstrate that our method outperforms current state-of-the-art methods.

研究动机与目标

• 解决现有基于2D和3D变换器的模型在医学图像分割中无法有效利用3D体数据中轴向（切片间）上下文信息的局限性。
• 在保持长距离建模能力的同时，降低3D体数据中3D自注意力机制带来的高GPU显存消耗。
• 通过融合多切片的上下文线索，提升细长及解剖结构多变器官的分割精度。
• 在参数数量更少、显存使用更低的前提下，保持模型效率，优于以往基于变换器的模型。
• 开发一种可扩展的端到端框架，兼容消费级GPU，适用于临床部署。

提出的方法

• 提出一种轴向融合机制，先在轴向上独立计算自注意力，再在每个2D切片内独立计算自注意力，最后融合输出结果。
• 将轴向融合变换器嵌入U-Net的编码器与解码器之间，保留U-Net的U形跳跃连接结构。
• 采用2D卷积神经网络作为编码器和解码器进行特征提取，而轴向融合变换器则用于捕捉跨切片的长距离依赖关系。
• 在轴向方向和每个2D切片内应用多头自注意力机制，并结合相对位置编码，以建模长距离上下文信息。
• 引入采样频率 $N_f$ 控制切片间注意力的密度，优先关注邻近切片，以提升效率与性能。
• 通过解耦轴向与空间注意力，将计算复杂度从 $O(HW D^2)$ 降低至 $O(HW D + H^2W^2)$，实现仅使用单张11GB GPU进行训练。

实验结果

研究问题

• RQ1 轴向融合的切片内与切片间注意力是否能超越纯2D或3D变换器模型，在3D医学图像分割中取得更优性能？
• RQ2 解耦轴向与空间注意力是否能在降低GPU显存使用的同时，保持长距离建模能力？
• RQ3 邻近轴向切片的数量如何影响分割性能，特别是在细长器官上的表现？
• RQ4 增加轴向融合变换器层数在多大程度上能提升分割精度？
• RQ5 轴向上采样频率的变化在多大程度上影响模型性能与效率？

主要发现

• 在Thorax-85数据集上，AFTer-UNet相比U-Net实现2.13%的平均Dice分数提升，相比CoTr提升0.69%。
• 对于挑战性较大的细长器官——食管，AFTer-UNet相比U-Net实现2.91%的Dice提升，相比CoTr提升1.45%。
• 在多器官数据集上，AFTer-UNet相比U-Net在肝脏上提升3.17%，左肾提升7.30%，右肾提升7.11%，胰腺提升6.64%。
• 相比CoTr，AFTer-UNet在肝脏上实现3.85%更高的Dice分数，左肾提升5.26%，右肾提升6.61%，胰腺提升5.68%。
• AFTer-UNet显著降低GPU显存使用，可在单张RTX-2080Ti（11GB）上训练，而此前的3D变换器模型需多张GPU。
• AFTer-UNet参数量为41.5M，与TransUNet（43.5M）和CoTr（41.9M）相当，证明其在不增加参数量的前提下实现高效性。