Expert Race: A Flexible Routing Strategy for Scaling Diffusion Transformer with Mixture of Experts | Summary

This article explains the key points of Expert Race: A Flexible Routing Strategy for Scaling Diffusion Transformer with Mixture of Experts.

• 2025-03-20 (arXiv)
• Yuan, Yike, Wang, Ziyu, Huang, Zihao, Zhu, Defa, Zhou, Xun, Yu, Jingyi, Min, Qiyang.
• ByteDance Seed, ShanghaiTech University
• Paper

Read this article in Korean

Summary

• The paper addresses how to allocate computation across image regions and denoising steps in diffusion transformers, which generate images by progressively removing noise with a Transformer network. On ImageNet 256 × 256, RaceDiT-XL/2-4in32 achieves Fréchet Inception Distance (FID), an image-generation quality metric where lower is better, of 2.06 versus 3.02 for the authors’ retrained DiT-XL/2*, with the same reported training exposure and similar activated parameter counts—the parameters used for computation—of 0.707B versus 0.669B.
• Expert Race selects the highest token–expert scores globally across the batch, image sequence, and expert dimensions; tokens represent image patches, and scores determine their expert assignments. Within a Mixture of Experts (MoE), a layer that routes inputs to selected feed-forward networks called experts, this permits variable computation per patch and denoising timestep. An exponential moving average (EMA), a running estimate that discounts older observations, supplies an inference threshold without competition between samples.
• Router similarity loss discourages experts from repeatedly processing identical token groups, while per-layer regularization directly supervises intermediate representations with the final denoising target. Cumulative component ablations reduce FID from 16.47 to 8.03 on ImageNet 256 × 256. The reported 7.2× acceleration measures iterations to a matched training loss, without establishing an equivalent reduction in elapsed training time.

1 Introduction

Image patches differ in information density, and denoising steps differ in reconstruction difficulty. The authors therefore propose allocating more experts to difficult tokens through a router, the module that scores token–expert assignments. Race-DiT replaces the Transformer’s multi-layer perceptron (MLP), its feed-forward subnetwork, with an MoE block and adds objectives for shallow-layer learning and expert diversity.

The routing diagram expands the selection pool from constrained slices to the full score tensor. The training curves report a 7.2× reduction in iterations to a matched loss, supporting faster optimization in iteration count without measuring elapsed training time.

Flattening the routing scores before selection allows multiple assignments to concentrate on some tokens while others receive none. This changes how feed-forward computation is allocated within the Transformer block.

2 Related Work

The related work connects sparse expert routing, which activates only selected experts, with methods that divide diffusion computation across timesteps, channels, or tokens. The paper positions its contribution as expanding the dimensions over which routing decisions compete.

2.1 Mixture of Experts

Token Choice lets each token select experts, usually with a fixed selection count; Expert Choice lets each expert select tokens. DYNMoE and ReMoE, earlier methods for variable expert allocation, already use thresholds and budget regularization to vary expert counts per token. Expert Race instead introduces global training-time competition across batch, sequence, and expert dimensions.

2.2 Multiple Experts in Diffusion

Earlier diffusion systems separate denoising stages into different models or partition computation through heterogeneous networks, channel groups, and masks. DiT-MoE, EC-DiT, and Raphael are diffusion architectures that learn expert assignments within the model. Expert Race extends this direction by relaxing constraints on which token–expert assignments can compete.

3 Preliminaries

The preliminaries introduce the denoising objective, weighted aggregation of selected expert outputs, and the rationale for adapting computation to spatial redundancy and timestep-dependent difficulty.

3.1 Diffusion Models

The forward process mixes a clean sample with Gaussian noise according to \(x_t=\sqrt{\bar\alpha_t}x_0+\sqrt{1-\bar\alpha_t}\epsilon\), where x_0 is the clean sample, ε is noise, and ᾱ_t decreases with timestep t. A conditioned network learns the reverse process by minimizing squared prediction error; the target may be noise, clean data, or velocity, a specified linear combination of noise and clean data. The diffusion Transformer (DiT) uses adaptive layer normalization (adaLN), normalization modulated by conditioning information.

3.2 Mixture of Experts

