K-Prism: A Knowledge-Guided and Prompt Integrated Universal Medical Image Segmentation Model
Authors: Bangwei Guo
Yunhe Gao
Meng Ye
Difei Gu
Yang Zhou
Leon Axel
Dimitris N. Metaxas
Venue: ICLR 2026 Poster (2026)
Keywords: Medical Image
Image Segmentation
Universal Model
Prompt Integration
Primary area: applications to computer vision, audio, language, and other modalities
TL;DR: We propose K-Prism, a unified segmentation framework that integrates semantic priors, in-context examples, and interactive feedback into a dual-prompt MoE decoder, achieving state-of-the-art performance across 18 diverse medical imaging datasets.

Abstract: Medical image segmentation is fundamental to clinical decision-making, yet existing models remain fragmented. They are usually trained on single knowledge sources and specific to individual tasks, modalities, or organs. This fragmentation contrasts sharply with clinical practice, where experts seamlessly integrate diverse knowledge: anatomical priors from training, exemplar-based reasoning from reference cases, and iterative refinement through real-time interaction. We present $\textbf{K-Prism}$, a unified segmentation framework that mirrors this clinical flexibility by systematically integrating three knowledge paradigms: (i) $\textit{semantic priors}$ learned from annotated datasets, (ii) $\textit{in-context knowledge}$ from few-shot reference examples, and (iii) $\textit{interactive feedback}$ from user inputs like clicks or scribbles. Our key insight is that these heterogeneous knowledge sources can be encoded into a dual-prompt representation: 1-D sparse prompts defining $\textit{what}$ to segment and 2-D dense prompts indicating $\textit{where}$ to attend, which are then dynamically routed through a Mixture-of-Experts (MoE) decoder. This design enables flexible switching between paradigms and joint training across diverse tasks without architectural modifications. Comprehensive experiments on 18 public datasets spanning diverse modalities (CT, MRI, X-ray, pathology, ultrasound, etc.) demonstrate that K-Prism achieves state-of-the-art performance across semantic, in-context, and interactive segmentation settings. Code is available at https://github.com/bangwayne/K-Prism.

Published as a conference paper at ICLR 2026
K-PRISM: A KNOWLEDGE-GUIDED AND PROMPT IN- TEGRATED UNIVERSAL MEDICAL IMAGE SEGMENTA- TION MODEL
Bangwei Guo', Yunhe Gao", Meng Ye², Difei Gu', Yang Zhou', Leon Axel', Dimitris Metaxas "Rttger University *Stanford University 'The University of Texas Arlington
"New York University
bg65 66 rutgers.. edu, dnmêcs. rew.ers.utrrrreeeec.
ABSTRACT
Medical image segmentation fundamental to clinical decision decision-making, yet ex- isting models remain fragmented. They are usually trained on single knowledge sources and specific to individual tasks, modalities, or organs. This fragmenta- tion contrasts sharply with clinical practice, where experts seamlessly integrate diverse knowledge: anatomical priors from training, exemplar-bases reasoning from reference cases, and iterative refinement through real-time interaction. We present K-Prism, a unified segmentation framework that mirrors this clinical flex- ibility by systematically integrating three knowledge paradigms: (i) semantic pri- ors learned from annotated datasets, (ii) in-context knowledge from few-shot ref- erence examples, and (iii) interactive feedbuck from user inputs like clicks scribbles. Our key insight is that these heterogeneous knowledge sources can be encoded into a dual-prompt representation: 1-D sparse prompts defining what to segment and 2-D dense prompts indicating where to attend, which are then dynamically routed through Mixture-of- of-eertttss (MoE) decoder. This design enables flexible switching between paradigms and joint training across diverse tasks without architectural modifications. Comprehensive experiments on 18 pub- lic datasets spanning diverse modalities (CT, MRI, X-ray, pathology, ultrasound, etc.) demonstrate that K-Prism achieves state-of-the-an art performance across se- mantic, in-context, and interactive segmentation settings. Code is available at https: /t/t/tttttt /ttthub.t.u/ttt/ com haaam//aagganneeeee eaeiinnne//-iinn-e-
INTRODUCTION
Medical image segmentation is a cornerstone of modern clinical workflows, supporting tasks such tumor delineation (Heller et al., 2019; Bilic al., 2023), organ quantification (Wasserthal et al., 2023), and vessel segmentation (Livne et al., 2019). While deep learning achieves strong results on individual benchmarks (Isensee et al., 2021; Hatamizadeh al., 2022), real-world deployment remains challenging: healthcare institutions must maintain dozens of task-specific models tailored different organs, modalities, and clinical scenarios, resulting in high deployment complexity and inconsistent performance (Zhou et al., 2021). This fragmentation stems from deeper limitation: most existing models are constrained to single knowledge type. They either depend on semantic priors learned from large labeled datasets (Liu et al., 2023a; Gao, 2024), adapt through in-context knowledge with few-shot examples (Butoi al., 2023), or rely on interactive feedback (Ma al., 2024; Isensee et al., 2025). Yet clinical practice is rarely confined to single knowledge paradigm.
Consider a radiologist examining rare pediatric tumor: they may leverage their semantic knowl- edge of anatomy, retrieve similar historical cases for reference, and iteratively refine boundaries through interactive feedback. Such flexible integration of diverse knowledge is routine for human experts, who dynamically adapt their strategies to the different clinical settings. In contrast, current Al models remain rigid, unable to seamlessly combine multiple knowledge sources. This forces clinicians to switch between separate models, disrupting workflows and limiting the potential of AI assistance. Recent efforts have attempted partial unification, combining two paradigms at a time (Gao et al., 2025; Guo al., 2025; Wong et al., 2024), but framework yet integrates all

---

Published as a conference paper at ICLR 2026
Unified Segmentation
K-Prism
Semantic prior knowledge Annotated training data Asalony pallem
In-contest knowledge
Internctive feedhack knowledge
Few laboled examples
Expert feedback
Reference imago-mask pairs
Postiveengallive clicks
Figure I: K-Prism integrates three forms of external knowledge, semantic priors (from annotated training datasets), in-context exemplars (from reference image-mase mask pairs), and interactive feedback (from user clicks and previous masks) into single framework, enabling robust segmentation across diverse modalities and targets.
three knowledge types within a single architecture, see Table 1. This gap persists due to a key tech- nical challenge: how to represent and process fundamentally different forms of knowledge within unified framework that delivers strong performance across all modes.
We propose that achieving true universality in medical segmentation requires simultancously ad- dressing three clinically essential knowledge sources: (1) semantic prior knowledge from large-scale annotated datasets, capturing common anatomical and modality patterns, (2) in-context knowledge from reference examples, critical for rare diseases new protocols where labeled data is scarce, (3) interactive feedback knowledge from user interactions such as clicks or scribbles, enabling iterative refinement. To realize this vision, we propose knowledge-guiee medical image segmentation model, K-Prism, which can integrate all three knowledge forms and adapt seamlessly across diverse inference scenarios and imaging domains (Figure 1).
and prompt-integrated universal
To support joint training and inference across heterogeneous tasks and knowledge paradigms, design novel dual-prompt representation coupled with a Mixture-of -fxture-of-Experts (MoE) decoder. Our key insight is that diverse knowledge sources can be encoded into two complementary prompt types: 1-D sparse prompts to encode what to segment and 2-D dense prompts to encode where to attend, which are then dynamically routed through specialized experts based the task requirements. This design enables K-Prism to not only match state-of-the- paradigm but also support fluid transitions between modes, mirroring clinical workflows where different knowledge sources are combined based on availability and task demands. Beyond perfor- mance, K-Prism reduces deployment complexity by unifying all three paradigms, positioning itself as both robust segmentation model and practical infrastructure for medical foundation models.
(SOTA) performance
each individ-
In summary, our key contributions are:
We propose K-Prism, unified framework and practical foundation for medical image segmenta- tion that integrates three clinically relevant knowledge types.
We design dual- mool-p representation that combines 1-D sparse and 2-D dense prompts with MoE decoder for dynamic routing, enabling joint training across diverse tasks and modalities.
We conduct comprehensive experiments on 18 datasets covering diverse anatomical targets, imag- ing modalities (CT, MRI. X-ray, pathology, etc.), and segmentation paradigms (semantic, in- context, interactive), achieving SOTA performance and strong cross-dataset generalization.
2

---

Published as a conference paper at ICLR 2026
2
RELATED WORK
2.1 MEDICAL IMAGE SEGMENTATION UNDER DIVERSE KNOWLEDGE
Medical image segmentation has conventionally relied task-specific, fully supervised models trained on well-annotated single-modality datasets (Isensee et al., 2021; Hatamizadeh et al., 2022). While effective in constrained settings, these models often fail to generalize to real-world scenarios where the available knowledge varies across cases (Zhou et al., 2022). To improve adaptability, recent studies have explored knowledge-driee ledeeeddivvien strategies across diverse segmentation paradigms. In semantic segmentation, Liu et al. (2023a) leverage CLIP to encode anatomical descriptions into semantic priors, while Hermes (Gao, 2024) introduces learnable task-specific embeddings to guide universal model. In in-context segmentation, Universeg (Butoi et al., 2023) and Tyche (Rakic et al., 2024) employ annotated support exemplars as visual references, enabling few-shot generalization across diverse tasks and datasets. In interactive segmentation, McdSAM (Ma et al., 2024) and nnInteractive (Isensee et al., 2025) refine model predictions by incorporating user prompts such as clicks and bounding boxes. Despite these advances, most existing methods are restricted to single segmentation paradigm, limiting their ability to integrate heterogeneous knowledge sources.
2.2 UNIFIED SEGMENTATION FRAMEWORKS AND GENERALIZATION
Recent work has explored unified med- ical segmentation frameworks capable of integrating multiple tasks and knowl- edge types. As shown in Table I, Iris (Gao al., 2025) encodes refer- ence image-label label pairs into 1-D to- kens, supporting both semantic and in-context segmentation. Verse (Guo al., 2025) unifies semantic and inter- active segmentation, enabling iterative refinement of initial predictions, while MultiverSeg (Wong et al., 2024) com- bines in-context and interactive seg- mentation to learn from a small num- ber of annotated examples and improve results through expert feedback. How- ever, no existing medical image seg- mentation framework integrates seman- tic, in-context, and interactive feedback knowledge within single architecture, leaving a gap for models that can adapt across all three paradigms.
Table Comparison of representative medical image seg- mentation methods across different paradigms.
Method
Semantic
In-context
Interactive
naU-Net (Isensee et 2021)
UNETR (llatamizadeh et al., 2022) Clip-driven Hermes (Gao, 2024)
2023a)
UniverSeg (Butoi et al., 2023)
Tyyhe (Rakic al., 2024)
MedSAM (Ma al., 2024)
maInteractive (Isensce et al., 2025)
MultiverSeg (Wong et al., 2024)
Iris (Gao al., 2025)
Verse (Gun et al., 2025)
Ours (K-Prism)
METHOD
3.1 PROBLEM DEFINITION
Conventional medical image segmentation methods follow a knowledge-specie paradigm, where a dedicated model fo, is trained for specific segmentation task relying solely single knowledge type such as semantic priors or in-context knowledge. K-Prism moves beyond this limitation by jointly integrating three complementary clinical knowledge sources, each aligned with specific operational mode within unified architecture:
Mode-1: Semantic segmentation leverages learned class-level priors. Given learnable em- bedding matrix P E R (xx()) and input image where Nes is the number of classes in the training set, and each entry P.. E C encodes semantic knowledge for class n E (1...., through p query vectors of dimension C, the model predicts Ý, = fe [I] P.)) Mode-2: In-context segmentation uses reference examples to guide segmentation. Given a support set S (re, M containing reference images and masks, the model predicts
Y, fo 8).

