Technical Summary: K-PRISM

Problem Statement

Medical image segmentation is critical for clinical workflows (e.g., tumor delineation, organ quantification), yet current deep learning approaches suffer from fragmentation. Existing models are typically constrained to a single knowledge paradigm:

1. Semantic Priors: Relying on large-scale labeled datasets for specific tasks or organs.
2. In-Context Knowledge: Utilizing few-shot reference examples for generalization.
3. Interactive Feedback: Depending on user inputs (clicks, scribbles) for iterative refinement.

This fragmentation forces healthcare institutions to maintain dozens of task-specific models, leading to high deployment complexity and inconsistent performance. Furthermore, clinical practice naturally integrates these diverse knowledge sources dynamically (e.g., a radiologist using anatomical knowledge, referencing similar cases, and refining boundaries interactively), a flexibility that current AI models lack. No existing framework successfully unifies all three knowledge types within a single architecture.

Methodology

The authors propose K-Prism, a unified segmentation framework designed to integrate semantic priors, in-context exemplars, and interactive feedback into a single model. The core innovation lies in a dual-prompt representation coupled with a Mixture-of-Experts (MoE) decoder.

1. Dual-Prompt Representation

K-Prism encodes heterogeneous knowledge sources into two complementary prompt types:

• 1-D Sparse Prompts: Define what to segment. These encode task-level or instance-specific queries (e.g., class embeddings for semantic segmentation, object queries for in-context, or click coordinates for interactive).
• 2-D Dense Prompts: Indicate where to attend. These modulate spatial feature maps to inject localization cues and refine structural details (e.g., affinity matrices derived from reference image-mask pairs or user click maps).

The framework supports three operational modes:

• Mode-1 (Semantic): Uses only 1-D sparse prompts (learnable class embeddings) to predict segmentation based on learned priors.
• Mode-2 (In-Context): Uses both 1-D sparse prompts (object queries derived from reference masks) and 2-D dense prompts (affinity matrices projecting reference features to the query image) to enable few-shot generalization.
• Mode-3 (Interactive): Uses 2-D dense prompts (encoding positive/negative clicks and previous masks) fused with image features. It can also refine predictions from Mode-1 or Mode-2 by incorporating user clicks as 1-D sparse queries.

2. Mixture-of-Experts (MoE) Decoder

To handle the distinct processing requirements of these different knowledge types, K-Prism employs a MoE decoder.

• Dynamic Routing: A gating network predicts query-specific expert weights, dynamically routing 1-D sparse prompts and 2-D fusion feature maps through specialized experts.
• Bidirectional Interaction: The decoder utilizes bidirectional cross-attention layers where 1-D prompts and 2-D feature maps interact, allowing the model to jointly reason over object semantics and spatial context.
• Architecture: The MoE is applied to both the Cross-Attention (CA) and Feed-Forward Network (FFN) layers, enabling task-aware specialization without architectural modifications when switching between paradigms.

Key Contributions

1. Unified Framework: K-Prism is the first medical image segmentation model to integrate semantic, in-context, and interactive knowledge types within a single architecture.
2. Dual-Prompt Design: The novel combination of 1-D sparse and 2-D dense prompts allows for the flexible encoding of diverse knowledge sources.
3. MoE Decoder: The introduction of a MoE decoder enables dynamic routing and joint training across diverse tasks and modalities, preventing performance degradation from conflicting optimization objectives.
4. Comprehensive Evaluation: The model was validated on 18 public datasets covering diverse modalities (CT, MRI, X-ray, pathology, ultrasound) and segmentation paradigms.

Experimental Results

Experiments demonstrate that K-Prism achieves State-of-the-Art (SOTA) performance across all three paradigms:

• Semantic Segmentation: On 12 in-distribution datasets, K-Prism achieved an average Dice score of 86.21%, outperforming universal models like Hermes (85.02%) and task-specific models like nnU-Net in cross-dataset generalization (83.45% on external datasets).
• In-Context Segmentation: K-Prism achieved an average Dice of 84.82% on in-distribution datasets and 82.49% on external datasets, significantly outperforming prior in-context methods (e.g., Iris, UniverSeg). It showed strong adaptability to unseen classes, achieving 31.91% average Dice compared to 26.07% for the next best method (Iris).
• Interactive Segmentation: K-Prism required the fewest clicks to reach target accuracy (NoC90: 1.95, NoC95: 3.51) on in-distribution datasets. It achieved a 95.50% Dice score after 5 clicks, surpassing baselines like SegNext (93.80%) and MedSAM (91.21%). Crucially, it maintained high efficiency and accuracy even on unseen-class datasets.

Significance and Claims

The paper claims that K-Prism addresses the critical gap between fragmented AI models and the flexible, multi-source nature of clinical practice. By unifying three knowledge paradigms, K-Prism:

• Reduces Deployment Complexity: Eliminates the need for institutions to maintain separate models for different tasks or knowledge types.
• Mimics Clinical Workflows: Enables seamless transitions between semantic, few-shot, and interactive modes, allowing clinicians to leverage available knowledge (e.g., switching from a general anatomical prior to a specific reference case or interactive refinement) without changing the underlying model.
• Serves as a Foundation: Positions itself not just as a segmentation tool but as practical infrastructure for medical foundation models and large-scale annotation pipelines, where clinicians can generate initial masks via priors or in-context learning and refine them with minimal interaction.

The authors acknowledge limitations, including the dependency of in-context performance on reference quality, computational costs from the MoE decoder, and challenges with severe domain shifts (unseen classes). However, they argue that the unified design allows interactive refinement to compensate for failures in other paradigms, offering a robust solution for real-world deployment.