A router maps inputs with batch size B, sequence length L, and hidden dimension D to scores over E experts. Selected expert outputs are combined using their gating scores, the weights assigned to their contributions. In an x-in-y configuration, y experts are available and each expert’s intermediate width is divided by x; for Expert Race, x specifies the average activation budget rather than an identical count for every token.

3.3 The Rationality of Using MoE in Diffusion Models

The paper treats diffusion as multiple related prediction tasks because the reconstruction problem changes with timestep. Unequal information density across image patches provides a second motivation for adaptive routing. These observations motivate flexible allocation, although the router learns scores without explicit labels of token difficulty.

4 Taming Diffusion Models with Expert Race

The routing framework expresses existing and proposed strategies as selection of the highest scores from differently reshaped versions of the same score tensor. Expert Race removes the separate selection pools imposed by conventional routing.

4.1 General Routing Formulation

The score tensor is reshaped into a matrix with D_A independent selection operations and D_B candidates per operation. For an expected k activated experts per token and N_E total experts, each operation selects \(\mathcal K=\frac{k}{N_E}D_B\) entries. Token Choice fixes a budget per token, Expert Choice fixes one per sample–expert pair, and Expert Race pools all B × L × E scores to select B × L × k assignments.

The three cubes identify which score dimensions compete within each selection operation. Expert Race selects B × L × k assignments from the full B × L × E pool, whereas conventional strategies impose separate local budgets.

4.2 Expert Race

Expert Race selects the globally largest routing scores in one pass, maximizing their sum under the specified assignment budget. Identity gating leaves scores unchanged, avoiding the cross-token ordering changes introduced by expert-wise softmax, which converts each token’s expert scores into probabilities summing to one. During training, the selection cutoff updates an EMA threshold; inference compares each score against that stored threshold, allowing independent samples and variable assignment counts without recomputing a batch-wide ranking.

5 Load Balancing via Router Similarity Loss

Global routing flexibility creates a balancing problem because a few experts or expert groups can dominate assignments. The proposed loss regulates individual utilization and pairwise co-selection, meaning two experts receiving the same token.

5.1 Mode Collapse in Balancing Loss.

Equal token counts do not guarantee distinct expert behavior. The toy example gives every expert two tokens, yet two pairs of experts always receive identical inputs, making a 2-in-4 configuration behave like two wider experts. Conventional balancing of individual expert loads does not directly penalize this repeated co-selection.

Both examples give every expert the same workload, but the left example repeatedly assigns identical token sets to pairs of experts. Equal individual loads can therefore coexist with limited diversity of expert combinations.

5.2 Router Similarity Loss.

The method computes M′ = MᵀM and P′ = PᵀP, where M records binary token assignments and P contains expert-wise softmax probabilities derived from router scores. Router similarity loss weights entries of P′ using observed assignment counts, normalizing diagonal and off-diagonal terms separately. Diagonal terms regulate individual utilization, while off-diagonal terms discourage concentration in repeatedly co-selected expert pairs; softmax serves the auxiliary loss even though routing uses identity gating.

6 Per-Layer Regularization for Efficient Training

The authors observe rapidly increasing block-output norms in deeper layers, reducing shallow layers’ relative contribution to the residual stream, the representation accumulated through shortcut additions. Per-layer regularization (PLR), auxiliary prediction of the final denoising target from intermediate hidden states, provides direct squared-error supervision through a projection head integrated into the router. Figure 5 shows a less extreme depth imbalance with PLR, and the cumulative component ablation improves FID from 11.56 to 8.95 when it is added.

Without PLR, the final block has a much larger output norm than most shallow blocks. PLR increases several earlier contributions and reduces the final block’s dominance, supporting the authors’ account of depth imbalance.

7 Experiments

The experiments evaluate routing, gating, load balancing, auxiliary supervision, model size, and expert segmentation on ImageNet. Separate studies examine individual design choices and the complete Race-DiT training and inference procedure.

7.1 Implementation Details

Default training uses the AdamW optimizer, batch size 256, constant learning rate 1e-4, and no weight decay; ablations use the Base DiT model, DiT-B, and its 2-in-8 variant for 500K iterations unless specified otherwise. PLR receives weight 1e-2 and router similarity loss 1e-4, and evaluation uses EMA model weights. Besides FID, the paper reports CMMD, an image-distribution discrepancy metric where lower is better, and CLIP Score, an image–conditioning alignment metric where higher is better.