---

fyeon fratur rop Published as a conference paper at ICLR 2026
Ouepee Isere reap
FFN
Mel: Decoder
Add No.s
Moll Cruss Amenicon
Intenztive
10
Asgespaio
Seftras
quea querios
MoFsFFN
Prompt Fision Module
Seif-Attemion
Imsaee Enciderr
Musk Ensoder
Interactive Prompt Eneedeer
Mol Cross Ameniionn
Irpar image
image-made
Festure map
Quance
Poskiveiasative clica
(a)
(b)
Figure 2: (a) Overview of the proposed K-Prism framework. Our model integrates three forms of external knowledge via the prompt fusion modules, encoding them into I-D sparse queries and 2-D dense prompts produce fusion feature maps. (b) The MoE decoder dynamically routes different prompts to specialized experts through cross-attentiin and gating, enabling task-aware specialization and robust segmentation across diverse scenarios.
Mode-3: Interactive segmentation incorporates user feedback through clicks, enabling iterative refinement. Given click set C = N. the model predicts Y = fI C). This mode can also refine the initial predictions from Mode-1 and Mode-2.
In the following sections, we deseribe how K-Prism's 's unified architecture realizes each mode and seamlessly integrates multiple knowledge sources to achieve efficient, accurate, and broadly gener- alizable medical image segmentation.
3.2 INTEGRATING DIVERSE K KNOWLEDGE INTO UNIFIED PROMPTS
Figure 2(a) shows the K-Prism framework. To process different knowledge types within a single architecture, we propose dual-prompt design that captures two complementary aspects segmen- tation by converting the inputs from each operational mode into a unified prompt representation: (i) I-D sparse prompts, which define what the model should focus on by encoding task- or instance-specie specific queries, and (ii) 2-D dense prompts, which indicate where the model should attend by modulating spatial feature maps to inject localization cues and refine structural details. Given an input image HW we first extract its feature map F through the encoder:
= Encoder(I), E Cxhxw
(I)
where C the feature dimension and h, w are the spatial resolutions of the encoded features.
3.2.1 SEMANTIC SEGMENTATION
In this mode, we use only 1-D sparse prompts, since semantic segmentation task relies solely high-level class knowledge. The semantic prior embedding matrix P E Xxxx)) is a learnable parameter, optimized through gradient backpropagation during training to learn class knowledge. To segment class n, the corresponding query set p.. is combined with the feature map F and fed into the decoder, without requiring any prompt fusion:
Output, = Decoder( F pm))
(2)
3.2.2 IN-CONTEXT SEGMENTATION
This mode requires both 1-D sparse prompts and 2-D dense prompts. Let the support set be S {Ire Mer)) where I E R Nxxx3x denotes the reference images and M E R Nax Xx denotes the corresponding segmentation masks.
4

---

