Source-linked AI summary
One-step Latent-free Image Generation with Pixel Mean Flows
Yiyang Lu, Susie Lu, Qiao Sun, Hanhong Zhao, Zhicheng Jiang, Xianbang Wang, Tianhong Li, Zhengyang Geng, Kaiming He
TL;DR
Modern diffusion and flow models often combine multi-step sampling with latent-space operation, leaving one-step latent-free generation as an open direction. pMF addresses this by predicting pixel-space denoised images while training through a MeanFlow velocity loss, reaching 2.22 FID at 256×256 and 2.48 FID at 512×512 on ImageNet.
Problem
Modern diffusion and flow-based image generation commonly uses multi-step sampling and latent spaces, motivating alternatives that combine one-step sampling with pixel-space generation.
Method
pMF predicts denoised images on a presumed low-dimensional pixel manifold while defining its training loss in instantaneous-velocity space through MeanFlow conversions.
Results
2.22 FID at 256×256 and 2.48 FID at 512×512 are achieved on ImageNet for one-step latent-free generation.
Takeaways & Limitations
The study reports that appropriately designed neural networks can learn direct end-to-end mappings from noise to pixels for one-step latent-free generation.
Takeaways & Limitations
The generalized manifold assumption is that the endpoint x(zt,0,t) is approximately on the image manifold, while some parameterizations can deviate from x-space and lie on higher-dimensional manifolds.
Abstract
from arXiv · showhide
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.
1. Introduction
Modern image-generation models commonly use multi-step sampling and latent spaces, while recent work has advanced each alternative separately. pMF combines one-step sampling with pixel-space generation through separate network-output and loss spaces, achieving strong ImageNet results.
- Modern diffusion and flow-based models typically rely on multi-step sampling and latent spaces to make complex generation more tractable.
- Recent advances in few-/one-step sampling and raw pixel-space generation provide ingredients for one-step latent-free image generation.
- Combining one-step sampling with pixel-space modeling is challenging because one network must model trajectories while learning image manifolds without latent tokenizers.
- pMF separates network output from loss space by predicting manifold-oriented denoised images while defining training loss in instantaneous-velocity space.
- pMF directly maps noisy inputs to pixels, enabling a what-you-see-is-what-you-get property and natural use of perceptual loss.
- 2.22 FID at 256×256 and 2.48 FID at 512×512 are achieved by pMF for one-step latent-free ImageNet generation.
2. Related Work
Diffusion and flow methods generate samples by following probability flows, commonly through multi-step numerical solving and latent representations. Related work separately explores pixel-space modeling and one-step trajectory prediction.
- Diffusion and Flow Matching learn probability flows between distributions and commonly require multiple numerical function evaluations during inference.
- Latent tokenizers reduce dimensionality and support high-resolution generation by emphasizing high-level semantics while relaxing low-level detail constraints.
- Earlier pixel-space diffusion methods generally use U-nets, while recent work investigates pixel-space Transformers and mechanisms for recovering patch-level details.
- JiT predicts denoised images hypothesized to lie on a low-dimensional manifold as an alternative strategy for pixel-space Transformers.
- Consistency Models and related approaches pursue one-step generation by mapping points along an ODE trajectory to its endpoint.
- Other one-step models characterize SDE/ODE trajectories by choosing different prediction quantities and loss formulations, often involving two trajectory time steps.
3. Background
The paper builds pMF from Flow Matching, MeanFlow, and pixel-space x-prediction. Its central design distinguishes the network’s direct prediction space from the regression loss space and converts between image, average-velocity, and instantaneous-velocity representations.
- pMF is built on Flow Matching, MeanFlow, and JiT as its background frameworks.
- The prediction space is the network’s direct output, whereas the loss space is the regression target; differing spaces require a conversion.
- Flow Matching learns a velocity field mapping a prior distribution to the data distribution and trains it with a velocity-space loss.
- The trained Flow Matching velocity approximates the marginal velocity v(z_t,t), and inference solves an ODE from noise toward data.
- JiT uses x-prediction, taking the network’s direct denoised-image output and converting it to velocity for training.
- JiT’s time convention differs from the paper’s, reversing which endpoint corresponds to t = 0 and changing the conversion coefficient.
- MeanFlow learns average velocity u while retaining an instantaneous-velocity v-loss through a MeanFlow identity and Jacobian-vector-product formulation.
4. Pixel MeanFlow
Pixel MeanFlow separates the network’s image-space output from its velocity-space training loss by converting x-predictions into MeanFlow fields. This design targets denoised, manifold-like quantities while retaining one-step modeling and supports direct perceptual supervision in pixel space.
- The Denoised Image Field: pMF directly outputs a generalized denoised-image field x indexed by two time steps, rather than a one-time-step trajectory quantity.The field is designed to resemble clean or mildly blurry images, unlike noisier velocity-space fields.
- The Generalized Manifold Hypothesis: At r = t, the generalized x-field reduces to the denoised-image target used by JiT, which is approximately low-dimensional even when noise makes it blurry.This boundary case identifies the direct image prediction target and motivates the manifold hypothesis.
- The Generalized Manifold Hypothesis: At r = 0, x corresponds to the ODE trajectory endpoint, which follows the data distribution and can therefore be assumed approximately on the image manifold.The endpoint is denoted z0 and is distributed according to pdata for a ground-truth trajectory.
- The Generalized Manifold Hypothesis: For intermediate r, x is not guaranteed to lie on the data manifold, but simulations show denoised-looking fields that are less noisy than velocity quantities.This empirical contrast suggests that x may be easier for a limited-capacity network to model.
- Algorithm: The network outputs x directly, then converts it through u into the compound velocity target V used by the MeanFlow v-loss.The implementation computes u from the network’s x-output and uses automatic differentiation to form the required velocity-space quantity.
- Pixel MeanFlow with Perceptual Loss: Pixel-space x-prediction enables perceptual losses such as LPIPS on the denoised output, optionally restricted to noise levels t ≤ t_thr to avoid excessive blur.The overall objective combines the pMF loss with a weighted perceptual term.
5. Toy Experiments
The toy experiment tests whether x-prediction remains effective when low-dimensional data are embedded in increasingly high-dimensional observation spaces. x-prediction remains reasonably successful, while u-prediction degrades as dimensionality grows.
- Experimental Setup: The experiment projects a 2D Swiss-roll distribution into D-dimensional observations with D ∈ {2, 8, 16, 512}.Both prediction variants use the same 7-layer ReLU MLP with 256 hidden units.
- Results: x-prediction produces reasonably good 1-NFE generations across the tested observation dimensions.The comparison uses the same network backbone for both prediction targets.
- Results: u-prediction degrades rapidly as observation dimension D increases, producing a widening gap relative to x-prediction.The models minimize the same v-loss, yet x-prediction also achieves lower training loss.
6. ImageNet Experiments
ImageNet experiments show that pMF’s x-prediction, MeanFlow time sampling, perceptual loss, Muon optimization, and scaling choices are important for strong one-step pixel-space generation. The method reaches competitive results at 256×256 and 512×512 resolutions.
- Prediction target: At 256×256, x-prediction performs reasonably well, whereas u-prediction fails catastrophically in the high-dimensional pixel space.The patch dimension is 768 at 256×256, compared with 48 at 64×64, where both targets work well.
- Optimizer: Muon converges faster than Adam and substantially improves FID in pMF training.Preliminary multi-step diffusion experiments showed faster convergence without a final improvement, indicating a setting-specific benefit.
- Perceptual loss: Perceptual loss improves FID from 9.56 to 5.62 with VGG-based LPIPS and to 3.53 with a ConvNeXt-V2 variant.The pixel-space x output makes perceptual loss directly applicable to the generated image.
- Time samplers: Restricting MeanFlow time sampling to r = t, r = 0, or both lines leads to failure, while sampling the full 0 ≤ r ≤ t region succeeds.The comparison suggests that MeanFlow uses relations across points in the (r, t)-coordinate plane.
- High-resolution generation: In high-resolution generation, predicting x maintains performance despite high per-patch dimensionality and provides a FLOP-efficient solution.The reported 512×512 result uses computational resources comparable to the 256×256 counterpart.
- System-level comparisons: pMF achieves 2.22 FID at 256×256 and 2.48 FID at 512×512 for one-step latent-free ImageNet generation.At 256×256, pMF is competitive with leading GAN results at substantially lower compute and narrows the gap with multi-step and/or latent-based methods.
7. Conclusion
The paper frames pMF as a direct noise-to-pixel approach that challenges multi-stage generative designs. Its conclusion presents end-to-end pixel generation as feasible and reports strong 512×512 performance.
- 7. Conclusion: pMF directly maps noise to image pixels, avoiding the multiple steps and stages commonly used to make generation more tractable.This supports an end-to-end formulation of image generation.
- 7. Conclusion: pMF suggests that appropriately designed neural networks can learn complex end-to-end mappings from noise to pixels.
- 7. Conclusion: The paper aims to encourage further exploration of direct, end-to-end generative modeling.
- 7. Conclusion: pMF achieves strong performance for ImageNet 512×512 generation with an aggressive patch size and computational cost similar to 256×256 generation.
A.1. Configurations
The appendix specifies the pMF implementation and training guidance, including iMF-based architecture, classifier-free guidance, EMA selection, perceptual-loss preprocessing, and longer-run sampling changes.
- A.1. Configurations: EMA maintains several decay rates and selects the best-performing rate during inference.
- A.1. Configurations: The network is an x-prediction model whose output is converted into the average velocity u.
- A.1. Configurations: Classifier-free guidance combines conditional and unconditional instantaneous velocities into a guided target, following iMF’s implementation.The network is conditioned on a classifier-free guidance scale interval during training.
- A.1. Configurations: The training procedure computes dudt with a JVP and forms a compound velocity V = u + (t - r) · stopgrad(dudt).
- A.1. Configurations: Longer training doubles the noise scale, samples (t, r) uniformly with 10% probability, and lowers t_thr to 0.6.
A.2. Visualization of the generalized denoised images
The visualization tracks an ODE trajectory from t = 1 to t = 0 and computes the average velocity and generalized denoised-image fields for different time pairs.
- A.2. Visualization of the generalized denoised images: The visualization uses a pretrained JiT-H/16 model to obtain instantaneous velocity and numerically solve the ODE trajectory.
- A.2. Visualization of the generalized denoised images: The fields u and x are computed for different (r, t) pairs from the simulated trajectory using Eq. (5) and Eq. (8).
B. Visualizations
The appendix presents uncurated one-function-evaluation, class-conditional ImageNet samples from pMF-H/16 at 256×256 resolution, together with their evaluation setting.
- B. Visualizations: The visualizations show uncurated 1-NFE pixel class-conditional samples from pMF-H/16 on ImageNet 256×256.Figures 3–6 include samples for listed ImageNet classes.
- B. Visualizations: Under CFG scale ω = 7.0 and interval [0.1, 0.7], the evaluation setting obtains 2.74 FID and 290.0 IS.
- B. Visualizations: The displayed classes include water ouzel, common iguana, and ocean liner.