Router learning improves mixed medical image generation across domains
Regime-Aware Router Learning for Multi-Domain Medical Image Generation with Diffusion Experts
Computer Vision and Pattern RecognitionMachine Learning
Summary
Medical images come from many different sources, and special AI models called diffusion experts focus on each type. But it's tricky to combine their predictions well, especially when the main system only has a small or unbalanced sample of images. The authors tested two ways to teach a global router to pick which expert to trust when generating images. One method uses expert performance data, and another mixes this with extra clues from the image generation process itself. They found that combining these clues helps the system create better medical images from diverse sources.
What this means in practice
- •For medical image software developers: Improve AI systems that generate synthetic medical images by better combining specialized models from different imaging sources.
- •For healthcare data engineers: Enhance server-side management of heterogeneous medical datasets by learning optimal routing for expert predictions during image synthesis.
Authors
Yaoshuai Ma, Wei Yao, Yongxiang Liu, Feifan Wang, Da Chang, Pengxiang Xu, Shixun Zhang, Yonghong Tian
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.