MeanFlowNFT: Bringing Forward-Process RL to Average-Velocity Generators

AI summary

Overview

Research area: Generative AI / computer vision — reinforcement learning (RL) fine-tuning of few-step image and video generation models, specifically average-velocity (MeanFlow) generative models.

Technical level: Advanced. The paper assumes familiarity with flow matching, diffusion reinforcement learning, and ordinary differential equation (ODE) samplers. The core idea, however, can be explained without that background.

Scope: The paper introduces MeanFlowNFT, the first forward-process RL framework for MeanFlow generators, and demonstrates it on text-to-image (SD3.5-M) and text-to-video (Wan2.1 1.3B) generation.

What This Paper Is About

Diffusion and flow models normally generate samples by integrating an instantaneous velocity over many sequential steps, which is slow. MeanFlow models remove this bottleneck by predicting the average velocity over a time interval, so only one or a few steps are needed. Separately, RL has become the standard way to align such models with human preferences, and a method called DiffusionNFT achieves this efficiently by training directly on the forward noising process instead of unrolling the reverse sampler. The problem is that DiffusionNFT optimizes instantaneous velocities, while MeanFlow samples with average velocities. This paper's goal is to close that gap so that MeanFlow generators can be RL-fine-tuned efficiently while keeping their few-step sampling speed.

Key Contributions

1. The first forward-process RL framework for MeanFlow generators. MeanFlowNFT applies a DiffusionNFT-style objective to an "induced instantaneous-velocity predictor" derived from the average-velocity network via the MeanFlow identity. Training stays likelihood-free and on the forward noising process, and inference keeps MeanFlow's native few-step sampler unchanged.

2. Theoretical guarantees. In an idealized setting, the authors prove a closed-form pointwise optimum for the induced predictor (Proposition 3.1), show that at guidance strength β = 2α it recovers the improved policy's marginal instantaneous velocity (Corollary 3.2), and prove via a MeanFlow consistency lemma (Lemma 3.3) that this improvement provably transfers to the average-velocity network (Theorem 3.4), yielding J(π_θ*) = J(π+) > J(π^old).

3. Practical implementation design. The total-derivative term is approximated by a central finite difference in t rather than a Jacobian-vector product (which is expensive and not fully compatible with FSDP), the same total derivative is shared by the trainable and reference predictors to prevent training collapse, and the already-available forward-process conditional velocity replaces the network's instantaneous velocity as the direction for displacement.

4. Comprehensive image and video experiments. MeanFlowNFT consistently improves MeanFlow baselines, outperforms prior few-step RL methods on most metrics, and can surpass multi-step RL-tuned diffusion while using far fewer sampling steps.

Main Findings

• Best on 6 of 8 metrics on SD3.5-M among few-step models. MeanFlowNFT achieves the best results on 6 of the 8 reported image metrics (ImageReward, CLIPScore, Aesthetic, PickScore, HPSv2, HPSv3, GenEval2, OCR) among few-step models, with ImageReward 1.4504, CLIPScore 0.2967, Aesthetic 5.9275, PickScore 23.5019, HPSv2 0.3269, HPSv3 13.8826, GenEval2 0.2375, and OCR 0.6534.

• Beats prior few-step distillation and RL baselines. It outperforms the few-step distillation baselines DMD, CDM, and AnyFlow, and also beats the recent few-step RL methods R_dm and RTDMD on most metrics — for example, OCR 0.65 vs. 0.30 against RTDMD.

• 4-step RL can match or exceed 40-step RL. With only 4 sampling steps, MeanFlowNFT matches or exceeds the 40-step forward-process RL method DiffusionNFT on several reward metrics (ImageReward 1.45 vs. 1.41, CLIPScore 0.297 vs. 0.289) at 10 times fewer function evaluations. DiffusionNFT retains the advantage on other metrics, such as GenEval2 (0.2613 vs. 0.2375) and OCR (0.9098 vs. 0.6534).

• Video result: 4-step MeanFlowNFT surpasses 50-step LongCat-Video RL. On Wan2.1, 4-step MeanFlowNFT reaches a VBench score of 84.33, surpassing 50-step LongCat-Video RL (82.57).

• Naively applying DiffusionNFT to few-step generators fails. AnyFlow + DiffusionNFT, DMD + DiffusionNFT, and CDM + DiffusionNFT fall far short of MeanFlowNFT and train very unstably, collapsing early; CDM + DiffusionNFT diverges within 400 steps. The authors attribute this to those networks not predicting the instantaneous velocity as a posterior mean, so DiffusionNFT's linearity-based improvement guarantee does not apply.

• Fewer reward-hacking artifacts qualitatively. The paper reports that MeanFlowNFT produces more faithful samples with fewer reward-hacking artifacts, such as over-saturated colors and implausible object scales, than DiffusionNFT and RTDMD.

• Theoretical inheritance of strict policy improvement. MeanFlowNFT is proven to inherit DiffusionNFT's strict policy-improvement guarantee under the stated idealized conditions.

Methodology in Plain English

The starting point is the MeanFlow identity, which links average velocity and instantaneous velocity: the instantaneous velocity equals the average velocity plus (t − s) times the total derivative of the average velocity.

The authors use this identity to construct a second, "induced" predictor that outputs an instantaneous velocity, built from the MeanFlow network's average-velocity prediction. Optimization happens in this induced instantaneous-velocity space, while sampling still uses the average-velocity network. This decoupling is the central trick: the reward objective is applied to a quantity where it is well-defined, but the deployed generator keeps its fast few-step sampler.

The training objective mirrors DiffusionNFT: roll out samples with a frozen reference policy, score them with a reward in [0, 1], then split predictions into implicit "positive" and "negative" versions (linear blends of the reference and trainable predictors). The positive version is fitted to the target conditional velocity weighted by the reward, and the negative version weighted by (1 − reward). The reference policy is an exponential moving average (EMA) of the trainable one.

