TP-Seg: Task-Prototype Framework for Unified Medical Lesion Segmentation

1. Problem Statement

Medical lesion segmentation (MLS) is critical for clinical diagnosis but faces significant challenges due to the diversity of lesion types (e.g., tumors, polyps, infections) and imaging modalities (e.g., MRI, CT, OCT, Ultrasound).

• Current Limitations:
  • Specialized Models: Training separate models for each task leads to parameter redundancy, high deployment costs, and limited knowledge transfer.
  • Existing Unified Models: Current approaches that use a single model for multiple tasks typically rely on shared encoders and task-agnostic decoders. This causes:
    • Feature Entanglement: Heterogeneous tasks interfere with each other, leading to conflicting gradients.
    • Modality Conflicts: Different imaging modalities are forced into a single representation space, degrading performance.
    • Semantic Confusion: Without explicit task awareness, decoders struggle to distinguish specific foreground-background relationships, resulting in boundary inaccuracies and inconsistent performance.

2. Methodology: TP-Seg Framework

The authors propose TP-Seg, a unified framework built upon the SAM 2 (Segment Anything Model 2) backbone. It introduces two core components to balance shared knowledge with task-specific adaptation:

A. Task-Conditioned Adapter (TCA) & Routing Block (TCRB)

Located before the frozen transformer blocks of the encoder, the TCA addresses feature interference through a dual-path expert structure:

• Dual-Path Design: It consists of a Shared Router (extracting generic visual priors) and Task-Specific Routers (capturing modality/lesion-specific features).
• Learnable Split Gate: A learnable vector dynamically determines the optimal depth ($b^*_\tau$) at which the network switches from shared processing to task-specific processing.
  • Early layers: Focus on shared features (high-level generalization).
  • Later layers: Focus on task-specific details.
• Mechanism: This "shared-first, specific-later" routing is learned end-to-end using a cumulative sigmoid function with a temperature decay, allowing the model to adaptively allocate representational capacity without manual tuning.

B. Prototype-Guided Task Decoder (PGTD)

To overcome the limitations of task-agnostic decoders, TP-Seg introduces learnable task prototypes as persistent semantic anchors:

• Prototype Initialization: Each task $\tau$ maintains a pair of foreground ($P^{fg}_\tau$) and background ($P^{bg}_\tau$) prototypes, initialized via an MLP from the task embedding.
• Cross-Attention Interaction: Prototypes interact with spatial encoder features via cross-attention to generate discriminative semantic descriptors ($Z_\tau$).
• Dynamic Feature Modulation:
  • Similarity Maps: A prototype-guided similarity map ($S_\tau$) is computed to highlight regions matching the foreground prototype.
  • Dynamic Experts: The decoder generates dynamic convolutional experts weighted by the task, allowing adaptive processing of features.
• Prototype Update: Prototypes are updated during training using an Exponential Moving Average (EMA) of the mean features from ground-truth regions. This ensures the prototypes evolve to capture the true distribution of foreground/background semantics over time.
• Reinforcement: The final prediction combines the expert output with a prototype-guided reinforcement term to sharpen boundaries.

3. Key Contributions

1. Unified Framework: A simple yet effective architecture that unifies 8 diverse medical lesion segmentation tasks into a single model with a single set of parameters.
2. Dual-Path Routing (TCA): A parameter-efficient module that explicitly models shared and task-specific representations, mitigating gradient interference and modality conflicts.
3. Prototype-Guided Decoding (PGTD): A novel decoder mechanism using learnable task prototypes as semantic memory, enabling fine-grained modeling of foreground-background relationships across tasks.
4. State-of-the-Art Performance: Demonstrated superior generalization and scalability across multiple imaging modalities without complex bells and whistles.

4. Experimental Results

The framework was evaluated on 8 medical lesion segmentation tasks covering 5 imaging modalities (OCT, MRI, Pathology, Ultrasound, Endoscopy, CT, Dermoscopy).

• Quantitative Performance:
  • TP-Seg-Unified achieved an average Dice score of 86.63% and mIoU of 86.44%.
  • It outperformed all existing unified baselines (e.g., SegGPT, Spider, SAM2-UNet, SR-ICL) and specialized models trained independently.
  • Notably, the unified model performed comparably to or better than models trained separately for each task (TP-Seg-General), proving effective knowledge transfer.
• Qualitative Analysis:
  • Visual results show TP-Seg produces smoother, more accurate boundaries, especially for small, complex, or occluded lesions (e.g., thyroid nodules, breast lesions, polyps).
  • It effectively suppresses background noise and handles diverse morphologies better than competitors.
• Ablation Studies:
  • Removing TCA or PGTD resulted in significant performance drops, confirming the necessity of both modules.
  • The dual-path design (shared + specific routes) was proven superior to standard LoRA or Adapter fine-tuning.
  • Cosine similarity analysis confirmed that the model successfully learns discriminative prototypes for different tasks, avoiding "task collapse."
• Efficiency:
  • Extending the model from 1 task to 8 tasks incurs a negligible parameter increase (~1.27% growth), demonstrating high scalability.

5. Significance

TP-Seg represents a significant step forward in Unified Medical AI. By solving the "feature entanglement" and "semantic confusion" problems inherent in multi-task learning, it offers a practical solution for clinical deployment where:

• Efficiency: A single model can replace multiple specialized models, reducing storage and computational costs.
• Generalization: The model leverages data from diverse modalities to improve robustness, particularly for tasks with limited data (acting as an implicit regularizer).
• Clinical Applicability: The ability to handle diverse lesion types and imaging modalities with high precision makes it a viable candidate for comprehensive computer-aided diagnosis systems.

The authors have made the code and pretrained models available, facilitating further research into generalized medical imaging frameworks.