Focusing on What Matters: Saliency-Harnessing Accurate Routing for Diffusion MoE

2026-06-25

AI Summary

This work addresses the challenge in existing diffusion-based Mixture-of-Experts (MoE) frameworks, where reliance on noisy latent features impedes accurate identification of salient visual tokens and efficient allocation of computational resources. To overcome this limitation, we propose SharpMoE—a post-training, plug-and-play routing optimization framework that leverages clean latent features obtained during the denoising process as noise-free guidance signals to precisely localize critical tokens. Furthermore, we introduce a trajectory routing loss to regulate the dynamic distribution of computational resources across multiple denoising steps. Without requiring any model retraining, SharpMoE achieves state-of-the-art performance across multiple visual generation benchmarks, effectively breaking through the bottleneck of inaccurate routing under noisy conditions.

Abstract

Mixture-of-Experts (MoE) architectures have emerged as a powerful paradigm for scaling diffusion models in visual generation. Recent advancements have focused on adaptively allocating computational resources across diverse tokens to improve efficiency and performance. However, we identify a routing assignment problem in existing diffusion MoE frameworks: the router fails to accurately allocate more computational resources to salient tokens. Our analysis attributes this failure to the router's reliance on noise-corrupted latent features throughout the denoising process. Such stochastic noise obscures the critical structural and textural information, thereby preventing the router from effectively distinguishing salient tokens. To address this, we propose SharpMoE, a post-training framework with a saliency-harnessing accurate routing mechanism, which utilizes clean latent features as a noise-free guidance signal for routing. By bypassing the noise-distorted inputs, SharpMoE provides the router with clear saliency guidance, enabling the identification of salient tokens even in high-noise stages. Furthermore, we introduce a trajectory routing loss to constrain the compute allocation throughout the multi-step denoising trajectory, ensuring precise resource allocation along the generation rollout. Extensive experiments demonstrate that SharpMoE serves as a versatile, plug-and-play solution that further enhances the pretrained, converged MoE models, achieving state-of-the-art performance in visual generation.

Problem

• Mixture-of-Experts
• diffusion models
• routing assignment
• saliency
• noise-corrupted features

Innovation

• saliency-aware routing
• clean latent features
• trajectory routing loss
• diffusion MoE
• noise-robust routing