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

Conference: ECCV 2026
arXiv: 2606.26938
Code: None
Area: Diffusion Models
Keywords: Mixture-of-Experts, Diffusion Models, Saliency-Aware Routing, Post-Training, Visual Generation

TL;DR

SharpMoE proposes a plug-and-play post-training framework that uses the clean latent variable \(x̂₀\) predicted from the previous denoising step instead of the noisy latent variable as the router input. This addresses the core issue in diffusion MoEs where "noisy routing" leads to a failure in identifying salient tokens. Augmented with a trajectory routing loss to globally align cumulative computation allocation, it significantly improves the generation quality of various converged pre-trained diffusion MoE models with only 100K post-training steps.

Background & Motivation

Diffusion Transformers (DiTs) transition diffusion models from U-Net to Transformer architectures, demonstrating excellent scalability. To overcome the efficiency bottleneck of dense parameter activation, Mixture-of-Experts (MoE) has been introduced into diffusion models, where a router dynamically allocates a sparse subset of experts to each token to support billion-parameter scales at controllable computational costs. Recent work (e.g., EC-DiT, DiffMoE) further recognizes that tokens from different image regions have varying demands for computing resources: salient tokens containing key structures and textures should receive more experts, whereas redundant background regions require only a few.

However, the authors identify a fundamental routing allocation issue in existing diffusion MoEs: the actual number of experts allocated by the router is almost unrelated to token saliency. By using a Laplacian operator to extract the texture information of each token as a saliency representation, and conducting statistical analysis on the number of experts allocated to tokens at different saliency levels, the results show that existing methods (e.g., DiffMoE) exhibit "saliency-insensitive" routing — differences in the number of experts allocated to high-saliency and low-saliency tokens are extremely small. The authors attribute the root cause to "noisy routing": routers in diffusion MoEs always take the currently noise-contaminated latent variable \(x_{t_k}\) as input. Residual noise during the denoising process — especially in early time steps with high noise — severely masks the semantic structural information of tokens, rendering the router unable to effectively distinguish salient regions from the background.

Core Idea: Use the clean latent variable \(x̂₀\) predicted from the previous denoising step as the input signal for the router, allowing the router to consistently make expert allocation decisions based on an estimate of "what the image should look like" rather than "how much noise remains." The clean latent variable naturally encodes the semantic skeleton of the image (the VAE latent space itself emphasizes structural saliency, and \(x̂₀\) is a projection onto the clean image manifold), stably capturing object contours and texture complexity even in the early high-noise stages.

Method

Overall Architecture

SharpMoE operates on converged pre-trained diffusion MoE models as a lightweight post-training enhancement framework. The inputs are a clean image \(x_0\) and Gaussian noise \(x_1\). Through a \(T\)-step recursive denoising rollout, the core operation of each step is a dual-router fusion: the current noisy latent variable \(x_{t_k}\) is fed into the preserved pre-trained router \(ℛ_{pre}\), and the clean latent variable \(x̂₀^{t_{k-1}}\) predicted from the previous step is fed into the newly added saliency router \(ℛ_{sal}\). Their outputs are summed to obtain the final gating score to drive sparse expert allocation. After the \(T\)-step rollout is complete, a trajectory routing loss is used to constrain the cumulative expert allocation across the trajectory to align with the Laplacian saliency distribution. The total loss is a weighted sum of the Flow Matching loss and the trajectory routing loss.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Input: x₀ clean image + x₁ pure noise"] --> B["Saliency-Aware Routing<br/>Dual-router fusion: ℛ_pre processes<br/>noisy latent x_t, ℛ_sal processes<br/>clean latent x̂₀, summed to get gating scores"]
    B --> C["Recursive Full-Trajectory Training<br/>T-step denoising recursive rollout<br/>Each step outputs x̂₀ passed to next step<br/>as input for ℛ_sal"]
    C -->|"x̂₀^{t_k} → next-step routing"| C
    C --> D["Trajectory Routing Loss<br/>Laplacian extracts saliency map<br/>KL divergence aligns full-trajectory<br/>cumulative expert allocation"]
    D --> E["Output: High-fidelity generated image"]

