RMFlow: Refined Mean Flow by a Noise-Injection Step for Multimodal Generation

AI summary

Overview

Research area: Generative modeling, specifically flow matching and mean flow methods for efficient high-fidelity sample generation (images, molecules, time series).

Technical level: Advanced. The paper builds on flow matching theory, Wasserstein distance, KL divergence, and evidence lower bound (ELBO) arguments, and assumes familiarity with diffusion models, ordinary differential equations, and neural network training objectives.

Scope: A single paper proposing RMFlow, a modification of the MeanFlow framework that adds one noise-injection step to improve 1-NFE (one function evaluation) generation across image, molecule, and time-series benchmarks.

What This Paper Is About

MeanFlow models can generate high-quality samples with far fewer neural network evaluations than standard flow matching, but collapsing generation all the way down to a single function evaluation (1-NFE) causes a large drop in quality. The authors show this concretely: 1-NFE MeanFlow produces distorted mixture-Gaussian samples and frequently generates fragmented, invalid molecules on QM9. RMFlow addresses this by pairing the coarse 1-NFE MeanFlow transport with a single, tailored noise-injection refinement step, and by training with a loss that both minimizes the Wasserstein distance between probability paths and maximizes sample likelihood.

Key Contributions

1. A 1-NFE multimodal generative model (RMFlow) that combines 1-NFE MeanFlow transport with a tailored noise-injection refinement step, and that integrates a guidance encoder for conditioned generation.
2. A theoretically principled training objective that jointly balances minimizing the Wasserstein distance between the target and learned probability paths (the conditional mean flow matching loss) with maximizing the likelihood of the learned target distribution, plus an optional guidance regularization term.
3. A proof (Theorem 4.1) that the noise-injection step's negative log-likelihood loss provides a lower bound on the expected log-likelihood of the target distribution, which in turn controls the KL divergence between the target and learned distributions.
4. Near state-of-the-art empirical results on benchmark text-to-image, context-to-molecule, and time-series generation, all at 1-NFE and at computational cost comparable to baseline MeanFlows. The authors state this is the first work to improve MeanFlows by introducing a noise-injection refinement for 1-NFE generation.

Main Findings

• 1D mixture Gaussian sampling: 1-NFE RMFlow reaches TV 0.7567 and KL 0.2332, versus 1-NFE MeanFlow at TV 1.4422 and KL 0.8074 and 8-NFE MeanFlow at TV 0.7977 and KL 0.4074. RMFlow beats 1-NFE and 8-NFE MeanFlow but is slightly worse than 32-NFE MeanFlow (TV 0.6737, KL 0.1017).
• 2D checkerboard sampling: 1-NFE RMFlow reaches TV 0.173 and KL 0.163, versus 1-NFE MeanFlow at TV 0.238 and KL 0.311, 8-NFE MeanFlow at TV 0.167 and KL 0.139, and 32-NFE MeanFlow at TV 0.155 and KL 0.118. RMFlow substantially outperforms 1-NFE MeanFlow and closes the gap to multi-NFE MeanFlow.
• QM9 molecule generation: RMFlow with contexts and reinforcement-learning policy feedback (RLPF) achieves 98.9 percent atomic stability and 93.5 percent molecule stability at 1 NFE. RMFlow with contexts alone gives 98.9 percent and 93.2 percent; RMFlow without contexts gives 98.8 percent and 90.1 percent. MeanFlow with contexts gives 98.4 percent and 84.3 percent; without contexts 98.2 percent and 79.3 percent. For reference, GeoBFN reports 99.0 percent and 93.9 percent and E-DM + RLPF reports 99.1 percent and 93.4 percent, both at NFE much greater than 1. The dataset itself is listed at 99 percent and 95.2 percent.
• Time series (Lorenz and FitzHugh-Nagumo): RMFlow at 1 NFE yields TV distances of 0.0332 (Lorenz, without event) and 0.0956 (Lorenz, with event), and 0.0289 and 0.1543 for FitzHugh-Nagumo. MeanFlow at 1 NFE gives 0.0469, 0.1250, 0.0398, and 0.2268 respectively; at 32 NFE it gives 0.0351, 0.0991, 0.0302, and 0.1723. KL divergences follow the same ordering, with RMFlow at 1 NFE reaching 0.0059 and 0.2866 (Lorenz) and 0.0287 and 0.2499 (FitzHugh-Nagumo).
• Text-to-image on COCO: RMFlow attains COCO FID-30K of 18.91 at 1 NFE with a 620M-parameter model, compared with 27.31 for MeanFlow at 1 NFE. RMFlow's FID is reported as comparable to the best single-step generators listed, such as InstaFlow (13.10), UFOGen (12.78), and StyleGAN-T (13.90), and it is teacher-free and discriminator-free, unlike those baselines.
• Training-resource gap: The authors note their models were trained on RTX 3090/4090 GPUs with 24 GB memory using mixed-precision bf16, while most state-of-the-art models in their comparison table were trained on multiple A100 80 GB GPUs with fp16, suggesting further gains with larger compute budgets.
• Guidance regularization behavior: The authors report that the guidance regularization term (Term III in the loss) can become very large and cause substantial performance degradation.

