Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Yuhao Huang, Shih-Hsin Wang, Andrea L. Bertozzi, Bao Wang

TL;DR

RMFlow addresses the quality gap of 1-NFE MeanFlow by introducing a noise-injection refinement after a coarse 1-NFE transport. It derives a joint training objective that combines Wasserstein-path alignment with likelihood maximization, enabling principled control of the learned distribution via a bound on $\mathbb{E}[\log p_{\theta}(\mathbf{x}_{\text{tgt}})]$ and a negative log-likelihood term $\mathcal{L}_{\rm NLL}$. The approach supports multimodal generation through a conditioning encoder and a small perturbation prior, while employing memory-efficient fine-tuning for scalability. Empirically, RMFlow achieves near-state-of-the-art performance across text-to-image, context-to-molecule, and time-series tasks using only 1-NFE, with computational cost comparable to baseline MeanFlows and potential for further improvements with larger budgets.

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.

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

TL;DR

RMFlow addresses the quality gap of 1-NFE MeanFlow by introducing a noise-injection refinement after a coarse 1-NFE transport. It derives a joint training objective that combines Wasserstein-path alignment with likelihood maximization, enabling principled control of the learned distribution via a bound on and a negative log-likelihood term . The approach supports multimodal generation through a conditioning encoder and a small perturbation prior, while employing memory-efficient fine-tuning for scalability. Empirically, RMFlow achieves near-state-of-the-art performance across text-to-image, context-to-molecule, and time-series tasks using only 1-NFE, with computational cost comparable to baseline MeanFlows and potential for further improvements with larger budgets.

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.
Paper Structure (33 sections, 3 theorems, 28 equations, 5 figures, 14 tables)

This paper contains 33 sections, 3 theorems, 28 equations, 5 figures, 14 tables.

Table of Contents

  1. Introduction
  2. Our Contributions
  3. Additional Related Works
  4. Organization
  5. Background
  6. The Design of RMFlow
  7. Motivation
  8. Noise Injection Refinement
  9. Multimodality
  10. The Training of RMFlow
  11. Likelihood Maximization
  12. Joint Training Objective
  13. Memory-Efficient Fine-Tuning
  14. Numerical Experiments
  15. Synthetic Tasks
  16. ...and 18 more sections

Key Result

Theorem 3.1

[boffi2025build] There exists a constant $M > 0$ such that: where $M$ is a constant, $W^2_2(p_{\rm tgt}, p_\theta)$ denotes the Wasserstein distance between $p_{\rm tgt}$ and $p_\theta$, and $\Pi(p_{\rm tgt}, p_\theta)$ is the set of all joint distributions with marginals $p_{\rm tgt}$ and $p_\theta$.

Figures (5)

  • Figure 1: Contrasting MeanFlow with RMFlow for mixture Gaussian sampling; see Section \ref{['subsec:density-estimation']} for experimental details and more results.
  • Figure 2: Contrasting MeanFlow with RMFlow, under the same context, for QM9 molecule generation.
  • Figure 3: Schematic of our proposed RMFlow: it first applies 1-NFE MeanFlow transport, then refines the result by a subsequent noise-injection step; see Section \ref{['sec:model-design']}. The average velocity $\hat{{\bm{u}}}_{0,1}({\bm{x}}_0;\theta)$ of RMFlow is trained by incorporating the maximum likelihood objective into the MeanFlow framework, as in equation \ref{['eq:RMFlow-objective']}.
  • Figure 4: A few randomly selected RMFlow-generated molecules, together with the corresponding contexts.
  • Figure 5: COCO dataset samples generated using 1-NFE RMFlow conditioned on different input prompts.

Theorems & Definitions (6)

  • Theorem 3.1
  • Remark 1
  • Theorem 4.1
  • Remark 2
  • Theorem A.1
  • proof : proof of Theorem \ref{['thm:KL']}