Key Designs

1. Saliency-Aware Routing: Replacing Noisy Latents with Clean Latents as Routing Input

Routers \(ℛ\) in existing diffusion MoEs directly take the current noisy latent variable \(x_{t_k}\) as input to calculate the token-expert affinity score \(S = ℛ(x)\). However, the residual noise in \(x_{t_k}\) masks the semantic saliency of the tokens, causing the routing outcome to be decoupled from token importance. The core insight of SharpMoE is that the clean latent variable \(x̂₀^{t_{k-1}}\) predicted during the denoising process (derived from the previous step via \(x̂₀^{t_{k-1}} = x_{t_{k-1}} - t_{k-1}·v_{k-1}\)) naturally contains the image's semantic structural information. The VAE latent space inherently encodes structural saliency, and \(x̂₀\) is a projection onto the clean image manifold, stably capturing object contours and texture complexity even in the early, highly noisy stages (as visually demonstrated in Fig. 1).

Specifically, SharpMoE introduces a dual-router mechanism at each MoE layer: keeping the original pre-trained router \(ℛ_{pre}\) (inputting \(x_{t_k}\) to capture the transient demands of the current denoising state) and adding a new saliency-aware router \(ℛ_{sal}\) (inputting \(x̂₀^{t_{k-1}}\) to provide clean saliency guidance). The final routing score is the sum of both: \(S = ℛ_{pre}(x_{t_k}) + ℛ_{sal}(x̂₀^{t_{k-1}})\). \(ℛ_{sal}\) is a two-layer MLP (with SiLU activation) whose weights are initialized to zero—this design ensures that the new router outputs no signal at the beginning of post-training, keeping the model's behavior identical to the original pre-trained model, and gradually incorporates saliency guidance during training to achieve a smooth transition. Gating matrix \(G\) is obtained after TopK selection, and the MoE output is the weighted sum of the selected experts.

2. Recursive Full-Trajectory Training: Moving from Single-Step Denoising to \(T\)-Step Recursive Rollout to Support \(x̂₀\) Dependency

Standard diffusion training uses a single-step paradigm — randomly sampling timestep \(t\), adding noise to obtain \(x_t\), and training the model to predict the velocity field \(v_θ(x_t, t)\). However, \(ℛ_{sal}\) in SharpMoE requires the previous step's \(x̂₀^{t_{k-1}}\) as input, which does not exist in single-step training. To address this recursive dependency, SharpMoE reformulates training into a \(T\)-step recursive full-trajectory rollout: each training iteration samples \(T\) continuous timesteps \(t_1 > t_2 > ... > t_T\), starting from \(t_1=0.999\) (approximate pure noise, as pure Gaussian noise at \(t=1\) would make the Flow Matching objective meaningless). The \(k\)-th step takes \(x_{t_k}\) and the stop-gradient-blocked \(x̂₀^{t_{k-1}}\) as inputs to predict velocity \(v_k\), then derives the next clean prediction via \(x̂₀^{t_k} = x_{t_k} - t_k·v_k\), and passes it to the \((k+1)\)-th step. At the first step \(t_1\) where there is no prior \(x̂₀\), the noisy latent variable \(x_{t_1}\) is directly used as a saliency approximation. The stop-gradient ensures that \(x̂₀\) only acts as a routing signal and is not involved in gradient computation for velocity prediction, avoiding training instability. \(T=10\) achieves optimal results in experiments and is highly robust to different choices of \(T\).

3. Trajectory Routing Loss: Globally Aligning Full-Trajectory Cumulative Expert Allocation with Saliency Distribution