Published as a conference paper at ICLR 2026
Given query image Ig and the reference images Let, the image encoder first extracts query and key features as shown in Figure 2(a):
Fg Encoder(f,) (T),, F2 Encoder( Teet;
(3)
and the reference image-mask pairs are encoded into value features via a lightweight mask encoder:
Fe MaskEncoder
Mel))
(4)
For 2-D dense prompts, these encoded features are flattened and utilized in the prompt fusion:
kf (Fa) E RCx v Flaa((F)) C
K Flat((F)) Rx
(5)
In our design, Kt and K" are extracted by the same encoder, placing them in shared feature space where the similarity matrix can be computed to align query and reference images. The value features Vief derived from reference image-mask pairs, are then projected via this similarity matrix into the query-aligned space, transferring mask semantics to the query image representation. Both Kt and V thus serve as 2-D dense prompts. Formally, for any similarity function e: gC x XBC - R, we compute the softmax-normalize affinity matrix W by:
exp(A..) Wi Σn, (n.))
A.j c K.))
A, W R R-Nadwukkw
(6)
We use negative squared Euclidean distance for c(..)) following Cheng et al. (2021; 2024). The affinity matrix W projects the value features V to align with the query features, producing the fusion feature map Fiuee = Ve Fiuse E gg which then passed to the decoder.
For the 1-D sparse prompt, given the encoded reference features K construct set of n, object queries Q E IR=,,C XC where each query encodes compact, high-level object information. The first half represents the foreground, while the second half represents the background. pooling mask M E [0, 1] Narww obtained by downsampling and flattening M and the n-th object query at location then derived via masked average pooling:
n 3 and M 0.5,
Nartow (K TT m Lim1
m
0,
n > Re and M.odl 2 0.5, Qº 2
0 E ERC
(7)
Eis1 m
otherwise.
In summary, Q" provides a 1-D sparse prompt encoding object-level foreground and background information, while K and Vrel act as 2-D dense prompts for spatial modulation. design enables the MoE decoder in the in-context setting to reason jointly over object semantics and spatial context. Further details and settings are provided Appendix.
This dual-prompt
3.2.3 INTERACTIVE SEGMENTATION
For interactive segmentation, given an image and user click set C channel prompt Lelick E x encoding positive clicks, negative clicks, and the previous mask prediction (Sofiukk et al., 2022; Liu et al., 2023b). This prompt serves as 2-D dense signal and is additively fused with the image features.
N: { we construct three-
PromptEnccoeer(anit (Teni))) + Encoder(T), Mode-3 or refine Mode-1
F Fruse:
Fluse E Cxbu
PrompLEncoder (Teic)) + Vid W,
Refine Mode-2
Each click also encoded as a I-D sparse query. Following Guo et al. (2025), a click at image coordinates (x, y) is mapped to feature map coordinates (x',y') = [[z/s], Lw/s]), with denot- ing the downsampling ratio. Features within a local window of size 2r+1 centered at (a', s/) are average-pooled linearly transformed, and combined with SAM-style positional embedding (Kir- illov et al.. 2023) to form the sparse query for that click. The N₂ clicks together form the query set Qo NCC. which is then fed into the decoder produce the segmentation.
5

---

Published as a conference paper at ICLR 2026
Decoder( Fruse Prc: Q"), Refine Mode-1 1
Output
Decoder( Fiuse 1", QQ):;
Refine Mode-2
(9)
Decoder( Fruse
Qa,,
Mode-3
3.3 MIXTURE-OF-EXPERTS
DECODER
The unified nature of K-Prism makes it challenging for single decoder to effectively process fun- damentally different knowledge types, as each requires distinct processing strategies. A standard decoder would struggle to optimize across all three paradigms simultaneously, potentially leading to suboptimal performance. To address this, we introduce a Mixture-of-Expeee (MoE) decoder that enables dynamic, task-aware specialization while maintaining shared representational capacity. Our decoder adopts bidirectional
design to effectively fuse 1-D sparse prompt and 2-1) fusion feature maps (Figure 2(b)). In the first layer, 1-D sparse prompts projected into queries Q E _axC_ while 2-D fusion feature maps are projected into keys and values K, V E ghwxC. MoE is applied both cross-attention (CA) and the feed-forward network (FFN). And each MoE- CA layer consists of M multi-head attention experts {Am mm= where each computes the cross- attention: Om = Am (Q+ Pq. K + Pk, V), Om E RTxC, with Pq and Px denoting positional embeddings. Pross-attention is given by:
QK VC
CrossAltn Q,K, softmax
V
(10)
To adaptively combine the outputs of multiple experts, a gating network G RRxC + R
predicts query-specific expert weights: Softmax(G(()), are stacked as O Stckk E R and the gating weights are broadcast to a E Rxx The final MoE output is then obtained via element-wise multiplication and summation:
E RTxM The expert outputs
O moe
b mm,
stack 1,m,3
O mose RAXC
(11)
We then apply residual addition and layer normalization to obtain the updated query representation Q', which is then passed to Self- Attention and MoE-FFN blocks to further update. In the second MoE cross-attention layer, the updated 1-D queries are projected into keys and values, while the 2-D fusion feature maps serve as queries, establishing bidirectional interaction between the sparse queries and fusion feature maps. More details of the decoder are provided Appendix.
EXPERIMENT
4.1 EXPERIMENTAL SETUP
Datasets. We train 12 publicly available datasets et al., 2022: Campello al., 2021; Bilic
et al., 2023; Heller et al., 2019; Li et al., 2020; 2021; 2022a;b; Al-Dhabyani et al., 2020; Jaeger al., 2013; Candemir al., 2013; Deng et al., 2025; Kovalyk et al., 2022: Ngoc Lan et al., 2021; Codella et al., 2019; Tschandl et al., 2018) spanning diverse imaging modalities (CT, MRI, pathol- ogy, ultrasound, etc.) and clinical targets (organs, tumors, lesions, etc.). For evaluation, we use (i) four external datasets: BTCV (Landman et al., 2015), ACDC (Bernard et al., 2018), UW-SC (Uni- versity of Waterloo), and BUS (Yap et al., 2020), to test cross-dataset generalization, and (ii) two unscen-class datasets: BraTS (Baid et al., 2021) and M&Ms-2 (Campello et al., 2021), to assess adaptation to novel structures. Details are provided in Appendix.
Baselines. We compare against SOTA methods in three categories: (I) Semantic segmentation: the task-specific model nnU-Net (Isensee et al., 2021), and universal models Clip-driven (Liu et al., 2023a), UniSeg (Ye et al., 2023), and Hermes (Gao, 2024): (2) In-context segmentation: Uni- verSeg (Butoi et al., 2023), Tyche (Rakic et al., 2024), MultiverSeg (Wong et al., 2024), and Iris (Gao al., 2025); (3) Interactive segmentation: nnInteractive (Isensee al., 2025), MedSAM (Ma al.,

---

Published as a conference paper at ICLR 2026
Comparison of semantic and in-context segmentation across different in-distribution datasets, measured by mean Dice scores (%). All in-context models use one-shot inference.
Table 2:
LITS
KITS
LASeaQQS
Breast
Chest X-my
Method CT
MRI
Tumor
Tumor
Cancer
POLY
Semantic (Task-specific Model)
nnU-Net
76.26 76.57 85.42 56.84
73.39
84.28
71.68
94.95
86.22
96.17
83.45
86.76
81.00
Semantic (Universal Models)
Clip-driven
84.37 83.20 84.10 10 60.12
75.99
80.94
76.47
95.45
85.82
95.81
86.32
86.04
84.31
UniSeg
84,91 85.36 85.34 56.97
73.41
81.57
75.71
96.02
80.88
91.40
82.57
84.24
83.96
Hermes
85.06
84.65
85.87
58.66
73.64
83.30
78.09
95.64
85.41
95.87
85.59
88.64
85.02
K-Prism
85.21
84.39
86.15
64.22
78.70
83.27
76.86
95.93
86.68
95.71
86.76
89.36
86.21
In-context
UniverSeg
60.51
59.02
60.51
49.02
58.60
72.18
54.82
94.65
81.07
95.45
73.93
80.01
71.53
Tyche
61.62
59.52
61.21
50.01
59.68
72.26
54.94
94.42
81.05
95.60
74.73
80.80
72.12
MultiverSeg
58.48
58.13
58.83
50.60
65.43
67.74
67.02
94.01
82.36
94.98
72.77
85.10
72.41
Iris
74.10
76.59
83.37
59.26
74.65
78.14
75.47
95.86
83.96
95.31
82.96
85.87
81.76
K-Prism
79.72
78.21
85.22
62.93
79.12
81.22
75.55
95.80
86.48
95.74
84.88
89.78
84.82
Table 3: Comparison of semantic and in-context segmentation across external and unscen-class datasets, measured by mean Dice scores (%). All in-context models use one-shot inference.
Method
External
Unseen-Clan
BTCV ACDC
UW-SC
BUS
AVG
BraTS
M&Ms-2
AVG
Semantic Task-speciin Model)
nnU-Net
67.25
84.64
79.78
68.00
74.92
Semantic (Universal Models)
Clip-driven
78.43
85.35
83.19
69.86
79.21
UniSeg
80.89
86.69
84.47
69.2 69.24
80.82
Hermes
79.25
85.87
87.64
72.49
81.81
K-Prism
80.24
87.26
87.54
78.75
83.45
In-context
UniverSeg
55.63
56.87
75.08
47.56
58.29
15.61
21.16
18.39
Tyche
57.07
57.62
74.60
43.68
58.74
15.92
21.95
18.94
MultiveSeg
48.69
51.15
78.76
62.08
60.67
18.23
26.21
22.22
Iris
74.97
83.36
86.82
68.93
78.52
25.83
26.30
26.07
K- K-Prism
76.82
85.89
87.66
79.59
82.49
22.20
41.61
31.91
2024), MultiverSeg (Wong et al., 2024), SAM2 (Ravi et al., 2024), and SegNext (Liu et al., 2024a).
All models are trained with 2-D slices extracted from our curated datasets under identical conditions. Implementation details are described in the Appendix.
4.2
RESULTS
4.2.1
SEMANTIC SEGMENTATION
Table 2 shows that K-Prism achieves the highest average Dice scores (86.21%) across 12 in- distribution datasets of diverse modalities, surpassing Hermes (85.02%). Clip-driven (84.31%) %), and UniSeg (83.96%). As task-specific model, nnU-Net performs well certain single-organ datasets (e.g., PAPILA 96.17%, LAScarQS 84.28%), consistent with its per-dataset optimization strategy (Isensee et al., 2021). However, its accuracy drops notably on multi-organ 2D slice segmen- tation (Liu et al., 2024b), such as AMOS-MRI (76.57%). and declines further on external datasets. In contrast, universal models achieve more stable performance across modalities, and K-Prism con- tinues this trend with even greater consistency. On challenging tasks such tumor segmentation, achieves notable performance improvements: 64.22% for LiTS and 78.70% for KiTS. These re- sults demonstrate the effectiveness of our framework in capturing anatomical context and delivering robust segmentation. Table presents results on external datasets. Our model yields the best gen- eralization across datasets with 83.45% average Dice score. These results establish K-Prism the new state-of of the-art for semantic medical segmentation, with strong generalization across diverse modalities and clinical settings.

---

Published as a conference paper at ICLR 2026
Table 4: Comparison of interactive segmentation performance across in-distribution, external and unseen-class datasets. To maintain clarity, only the mean values of key metrics are reported across multiple datasets. NoC90/NoC95 denote the average number of clicks required reach 90% and 95% Dice scores, while Dice( \/Dice(5) refer to Dice scores after and 5 clicks, respectively.
In-distribatiio
External
Unscen-Clase
Method
NoC90L NoC95 Dice(1)f Dice(S)) NoC90. NoC95 Dice(1)f Dico(5)t NoC90, NoC95 Diec(1)f Dice(5y7
Interactive
onlotcractive
5.08
73.36
88.44
3.52
5.58
68.43
85.53
4.96
7.42
47.60
86.77
McdSAM
3.17
5.16
88.23
91.21
3.35
6.40
87.26
90.93
6.69
8.88
65.44
80.13
MultiverSog
5.19
70.57
92.97
3.27
5.64
73.45
92.80
5.47
7.98
53.15
87.95
SAM2
3.94
6.27
86.46
87.96
3.59
7.34
87.57
89.00
8.71
9.84
59.88
66.81
SegNext
2.50
4.08
89.53
93.80
2.63
5.00
88.43
92.96
4.77
7.22
71.99
87.72
K. E.Prism
1.95
3.51
89.55
95.50
2.01
4.24
88.67
94,92
4.32
6.62
68.67
90.67
4.2.2 IN-CONTEXT SEGMENTATION
Table 2 (bottom) summarizes in-context segmentation performance across 12 in-distribution datasets. Our method achieves the highest average Dice score (84.82%) and ranks first on 11 out of 12 datasets. Table 3 further evaluates generalization to external and unseen-class datasets. Across four external datasets (BTCV, ACDC, UW-SC, and BUS), our method consistently outperforms all prior in-context segmentation models, achieving average Dice score of 82.49%. In the unseen- class setting, our model attains 31.91% on average, including a notable 15% improvement over Iris on M&Ms-2 (41.61% vs. 26.30%), demonstrating strong adaptability novel anatomical structures under limited supervision.
Compared to semantic segmentation, the in-context segmentation mode of K-Prism shows highly competitive performance- even surpassing Mode-1 (semantic segmentation) on challeng- ing datasets such as KiTS (79.129 vs. 78.70%) and ISIC (89.78% vs. 89.36%), though perfor- mance drops are observed on AMOS CT (85.21% + 79.72%). AMOS MRI (84.39% + 78.21%) and BTCV (80.24% - 76.82%) due to the complexity of multi-organ segmentation tasks where 2-D reference slices provide limited anatomical context. Despite this, K-Prism sets new SOTA for in-context medical segmentation and remains competitive with fully supervised counterparts, highlighting its effectiveness and potential for future in-context learning frameworks.
4.2.3 INTERACTIVE SEGMENTATION
We further evaluate our model in the interactive segmentation setting, where user clicks iteratively refine predictions (Table 4). Our method consistently achieves the best performance across in- distribution, external, and unseen-class datasets. With five clicks, our model reaches 95.50%
Dice score on in-distribution datasets, surpassing strong baselines like SegNext (93.80%), SAM2 (87.96%) MedSAM (1.21%) and MultiverSeg (92.97%). Notably, method also achieves the lowest NoC90 and NoC95 (1.95 and 3.51, respectively), indicating high-accuracy predictions with fewer interactions. On external datasets, 94.92% Dice score at five clicks, along with the best NoC scores (2.01 and 4.24). For unseen-class datasets, our method delivers clear advantages, reaching 90.67% Dice score at five clicks, while maintaining efficient convergence (4.32 for NoC90).
model maintains strong generalization, attaining
To further examine model behavior during interactive segmentation, analyze the convergence curves (Figure 3). CNN-based methods, nninteractive and MultiverSeg, that emphasize local struc- tural biases, show low initial Dice scores but improve rapidly with more clicks. SAMM-amily models, SAM2 and MedSA.M, achieve relatively high initial Dice scores but show only marginal gains with additional clicks. Because clicks are encoded only as 1-D sparse points without any 2-D spatial point map, these models show very limited improvement even with more user feedback. SegNext achieves a better balance with strong initialization and steady improvement. model which fuses 1-D sparse and 2-D dense prompts, achieves the best of both: high starting Dice and click- efficciff,, demonstrating the effectiveness of unified prompt integration for precise and efficient interaction.
comparison, our
8

---

Published as a conference paper at ICLR 2026
Datasets
Datasets
900
Setkon
Segledt
Clic
Click
Figure 3: Convergence curves of interactive segmentation on in-distribution, external, and unseen- class datasets. K- K-Prism consistently achieves higher Dice scores and faster convergence compared to all baselines.
Table 5: K-Prism ablations.
(a) MoE-based Cross-Attention (CA) and FFN layers contribute perfor- mance gains. We remove one component at a time.
(b) Evaluates the impact of dif- ferent components the in-
context segmentation mode.
Semantic In-context
Interactive
Method
Params
Dice
Dice
NoCo0q
NoC95,
Dioe())1 Dicet5)1
Method
Dice
Ours
43.29M
81.28
79.21
2.31
4.80
86.75
93.79
Ours
80. 84
wio MoE CA
30.6-M
77.38
77.11
2.47
5.11
86.55
93.23
2
64.65
w/o MoE FPN
40.13M
78.57
78.37
2.37
4.84
86.50
93.67
w/o I-D queries
77.19
w/o MoE FFN CA
27.48M
76.77
75.10
2.62
5.04
84.16
92.22
(c) Evaluates the impact of different components the interactive segmentation mode.
(d) Impact of refinement in interactive segmentation on the ACDC and BUS dataset.
Method
NoC NoCV5 Dicetty" Dicet5jf
Method
ACDC
BUS
NoC9o0 NoCoS, NoCCsO NoC95
2.34
86.55
95.68
w/o 2-D fusion
5.07
7.84
77.65
79.21
Mode
2.57
4.97
1.57
4.05
Mode-1 refine
1.67
4.52
1.10
3.47
W/o I-D queries
2,57
5.01
85.44
93.39
Mode 2 reline
1.77
4.59
1.17
3.53
4.3 ABLATION
This section covers component ablations (Tables 5(a)-(c) and separate refinemeentefffccieney study (Table 5(d)). Parts (a)-(c) are based on models trained on M&Ms and Breast Cancer and evaluated on ACDC and BUS. Part (a) examines the impact of MoE components across all modes: remov- ing either MoE cross-attention or the MoE FFN noticeably degrades performance, while removing both leads the largest drop. Parts (b) and (c) evaluate single-mode settings (in-context and inter- active, respectively). In both cases, 2-D fusion proves indispensable-its removal nearly collapses performance- performance-whereat removing 1-D queries causes only modest decline, showing that they are complementary but less critical.
Table 5(d) evaluates refinement efficiency in the interactive setting using the model trained in our main experiment. Mode-3 serves as a pure interactive baseline, while Mode-1 and Mode-2 leverage semantic or in-context knowledge to provide an initial mask, which then refined with user clicks. With such initialization, the required clicks to reach 90% Dice scores decrease by over 30%. This initialization strategy transforms the clinical workflow from manual delineation to efficient refine- ment, substantially reducing annotation burden while maintaining flexibility through seamless mode switching. Further ablation studies and analyses are included in Appendix.
5 DISCUSSION
Our work presents unified framework for medical image segmentation that integrates semantic priors, in-context, and interactive feedback knowledge. While the results demonstrate strong perfor- mance and generalization, several aspects merit further discussion.
K-Prism provides a flexible framework that unifies different knowledge sources, enabling seamless switching across paradigms. This positions it foundation for universal medical segmentation.

---

Published as a conference paper at ICLR 2026
Future work may refine components like semantic priors, in-context retrieval, interactive refine- ment, making K-Prism not just a model but a platform for ongoing innovation. Our model can also serve as an efficient annotation tool for building large-scale medical image datasets. Clini- cians can generate initial segmentation from semantic priors or in-context inference and refine them with minimal clicks, creating a human-in-thee -ooo--oo-- pipeline that highlights K-Prism's value both a segmentation framework and practical infrastructure for large-scale medical modeling.
Despite these strengths, several challenges remain. The accuracy of in-context segmentation still lies on the quality of reference examples, and the MoE decoder introduces additional computational cost that may hinder real-time deployment. Furthermore, large domain shifts, such as unseen unseen-class datasets or new imaging protocols, continue to pose difficulties, often leading to degraded perfor- mance. unseen-class scenarios, we observe that K-Prism's 's interactive mode remains highly ro- bust, as olick-based refinement relies on local boundary and texture cues rather than global semantic priors, enabling reliable adaptation even to anatomy never seen during training. In contrast, in- context segmentation exhibits larger performance degradation under severe anatomical or modality shifts, reflecting the intrinsic difficulty of exemplar-based matching when the reference and query differ substantially. These observations highlight both the strengths and natural limitations of dif- ferent knowledge paradigms, and they further demonstrate the strength of our unified design, which maintains reliable performance by allowing interactive refinement to compensate when exemplar- guided inference becomes unreliable. Furthermore, our current 2D slice-based formulation inher- ently limits the full utilization of 3D volumetric context. Although 2D interactive refinement aligns well with real clinical workflows, where radiologists typically correct boundaries on a slice-by- slice basis, future work could explore efficient 2D-to-3D propagation strategies that transfer refined slice-level predictions to full volumetric outputs. Additional research may also focus on (i) optimiz- ing expert specialization across segmentation paradigms and (ii) enhancing robustness for unseen anatomical classes and domain-shifted scenarios.
CONCLUSION
We introduced K-Prism, unified segmentation framework that integrates semantic priors, in- context, and interactive feedback knowledge into unified dual-prompt representation, coupled with a Mfxture-of-Expeer decoder for dynamic routing. This design supports flexible adaptation across clinical scenarios, from supervised segmentation few-shot adaptation and interactive refinement. Experiments on 18 public datasets show consistent state-of-the-an performance with strong robust- ness and transfcrability to unseen domains. We view this work as a step toward universal medical image segmentation models that can serve reliable backbones for diverse clinical applications, narrowing the gap between algorithmic advances and real-world deployment.
ACKNOWLEDGMENTS
This research has been partially funded by research grants to D. Metaxas through NSF: 2310966, 2235405, 2212301, 2003874, 1951890 and NIH 2101HL127661.
REFERENCES
Walid Al-Dhabyani, Mohammed Gomaa, Hussien Khaled, and Aly Fahmy. Dataset of breast ultra- sound images. Data in brief. 28: 104863, 2020.
Ujjwal Baid, Satyam Ghodasara, Suyash Mohan, Michel Bilello, Evan Calabrese, Errol Colak, Key- van Farahani, Jayashree Kalpathy-Crran Cramer, Felipe C Kitamura, Sarthak Pati, al. The rsna-asnr- miccai brats 2021 benchmark on brain tumor segmentation and radiogenomic classification. arXiv preprint ar Xrv: 20t.:2207.002//42 07, 2021.
Olivier Bernard, Alain Lalande, Clement Zotti, Frederick Cervenansky, Xin Yang, Pheng-Ann Heng, Irem Cetin, Karim Lekadir, Oscar Camara, Miguel Angel Gonzalez Ballester, al. Deep learning techniques for automatic mri cardiac multi-structures segmentation and diagnosis: is the problem solved? IEEE transactions on medical imaging, 37(1 1):2514-225 2525, 2018
10

---

Published as a conference paper at ICLR 2026
Patrick Bilic, Patrick Christ, Hongwei Bran Li, Eugene Vorontsov, Avi Ben-Cohen, Georgios Kaissis, Adi Szeskin, Colin Jacobs, Gabriel Efrain Humpire Mamani, Gabriel Chartrand, et al. The liver tumor segmentation benchmark (lits). Medical image analysis, 84:102680, 2023.
Victor Ion Butoi, Jose Javier Gonzalez Ortiz, Tianyu Ma. Mert R Sabuneu, John Guttag, and Adrian V Dalca. Universeg: Universal medical image segmentation. In Proceedings of the IEEE/CVF International Conference Computer Vision, pp. 21438-21451, 2023.
Victor M Campello, Polyxeni Gkontra, Cristian Izquierdo, Carlos Martin-Isla, Alireza Sojoudi, Pe- ter M Full, Klaus Maier-Hein, Yao Zhang, Zhiqiang He, Jun Ma, et al. Multi-centre, multi-vendor and multi-disease cardiac segmentation: the m&ms challenge. IEEE Transactions on Medical Imaging, 40(12):3543-355 2021.
Sema Candemir, Stefan Jaeger, Kannappan Palaniappan, Jonathan Jonathan Musco, Rahul K Singh, Zhiyun Xue, Alexandros Karargyris. Sameer Antani, George Thoma, and Clement J McDonald. Lung segmentation in chest radiographs using anatomical atlases with nonrigid registration. IEEE trans- actions medical imaging, 33(2):777-590 2013.
Bowen Cheng, Ishan Misra, Alexander G Schwing, Alexander Kirillov, and Rohit Girdhar. Masked- attention mask transformer for universal image segmentation. In Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pp. 1290-1299, 1299, 2022.
Ho Kei Cheng, Yu-Wing Tai, and Chi- Keung Tang. Rethinking space-tima networks with improved memory coverage for efficient video object segmentation. Advances in neural information pro- cessing systems, 34:1 1777-117999 2021.
Ho Kei Cheng, Seoung Wug Oh, Brian Price, Youn-Yung Lee, and Alexander Schwing. Putting the object back into video object segmentation. Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 3151-3161, 2024.
Noel Codella, Veronica Rotemberg, Philipp Tschandl, M Emre Celebi, Stephen Dusza, David Gut- man, Brian Helba, Aadi Kalloo, Konstantinos Liopyris, Michael Marchetti, et al. Skin lesion analysis toward melanoma detection 2018: A challenge hosted by the international skin imaging collaboration (isic). arXiv preprint arXiv: /90.0036,, 2019.
Ruining Deng, Tianyuan Yao, Yucheng Tang, Junlin Guo, Siqi Lu, Juming Xiong, Lining Yu, Quan Huu Cap, Pengzhou Cai, Libin Lan, et al. Kpis 2024 challenge: Advancing glomerular segmentation from patch-to slide-level. arXiv preprint arXiv:2502. 07288, 2025.
Yunhe Gao. Training like medical resident: Context-prior learning toward universal medical image segmentation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11194-11200 2024.
Yunhe Gao, Di Liu, Zhuowei Li, Yunsheng Li, Dongdong Chen, Mu Zhou, and Dimitris N Metaxas. Show and segment: Universal medical image segmentation via in-context learning. In Proceed- ings of the Computer Vision and Pattern Recognition Conference, pp.
2025.
Bangwei Guo, Meng Ye, Yunhe Gao, Bingyu Xin, Leon Axel, and Dimitris Metaxas. Verse: In- tegrating multiple queries as prompts for versatile cardiac mri segmentation. In International Conference on Information Processing in Medical Imaging, pp. 373-387. Springer, 2025.
Ali Hatamizadeh, Yucheng Tang, Vishwesh Nath, Dong Yang, Andriy Myronenko, Bennett Land- man, Holger R Roth, and Daguang Xu. Unetr: Transformers for 3d medical image segmentation. In Proceedings of the IEEE/CVF winter conference applications of computer vision, pp. 574- 584, 2022.
Nicholas Heller, Niranjan Sathianathen, Arveen Kalapara, Edward Walczak, Keenan Moore, Heather Kaluzniak, Joel Rosenberg, Paul Blake, Zachary Rengel, Makinna Oestreich, et al. The kits19 challenge data: 300 kidney tumor cases with clinical context, et semantic segmentations, and surgical outcomes. arXiv preprint urXiv:1904, 004400 2019.
Fabian Isensee, Paul F Jaeger, Simon AA Kohl, Jens Petersen, and Klaus H Maier-Hein. mnu- net: self-configuring method for deep learning- based biomedical image segmentation. Nature methods, 1822)200-211, 2021.
beaseng--based
11

---

Published as a conference paper at ICLR 2026
Fabian Isensee, Maximilian Rokuss, Lars Krämer, Stefan Dinkelacker, Ashis Ravindran, Florian Stritzke, Benjamin Hamm, Tassilo Wald, Moritz Langenberg, Constantin Ulrich, et al. nninterac- tive: Redefining 3d promptable segmentation. arXiv preprint arXiv:2503.088733 2025.
Stefan Jaeger, Alexandros Karargyris, Sema Candemir, Les Folio, Jenifer Siegelman, Fiona Callaghan, Zhiyun Xue, Kannappan Palaniappan, Rahul K Singh, Sameer Antani, et al. Au- tomatic tuberculosis screening using chest radiographs. IEEE transactions on medical imaging, 33(2):233-244 245, 2013.
Yuanfeng Ji, Haotian Bai, Chongjian Ge, Jie Yang, Ye Zhu, Ruimao Zhang, Zhen Li, Lingyan Zhanng. Wanling Ma, Xiang Wan, et al. Amos: A large-scale abdominal multi organ benchmark for versatile medical image segmentation. Advances in neural information processing systems, 35:36722-36732. 2022.
Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander Berg, Wan-Yen Lo, et al. Segment anything. Proceed- ings of the IEEE/CVF International Conference on Computer Vision, pp. 4015 4015,402,222 2023.
Oleksandr Kovalyk, Juan Mornles-Sánchez, Rafael Verdū-Monedero, Inmaculada Ana Palazón-Cabanes and José-Luis Sancho-Gómez. Papila: Dataset with fundus images and clinical data of both eyes of the same patient for glaucoma assessment. Scientific Data, 91):291, 2022.
Bennett Landman, Zhoubing Xu, Igelsias, Martin Styner, T Langerak, and Amo Klein. Miccai multi-atlas labeling beyond the cranial vault-workshop and challenge. In Proc. MICCAI Multi- Atlas Labeling Beyond Cranial Works-Wrkksop Challenge, volume 5, pp. 12, 2015.
Lei Li, Fuping Wu, Guang Yang, Lingchao Xu, Tom Wong, Raad Mohiaddin, David Firmin, Jen- nifer Keegan, and Xiahai Zhuang. Atrial scar quantification via multi-scale enn in the graph-cuts framework. Medical image analysis, 60: 60:101595, 2020.
Lei Li, Veronika A Zimmer, Julia A Schnabel, and Xiahai Zhuang. Atrialgeneral: domain gen- eralization for left atrial segmentation multi-center Ige mris. In International Conference on Medical Image Computing and Computer-ssiisted Intervention, pp. 557-566. Springer, 2021.
Lei Li, Veronika A Zimmer, Julia A Schnabel, and Xiahai Zhuang. Atrialjsqnet: new framework for joint segmentation and quantification of left atrium and scars incorporating spatial and shape information. Medical image analysis, 76: 76:102303 2022a.
Lei Li, Veronika A Zimmer, Julia A Schnabel, and Xiahai Zhuang. Medical image analysis on left atrial Ige mri for atrial fibrillation studies: A review, Medical image analysis, 77:102360, 2022b.
Jie Liu, Yixiao Zhang, Jie-Neng Chen, Junfei Xiao, Yongyi Lu,, Bennett A Landman, Yixuan Yuan, Alan Yuille, Yucheng Tang, and Zongwei Zhou. Clip-driven universal model for organ segmenta- tion and tumor detection. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 21152-21164 2023a.
Qin Liu, Zhenlin Xu, Gedas Bertasius, and Mare Niethammer. Simpleclick: Interactive image segmentation with simple vision transformers. In Proceedings of the IEEF/CVF International Conference on Computer Vision, pp. 22290-22300, 2023b.
Qin Liu, Jaemin Cho, Mohit Bansal, and Mare Niethammer. Rethinking interactive image seg- mentation with low latency high quality and diverse prompts. In Proceedings of the JEEF/CVF Conference on Computer Vision and Pattern Recognition, pp. 3773-3782, 2024a.
Xiaoyu Liu, Linhao Qu, Ziyue Xie, Jiayue Zhao, Yonghong Shi, and Zhijian Song. Towards more precise automatic analysis: a systematic review of deep learning -aanid multi-organ segmentation. BioMedical Engineering Online, 23(1):52, 2024b.
Michelle Livne, Jana Rieger, Orhun Utku Aydin, Abdel Aziz Taha, Ela Marie Akay, Tabea Kossen, Jan Sobesky, John D Kelleher, Kristian Hildebrand, Dietmar Frey, et al. A u-net deep learning framework for high performance vessel segmentation in patients with cerebrovascular disease. Frontiers neuroscience, 13:97, 2019.
12

