Functional Cross-layer Experts Routing for One-step Diffusion Distillation

Zile Huang; Ser-nam Lim

cs.CVOct 19, 2026

Functional Cross-layer Experts Routing for One-step Diffusion Distillation

Authors: Zile Huang, Ser-nam Lim

Organizations: University of Central Florida

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.

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

R_{1: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.

GenEval
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
Overall $\uparrow$	0.8245	0.8035	0.6888	0.7865	0.7771	0.7415	0.7812	0.8220
Single obj. $\uparrow$	0.9950	0.9980	0.9910	0.9980	0.9950	0.9969	1.0000	1.0000
Two obj. $\uparrow$	0.9180	0.9080	0.7300	0.9120	0.8920	0.8788	0.9190	0.9310
Counting $\uparrow$	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 .

Figure 3: Functional allocation evidence. Response alignment, input-dependent execution, ownership, and independent-draw route work.

Figure 4: Routing demand and controller evidence. Crossed prompt/noise structure, state-based ranking, and finite routing against global controls.

Registered full-population stack (same evaluation pipeline)
Variant	Scope	Over.	Count	Two obj.	Colors	Pos.	Col. attr.	Width
Dense DMD2-1, 5K	anchor	0.7865	0.6150	0.9120	0.9310	0.6840	0.5790	1.000
Earlier sparse v2	3E1A1S , $1/4$ w.	0.7415	0.3156	0.8788	0.9654	0.6925	0.6000	0.500
Progressive sparse	3E2A1S $\rightarrow$ 3E1A1S	0.7054	0.4520	0.8120	0.8980	0.5540	0.5244	0.500
DeMoE	7E3A1S	0.7812	0.6000	0.9190	0.9260	0.6720	0.5700	0.740
DeMoE	7E4A1S	0.8220	0.6720	0.9310	0.9420	0.7220	0.6650	0.805

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

\geq 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

Alternative one-step paradigms
Method	Backbone	Steps	Training data	MJHQ FID $\downarrow$	MJHQ CLIP $\uparrow$	GenEval $\uparrow$
EMF (MeanFlow)	BLIP3o-NEXT	1	$\sim$ 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.

GenEval (Text-to-Image Alignment)
	50 NFE	4 NFE	1 NFE
Metric	Dense Teacher	Dense DMD2	Dense DMD2	DeMoE 3E1A1S	DeMoE 7E4A1S
Overall $\uparrow$	0.564	0.548	0.532	0.468	0.541
Two obj. $\uparrow$	0.625	0.612	0.598	0.524	0.615
Counting $\uparrow$	0.382	0.365	0.342	0.185	0.358
Colors $\uparrow$	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 (

\sim

0.82B) and reflect UNet computation per step.

Panel of Figure 5 — 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 $\uparrow$	FID $\downarrow$	Latency (s) $\downarrow$
Shared fraction: total $1\times$ , active $0.5\times$ (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.

Panel of Figure 7 — Figure 7: Executed write shares and soft preferences at all 13 routed depths.

Panel of Figure 8 — Figure 8: Executed write shares and soft preferences (continued).

Panel of Figure 9 — Figure 9: Executed write shares and soft preferences (continued).

Panel of Figure 10 — Figure 10: Executed write shares and soft preferences (continued).

Panel of Figure 11 — Figure 11: Executed write shares and soft preferences (continued).

Panel of Figure 12 — Figure 12: Executed write shares and soft preferences (continued).

Panel of Figure 13 — Figure 13: Executed write shares and soft preferences (continued).

Diffusion Model Distillation

Diffusion model distillation transfers a large model’s generative behavior to faster or smaller models that need fewer sampling steps.

See all 183 papers

MoECa: Aligning Feature Reuse with Expert Decomposition in Diffusion Transformers
Jun 14, 2026Maoliang Li, Haojing Chen, Jiayu Chen +4Diffusion Transformer Diffusion Model Caching
School of Computer Science, Peking University, Beijing, China · School of Software Engineering, University of Electronic Science and Technology of China, Chengdu, China
Teacher-Feature Drifting: One-Step Diffusion Distillation with Pretrained Diffusion Representations
May 8, 2026Yuan Zhang, Chenyi Li, Guoqing Ma +7Diffusion Model Distillation Diffusion Sampling
JD Explore Academy, China · The Hong Kong University of Science and Technology · Tsinghua University +1
E-MoE: Enhanced Mixture-of-Experts for Non-Factorized Diffusion Language Models
Sep 29, 2026Arseny Ivanov, Alexander Kolesov, Alexander Korotin +2Masked Diffusion Models Few-Step Diffusion Sampling
AXXX, Moscow, Russia · Applied AI Institute, Moscow, Russia

Jun 14, 2026J/K move · Enter open · S save

MoECa: Aligning Feature Reuse with Expert Decomposition in Diffusion Transformers

Maoliang Li, Haojing Chen, Jiayu Chen, Zihao Zheng+3

School of Computer Science, Peking University, Beijing, China · School of Software Engineering, University of Electronic Science and Technology of China, Chengdu, China

Diffusion Transformer Diffusion Model Caching Mixture-of-Experts Models Diffusion Model Inference Acceleration

Open paper AskPDF arXiv

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

\times

inference speedup and minimal quality degradation.

Abstract

Diffusion Model Distillation

Explore similar work

MoECa: Aligning Feature Reuse with Expert Decomposition in Diffusion Transformers

Teacher-Feature Drifting: One-Step Diffusion Distillation with Pretrained Diffusion Representations

E-MoE: Enhanced Mixture-of-Experts for Non-Factorized Diffusion Language Models