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.
Figures & tables
Figure 1: Alternative ways to compress generative computation to one neural function evaluation. Few-step DMD retains an explicit sequence of state-dependent generator evaluations, whereas the remaining paradigms target one-step inference but introduce different forms of training structure. One-step DMD removes the explicit sampling trajectory and matches the target distribution directly. Progressive Distillation compresses multiple teacher transitions into one student transition; Consistency Distillation transfers cross-time endpoint consistency; Shortcut Models construct long transitions by composing shorter stop-gradient shortcuts; and MeanFlow directly parameterizes the average velocity over a finite interval. Thus, these methods differ primarily in how the one-step mapping is supervised , even when their deployed inference graph contains only one network evaluation. Red denotes the trainable one-step mapping, blue denotes frozen/EMA/stop-gradient computation, gray denotes teacher or data quantities, and dashed arrows denote losses or target construction.
Figure 2: DeMoE block and depth-recurrent routing. (a) Standard Image DiT block. (b) DeMoE preserves the original modulation, attention, and residual structure, but partitions the dense image FFN into an always-active frozen shared slice S and seven routed slices R1:7 , of which four are selected per token. Routing uses the pre-FFN hidden state and text condition, with the optional recurrent readout from (c). (c) The history variant maintains sample-private fast weights across depth: an aligned inter-layer residual provides keys and values for the inner update, while the current hidden state provides the query used to read the routing state.
Inference budget
50 NFE
4 NFE
1 NFE
Dense Teacher
Dense DMD2
Dense Cold start
Dense DMD2
Dense Progressive
DeMoE 3E1A1S
DeMoE 7E3A1S
DeMoE 7E4A1S
GenEval
Overall ↑
0.8245
0.8035
0.6888
0.7865
0.7771
0.7415
0.7812
0.8220
Single obj. ↑
0.9950
0.9980
0.9910
0.9980
0.9950
0.9969
1.0000
1.0000
Two obj. ↑
0.9180
0.9080
0.7300
0.9120
0.8920
0.8788
0.9190
0.9310
Counting ↑
0.6550
0.6320
0.5750
0.6150
0.5840
0.3156
0.6000
0.6720
Table 1: Main generation results. Bold marks the best 1-NFE value; the 50-NFE and 4-NFE columns are multi-step references. External published anchors are listed separately in Table 3 .
Table 2: Every recorded variant and the scope it was measured in. Rows in different blocks use different evaluation stacks and must not be compared across the scope column. Sub-metrics report GenEval components (Single object maintains ≥0.990 across valid runs and is omitted for width). Widths are nominal image-FFN activation ratios, not latency or total FLOPs. Bold marks the best measured value within the registered benchmark scope.
Appendix figures & tables15 assets
Supplementary material from the paper’s appendix.
Appendix
Method
Backbone
Steps
Training data
MJHQ FID ↓
MJHQ CLIP ↑
GenEval ↑
Alternative one-step paradigms
EMF (MeanFlow)
BLIP3o-NEXT
1
∼ 170K pairs
—
—
0.74
SANA-Sprint (sCM+LADD)
SANA (1.6B)
1
SANA dist.
7.69
28.27
0.76
Distribution-matching and adversarial distillation
DMD2 (SDXL)
SDXL
1
3M LAION-A.
7.10
28.93
0.59
SDXL-Turbo (ADD)
SDXL (2.6B)
1
SA-1B / LAION
—
—
—
Appendix
Table 3: Reported performance anchors from the one- and few-step text-to-image generation literature on MJHQ-30K and GenEval, as published by the respective authors. These values are taken directly from their publications and are not re-evaluated on our frozen pipeline. Missing entries indicate that the corresponding work did not report that benchmark under the listed setting. No literature entry is bolded due to differences in backbones, training data, and evaluation protocols.
50 NFE
4 NFE
1 NFE
Metric
Dense Teacher
Dense DMD2
Dense DMD2
DeMoE 3E1A1S
DeMoE 7E4A1S
GenEval (Text-to-Image Alignment)
Overall ↑
0.564
0.548
0.532
0.468
0.541
Two obj. ↑
0.625
0.612
0.598
0.524
0.615
Counting ↑
0.382
0.365
0.342
0.185
0.358
Colors ↑
0.784
0.771
0.760
0.775
0.768
Appendix
Table 4: SDXL Performance. Bold marks the best 1-NFE result; multi-step references provide competitive baselines. Active parameters exclude the static text encoders ( ∼ 0.82B) and reflect UNet computation per step.
Figure 5: Qualitative comparison of text-to-image generation results. We compare our proposed method ( DeMoE ) with the multi-step Dense Teacher and the one-step distilled baseline Dense DMD2 on various challenging compositional prompts (categorized into spatial position , object counting , multi-object attribute binding, and color ). Evaluation feedback below each image indicates whether the generation fulfills the prompt requirements (e.g., detected object counts or spatial relations). While baseline models frequently suffer from missing objects, incorrect counts, and attribute leakage, DeMoE (Ours) reliably generates images with precise semantic alignments, accurate object quantities, and correct spatial configurations in a single step.
Control minus fitted
Mean
Marginal 95% interval
Fit: natural
0.084
[−0.327,0.487]
Eval: natural
0.536
[0.070,1.037]
Eval: global raw
0.957
[0.505,1.433]
Eval: global matched
1.197
[0.733,1.725]
Eval: prompt-shuffled
0.269
[−0.316,0.814]
Eval: matched-random
−0.149
[−0.703,0.474]
Appendix
Table 5: All registered route-control contrasts, in units of 10−3 . Positive means the named control has larger recomputed pseudo-loss than the fitted route, so positive values favour the fitted route.
Bank
Top-two centered energy
Minimum principal cosine
Construction
93.124%
1.000000
Confirmation
93.082%
0.999945
Second noise
92.942%
0.999956
Appendix
Table 6: Stability of the construction-fit two-axis teacher-response coordinate. Each bank contains 128 prompts; the second-noise bank reuses the confirmation prompts.
Figure 6: Numerical trainability and soft-to-hard switch checks. (a) BF16-visible update fraction and local response loss; (b) GenEval change at one earlier checkpoint. Both are implementation diagnostics outside the main configuration.
FFN Width Ratio
1-NFE Performance
Design
Shared S
Routed R
Active k
Width q
Total
Active
GenEval ↑
FID ↓
Latency (s) ↓
Shared fraction: total 1× , active 0.5× (ablation on shared capacity)
8E4A0S
0
8
4
1/8
1
1/2
0.7180
23.4500
0.2185
7E3A1S
1
7
3
1/8
1
1/2
0.7460
21.1500
0.2142
6E2A2S
2
6
2
1/8
1
1/2
0.7510
20.8200
0.2124
5E1A3S
3
5
1
1/8
1
1/2
0.7390
21.8000
0.2096
Appendix
Table 7: Expert configurations under explicit capacity control and measured 1-NFE performance. An expert’s intermediate width is q times the parent FFN width, yielding total FFN ratio q(S+R) and active FFN ratio q(S+k) . Performance columns report text alignment (GenEval Overall), distribution fidelity (COCO-10K FID). Bold marks the best performance within each ablation section.
Router
Shared response
History state
Action-response objective
Feedback objective
History in expert
Current hidden
—
—
—
—
—
Current + shared
Yes
—
—
—
—
GRU execution observer
Yes
GRU
—
—
—
Execution-aware TTT
Yes
Fast weights
—
—
—
AR-CURRENT
Yes
—
Yes
—
—
AR-HISTORY
Yes
Fast weights
Yes
—
—
Appendix
Table 8: Router design comparison on an identical expert bank and training protocol. Entries record which mechanism a design contains; they are not results. The lower block is a prospective extension and is not part of the validated configuration.
Figure 7: Executed write shares and soft preferences at all 13 routed depths.
Figure 8: Executed write shares and soft preferences (continued).
Figure 9: Executed write shares and soft preferences (continued).
Figure 10: Executed write shares and soft preferences (continued).
Figure 11: Executed write shares and soft preferences (continued).
Figure 12: Executed write shares and soft preferences (continued).
Figure 13: Executed write shares and soft preferences (continued).
Diffusion Transformers with Mixture-of-Experts (DiT-MoE) improve model capacity under sparse activation, but diffusion inference is still bottlenecked by redundant computation across timesteps. Existing caching methods mainly operate at the token level, which becomes suboptimal in DiT-MoE because each token update is internally decomposed into multiple routed expert branches. Our analysis shows that cross-timestep redundancy in DiT-MoE is better characterized at the expert-branch level than at the whole-token level. Based on this observation, we propose MoECa, a fine-grained caching framework that performs branch-level feature reuse across timesteps. MoECa further introduces expert-aware adaptive control and synchronized cache updates across MoE and attention paths to maintain stable intermediate states. Experiments on multiple DiT-MoE models show that MoECa consistently achieves a better speed-quality trade-off than prior caching methods, with up to 2.83× inference speedup and minimal quality degradation.
Maoliang Li, Haojing Chen, Jiayu Chen +4
School of Computer Science, Peking University, Beijing, China · School of Software Engineering, University of Electronic Science and Technology of China, Chengdu, China
Sampling from pretrained diffusion and flow-matching models typically requires many forward passes to generate diverse and high-fidelity images. Existing distillation methods often rely on multiple auxiliary networks, carefully designed training stages, or complex optimization pipelines. In this work, we revisit the recently proposed Drifting Model objective and show that a single drifting loss can be directly used to simplify one step distillation. A key observation is that the pretrained diffusion teacher itself already provides a strong representation space. Unlike the original Drifting Model, which relies on an additional pretrained feature extractor, we use intermediate hidden states of the pretrained teacher model as the feature representation. This removes the need for training or introducing an extra representation network while preserving a semantically meaningful feature geometry for drifting. Furthermore, we introduce a lightweight mode coverage loss to mitigate mode collapse during distillation and encourage the student generator to cover diverse teacher-supported regions. Extensive experiments on ImageNet and SDXL demonstrate that our method achieves efficient one step generation with competitive image quality and diversity, achieving FID scores of 1.58 on ImageNet-64×64 and 18.4 on SDXL, while substantially simplifying the overall distillation framework.
Yuan Zhang, Chenyi Li, Guoqing Ma +7
JD Explore Academy, China · The Hong Kong University of Science and Technology · Tsinghua University +1
Masked diffusion models (MDMs) generate sequences by progressively unmasking several tokens per denoising step, but their reverse process is typically factorized over positions, limiting sample quality in the few-step regime where diffusion's speed advantage over autoregressive decoding matters most. A recent line of work introduces a continuous Gaussian latent, trained as a variational autoencoder, to capture correlations across positions, but such approaches are prone to posterior collapse, where the latent is silently ignored. We propose Enhanced Mixture-of-Experts (E-MoE), which builds the reverse process as a mixture of factorized distributions over a discrete shared latent given by the expert-routing decisions of a Mixture-of-Experts (MoE) backbone, without increasing active parameters over the factorized baseline. Across synthetic multi-modal benchmarks, binarized MNIST, and LM1B, E-MoE improves few-step generation over factorized baselines.
Arseny Ivanov, Alexander Kolesov, Alexander Korotin +2
AXXX, Moscow, Russia · Applied AI Institute, Moscow, Russia