Pixel-space diffusion models have long appealed as an alternative to latent-space approaches, offering direct modeling of raw images without dependence on a learned tokenizer. Yet scaling these models has followed dense strategies, binding model capacity tightly with active computation. The authors of this work demonstrate that standard LLM-style Mixture-of-Experts architectures transfer effectively to pixel-space diffusion, introducing PixelMoE—a framework that sparsely scales capacity without diffusion-specific routing mechanisms or auxiliary routing objectives.

Token-Choice Routing for Input-Adaptive Sparse Computation

PixelMoE's routing mechanism combines standard Token-Choice (TC) selection with fine-grained experts that compute on sparsely selected image patches. Unlike custom routing designs used in prior diffusion work, PixelMoE reuses the TC strategy familiar from LLM MoE systems. Each token—here, each pixel patch—is assigned to one or more experts based on a simple choice function, and the resulting sparse activations enable parameter-efficient computation. A shared expert runs in parallel with the routed experts, capturing representations common across inputs and ensuring that frequent patterns remain accessible without increasing the active parameter count.

The load-balancing strategy relies on a bias term added to each expert's routing logits. This bias adjusts dynamically based on per-expert usage statistics, encouraging underutilized experts to receive more traffic while preventing any single expert from dominating. The approach requires no auxiliary loss terms or auxiliary classifier free guidance; the bias term alone suffices to maintain balanced expert load across training.

Pixel-Space Hierarchical Flow Matching

Complementing the sparse capacity scaling, PixelMoE introduces Pixel-Space Hierarchical Flow Matching, a schedule that varies spatial resolution along the denoising trajectory. Early steps operate at lower resolution, allowing the model to form global structure before progressively increasing resolution for fine-grained detail refinement. This resolution annealing mirrors the way coarse-to-fine refinement works in traditional image synthesis, but is embedded directly into the flow-matching objective without requiring separate super-resolution or up-sampling modules.

The hierarchical schedule is defined by a resizing function that maps each timestep to a spatial grid size. At large timesteps (early denoising), the grid corresponds to a downsampled representation; as timesteps decrease, the grid resolution increases. This gradual escalation permits the model to allocate its sparse computation budget where it matters most—coarse layout in early steps, pixel-level texture in later steps.

Sparse Scaling Results on ImageNet-1K

On class-conditional ImageNet-1K, PixelMoE achieves an FID score of 1.47 at 256×256 and 1.53 at 512×512. These results compare favorably with leading pixel-space and latent-space generative models, all while operating at comparable active parameter counts. The paper reports results across model scales, showing that PixelMoE's sparse strategy consistently outperforms dense counterparts when both are constrained to the same FLOP budget.

The improvement stems from two factors. First, the MoE architecture expands total parameter count without proportionally increasing active parameters during inference or training sampling. Second, the hierarchical flow matching schedule ensures that sparse experts are engaged where they contribute most, avoiding the wasted computation that plagues uniformly dense models at high resolution.

Limitations and Trade-offs

The authors acknowledge that PixelMoE introduces additional complexity in expert initialization and fine-tuning. While the bias-based load-balancing eliminates the need for auxiliary routing losses, it does require careful tuning of the bias update coefficient to avoid instability in very wide expert panels. The paper also notes that the hierarchical resolution schedule adds a modest overhead in schedule coordination, though this is amortized by the reduction in active parameters per step.

Another trade-off involves the shared expert. By concentrating common representations in a single expert, the model risks creating a bottleneck if the shared path becomes overloaded during certain data distribution shifts. The authors mitigate this by periodically resetting the shared expert's bias, but the approach may not scale to extremely large expert counts without additional modifications.

Practical Takeaways for Working Developers

Developers familiar with LLM MoE systems can transfer many of the same patterns directly to pixel-space diffusion. The Token-Choice routing schema, the bias-based load balancer, and the shared expert pattern all map cleanly from token-level to patch-level computation. The hierarchical flow matching schedule can be implemented as a simple resizing function appended to the existing denoising loop, requiring no changes to the core U-Net or diffusion objective.

For teams looking to scale diffusion models beyond the parameter limits of dense architectures, PixelMoE demonstrates that LLM-style sparsity works in pixel space. The reported FID improvements suggest that sparse capacity scaling does not come at the cost of quality, and the absence of perceptual losses or representation alignment simplifies the training pipeline. Code and model checkpoints are slated for release, which will lower the barrier for experimentation.