From Adaptation to Generalization: Adaptive Visual Prompting for Medical Image Segmentation

TL;DR

APEX introduces an adaptive, input-specific prompting framework for medical image segmentation, enhancing model generalization across diverse and unseen domains by leveraging a learnable prompt memory and domain feature contrastive learning.

Contribution

The paper proposes APEX, a novel adaptive prompt extraction method that improves domain generalization in medical image segmentation using a learnable prompt memory and contrastive learning.

Findings

APEX significantly improves generalization on unseen domains.

It enhances performance across multiple medical segmentation tasks.

The method is compatible with various backbone models.

Abstract

Visual prompting has emerged as a powerful method for adapting pre-trained models to new domains without updating model parameters. However, existing prompting methods typically optimize a single prompt per domain and apply it uniformly to all inputs, limiting their ability to generalize under intra and inter-domain variability, which is especially critical in the medical field. To address this, we propose APEX, an Adaptive Prompt EXtraction framework that retrieves input-specific prompts from a learnable prompt memory. The memory stores diverse, domain-discriminative prompt representations and is queried via domain features extracted from the Fourier spectrum. To learn robust and discriminative domain features, we introduce a novel Low-Frequency Feature Contrastive (LFC) learning framework that clusters representations from the same domain while separating those from different domains. Extensive experiments on two medical segmentation tasks demonstrate that APEX significantly improves generalization across both seen and unseen domains. Furthermore, it complements any existing backbones and consistently enhances performance, confirming its effectiveness as a plug-and-play prompting solution in medical fields. The code is available at https://github.com/cetinkayaevren/apex/