One-step Latent-free Image Generation with Pixel Mean Flows

AI summary

Overview

Research area: Generative modeling for computer vision — specifically diffusion/flow-based image generation, one-step sampling, and pixel-space (latent-free) Transformers.

Technical level: Advanced. The paper assumes familiarity with Flow Matching, MeanFlow, consistency models, the manifold hypothesis, and Transformer-based diffusion architectures such as DiT.

Scope: The paper introduces "pixel MeanFlow" (pMF), a formulation that separates the network's output space from its loss space to enable one-step image generation directly in raw pixel space, reporting 2.22 FID at 256×256 and 2.48 FID at 512×512 on ImageNet.

What This Paper Is About

Modern diffusion and flow-based image generators typically rely on two things: multi-step sampling and operation in a compressed latent space. Recent work has tackled these two properties separately, but combining them — one-step generation without latents — places a much heavier burden on the neural network. This paper asks how to design a network that can do both at once, proposing a formulation in which the network directly predicts clean image pixels while the training loss is still defined in a velocity space.

Key Contributions

1. A new field definition linking average velocity to denoised images. The authors induce the field x(z_t, r, t) ≜ z_t − t · u(z_t, r, t), which behaves like a denoised image and is hypothesized to lie approximately on a low-dimensional data manifold.

2. Separation of prediction space and loss space. pMF makes the network output an x-prediction (image-manifold quantity) while the training objective remains the MeanFlow v-loss, requiring a conversion chain of x → u → v. This is contrasted against DiT (ε-pred, ε-loss), SiT (v, v), MF (u, u), iMF (u-pred, v-loss), and JiT (x-pred, v-loss).

3. A toy-experiment validation of the manifold hypothesis. Using a Swiss roll projected into a D-dimensional observation space for D ∈ {2, 8, 16, 512}, the authors show x-prediction works reasonably well while u-prediction fails in high-dimensional observation spaces.

4. Demonstration that one-step, latent-free generation is competitive. pMF reaches 2.22 FID at 256×256 and 2.48 FID at 512×512 on ImageNet, and the authors note it fills a "key missing piece" in the one-step, latent-free regime.

Main Findings

• x-prediction is essential at high dimensionality. At 64×64 with a B/4 model and patch dimension of 48, x-prediction and u-prediction perform similarly (3.80 vs. 3.82 FID). At 256×256 with B/16 and patch dimension 768, x-prediction reaches 9.56 FID while u-prediction fails catastrophically at 164.89 FID.

• Muon beats Adam in the one-step setting. At 320 epochs, Adam reaches 11.86 FID while Muon achieves 8.71 FID. The authors note that on multi-step diffusion they saw faster convergence with Muon but no final improvement, suggesting the benefit is amplified in the single-step setting because the stop-gradient target depends on network quality in early epochs.

• Perceptual loss gives a large boost. Adding VGG-based LPIPS improves FID from 9.56 to 5.62, and a ConvNeXt-V2 variant further improves it to 3.53 — roughly a 6-point improvement overall.

• Pre-conditioners hurt here. Compared with pMF's plain x-prediction (3.53 FID), linear (14.43), EDM-style (13.81), and sCM-style (34.61) pre-conditioners all perform worse, because any non-zero skip coefficient moves the network away from the x-space.

• Restricting the time sampler breaks the method. Using only r = t gives 194.53 FID, only r = 0 gives 389.28 FID, and using only those two lines gives 106.59 FID, versus 3.53 FID for the full 0 ≤ r ≤ t sampling used by pMF.

• Aggressive patch sizes scale to high resolution cheaply. Holding sequence length at 16², pMF achieves 3.53 FID at 256×256 (patch dim 768), 4.06 FID at 512×512 (patch dim 3072), and 4.58 FID at 1024×1024 (patch dim 12288), all at hidden dimension 768 and B-scale.

• Scaling helps along both axes. Going from B/16 (16 depth, 768 width, 119M params, 34 Gflops) to L/16 (32, 1024, 411M, 117 Gflops) to H/16 (48, 1280, 956M, 271 Gflops) improves FID from 3.53 to 2.85 to 2.57 at 160 epochs, and from 3.12 to 2.52 to 2.29 at 320 epochs.

• System-level results. Final pMF models: pMF-B/16 at 118M params, 33 Gflops, 3.12 FID, 254.6 IS; pMF-L/16 at 410M, 117 Gflops, 2.52 FID, 262.6 IS; pMF-H/16 at 956M, 271 Gflops, 2.22 FID, 268.8 IS. The 2.22 figure is reported at 360 epochs.

