[论文解读] UCTransNet: Rethinking the Skip Connections in U-Net from a Channel-wise Perspective with Transformer

Haonan Wang, Peng Cao+2 more|arXiv (Cornell University)|Sep 9, 2021

Advanced Neural Network Applications被引用 44

一句话总结

UCTransNet 将普通 U-Net 跳跃连接替换为通道变换器（CTrans），它将多尺度编码器特征（CCT）融合并与解码器特征（CCA）对齐，从而在不同数据集上提升医学图像分割性能。

ABSTRACT

Most recent semantic segmentation methods adopt a U-Net framework with an encoder-decoder architecture. It is still challenging for U-Net with a simple skip connection scheme to model the global multi-scale context: 1) Not each skip connection setting is effective due to the issue of incompatible feature sets of encoder and decoder stage, even some skip connection negatively influence the segmentation performance; 2) The original U-Net is worse than the one without any skip connection on some datasets. Based on our findings, we propose a new segmentation framework, named UCTransNet (with a proposed CTrans module in U-Net), from the channel perspective with attention mechanism. Specifically, the CTrans module is an alternate of the U-Net skip connections, which consists of a sub-module to conduct the multi-scale Channel Cross fusion with Transformer (named CCT) and a sub-module Channel-wise Cross-Attention (named CCA) to guide the fused multi-scale channel-wise information to effectively connect to the decoder features for eliminating the ambiguity. Hence, the proposed connection consisting of the CCT and CCA is able to replace the original skip connection to solve the semantic gaps for an accurate automatic medical image segmentation. The experimental results suggest that our UCTransNet produces more precise segmentation performance and achieves consistent improvements over the state-of-the-art for semantic segmentation across different datasets and conventional architectures involving transformer or U-shaped framework. Code: https://github.com/McGregorWwww/UCTransNet.

研究动机与目标

• 证明 U-Net 中简单跳跃连接并非普遍有益，且在某些数据集上可能降低性能。
• 提出基于通道的Transformer跳跃机制（CCT + CCA），以更好地将多尺度编码器特征与解码器特征融合。
• 表明通道级融合能减小语义和分辨率差异并提高跨数据集的分割精度。
• 在 GlaS、MoNuSeg 与 Synapse 数据集上评估 UCTransNet，并与强基线 U-Net 和基于Transformer 的方法进行比较。

提出的方法

• 用通道变换器（CTrans）替换标准跳跃连接，其由 CCT（Channel-wise Cross Fusion Transformer，通道级跨融合 Transformer）用于多尺度编码器特征融合，以及 CCA（Channel-wise Cross Attention，通道级跨注意力）用于解码器特征融合，组合 Transformer 输出。
• CCT 通过将四个跳跃层分成补丁进行标记，将多头通道级跨注意力与拼接的 Key/Value 进行计算，并在 L 层中应用带残差连接的 MLP 来融合多尺度编码器特征。
• CCA 通过对 O_i 和 D_i 做池化以计算通道注意力图，对 O_i 进行重新校准，然后与上采样的解码器特征拼接。

实验结果

研究问题

• RQ1 基于通道的 Transformer 跳跃连接能否在医学图像分割中优于传统的 U-Net 跳跃连接？
• RQ2 多尺度通道级融合（CCT）如何与解码器感知融合（CCA）交互，以弥合编码器与解码器之间的语义与分辨率差距？
• RQ3 与最先进基线相比，UCTransNet及其 CTrans 模块是否在多个医疗影像数据集和结构上提供一致的改进？

主要发现

方法	GlaS Dice (%)	GlaS IoU (%)	MoNuSeg Dice (%)	MoNuSeg IoU (%)	Synapse Dice (%)	Synapse HD (mm)
U-Net	85.45	74.78	76.45	62.86	-	-
UNet++	87.56	79.13	77.01	63.04	-	-
AttUNet	88.80	80.69	76.67	63.47	-	-
MRUNet	88.73	80.89	78.22	64.83	-	-
TransUNet	88.40	80.40	78.53	65.05	-	-
Swin-Unet	89.58	82.06	77.69	63.77	-	-
Ours (UCTransNet w/o CCA)	78.99	30.29	78.23	26.75	-	-
Ours (UCTransNet)	90.18	82.96	79.08	65.50	-	-

• 在 GlaS 和 MoNuSeg 上，UCTransNet 相对于基线取得显著改善（Dice 与 IoU 提升），在 Synapse 上也报告了 Dice 和 Hausdorff Distance 的改进。
• 消融结果显示 Baseline+CCT+CCA 在各数据集上持续优于 Baseline、Baseline+CCT 和 Baseline+CCA。
• 增加 CCT 中跳跃尺度输入的数量可提升性能，表明多尺度特征融合的价值。
• 跨注意力可视化显示了哪些编码器层对分割贡献最大，与跳跃连接分析结果一致。
• 对 UCTransNet 进行预训练可进一步提升在 MoNuSeg 与 Synapse 上的收敛速度和最终性能。