[论文解读] More than Encoder: Introducing Transformer Decoder to Upsample

Yijiang Li, Wentian Cai+3 more|arXiv (Cornell University)|Jun 20, 2021

Radiomics and Machine Learning in Medical Imaging参考文献 52被引用 6

一句话总结

本文提出了一种新型基于Transformer的上采样方法——窗口注意力上采样（WAU），该方法利用Transformer解码器增强医学图像分割中的特征重建。通过在下采样路径的特征上应用局部窗口自注意力机制，WAU提升了空间细节和语义一致性，在Synapse、MSD Brain和ACDC数据集上，多种架构均实现了最先进性能。

ABSTRACT

Medical image segmentation methods downsample images for feature extraction and then upsample them to restore resolution for pixel-level predictions. In such a schema, upsample technique is vital in restoring information for better performance. However, existing upsample techniques leverage little information from downsampling paths. The local and detailed feature from the shallower layer such as boundary and tissue texture is particularly more important in medical segmentation compared with natural image segmentation. To this end, we propose a novel upsample approach for medical image segmentation, Window Attention Upsample (WAU), which upsamples features conditioned on local and detailed features from downsampling path in local windows by introducing attention decoders of Transformer. WAU could serve as a general upsample method and be incorporated into any segmentation model that possesses lateral connections. We first propose the Attention Upsample which consists of Attention Decoder (AD) and bilinear upsample. AD leverages pixel-level attention to model long-range dependency and global information for a better upsample. Bilinear upsample is introduced as the residual connection to complement the upsampled features. Moreover, considering the extensive memory and computation cost of pixel-level attention, we further design a window attention scheme to restrict attention computation in local windows instead of the global range. We evaluate our method (WAU) on classic U-Net structure with lateral connections and achieve state-of-the-art performance on Synapse multi-organ segmentation, Medical Segmentation Decathlon (MSD) Brain, and Automatic Cardiac Diagnosis Challenge (ACDC) datasets. We also validate the effectiveness of our method on multiple classic architectures and achieve consistent improvement.

研究动机与目标

• 为解决现有上采样技术在医学图像分割中对下采样路径中局部和细节特征利用不足的问题。
• 探索Transformer解码器在特征上采样中的潜力，突破其在序列生成中的典型应用。
• 开发一种内存和计算效率高的上采样机制，以保留医学影像中至关重要的细粒度细节。
• 设计一种通用的上采样模块，可与任何使用跳跃连接的分割模型兼容。
• 在多种架构（包括2D和3D U-Net、FCN和DeepLabV3）上展示一致的性能提升。

提出的方法

• 提出注意力上采样（AU），一种双流模块，结合基于Transformer的注意力解码器（AD）与双线性上采样作为残差连接。
• 设计窗口注意力解码器（WAD），将自注意力计算限制在局部空间窗口内，降低计算成本，同时保持长程建模能力。
• 使用逐通道可分离卷积，在每个窗口内实现高效的特征投影，提升在小型医学数据集上的性能。
• 对局部特征窗口计算查询、键和值矩阵，应用窗口化自注意力机制，实现上下文感知的上采样。
• 将WAU集成到类似U-Net的架构的跳跃连接路径中，支持与标准分割头端到端训练。
• 采用残差学习策略，其中双线性上采样与基于注意力的优化互补，稳定训练并提升特征质量。

实验结果

研究问题

• RQ1 Transformer的解码器能否被有效重用于医学图像分割中的特征上采样？
• RQ2 与标准转置卷积或双线性上采样相比，将下采样路径中的局部细节特征整合进来，如何提升上采样质量？
• RQ3 在窗口化注意力机制中，实现局部细节与长程上下文建模之间平衡的最优窗口尺寸是什么？
• RQ4 所提出的WAU方法在不同分割架构（包括2D和3D模型）上的泛化能力如何？
• RQ5 在上采样中使用全局注意力与局部窗口化注意力时，计算成本与性能增益之间的权衡是什么？

主要发现

• WAU在Synapse多器官分割数据集上实现了最先进性能，优于现有方法。
• 在医学分割十项全能（MSD Brain）数据集中，WAU相较于基线ResUNet实现了平均DSC提升1.47分。
• WAU使3D U-Net的性能提升0.36 DSC，证明其在3D医学体积分割中的强大泛化能力。
• WAU的最优窗口尺寸为4，窗口过大或过小均会导致性能下降，原因分别为注意力焦点丢失或上下文不足。
• 在卷积类型中，逐通道可分离卷积表现最佳，尤其在参数受限的小型医学数据集中优势明显。
• 可视化结果表明，WAU的注意力权重在相关解剖结构上具有空间聚焦性，增强了特征区分度和病灶可见性。