When Adaptation Hurts: Connecting Representational Drift to OOD Failures in MedSAM Fine-Tuning
This study systematically compares six fine-tuning strategies for MedSAM in OOD settings, revealing encoder LoRA enhances robustness by stabilizing decoder representations.
Key Findings
Methodology
Six fine-tuning methods (full, encoder-only LoRA, shallow/deep VPT, decoder-only, full fine-tuning) were trained on ISIC 2018. Evaluation included in-domain, close-OOD (PH2), and far-OOD (BUSI, CBIS-DDSM). ROI prompts with 0-100 pixel jitter simulated noise. Internal representations were analyzed via Centered Kernel Alignment (CKA) to relate drift in decoder layers to performance degradation. Experiments measured Dice scores, parameter counts, and computational costs, emphasizing the robustness of encoder LoRA under domain shifts.
Key Results
- Full fine-tuning achieved highest in-domain and close-OOD Dice scores (average +0.02), but performance dropped significantly on far-OOD datasets (-0.05). Encoder-only LoRA maintained better robustness in far-OOD (average -0.043), outperforming standard LoRA (-0.176) and VPT. Introducing ROI jitter during training improved noise robustness, especially under severe perturbations. CKA analysis showed decoder layer representations' drift (correlation >0.75) strongly associated with far-OOD performance decline, while encoder similarity was weakly related.
- Limiting LoRA to encoder layers preserves decoder stability, leading to improved generalization. Although full fine-tuning remains the top performer overall, encoder LoRA offers a promising efficiency-performance tradeoff, especially for deployment in resource-constrained settings.
Significance
This work advances understanding of how fine-tuning strategies influence model robustness and generalization in medical image segmentation. By linking internal representation drift to OOD performance, it provides a theoretical basis for designing more resilient models. The findings suggest that stabilizing decoder representations while adapting the encoder can mitigate performance drops across modalities and institutions, addressing critical barriers for clinical translation. The study also highlights the importance of prompt noise simulation in training, offering practical insights for deploying models in real-world, noisy environments.
Technical Contribution
The paper introduces a CKA-based framework to quantify internal representation stability across fine-tuning methods. It systematically compares six strategies, demonstrating that encoder-only LoRA effectively reduces decoder drift, thus enhancing OOD robustness. The integration of prompt jitter during training and the analysis of layerwise representation similarity provide new insights into the mechanisms underpinning domain generalization. These contributions offer a novel, scalable approach for robust foundation model adaptation in medical imaging.
Novelty
This is the first comprehensive comparison of multiple fine-tuning strategies in medical image segmentation under domain shift, explicitly linking internal representation drift with OOD performance. The innovative use of encoder-only LoRA to stabilize decoder representations and improve robustness is a key novelty. Additionally, the combined analysis of prompt noise and internal features offers fresh perspectives on enhancing model resilience, setting a new benchmark for future research in medical foundation models.
Limitations
- The experiments focus on dermoscopic and ultrasound datasets; applicability to other modalities like MRI or CT remains to be validated. The prompt jitter simulation is simplified and may not fully capture real clinical prompt variability. The study primarily assesses static datasets; real-time deployment scenarios could introduce additional challenges. Although encoder LoRA improves robustness, the approach may still face limitations with extreme domain shifts or highly noisy prompts. Computational costs, while reduced compared to full fine-tuning, could still hinder large-scale clinical deployment without further optimization.
Future Work
Future research should extend evaluations to multi-modal datasets, including MRI and CT, to validate generalizability. Developing adaptive prompt augmentation techniques that mimic clinical variability more accurately will enhance robustness. Combining self-supervised learning with representation stabilization could further improve domain transfer. Exploring dynamic, real-time fine-tuning methods and model compression techniques will facilitate deployment in resource-limited settings. Additionally, deeper analysis of layerwise representation drift mechanisms may inform the design of inherently robust architectures.
AI Executive Summary
Foundation models like MedSAM have revolutionized medical image segmentation, offering zero-shot capabilities across diverse modalities. However, their deployment in real-world clinical settings faces significant hurdles: domain shifts, noisy prompts, and computational costs. This study systematically compares six fine-tuning strategies—full, decoder-only, encoder LoRA, shallow/deep VPT, and zero-shot—evaluating their robustness across in-domain, close-OOD, and far-OOD datasets. The experiments reveal that while full fine-tuning achieves the best in-domain performance, it suffers notable degradation under domain shifts, especially in far-OOD scenarios. Conversely, encoder-only LoRA, a parameter-efficient method, maintains superior robustness by stabilizing decoder representations, as confirmed by Layerwise Centered Kernel Alignment (CKA) analysis. The study further demonstrates that introducing ROI prompt jitter during training significantly enhances noise robustness, a critical factor for clinical applicability. The key insight is that the primary driver of performance decline under domain shift is the drift in decoder layer representations, not encoder similarity. This understanding guides the recommendation that stabilizing decoder features—via encoder-only adaptation—strikes an optimal balance between performance and robustness. The findings have profound implications for deploying reliable, efficient medical segmentation models in diverse clinical environments, paving the way for more resilient foundation models. Limitations include the focus on dermoscopic and ultrasound data, with future work needed to extend to other modalities and real-time scenarios, integrating self-supervised learning and model compression for scalable clinical translation.
Deep Dive
Abstract
Foundation models for medical image segmentation, like prompt-based MedSAM, generalize well across domains and modalities, often in zero or few-shot setups. However, their performance depends on the quality of prompts and the adaptation of the models to custom datasets. This work systematically examines how MedSAM generalizes across diverse medical imaging benchmarks, with six adaptation strategies: full-model and encoder-only LoRA, shallow and deep visual prompt tuning (VPT), and decoder-only and full fine-tuning. Models are trained on the International Skin Imaging Collaboration Challenge (ISIC 2018) dataset and evaluated under clean and increasingly noisy prompts on IN and Out-of-Distribution (OOD) datasets: close-OOD PH2 (dermoscopy), far-OOD BUSI (Breast Ultrasound Images Dataset) and CBIS-DDSM (Curated Breast Imaging Subset of the Digital Database for Screening Mammography). We show that adaptation improves performance on IN and close-OOD data but often reduces performance on far-OOD data. Full fine-tuning provides the best tradeoff, while encoder-only LoRA is the strongest parameter-efficient alternative, outperforming standard LoRA and VPT under far-OOD shifts. Using Centered Kernel Alignment (CKA), we show that far-OOD degradation is strongly associated with drift in decoder representations, whereas encoder similarity alone does not explain robustness. This suggests encoder-only LoRA provides stronger robustness than standard LoRA by adapting the encoder to distribution shift in visual features, while preserving the decoder pathway. We further show that random 0-100 pixel jitter on prompts produces more robust and better performing models. We thus conclude that robust MedSAM adaptation requires the combined consideration of prompt noise exposure, domain shift, and representation preservation. We release our code: https://github.com/ImSounic/medsam-vpt