Regime-Aware Router Learning for Multi-Domain Medical Image Generation with Diffusion Experts

Yaoshuai Ma; Wei Yao; Yongxiang Liu; Feifan Wang; Da Chang; Pengxiang Xu; Shixun Zhang; Yonghong Tian

cs.CVOct 12, 2026

Regime-Aware Router Learning for Multi-Domain Medical Image Generation with Diffusion Experts

Authors: Yaoshuai Ma, Wei Yao, Yongxiang Liu, Feifan Wang, Da Chang, Pengxiang Xu, Shixun Zhang, Yonghong Tian

Organizations: Peng Cheng Laboratory, Shenzhen, China · Southern University of Science and Technology, Shenzhen, China

Abstract

Heterogeneous medical image data motivate domain-specialized diffusion experts, but expert specialization alone does not determine how their predictions should be combined during sampling. Learning a global router is particularly challenging when the server has limited, potentially unrepresentative images. We study two router-learning strategies for a fixed expert pool. Server-Only learns input- and timestep-dependent expert assignments from server images with component labels. Anchor-Delta (AD) uses the same images and frozen experts to obtain additional evidence from denoising losses. It distills targets that combine this evidence with the current router's assignments, while a proxy gate controls a subsequent refinement step. Component labels describe data origin; expert losses measure empirical prediction performance. We evaluate these distinct signals through routing analyses and paired medical image generation. An AD-trained router improves generation metrics over Server-Only in a heterogeneous four-domain setting, and a separate server-data study shows gains in held-out routing over a corrected-routing comparator. Additional paired evaluations show that the effect of inference-time evidence mixing varies across benchmarks. These findings establish server-data availability and expert performance as distinct considerations in learning routers for multi-domain diffusion models.

Figures & tables

Method	FID $\downarrow$	KID $\times 10^{3}\downarrow$
Server-Only	$35.54\pm 0.53$	$1.491\pm 0.186$
VQ-FedDiff, mixed-client	$38.51\pm 0.34$	$3.416\pm 0.188$
FedPer	$38.99\pm 0.33$	$3.206\pm 0.025$
FedBN	$40.27\pm 0.39$	$3.276\pm 0.239$
FedDDPM-style	$45.61\pm 1.60$	$4.514\pm 0.540$
Fixed- $\alpha$ AD ( $\alpha=0.05$ )	$\mathbf{34.19\pm 0.92}$	$\mathbf{1.349\pm 0.140}$

TABLE I: MedMNIST2D paired generation (

t=180

). Mean

\pm

sample standard deviation over three training seeds; 1000 test targets and outputs per seed.

Method	FID $\downarrow$	KID $\times 10^{3}\downarrow$
Server-Only	$56.99\pm 0.37$	$9.48\pm 0.18$
VQ-FedDiff, inferred-client	$108.82\pm 1.93$	$27.83\pm 0.79$
FedPer	$104.71\pm 0.59$	$26.05\pm 0.54$
FedBN	$95.51\pm 0.82$	$21.57\pm 0.17$
FedDDPM-style	$108.21\pm 1.92$	$27.20\pm 0.83$
Fixed- $\alpha$ AD ( $\alpha=0.25$ )	$\mathbf{53.69\pm 0.21}$	$\mathbf{7.01\pm 0.28}$

TABLE II: Medical-4 paired generation (

t=850

). Mean

\pm

sample standard deviation over three seeds, with 822 real references and 1000 outputs per seed. For AD,

\alpha

is the router-training correction coefficient; sampling uses the trained router directly.

Panel of Figure 3 — Fig. 1: Paired reference-assisted generation across Medical-4. Each method panel contains a $2\times 2$ grid in the same order: ISIC skin (upper left), BreakHis histopathology (upper right), Messidor fundus (lower left), and MURA radiograph (lower right). Each domain uses the same selected test case across methods. AD uses its trained router directly during sampling, with $\alpha=0.25$ specifying the training correction coefficient. The examples illustrate differences in tissue appearance, vascular structure, and bone contours rather than diagnostic validity. Aggregate generation quality is assessed in Table II .

Method	3D SSIM $\uparrow$	MS-SSIM3D $\uparrow$	PSNR (dB) $\uparrow$	MAE $\downarrow$
Server-Only	$\mathbf{0.9231\pm 0.0035}$	$\mathbf{0.9588\pm 0.0031}$	$28.079\pm 0.210$	$\mathbf{0.03352\pm 0.00038}$
VQ-FedDiff, inferred-client	$0.8980\pm 0.0065$	$0.9256\pm 0.0064$	$28.127\pm 0.122$	$0.03605\pm 0.00049$
FedPer	$0.7210\pm 0.0022$	$0.7917\pm 0.0028$	$23.302\pm 0.077$	$0.08550\pm 0.00018$
FedBN	$0.7198\pm 0.0059$	$0.7903\pm 0.0060$	$23.311\pm 0.105$	$0.08544\pm 0.00061$
FedDDPM-style	$0.8642\pm 0.0337$	$0.9020\pm 0.0215$	$26.401\pm 1.572$	$0.04435\pm 0.00859$
Fixed- $\alpha$ AD ( $\alpha=0.25$ )	$0.9226\pm 0.0038$	$0.9583\pm 0.0035$	$28.108\pm 0.252$	$0.03386\pm 0.00045$