---

Published as a conference paper at ICLR 2026
Jun Ma, Yuting He, Feifei Li, Lin Han, Chenyu You, and Bo Wang. Segment anything in medical images. Nature Communications, 15(1):654, 2024.
Phan Ngoc Lan, Nguyen Sy An, Dao Viet Hang. Dao Van Long, Tran Quang Trung, Nguyen Thi Thuy, and Dinh Viet Sang. Neounet: Towards accurate colon polyp segmentation and neoplasm detection. In International Symposium on Visual Computing, pp. 15-28. Springer, 2021.
Marianne Rakic, Hallee Wong, Jose Javier Gonzalez Ortiz, Beth A Cimini, John V Guttag, and Adrian Dalca. Tyche: Stochastic in-context learning for medical image segmentation. In Pro- ceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11159- 11173, 2024.
Nikhila Ravi, Valentin Gabeur, Yuan- Hu, Ronghang Hu, Chaitanya Ryali, Tengyu Ma, Haitham Khedr, Roman Rädle, Chloe Rolland, Laura Gustafson, al. Sam 2: Segment anything in images and videos. arXiv preprint vrXiv:2408.007: 14, 2024.
Konstantin Sofiiuk, Ilya A Petrov, and Anton Konushin. Reviving iterative training with mask guid- ance for interactive segmentation. In 2022 IEEE International Conference Image Processing (ICIP), pp. 3141-3145. IEEE, 2022.
Philipp Tschandl, Cliff Rosendahl, and Harald Kittler. The ham10000 dataset, large collection of multi-source dermatoscopic images of common pigmented skin lesions. Scientific data, 5(1):1-9, 2018.
University of Waterloo.
Skin
cancer detection.
https: //
uwaterloo wwtvisiaaa ima---maaaeeee wwwp---iaae---pesssseeeee lab/ research-demost skin-cancen www.ernneeeeeetioonn. detection.
Jakob Wasserthal, Hanns-Christian Breit, Manfred T Meyer, Maurice Pradella, Daniel Hinck, Alexander W Sauter, Tobias Heye, Daniel T Boll, Joshy Cyriac, Shan Yang, et al. Totalseg- mentator: robust segmentation of 104 anatomic structures in images. Radiology: Artificial Intelligence, 233)222 2023.
Hallee Wong, Jose Javier Gonzalez Ortiz, John Guttag, and Adrian V Dalca. Multiverseg: Scalable interactive segmentation of biomedical imaging datasets with in-context guidance. urXiv preprint 1X.i:2412.110588 2024.
Moi Hoon Yap. Manu Goyal, Fatima Osman, Robert Martí, Erika Denton, Ame Juette, and Reyer Zwiggelaar. Breast ultrasound region of interest detection and lesion localisation. Artificial intel- ligence in medicine, 107:101880, 2020.
Yiwen Ye, Yutong Xie, Jianpeng Zhang, Ziyang Chen, and Yong Xia. Uniseg: A prompt-driven uni- versal segmentation model as well as strong representation learner. In International Conference on Medical Image Computing and Computer-Assisted Intervention, pp. 508-518. Springer, 2023.
Rammah Yousef, Shakir Khan, Gaurav Gupta, Bader M Albahlal, Saad Abdullah Alajlan, and Alcem Ali. seddged-u-net-sspp-ee and deep learning optimization for brain tumor segmenta- tion. Diagnostics, 13(16):2633. 2023.
Hong- Yu Zhou, Xiaoyu Chen, Yinghao Zhang, Ruibang Luo, Liansheng Wang, and Yizhou Yu. Generalized radiograph representation learning via cross cross- superviseni between images and free- text radiology reports. Nature Machine Intelligence, 4(1):32-40) 2022.
Kevin Zhou, Hayit Greenspan, Christos Davatzikos, James S Duncan, Bram Van Ginneken, Anant Madabhushi, Jerry L Prince, Daniel Rueckert, and Ronald M Summers. A review of deep learning in medical imaging: Imaging traits, technology trends, case studies with progress highlights, and future promises. Proceedings of the IEEE, 109(5):820- 838, 2021.
13