Relying solely on \(ℛ_{sal}\) to provide clean routing signals is insufficient to optimize and guarantee that expert allocation is truly proportional to saliency. Full-trajectory training provides a global perspective not available in single-step approaches: it allows tracking the cumulative expert allocation for each token throughout the entire denoising process rather than relying on a potentially biased single-step snapshot. SharpMoE proposes a trajectory routing loss \(ℒ_{routing}\) to enforce this constraint: for a \(T\)-step rollout sequence, the trajectory-level cumulative allocation score for the \(i\)-th token is \(A_i = \sum_k \sum_l \sum_e ℐ(k,l,e,i) · S_{k,l}(e,i)\), which is the sum of the routing scores allocated to this token across all timesteps, all MoE layers, and all experts. Concurrently, a Laplacian operator \(\nabla^2\) is applied to the clean image \(X_0\) to extract second-order derivatives representing edge and texture features, which are then downsampled via AvgPool to token granularity to obtain the saliency level \(M_i\) for each token. Laplacian responses are naturally sensitive to high-frequency textures, sharp edges, and foreground boundaries, serving as an effective computational proxy for visual saliency. Finally, \(ℒ_{routing} = D_{KL}(softmax(A) \parallel softmax(M))\), which minimizes the KL divergence between the normalized cumulative allocation and the normalized saliency, forcing the model to allocate more experts to highly salient tokens across the full trajectory.

Loss & Training

The total training objective is \(ℒ_{total} = ℒ_{fm} + \lambda_{routing}·ℒ_{routing}\), where \(ℒ_{fm} = E[||(x₁-x₀) - v_θ(x_t, t)||²]\) is the standard Flow Matching MSE velocity field regression loss, \(ℒ_{routing}\) is the trajectory routing KL divergence loss, and \(\lambda_{routing} = 0.001\). The full-trajectory training step size is \(T=10\), and the first step is \(t_1=0.999\). The optimizer is AdamW, with a learning rate of \(1\times10^{-4}\), a batch size of 256, and an EMA decay rate of 0.9999 (all reported metrics are computed based on EMA weights). \(ℛ_{sal}\) is zero-initialized, and all other parameters are loaded from the pre-trained checkpoint. Post-training requires only 100K steps. Note that since TC-DiT employs uniform token computation allocation, \(ℒ_{routing}\) is meaningless for it; thus, SharpMoE on TC-DiT trains \(ℛ_{sal}\) under \(ℒ_{fm}\) alone.

Key Experimental Results

Main Results

Evaluated on ImageNet 256×256 class-conditional generation. All models are pre-trained for 500K steps, followed by 100K steps of post-training. 250-step Flow Matching Euler sampling is used to generate 50K images for calculating FID and IS metrics. The table below presents the core comparison on the L scale (conclusions on S and B scales are consistent, see Tab. 1 of the original paper for details).

Method	Active Params	Total Params	FID₅₀K (cfg=1.0)↓	IS (cfg=1.0)↑	FID₅₀K (cfg=1.5)↓	IS (cfg=1.5)↑
Dense-DiT-L	458M	458M	19.02	71.61	4.74	183.54
TC-DiT-L	458M	1.16B	20.53	69.01	5.07	174.98
TC-DiT-L + SharpMoE	484M	1.18B	16.63	81.29	3.72	206.93
EC-DiT-L	458M	1.16B	16.61	78.20	4.09	195.12
EC-DiT-L + SharpMoE	484M	1.18B	14.82	88.02	3.27	221.36
DiffMoE-L	458M	1.98B	15.13	83.62	3.86	203.00
DiffMoE-L + SharpMoE	470M	1.99B	13.93	93.66	3.10	228.88

SharpMoE achieves consistent and substantial performance gains across all baselines (TC-DiT, EC-DiT, DiffMoE) and all model scales (S, B, L, see original paper for details). On the strongest baseline, DiffMoE-L, SharpMoE reduces the FID from 3.86 to 3.10 (cfg=1.5) and improves the IS from 203.00 to 228.88. Even on TC-DiT, which deploys uniform token allocation, SharpMoE still significantly improves routing by introducing saliency awareness. The parameter increment is minimal (around 2%-3%), demonstrating the high efficiency of the post-training framework.

Ablation Study

