MeanFlow Transformers with Representation Autoencoders

TL;DR

This work tackles the inefficiency and instability of MeanFlow (MF) training in high-dimensional latent spaces by introducing MF-RAE, which operates in the Representation Autoencoder (RAE) latent space with a lightweight DiT$^{\text{DH}}$-based decoder. The approach stabilizes training via Consistency Mid-Training (CMT), accelerates convergence through distillation from a pre-trained flow matching teacher (MFD), and optionally bootstraps with a one-point velocity estimator (MFT), while replacing Jacobian-vector products with finite-difference approximations. It yields state-of-the-art 1-step FID on ImageNet-256 (2.03) and competitive results on ImageNet-512 (3.23) with significantly lower GFLOPS and training cost, all without guidance. The method offers a robust, hyperparameter-light recipe for efficient few-step generation in high-dimensional latent spaces, with strong practical impact for scalable, fast diffusion-based generation.

Abstract

MeanFlow (MF) is a diffusion-motivated generative model that enables efficient few-step generation by learning long jumps directly from noise to data. In practice, it is often used as a latent MF by leveraging the pre-trained Stable Diffusion variational autoencoder (SD-VAE) for high-dimensional data modeling. However, MF training remains computationally demanding and is often unstable. During inference, the SD-VAE decoder dominates the generation cost, and MF depends on complex guidance hyperparameters for class-conditional generation. In this work, we develop an efficient training and sampling scheme for MF in the latent space of a Representation Autoencoder (RAE), where a pre-trained vision encoder (e.g., DINO) provides semantically rich latents paired with a lightweight decoder. We observe that naive MF training in the RAE latent space suffers from severe gradient explosion. To stabilize and accelerate training, we adopt Consistency Mid-Training for trajectory-aware initialization and use a two-stage scheme: distillation from a pre-trained flow matching teacher to speed convergence and reduce variance, followed by an optional bootstrapping stage with a one-point velocity estimator to further reduce deviation from the oracle mean flow. This design removes the need for guidance, simplifies training configurations, and reduces computation in both training and sampling. Empirically, our method achieves a 1-step FID of 2.03, outperforming vanilla MF's 3.43, while reducing sampling GFLOPS by 38% and total training cost by 83% on ImageNet 256. We further scale our approach to ImageNet 512, achieving a competitive 1-step FID of 3.23 with the lowest GFLOPS among all baselines. Code is available at https://github.com/sony/mf-rae.

Figures (8)

• Figure 1: Overview of our method’s advantages. On ImageNet 256, vanilla MF (a) employs the slow SD-VAE decoder, which accounts for 73% of the total generation cost and thus bottlenecks the few-step generation speed. In contrast, our MF-RAE (b) leverages a higher-dimensional RAE latent space with semantically rich features and an efficient decoder. The DiT$^\text{DH}$ architecture is adopted to effectively process the high-dimensional latent space. As a result, while the converged 1-step FID of vanilla MF is 3.43 after more than 600 H100 GPU-days, our MF-RAE achieves a superior FID of 2.03 in only 100 H100 GPU-days. Additionally, our total generation cost is reduced by 38% compared to vanilla MF in terms of GFLOPS, even with the same 1-step generation setting.
• Figure 2: ImageNet 256 MF-RAE 1-step samples on class 437: beacon, lighthouse, beacon light, pharos.
• Figure 3: ImageNet 256 MF-RAE 1-step samples on classes 288 and 290: leopard and snow leopard.
• Figure 4: ImageNet 256 MF-RAE 1-step samples on classes 13, 14, 94, and 134: snowbird, indigo bird, hummingbird, and crane bird.
• Figure 5: ImageNet 256 MF-RAE 1-step samples for various dogs.

Theorems & Definitions (3)

• Proposition 3.1
• Proposition A.1
• proof