Three practical simplifications make this workable. First, the total derivative is approximated with a central finite difference in t, displacing the sample along the forward-process conditional velocity, rather than computing an exact Jacobian-vector product. Second, both the trainable and reference predictors use the same finite-difference derivative, so the derivative-difference term cancels exactly — otherwise training collapses. Third, the direction passed into the derivative is the forward-process conditional velocity, which is already available at no extra cost, rather than an extra network evaluation.

Experimental setup: SD3.5-M is used for image generation at 512×512 for training and 1024×1024 for evaluation; Wan2.1 1.3B at 480p with 81 frames for video. MeanFlow policies come from distilling the base model with AnyFlow, and the released AnyFlow checkpoint is used directly for video. The pipeline is CFG-free. Training uses LoRA (rank 32, scaling factor 64) on all linear layers in attention blocks, guidance strength β = 0.1, the DiffusionNFT EMA schedule η_i = min(0.001 i, 0.5), KL regularization weighted 10⁻⁴, and AdamW with a constant learning rate of 3×10⁻⁶. Rollouts use 4 steps. For SD3.5-M, group size 24 with 48 groups per update on 8 NVIDIA H20 GPUs; for Wan2.1, group size 16 with 8 groups per epoch on 32 NVIDIA H20 GPUs. Image rewards are CLIPScore, PickScore, and HPSv2 on the PickScore prompt set; video rewards are HPSv3-general, HPSv3-percentile, and the motion-quality and text-alignment scores of VideoAlign, with training prompts from DanceGRPO.

Why This Matters

Impact on research. This is the first work to bring forward-process RL to average-velocity generators, and it supplies a proof that the improvement carries over from the induced instantaneous predictor to the deployed average-velocity model. That suggests likelihood-free RL is applicable to a broader class of few-step generators than previously assumed, while the negative results for DMD/CDM + DiffusionNFT clarify when DiffusionNFT's guarantee does and does not transfer.

Real-world applications:

• Fast text-to-image generation on consumer hardware, where few-step sampling reduces inference cost dramatically (the paper's 4-step vs. 40-step comparison implies 10× fewer function evaluations).
• Text-to-video generation, where multi-step sampling is especially expensive given the large number of frames involved.
• Preference-aligned creative tools that need to follow human aesthetic or instruction-following objectives without giving up interactive latency.
• Deployment settings with tight compute budgets, where the same compute can be spent on more samples rather than more denoising steps.

Industry relevance. The efficiency claims (few-step generation plus 10× fewer function evaluations) and the open resources (project page, GitHub, Hugging Face) point directly at production inference costs. The reward stack used — CLIPScore, PickScore, HPSv2 for images and HPSv3 plus VideoAlign for video — is the same kind of preference/metric stack industry alignment pipelines already use.

Future Directions

• Extending the guarantees beyond the idealized setting. The policy-improvement result assumes the induced optimum is attained for all intervals s ≤ t; understanding behavior when this does not hold, and how the finite-difference approximation of the total derivative affects the guarantee, is left open.
• Improving metrics where MeanFlowNFT still trails. DiffusionNFT retains advantages on GenEval2 (0.2613 vs. 0.2375) and OCR (0.9098 vs. 0.6534), and RTDMD leads on Aesthetic (6.1290 vs. 5.9275) — closing these gaps is a natural next step.
• Scaling to larger models and more modalities. The experiments cover SD3.5-M and Wan2.1 1.3B; whether the approach holds for larger backbones and other generators (the paper notes it "scales gracefully at test time") is not established here.
• Reducing the remaining training cost. The finite-difference derivative introduces an approximation and extra evaluations, and the GPU budgets reported (8 NVIDIA H20 GPUs for SD3.5-M, 32 for Wan2.1) remain substantial.

Target Audience

Researchers and engineers working on diffusion and flow-based generative models, particularly those interested in RL fine-tuning, few-step distillation, and efficient sampling. It is most useful to readers with some background in flow matching and policy-gradient or DiffusionNFT-style methods, though the core idea — optimize in instantaneous-velocity space, sample in average-velocity space — is accessible to a broader machine learning audience. Practitioners concerned with inference cost for image and video generation will also find the empirical comparisons relevant.

Authors’ abstract

MeanFlow generators achieve fast few-step sampling by predicting average velocities over time intervals, making them attractive for efficient generation. Reinforcement learning (RL) has become a powerful way to align diffusion and flow models with human preferences and task-specific objectives. In particular, DiffusionNFT offers an efficient forward-process RL framework that does not require reverse-process trajectories or likelihood estimation. However, applying such RL methods to MeanFlow remains underexplored. DiffusionNFT optimizes instantaneous velocities, whereas MeanFlow samples with average velocities. To bridge this gap, we introduce MeanFlowNFT. Inspired by the MeanFlow identity, which bridges average and instantaneous velocities, we construct an induced instantaneous-velocity predictor. We apply the DiffusionNFT objective to this predictor, making reward optimization well-defined for MeanFlow. Sampling remains based on the average velocity, preserving MeanFlow's fast few-step generation. We further prove that MeanFlowNFT inherits DiffusionNFT's strict policy-improvement guarantee. Experiments on image and video generation show that MeanFlowNFT consistently improves baselines. Moreover, it outperforms prior state-of-the-art RL-tuned few-step generators on most metrics ($6$ of $8$ on SD3.5-M), and can even surpass multi-step RL-tuned diffusion while using only a few sampling steps. For instance, on Wan 2.1, $4$-step MeanFlowNFT reaches a VBench score of $84.33$, surpassing $50$-step LongCat-Video RL ($82.57$).