This paper introduces MedSAMix, a training-free model merging approach designed to enhance medical image segmentation by integrating the strengths of generalist models like SAM and specialist models like MedSAM. MedSAMix employs a zero-order optimization method to automatically discover optimal layer-wise merging configurations, aiming to improve both domain-specific accuracy and generalization without requiring retraining. Extensive evaluations on 25 medical segmentation tasks demonstrate significant performance improvements.
Problem
The paper addresses several bottlenecks in current medical image segmentation approaches:
Unbalanced performance of fine-tuned models: Fine-tuned models, while showing promise, often exhibit unbalanced performance across different medical segmentation tasks, sometimes underperforming even the original generalist model.
Limited adaptability and generalization: These models can be susceptible to suboptimal generalization due to the inherent complexities of medical imaging data, such as heterogeneity, domain shifts, and class imbalance, leading to convergence on suboptimal local minima.
Catastrophic forgetting: Fine-tuned variants can suffer from catastrophic forgetting during adaptation, potentially losing their original broad segmentation abilities.
Inefficiency and suboptimality of existing model merging methods: Traditional model merging approaches often rely on manual configuration, leading to suboptimal outcomes, or require computationally intensive merging during training, which is costly for large foundation models. They also typically lack support for multi-objective optimization, hindering the generalization ability of merged models.
Key Contributions
MedSAMix, a training-free model merging method that leverages the strengths of general and expert knowledge for medical image segmentation.
Flexible merging regimes within MedSAMix: single-task merging for domain-specific tasks and multi-task merging for improved generalization across diverse tasks.
Empirical demonstration that MedSAMix enhances expert-level performance and generalization without retraining, achieving significant improvements on specialized and multi-task evaluations.
Method
MedSAMix is a training-free model merging framework that operates at the model level. It aims to balance domain-specific capabilities with generalization by integrating multiple SAM-based models. The framework consists of three core components:
Search Space: The search space is designed for SAM-series models, which are based on the Vision Transformer architecture. It allows for fine-grained control over merging strategies by supporting multiple merging methods at varying layer granularities across the image encoder, prompt encoder, and mask decoder. This space is defined by merge methods and associated hyperparameters for each layer group.
Optimization Objectives: MedSAMix accommodates both single-task optimization, focusing on maximizing performance on a specific target task using metrics like the Dice coefficient, and multi-objective optimization using the ParEGO algorithm to identify Pareto-optimal configurations across multiple tasks.
Search Algorithm: The framework utilizes the SMAC (Sequential Model-based Algorithm Configuration) optimization algorithm with a Random Forest surrogate model. This Bayesian optimization approach efficiently navigates the search space by selecting promising configurations based on their empirical performance, aiming to find optimal merging configurations. The process involves initializing a search space, evaluating configurations, and iteratively updating a history of evaluated configurations to guide the search towards the optimum.
Technical Details
Architecture: Candidate models are based on the Segment Anything Model (SAM) architecture, which uses a Vision Transformer. The SAM architecture includes an image encoder with l transformer layers, a prompt encoder with k convolutional downsampling layers, and a mask decoder with z transformer layers and p transposed convolutional layers. The merging framework defines granularity for the image encoder (genc), prompt encoder (gprompt), and mask decoder (gdec), dividing layers into groups for merging.
Key Equations:
Single-task optimization objective: fsingle (Mw) = L(Mw, DT) where L is the segmentation loss and DT is the target dataset.
Multi-task optimization objective: fmulti (Mw, λ) = max {fsingle,i(Mw)} for i=1,...,m + α Σ λᵢ · fsingle,i(Mw) where fsingle,i is the i-th objective function, λᵢ are weights, and α is a small positive constant.
SMAC Expected Improvement: EI(w) = E[max(f* − f(w), 0)] where f* is the best observed value and f(w) is the Random Forest prediction.
Loss/Objective: The primary objective is to maximize segmentation performance, measured by the Dice coefficient or other segmentation loss functions. For multi-task optimization, Pareto Efficient Global Optimization (ParEGO) is used.
Optimization: SMAC (Sequential Model-based Algorithm Configuration) with a Random Forest surrogate model is used for Bayesian optimization.
Augmentation: Not stated in the paper.
Data Handling: Calibration datasets (DT) are used for evaluating merging configurations. For multi-task optimization, a subset of eight tasks is selected.
Implementation Details: The base architectures of SAM, MedSAM, and MedicoSAM are used. The model has l=12 Transformer layers, 768 hidden features, 12 heads, k=4 prompt encoder layers, and a mask decoder with z=2 Transformer layers and p=3 transposed convolutional layers. Candidate merging techniques include TIES, task arithmetic, linear combination, and SLERP. For single-task optimization, 120 trials are performed using two GPUs. For multi-task optimization, 200 trials are conducted on four GPUs. Layer granularity is searched within the range [1, 4].
Results
Datasets: 25 publicly available medical image segmentation datasets are used, covering a wide range of organs and imaging modalities, including brain tumor, vascular, optic disc/cup, liver, kidney, spleen, pancreas, maxillary sinus, nasal cavity, nasal pharynx, cardiac structures, and prostate segmentation. Datasets include BraTS, Topcow, Fundus, FLARE22, RAOS, Nasal, MSD, PROMISE, Mice, and ADNI.
Metrics: Dice coefficient score is used for evaluation.
Baselines: Comparisons are made against ICL-based models (UniverSeg, SegGPT, Neuroverse3D), SAM-based models (SAM-Med2D, SAM, MedSAM, MedicoSAM), and the fully supervised nnU-Net as an upper bound. Other model merging baselines include Task-Arithmetic (TA), TIES, SLERP, and linear weighted combination.
Numeric Results:
Single-task (MedSAMix-S): Achieves significant improvements over SAM and other fine-tuned models. For example, it outperforms MedicoSAM by 5.19% on brain tumor and 4.33% on Nasal Pharynx tasks. It also shows improvements of around 1.2% on optic disk segmentation.
Multi-task (MedSAMix-M): Outperforms baselines on 18 out of 25 tasks. It achieves the best overall performance among universal segmentation models with an average improvement of 4.37% over the second-best model.
Model Merging Baselines (Table 2): MedSAMix-S achieves a 6.67% average Dice coefficient improvement, and MedSAMix-M achieves a 4.37% improvement over the MedicoSAM baseline across 25 tasks.
Runtime/Training Cost: MedSAMix is training-free. For optimization, 120 trials are performed using two GPUs for single-task and 200 trials on four GPUs for multi-task optimization.
Weaknesses / Limitations
The study, despite covering 25 diverse tasks, may not fully represent the entire landscape of medical image segmentation.
Future work is suggested to expand task diversity and develop systematic methods to assess generalization and robustness.
Evidence
Method evidence: "We propose MedSAMix, a training-free model merging method that integrates the strengths of both generalist models (e.g., SAM) and specialist models (e.g., Med-SAM) for medical image segmentation. In contrast to tradi- tional model merging approaches that rely on manual con- figuration and often result in suboptimal outcomes, we pro- pose a zero-order optimization method to automatically dis- cover optimal layer-wise merging solutions." (Abstract)
Result evidence: "Extensive evaluations on 25 medical segmentation tasks demonstrate that MedSAMix effectively mitigates model bias and consistently improves performance in both domain-specific accuracy and generalization, achieving improvements of 6.67% on specialized tasks and 4.37% on multi-task evaluations." (Abstract)
Limitation evidence: "Limitation. We acknowledge that, despite covering 25 di- verse tasks, our study cannot fully represent the entire land- scape of medical image segmentation. Future work will ex- pand task diversity and develop systematic methods to assess generalization and robustness." (Discussion and Conclusion)
Additional Figures
Figure 1: Segmentation performance of different SAM- based models across diverse medical tasks, evaluated using the Dice coefficient. Blue lines indicate the superior results between MedSAM and MedicoSAM, while green lines de- note the performance of the original SAM. Although Med-
Figure 2: Overview of our model merging framework. Given a pool of models, MedSAMix searches for optimal merging con- figurations by using single-task or multi-task performance as rewards. While individual models may exhibit varying behaviors across tasks, MedSAMix adaptively combines them to optimize performance for the target task. training coverage, which contribute to their superior gener-
Figure 3: Visual examples of segmentation results of implemented models on pancreas, mouse lung, and brain vascular datasets. with coarse granularity. In contrast, multi-task performance is more dependent on fine-grained control over layer merg- ing. Notably, the best performance in the multi-task setting
Fig. 4 (3). We find that merging MedicoSAM and SAM yields the best pairwise performance, consistent with their strong individual results. Although MedSAM achieves the worst performance among the three models, combining all three models can lead to further improvements, demon-