Methodology in Plain English

The starting point is MeanFlow, which learns an averaged velocity field that transports a noise sample directly to data in one or a few steps. The authors split generation into two stages.

In stage one, a 1-NFE MeanFlow moves the prior sample to an intermediate state that is the data corrupted by a small amount of noise, with a noise level smaller than the small floor value used during training. In stage two, a single extra noise injection adds back the remaining noise, using a standard deviation chosen so that the final sample's noise level matches the training-time floor.

The key insight is that this final noise injection turns the generator into a Gaussian conditional distribution: given the prior sample, the output is normally distributed around the flow's prediction with a known variance. That means a simple squared-error term between the generated sample and the noisy target acts as a negative log-likelihood, which the authors prove gives a lower bound on the expected log-likelihood and hence bounds the KL divergence to the target distribution.

RMFlow is trained with three terms summed together: the original MeanFlow objective (Wasserstein control), the likelihood term scaled by a hyperparameter, and a regularization term on the guidance embedding (used only for conditioned generation). For large tasks, the authors first train with the MeanFlow loss and then fine-tune with the full RMFlow loss using parameter-efficient fine-tuning; for molecule generation they additionally add a policy-gradient signal that uses molecule stability as a reward. Conditioning works by embedding the context (text, molecular properties, or event constraints) into the prior sample itself.

Why This Matters

The work shows that a single extra noise-injection step, applied in parallel with the flow evaluation, can recover much of the quality lost when collapsing generative sampling to one neural network evaluation. Because the extra noise is nearly free computationally, this is a practical route to high-quality one-step generation without distillation, teachers, or discriminators.

Real-world applications implied by the tasks studied:

• Fast text-to-image generation for interactive creative tools and low-latency image synthesis on consumer GPUs, where multi-step diffusion is too slow.
• Drug and materials discovery, where conditioned molecule generation must produce chemically valid structures under property constraints.
• Scientific forecasting and rare-event sampling, such as generating plausible trajectories of systems like Lorenz or FitzHugh-Nagumo dynamics under event constraints, relevant to climate and extreme-event modeling.
• Edge and on-device deployment of generative models where memory and compute are tightly constrained.

Industry relevance: the method requires no pre-trained teacher model and no discriminator, unlike distillation- or GAN-based one-step approaches, so it can be trained and deployed with a simpler pipeline. Its stated training on 24 GB consumer GPUs makes high-fidelity 1-NFE generation accessible to organizations without large compute clusters.

Future Directions

• Extending RMFlow to multiple mean flow transport steps, applying a noise-injection step after each transport step. The authors note this requires designing a corresponding loss function to preserve likelihood maximization and presents additional challenges relative to the current formulation.
• Adaptive selection of the noise-injection magnitude. The current method fixes the value of the square root term at every step, which the authors describe as possibly suboptimal. They propose making it learnable or following a dynamic schedule.
• Scaling training compute. Since the reported experiments used consumer GPUs, the authors suggest RMFlow has potential for further improvement under larger computational budgets.
• Refining the guidance regularization term, which the authors observed can become very large and hurt performance.

Target Audience

Researchers and graduate students working on generative modeling, particularly those focused on efficient sampling, one-step generation, flow matching, and consistency or mean flow methods. It is also relevant to practitioners in computational chemistry, scientific machine learning, and computer vision who need fast conditioned generation. The paper is written at an advanced level: readers should already be comfortable with flow matching objectives, Wasserstein and KL divergences, and the structure of diffusion-style training losses.

Authors’ abstract

Mean flow (MeanFlow) enables efficient, high-fidelity image generation, yet its single-function evaluation (1-NFE) generation often cannot yield compelling results. We address this issue by introducing RMFlow, an efficient multimodal generative model that integrates a coarse 1-NFE MeanFlow transport with a subsequent tailored noise-injection refinement step. RMFlow approximates the average velocity of the flow path using a neural network trained with a new loss function that balances minimizing the Wasserstein distance between probability paths and maximizing sample likelihood. RMFlow achieves near state-of-the-art results on text-to-image, context-to-molecule, and time-series generation using only 1-NFE, at a computational cost comparable to the baseline MeanFlows.