Mixture-of-Shape-Experts (MoSE): End-to-End Shape Dictionary Framework to Prompt SAM for Generalizable Medical Segmentation

TL;DR

This work tackles the challenge of single-domain generalization in medical image segmentation by leveraging semantic shape priors through a novel Mixture-of-Shape-Experts (MoSE) framework. MoSE treats each dictionary atom as a shape expert and uses a lightweight gating network to produce sparse, per-pixel combinations, forming a shape map that is then used as a prompt in the Segment Anything Model (SAM) to guide segmentation across unseen domains. The approach enables end-to-end training and scales the shape dictionary up to 1024 experts, achieving state-of-the-art generalization on multi-domain liver CT datasets, with notable improvements in Dice and Hausdorff distance. A bidirectional integration with SAM and the prompt-based usage of the shape map enhances SAM’s generalization, making MoSE a practical pathway for robust SDG in medical imaging and a foundation for future multi-class and 3D extensions.

Abstract

Single domain generalization (SDG) has recently attracted growing attention in medical image segmentation. One promising strategy for SDG is to leverage consistent semantic shape priors across different imaging protocols, scanner vendors, and clinical sites. However, existing dictionary learning methods that encode shape priors often suffer from limited representational power with a small set of offline computed shape elements, or overfitting when the dictionary size grows. Moreover, they are not readily compatible with large foundation models such as the Segment Anything Model (SAM). In this paper, we propose a novel Mixture-of-Shape-Experts (MoSE) framework that seamlessly integrates the idea of mixture-of-experts (MoE) training into dictionary learning to efficiently capture diverse and robust shape priors. Our method conceptualizes each dictionary atom as a shape expert, which specializes in encoding distinct semantic shape information. A gating network dynamically fuses these shape experts into a robust shape map, with sparse activation guided by SAM encoding to prevent overfitting. We further provide this shape map as a prompt to SAM, utilizing the powerful generalization capability of SAM through bidirectional integration. All modules, including the shape dictionary, are trained in an end-to-end manner. Extensive experiments on multiple public datasets demonstrate its effectiveness.

Figures (6)

• Figure 1: Overview of the proposed Mixture-of-Shape-Experts (MoSE) framework for Single Domain Generalization (SDG). The framework leverages a dictionary of shape experts to store diverse shape priors, which are dynamically combined into a shape map via sparse coding generated by SAM encoder and gating network. The shape map serves as a prompt integrated into the SAM pipeline.
• Figure 2: Visualization of the shape maps and segmentation results. The shape map is a heatmap of shape priors formed by combining shape experts and coefficients generated by the gating network. It is then used as a prompt in the subsequent SAM prompt encoder to assist in segmentation.
• Figure 3: Sensitivity analysis of hyperparameters. (a) With or without MoE for different dictionary sizes $n$, where the number of selected experts is set as $k=\frac{n}{2}$ for each case. (b) and (c) Dice and HD of different (b) $\beta$ and (c) $T_{warm-up}$. All results are averaged across all target domains.
• Figure 4: Visualization of segmentation results on unseen domains compared with other state-of-the-art SDG methods.
• Figure 5: Visualization of the learned shape experts.