7.2 Routing Strategy

The allocation visualization shows more expert assignments at lower timestep indices, corresponding to later denoising stages that reconstruct finer detail. Spatial allocations also change with the emerging object. Token Choice and Expert Choice maintain fixed average allocations per timestep under the compared formulations, whereas Expert Race reallocates the budget across timesteps.

Expert Race changes both the spatial assignment map and the average number of experts across timesteps. Its larger allocation at lower timestep indices agrees with the authors’ interpretation that later denoising requires more detailed reconstruction.

7.3 Gating Function

With threshold and regularization components isolated, Expert Race obtains FID 16.47 with softmax, 13.85 with sigmoid, a score transformation bounded between zero and one, and 13.66 with identity gating. Its identity variant also outperforms identity-gated Token Choice at 15.98 and Expert Choice at 15.70. The authors attribute this advantage to preserving score ordering across tokens; identity is best for Expert Race, although sigmoid gives the best Token Choice FID in this table.

Expert Race performs best within each gating group, and identity-gated Expert Race achieves FID 13.66. Identity is not universally best across routing methods: Token Choice achieves its lowest listed FID, 15.25, with sigmoid.

7.4 Load Balance

For 4-in-32 MoE, router similarity loss achieves FID 10.77 versus 11.67 with conventional balance loss and 11.38 without a constraint. MaxVio, the largest expert-capacity violation, falls to 0.850 from 2.052 with balance loss, while Combination Usage, the percentage of expert pairs counted as active under the appendix’s cumulative-frequency criterion, rises from 72.36 to 83.10. Across 16, 32, 64, and 128 experts with four average activations, router similarity retains higher pair usage than the alternatives.

Router similarity improves all three reported measures over conventional balancing: FID falls from 11.67 to 10.77, MaxVio from 2.052 to 0.850, and Combination Usage rises from 72.36 to 83.10. Conventional balance loss improves utilization measures but slightly worsens FID relative to no constraint.

Pair usage declines as the candidate expert count increases for all three approaches. Router similarity maintains the highest usage throughout the tested range, while unconstrained routing becomes concentrated in very few combinations.

7.5 Core Components

The cumulative ablation reports FID 16.47 for softmax Expert Race, 13.66 after identity gating, 11.56 after the inference threshold, 8.95 after PLR, and 8.03 after router similarity loss. CMMD decreases from 0.7104 to 0.4587 and CLIP Score increases from 21.97 to 23.09 across the same sequence. Each gain is conditional on the components already present; the study does not test every component combination.

Every successive addition improves all reported metrics, with the largest FID reductions occurring after identity gating and PLR. Because the rows accumulate changes, their differences describe gains conditional on the components already present.

7.6 Scaling Law

Across Base, Medium, Large, and Extra-large backbones, 4-in-32 models improve FID from dense-model values of 18.03, 11.18, 7.88, and 6.31 to 7.35, 5.16, 4.04, and 3.31, respectively, with similar but slightly higher activated parameter counts. The Medium MoE uses 0.281B activated parameters and achieves FID 5.16 versus 6.31 for the 0.669B-parameter dense XL model, although its total parameter count is larger at 1.106B. At a fixed base-backbone activation budget, finer expert segmentation, which divides expert intermediate widths, and more candidate experts improve FID from 9.70 for 1-in-4 to 6.28 for 8-in-128.

MoE improves generation metrics across all four backbone sizes while using slightly more activated parameters and substantially more total parameters. The Medium MoE’s 1.106B total parameters exceed dense XL’s 0.669B, clarifying the distinction between active computation and stored model capacity.

The quality plot places MoE models below dense counterparts at comparable activated sizes, while bubble sizes show their larger total capacity. Training curves also favor MoE, supporting an optimization advantage alongside improved final quality.

At comparable total sizes, finer expert segmentation improves FID; adding candidates within each segmentation group also lowers FID. The best listed configuration, 8-in-128, reaches 6.28, although the other metrics are not strictly monotonic in every intermediate comparison.

The plotted configurations connect finer segmentation and larger candidate pools with lower FID, while training losses remain relatively close. These finite comparisons demonstrate a scaling trend without independently establishing a general linear law.

7.7 Extended Routing Strategy

