UCTransNet: Rethinking the Skip Connections in U-Net from a Channel-wise Perspective with Transformer (2109.04335v3)

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.

• The paper introduces a novel Channel Transformer module that replaces traditional skip connections to enhance feature fusion in U-Net architectures.
• It integrates Channel-wise Cross Fusion and Cross Attention to effectively bridge semantic gaps and capture multi-scale contextual information.
• Experimental results show Dice score improvements of up to 9% over U-Net on multiple medical imaging datasets, demonstrating practical segmentation gains.

UCTransNet: Rethinking Skip Connections in U-Net with Transformers

This paper introduces UCTransNet, a novel medical image segmentation model that leverages Transformer architecture to enhance the traditional U-Net framework. The paper addresses inherent limitations in the U-Net's skip connections, emphasizing the challenges in modeling global multi-scale context without exacerbating semantic gaps.

Key Contributions

1. Revised Skip Connections: UCTransNet replaces conventional skip connections with a Channel Transformer (CTrans) module, redefining the connectivity between encoder and decoder stages. The CTrans consists of Channel-wise Cross Fusion Transformer (CCT) and Channel-wise Cross Attention (CCA) for effective multi-scale feature fusion.
2. Channel-Wise Perspective: The CCT module captures local cross-channel interactions, while the CCA module integrates these transformed features with decoder outputs. This approach aims to bridge semantic gaps and optimize feature reconciliation across different network scales.
3. Empirical Validation: Experimental results demonstrate notable segmentation enhancements over state-of-the-art methods across various datasets. UCTransNet achieved absolute improvements of 4.05%, 7.98%, and 9.00% Dice scores over U-Net on GlaS, MoNuSeg, and Synapse datasets, respectively.

Methodological Insights

The authors propose a CCT module that employs multi-scale feature embedding augmented by multi-head channel-wise cross attention, differing from traditional Transformers by focusing attention along the channel axis. This strategy capitalizes on channel-wise semantics which are critical for medical images characterized by complex structures.

CCA further refines feature integration by recalibrating channel importance via an attention mechanism, addressing inconsistencies between encoder and decoder feature semantics. This dual strategy forms a robust framework that implicitly tackles the traditional U-Net's simplistic skip connection model.

Implications and Future Work

UCTransNet's framework suggests a paradigm shift in medical image segmentation, indicating the value of integrating channel-wise attention mechanisms within Transformer-enhanced architectures. This model's success exemplifies the importance of detailed multi-scale feature integration and offers pathways for incorporating similar mechanisms in other deep learning models.

Future research could explore extending the CTrans module to other segmentation architectures, evaluating its utility in diverse image domains beyond medical imaging. Further investigation is warranted to compare computational trade-offs in different Transformer configurations and optimize them for real-world clinical applicability.

In conclusion, UCTransNet presents a promising advancement in medical image segmentation, demonstrating the efficacy of Transformer-based enhancements in overcoming the longstanding limitations of U-Net.