PixelMoE: the LLM mixture-of-experts recipe works in pixel space

Mixture-of-experts is how frontier language models grow total capacity while keeping the compute per token roughly flat. Diffusion has not had the same experience. Earlier attempts to drop LLM-style MoE into diffusion transformers often failed to clearly beat dense baselines, and the field reacted by inventing diffusion-specific machinery: expert-choice routing, timestep-aware routers, semantic routing objectives. PixelMoE (Dong, Bao, et al., arXiv:2610.18857), from USTC, KlingAI and Peking University, makes the opposite argument. Take the standard LLM recipe as it is, apply it to raw pixels, and it works. I think that is the more useful result, and the ablations are what persuade me.

Why pixel space, and why sparsity

Most production image models denoise in a VAE latent space. That is cheap, but the autoencoder is lossy, so it sets a ceiling on fidelity and can introduce its own artifacts. Pixel-space diffusion skips the tokenizer and models raw images directly. Recent work such as JiT's clean-image prediction and several hierarchical architectures has made it competitive. Every one of those models is dense, though. If you want more capacity, you pay for it in active compute.

MoE separates the two. You can add experts, and with them total parameters, while each token still only passes through a fixed number of them. The question the paper asks is narrow and testable: does the off-the-shelf LLM version of MoE transfer to pixel-space diffusion, or does it need special handling?

The architecture is deliberately boring

PixelMoE follows the DeepSeek-style layout. The first transformer blocks keep dense FFNs, and the FFNs in the remaining blocks become sparse MoE layers. Each MoE layer has three parts:

• Token-choice routing. A bias-free linear router produces sigmoid scores over E routed experts. Each pixel token picks its top-A. The default configuration is E32A4: 32 fine-grained experts, 4 active per token.
• A shared expert. It is always on and handles representations common to all tokens. Its output is added to the weighted sum of the routed experts.
• Bias-based load balancing. Each expert gets a bias that is added to its score only when experts are being selected. The bias is nudged after each batch toward a uniform load, and there is no auxiliary loss.

The routed and shared experts use narrower FFNs, so the active parameter count stays close to the dense baseline's. The router sees only token hidden states. It gets no timestep embedding and no class signal.

Hierarchical flow matching

The second contribution takes advantage of something latents do not offer: a downsampled pixel image is still a meaningful image. PixelMoE splits denoising into two stages that share one backbone and one DDT prediction head. Stage 0 denoises partway at half resolution to set the global structure. The coarse endpoint is then upsampled, rescaled, and given block-correlated noise so that its statistics match the starting distribution of Stage 1. Stage 1 finishes at full resolution.

The two stages see tokens at different resolutions and noise levels. That heterogeneity suits input-adaptive routing, and the authors argue the MoE and the hierarchy complement each other. Training uses stage-aware timestep sampling. One logit-normal distribution decides which stage a sample is drawn for, and a second sets the timestep within that stage. The only generation loss is a velocity-space endpoint objective.

What the numbers say

The controlled comparisons are the core of the paper. Each dense baseline shares the hierarchical flow and the training recipe, and only the FFNs change. After 100 epochs on ImageNet-1K at 256:

The gap holds at every scale. Holding active parameters at 632M and going from 8 to 16 to 32 experts moves FID from 2.28 to 2.03 to 1.78. PixelMoE-L also reaches Dense-L's 200-epoch FID after about 71 epochs, which the authors put at roughly 2.8x fewer epochs.

After a full 320 epochs, PixelMoE-H-E32A4 (3.68B total, 1.47B active) reaches FID 1.47 at 256x256 and 1.53 at 512x512. In the paper's comparison tables, that is ahead of strong dense pixel-space models such as HyperDiT-H (1.56) and JiT-G (1.82), and ahead of the latent-space DiT-MoE-XL (1.72). REPA, a latent model, still reports 1.42 at 256. The L model is a useful reference point: FID 1.54 with 632M active parameters.

The ablations carry the argument

Several ablations go against the diffusion-specific designs that came before:

• Token-choice vs expert-choice. The authors reimplemented DiffMoE's expert-choice routing on the same backbone. It scored 2.62 against PixelMoE's 1.78, so most of the dense-to-sparse gain disappears.
• Explicit routing guidance hurts or does nothing. Feeding the timestep to the router gives 2.10, and ProMoE-style semantic routing gives 1.85. The default router, which sees neither, gets 1.78.
• Bias balancing beats an auxiliary loss (1.78 vs 2.05) at identical parameter counts.
• The shared expert matters. Removing it costs FID (2.22 vs 1.78) even though the model then has more total parameters.
• Granularity matters. E8A1 and E32A4 have the same parameter budget, but the second allows 35,960 expert combinations instead of 8. FID goes from 1.95 to 1.78.

The routing analysis explains why explicit guidance is unnecessary. Experts specialize by class in the middle layers and by position along the denoising trajectory in the shallow and final layers, and global load stays close to uniform. The hidden states already carry the timestep and semantic information. A router that is free to read them will find it.

My read

The pattern I take from this paper is common in ML engineering. When a well-understood method fails in a new domain, the cause is often the surrounding configuration and not the method. This paper attributes the earlier disappointing results to missing pieces (shared experts, fine granularity, bias balancing), and none of them is new. That matters in practice. A team already running MoE language models can reuse its routing kernels, load-balancing code and serving instincts for image generation without maintaining a separate set of diffusion-specific tricks.

The caveats are the usual ones. Every result is class-conditional ImageNet. There is no text-to-image or video yet, which is where KlingAI's interest presumably lies. The paper reports active parameters but not wall-clock throughput or memory, and MoE inference has real costs in all-to-all communication and in holding 3.68B parameters in memory to use 1.47B. The hierarchical flow gain is small in isolation (2.54 to 2.46). Code is listed as "coming soon."

Still, the main claim is clean and well controlled: sparse scaling transfers to pixel space with no diffusion-specific routing. If it holds up on text-conditioned generation at scale, the best pixel-space model in a year or two will probably look like an LLM in its MoE layers.