AI Research Portal

SegMoTE is an efficient and adaptive framework for medical image segmentation that extends the Segment Anything Model (SAM) with a token-level Mixture of Experts (MoE) mechanism and Progressive Prompt Tokenization (PPT). The method addresses challenges in medical image segmentation such as modality- and anatomy-specific adaptation and the high cost of pixel-level annotations. SegMoTE preserves SAM's zero-shot generalization and efficient inference while introducing minimal learnable parameters. Trained on MedSeg-HQ, a curated dataset with 0.15M high-quality masks, SegMoTE achieves state-of-the-art performance across diverse imaging modalities and anatomical tasks, improving 1% to 6% over the second best approach.

The paper identifies two key bottlenecks in transferring general interactive segmentation models like SAM to medical imaging:

Lack of adaptive mechanisms: Existing methods lack adaptive mechanisms for modality- and anatomy-specific tasks, which limits generalization in out-of-distribution medical scenarios.

SegMoTE Framework: Proposed an efficient and adaptive framework that preserves SAM's strong zero-shot capability and flexibility while achieving modality-adaptive and precise medical image segmentation through dynamic selection and updating of expert tokens, using minimal additional parameters and training data.
MedSeg-HQ Dataset: Introduced MedSeg-HQ, a curated multimodal medical segmentation dataset of about 0.15M high-quality masks, which boosts model generalization and transferability with minimal supervision.
Progressive Prompt Tokenization (PPT): Proposed a mechanism that uses randomly sampled mask and text prompts to guide prompt tokens progressively toward foreground and background regions, enabling efficient, interaction-free few-class segmentation.
Empirical Superiority: Demonstrated through extensive experiments on both in-domain and out-of-domain datasets that SegMoTE improves performance by 1% to 6% over the second best approach, training only 17M parameters on a 0.15M dataset.

SegMoTE extends SAM with token-level expert routing for adaptive multimodal medical image segmentation.

Architecture The framework utilizes a frozen SAM encoder to extract modality-agnostic image embeddings. These embeddings are then processed by the Progressive Prompt Tokenization (PPT) mechanism, which converts latent features into semantically aligned feature tokens. These feature tokens are concatenated with original SAM output tokens and prompt tokens, then fed into the mask decoder. The mask decoder layers perform self-attention and token-to-image attention, passing results to the Mixture of Token Experts (MoTE) for dynamic expert selection and token updating. The updated expert tokens are used for final mask prediction.

Expert Token Selection A set of learnable expert tokens with a dimension of [N×256] is introduced, where N depends on the number of input modalities or task complexity. These tokens are concatenated with original SAM output tokens [4×256] and prompt tokens. The MoTE mechanism provides adaptive token selection. Given an input expert token x ∈ RB×T×D, logits for each token across all experts are computed using a router: L = XWg ∈ RB×T×E

During training, noisy top-k gating is applied: L = L+ (softplus(XWn) +ε) Z, Z ~ N(0, 1)

The top-k expert score for the t-th token in the b-th image is: Sb,t = (Lb,t,i1, Lb,t,i2, ..., Lb,t,ik) ∈ Rk, ij∈ Sb,t

The token confidence score and optimal expert index are determined: Cb,t = max sb,t[j], idxb,t = arg max sb,t[j] j=1,...,k j=1,...,k

The token weight is: G()b,t = softmax(cb,t), ab,t ∈ (0,1]

The updated token representation is: Zb,t = G()b,th(idxb,t)

Progressive Prompt Tokenization (PPT) PPT reduces reliance on manual annotations by treating mask and text prompts as foreground information. It randomly samples mask and text prompts to guide a learnable query Q via multi-head attention toward normalized image features. This process gradually learns to distinguish foreground from background, generating feature-conditioned prompt tokens for fully automatic segmentation.

Loss Function The total loss combines segmentation loss and load balancing loss: Ltotal = Lseg + Abalance. Lbalance

The segmentation loss is computed using Dice loss: Lseg(y,y) = 1 - 2 Σi ŷi yi / (Σi ŷi + Σi yi)

The load balancing loss is based on the coefficient of variation (CV²), encouraging balanced expert utilization: impe = ∑ Gb,t,e b,t loade = ∑ 1 (Gb,t,e > 0) b,t Lbalance = CV² ({impe}Ee=1) + CV² ({loade}Ee=1) CV² = std(x)² / (Σi=1 xi)²

