PixelMoE: Scaling Pixel-Space Diffusion with LLM-Style Mixture-of-Experts

What it claims

Standard LLM-style MoE, used as-is, gives clear gains in pixel-space diffusion. Earlier work had found otherwise. Fei et al. 2024 saw no consistent gains, and JiT-MoE reported only modest gains over dense JiT, so later methods moved to Expert-Choice routing or diffusion-specific routing objectives. The authors report three further findings:

• Token-Choice routing "substantially outperforms" Expert-Choice.
• Adding more routed experts consistently improves FID.
• Experts specialise by semantic class and by denoising stage, even though the router never sees the class or the timestep.

The headline model is PixelMoE-H-E32A4, with 32 routed experts and the top 4 active per token.

How it works

Backbone. The first Transformer blocks keep dense FFNs, and later blocks swap their FFN for an MoE layer, following the DeepSeek/Kimi pattern. A DDT (Decoupled Diffusion Transformer) prediction head sits on top.

Router. A bias-free linear router produces sigmoid scores over E experts.

• A per-expert balancing bias b_e is added to the scores only when picking the top-A experts.
• The weights actually used are the raw scores, renormalised over the selected experts.
• The routed sum is multiplied by λ_r = √A, and the shared expert's output is added on top.
• Routed and shared experts are narrowed FFNs, sized so that active parameters roughly match the dense baseline.

flowchart LR
  T[pixel token] --> R[sigmoid router]
  R -->|score + bias, select only| K[top-A of E routed experts]
  K -->|renormalised scores × √A| S((sum))
  T --> SH[shared expert, always on] --> S

Load balancing. There is no auxiliary loss. After each batch the bias is updated as b ← b − γ (F − Q) / RMS(F − Q), with γ = 0.001. Here F is the batch's expert-load distribution and Q is uniform.

Hierarchical flow matching. One shared model is trained on two stages.

1. Stage 0 denoises from noise at half resolution (2× bilinear downsample) and stops at a noisy coarse endpoint t_e0.
2. Stage 1 starts at full resolution from the coarse image (down then 2× nearest-neighbour up), re-noised to level t_s1, and denoises to the clean image.
3. Noise scale δ_n is 1 at 256×256 and 2 at 512×512.

Training objective. Each stage predicts its own endpoint, JiT-style clean prediction, converted to a velocity. The loss is MSE on velocity, which equals an endpoint MSE weighted by 1/σ². σ is clamped at σ_min = 0.05 both in training and when computing velocity during sampling. This is the only loss used: no REPA and no perceptual loss.

Evidence

• Dense vs. MoE at matched active size. Dense-B (223M parameters, all active) scores FID 4.40. PixelMoE-B-E32A4 (461M total, 225M active) scores 2.46. The fetched table is cut off after this row, so the L and H comparisons are not visible as numbers.
• Figure 2 (100 training epochs). MoE has lower FID than dense at B, L and H sizes for comparable active parameters. The figure plots this over roughly 0.2–1.4B active parameters.
• Headline FID. Class-conditional ImageNet-1K: 1.47 at 256×256 and 1.53 at 512×512. The authors compare against JiT-G, PixelDiT-XL, PixelREPA-H, PixelFlow-XL and Simple Diffusion. Exact competitor numbers, sampler settings and CFG for these results are not in the fetched text.

Limits

• Truncated page. The text ends partway through the first results table. Epochs, compute, sampling steps, t_e0 / t_s1 values and the ablation numbers (TC vs. EC, expert count) are not available here.
• Memory cost. Total parameters are roughly double the active count (461M vs. 225M at B scale), so memory grows even though per-token compute does not.
• Narrow evaluation. Everything reported is class-conditional ImageNet. No text-to-image or other conditioning is tested.
• Nothing to run yet. Code and weights are "will be released soon". There is no repository, licence or run instructions.