---

Published as a conference paper at ICLR 2026
A APPENDIX
A.1 DATASET
A.1.1 DATASET DETAILS
Abdominal Multi-Organ Segmentation (AMOS) (Ji et al., 2022). The AMOS dataset comprises 500 CT and 100 MRI abdominal scans from 600 patients at Longgang District People's Hospital. provides annotations for 13 anatomical structures, including spleen, right kidney, left kidney, gallbladder, esophagus, liver, stomach, aorta, inferior vena cava, pancreas, right adrenal gland, left adrenal gland, duodenum, bladder, and prostate/uterus. The AMOS_C set includes 200 training and 100 validation cases, while the AMOS .RRI set offers 40 training and 20 validation cases. We use both modalities for upstream training with a 95%/5% split on the official training set to train and validate, then test on the provided test set.
Multi-Centre Multi- Vendor
Multi-Disease Cardiac Image Segmentation Challenge (M&Ms) (Campello et al., 2021). The M&Ms dataset from the MICCAI 2020 Challenge includes multi-center, multi-vendor cardiac MRI scans from patients with cardiomyopathies and healthy con- trols. All scans are short-axis cine images and have expert annotations for left/right ventricles and left ventricular myocardium at end-diastolic and end-systolic phases. The official train and test sets used.
The Liver Tumor Segmentation Benchmark (LITS) (Bilic et al., 2023). This dataset includes 201 contrast-enhanced abdominal CT scans (131 training, 70 testing) from multiple international medical centers. It covers patients with diverse liver tumors, including hepatocellular carcinoma and metastases from colorectal, breast, and lung cancers. All scans contain expert annotations for liver and tumor regions. We use the 131 public training cases, split 80%/20%, for upstream training and testing. In our experiments, only the tumor labels from each CT volume utilized.
The Kidney Tumor Segmentation 2019 dataset (KiTS) (Heller et al., 2019), This dataset is collected at the University of Minnesota Medical Center between 2010 and 2018, includes CT scans and clinical treatment outcomes from 300 patients who underwent nephrectomy for kidney tumors. Of these, 210 cases are publicly available, while the remaining 90 are reserved for evaluation. In our upstream training, we utilize the public portion by splitting it into 80% for training, and 20% for testing. In our experiments, only the tumor labels from each CT volume are utilized.
Left Atrial and Scar Quantification & Segmentation Challenge (LASSarQSS) (Li et al., 2020; 2021; 2022a;b). The LAScarQS dataset was released as part of a MICCAT challenge dedicated to left atrial (LA) cavity and scar segmentation from late gadolinium enhancement (LGE) cardiac MRI. comprises 130 LGE MRI volumes from patients with atrial fibrillation (AF), acquired across mul- tiple clinical centers. Each volume contains expert annotations for both the LA cavity and atrial scar regions. The dataset reflects substantial real-world variability in image quality, atrial morphology, and scar patterns, with many cases presenting significant segmentation challenges. Following the official protocol, we adopt the 80%/20% split of the training set for model training and testing. In our experiments, only the LA cavity annotations used.
Dataset of Breast Ultrasound Images (Breast Cancer) (Al-Dhabyani et al., 2020). This dataset comprises 780 grayscale ultrasound images collected from 600 female patients (aged 25-75) at Baheya Hospital for Early Detection and Treatment of Women's Cancer, Egypt. Each image is categorized into one of three classes: normal, benign, malignant. Expert-annotated segmentation masks are provided for all desion-containg images. The dataset reflects real-world variability in breast anatomy, lesion characteristics, and image quality, making valuable for developing and evaluating models for breast cancer classification, detection, and segmentation. We use only the 647 images labeled as benign or malignant and adopt an 80% split for training and testing.
Chest X-ray Masks and Labels (Chest X-ray) (Al-Dhabyani et al., 2020; Jaeger et al., 2013). Chest X-ray is a chest radiograph dataset provided by the National Library of Medicine, National Institutes of Health (Bethesda, USA) and Shenzhen No.3 People's Hospital (Guangdong, China). This dataset includes a wide spectrum of abnormalities such as effusions and miliary patterns and is widely used for tuberculosis screening, lung segmentation, and domain adaptation studies. In this study, we only use the lung segmentation mask to train our models. We adopt an 80%/20% split of the training set for model training and testing.
14

---

Published as a conference paper at ICLR 2026
Kidney Pathology Image Segmentation (KPls) Challenge (Deng et al., 2025). The KPIs dataset was released as part of the MICCAI 2024 Challenge to benchmark glomeruli segmentation per- formance in chronic kidney disease (CKD) pathology. We focus the patch-level segmentation task, which involves pixel-wise identification of glomeruli within PASS-tained image patches. The patches exhibit variations in glomeruli morphology and surrounding tissue structures due to differ- ences in disease states and slide preparation. We adopt a 75%/5% split on the official training set for training, validation, and testing, The original images are then uniformly partitioned into non- overlapping patches of size 512 x 512, and patches without any glomeruli annotations are discarded.
PAPILA Dataset (Kovalyk et al., 2022). The PAPILA dataset provides fundus photographs from both eyes of individual patients, accompanied by expert annotations. Each image is annotated with optic disc and optic cup segmentations, while patient-level clinical labels are available for disease evaluation. In this study, we focus exclusively on the optic disc segmentation task, using the provided disc masks to train and evaluate our models. We adopt a 80%/20% split on the official dataset with 488 image-mask pairs for training and testing.
POLY) (Ngoc Lan et al., 2021). The BKAI-IGH NeoPolyp dataset, released by the BKAI Research Center (Hanoi University of Science and Technology) in collaboration with the Institute of Gastroenterology and Hepatology (IGH), Vietnam, consists of 1,200 colonoscopy images (1,000 white-light light imaging (WLI) and 200 flexible spectral imaging color enhancement (FICH) images). The dataset is split into 1,000 training and 200 test images. Each polyp is annotated with both segmentation masks and binary labels indicating neoplastic (red) or non-ncoplastic (green) classes, verified by two experienced endoscopists. We adopt an 80%/20% split on the official training set for training and testing.
BKAI-IGH NeoPolyp Dataset
International Skin Imaging Collaboration (ISIC) Dataset (Tschandl et al., 2018; Codella et al., 2019). The ISIC dataset is a large-scale dermoscopic image collection, introduced through the ISIC Skin Lesion Analysis Challenges. It contains high-resolution skin lesion images accompanied by expert-annotated segmentation masks delineating lesion boundaries. In this study, we use the ISIC 2018 segmentation subset. The dataset encompasses diverse lesion appearances and acquisition conditions, capturing real-world variability in skin tone, lighting, and lesion morphology. We adopt the official dataset split, using it for model training, validation, and testing.
Beyond the Cranial Vault (BTCV) Dataset (Landman et al., 2015). The BTCV dataset, released as part of the MICCAI 2015 Multi-Atlas Labeling Beyond the Cranial Vault challenge, is a widely used benchmark for abdominal organ segmentation. It consists of 50 contrast-enhanced abdominal CT scans provided by Vanderbilt University Medical Center, acquired in the portal venous phase from patients with either metastatic liver cancer or postoperative abdominal wall hernia. Each scan is annotated with 13 abdominal organs, including the liver, spleen, pancreas, kidneys, stomach, gallbladder, esophagus, aorta, inferior vena cava, and duodenum, among others. The scans exhibit variable field of view and resolution, with in-plane spacing ranging from 0.54 0.54 mm² to 0.98 x 0.98 mm² and slice thickness between 2.5 mm and 5.0 mm. In this study, we adopt 30 scans from the official training set as an external test set to evaluate the generalization performance of our model.
Automatic Cardiac Diagnosis Challenge (ACDC) Dataset (Bernard et al., 2018). The ACDC dataset comprises cardiac MRI scans collected over six years at the University Hospital of Dijon, acquired with 1.5T and 3.OT Siemens scanners. Each case includes short-axis cine sequences with expert annotations of the same anatomical structures as in M&Ms, at both end-systolic and end- diastolic phases. In our study, ACDC serves an external benchmark to assess the generalization capability of our model, where we adopt the 200 official training cases as validation data and the 100 official test cases for final evaluation.
University of Waterloo Skin Cancer (UW-SC) Dataset (University of Waterloo). The UWaterloo Skin Cancer Dataset comprises 167 dermoscopic images of skin lesions, with manual segmentation masks provided for each lesion and verified by experts. In our work, we leverage this dataset to evaluate the segmentation performance in a challenging and heterogeneous clinical image collection.
Breast Ultrasound Dataset B (BUS) (Al-Dhabyani et al., 2020). BUS is a breast ultrasound collection designed for offieet-o-interest (ROI) detection and lesion localization. It contains 163 ultrasound images with expert-annotated lesion regions. In our study, all benign and malignant lesions merged into single lesion class simplify the task. We use this dataset solely as an external validation set to assess the generalization capability of our model.
15

