Focusing on What Matters: Saliency-Harnessing Accurate Routing for Diffusion MoE
SharpMoE leverages predicted clean latent features for saliency-aware routing, boosting diffusion MoE image generation by ~20% in FID.
Key Findings
Methodology
This work analyzes the failure of existing diffusion MoE routers, which rely on noisy latent features, leading to poor saliency discrimination. The authors propose using predicted clean latent features (x̂0) as guidance signals for routing, combined with a recursive full-trajectory training scheme and a trajectory routing loss. The approach involves estimating saliency via the Laplacian operator, and enforcing resource allocation alignment across multiple denoising steps, ensuring saliency-aware compute distribution. The core algorithm integrates a dual-router system—pretrained and saliency-harnessing—fusing their outputs for expert assignment, and optimizes the entire process through a novel loss function that considers the cumulative expert scores over the full generation trajectory.
Key Results
- On ImageNet 256×256, SharpMoE improves FID scores by approximately 20% across multiple architectures (e.g., TC-DiT, EC-DiT, DiffMoE). It achieves state-of-the-art visual quality, with expert allocations highly correlated with saliency maps, leading to finer details and structural accuracy. The method demonstrates robustness in high-noise stages, maintaining salient object recognition even during early denoising steps. Quantitative improvements include reduced FID from 31.87 to 16.63 on TC-DiT-B, and similar gains across other models.
- The introduction of the trajectory routing loss ensures resource distribution aligns with saliency throughout the entire denoising process, resulting in more coherent and detailed images. Ablation studies confirm that using predicted clean latents and the full-trajectory loss significantly outperform baseline noisy-routing methods, with improvements exceeding 15% in FID and notable gains in IS scores.
- The approach is computationally efficient, requiring only limited post-training steps to enhance already converged models. It offers a plug-and-play solution that can be integrated into existing diffusion MoE architectures without retraining from scratch, making it highly practical for industry deployment.
Significance
This research addresses a fundamental bottleneck in scaling diffusion models: ineffective resource allocation due to noisy saliency signals. By introducing a noise-free guidance mechanism, it elevates the fidelity and efficiency of visual generation, bridging the gap between large-scale capacity and precise, saliency-aware resource management. The method’s plug-and-play nature facilitates rapid adoption, promising significant impact in high-fidelity image synthesis, virtual content creation, and beyond. It paves the way for more intelligent, resource-efficient generative systems that can focus computational power where it matters most, leading to richer, more detailed outputs with less waste.
Technical Contribution
The paper introduces a novel saliency-guided routing framework that leverages predicted clean latent features, combined with a recursive full-trajectory training scheme and a trajectory routing loss. This setup ensures the expert assignment process is robust against noise and aligns with the true importance of image regions. The dual-router architecture, integrating a pretrained router with a saliency-harnessing router, offers a flexible, plug-and-play enhancement for existing models. The theoretical guarantee of resource allocation matching saliency distribution across multiple steps is a key advancement, enabling more precise and efficient diffusion-based image synthesis.
Novelty
This is the first work to incorporate predicted clean latents as guidance for dynamic expert routing in diffusion MoE, effectively mitigating noise interference. The recursive full-trajectory training and the trajectory routing loss together form a new paradigm for holistic, saliency-aware resource management across multiple denoising steps, setting this apart from prior noisy or single-step routing strategies.
Limitations
- The method's effectiveness depends on the accuracy of the predicted clean latent features; errors here can reduce saliency discrimination. In scenarios with extreme noise or highly complex scenes, the predictions may be less reliable, affecting resource allocation.
- While the post-training scheme is efficient, applying it to very large models or high-resolution images still incurs significant computational overhead. Further optimization is needed for real-time applications.
- The approach assumes that saliency correlates with high-frequency components estimated via Laplacian, which may not hold in all cases, especially for abstract or textureless images.
Future Work
Future directions include developing more robust latent prediction methods, extending the framework to video and 3D scene generation, and exploring adaptive hyperparameter tuning for different data domains. Additionally, integrating multimodal saliency cues could further enhance resource allocation precision.
AI Executive Summary
Diffusion models have revolutionized visual generation, yet scaling them efficiently remains a challenge due to ineffective resource allocation driven by noisy saliency signals. Traditional routing mechanisms rely on corrupted latent features during the denoising process, which obscures critical structural and textural information, leading to suboptimal expert assignment. This results in generated images lacking detail fidelity, especially in salient regions. To address this, the authors propose SharpMoE, a novel framework that leverages predicted clean latent features as guidance signals for dynamic routing.
The core innovation involves integrating a saliency-harnessing router that uses these noise-free features, combined with a recursive full-trajectory training scheme and a trajectory routing loss. This setup ensures that resource allocation aligns with the true importance of image regions across all denoising steps. The approach employs a dual-router architecture—retaining the pretrained router and adding a saliency-guided router—fusing their outputs for expert assignment. The trajectory routing loss further enforces the alignment of cumulative expert scores with saliency maps estimated via the Laplacian operator.
Extensive experiments on ImageNet 256×256 demonstrate that SharpMoE consistently improves FID scores by around 20%, achieving state-of-the-art results across multiple architectures. The method effectively enhances the model’s ability to identify and prioritize salient regions, leading to images with richer details and better structural coherence. Its plug-and-play nature allows seamless integration into existing pretrained models, requiring only limited post-training steps. While promising, the approach still relies on accurate latent predictions, and future work will focus on robustness in diverse scenarios and extending to video and 3D generation. Overall, SharpMoE marks a significant step toward more intelligent, resource-efficient diffusion-based image synthesis, with broad implications for both academia and industry.
Deep Dive
Limitations & Outlook
What gaps remain?
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.