All extended-routing variants use the full DiT-B/2-MoE-2-in-8 procedure, 500K training iterations, and inference thresholds. BL Choice pools batch and sequence positions, BE Choice pools batch and experts, and LE Choice pools sequence positions and experts. Expert Race achieves FID 8.03, compared with 9.50 for Token Choice, 10.13 for Expert Choice, 9.08 for BL Choice, 8.28 for BE Choice, and 8.89 for LE Choice, supporting the benefit of pooling all three dimensions in this setting.

The shaded selection regions distinguish one-dimensional, pairwise, and fully pooled routing choices. They specify which budget constraints are relaxed by each variant tested in Table 6.

With the full procedure and inference thresholds shared, Expert Race gives the best FID, CMMD, and CLIP Score. BE Choice is the closest alternative at FID 8.28, showing that partial pooling recovers much of the improvement in this setting.

7.8 More Results on ImageNet 256 × 256

For the longer XL experiment, the prose specifies batch size 1024, 1.75M steps, and learning rate 1e-4, while Table 7 displays training exposure as 1.7M × 1024 for both retrained models. Classifier-free guidance, a sampling control that strengthens conditioning, is set to 1.5 for metrics and 4.0 for displayed images. RaceDiT-XL/2-4in32 improves FID from 3.02 to 2.06 against DiT-XL/2*, but Table 7 also contains lower FID values from differently trained models, including DiT-MoE-XL/2 at 1.72.

The samples span animals, objects, and scenes, illustrating the range of generated ImageNet categories. Their guidance scale differs from the scale used for quantitative evaluation, so the montage provides qualitative evidence.

The paired retraining rows show FID improving from 3.02 to 2.06 at the same displayed exposure of 1.7M × 1024. Other rows use different architectures and training budgets; DiT-MoE-XL/2 reaches FID 1.72 with 1.530B activated parameters and 7.0M × 1024 training exposure.

8 Conclusion

The paper concludes that global routing, an EMA inference threshold, intermediate supervision, and pairwise expert regularization support scaling diffusion transformers on ImageNet 256 × 256. Extension to other diffusion-based visual tasks is presented as future work.

Appendix

• A Implemention of the Per-Layer Regularization: A two-layer MLP router first transforms the hidden state while retaining its width, then branches into a gating head and a target-prediction head. Each layer’s target prediction receives squared-error supervision against the final denoising target.
• B Analysis of Router Similarity Loss: The appendix interprets the loss as extending individual expert balancing to pairwise interactions. It derives the diagonal contribution using squared routing probabilities and explains separate normalization of diagonal and off-diagonal terms because actual expert selection occurs without replacement.
• C Combination Usage: Counts for all unordered expert pairs are sorted in descending order and normalized. The metric divides the number of bins whose cumulative mass remains below 95% by the total number of possible pairs, measuring how broadly pairwise assignments are distributed.
• D Additional Comparisons with DiT-MoE: At 500K iterations, Table 8 reports FID 8.02 for RaceDiT-B/2-2in8 versus 9.06 for DiT-MoE-B/2-8E2A, with 0.135B versus 0.286B activated parameters. The comparison also changes the expert implementation and shared-expert design, so it evaluates complete architectures rather than routing alone; Table 8’s 8.02 differs slightly from the 8.03 reported in the main ablations.
• E Additional Image Generation Results: Figures 15–22 show RaceDiT-XL/2-4in32 samples at 256 × 256 with classifier-free guidance 4.0 for Macaw (88), loggerhead turtle (33), lion (291), lesser panda (387), convertible (511), balloon (417), lakeside (975), and volcano (980). They illustrate within-class variation but do not provide a controlled comparison with other methods.

Brief Thoughts

Table 6 provides the clearest comparison of routing strategies: the methods share the full procedure, training duration, and inference-threshold mechanism, and Expert Race performs best on all three reported metrics. Table 2 separately supports regulating expert-pair usage in addition to individual expert loads.

The efficiency evidence is limited to approximately matched activated parameter counts and faster optimization in iteration count; the paper does not benchmark elapsed-time throughput or communication costs. Table 4 also contradicts the nearby prose claiming better performance with less than half the total parameters of dense XL: the Medium MoE has fewer activated parameters but more total parameters. Broader task transfer and variability across repeated training runs remain unestablished by the reported experiments.