Ablating each component of SharpMoE on DiffMoE-B (ImageNet 256×256).

Configuration	FID₅₀K (cfg=1.0)↓	IS (cfg=1.0)↑	FID₅₀K (cfg=1.5)↓	IS (cfg=1.5)↑
DiffMoE-B (Baseline)	27.50	54.45	8.03	138.49
+ Saliency-Aware Routing (\(ℛ_{sal}\) only)	25.93	58.33	6.95	152.65
+ Trajectory Routing Loss (Full SharpMoE)	25.71	59.60	6.66	155.64

Introducing only saliency-aware routing (\(ℛ_{sal}\) to guide the router with \(x̂₀\)) reduces the FID from 8.03 to 6.95, confirming that replacing noisy routing with clean routing signals is the core source of improvement. Further adding the trajectory routing loss drops the FID to 6.66, demonstrating that aligning global cumulative allocation yields additional gains.

Key Findings

• Clean routing is the primary driver: Simply replacing the routing input signal (without changing the loss function) yields a significant improvement, demonstrating that the primary bottleneck in existing diffusion MoEs is indeed the misalignment caused by noisy routing.
• The number of full-trajectory training steps \(T\) is highly robust: Varying \(T\) from 5 to 20 causes negligible fluctuations in FID, demonstrating that SharpMoE can be deployed without meticulous parameter tuning, making it highly practical.
• Insensitive to pre-training stages: Whether initiated from a 400K-step or a 700K-step pre-training checkpoint, SharpMoE consistently yields improvements within 100K post-training steps, showcasing its robustness as a general post-training enhancement.
• Visualization of expert allocation directly verifies the core hypothesis: The number of experts allocated by SharpMoE shows a statistically significant monotonic positive correlation with token saliency (with the most pronounced improvement witnessed in the highly noisy early steps), whereas the DiffMoE baseline displays near-zero correlation (Fig. 5), directly validating the effectiveness of clean routing.

Highlights & Insights

• The discovery and proof of the "noisy routing" issue are compelling: Rather than just pointing out that current methods have "inaccurate routing," the authors back this claim with quantitative evidence using Laplacian saliency analysis + expert allocation statistics (Fig. 1, Fig. 5) — demonstrating that the expert allocation curve across different token saliency levels is virtually flat. This visualization is highly convincing.
• Using \(x̂₀\) as the routing signal is a simple yet profound idea: Diffusion models naturally derive \(x̂₀\) from the velocity field anyway. SharpMoE simply repurposes this existing "byproduct" as the input to the router, entailing almost zero extra inference overhead (\(ℛ_{sal}\) is merely a lightweight two-layer MLP). This design philosophy of "leveraging existing but overlooked signals" is highly referenceable for other tasks.
• stop-gradient + zero initialization = true plug-and-play: These two seemingly minor implementation choices are crucial for SharpMoE to perform non-destructive updates on pre-trained models — the new router learns from scratch, emitting no signals at the start of training so that the model behaves exactly like the original pre-trained version, before smoothly transitioning to saliency-aware routing. This seamless integration strategy can be migrated to any scenario requiring the insertion of new components into pre-trained models.
• Trajectory-level rather than single-step global constraint: Expert allocation viewed from a single step might be biased (higher allocation at one specific step does not imply global rationale). Aligning after accumulating over the full trajectory is a much fairer constraint. This "local-to-global" perspective can be extended to other sequential generation tasks requiring temporal consistency constraints.

Limitations & Future Work