TABLE III: MedMNIST3D paired volumetric generation (

t=100

, four variants per source). Mean

\pm

sample standard deviation over three seeds; 384 source volumes per seed.

Method	Slice FID $\downarrow$	SSIM $\uparrow$ [95% CI]
Server-Only	$6.716$\pm$0.124$	$0.5945$ $[$ 0.5906 $,$ 0.5984 $]$
VQ-FedDiff inferred-client	$292.299$\pm$14.306$	$0.0893$ $[$ 0.0861 $,$ 0.0932 $]$
FedPer	$11.327$\pm$1.219$	$0.5737$ $[$ 0.5650 $,$ 0.5819 $]$
FedBN	$11.412$\pm$2.055$	$0.5759$ $[$ 0.5669 $,$ 0.5853 $]$
FedDDPM-style	$74.200$\pm$25.838$	$0.3535$ $[$ 0.2180 $,$ 0.5067 $]$
Fixed- $\alpha$ AD ( $\alpha=0.05$ )	$6.388$\pm$0.140$	$0.6132$ $[$ 0.5960 $,$ 0.6238 $]$

TABLE IV: TotalSeg-CT4 paired generation (

t=160

). FID: mean

\pm

sample standard deviation across three seeds. SSIM: volume-level domain-macro mean [hierarchical 95% CI], derived from slice metrics. Values are displayed to two decimal places.

Fig. 2: Server-supervision ablation on Medical-4: corrected implicit routing (CI) versus proxy-gated AD. Higher global and proxy router scores are better. The dotted line indicates the three-of-four strong-domain criterion. Markers represent the seven evaluated ratios; connecting lines are visual guides, not additional observations. No uncertainty bands are inferred from these aggregate scores.

Fig. 3: MedMNIST2D training dynamics from complete checkpoint histories. Lines show three-seed means and shading shows one sample standard deviation after 1000-update window averaging within each seed. All vertical axes are logarithmic. Panels use different objectives: true-noise prediction, component classification, and teacher-prediction matching. Loss magnitudes are not comparable across objectives.

Fig. 4: Measured MedMNIST2D deployment costs. All models use the same NPU and sampling protocol. Bars show sample standard deviations across three seeds; latency uses the per-seed median of three timed batches. Both teacher variants retain all four experts in memory. These measurements quantify deployment cost, not equivalence in generation quality.

Expert Routing

Expert routing in mixture-of-experts models selects which subnetworks process each input, balancing accuracy, computation, and expert workload.

See all 119 papers

Expert-Choice Routing Enables Adaptive Computation in Diffusion Language Models
Apr 2, 2026Shuibai Zhang, Caspian Zhuang, Chihan Cui +8Expert Routing Mixture-of-Experts Language Models
University of Wisconsin-Madison · Scitix · Cornell University +4
Functional Cross-layer Experts Routing for One-step Diffusion Distillation
Oct 19, 2026Zile Huang, Ser-nam LimDiffusion Model Distillation Diffusion Distillation
University of Central Florida
Beyond Routing: Decoupling Expert Dispatch and Aggregation in Sparse Mixture-of-Experts
Aug 9, 2026Zongfei LiMixture-of-Experts Language Models Sparse Mixture-of-Experts

Apr 2, 2026J/K move · Enter open · S save

Expert-Choice Routing Enables Adaptive Computation in Diffusion Language Models

Shuibai Zhang, Caspian Zhuang, Chihan Cui, Zhihan Yang+7

University of Wisconsin-Madison · Scitix · Cornell University+4

Expert Routing Mixture-of-Experts Language Models Diffusion Language Models Adaptive Computation

Open paper AskPDF arXiv

Diffusion language models (DLMs) enable parallel, non-autoregressive text generation, yet existing DLM mixture-of-experts (MoE) models inherit token-choice (TC) routing from autoregressive systems, leading to load imbalance and rigid computation allocation. We show that expert-choice (EC) routing is a better fit for DLMs: it provides deterministic load balancing by design, yielding higher throughput and faster convergence than TC. Building on the property that EC capacity is externally controllable, we introduce timestep-dependent expert capacity, which varies expert allocation according to the denoising step. We find that allocating more capacity to low-mask-ratio steps consistently achieves the best performance under matched FLOPs, and provide a mechanistic explanation: tokens in low-mask-ratio contexts exhibit an order-of-magnitude higher learning efficiency, so concentrating compute on these steps yields the largest marginal return. Finally, we show that existing pretrained TC DLMs can be retrofitted to EC by replacing only the router, achieving faster convergence and improved accuracy across diverse downstream tasks. Together, these results establish EC routing as a superior paradigm for DLM MoE models and demonstrate that computation in DLMs can be treated as an adaptive policy rather than a fixed architectural constant. Code is available at https://github.com/zhangshuibai/EC-DLM.

Abstract

Expert Routing

Explore similar work

Expert-Choice Routing Enables Adaptive Computation in Diffusion Language Models

Functional Cross-layer Experts Routing for One-step Diffusion Distillation

Beyond Routing: Decoupling Expert Dispatch and Aggregation in Sparse Mixture-of-Experts