UniMedSeg: Unified In-Context Learning for Multi-Paradigm 2D/3D Medical Image Segmentation

2026-07-14

AI Summary

This work addresses the fragmentation in existing medical image segmentation approaches across prompting paradigms—such as visual exemplars, interactive annotations, and language instructions—and spatial dimensions (2D/3D), which hinders unified modeling and knowledge transfer. To overcome this limitation, we propose UniMedSeg, a Transformer-based universal segmentation framework that, for the first time, enables joint learning across multiple prompting paradigms and dimensionalities. Our method employs a unified sequential representation to map heterogeneous inputs into a shared latent space and introduces a decoupled chunked attention mechanism to efficiently handle long sequences. Evaluated on 27 public datasets, UniMedSeg achieves state-of-the-art performance without task-specific fine-tuning, demonstrating substantially improved generalization capability.

Abstract

Medical image segmentation foundation models are expected to generalize across diverse clinical scenarios, yet existing universal methods remain fragmented by prompt paradigms and spatial dimensions. Visual in-context learning, interactive segmentation, and language-guided segmentation are typically handled by paradigm-specific models, while 2D and 3D images are also modeled separately. Such isolation prevents heterogeneous annotations and data from being jointly absorbed by a single scalable model and limits cross-paradigm knowledge transfer. To address this bottleneck, we propose UniMedSeg, a Transformer-centric universal segmentation framework that maps visual examples, geometric interactions, language instructions, and 2D/3D images into a shared sequence space, enabling heterogeneous medical supervision to be jointly learned through a unified in-context interface without prompt- or dimension-specific branches. To overcome the long-sequence memory bottleneck caused by visual contexts, we introduce Decoupled Split Attention, which reduces attention complexity to linear while preserving hardware-friendly computation and focused context-target interaction. Extensively trained and evaluated on a large corpus curated from 27 public datasets, UniMedSeg achieves state-of-the-art performance across visual in-context, interactive, and language-guided segmentation without task-specific fine-tuning, demonstrating strong generalization on diverse held-out tasks. The code and model weights are publicly available at https://github.com/Lii1228/UniMedSeg

Problem

• medical image segmentation
• in-context learning
• multi-paradigm
• 2D/3D unification
• foundation model

Innovation

• in-context learning
• unified medical segmentation
• 2D/3D image fusion
• decoupled split attention
• foundation model