• Evaluation limited to ImageNet class-conditional generation: The paper does not evaluate on text-to-image generation (such as MoE variants of Stable Diffusion) or higher resolution scenarios, leaving the generalizability of saliency-aware routing in open-domain, multimodal conditional generation to be verified.
• Semantic limitations of the Laplacian saliency proxy: The Laplacian operator mainly captures low-level features like edges and textures, which may be insensitive to high-level semantic saliency (e.g., "human faces are inherently more salient than background"). Replacing the Laplacian with a stronger pre-trained saliency detector could yield further improvements.
• Insufficient discussion on the memory overhead of full-trajectory training: A recursive rollout of \(T=10\) steps implies maintaining multi-step computation graphs during training. Although \(x̂₀\) at each step disconnects its gradient chain via stop-gradient, the memory and computational overhead remain significantly higher than single-step training. The paper does not report concrete training costs.
• Dependence on the quality of \(x̂₀\) from the pre-trained model: If the base model's predictions of \(x̂₀\) in the early steps are of poor quality, the advantage of clean routing will diminish. The paper does not explore this lower bound, nor does it analyze the relationship between the quality of the base model and the gains of SharpMoE.
• Explorable avenues for improvement: \(ℛ_{pre}\) and \(ℛ_{sal}\) are currently combined via simple addition; more intimate interaction schemes (e.g., cross-attention) could be attempted. The trajectory routing loss could also be extended to video generation with longer temporal horizons to serve as a temporal consistency constraint, and one could investigate introducing clean routing directly during the pre-training stage rather than post-training.

Related Work & Insights

• vs DiffMoE / EC-DiT: These methods also strive for saliency-aware expert allocation but rely on noisy latent variables for routing decisions. The fundamental difference of SharpMoE is using \(x̂₀\) instead of \(x_t\) as the routing signal source, thereby bypassing noise interference at its root. Rather than replacing these methods, SharpMoE functions as a universal post-training enhancement layer applied on top of them, as demonstrated by the consistent improvements across three distinct architectures in the experiments.
• vs TC-DiT: TC-DiT employs fixed Top-K routing where each token is allocated the same number of experts, completely ignoring differences in saliency across tokens. SharpMoE's improvement over TC-DiT reveals that even the most basic uniform routing setup can benefit from saliency signals, hinting that "knowing which tokens are important" is more fundamental than "designing complex dynamic allocation strategies."
• vs MoEs in LLMs (e.g., DeepSeek-V3): Routing in LLM MoEs also faces challenges like load balancing, but the discreteness and high semantic density of textual tokens are starkly different from the spatial redundancy of visual tokens. The "clean routing" concept in SharpMoE has no direct counterpart in the text domain (as "denoising \(\rightarrow x̂₀\)" does not exist), but its meta-idea — "using auxiliary signals to replace primary signals for routing to bypass noise/bias in primary signals" — holds potential application scenarios in LLMs (e.g., leveraging retrieval-augmented signals to assist MoE routing).

Rating

• Novelty: ⭐⭐⭐⭐ Revealing the "noisy routing" issue and using \(x̂₀\) as the clean routing signal is simple and powerful. The plug-and-play post-training strategy with dual-routers + zero initialization is highly clever, though the core mechanism is not a paradigm-shifting breakthrough.
• Experimental Thoroughness: ⭐⭐⭐⭐ It covers 3 MoE architectures \(\times\) 3 model scales, with comprehensive ablation studies, robustness analyses, and expert allocation visualizations, though it lacks evaluations on text-to-image generation and broader datasets.
• Writing Quality: ⭐⭐⭐⭐ The problem is clearly formulated, with intuitive and convincing visualizations in Fig. 1 and Fig. 5. The motivational thread is complete, and technical details are thoroughly explained.
• Value: ⭐⭐⭐⭐ The plug-and-play post-training framework is highly practical, yielding improvements on matured models in just 100K steps. It holds direct reference value for industrial deployment of diffusion MoEs. The insights of clean routing are likely to influence the design of future diffusion MoE routers.

Related Papers

• [ICLR 2026] Routing Matters in MoE: Scaling Diffusion Transformers with Explicit Routing Guidance
• [ECCV 2026] Symbiotic-MoE: Unlocking the Synergy between Generation and Understanding
• [ICLR 2026] What Matters for Representation Alignment: Global Information or Spatial Structure?
• [CVPR 2026] Seeing What Matters: Visual Preference Policy Optimization for Visual Generation
• [ECCV 2026] Accelerated Likelihood Maximization for Diffusion-based Versatile Content Generation