Source-linked AI summary
Focusing on What Matters: Saliency-Harnessing Accurate Routing for Diffusion MoE
Haoyou Deng, Keyu Yan, Chaojie Mao, Xiang Wang, Yu Liu, Changxin Gao, Nong Sang
TL;DR
Existing diffusion MoE routers allocate compute insensitively to token saliency because noisy latents obscure salient features. SharpMoE uses clean latent guidance and trajectory-aware routing, outperforming competing methods across evaluated settings and achieving FID 3.10 and IS 228.88 on DiffMoE-L at cfg = 1.5.
Problem
Existing diffusion MoE routers struggle to distinguish salient tokens because noise-corrupted latents produce saliency-insensitive expert allocation.
Method
SharpMoE uses preceding-step clean latent predictions for saliency-guided routing and a trajectory routing loss to regulate allocation across denoising.
Results
FID 3.10 and IS 228.88 are achieved on DiffMoE-L at cfg = 1.5, while SharpMoE outperforms competing diffusion MoE methods across reported settings.
Takeaways & Limitations
SharpMoE is a plug-and-play post-training enhancement that improves pretrained diffusion MoE models across architectures, scales, metrics, and guidance settings.
Takeaways & Limitations
At the initial inference timestep, SharpMoE lacks a prior clean prediction and therefore uses the noisy latent as a saliency-guidance proxy.
Abstract
from arXiv · showhide
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.
1 Introduction
Diffusion MoE models scale capacity efficiently but struggle to allocate computation according to token saliency because noisy routing inputs obscure structural and textural cues. SharpMoE addresses this with clean-latent routing and trajectory-level supervision for accurate, stable resource allocation.
- Motivation: Dense parameter activation makes scaling diffusion transformers to larger models increasingly inefficient, motivating sparse Mixture-of-Experts architectures.MoE dynamically routes each token to a sparse subset of experts, expanding model capacity while maintaining computational efficiency.
- Routing problem: Existing diffusion MoE routers fail to allocate more computational resources to salient tokens, despite differing computational needs across image regions.Salient regions contain critical details and require greater computation than less significant regions.
- SharpMoE: SharpMoE uses the preceding timestep’s predicted clean latent, x̂_0, as the current router input to provide noise-free saliency guidance.Clean features preserve structural and textural information and can identify salient regions even during heavily noisy early timesteps.
- Trajectory Routing Loss: The Trajectory Routing Loss aggregates each token’s expert assignments across the denoising trajectory and aligns cumulative computational effort with saliency.The loss regulates the computational budget so capacity concentrates on regions with high structural and textural complexity.
- Results: Experiments across multiple pretrained architectures show that SharpMoE is an effective plug-and-play post-training enhancement for diffusion MoE generation.The results emphasize clean routing as important for improving generative fidelity.
2 Related Work
Related work traces diffusion models’ progression from U-Net/DDPM systems to scalable DiT models trained with Rectified Flow, alongside MoE methods for efficient sparse activation. In diffusion MoE, noisy latents can obscure saliency and hinder precise expert assignment, motivating SharpMoE’s use of predicted clean latents for routing.
- Diffusion Models: Diffusion models evolved from U-Net backbones optimized with DDPM objectives toward scalable Diffusion Transformer architectures combined with Rectified Flow training.The cited DiT-based models include recent diffusion systems built for model scaling.
- Mixture of Experts: MoE architectures efficiently expand model capacity through sparse activation, activating only a subset of parameters for each token.MoE has achieved considerable success in large language models, including DeepSeek-V3 and MiniMax-01.
- Mixture of Experts: Noisy latents during denoising obscure saliency representations, making precise expert assignment difficult in diffusion MoE.SharpMoE addresses this limitation by routing with predicted clean latents to provide robust saliency representations for dynamic expert assignment.
3 Preliminary
The section introduces Rectified Flow diffusion, Flow Matching denoising, and sparse Mixture-of-Experts routing. It identifies noise-corrupted inputs as causing inaccurate salient-token assignments and motivates clean-latent guidance with trajectory-level routing alignment.
- Diffusion Models: Rectified Flow defines the noised data as x_t = t x_1 + (1-t) x_0, interpolating between a data sample and a noise sample.Here, x_0 is sampled from the data distribution and x_1 from a noise distribution.
- Diffusion Models: Flow Matching trains a denoising model to regress the velocity field v_θ(x_t, t) toward the target v = x_1 − x_0.The objective minimizes the expected squared error between the target and predicted velocity.
- Diffusion Models: During generation, the state updates as x_{t+dt} = x_t + dt · v_θ(x_t, t), while the predicted clean latent is ˆx_0 = x_t − t v_θ(x_t, t).dt denotes the timestep gap.
- Mixture of Experts: A standard MoE router scores token–expert affinities, selects the top-k experts, and combines their outputs through a gating tensor.The layer contains N_E experts, each implemented as a Feed-Forward Network.
- Routing Challenge and SharpMoE: Because diffusion inputs retain noise, the router cannot reliably distinguish salient tokens, producing incorrect routing assignments.SharpMoE addresses this with predicted clean latents for saliency guidance and trajectory routing loss for cumulative compute alignment across denoising steps.
4 SharpMoE
SharpMoE replaces noisy routing with saliency-guided routing using predicted clean latents, enabling computation to focus on salient regions during denoising. Recursive full-trajectory training and a trajectory routing loss support clean guidance and regulate cumulative expert allocation across generation.
- Framework overview: SharpMoE replaces standard FFNs in DiT with SharpMoE blocks that dynamically allocate computation toward salient regions during denoising.The framework is designed for scalable, efficient modeling.
- Saliency-Harnessing Router: The saliency-harnessing router uses the predicted clean latent from the preceding timestep instead of the current noisy latent for routing guidance.This clean latent provides stable, noise-free information that captures structurally significant regions and helps prioritize salient tokens despite residual noise.
- Saliency-Harnessing Router: SharpMoE combines a pretrained router with a saliency-harnessing router in a plug-and-play post-training enhancement for established diffusion MoE models.The saliency router is initialized with zero weights so saliency guidance is incorporated progressively without disrupting pretrained denoising capabilities.
- Recursive Full-Trajectory Training: Recursive full-trajectory training simulates consecutive denoising timesteps because the saliency router requires the preceding step’s predicted clean latent, which single-step training does not compute.At the initial timestep, the noise latent is used as a saliency proxy because structure and saliency have not yet been determined.
- Trajectory Routing Loss: Trajectory Routing Loss aligns cumulative expert allocation with saliency across the entire generation process, reducing background redundancy while enhancing crucial foreground detail fidelity.The overall objective combines Flow Matching loss with the trajectory routing loss.
5 Experiment
Experiments show that SharpMoE consistently improves pretrained diffusion MoE models across scales, guidance settings, and evaluation metrics. Ablations and visualization further attribute these gains to clean-latent saliency guidance and robust trajectory-aware routing, especially during high-noise stages.
- Main Results: SharpMoE consistently outperforms TC-DiT, EC-DiT, and DiffMoE across S, B, and L scales, all reported metrics, and CFG scales after 100K post-training steps.Models are initialized from 500K-step pretrained checkpoints.
- Main Results: 3.10 FID and 228.88 IS are achieved by SharpMoE on DiffMoE-L with cfg = 1.5, its strongest reported configuration.The result is obtained after 100K post-training steps from a 500K-step pretrained model.
- Ablation Studies: Reducing FID from 8.03 to 6.95 at cfg = 1.5 demonstrates the saliency-harnessing routing mechanism’s contribution over conventional noisy routing.The ablation identifies clean latent guidance as critical for improving saliency awareness.
- Ablation Studies: SharpMoE consistently improves DiffMoE-B checkpoints initialized at both 400K and 700K training steps, demonstrating plug-and-play adaptability.The improvement occurs within 100K post-training steps.
- Routing Analysis: SharpMoE remains robust across rollout counts T = 5 to 20 and establishes a monotonic saliency-to-expert relationship, with the largest routing gains during high-noise stages.Higher-saliency tokens receive more experts through noise-free latent guidance, prioritizing structurally or texturally rich regions.
6 Conclusion
SharpMoE is a post-training framework for diffusion MoE that addresses saliency-insensitive compute allocation by routing with clean predictions instead of noisy latents. It introduces saliency-harnessing accurate routing and builds upon full-trajectory training.
- Routing problem: SharpMoE addresses diffusion MoE routing failures in which noisy-latent-conditioned routers struggle to recognize salient tokens.This causes saliency-insensitive compute allocation.
- Accurate routing: Its saliency-harnessing accurate routing mechanism uses clean predictions as saliency representations for noise-free routing.The clean prediction provides guidance without the noise affecting latent features.
- Training framework: SharpMoE is presented as a post-training framework that builds upon a full-trajectory training scheme.