- The paper introduces APEX, a framework that generates input-specific prompts from low-frequency image features to robustly adapt segmentation models across diverse domains.
- It employs a novel prompt memory module and Low-frequency Feature Contrastive learning to enhance the discrimination of domain features without re-training backbone weights.
- Empirical evaluations on polyp and optic segmentation tasks reveal significant improvements in Dice and IoU metrics over traditional adaptation and static prompting methods.
Adaptive Visual Prompting for Robust Medical Image Segmentation Across Domains
Introduction
Domain shift remains a central hindrance in deploying DNN-based medical image segmentation methods in clinical practice due to significant intra- and inter-domain data variability, often introduced by different imaging devices, acquisition protocols, or pathology-specific appearances. While domain adaptation (DA) strategies using fine-tuning or domain transfer attempt to address these challenges, they suffer from catastrophic forgetting of source knowledge and are regularly impractical due to stringent data access and privacy constraints. Visual Prompting (VP) has recently emerged as an alternative, enabling adaptation by introducing learnable image-space parametersโso-called promptsโoptimized for the target domain without updating backbone weights. However, existing VP strategies mostly utilize a single, domain-level prompt per domain, thus inadequately accounting for intra-domain heterogeneity and showing weak generalization to unseen domains.
This work introduces APEX (Adaptive Prompt EXtraction), a framework for adaptive visual prompting in medical image segmentation, which leverages a learned prompt memory to dynamically generate input-specific prompts. The framework advances prior VP paradigms by addressing both intra-domain and inter-domain variability through frequency-domain conditioning and a novel Low-frequency Feature Contrastive (LFC) learning objective.
Figure 1: Comparison of conventional domain-level prompting versus the proposed input-adaptive prompting workflow in APEX.
Adaptive Prompt Extraction: Architecture and Rationale
APEX is composed of a domain feature encoder, a learnable prompt memory, and a prompt decoder. The fundamental innovation is the retrieval of input-specific prompt features by querying the prompt memory with image-dependent domain features, extracted from the low-frequency spectrum of the image using FFT. This is motivated by the finding that low-frequency amplitude components in medical images predominantly encode global style and domain shifts, whereas high-frequency information embodies anatomical content and spatial structure. Consequently, prompting is applied in the low-frequency spectral band, ensuring effective adaptation without distorting anatomy-critical details.
During inference, for each input, the domain feature encoder produces a feature vector by processing the image's low-frequency amplitude. This representation is then used to generate a cosine-similarity-based attention vector over the slots of prompt memoryโa bank of orthogonally initialized prompt vectors. The final prompt is synthesized as a linear combination of these memory vectors, weighted by the attention, and then decoded to the prompt applied in the frequency domain.
Figure 2: Schematic of the APEX framework, showing the retrieval pipeline from frequency domain input through prompt generation and application.
Low-frequency Feature Contrastive (LFC) Learning
To optimize the capability of the domain encoder in capturing discriminative, domain-relevant information, APEX integrates LFC learning. This objective clusters low-frequency features within the same domain while maximizing distance between features from different domains, leveraging a supervised contrastive loss with temperature scaling in the latent space. To prevent representational collapse, an auxiliary projection head is employed during training.
Figure 3: Overview of LFC learning, illustrating clustering of domain features and inter-domain separation in the frequency domain.
The integration of LFC substantially enhances the fidelity of domain representations, which directly contributes to instance-adaptive retrieval from memory and diverse prompt composition.
Empirical Evaluation
APEX was evaluated on polyp and optic cup/disk segmentation tasks, encompassing diverse datasets reflecting significant real-world domain discrepancy. The framework was benchmarked using multiple segmentation backbones (PraNet, UNet, ResUNet, SwinUNet, TransUNet). Compared to SOTA prompting (VPT, FVP, A2XP, VPAD) and source-free adaptation baselines, APEX demonstrates the following:
- Consistent improvements in DICE and IoU across both seen and unseen domains for all backbone networks.
- For PraNet on the polyp task, APEX achieved gains of up to 2.70% (seen) and 3.33% (unseen) in DICE over baseline, and outperformed VPAD by 1.91%โ3.33%.
- For UNet on the OC/OD task, APEX reached DICE improvements of 4.84% (seen) and over 41% (unseen) compared to baseline.
These results provide quantitative evidence that instance adaptation via memory-composed prompting robustly generalizes beyond the seen/known distribution, outperforming both static VPs and continuous test-time adaptation methods, while requiring minimal additional computational overhead.
Figure 4: Qualitative segmentation results before and after prompt application, highlighting completion of ambiguous regions post-prompting.
Ablation and Analytical Insights
Detailed ablation underscores the contributions of each component:
- Prompt Memory: Removal of memory results in an average Dice drop of >3% on seen and unseen domains, confirming the advantage of diverse, input-specific prompt composition.
- LFC Loss: Excluding LFC learning consistently reduced generalizability, especially to unseen domains, corroborating its role in learning robust, transferable domain features.
- Prompt Memory Size: Increasing slots improves DICE until saturation at 150 slots; over-parameterization leads to distributed, less discriminative memory utilization.
Figure 5: Performance as a function of prompt memory size, demonstrating optimality and saturation at 150 slots.
Prompt memory slot activation analysis reveals that samples from the same domain typically activate overlapping memory regions, reflecting domain-level structure, but also exhibit unique activation for individual samples, supporting flexible intra-domain adaptation.
Figure 6: Visualization of top-activated slots for samples from unseen domains, demonstrating both shared domain-level and sample-specific retrieval patterns.
Practical Implications and Future Directions
APEX achieves adaptation without access to source data, with negligible parameter and inference overhead, satisfying key constraints in clinical deployment. In practical comparison with test-time adaptation paradigms (e.g., VPTTA), APEX shows comparable or superior accuracy at orders-of-magnitude faster speeds (20ms vs. 3000ms per image CPU), without accumulation risk. This validates its utility as a plug-and-play module for robust segmentation across distributed medical institutions.
Theoretically, conditioning prompts with frequency-domain discriminative features delivers a general framework for instance-based adaptation, extensible to other visual tasks inherently sensitive to style/appearance variation. Notably, the clear separation of prompting in low-frequency (domain) versus high-frequency/phase (structure) also suggests a pathway for disentangled adaptation to anatomical versus device-induced domain shifts.
A future research trajectory involves extending APEX to adaptation across shifts in anatomical structure by leveraging prompt design in phase or high-frequency bands, as well as integrating more advanced memory update and retrieval strategies possibly informed by self-supervision or continual learning principles.
Conclusion
APEX establishes a new paradigm for adaptive, memory-based visual prompting in medical image segmentation by generating input-specific prompts from domain-discriminative, frequency-domain features. The framework achieves robust adaptation without catastrophic forgetting or source data exposure andโthrough LFC learningโenables improved clustering and discrimination of intra- and inter-domain variation. Empirical results show consistent, strong increases in both accuracy and generalization across a range of datasets and architectures. APEX advances the state of plug-and-play domain adaptation while maintaining strict operational constraints critical to translational medical AI.