---

Published as a conference paper at ICLR 2026
BraTS Dataset (Baid et al., 2021). The BraT'S series of datasets is widely used benchmark in brain tumor analysis. It provides multimodal MRI scans with expert annotations of glioma substructures. In our experiments, we test method on the same 369 slices as in (Liu al., 2023b).
M&Ms-2 Dataset (Campello et al., 2021). The Multi-Centre, Multi-Vendor & Multi-Disease (M&Ms-2) dataset released as part of the M&Ms-2 challenge at MICCAI 2021. For the M&Ms- 2 dataset, we utilize only the long-axis (LA) cine MRI images, which are not present in the M&Ms-1 dataset. This setting enhances data diversity and allows us to assess the generalization ability of our model under an unseen-class scenario. In our experiments, we exclusively use the 320 long-axis (LAX) cine MRI images from the official M&Ms-2 test set as an external test set to evaluate model generalization. The 400 LAX images from the official training set are adopted as validation set, while none of the M&Ms-2 images are used during model training.
A.1.2 DATASET STATISTICS
As shown in Table 5, we train and evaluate our framework on a diverse set of publicly avail- able medical image segmentation datasets spanning multiple imaging modalities and anatomical regions. The training pool covers 12 datasets: abdominal CT/MRI (c.g., AMOS, LiTS, KiTS), car- diac MRI (M&Ms, ACDC, LAScarQS), ultrasound (Breast cancer), dermoscopy (ISIC), endoscopy (BKAI POLY), pathology (KPIs), fundus imaging (PAPILA), and chest X-rays. These datasets vary substantially in size, with sample counts ranging from few hundred several thousand 3D volumes or 2-D slices, and in annotation granularity, from single-organ labels to multi-organ delin- eations. For evaluation, adopt both external and unseen-class test to assess generalization across modalities, anatomical structures, and imaging centers.
Table 6: Datasets statistics.
3D Volumes
2-D Slices
Dataset
cls
Modality
Train
Validate
Test
Train
Validate
Test
AMOS_CT
CT
33944
1689
AMOS.MRI
13
MRI
38
2
20
5160
263
2614
M&Ms
3
MRI
300
340
2475
2821
LITS
CT
104
5687
1625
KITS
CT
168
42
5269
1609
LAScarQS
MRI
3561
Chest r-y
X ray
450
116
Brenst Cancer
Ultrasound
518
KPis
Pathology
8261
500
4055
PAPILA
Fundus
390
98
BKALPOLY
Endoscopy
800
200
ISCC
Demeescopy
2594
100
1000
BTCV
13
CT
30
3791
ACDC
3
MRI
200
100
1841
1001
UW-SC
Demmescugy
167
BUS
Uhraasoun
163
BraTS
MRI
369
MRM
MRI
400
A.2 IMPLEMENTATION
Training details. Our models are trained for 75 epochs with a batch size of 16 on 8 Quadro RTX 8000 GPUs using the AdamW optimizer with a base learning rate of x 10 A cosine annealing scheduler with 10 warm-up epochs is applied, where the minimum learning rate is scaled by 1 x 10-5 Images are resized to 512 x 512 during training, with augmentations including random flips, affine transforms (shift, scale, rotation), brightness/contrast adjustments, Gaussian blur/noise, and grid distortions. To encompass all operational modes, each training batch randomly assigned to one mode, with probabilities of 0.3, 0.3, and 0.4 for Mode-1, Mode-2, and Mode-3, respectively. During inference, input images are resized with the long side fixed to 512 while preserving aspect ratio, and Dice scores are computed after remapping predictions to the original resolution. We use combination of binary cross cross-entroyy entropy loss and Dice loss to compute the mask loss L = Loc + Laiinn The same settings are applied on ablation study.
16

---

Published as a conference paper at ICLR 2026
Click representation. In both training and inference, user clicks are encoded as disk-shaped maps with a fixed radius of pixel. Consistent with previous studies (Sofiiuk et al., 2022; Liu et al., 2023b), simulated clicks are generated by comparing the predicted segmentation against the ground truth. Differing from prior strategies, however, we place each new click at the centroid of the largest misclassified connected component, which more closely mimics practical user interactions in medi- cal image analysis. We consider three operating modes: in Mode-1 and Mode-2, the model produces an initial mask and then refines it using two clicks, whereas in Mode-3 the model iteratively applies three clicks, yielding three successive segmentation masks.
Additional implementation details for baseline models. For SAM2, which is pretrained on large-scale natural images, we initialize the model using the official sam2 1 hiera base plus checkpoint and fine-tune on our curated medical datasets under the same training settings as all other methods. For SegNext and MedSAM, the ViT-based encoder is first initialized using MAE- pretrained ViT-Base weights. After loading these pretrained encoder weights, the entire network is jointly trained end-to-end on our curated medical datasets. For other medical segmentation models such as MultiverSeg, to avoid any potential data overlap with their original training data sources, train them from scratch User interactions (positive/negative clicks) are simulated consistently across all methods following the same policy described above.
curated datasets to ensure fair comparison and full convergence.
A.3 MODEL ARCHITECTURE
Image encoder. We employ UNet (Isensee et al., 2021) as the image encoder, a widely adopted lightweight backbone for medical image segmentation. The encoder extracts hierarchical features three different resolutions, namely 1/16, 1/8, and 1/4 of the original image, yielding comprehen- sive multi-scale representation (S 3). In our experiments, the feature maps have spatial sizes of 32 x 32, 64 x 64, and 128 x 128, with channel dimensions of 384, 192, and 96, respectively.
Mask encoder Interactive prompt encoder. We utilize a lightweight mask encoder extract hierarchical representations from the reference masks. The mask encoder employs simple consecu- tive convolutional blocks to generate hierarchical features that are spatially aligned with those from the image encoder, producing reference mask features at the same scales (1/4, 1/8, and 1/16 of the input resolution). For the interactive prompt encoder, we adopt an identical architectural design to ensure consistency.
Fusion of image features and 2-D dense prompts. After extracting features from the image en- coder and either the mask encoder or the interactive prompt encoder, employ unified prompt fusion module to combine them. To fully exploit contextual cues, our model operates on multi- scale features. For elarity, the main paper illustrates the architecture with single-scale example, while the appendix provides the details of the multi-scale processing. Specifically, the feature maps FF22, Fo4, Fi28) ) are obtained from the image encoder different resolutions.
For Mode-3 (interactive segmentation), fusion is straightforward: the image features are directly added to the interactive click and mask features, as they share the same dimensionality. For Mode-2, as discussed earlier, fusion is applied only at the lowest resolution F32 (384 x x 32). The resulting fused features are then propagated across higher scales using the residual connections described in the following section.
After flattening, we define the reference keys, values, and query keys as: Tf Flat (F&) E N re (lat(FF)) E Cx Nathon
K Flaa Cx
(12)
where h 32, C 384.
For any similarity function e CCC x RC + R, the pairwise affinity matrix A and its softmax- normalized form W are computed as:
Ai c K KG)),
W Σ,, exp (An.))
A. W E R Nahwwhw
(13)
Following Cheng et al. (2021; 2024), we adopt the negative squared Euclidean distance to compute A:
KTd
2K K3 - k
||
(14)
17

---

