Research map / Architectures and efficiency / Mixture of experts
Rethinking U-Net: Task-Adaptive Mixture of Skip Connections for Enhanced Medical Image Segmentation
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes TA-MoSC, a sparse mixture-of-experts module that replaces U-Net's one-to-one skip connections with routed, task-adaptive combinations of multi-stage encoder features, plus a Balanced Expert Utilization strategy.
Problem
U-Net skip connections are fixed one-to-one, but decoder stages prefer different supplementary encoder information depending on the task/dataset, limiting flexibility and generalization in medical image segmentation.
Approach
Reinterpret skip connections as a task allocation problem: shared lightweight convolutional experts are combined by per-stage top-K routers over an aggregated multi-scale encoder feature, with an expert variance loss and unused-expert handling to keep all experts trained.
Key results
- UTANet reports 92.00% Dice on GlaS vs 91.14% for R34UNet and 85.45% for U-Net under 5-fold cross-validation.
- On Synapse, UTANet achieves 88.23% Dice and 28.13 HD95, versus 87.26% Dice for the next best R34UNet.
- On MoNuSeg, UTANet reaches 79.10% Dice, a 2.65-point gain over U-Net's 76.45%, with ablation showing TA-MoSC and BEU each add gains.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace U-Net skip connections with TA-MoSC and add BEU (ablation) | U-Net · GlaS | Dice (%): 85.45±1.26 -> 92.00±0.30 | Helped |
| Replace U-Net skip connections with TA-MoSC and add BEU (ablation) | U-Net · MoNuSeg | Dice (%): 76.45±1.86 -> 79.10±0.49 | Helped |
| Add TA-MoSC to U-Net (ablation) | U-Net · GlaS | Dice (%): 85.45±1.26 -> 89.77±2.59 | Helped |
| Add expert variance loss to TA-MoSC (ablation) | UTANet · MoNuSeg | Dice (%): 77.81±0.98 -> 78.61±1.05 | Helped |
| Add unused experts handling to TA-MoSC (ablation) | UTANet · MoNuSeg | Dice (%): 77.81±0.98 -> 78.78±1.01 | Helped |
| Use TopK=3 at inference instead of TopK=1 (ablation) | UTANet · GlaS | Dice (%): 89.05±0.77 -> 92.00±0.20 | Helped |
| Use TopK=3 at inference instead of TopK=1 (ablation) | UTANet · MoNuSeg | Dice (%): 77.87±0.83 -> 79.10±0.49 | Helped |
| Use TopK=3 at inference instead of TopK=1 (ablation) | UTANet · ISIC16 | Dice (%): 78.89±6.64 -> 91.87±0.07 | Helped |
| Use TopK=4 at inference instead of TopK=1 (ablation) | UTANet · Synapse | Dice (%): 45.41±7.49 -> 88.09±0.92 | Helped |
| Increase inference TopK from 3 to 4 (ablation) | UTANet · Synapse | Dice (%): 86.83±0.65 -> 88.09±0.92 | Helped |
| Increase inference TopK from 3 to 4 (ablation) | UTANet · GlaS | Dice (%): 92.00±0.20 -> 91.18±0.65 | Hurt |
| Increase inference TopK from 3 to 4 (ablation) | UTANet · MoNuSeg | Dice (%): 79.10±0.49 -> 78.90±0.63 | Hurt |
Surprising findings
- The optimal skip-connection order on GlaS (3f or 3421) deviates from the original one-to-one configuration, suggesting decoders prefer information from multiple stages.
Limitations
- Efficiency and FLOPs are higher than U-Net (64.41G vs 50.3G) and the authors defer computational optimization to future work.
- Two-stage training with frozen encoder and decoder complicates the training procedure.
- Evaluated only on four medical datasets; generalization to natural images is claimed but not tested.
Experimental setup
Datasets: GlaS MoNuSeg Synapse ISIC16
Benchmarks: GlaS MoNuSeg Synapse ISIC16
Metrics: Dice IoU HD95
Baselines: U-Net UNet++ Attention U-Net Swin-UNet TransUNet R34UNet UNet-v2 SelfReg-UNet UDTransNet
Builds on: U-Net Mixture of Experts UDTransNet Dice-BCE loss
Closest papers by meaning
- Upping the Game: How 2D U-Net Skip Connections Flip 3D Segmentation (NeurIPS 2024 · similarity 0.90)
- FuseUNet: A Multi-Scale Feature Fusion Method for U-like Networks (ICML 2025 · similarity 0.88)
- U-KAN Makes Strong Backbone for Medical Image Segmentation and Generation (AAAI 2025 · similarity 0.87)
- nnWNet: Rethinking the Use of Transformers in Biomedical Image Segmentation and Calling for a Unified Evaluation Benchmark (CVPR 2025 · similarity 0.86)
- Mamba Goes HoME: Hierarchical Soft Mixture-of-Experts for 3D Medical Image Segmentation (NeurIPS 2025 · similarity 0.86)
- ConDSeg: A General Medical Image Segmentation Framework via Contrast-Driven Feature Enhancement (AAAI 2025 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.