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

Abstract: Building a unified model with a single set of parameters to efficiently handle diverse types of medical lesion segmentation has become a crucial objective for AI-assisted diagnosis. Existing unified segmentation approaches typically rely on shared encoders across heterogeneous tasks and modalities, which often leads to feature entanglement, gradient interference, and suboptimal lesion discrimination. In this work, we propose TP-Seg, a task-prototype framework for unified medical lesion segmentation. On one hand, the task-conditioned adapter effectively balances shared and task-specific representations through a dual-path expert structure, enabling adaptive feature extraction across diverse medical imaging modalities and lesion types. On the other hand, the prototype-guided task decoder introduces learnable task prototypes as semantic anchors and employs a cross-attention mechanism to achieve fine-grained modeling of task-specific foreground and background semantics. Without bells and whistles, TP-Seg consistently outperforms specialized, general and unified segmentation methods across 8 different medical lesion segmentation tasks covering multiple imaging modalities, demonstrating strong generalization, scalability and clinical applicability.

• The paper presents a task-prototype framework that disentangles task-general and specific features using a dual-path Task-Conditioned Adapter and a Prototype-Guided Task Decoder.
• It achieves high segmentation accuracy with an average Dice of 86.63% and mIoU of 86.44% across eight diverse medical lesion segmentation benchmarks.
• Ablation studies confirm that both TCA and PGTD are crucial for maintaining semantic specificity, reducing feature entanglement, and ensuring robust multi-task generalization.

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

Introduction and Motivation

Unified medical lesion segmentation (MLS) seeks to develop a single model capable of addressing diverse lesion types and imaging modalities, circumventing the inefficiency and redundancy of training and deploying isolated task-specific models. Recent unified approaches often use shared encoders and task-agnostic decoders, which are prone to feature entanglement, gradient interference, and poor discrimination of lesion boundaries due to modality and anatomical heterogeneity. The "TP-Seg: Task-Prototype Framework for Unified Medical Lesion Segmentation" (2604.00684) proposes a framework that explicitly disentangles task-general and task-specific information in feature extraction and introduces learnable task prototypes to guide decoder semantics, addressing the limitations of existing unified segmentation models.

Figure 1: TP-Seg’s dual-path encoder and prototype-guided decoder achieve superior average Dice (86.63%) and mIoU (86.44%) over previous state-of-the-art unified models.

Framework Overview

TP-Seg integrates two principal innovations: a Task-Conditioned Adapter (TCA) for balancing shared and task-specific representations, and a Prototype-Guided Task Decoder (PGTD) which leverages learnable prototypes as semantic anchors for foreground–background differentiation. The overall pipeline is illustrated in Figure 2.

Figure 2: TP-Seg architecture: the input image and task embedding are processed by TCRB-TCA blocks before backbone encoding; PGTD decodes with prototype-based cross-attention for lesion prediction.

Task-Conditioned Routing Block (TCRB) with Dual-Path Adapter

TP-Seg introduces TCRBs, each containing both a shared router extracting task-general features and multiple task-specific routers for capturing unique task-specific increments. The fusion of these branches is learned with a selective gate, and a cumulative sigmoid-based split gate vector $L_\tau$ for each task establishes a dynamic “shared-first, task-specific-later” separation policy, which is differentiable and adaptively determined during end-to-end training, thereby optimizing the trade-off between generalization and task specificity.

Figure 3: The TCRB integrates shared and task-specific routers per adapter, dynamically adapting features pre-encoder via a learnable split gate.

Prototype-Guided Task Decoder (PGTD)

The PGTD targets decoder-stage feature collapse due to task-agnostic processing by introducing foreground (FG) and background (BG) task prototypes, which are trainable and maintained via exponential moving average updates. These prototypes act as persistent semantic centroids, interacting with spatial features via cross-attention and informing a prototype-guided similarity map. The map modulates features prior to convolution by a dynamic ensemble of expert heads re-weighted according to the semantic descriptor extracted over prototype responses. This design ensures robust modeling of task-aware FG–BG separation, adapting to anatomical and morphological variability inherent across tasks.

Figure 4: The PGTD module fuses encoder features, attends with learned FG/BG prototypes, and modulates decoding via similarity maps.

Experimental Evaluation

The effectiveness and scalability of TP-Seg are validated on 8 MLS benchmarks covering ophthalmology, neurology, pathology, thyroid, endoscopy, thoracic CT, breast ultrasound, and dermatology modalities. Evaluation is performed with Dice and mIoU metrics, comparing task-specialized, general-purpose, and unified models.

Quantitative Results

TP-Seg-UNIFIED achieves top Dice (86.63%) and mIoU (86.44%) scores, exceeding previous unified baselines (e.g., SR-ICL, SAM2-UNet, SegGPT) by 1–3 percentage points on average over the task suite. Notably, the framework's unified model does not underperform versus task-wise training; in some categories, unified training shows improved generalization, highlighting effective aggregation and implicit regularization. The model demonstrates strong boundary delineation, sensitivity to lesion scale, and robustness across all modalities.

Qualitative Results

Figure 5: Visual comparison indicates TP-Seg’s predictions are consistently more precise, with sharper lesion boundaries and better small-object sensitivity relative to Spider, SegGPT, SR-ICL, and SAM2-UNet.

Ablation and Mechanistic Analysis

Ablation experiments demonstrate both TCA and PGTD are essential. Exclusion of either module leads to significant degradation (up to 8% Dice drop). TCA’s dual-route structure outperforms classical LoRA and adapter-based routes, indicating the necessity of independently learnable shared and task-specific feature flows within high-capacity encoders.

Figure 6: Removing TCA or PGTD components results in blurred or inconsistent lesion segmentation, confirming their critical roles.

Further analysis of prototype space shows that learned FG prototypes across tasks are well-separated—for all PGTD modules, inter-task cosine similarity remains below 0.9 (i.e., avoids task collapse) and maintains semantic specificity even for anatomically related tasks (Figure 7).

Figure 7: Foreground prototype inter-task similarities remain far from collapse, validating task discriminability.

Foreground–background prototype separation is well-maintained (all scores >0.85), and separation degree demonstrably correlates with better segmentation (Figure 8).

Figure 8: Strong FG-BG prototype separation is positively correlated with improved Dice and mIoU scores across tasks.

Parameter growth per task is negligible: scaling from 1 to 8 tasks increases parameters by ~8.9%, confirming the efficiency and scalability of the architecture.

Practical and Theoretical Implications

TP-Seg establishes a compelling paradigm for unified MLS by demonstrating that incorporating task-adaptivity at both feature extraction and semantic decoding stages eliminates key bottlenecks in cross-task learning, namely representation entanglement and semantic confusion. The design is robust to hyperparameters and compatible with foundation model backbones (e.g., SAM2, Hiera-L), facilitating downstream clinical deployment. Prototype-driven FG–BG guidance may also inform further research in segmentation under multi-domain and few-shot regimes, where explicit task semantics and class anchors are paramount. This framework may generalize to other multi-task biomedical contexts, e.g., organ segmentation, disease grading, or population-level image analysis.

Conclusion

TP-Seg, via its dual-path TCA and prototype-guided PGTD, achieves state-of-the-art performance on unified MLS benchmarks, outperforming both established unified and task-specialized segmentation networks. Its strong empirical results validate the necessity of architectural explicitness in both feature routing and semantic conditioning for unified medical image analysis. The approach presents a strong foundation for generalized, scalable, and semantically robust segmentation in future clinical AI systems.

---

Reference:

“TP-Seg: Task-Prototype Framework for Unified Medical Lesion Segmentation” (2604.00684)