PixelMoE: LLM-Style MoE Sparse Scaling for Pixel-Space Diffusion Models

The field of diffusion models has long faced a choice of approach: model in latent space, or model directly in pixel space. The latent-space approach compresses images with a pretrained tokenizer, offering high computational efficiency, but introduces a lossy intermediate representation. The pixel-space approach models raw images directly, which is theoretically cleaner, yet has consistently struggled with "scaling difficulty"—model capacity and computation are tightly coupled, so getting stronger means burning more compute.

This is exactly where PixelMoE comes in. It demonstrates one thing: the mature Mixture-of-Experts (MoE) architecture from the LLM world can be directly transferred to pixel-space diffusion models, without any diffusion-specific routing mechanism.

The Scaling Bottleneck of Pixel-Space Diffusion

The core advantage of latent-space diffusion (such as the Stable Diffusion series) lies in dimensionality reduction: compressing a 512×512 image into a 64×64 latent representation, drastically shortening the Transformer's sequence length and thereby reducing computational cost. The cost is that the tokenizer itself loses high-frequency details, capping generation quality at the ceiling of reconstruction quality.

Pixel-space diffusion skips this bottleneck, performing denoising directly at high-resolution pixels. But existing pixel-space models (such as JiT) adopt a dense scaling strategy—every increase in model capacity comes with a proportional increase in inference computation. This means that under a fixed inference budget, it's hard to obtain better generation quality simply by enlarging parameter count.

The LLM field has long had a mature solution to this problem: MoE. Through sparse activation, the total parameter count can far exceed the parameters actually used in each forward pass. But no one had previously verified whether standard LLM-style MoE could be directly applied to pixel-space diffusion.

What Is Dense Scaling vs. Sparse Scaling

Dense scaling means every parameter of the model is used in each forward pass, so total parameter count is roughly proportional to single-inference computation (FLOPs). This has a direct consequence: under a fixed inference budget, improving model capacity requires increasing computational cost—the two cannot be decoupled.

The core idea of sparse scaling is to introduce conditional computation: the model has a large number of parameters, but each input activates only a small subset. MoE is the most typical implementation. At this point "total parameter count" and "active parameter count" become two independent dimensions—the former determines the upper bound of model expressiveness, the latter determines inference cost. PixelMoE leverages exactly this decoupling to expand total capacity under a fixed active parameter count, thereby breaking through the scaling bottleneck that has long trapped pixel-space diffusion.

PixelMoE's Core Design

Standard Token-Choice Routing, No Diffusion Specialization

The most noteworthy design decision of PixelMoE is "what not to do." It does not design a diffusion-specific routing mechanism, does not introduce auxiliary routing objectives, but directly adopts the standard Token-Choice (TC) routing from LLMs.

Specifically, each token selects the most matching subset of experts for sparse computation based on its input content, while retaining a shared expert to capture the general representations needed by all inputs. Load balancing uses a bias-based strategy to avoid expert utilization collapse.

This combination has been repeatedly validated in LLMs, and PixelMoE's contribution is proving it works equally well in pixel-space diffusion—there's no need to reinvent routing logic just because "the denoising process of diffusion models has a special timestep structure."

Token-Choice Routing and Load Balancing

Token-Choice (TC) routing is the most common form in LLM MoE: the routing network generates scores for each token across all experts, then selects the Top-K experts to process that token. The counterpart is Expert-Choice routing, where each expert chooses the tokens it will process. TC routing is simple to implement and naturally supports autoregressive and variable-length inputs, making it the de facto standard.

Its inherent risk is imbalanced expert utilization—a few experts get repeatedly selected while the rest sit idle, leading to wasted capacity and unstable training. Common countermeasures include: additive bias adjusting routing scores, auxiliary load-balancing losses, expert capacity limits, etc. PixelMoE adopts a bias-based strategy, correcting routing logits through learnable or dynamically adjusted biases to pull tokens toward underutilized experts, thereby maintaining balance without interfering with the main task loss. This shows the mechanism also applies to tokens distributed by spatial position and timestep in diffusion denoising.

Pixel-Space Hierarchical Flow Matching

Sparse scaling solves the capacity problem, but pixel-space modeling has another challenge: different denoising stages have different spatial resolution requirements.

