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

arXiv: https://arxiv.org/abs/2504.09601 · HTML: https://mlsanity.com/paper/2504.09601/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Placed in Medical and biomedical MoE, one of 12 directions in Mixture-of-experts.
- Builds on 45 of its 64 references that are in the map.
- First posted on arXiv 13 Apr 2025.

## Placed in

- Mixture-of-experts / [Medical and biomedical MoE](https://mlsanity.com/moe/medical-and-biomedical-moe/)

## Builds on (most cited first, 30 of 45)

- [Segment Anything](https://mlsanity.com/paper/2304.02643/) (2023)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Visual Prompt Tuning](https://mlsanity.com/paper/2203.12119/) (2022)
- [AdaptFormer: Adapting Vision Transformers for Scalable Visual Recognition](https://mlsanity.com/paper/2205.13535/) (2022)
- [The Liver Tumor Segmentation Benchmark (LiTS)](https://mlsanity.com/paper/1901.04056/) (2019)
- [Generalizing to Unseen Domains via Adversarial Data Augmentation](https://mlsanity.com/paper/1805.12018/) (2018)
- [AMOS: A Large-Scale Abdominal Multi-Organ Benchmark for Versatile Medical Image Segmentation](https://mlsanity.com/paper/2206.08023/) (2022)
- [CHAOS Challenge -- Combined (CT-MR) Healthy Abdominal Organ Segmentation](https://mlsanity.com/paper/2001.06535/) (2020)
- [Learning to Learn Single Domain Generalization](https://mlsanity.com/paper/2003.13216/) (2020)
- [Customized Segment Anything Model for Medical Image Segmentation](https://mlsanity.com/paper/2304.13785/) (2023)
- [Learning to Diversify for Single Domain Generalization](https://mlsanity.com/paper/2108.11726/) (2021)
- [Conditional Computation in Neural Networks for faster models](https://mlsanity.com/paper/1511.06297/) (2015)
- [Semantic Segmentation with Generative Models: Semi-Supervised Learning and Strong Out-of-Domain Generalization](https://mlsanity.com/paper/2104.05833/) (2021)
- [Exact Feature Distribution Matching for Arbitrary Style Transfer and Domain Generalization](https://mlsanity.com/paper/2203.07740/) (2022)
- [Causality-inspired Single-source Domain Generalization for Medical Image Segmentation](https://mlsanity.com/paper/2111.12525/) (2021)
- [Deep Unsupervised Domain Adaptation: A Review of Recent Advances and Perspectives](https://mlsanity.com/paper/2208.07422/) (2022)
- [A Comprehensive Survey of Mixture-of-Experts: Algorithms, Theory, and Applications](https://mlsanity.com/paper/2503.07137/) (2025)
- [CLIP the Gap: A Single Domain Generalization Approach for Object Detection](https://mlsanity.com/paper/2301.05499/) (2023)
- [Adversarially Adaptive Normalization for Single Domain Generalization](https://mlsanity.com/paper/2106.01899/) (2021)
- [A Closer Look into Mixture-of-Experts in Large Language Models](https://mlsanity.com/paper/2406.18219/) (2024)
- [Generalizable Cross-modality Medical Image Segmentation via Style Augmentation and Dual Normalization](https://mlsanity.com/paper/2112.11177/) (2021)
- [Rethinking Data Augmentation for Single-source Domain Generalization in Medical Image Segmentation](https://mlsanity.com/paper/2211.14805/) (2022)
- [Model-Based Image Signal Processors via Learnable Dictionaries](https://mlsanity.com/paper/2201.03210/) (2022)
- [Devil is in Channels: Contrastive Single Domain Generalization for Medical Image Segmentation](https://mlsanity.com/paper/2306.05254/) (2023)
- [DeSAM: Decoupled Segment Anything Model for Generalizable Medical Image Segmentation](https://mlsanity.com/paper/2306.00499/) (2023)
- [Single-domain Generalization in Medical Image Segmentation via Test-time Adaptation from Shape Dictionary](https://mlsanity.com/paper/2206.14467/) (2022)
- [ACT: Semi-supervised Domain-adaptive Medical Image Segmentation with Asymmetric Co-training](https://mlsanity.com/paper/2206.02288/) (2022)
- [From Denoising Training to Test-Time Adaptation: Enhancing Domain Generalization for Medical Image Segmentation](https://mlsanity.com/paper/2310.20271/) (2023)
- [Memory Consistent Unsupervised Off-the-Shelf Model Adaptation for Source-Relaxed Medical Image Segmentation](https://mlsanity.com/paper/2209.07910/) (2022)
- [Text-Driven Tumor Synthesis](https://mlsanity.com/paper/2412.18589/) (2024)

## Built on by (most cited first, 4 of 4)

- [SegMoTE: Token-Level Mixture of Experts for Medical Image Segmentation](https://mlsanity.com/paper/2602.19213/) (2026)
- [SplatWeaver: Learning to Allocate Gaussian Primitives for Generalizable Novel View Synthesis](https://mlsanity.com/paper/2605.07287/) (2026)
- [GC-MoE: Genomics-Guided Cell-Type-Specific Mixture of Experts for Histology-Based Single-Cell Spatial Transcriptomics](https://mlsanity.com/paper/2606.02424/) (2026)
- [HadBalance: A Plug-and-Play Unified Global Geometric Prior Framework for Generalizable Biomedical Segmentation](https://mlsanity.com/paper/2606.15976/) (2026)

## Abstract (arXiv)

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.