Published as a conference paper at ICLR 2026
where the last term can be omitted as shown in (Cheng et al., 2021), which improves efficiency and reduces computational cost. Finally, the 2-D fusion feature map is aggregated as:
Fuse- V
Fiuse E CRC.
(15)
and is passed into the decoder for mask prediction. To facilitate subsequent operations, the feature channels at all scales are projected to 256 via a linear layer before being fed into the decoder.
Generating the click queries. In this section, we provide additional details on how user clicks are encoded into 1-D query embeddings for interactive segmentation. Given a user click at pixel coordinates Po (z, y) in the original image I p3xH we first map it to the downsampled feature space F E (BX(((()))((W))) where denotes the stride of the encoder. The corresponding coordinates are:
x'
(16)
Around the mapped point (x',y'), we extract (2r + 1) (2r + 1) local window to capture neigh- borhood context. We set r in our experiments. The pooled feature vector is then obtained via average pooling:
f pooledd
Σ Σ
F F24144
(17)
(2r 1)²
The pooled feature projected into the query space using multilayer perceptron (MIP))
x f MLP Ty R'
(18)
To obtain the final single click query embedding, we incorporate the positional encoding from SAM (Kirillov al., 2023):
of = + PosEmbed((r, v), E
(19)
We generate two separate groups of click queries, QT for positive clicks and Q"n for negative clicks which correspond to the positive and negative attention masks. Since the number of clicks in the groups may differ during interaction, the smaller group is padded with dummy queries (zero vectors) to maintain balance. This design ensures stable training and inference when combining click queries across interaction steps. The two groups are then concatenated into the final 1-D click queries Q
Decoder. As shown in Figure 2 (b), we use L 6 layers in total in the transformer decoder. Similar to Fask2Former (Cheng et al., 2022), we adopt a round-robin strategy for multi-scale interaction between image features and integrated queries. The blocks are scheduled in a cyclic order across scales (i.e., 1 + 2 + 3 + + 2 + 3), ensuring balanced cross-scale information exchange throughout the decoding process. Finally, the decoded features are passed through a lightweight mask decoder that projects them back to the spatial resolution of the input, producing the final segmentation masks.
Residual connections across scales. Following Verse (Guo et al., 2025), we adopt residual resam- pling connections to enhance the interaction between 2-D fusion features and I-D queries across multiple scales. Concretely. for feature map Fi layer we first resample it to match the resolu- tion of the next layer F1+1:
presampled
Resample (Fi).
(20)
where Resample(-) denotes upsampling or downsampling depending on the relative scales. We then compute a residual representation via a convolutional layer and add it to the next-layer features:
uupdated H+
F + Res [itesamplld
(21)
By iteratively resampling and aggregating residuals across layers, the model improves the interaction between features and queries over different scales.
masked attention. Attention masks have been shown to improve the efficiency of attention by constraining it to more relevant regions (Cheng et al., 2022; 2024). There- fore, employ attention masks in every decoder block consistently guide query updating and
dareground-backgron
18

---

Published as a conference paper at ICLR 2026
feature refinement. We take in-context segmentation are used. In this setting, we obtain a 1-D query set Q" E R=, where the first 2x queries corre- spond to the foreground and the remaining 2 queries correspond the background. The queries are then updated through a foreground brroouundd--akrruudddd masked attention mechanism. This design en- ables multiple queries to focus target-related image features while also leveraging complementary background information.
example to illustrate how attention masks
2
The foreground-backgrounn masked cross-attention at the (-th decoder layer is formulated as:
Q (oftmax(M; Q V + Qi
(22)
where Qi denotes the transformed queries at layer 1, and K, Vi are the linear projections of the key features Fi. We omit the standard vc scaling term for brevity. The binary mask
Mi E {0,- I is used to enforce the foreground-backgrruun 1, we first predict a probability mask Pt E [0. 10.14xxx,, where each value represents the foreground probability at spatial location (), which is obtained from the pixel features Fi-1 of the previous layer using a lightweight mask decoder followed by resizing. The n-th mask M) at spatial location (f,j) is defined as:
separation. At each layer
0, M((,))
P((i,) 0.5 and n 2
2
O, P((i,) < 0.5 and n 2 2
(23)
2
-00, otherwise.
The same procedure is applied in mode-3 interactive segmentation, where positive and negative click queries represent the foreground and background masks, respectively. When the number of positive and negative clicks is unbalanced, dummy clicks are added to equalize the query set. For semantic segmentation, only the foreground (positive) mask is applied.
Other hy pyeepaarameters. For semantic segmentation, the number of 1-D queries is set to p = 2, while in context segmentation we use nt, = 6. We observe only minor performance differences when varying the query number. In the transformer decoder, all feature channels are fixed at 256, and each attention block uses 8 heads. We set the number of experts M 5 in main experiments. During training, single image-mask pair is employed as the reference.
A.4 MORE EXPERIMENTS
Analysis of expert weights Figure illustrates the distribution of softmax expert weights across different modes when tested on external ACDC dataset. We observe that the weighting patterns vary substantially by mode, reflecting the task-specific routing behavior of the MoE decoder. In Mode-1, 1, Expert 5 receives the highest weight, suggesting specialization toward that expert, while the other experts contribute more evenly. In contrast, Mode-2 strongly favors Expert 1, with minimal reliance on Experts 4 and 5, indicating a different specialization strategy. Mode-3 shows more balanced distribution across all experts, with moderate weights assigned consistently. These results confirm that the model dynamically allocates expert capacity depending on the input mode, which supports the effectiveness of the MoE design in capturing heterogeneous segmentation requirements.
Analysis of convergence curves of all datasets. Figure shows how interactive segmentation (Mode-3) performance of K-Prism changes as the number of clicks increases across all datasets. In most cases, accuracy improves quickly with just few clicks and then gradually levels off, which highlights the benefit of interactive refinement. Our model consistently achieves the strongest results on the majority of datasets, keeping clear advantage over baselines across the entire interaction range. We also note some variation: certain datasets reach high accuracy after only few clicks, while others need more interactions to stabilize, likely due to differences in image quality, structural complexity, segmentation difficulty.
Analysis of molti-mode training. Table 7(a) compares joint training with single-mode training under the same ablation setting as used in the main paper, where models are trained on the M&Ms and Breast Cancer dataset and evaluated on ACDC and BUS. We find that our unified framework achieves performance on par with or better than training each mode individually, while being much more efficient since a single model handles all paradigms simultaneously. This indicates that knowl- edge sharing across modes is beneficial: semantic and in-context learning provide complementary
19

---

Published as a conference paper at ICLR 2026
Expert Weights by Mode
0.6
Mode-1
Mode-2
Mode-3
0.5
Weight
0.3
0.2
0.1
0.0
2
Expert Index
Distribution of softmax expert weights across different modes on the external ACDC dataset.
Figure
Table 7: Extensive ablation studies.
(a) Joint training achieves performance comparable to single single-modc mode train-
ing, while being three times more efficient.
Semanntc
In-context
Interactive
Method
Dice
Dioe
NoC90 | NoC95 & Dice(1) f Dice(5)
81.28
Mode-1
80.57
Mode-2 only
80.84
Mode only
2.34
4.97
86.55
93.68
(b) As the number of experts increases. model performance improves. albeit at the cost of larger number of parameters.
Semantic In-context
Interactive
Expert # Params
Dice
Dice
NoC90 J NoC95 Dice(1) 1 Diee(5) t
(Ours)
43.29M
81.28
79.21
2.31
4,80
86.76
93.79
39.34M
79,80
77.97
2.38
4.96
86.02
93.43
35.39M
79.24
76.81
2.58
5.03
84.24
92.70
31.43M
77.62
75.55
2.64
5.12
84.16
92.22
supervision that improves overall representation quality, and interactive refinement benefits from the shared backbone. These results highlight the practicality of our unified design, which avoids train- ing separate models yet still delivers competitive or superior accuracy across different segmentation settings.
Analysis of number of experts. Table 7(b) analyzes the effect of varying the number of experts the MoE decoder. We observe consistent performance improvement as the number of experts increases from 2 to 5 across semantic, in-context, and interactive settings. Dice scores steadily rise while the number of clicks required to reach target accuracy decreases. These gains, however, come with larger parameter count and additional computational overhead. Due resource constraints, were unable to scale beyond five experts, but the trend suggests that larger expert capacity may yield further benefits. This indicates strong potential for scaling K-Prism into even powerful seg-
20

---

Published as a conference paper at ICLR 2026
Table 8: Throughput (FPS) with different numbers of experts in MoE decoder on single A GPU.
# Experts
5
4
3
2
1 (Plain)
FPS
3.63
4.24
4.61
5.10
5.38
Table 9: Sensitivity of in-context segmentation to reference exemplar selection. Results are averaged over 10 independent 1-shot evaluations per dataset.
Dataset
Mean
Std
Min
Max
ACDC
83.68
2.38
79.10
86.18
BUS
78.43
0.72
77.09
79.15
mentation models, potentially improving initialization robustness and reducing annotation workload in clinical workflows. To further quantify inference efficiency. we evaluate throughput on the ACDC dataset using a single A100 GPU across different numbers of experts (Table 8). The results show that increasing the number of experts introduces only moderate overhead: the throughput decreases by roughly 1.7 FPS (frames per second) when moving from a plain decoder (1 expert) to 5-expert MoE decoder. Importantly. in practical clinical deployments that typically operate on multi-GP! servers, throughput scales nearly linearly with available GPUs, meaning that adding more experts does not compromise real-time usability.
Analysis of reference exemplar sensitivity in in-context segmentation. To further assess the bustness the in-context segmentation mode (Mode-2), we evaluate how performance varies with different choices of reference exemplars. We conduct experiments datasets: BUS (2D ultrasound) and ACDC (3D cardiac MRI). For each dataset, we randomly select 10 distinct reference exemplars (independent subjects/images) Using each exemplar as the sole ref- erence, we run one-shot in in-contex segmentation on all query images, yielding 10 independent eval- uations per dataset. Table reports the resulting mean, standard deviation, and ciciimmmmmxiimin range across these evaluations. The BUS dataset shows extremely low variance (78.43 + 0.72), in- dicating that Mode-2 predictions remain highly stable across different reference choices in typical 2D settings. For ACDC, where anatomical and positional variability across subjects is inherently larger, the variance somewhat higher (83.68 133.68-22 183.68812.38 2.38) but still falls within moderate and acceptable range. Overall, these results demonstrate that K-Prism is robust reasonable variations in reference exemplar quality and does not depend on a highly specific exemplar to achieve strong in-context segmentation performance.
two representative external
Analysis of reference quantity for few-shot in-context segmentation. To further examine the Mode-2 bottleneck in complex multi-organ scenarios, conduct few-shot in-context study on the AMOS MRI dataset. Because axial slices in 3D multi-organ volumes exhibit substantial anatomical variability, a single reference exemplar may misalign with many target slices. To assess whether additional exemplars can mitigate this issue, we evaluate 1-, 3-, 5-, and 9-shot settings, where refer- ence slices are sampled from fixed axial locations corresponding to typical anatomical positions of the target organ within the 3D volume (1-shot: 50%; 3-shot: 20/50/80%; 5-shot: 20/40/50/60/805; 9-shot: 10-90% at 10% intervals). As shown in Table 10, accuracy improves consistently with more exemplars, rising from 78.21 % (11%%)111)))111 (1-shot) to 79.90% (9-shot). This shows that the performance drop observed under strict I-shot conditions is not an inherent limitation of our framework: when ad- ditional exemplars are provided, K-Prism can effectively exploit them. These results highlight the adaptability and robustness of K-Prism's in-context mode in challenging multi-organ segmentation scenarios.
A.S VISUALIZATION
Figures 6,7 and 8 present qualitative comparisons of semantic, in-context, and interactive segmen- tation across diverse datasets and modalities. For semantic segmentation (Figure 6), K-Prism pro- duces more accurate and consistent results than competing methods such as Clip-driven, UniSeg, and Hermes, particularly on challenging tumor and pathology cases. For in-context segmentation (Figure 7), K-Prism achieves clearer boundaries and higher Dice scores across both CT/MRI (e.g..
21

---

Published as a conference paper at ICLR 2026
Table 10: Few-shot in-context segmentation results on AMOS MRI. Increasing the number of ref- erence exemplars improves performance in multi-organ scenarios.
Shots
5
Dice (%) 78.21 79.22
79.70
79.90
AMOS, M&Ms) and endoscopic datasets (c.g., BKALPOLY). For interactive segmentation (Fig- ure 8), we visualize results at the fifth click, showing that K-Prism converges faster and yields more precise masks compared with strong baselines such nnInteractive, MultiverSeg, SAM2, and Seg- Next. Overall, these visualizations demonstrate the robustness and versatility of K-Prism across three segmentation paradigms.
A.6 ANALYSIS OF FAILURE CASES FOR K-PRISM
In the BraTS example (Figures 9, top-left), the model mistakenly segments the bright peritumoral edematous/invaded tissue (Yousef et al., 2023) as tumor. Because the reference image contains rounded, high-intenssty tumor with a compact mass-like morphology and the peritumoral edema- tous/invaded tissue signal in the query slice exhibits superficially similar intensity distribution, the appearance-driven matching in Mode-2 incorrectly aligns this region with the tumor exemplar. Cru- cially, brain MRI is completely unseen domain for model: brain anatomy, texture patterns, or disease manifestations appear in the training datasets. As result, the model lacks any semantic prior to distinguish true oncologic tissue from peritumoral edematous/invaded tissue hyperintensi- ties. Under such severe domain and anatomical shifts, Mode-2 naturally over-relies on low-level intensity correspondences in the absence of domain knowledge, causing misleading but internally consistent matches.
In the M&Ms-2 example (Figures bottom-left), the model incorrectly segments part of the left atrial wall as left ventricle myocardium. M&Ms-2 long-axis cardiac MRI has never appeared in the training data, and its anatomical geometry differs substantially from the short-axis cardiae MRI datasets used for model development. Importantly, the left atrial wall and the left ventricular my- ocardium exhibit highly similar shaped appearances in the long-axis view. Because Mode-2 relies heavily appearance-based alignment between the reference exemplar and the query slice, the model misled by this shape similarity. As a result, the model aligns the atrial wall with the "my- ocardium" region from the exemplar, despite the anatomical mismatch.
A.7 THE USE OF LARGE LANGUAGE MODELS (LLMs)
In this work, large language models (LLMs) were used solely as general-purpose writing assistants for text polishing and formatting adjustments. They did not contribute research ideation, experi- mental design, analysis, or result interpretation.
A.8 ETHICS STATEMENT
This research adheres to the ICLR Code of Ethics and its guiding principles of responsible steward- ship, fairness, and transparency. All experiments are conducted exclusively on publicly available, de-identified datasets. No identifiable patient data were collected or generated, and all datasets in- clude appropriate ethical approvals from their original organizers. Our goal to advance trustworthy and socially beneficial Al for medical image analysis. While the proposed framework demonstrates strong performance, it is intended for research use only and not for direct clinical deployment with- out further validation. We acknowledge that premature use of automated segmentation could in- troduce risks, including potential bias across populations, imaging protocols, or disease types. To mitigate these concerns, we evaluate extensively across diverse datasets and modalities and high- light limitations in disttriibitiitiinn generalization. We also stress the importance of faimess, reproducibility, and responsible application in any future clinical translation.
22

---

Published as a conference paper at ICLR 2026
A.9 REPRODUCIBILITY STATEMENT
We take reproducibility seriously and provide extensive details throughout the main text and Ap- pendix, including dataset descriptions, preprocessing steps, model architecture, training hyperpa- rameters, and evaluation protocols. Comprehensive ablation studies are included to highlight the contribution of each component. All experiments are conducted on publicly available datasets. The full codebase is publicly available to support replication and further research.
23

---

Published as a conference paper at ICLR 2026
AMOS MAI
AMOS
Med
Clic
LAScar0S
diesdiv SMA
Auriee
Xray
Click Kurboy
Clici
RAPILA
DKAI POLY
Ciic Number
Click
Modilal
gand
Rumbor
8US
BrOTS
SM4
Click Aumbee
Figure 5: Convergence curves of K-Prism's interactive segmentation (Mode-3) all datasets.
24

---

Published as a conference paper at ICLR 2026
Input inage
Clig-driven
UniSeg
Hemnes
K-Prisem
Ground truth
(Tumor)
KITS
Dioe: 91, 1.86%
Dice: 85 22%
Dier: $7 68%
Dicu: 93.27%
KPI (Glomenili)
Diez: 94%
Diee: 92, 90%
Dice: 94.1 14.1%%
Dice: 94,47%
Breast cancer (Tunce)
Dice: 89.87%
43%4%%
88.10%
Figure 6: Qualitative comparison of semantic segmentation (tumor), KiTS (tumor), KPis (glomeruli), and Breast cancer (tumor). From left to right: input image, predictions from universal models Clip-driven, UniSeg, Hermes, and K-Prism, and the ground truth. Reported Dice scores highlight that K-Prism produces more accurate and consistent results across diverse modalities and targets.
representative datasets: LiTS
MultivedSog
(Right Xidneyy
83,49%
83%
90 80%
M&M& (Myocardhum)
Dicr: M0 MMS
Dice: 83.22%
Dice: 61.31%
Dice: 53%
90.94%
RKAL POLY (Tumor)
Dioee 69 00%
Dioe: 68 90%
Dooe 69 20%
Dice: 67.65%
Dice: 93.23%
HTCV (Liven
84.27%
85.38%
87%%%
Dise: 39%
Figure 7: Qualitative comparison of in-context segmentation across four representative datasets: AMOS MRI (right kidney), M&Mss (myocardium). BKAI POLY (tumor), and BTCV (liver). From left to right: input image, predictions from UniverSeg, Tyche, MultiverSeg, Iris, and K-Prism, and the ground truth. Reported Dice scores show that K-Prism achieves the most accurate and consistent results across diverse modalities and targets.
25

---

Published as a conference paper at ICLR 2026
lepet image
MadsAM
SAMZ
Siahhen
K-frion
Grcend tratt
(tunee)
WeMled
(Leff Vantrole)
$4.9%
Dioz 94.58%
93.75%
20%
50%
Figure 8: Qualitative comparison of interactive segmentation results the fifth click across four representative datasets: BraTS (tumor), BUS (tumor), M&Ms-2 (left ventricle), and UW-SC (skin cancer). From left to right: input image, predictions from nnInteractive, MedSAM, MultiverSeg, SAM2, SegNext, and K-Prism, and the ground truth. Reported Dice scores show that K-Prism consistently produces the most accurate and reliable segmentations across diverse modalities and targets.
Query image
Reference image
Imuge
Ground truth
Prediction
Image
Ground truth
BraTS (Tumor)
Dice: 6.67%
Image
Ground truth
Prediction
Image
Ground truth
M&Ms-2 (Myocardium)
Figure 9: Representative failure cases of in-context segmentation (Mode-2) unseen datasets: BraTS and M&Ms-2.
26


Reviews:
[{"review_id": "si7V8h1zjm", "reviewer": "Reviewer_ZEBx", "summary": "This paper presents K-Prism, a unified framework for medical image segmentation that systematically integrates three clinically relevant knowledge paradigms: (i) semantic priors learned from annotated datasets, (ii) in-context knowledge from few-shot reference examples, and (iii) interactive feedback from user inputs.\nThe key design lies in a dual-prompt representation—1-D sparse prompts encoding “what to segment” and 2-D dense prompts encoding “where to attend”—together with a Mixture-of-Experts (MoE) decoder for dynamic, task-aware routing.\nK-Prism supports three operational modes (semantic, in-context, interactive) within a single architecture and achieves state-of-the-art results across 18 datasets spanning CT, MRI, X-ray, ultrasound, pathology, and fundus imaging. The model demonstrates strong generalization to external and unseen-class datasets and enables efficient refinement through user interaction.", "questions": "- Could the framework be extended to 3D volumetric segmentation without major architectural changes?\n- How significant is the computational cost of the MoE decoder compared to standard transformer-based decoders (e.g., in inference latency)?\n- Is there any strategy to automatically balance or select the optimal knowledge source (semantic / in-context / interactive) during inference, especially when user feedback is unavailable?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper addresses a clear and practical gap: the fragmentation of current medical segmentation models across different knowledge paradigms.\n- The dual-prompt + MoE design is conceptually elegant and technically sound, enabling unified training and inference across tasks without architectural modifications.\n- Experiments are comprehensive and convincing, covering 18 public datasets, multiple modalities, and both in-distribution and cross-dataset evaluations.\n- K-Prism achieves consistent SOTA performance across semantic, few-shot, and interactive settings, with particularly impressive efficiency in interactive segmentation (lowest NoC90 and NoC95).\n- Ablation and analysis are thorough, confirming the contribution of each component and the dynamic specialization of experts.\n- The work has practical clinical potential, as it can streamline deployment and annotation pipelines.", "weaknesses": "- All experiments are conducted on 2D slices, while many clinical workflows require 3D volumetric segmentation.\n- The method relies on the quality and availability of reference examples for in-context mode, which may constrain performance in data-scarce or noisy scenarios.\n- Discussion on model interpretability or failure cases is relatively limited.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VobQXM9KKM", "reviewer": "Reviewer_MRgi", "summary": "The authors propose K-Prism, a unified segmentation framework capable of performing (1) standard segmentation on in-class labels, (2) in-context learning–based segmentation using reference exemplars, and (3) interactive segmentation through user feedback. This is achieved by learning a 1-D sparse prompt that encodes what to segment and a 2-D dense prompt that encodes where to attend, which are jointly processed through a Mixture-of-Experts decoder for dynamic, task-aware feature routing. Experiments conducted on 18 public datasets spanning CT, MRI, X-ray, pathology, and ultrasound demonstrate that K-Prism achieves state-of-the-art performance across all three segmentation paradigms, with generalization to unseen classes and external datasets.", "questions": "1. The model shows noticeably lower accuracy on unseen anatomical classes compared to in-distribution results. Could the authors provide an analysis explaining why this gap occurs?\n2. To better contextualize the proposed framework, can the authors include experiments comparing K-Prism to nnU-Net trained (a) jointly on all datasets as a generalist model, and (b) separately on each dataset as task-specific specialists? \n3. The paper mentions that all competing methods are trained on the same dataset, but it is unclear how interactive baselines (e.g., SAM2, MultiverSeg) were handled. Were these models trained or fine-tuned from scratch, or initialized from existing pretrained weights?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. K-Prism unifies semantic, in-context, and interactive segmentation within a single framework through a novel design that leverages 1-D and 2-D prompts to encode complementary information about what and where to segment. \n2. The use of a Mixture-of-Experts decoder is well-motivated, enabling task-aware specialization and effective fusion of different prompt types while maintaining a shared representation space.\n3. The paper is well-written and easy to follow, with clear methodological explanations and thorough experimental validation that make the work accessible and reproducible.", "weaknesses": "1. While K-Prism performs well on in-distribution and cross-dataset tests, its accuracy drops notably for unseen anatomical classes, suggesting limited generalization beyond trained label types.\n2. The paper does not compare against nnU-Net trained under equivalent conditions—either as a single generalist model trained on all datasets or as individual task-specific models. Such comparisons are important to validate that the proposed unification truly offers an advantage over established supervised baselines.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xn4cdPgHX5", "reviewer": "Reviewer_3nT9", "summary": "The paper proposes K-Prism, a unified medical image segmentation framework that integrates three knowledge sources—semantic priors, few-shot in-context examples, and interactive user feedback. Using a dual-prompt and Mixture-of-Experts design, K-Prism flexibly adapts across tasks and modalities. Tested on 18 public datasets, it achieves state-of-the-art performance in semantic, in-context, and interactive segmentation.", "questions": "1. Authors should consider adding MedSAM for comparison\n2. As shown in Table 2 and 3, the improvement is limited", "limitations": "", "rating": 2, "confidence": 5, "soundness": 1, "presentation": 2, "contribution": 1, "strengths": "1. Unified design integrating semantic, in-context, and interactive segmentation.\n2. Consistent performance gain across diverse datasets and modalities.", "weaknesses": "The novelty of the paper is somewhat unclear. It appears to mainly combine existing ideas from interactive segmentation models (e.g., MedSAM) and in-context segmentation models (e.g., UniverSeg). It is not fully evident what new conceptual or methodological insight the proposed framework contributes beyond this integration. The authors should clarify the unique innovation or theoretical advancement that distinguishes K-Prism from prior works.\n\nThe three proposed “knowledge paradigms” — semantic priors, in-context examples, and interactive feedback — are somewhat loosely defined and may not convincingly qualify as distinct forms of knowledge. They are more simply like working mode instead of knoledge.The authors should better justify why these modes are conceptualized as “knowledge sources”.\n\nThe motivation for adopting a Mixture-of-Experts (MoE) design is insufficient. In medical segmentation, computational efficiency is typically not a critical bottleneck, so replacing a dense model with MoE seems unnecessary and possibly superficial.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Oyz7Wt9moE", "reviewer": "Reviewer_GsDa", "summary": "This work propose K-Prism, which is a unified framework for medical image segmentation that innovatively integrates three knowledge paradigms into a single model: semantic priors (Mode-1), in-context exemplars (Mode-2), and interactive feedback (Mode-3). Its core technique encodes heterogeneous knowledge into a unified dual-prompt representation (1D sparse prompts for \"what\" and 2D dense prompts for \"where\") , which is then processed by a Mixture-of-Experts (MoE) decoder with dynamic routing. Experiments on 18 public, multi-modal datasets demonstrate SOTA performance across all three segmentation settings.", "questions": "- 3D Extension: What are the main challenges in extending this framework to 3D? Specifically, can the dual-prompt and MoE mechanisms be efficiently scaled computationally?\n\n- In-context Bottleneck: Could using k-shot (k>1) references in Mode-2 help alleviate the performance drop observed in multi-organ tasks?\n\n- MoE Expert Specialization: What specific functions have the MoE experts specialized in? (e.g., are some experts dedicated to positive clicks and others to negative clicks?)\n\n- Reference Quality: How sensitive is the performance of Mode-2 (in-context) to the quality of the reference exemplar (e.g., inaccurate mask) or domain shift (e.g., different scanner)?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Significant Problem: Directly addresses the long-standing problem of model fragmentation in medical segmentation.\n\n- Novel Architecture: The proposed \"dual-prompt\" and MoE decoder design is elegant, successfully unifying three heterogeneous knowledge sources.\n\n- Comprehensive Experiments: Extensively validated on 18 public datasets covering diverse modalities.\n\n- SOTA Performance: Outperforms existing state-of-the-art methods in all three tracks: semantic, in-context, and interactive segmentation.", "weaknesses": "- 2D Limitation: The model currently processes 3D medical data (like CT/MRI) based on 2D slices, which loses 3D contextual information.\n- In-context Bottleneck: For complex multi-organ segmentation tasks, the in-context mode (Mode-2) performs noticeably worse than the semantic mode (Mode-1).\n- Computational Cost: The MoE architecture, while effective, introduces a higher parameter count and computational overhead, potentially hindering real-time deployment.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)