[论文解读] Narrowing the semantic gaps in U-Net with learnable skip connections: The case of medical image segmentation

Haonan Wang, Peng Cao+3 more|arXiv (Cornell University)|Dec 23, 2023

COVID-19 diagnosis using AI被引用 4

一句话总结

本文提出UDTransNet，一种基于U-Net的医学图像分割框架，通过可学习的注意力机制替代标准的跳跃连接，以缩小编码器与解码器特征之间的语义差距。通过引入用于多尺度特征融合的双注意力变换器（DAT）以及用于跨阶段一致性的解码器引导重校准注意力（DRA）模块，UDTransNet在多个医学影像基准上实现了参数更少的最先进性能。

ABSTRACT

Most state-of-the-art methods for medical image segmentation adopt the encoder-decoder architecture. However, this U-shaped framework still has limitations in capturing the non-local multi-scale information with a simple skip connection. To solve the problem, we firstly explore the potential weakness of skip connections in U-Net on multiple segmentation tasks, and find that i) not all skip connections are useful, each skip connection has different contribution; ii) the optimal combinations of skip connections are different, relying on the specific datasets. Based on our findings, we propose a new segmentation framework, named UDTransNet, to solve three semantic gaps in U-Net. Specifically, we propose a Dual Attention Transformer (DAT) module for capturing the channel- and spatial-wise relationships to better fuse the encoder features, and a Decoder-guided Recalibration Attention (DRA) module for effectively connecting the DAT tokens and the decoder features to eliminate the inconsistency. Hence, both modules establish a learnable connection to solve the semantic gaps between the encoder and the decoder, which leads to a high-performance segmentation model for medical images. Comprehensive experimental results indicate that our UDTransNet produces higher evaluation scores and finer segmentation results with relatively fewer parameters over the state-of-the-art segmentation methods on different public datasets. Code: https://github.com/McGregorWwww/UDTransNet.

研究动机与目标

• 探究固定跳跃连接在U-Net用于医学图像分割时的局限性。
• 识别出并非所有跳跃连接都同样有效，且最优组合因数据集而异。
• 解决多尺度编码器特征之间以及编码器-解码器阶段之间的语义差距问题。
• 设计一种可学习的、基于注意力的机制，以提升U-Net架构中特征的兼容性。
• 在减少参数量的同时实现最先进水平的分割性能。

提出的方法

• 提出双注意力变换器（DAT）模块，结合通道融合注意力（CFA）与空间选择注意力（SSA），以建模多尺度编码器特征中的通道与空间依赖性。
• 在CFA中使用自注意力机制对不同编码器阶段的通道进行对齐，在SSA中使用交叉注意力机制突出显示跨尺度的重要空间区域。
• 引入解码器引导重校准注意力（DRA）模块，利用解码器特征对DAT输出进行重校准，以减少语义不一致性。
• 将DAT与DRA模块集成到U-Net主干网络的跳跃连接中，实现基于注意力的特征融合的端到端训练。
• 采用混合架构，将基于CNN的编码器与基于Transformer的注意力模块结合在跳跃连接中，而非主编码器或解码器中。
• 对超参数（如注意力头数量与变换器层数）进行优化，发现4个头与4层为最优设置。

实验结果

研究问题

• RQ1 在不同医学图像分割任务中，U-Net中的所有跳跃连接是否都具有同等有效性？
• RQ2 多尺度编码器特征之间以及编码器-解码器阶段之间的语义差距如何影响分割性能？
• RQ3 跳跃连接中可学习的注意力机制是否能更好地对齐特征并提升分割准确率？
• RQ4 将注意力机制集成到跳跃连接中是否比单独应用于编码器或解码器更有效？
• RQ5 所提出的注意力模块中，注意力头与层数的最优配置是什么？

主要发现

• UDTransNet在三个公开数据集上实现了最先进水平的Dice分数：在GlaS数据集上为91.03±0.56，在Synapse数据集上为89.28±0.58，优于现有方法。
• 当使用相同主干网络（ResNet-34）时，UDTransNet在GlaS数据集上达到91.03±0.56的Dice分数，显著优于MCTrans（88.44±0.74）与TransUNet（87.41±1.15）。
• 当使用Swin Transformer作为主干网络时，UDTransNet在GlaS数据集上仍取得88.03±1.27的Dice分数，优于使用相同主干网络的Swin-UNet（89.04±0.88）。
• 该模型对超参数变化具有鲁棒性，最优性能出现在4个注意力头与4层时，表明对这些设置的敏感度较低。
• 消融实验证实，DAT与DRA模块均不可或缺，其中DRA显著提升了编码器与解码器之间的特征一致性。
• 与最先进模型相比，UDTransNet以更少的参数实现了更高的性能，展现出卓越的效率与有效性。