[论文解读] DAE-Former: Dual Attention-guided Efficient Transformer for Medical Image Segmentation

Reza Azad, René Arimond+3 more|arXiv (Cornell University)|Dec 27, 2022

Advanced Neural Network Applications被引用 5

一句话总结

DAE-Former 提出了一种计算高效、无需卷积的 U-Net 类 Transformer 架构，用于医学图像分割，通过引入双注意力机制：先使用高效空间注意力，再使用转置（通道）注意力，以建模空间和通道维度上的长距离依赖关系。该方法在无需预训练的情况下，在心脏和皮肤病变分割基准上实现了最先进性能，通过引入一种新颖的跳跃连接交叉注意力（SCCA）模块，增强了特征融合与定位能力。

ABSTRACT

Transformers have recently gained attention in the computer vision domain due to their ability to model long-range dependencies. However, the self-attention mechanism, which is the core part of the Transformer model, usually suffers from quadratic computational complexity with respect to the number of tokens. Many architectures attempt to reduce model complexity by limiting the self-attention mechanism to local regions or by redesigning the tokenization process. In this paper, we propose DAE-Former, a novel method that seeks to provide an alternative perspective by efficiently designing the self-attention mechanism. More specifically, we reformulate the self-attention mechanism to capture both spatial and channel relations across the whole feature dimension while staying computationally efficient. Furthermore, we redesign the skip connection path by including the cross-attention module to ensure the feature reusability and enhance the localization power. Our method outperforms state-of-the-art methods on multi-organ cardiac and skin lesion segmentation datasets without requiring pre-training weights. The code is publicly available at https://github.com/mindflow-institue/DAEFormer.

研究动机与目标

• 解决标准自注意力机制在高分辨率医学图像中计算复杂度为二次方的问题。
• 克服 CNN 在捕捉医学图像分割中长距离空间与结构依赖关系方面的局限性。
• 设计一种纯 Transformer 基础的 U-Net 架构，保持高效率与强定位能力，且不依赖 CNN 主干网络。
• 通过引入跳跃连接交叉注意力（SCCA）模块，提升编码器与解码器之间的特征融合，增强表征复用性与分割精度。
• 在无需预训练权重的情况下，实现在多器官与皮肤病变分割任务上的最先进性能。

提出的方法

• 采用重叠的 4×4 小块进行嵌入，将输入图像标记化为序列标记，消除卷积操作。
• 采用分层的编码器-解码器结构，包含三个堆叠模块，通过小块下采样与上采样层降低或增加空间维度，同时将通道数加倍。
• 引入双注意力模块，包含两个阶段：(1) 高效注意力通过归一化查询与键计算空间重要性，将 O(N²) 复杂度降低至 O(d²N)；(2) 转置注意力通过转置键与值矩阵，建模通道间依赖关系。
• 使用 softmax 归一化的查询与键实现高效注意力机制，随后与值矩阵相乘，生成全局上下文向量，实现高表征能力的同时降低计算量。
• 设计跳跃连接交叉注意力（SCCA）模块，通过同时关注空间与通道表征，融合编码器与解码器路径的特征，提升特征复用性与定位精度。
• 在最后阶段使用线性投影层，从优化后的特征生成分割图。

实验结果

研究问题

• RQ1 纯 Transformer 基础的 U-Net 架构是否能在保持计算效率的同时，实现在医学图像分割中的最先进性能？
• RQ2 如何重构自注意力机制，以在不引入二次方复杂度的前提下，高效建模空间与通道依赖关系？
• RQ3 引入跳跃连接交叉注意力（SCCA）模块在多大程度上提升了 Transformer 基础分割网络中的特征融合与定位能力？
• RQ4 仅基于双注意力机制（高效空间注意力 + 转置通道注意力）的模型是否能在无需预训练的情况下，超越 CNN-Transformer 混合模型与现有纯 Transformer 模型在医学图像分割中的表现？
• RQ5 所提出的架构在心脏与皮肤病变分割等多样化医学影像任务中具有多大程度的泛化能力？

主要发现

• DAE-Former 在无需预训练权重的情况下，于多器官心脏分割与皮肤病变分割数据集上实现了最先进性能。
• 该模型展现出强大的泛化能力，在心脏与皮肤病变分割基准上均优于现有最先进方法。
• 双注意力机制——先高效空间注意力，后转置通道注意力——以较低计算成本，有效建模了空间与通道维度上的长距离依赖关系。
• 跳跃连接交叉注意力（SCCA）模块通过实现编码器与解码器特征之间的跨模态注意力，显著提升了特征复用性与定位精度。
• 由于采用高效注意力与纯 Transformer 设计且无卷积主干，该架构保持了高效率，适用于高分辨率医学图像。
• 消融实验验证了双注意力机制与 SCCA 模块对性能提升的显著贡献，证实了其设计选择的有效性。