Functional routing improves single-step image diffusion generation
Functional Cross-layer Experts Routing for One-step Diffusion Distillation
Computer Vision and Pattern Recognition
Summary
Diffusion models generate images by gradually refining them through many steps, but some approaches try to do this in just one step to save time. The authors explore how to organize the model’s internal parts to handle both general image features and specific prompt details at once. They introduce DeMoE, which keeps a shared backbone always active while activating only parts of additional 'expert' modules depending on the input. This method improves image quality and diversity in one-step generation, matching or surpassing longer multi-step processes on some tasks.
What this means in practice
- •For machine learning engineers: Create faster image generation models that efficiently balance general features and prompt-specific details in one evaluation pass.
- •For graphics software developers: Build real-time image synthesis tools that generate high-quality images with fewer computational steps using structured expert routing.
Authors
Zile Huang, Ser-nam Lim
Abstract
One-step distillation compresses diffusion trajectories into a single forward evaluation, yet how internal capacity should be structured within that remaining pass remains largely unexplored. In this setting, identical parameters must simultaneously preserve general generative priors and resolve prompt-specific details without iterative refinement. We formalize this challenge as functional allocation. Our analysis reveals that decoupling representations is beneficial only when expert interference outweighs shared transfer, favoring an always-active shared backbone augmented with conditionally routed residuals over standard isolated MoE architectures. We propose DeMoE, which pairs a persistent shared expert with conditionally routed expert slices across image FFNs, trained under a one-step distribution-matching objective with executed-route supervision. On one-step FLUX.2-4B generation, DeMoE achieves a GenEval score of 0.8220 while activating only 80.5% of the image-FFN width, substantially outperforming dense one-step baselines, matching the 50-step teacher overall, and surpassing it on complex composition and counting. Across COCO-10K, DeMoE establishes state-of-the-art fidelity and diversity among single-step models, demonstrating that structured internal computation reallocation provides a superior compression paradigm for diffusion models. Code is available at https://github.com/toffeecat/demoe-ttt-router.