PixelMoE introduces Pixel-Space Hierarchical Flow Matching, allowing spatial resolution to vary along the denoising trajectory. Early stages focus on global structure formation, requiring larger receptive fields and lower resolution; later stages focus on detail refinement, requiring higher resolution. This coarse-to-fine strategy complements sparse capacity scaling, together supporting the model's scaling efficiency in pixel space.

Flow Matching and Hierarchical Generation

Flow Matching is a class of methods for training continuous-time generative models. It learns a time-dependent velocity field that transports a noise distribution along a straight or approximately straight probability path to the data distribution, and can be viewed as a more concise, more training-stable generalization of diffusion models, with sampling carried out by ODE solvers.

Hierarchical generation draws on the classic coarse-to-fine idea: first generate low-resolution results focused on global structure, then progressively increase resolution to refine details. PixelMoE combines the two, letting resolution vary along the denoising time axis—early on using low resolution and large receptive fields to capture layout, later using high resolution to focus on texture. This dynamic resolution design not only fits the coarse-to-fine nature of the diffusion process but also complements sparse capacity scaling—the former optimizes computation allocation at each stage, the latter expands model capacity, and together they improve the scaling efficiency of pixel-space diffusion.

Experimental Results

On the class-conditional ImageNet-1K benchmark, PixelMoE achieves FID 1.47 at 256×256 resolution and FID 1.53 at 512×512. The significance of these two numbers lies in: no representation alignment was used, and no perceptual losses were used.

Compared with dense versions, at equivalent active parameter counts, PixelMoE consistently outperforms the dense counterparts across all model scales. This directly validates the effectiveness of sparse capacity scaling in pixel space—increasing total parameter count does not mean a proportional increase in computation, yet generation quality does improve.

Compared with current leading pixel-space and latent-space generative models, PixelMoE's FID results are also competitive. Considering that pixel-space models are usually thought to be at a disadvantage in efficiency, this result suggests that sparse architectures may change the fundamental economics of this route.

FID and What It Means for Evaluation

FID (Fréchet Inception Distance) is a common metric for measuring the difference between the generated image distribution and the real image distribution; lower values indicate that generation quality and diversity are closer to real data. It extracts features through the Inception network and computes the Fréchet distance between the two feature distributions. FID is relatively sensitive to whether a model uses training tricks such as perceptual losses and representation alignment; therefore, PixelMoE achieving 1.47 (256×256) and 1.53 (512×512) under the premise of "no representation alignment, no perceptual losses" better reflects the gains from the architecture itself. It should be noted that FID is affected by sampling count, preprocessing, and Inception version, so cross-paper comparisons should be based on the same evaluation protocol.

Why This Matters

PixelMoE's value is not in refreshing some benchmark, but in opening a path: the scaling of pixel-space diffusion models can reuse the MoE engineering experience and best practices already accumulated by the LLM community.

MoE training infrastructure, routing strategies, load-balancing schemes, parallelization implementations—these have been refined in the LLM field for years. If pixel-space diffusion can directly borrow this toolchain, its scaling curve will be far steeper than previously expected.

Another implicit signal is architectural convergence. Transformers moved from NLP to vision, MoE is moving from LLMs to diffusion models, and this cross-domain component reuse is accelerating. For those working on model architectures, what's worth paying attention to is not the FID number of a specific model, but "which designs are truly domain-agnostic"—PixelMoE's answer is: standard MoE routing is one of them.

One-line Summary

PixelMoE proves that LLM-style MoE routing can be directly used for pixel-space diffusion, enabling efficient scaling.

Paragraph Summary

PixelMoE directly transfers the mature Mixture-of-Experts architecture from LLMs to pixel-space diffusion models, without needing diffusion-specific routing mechanisms. It adopts standard Token-Choice routing, shared experts, and bias-based load balancing, combined with Pixel-Space Hierarchical Flow Matching to vary resolution from coarse to fine along the denoising trajectory, solving the problem of pixel-space models coupling capacity with computation and struggling to scale. On ImageNet-1K, it achieves FID 1.47 and 1.53 at 256×256 and 512×512 resolutions respectively, without using representation alignment or perceptual losses. At equivalent active parameter counts, PixelMoE consistently outperforms dense solutions across all scales, validating the effectiveness of sparse scaling in pixel space. Its core significance lies in opening a path for pixel-space diffusion to reuse the LLM community's MoE engineering experience, accelerating cross-domain architectural convergence.