• Comparison context. The only other method in the one-step, latent-free diffusion/flow category reported is EPG-L/16 at 8.82 FID (with self-supervised pre-training). Among 1-NFE pixel-space GANs, StyleGAN-XL reaches 2.30 FID but costs 1574 Gflops per forward, 5.8× more than pMF-H/16, while BigGAN-deep reaches 6.95 FID and GigaGAN 3.45 FID.

• Latent decoder overhead matters. The paper argues this has been overlooked: the standard SD-VAE decoder takes 310G and 1230G flops at resolutions 256 and 512, which alone exceeds the cost of the entire pMF generator.

Methodology in Plain English

Flow Matching trains a network to predict a velocity that moves noise into data; generating an image then means solving a differential equation over many steps. MeanFlow instead learns the average velocity over a time interval, which allows one-step generation, and improved MeanFlow (iMF) trains that average-velocity network with a loss defined on the instantaneous velocity using a Jacobian-vector product with a stop-gradient.

The authors' key insight is that a network predicting velocity in pixel space is being asked to model something extremely noisy and high-dimensional. So they keep the iMF loss, but change what the network itself outputs: instead of average velocity, it outputs a denoised-image-like quantity x. A simple algebraic conversion, u = (z_t − x)/t, turns the network's output back into the average velocity needed for the loss. The authors argue this x quantity sits on a lower-dimensional image manifold, making it a more learnable target for a network of finite capacity.

Because the network's raw output is now literally a denoised image at training time — a "what-you-see-is-what-you-get" property — the perceptual loss (LPIPS) can be applied directly to it, which latent-based pipelines normally only use inside the tokenizer. Training samples two time steps r and t from the 0 ≤ r ≤ t region, and the architecture follows the iMF variant of DiT, with ablation models trained from scratch for 160 epochs and FID evaluated on 50,000 generated samples.

Why This Matters

The work suggests that the two standard crutches of modern image generation — many sampling steps and a learned latent space — may not be necessary. If a single network pass can produce raw pixels at competitive quality, generation pipelines become simpler and potentially far cheaper at inference.

Real-world applications:

• Interactive image generation tools where per-sample latency matters, since generation needs only one function evaluation.
• High-resolution content creation, where the paper shows cost can stay roughly constant from 256×256 to 512×512 by using large patches.
• Deployment on constrained hardware, given the removal of a separate latent decoder whose FLOP cost can exceed the generator itself.
• Large-scale image synthesis with simple, end-to-end pipelines — no tokenizer training stage.

Industry relevance: Efficiency-focused deployment, generative model serving, and any product where a diffusion model's multi-step sampling loop is the dominant cost. The comparison to GANs is notable because GANs have historically been the one-step option, and pMF offers Transformer-based scalability alongside fast sampling.

Future Directions

• Extending beyond ImageNet class conditioning to text-to-image and other open-ended generation settings, which the paper does not report.
• Investigating adversarial losses, which the authors explicitly state they do not study, even though latent pipelines typically pair them with perceptual loss in the tokenizer.
• Pushing higher resolutions further, where the paper already reports 1024×1024 at 4.58 FID with patch dimensionality of 12288 and notes the underlying dimensionality of x does not grow proportionally.
• Understanding why pre-conditioners fail so badly here, and whether better parameterizations can be designed for the very high-dimensional regime the paper targets.

Target Audience

Generative-modeling researchers, particularly those working on diffusion and flow matching, few-/one-step sampling, and consistency and MeanFlow-style formulations. It is also relevant to computer-vision graduate students studying pixel-space Transformers and the manifold hypothesis, and to ML engineers who care about inference efficiency and want to understand the trade-offs between latent and pixel-space generation pipelines.

Authors’ abstract

Modern diffusion/flow-based models for image generation typically exhibit two core characteristics: (i) using multi-step sampling, and (ii) operating in a latent space. Recent advances have made encouraging progress on each aspect individually, paving the way toward one-step diffusion/flow without latents. In this work, we take a further step towards this goal and propose "pixel MeanFlow" (pMF). Our core guideline is to formulate the network output space and the loss space separately. The network target is designed to be on a presumed low-dimensional image manifold (i.e., x-prediction), while the loss is defined via MeanFlow in the velocity space. We introduce a simple transformation between the image manifold and the average velocity field. In experiments, pMF achieves strong results for one-step latent-free generation on ImageNet at 256x256 resolution (2.22 FID) and 512x512 resolution (2.48 FID), filling a key missing piece in this regime. We hope that our study will further advance the boundaries of diffusion/flow-based generative models.