Flow Straighter and Faster: Efficient One-Step Generative Modeling via MeanFlow on Rectified Trajectories

TL;DR

This work tackles the inefficiency of sampling in flow-based generative models by merging trajectory rectification with mean-velocity learning. It introduces Re-MeanFlow, which trains a MeanFlow model on rectified trajectories produced by a single reflow, augmented with a simple distance truncation to reduce extreme curvature. Across ImageNet variants, Re-MeanFlow delivers superior one-step sample quality (FID) and substantially improved training efficiency compared to prior one-step methods. The approach shifts most compute to the inference-inspired reflow stage, improving accessibility and scalability of high-quality one-step generation. The method is architecture-agnostic and paves the way for broader applications, including potential real-data integration and extensions beyond image synthesis.

Abstract

Flow-based generative models have recently demonstrated strong performance, yet sampling typically relies on expensive numerical integration of ordinary differential equations (ODEs). Rectified Flow enables one-step sampling by learning nearly straight probability paths, but achieving such straightness requires multiple computationally intensive reflow iterations. MeanFlow achieves one-step generation by directly modeling the average velocity over time; however, when trained on highly curved flows, it suffers from slow convergence and noisy supervision. To address these limitations, we propose Rectified MeanFlow, a framework that models the mean velocity field along the rectified trajectory using only a single reflow step. This eliminates the need for perfectly straightened trajectories while enabling efficient training. Furthermore, we introduce a simple yet effective truncation heuristic that aims to reduce residual curvature and further improve performance. Extensive experiments on ImageNet at 64, 256, and 512 resolutions show that Re-MeanFlow consistently outperforms prior one-step flow distillation and Rectified Flow methods in both sample quality and training efficiency. Code is available at https://github.com/Xinxi-Zhang/Re-MeanFlow.

Figures (12)

• Figure 1: Re-Meanflow (Ours) leverages the synergy of trajectory rectification and mean-velocity modeling, achieving the best compute–quality trade-off—reaching strong FID. This synergy yields efficiency and quality that neither rectification nor mean-velocity modeling can achieve alone. All methods are initialized with pretrained EDM2-S karras2024analyzing.
• Figure 2: Why trajectory rectification and mean-velocity modeling reinforce each other.(a) A 1-rectified flow () still follows highly curved trajectories, requiring many ODE steps. Applying two rounds of rectification () straightens the paths and reduces NFEs, but one-step sampling () remains unreliable unless trajectories are nearly straight. (b) MeanFlow estimates the average velocity $u(\mathbf{z}_t, r, t)$ over all intervals $(r, t)$, which in principle avoids the need for straight paths. However, when the underlying velocity field is highly curved, the induced averages become complex and hard to learn. (c) Training MeanFlow on trajectories from a 2-rectified flow yields a significantly smoother average-velocity field, making estimation easier and enabling faster convergence and high-quality one-step generation.
• Figure 3: A Class-imbalanced 2D Toy Example. We consider a controlled 2D setup where a flow model transports a balanced two-component Gaussian mixture on the left to an imbalanced mixture on the right (weights $0.4$ for the upper mode and $0.6$ for the lower mode). Panels (a-c):(a) Linear interpolation of independently sampled couplings $p_{\mathbf x} \times p_{\mathbf z}$, which serve as the training signal for the first velocity model. (b) The resulting 1-rectified flow learned from these independent couplings; the learned velocity field remains noticeably curved. (c) Using the velocity field from (b), we generate a new set of couplings and train a second velocity model on their linear interpolations, yielding the 2-rectified flow. Panel (d): Due to imperfect straightening, one-step Euler sampling on the 2-rectified flow still yields noticeable outliers. Panel (e): MeanFlow trained directly on independent couplings fails to converge within the training budget because high-variance conditional velocities destabilize learning. Panel (f): Re-Meanflow combines trajectory rectification with MeanFlow, eliminating most outliers and achieving more accurate one-step generation.
• Figure 4: Distance-Error Correlation. Histogram of data-noise $\ell_2$ distances on ImageNet $512^2$, colored by angular error (computed as $\arccos ( \tfrac{\langle \mathbf z - \mathbf x,\; u_\theta(\mathbf z,0,1)\rangle} {\|\mathbf z - \mathbf x\|_2\,\|u_\theta(\mathbf z,0,1)\|_2} )$, the angle between the predicted and ground-truth velocities). Errors are computed using a MeanFlow model trained on untruncated 1-rectified couplings (config (d) in Tab. \ref{['tab:uncond-fid-only']}). A clear high-distance high-error tail (90th percentile marked) motivates our Distance Truncation.
• Figure 5: Qualitative results for Re-Meanflow (NFE=1) on ImageNet$64^2$ (Left), $256^2$ (Middle), and $512^2$ (Right).