Architecture: SAM encoder is frozen. MoTE and PPT are newly introduced and updated.
Key Equations:
- Router logits: L = XWg ∈ RB×T×E
- Noisy top-k gating: L = L+ (softplus(XWn) +ε) Z, Z ~ N(0, 1)
- Token confidence score: Cb,t = max sb,t[j], idxb,t = arg max sb,t[j]
- Token weight: G()b,t = softmax(cb,t), ab,t ∈ (0,1]
- Updated token representation: Zb,t = G()b,th(idxb,t)
- Expert importance: impe = ∑ Gb,t,e
- Expert load: loade = ∑ 1 (Gb,t,e > 0)
- Load balancing loss: Lbalance = CV² ({impe}Ee=1) + CV² ({loade}Ee=1)
- Coefficient of Variation: CV² = std(x)² / (Σi=1 xi)²
- Segmentation loss (Dice): Lseg(y,y) = 1 - 2 Σi ŷi yi / (Σi ŷi + Σi yi)
- Total loss: Ltotal = Lseg + Abalance. Lbalance
Loss/Objective: Dice loss for segmentation, load balancing loss (Lbalance) based on coefficient of variation.
Optimization: Adam optimizer with an initial learning rate of 1e-4. Learning rate is reduced by half after the 7th and 12th training epochs.
Augmentation: Not stated in the paper.
Data Handling: Trained on MedSeg-HQ, a curated multimodal dataset consisting of 154,569 high-quality annotations. Images are resampled to a resolution of 512×512. The dataset is split into training and testing sets with a 9:1 ratio, ensuring patient-level independence.
Implementation Details: Model training is conducted on 8 NVIDIA RTX 4090 GPUs with a total batch size of 10. The original SAM architecture remains frozen, and only MoTE and PPT are updated. The hyperparameter Abalance is set to 0.01. Experiments are conducted on the SAM-Base. SegMoTE contains 17M trainable parameters (MoTE: 10M, PPT: 7M).

Datasets:
- In-domain: MedSeg-HQ
- Out-of-domain: TotalSegmentator (MRI), SegThor, ISLES
- Ablation Study Datasets: ISIC2017, ISIC2018, SZ-CXR, ISLES
Metrics: Dice coefficient.
Baselines: SAM, SAM 2, MedSAM, Huang et al., SAM-Med2D, IMIS.
Numeric Results:
- Out-of-domain (Bounding-box interaction, Table 1):
  - ISLES (Ischemic Stroke Lesion): SegMoTE achieves 77.30 Dice, outperforming the second best (IMIS) at 71.24.
  - SegThor (Average): SegMoTE achieves 83.39 Dice, outperforming the second best (IMIS) at 80.52.
  - TotalSegmentator (MRI) (Average): SegMoTE achieves 71.48 Dice, outperforming the second best (IMIS) at 70.62.
- Mask decoder fine-tuning (Table 2):
  - AMOS (CT): SegMoTE achieves 85.16 Dice, outperforming IMIS at 83.56.
  - ISLES: SegMoTE achieves 78.05 Dice, outperforming IMIS at 71.20.
  - SegThor: SegMoTE achieves 83.91 Dice, outperforming IMIS at 80.96.
  - TotalSegmentator (MRI): SegMoTE achieves 73.67 Dice, outperforming IMIS at 71.01.
- Model Size Comparison (Table 3):
  - SegMoTE (Base): 17M learnable parameters.
  - MedSAM (Base): 93M learnable parameters.
  - IMIS (Base): 29M learnable parameters.
  - SAM (Large): 1191M learnable parameters.
- Ablation Study on Query Token Number (Table 4):
  - ISIC2017 (Q=2): 77.62 Dice.
  - ISIC2018 (Q=2): 87.68 Dice.
  - SZ-CXR (Q=2): 77.63 Dice.
  - ISLES (Q=2): 65.28 Dice.
  - w/o PPT (using points as prompt): ISIC2017 76.66, ISIC2018 84.87, SZ-CXR 75.20, ISLES 59.00.
Runtime/Training Cost: Not stated in the paper.

PPT for Multi-class Segmentation: The Progressive Prompt Tokenization (PPT) method performs well on binary classification datasets with clear foreground-background distinction, but in multi-class segmentation tasks (e.g., multi-organ segmentation), class interference complicates accurate prompt token mapping. Therefore, the paper focuses on applying PPT to binary classification tasks.

Method evidence: "To address these challenges, we propose SegMOTE (Segmentation with Mixture of Token Experts), an interactive medical image segmentation framework built upon the Mixture of Experts (MoE) paradigm [12, 22, 36]. As illustrated in Fig.1 (b), SegMoTE introduces token-level expert selection, dynamically activating the most suitable expert tokens for each imaging modality to achieve modality specific representation and adaptive processing." (section: Method)
Result evidence: "Trained on MedSeg-HQ, a curated dataset less than 1% of existing large-scale datasets, SegMoTE achieves SOTA performance across diverse imaging modalities and anatomical tasks. It represents the first efficient, robust, and scalable adaptation of general segmentation models to the medical domain under extremely low annotation cost, advancing the practical deployment of foundation vision models in clinical applications." (section: Abstract)
Limitation evidence: "This method performs well on binary classification datasets with clear foreground-background distinction, such as ISIC and SZ-CXR, reducing reliance on explicit visual prompts. However, in multi-class segmentation tasks (e.g., multi-organ segmentation), class interference complicates accurate prompt token mapping. Therefore, this paper focuses on applying the PPT method to binary classification tasks to maximize its simplicity and inference efficiency." (section: 3.4. Progressive Prompt Tokenization)

SegMoTE: Token-Level Mixture of Experts for Medical Image Segmentation

Summary

Problem

Key Contributions

Method

Technical Details

Results

Weaknesses / Limitations

Evidence

Additional Figures