MeanFlowNFT: Bringing Forward-Process RL to Average-Velocity Generators
By Yushi Huang, Xiangxin Zhou, Jun Zhang, Liefeng Bo, Tianyu Pang
"Bridges DiffusionNFT's forward-process RL to MeanFlow generators via an induced instantaneous-velocity predictor, enabling RL alignment while preserving few-step sampling."
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$).
Technical Analysis & Implementation
MeanFlowNFT: Bringing Forward-Process RL to Average-Velocity Generators§
Core Methodology§
MeanFlow generators model the probability flow ODE: $dx_t = v_t(x_t) dt$, where $v_t$ is the instantaneous velocity. In practice, MeanFlow predicts the average velocity over a time interval $[t, s]$:
$$ \bar{v}_{t \to s}(x_t) = \mathbb{E}_{x_s \sim p_{s|t}(\cdot | x_t)} \left[ \frac{x_s - x_t}{s - t} \right]. $$
Sampling iterates $x_{t+\Delta t} = x_t + \bar{v}_{t \to t+\Delta t}(x_t) \Delta t$, enabling few-step generation.
DiffusionNFT optimizes a reward $R(x_1)$ for predictions of instantaneous velocity at single timepoints. To adapt this to MeanFlow, the authors leverage the MeanFlow identity:
$$ \bar{v}_{t \to s}(x_t) = \mathbb{E}_{x_s \sim p_{s|t}(\cdot | x_t)} \left[ \frac{x_s - x_t}{s - t} \right] = \int_0^1 v_{t + u(s-t)}(x_{t+u(s-t)}) \, du, $$
which relates average and instantaneous velocities. However, evaluating the integral is intractable. Instead, they introduce an induced instantaneous-velocity predictor $\tilde{v}_\theta(x_t, \tau)$ that approximates the instantaneous velocity at time $\tau \in [t, s]$ given state $x_t$. This predictor is trained via a denoising objective on the same trajectories used for MeanFlow training.
Training Objective§
The induced predictor is trained with:
$$ \mathcal{L}_{\text{ind}} = \mathbb{E}_{t, s, x_t, x_s, \tau \sim \mathcal{U}[t,s]} \left[ \| \tilde{v}_\theta(x_t, \tau) - v_{\tau}(x_\tau) \|^2 \right], $$
where $x_\tau$ is interpolated linearly between $x_t$ and $x_s$ (or sampled from the true forward process). Once $\tilde{v}_\theta$ is trained, the DiffuisonNFT objective can be applied:
$$ \mathcal{L}_{\text{DRL}} = - \mathbb{E}_{t \sim \mathcal{U}[0,1], \, x_t \sim p_t, \, \tau \sim \mathcal{U}[t, t+\Delta t]} \left[ R(\text{sg}(x_\tau + (1-\tau) \tilde{v}_\theta(x_t, \tau))) \right], $$
where sg denotes stop-gradient to avoid degenerate solutions. Notably, sampling during inference still uses the original MeanFlow predictor, preserving the few-step generation advantage.
Theoretical Guarantee§
The paper proves that MeanFlowNFT inherits DiffusionNFT's strict policy improvement: if $\tilde{v}_\theta$ is trained sufficiently well, the induced policy improves monotonically in expectation over the reward.
Implementation Details§
A simplified training loop using PyTorch:
import torch
import torch.nn as nn
class InducedVelocityPredictor(nn.Module):
def __init__(self, backbone):
super().__init__()
self.backbone = backbone # e.g., U-Net
def forward(self, x_t, tau):
# Concatenate tau as additional channel or embed
return self.backbone(torch.cat([x_t, tau.expand_as(x_t[:,:1])], dim=1))
# Given a batch of (x_t, x_s) from diffusion process
# Sample tau uniformly in [t, s]
tau = torch.rand_like(t) * (s - t) + t
# Interpolate x_tau (simplified linear interpolation):
x_tau = x_t + (x_s - x_t) * (tau - t).unsqueeze(1) / (s - t).unsqueeze(1)
# Induced predictor loss
v_pred = induced_predictor(x_t, tau)
v_true = some_instantaneous_velocity_model(x_tau, tau) # or use score function
loss_ind = nn.MSELoss()(v_pred, v_true)
# RL fine-tuning (DiffusionNFT style)
with torch.no_grad():
x_hat = x_tau + (1 - tau) * v_pred # approximate x_1
reward = reward_model(x_hat)
loss_rl = -reward.mean()Experiments§
MeanFlowNFT was evaluated on image generation (SD3.5-M) and video generation (Wan 2.1). It outperformed prior RL-tuned few-step generators (e.g., LADD, LORA-finetuned) on 6 out of 8 metrics for SD3.5-M. For Wan 2.1, 4-step MeanFlowNFT achieved VBench 84.33, surpassing 50-step LongCat-Video RL (82.57). The induced predictor adds minimal overhead and does not affect sampling speed.