UTANet: Task-Adaptive Mixture of Skip Connections for Enhanced Medical Image Segmentation

Zichen Luo¹, Xinshan Zhu¹, Lan Zhang¹, Biao Sun¹

¹Tianjin University

• 2025.05.22 We update the TA-MoSC code in the utanet document, a plug-and-play module for U-shaed Architecture!

UTANet is a U-Net-based architecture enhanced with the Task-Adaptive Mixture of Skip Connections (TA-MoSC) module. By integrating the Mixture of Experts (MoE) framework into skip connections, our method dynamically aligns encoder-decoder features, addressing dataset-specific segmentation challenges. Key contributions include:

• TA-MoSC: Adaptive routing mechanism for multi-scale feature fusion.
• Balanced Expert Utilization (BEU): Ensures balanced training of lightweight convolutional experts.
• State-of-the-art Performance: Outperforms existing methods on GlaS, MoNuSeg, Synapse, and ISIC16 datasets.

The following features and improvements are planned or in progress:

• [2025.05.22] Add Model scripts.
• Training framework
• [2025.05.22] Implement TA-MoSC module
• Upload weight
• Add Jupyter notebook examples
• Add data preprocessing scripts
• [2025.01.09] Add README.md.

• Dynamic Skip Connections: Replaces fixed skip connections with task-adaptive expert mixtures.
• Lightweight Design: Minimal parameter overhead compared to vanilla U-Net.
• Multi-Dataset Generalization: Robust performance across diverse medical imaging tasks.
• Code Efficiency: Compatible with PyTorch and easy to integrate into existing U-Net variants.

TA-MoSC learns to select relevant skip features dynamically across different tasks, allowing better specialization in edge cases such as:

• 🔬 Tiny tumor detection
• 🧩 Irregular lesion boundaries
• 🧠 Low-contrast or noisy regions

1. Clone the repository

git clone https://github.com/AshleyLuo001/UTANet.git  
cd UTANet

2. Install dependencies

pip install -r requirements.txt

3. Training

waiting...

waiting...

We welcome contributions! Please see CONTRIBUTING.md for instructions.

This project is licensed under the MIT License - see the LICENSE file for details.

@article{luo2025rethinking,  
  title={Rethinking U-Net: Task-Adaptive Mixture of Skip Connections for Enhanced Medical Image Segmentation},  
  author={Luo, Zichen and Zhu, Xinshan and Zhang, Lan and Sun, Biao},  
  journal={AAAI},  
  year={2025}  
}