RMFlow: Refined Mean Flow by a Noise-Injection Step for Multimodal Generation
Authors: Yuhao Huang
Shih-Hsin Wang
Andrea L. Bertozzi
Bao Wang
Venue: ICLR 2026 Poster (2026)
Keywords: Mean Flow
Flow Matching
Noise-injection
Likelihood Maximization
Multimodal Generation
Primary area: generative models
TL;DR: We RMFlow, an efficient multimodal generative model that integrates a coarse 1-NFE MeanFlow transport with a subsequent tailored noise-injection refinement step.

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.

Published as a conference paper at ICLR 2026
RMFLOW: REFINED MEAN FLOW BY A NOISE- INJECTION STEP FOR MULTIMODAL GENERATION
Yuhao Huang', Shih-Hsin Wang', Andrea L. Bertozzi", Bao Wang'
DDparrmmett of Mathematics and Scientific Computing and Imaging (SCT) Institute University of Utah, Salt Lake City, UT. 84102, USA *Department of Mathematics, UCLA, Los Angeles, CA, 90095, USA
ABSTRACT
Mean flow (MeanFlow) enables efficient, high-fidelity image generation, yet its single-function evaluation (I-NFE) generation often cannot yield compelling re- sults. We address this issue by introducing RMFlow, an efficient multimodal gen- erative model that integrates a coarse 1-NFE McanFlow transport with subse- quent tailored noise-injection refinement step. RMFlow approximates the average velocity of the flow path using neural network trained with that balances minimizing the Wasserstein distance between probability paths and maximizing sample likelihood. RMFlow achieves near state-of-the-ee results on text-to-image, context-to-moledd and time-series generation using only 1-NFE, a computational cost comparable the baseline MeanFlows.
loss function
INTRODUCTION
Flow matching (FM), closely related to dif- fusion models (DMs) (ickstDic al.]
Sample
Data
2015 Ho et al. 2020 Song et al. 2020). has emerged as a flexible framework for generative modeling. offering principled way to learn transport between two distri- butions (cf. [Lipman et al. (2023; ; Liu et al. (2023a): Albergo & Vanden-Fijnden (2023)). By approximating the instantancous veloc- ity field of this transport with a neural net- work, FM enables high-fidelity multimodal generation by solving the ordinary differen- tial equation (ODE) with the neural network- approximated vector field as its forcing term (Esser et al.] 2024 Ma al. 2024; Polyak et al. 2024:Jing et al. [2224] Campbell et al. 2024). Nevertheless, this high-fidelity gener- ation requires multiple expensive neural net- work evaluations, counted by the number of function evaluations (NFEs) (Chen et al.
0.7
0.7
0.5
0.4
0.4
0.3
0.3
0.1
0.1
0.0
0.0
-2
X
I-NFE McanFlow
8-NFE McanFlow
0.8
0.8
0.7
0.7
0.5
0.5
0.4
0.4
0.1
0.1
X
32-NFE MeanFlow
1-NFE RMFlow (ours)
2018).
Figure 1: Contrasting MeanFlow with RMFlow for mixture Gaussian sampling: see Section 5.1 for ex- perimental details and more results.
Several approaches aim to accelerate diffusion- and flow-based models for high- fidelity generation with only a few NFEs. Among these, consistency models (CMs) Song et al. 2003);; Geng et al. (2024); [Song & Dhariwal (2023); Lu & Song (2025) achieve remarkable performance and efficiency. Distillation is noticeable idea; for instance, local FM (Xu et al.
2024a) breaks the flow into local sub-flows, enabling smaller models and easier distillation.
Recently, flow maps (Boffi et al. 2024 2025) and mean flows (MeanFlows) (Geng al. (2023);
cf. Section 2 have been proposed to enable aggressive 1-NFE generation, and a prominent advan-
"Correspond to wangbaon wangbaonjegmai c.c
j@gmail

---

Published as a conference paper at ICLR 2026
tage of flow maps and MeanFlows is that they require pre-training, distillation, curriculum learning. Empirically, MeanFlows achieve high-quality image generation with fewer transport steps than FM models. However, preserving this quality typically requires multiple evaluations of the mean velocity field, as collapsing the process to 1-NFE often causes significant performance degra- dation. We showease this issue by sampling mixture Gaussian distribution using MeanFlow; Section 5.1 for experimental details. Figure ] shows the significant gap between exact (data) and sampled distributions when using I-NFE MeanFlow, and this gap reduces NFE increases.
We further showcase the signif- icant generation етrоr of 1-NFE MeanFlow for the benchmark QM9 molecule generation (Ramakrishnan et al. 2014); see Section [.2 for ex- perimental details and additional re- sults. Figure 2 illustrates that 1- NFE MeanFlow produces an invalid structure, where the molecule is frag- mented into multiple disconnected pieces. Indeed, in experiments, we consistently observed that 1-NFE McanFlow frequently generates in- valid structures. Additional quantita- tive results in Section (.22) further con- firm the significant errors associated
1-NFE MFM
1-NFE RMFlow (ours)
Figure 2: Contrasting MeanFlow with RMFlow, under the
same context, for QM9 molecule generation.
with 1-NFE MeanFlow generation.
The above numerical results motivate us to study the following problem:
Can we improve the performance of I-NFE MeanFlows for multimodal generation?
Mean Flow Transport
Noise Injection
3o + u0 (Io; 0)
2o + ú (No; 0) (omin o)E v
Fo
1
E ~ N(O, I)
Etg
0.8
0.8
Sample
0.7
0.7
0.7
0.6
0.6
0.5
0.5
x 0.4
0.4
0.4
0.3
0.3
0.3
0.2
0.2
Data
0.1
0.1
0.1
0.0
-2
-1
1
0.0
X
Figure 3: Schematic of proposed RMFlow: it first applies 1-NFE MeanFlow transport, then
fines the result by subsequent noise-injection step; see Section The average velocity úo. (Zo; 0) of RMFlow is trained by incorporating the maximum likelihood objective into the MeanFlow frame-
work, as in equation I2
1.1 OUR CONTRIBUTIONS
We propose RMFlow- improved 1-NFE MeanFlow model for multimodal generation. RM- Flow leverages 1-NFE MeanFlow for coarse transport, accompanied by subsequent tailored noise- injection step refine the generation; Fig. depicts the idea of RMFlow. The results in Figs [ and demonstrate that RMFlow achieves substantial improvement in generation quality over 1-NFE MeanFlow. In particular, it effectively mitigates invalid structures, producing coherent and valid molecular graphs. Our key contributions are:
2
We propose RMFlow to enable I-NFE sigh-fiielity multimodal generation by integrating the guid- ance encoding with tailored noise-injection refinement strategy; Section
2

---

Published as a conference paper at ICLR 2026
We design theoretically principled training objective for RMFllow that balances minimizing the Wasserstein distance between probability paths and maximizing the likelihood of the learned target distribution; see Section 4
We show the compelling. often (near) state-of-the-art results of RMFlow for benchmark text-to- image, text-to-structure, and time series generation (see Section 5
1.2 ADDITIONAL RELATED WORKS
To our knowledge, this work is the first to improve MeanFlows by introducing noise-injection refinmeett for 1-NFE generation. This differs from existing couplings of flow and diffusion models, such as Diff2Flow (Schusterbauer et al. (2025). which transfers knowledge from pretrained diffu- sion models to flow matching models, and generator matching (Patel et al., 2024). which connects diffusion and flow matching under Markov generative processes.
Another line of work studies emoo control in FM. Prior analyses of probability flow ODEs and FM (Song et al. 2021; Lu et al. 2022, Lai et al. 2023 Albergo et al., (2023) show that FM alone can- not guarantee likelihood maximization or KL divergence minimization between target and learned distributions.
1.3 ORGANIZATION
We organize the rest of this paper as follows: We provide necessary background materials in Sec- tion 2 We present our proposed 1-NFE RMFlow in Section We present our training loss function for R.Fllow in Section 4 Our numerical results for RMFlow in Section Technical proofs and additional experimental details and results are provided in the appendix.
BACKGROUND
In this section, we provide brief review of flow-based generative models, especially MeanFlows.
For detailed exploration of FM, we refer the reader to (Lipman et al.) 2023 Liu et al. 2023a
Fukumizu al., 2024). For given data æ = &data ~ p and prior sample wo ~ (e-g. standard Gaussian (00, D), (conditional) flow path canh-ccnn the two samples-car cane be constructed as x; = 0g + b with a, and b, being predefined schedules. common choice is a₂ = 1 and b. = f, which corresponds to rectified flow (Liu et al. 2023a). This interpolation can be equivalently expressed as the solution to the ODE é; = (z, where = (Io. æ1) denotes the coupling of start and end points, and uy (xxx)22 = 0 + b,a0 is the conditional vector field. FM learns an unconditional vector field (x)) E.g (u, (ar[=)lar, æ:), which does not require knowledge of the pair z (Io, This is achieved by training a neural network (z 0) to minimize the objective:
(C E,,; [xú (x(x((2)1
After training, we generate data by integrating da ãz (x,; from to with wo q.
Although FM is conceptually simple, sample generation requires multiple evaluations of úr 0), which can be computationally intensive. To address this inefficiency issue, MeanFlow learns an averaged velocity field based on the instantancous velocity field u, (xx),, defined as:
J. (x.))dd..
Ur (x:)
This allows data generation by transporting x; to Er using the approximate average velocity field بz.;
x; z; + (r t)ú,, (xt; 0).
In particular, 1-NFE generation corresponds to I2 = 20g + u (Z:: 0). For Multi-NFE generation, (x; ()) is evaluated sequentially on a chosen grid 0 consecutive grid points to transport samples. This approach achieves high-fidelln generation with significantly fewer NFEs compared FM models that rely the instantaneous velocity field (ar).
T. and is applied between
3

---

Published as a conference paper at ICLR 2026
In MeanFlows, the mean velocity field (x) is approximated by neural network (x; 0), with the weights being calibrated by minimizing the following conditional mean flow matching CCMFM loss function:
LCMFM (0) = rR,, [ll (x: 0) sg (#4,
tgt rrr (a;; 0) []
(4)
where 0 S t S r S 1 are uniform samples from the interval [0, 1] and uirr is the target defined as: 1gt ,, a 0) un (z]z)) (r V ux( 2) (z;
with sq denoting stop-gradient operation. This stop- gradient approach prevents higher-oud order opti- mization while ensuring that zero loss guarantees dynamical consistency. The target velocity utr is efficiently computed using Jacobian-vector products (ivp) in autodiff libraries such PyTorch (Paszke et al. 2019) or JAX (Bradbury et al. 2018).
tgi
3 THE DESIGN OF RMFLOW
In this section, we describe the design of 1-NFE RMFlow for high-fidelity generation, with or with- out multimodal guidance.
3.1 MOTIVATION
In practice, the true data distribution Edata is unavailable due to its complexity and the limited nature of observed data. Following the standard practice in the field, as established in works such as (Ho) ett al.] [2020) (Song et al., 2021; Lu et al. [022] Lipman et al. 2023), we approximate it with noisy, smoothed version Figt = dana + Tmin ~ Pegt, where e ~ (O. 1) and amin is small (e.g.. 10-3). This approach ensures stability and robust learning of the data distribution.
MeanFlows learn a neural network, by minimizing LCMPM in equation to transport a prior sample 2o directly to the noisy target, i.e., œ Ptgt. Boffi et al. (2024) (2025) showed that this approach reduces the Wasserstein distance between the target distribution Pegt and the learned distribution pa:
Theorem 3.1. (Boffi al. 2025 There exists a constant M > 0 such that:
M (CMFM(O) 2 w2 (Pigt. pe) [))= == inf (EIl(PegtPo)
E ((aa))> [||= y|
(5)
where M is constant, W2 (Ptgt. pe) denotes the Wasserstein distance between Pigt and pa. and [pitt, Po) is the set of all joint distributions with marginals Pigt and po-
While controlling the Wasserstein distance provides a meaningful measure of distributional align- ment, empirical evidence indicates that FM enforces additional constraints, such as KI. divergence (Lu et al. 2022). often achieves superior generative performance over the FM baseline. With this in mind, aim to enhance the fidelity of I-NFE MeanFlows, pushing beyond current limitations in manner analogous to improvements seen in FMs.
3.2 NOISE INJECTION REFINEMENT
We decompose the generation process into two stages. In the first stage, a I-NFE MeanFlow trans- ports the prior 2g to an intermediate noisy state
3 Edata + de1; with e ~ NN I) and a Fmin-
(6)
In the second stage, single noise injection step is applied:
iigt= = σ1
P 0² -2. 6 I(0,I,
to generatte the final sample. This additional noise injection aligns with the designs of VABs (Kingma & Welling 2013). allowing principled likelihood maximization via a loss term derived from the evidence lower bound (ELBO) (Wainwright et al. (2088) to optimize the MeanFlow parameters. We will prove Theorem 4.T that this formulation enables control over the KL divergence between the target distribution Ptgt and the learned distribution po-
4

---

Published as a conference paper at ICLR 2026
In summary, data generation process is defined as
Ligt wo + Uo,1 (2o; 0) + min 0² e2. 62 (O, I),
where do (20; 0) denotes the learned average velocity field. Although RMFlow is conceptually a two-stage framework, equation 8 demonstrates that generation is performed in single step: the learned flow is evaluated once (I-NFE), and noise term is added in parallel to produce the output.
3.3 MULTIMODALITY
To support cross-modality generation, we incorporate an encoder (c) that embeds conditioning signals (e.g., text prompts). The prior samples for both guided (potentially multimodal) and un- guided generation are defined as
(c) + act, for guided generation,
370
for unguided generation,
where MO.,)) I). This design allows the flow to incorporate multimodal guidance if available, while defaulting to unconditional generation otherwise. Here, (-) an encoder chosen following common practice (see Section 5). and o. (e.g., 10-3) is pre-chosen control perturbations.
Specifically, for a given data pair (Idata: c), we train the MeanFlow to transport the prior sample 20 &w (c) + oe to the intermediate target 2 ddtta dcl, where ç NN( I). The encoder and Mean are optimized jointly, and we will discuss the training objective in Section 4
Remark 1. Our proposed RMFlow differs from MeanFlow two aspects: (1) We apply tailored encoder to the guidance, (2) add noise injection step to refine the generation result.
4 THE TRAINING OF RMFLOW
In this section, we present the training procedure for RMFlow. We first establish the theoretical foundation of noise-injection refinement, showing that enables likelihood maximization of the leamed distribution with respect to the target distribution. Building on this, we introduce joint training objective that combines LCMFM (for Wasserstein control) with likelihood maximization and optional guidance regularization, ensuring both fidelity and flexibility in guided generation. Finally, we adopt parameter- efficceett fine-tuning (PEFT; cf. (Hu et al. 2022 to implement RMFlow for large-scale tasks.
Dettmers et al. 2023)
4.1 LIKELIHOOD MAXIMIZATION
In this section, we show that the noise-injection step in equation R enables likelihood maximiza- tion during RMFlow training. Specifically, for a given prior sample 2p. the intermediate sample generated by the MeanFlow is
Φ1 3u ú (Zo; 0).
By equation g the conditional distribution of the final generated sample given the prior is
figt
/V
Øp + uo (zo: f), (o - o²) o
This specifies parametric conditional distribution. Given an observed target Zegt, the correspond- ing conditional log-likelihood is
log (o xo) 2( o2)) 1| (xo + ú 0))1 + +
(9)
where C- (o o²)) and d is the dimensionality of the data. Notice that equation and equation 7 indicate the observed target Zigt = Edata + Jminc, where c ~ NNO,0
Therefore, define the following loss term to maximize the likelihood:
L == Emg,
[I + Amine) - (wo + 0o, (20o; 0)) 0)
(10)
5

---

Published as a conference paper at ICLR 2026
The following theorem formalizes the theoretical guarantee of the noise-injection refinement. In particular, demonstrates that minimizing the loss LNLL maximizes the expected log-likelihood, thereby reducing the KL divergence between the target and learned distributions.
Theorem 4.1. The negative log- sellikeiihoodd loss LNLL provides lower bound on the expected log- likelihood of the target distribution:
-A LNL. + C < Emage [log (( ) (Pegt) - DET. (Prgt po),
(11)
where H (Pogs) == En.. [log Pigs] is the entropy of Pegt, DKL. (Pegesl lpa) Emg Dlo Pug denotes
pe
the KL divergence between the target and the learned distributions, and A, B > 0 are constants.
4.1.1 JOINT TRAINING OBJECTIVE
RMFlow is trained by jointly optimizing the original MeanFlow loss (Wasserstein control) and like- lihood maximization, resulting in the following objective function:
CRMFlOW (0,w) = LCMPM + AILLLL + A2lE e) (c)lI)
(12)
where A2 2 are two hyperparameters. We remark that Term controls the gap between the probability flows of the exact and approximated mean velocities in intermediate states, Term II for likelihood maximization, and Term III is designed for guided generation and is set for unguided generation. Here, the expectation in III is taken over all data aita-guidance pairs (Idata: c).
Remark 2. Term III in equation 12 can be considered as regularization on the prior distribution, and a similar term is used in training VAE (Kingma & Welling 2013). Empirically, we observe that term III can be very large, resulting in substantial performance degradation.
4.2 MEMORY-EITICIEN FINE-TUNING
For relatively small-scale tasks, we train our RMFlow by directly minimizing LHMFlow- Compared to LCMFM, our new objective CRMFlOW introduces additional gradient pathways, increasing memory footprint. To balance efficiency and performance for large-scale tasks, first train the MeanFlow model by minimizing LCMFM, and then fine-tune it using PEFT (Hu et al.) 2022 Dettmers et al., 2023), with LRMFlOW supervised objective in large-scale experiments
text-to-image and molecule generation tasks. During fine-tuning, we further strengthen training by integrating 1-NFE sampling with policy-gradient objective that incorporates physical feedback on sample quality for molecule generation tasks, as described in Zhou et al. (2025)
NUMERICAL EXPERIMENTS
In this section, we validate the efficacy and efficiency of RMFlow for both guided and unguided sample generation. We consider two synthetic tasks: sampling ID mixture Gaussian distribution and 2D checkerboard density (Section 5.I.. We also consider several benchmark tasks, including context-to-molecull structure generation (Section 5.2). sampling trajectories of dynamical systems (time series; Section (.3). and text-to-image generation (Section 5.4).
Software and Equipment. We implement syntheei tasks, context-to-molecnt generation, and text- to-image generation using PyTorch. We implement the time series generation task using JAX. Additionally, we use Torch DDP and torch. compile to optimize the model execution for context-to-molte molecule and text-to-image generation. All the experiments are carried out on multiple NVIDIA RTX 3090/4090 GPUs.
Training Setups. See Appendix B for the details of training setups.
Evaluation Metrics:
For synthetic tasks and time-series generation, evaluate performance using the estimated KL divergence and total variation (TV) distance between the generated samples and the ground-trutie Both KL and TV computed from densities obtained via histogram- based esti- mation the sample and ground-truth distributions. For molecule generation, we predict bond types from pairwise interatomic distances and atom types, and then compute atom and molecule stability,

---

Published as a conference paper at ICLR 2026
following Hoogeboom et al. (2022). For the image generation task, assess sample quality using the Fréchet Inception Distance (FID) (Heusel et al. 2017).
We use NFE to measure generation efficiency following (Geng et al. 2025). Notice that the Gaussian noise injection step takes negligible time compared to the neural network function evaluation.
I-NFE McanFlow
8-NFE McanFlow 32-NFE McanFlow
1-NFE RMFlow (ours)
TV
1,4422
0.7977
0.6737
0.7567
K1
0,8074
0.4074
0.1017
0.2332
Table I: Contrasting 1-NFE RMFlow with 1/8/32-NFE MeanFlow for mixture Gaussian sampling. I-NFE RMFlow outperforms both 1- and 8-NFE MeanFlows, while slightly worse than 32-NFE MeanFlow.
1-NFE McanFlow
8-NFE McanFlow
32-NFE McanFlow
1-NFE RMFlow (ours)
TV 0.238
0.167 0.139
0.155
0.173
KI
0.311
0.118
0.163
Table 2: Contrasting 1-NFE RMFlow with 1/8/32-NFE McanFlow for checkerboard density sam- pling. 1-NFR RMFlow significantly outperforms 1-NFE MeanFlow, closing the performance gap to multi-NFE MeanFlow.
5.1 SYNTHETIC TASKS
In this experiment, we train a simple ResNet-based model under both MeanFlow and RMFlow frameworks for 105 iterations using batch size of 256 to sample (1) ID Gaussian mixture Pdata (1.5, + 0. 0.04), and (2) 2D checkerboard where
the probability density resembles a checkerboard pattern. We consider 1/8/32-NFF MeanFlow and and / show that 1-NFE RMFlow significantly outperforms 1-NFE MeanFlow, closing the performance gap multi-NFE MeanFlow.
1-NFE RMFlow for sample generation. Tables [
5.2 CONTEXT-TO-MOLECCLEE QM9 GENERATION
We train MeanFlow and RMFlow for context-to-molecule generation the QM9 dataset Ramakr- ishnan et al.) (2014), a benchmark containing atomic coordinates and quantum-chemical properties for 130k small molecules with up 9 heavy atoms (up to 29 atoms including hydrogens). Fol- lowing Hoogeboom et al. (022), we perform condition generation on seven molecular properties: (1) number of atoms, (2) HOMO, (3) LUMO, (4) a (isotropic polarizability), (5) gap, (6) l (dipole moment), and (7) C, (heat capacity). These properties are concatenated into context vector and mapped the data space using (c). parameterized by a single EGNN block Gareia Satorras al. 2021
Our model backbone follows the EGNN architecture (Garcia Satorras et al. 2021: Hoogeboom et al. 2022), augmented with time-embedding module for the additional scalar time variable In addition, molecule stability used as the reward within the RL framework, following the approach of Zhou al. (2025). to provide feedback during training (see Section 4.2 and Appendix [.8). We adopt the train/val/tes splits of Anderson et al. (2019), comprising 100%; 18k/13/ molecules, respectively. B that I-NFE RMFlow attains ofate-of-the-art performance, whereas com- peting SOTA methods require n-NFE with n > 1. Figure depicts few randomly generated molecules and the corresponding contexts.

---

Published as a conference paper at ICLR 2026
(I)
(2)
(3)
(4)
(5)
(7)
(9)
(10)
(11)
(12)
Contexts:
(1) atom:
homo: -6.34
lumo: 1.82
alpha: 77.77
gap: 8.16
mu: 1.55
Cv: -27.27
(2) N atom: 18
homo: -6.62
lumo: -0,85
alpha: 74.21
gap: 5.77
mu: 3.27
Cv: -19.37
3 N atom: 22
homo: -6,44
lumo: 1.99
alpha: 83.22
gap: 8.43
mu: 1.52
-2:129.10
(4) atom:
homa: -6.56
lumn: 2.01
alpha: 81.16
gap: 8.57
mu: 0.31
Cv:-30.52
(5) atom:
homo: -5.79
lumo: -0.02
alpha: 77.48
gap: 5.77
mu: 2.00
Cv: -2.723.75
(6) N
atom: 20
homo: -5.30
lumo: 0.34
alpha: 91.41
gap: 5.64
1.54
Cv:
(7) N
atom: 16
homa: -7.05
lumo: -1 -1,20
alpha: 68.35
35
gap: 5.84
mu: 1.37
Cvv: -18.6
(8) atom:
homo: -6.94
lumo: 0.94
alpha: 78.01
gap: 7.88
mu: 2.36
Cv:099
(9)
16
homa: -6.74
lumo: 0.60
alpha: 58.89
gap: 7.34
4.43
Cv:-18.93
(10) N atom: 18
homo: -6.76
lumo: 0.77
alpha: 73.96
gap: 7,53
mu: 1.95
Cv: -24.84
(11) N atom: 16
homo: -6.09
lumo: -0.10-
aalha: 74.33
gap: 5.98
mu: 3.87
CC: -22.98
(12) N atom: 20
homo: -6.82
lumo: 0.53
alpha: 76.55
gap: 7.36
mu: 1.12 Cv: -26.04
Figure 4: A few randomly selected RMFlow-generated molecules, together with the corresponding
contexts.
Metrics
No Tuncation Error (Discretization)
Atomic Stab. (f)
Mol Stab. (t) NFE(L)
ENF
85.10.1 %
4.9...2%%
E-DM Hoogehoom
2022
98. 98.73++0.1
62.11
>
Bridge Wu Al 2022
98.7=0.1 %
81.8+..2%
Bridge + Foree Wi er al.
2022
98.8=0.11 %
84.6
>
GeoLDM Xu cl ul. 2023
98.
89.40+0.5%
GeoBEN Song
2024
99.0%
93.9%
E-DM RLPF Zhou et il. (025
99.1%
93.4%
>
MeanFlow w/o contexts
98. 08.22+0.777
7..3
MeanFlow w/ contexts
98 44..00 %
84.3
RMFlow
(ours)
98.8 0.0.5
0..5
RMFlow w/ contexts (ours)
98.94 %6
03. Ro
RMFlow w/ contexts RLPF (ours)
98 98.940..0 %
93.5+n.3%
Data
99%
95.2%
Table 3: Contrasting the performance of different models for QM9 molecule generation. We run
RMFlow with contexts by randomly selecting 10 contexts in the test dataset of QM9 five times.
5.3 TIME SERIES: DYNAMICAL SYSTEM
Sampling trajectories in dynamical systems under event guidance is a key challenge for predicting and understanding complex phenomena such as climate and extreme events (Perkins & Alexander, 2013 Mosavi et al. 2018). Recent works (Finzi et al.) (023) and (Huang et al.) 2025) have intro- duced diffusion and FM models specifically designed for event-guided sampling.
In this experiment, we perform dynamical system trajectory forecasting with MeanFlow and RM- Flow, formulating it as time series problem by discretizing the time variable into uniform tervals. Each trajectory (either from dataset sampled) is a discrete time series of vectors con- catenated into Tdata [an (m }] E Mdd where M the total number of time steps, the dimension of the system, and x (Tm) E Id denotes the discretized trajectory at time Tm Our goal is to generate Edata = (x((m) [=(Tm) E Mdd with I-NFE using McanFlow and RMFlow.
We train our models the Lorenz and FitzHugh-Nagumo dynamical systems (see (Huang et al. 2025 Appendix B. for brief review of these two models); using a U-Net backbone. For event guidance, where events are defined by a constraint function E {Zdata C(Idata) > 0), adopt a simple but effective design: the event- ween-guudancee vector and the first three states (TTT), x (22))an x x((T)) embedded through an MLP du into the target data space TR Md This avoids reliance Tweedie's
8

---

Published as a conference paper at ICLR 2026
formula used in (Finzi al. 2023 Huang al. 2025
Tables 4
and 5
show that RMFlow
yields significanlly better 1-NFE generation than MeanFlov Flow, while achieving accuracy comparable to multi-NFE methods.
Lorenz
FizzHugh-Nagunn
Model
w/o (4)
w/ E (L)
w/o (4)
E (4)
NFE (4)
Diffusion Huang et al. 2025
0.0314
0.1001
0.0277
0.1192
128
Huang el an (202
0.0348
0.0972
0.0314
0,2164
FDM Huang el 2025
0.0306
0.0914
0.0266
0.1168
128
McanFlow
0.0469
0.1250
0.0398
0.2268
MeanFlow
0.0366
0.1011
0.0345
0.1988
8
McanFlow
0.0351
0.0991
0,0302
0.1723
32
RMFlow (ours)
0.0332
0.0956
0.0289
0.1543
1
Table 4: TV distance between the generated (by different models) and test trajectory distributions,
estimated from histogram-based density approximations, with/without conditioning the event.
Lorenz
FiizHag--aagumo
Model
w/o E (L)
w/ E (L)
w/o E (U)
w/ E (1)
NFF (1)
Diffusion Huane (2025)
0.0056
0.2774
0.0260
0,3011
128
FM Huane el 12025
0,0081
0.2560
0,0280
0.3468
128
FDM Huang et al. 2025
0.0049
0.3045
0,0280
0.2084
128
MeanFlow
0.0109
0.3887
0.0347
0.3921
1
MeanFlow
0.0091
0.3163
0,0297
0.2422
8
McanFlow
0.0054
0.2722
0.0281
0,2490
32
RMFlow (ours)
0.0059
0.2866
0.0287
0.2499
Table KL divergence between the generated (by different models) and trajectory distributions,
estimated from histogram based density approximations, with/without conditioning on the event.
5.4 TEXT-TO-IMAGE
In this experiment, train MeanFlow and RMFlow for text-to-image generation the COCO dataset Chen al. (015)) Following Stable Diffusion Rombach et al. (2022), all operations are performed in the latent space IR 4x32x332 The mapping (c) converts the text conditions into initial latent states. Concretely, we fine-tune the pretrained text-embedding model e5-base (Wang et al.) 2022) and attach an MLP project the embeddings into the latent space. Additionally,
fine- tune the Stable Diffusion VAE decoder on COCO using PEFT (Hu et al., 2022] Dettmers et al.) 2023) that it can decode the final latent state into images. Both MeanFlow and RMFlow --MM-aaaretter U-Net as the latent-space backbone.
We adopt the Karpathy split (Karpathy & Fei-Feil 2015) for training and validation, and evaluation is performed with COCO FID-30K Tollowing Rombach et al. (2022); He et al. (2025) (details in Appendix B.4). As shown in Table 6 RMFlow attains FID comparable to the best single-step generators on COCO, such as Distilled Stable Diffusion (Liu et al. (023b), StyleGAN-T (Sauer al. 2023). Importantly, RMFlow (and MeanFlow) orthogonal to the other methods listed in Table it does not rely auxiliary models for training. In contrast, GAN-based approaches require a discriminator, and distilled models depend on a pretrained teacher. Moreover, our models trained under limited computational resources (e.g., RTX 3090/4090 GPUs with 24 GB memory) using mixed-precine precision bf16, whereas most state-of-the-ar models listed in Table 6 are trained on multiple A100 80 GB GPUs with full-precision fp16. These results indicate that RMFlow has strong potential for further improvement if trained with larger computational budgets. We also report the performance on CLIP score, see Appendix B.7
CONCLUDING REMARKS
In this work, introduce RMFlow, refinement of MeanFlow with minimal computational and memory overhead. The central innovation lies in augmenting the 1-NFE MeanFlow with subse- quent noise injection step, which facilitates likelihood maximization. To support this mechanism,

---

Published as a conference paper at ICLR 2026
CHURCH
GO
(1) The dining table near the kitchen has
(1) A man making sandwich lunch truck.
(2) A woman ricing surfboard
wave in the
(2) A skateboarder performing trick on an indoor ramp.
(3) A woman sitting on wooden park bench
large sign that says street name (3) There
(4) A stack of old trunks and luggage against
(4) A white plate lopped with onion rings and
(5) Rows of unripe bananas on display in
(5) A big zebra and little zebra standing and looking.
(6) stop sign with spray painted words on it.
(6) The meal ready the tray be caten.
(7) couple of horses that are next to fence.
(7) A bike and dog on the sidewalk outside a
(8) red and white street sign mounted on...
(8) A cat up on desk drinking milk from glass.
Figure 5: COCO dataset samples generated using 1-NFE RMFlow conditioned on different input prompts.
Type
params
NFE
Teacher-free discriminator free)
Resolution
Stuble Didfusion v1.5 Rombuch et
2022
Diff
S60M
9.62
255
Stable Diffasion v2
Diff
860M
13.45
256
10.1
Show-o
PixAnt hen af 2414
ODE
630M
20
7,32
256 x 256
LDM
Krmmmiie or 24122
Diff
1.48
26
12.63
254 X 256
VQGAN,
GAN
75M
26.94
258 X 256
StyleGA T 21
GAN
IB
13.90
256 x 256
InstaPlow Lirreni sureu
ODE
90DM
13:10
512 X 512
UFOGs
Cer0aaaKuurr Xira 20271
Diff
900M
12.78
512 X 512
Stable Dilfuston bisin Lin 2024
Diff
900M
34.6
256 256
Rectilied How Distill Lau
ODE
9UMM
20.0
256 256
MeanHlow
ODE Duff
620M
27:31
256 X 256
RMFlow (ours)
620M
18.91
256
256
Table FID of the generated images on the benchmark COCO2014 dataset using different models.
propose novel loss function that jointly minimizes the discrepancy between the exact and learned probability paths while maximizing likelihood. Empirical results demonstrate that 1-NFE RMFlow achieves strong performance in multimodal generation tasks.
promising direction for future research is to extend RMFlow to support multiple mean flow trans- port steps. Specifically, envision applying noise-injection step after each transport step, which would require the design of a corresponding loss function to maintain likelihood maximization. This extension presents additional challenges compared the current formulation and opens avenues for more expressive and accurate generative modeling. Another limitation of RMFlow is that it uses fixed parameter yo - na during the noise injection step, which may be suboptimal. As future work, we plan explore more adaptive strategies for selecting this parameter, such making leamable or following dynamic schedule.
10

---

Published as a conference paper at ICLR 2026
7 ACKNOWLEDGEMENT
This material is based on research sponsored by NSF grants DMS-2152711 DMS-2208361, DMS- 2219956, DMS-2436343, and DMS-2436344, and DOE grants DE-SC0023490, DE-SC0025589, and DE-SC0025801.
ETHICS STATEMENT
In this paper, we propose a new framework to improve MeanFlow for efficient data generation. The new model can generate high-fidelity data efficiently. Our work belongs to fundamental research and is expected to improve existing models for generative modeling. Our work is methodological and validate our proposed approaches on the benchmark datasets. We do not expect to cause negative societal problems. Furthermore, we do not see any issues with potential conflicts of interest and sponsorship, dissriminaaioo/bias/aiansss concerns, privacy and security issues, legal compliance, and research integrity issues (e.g., IRB, documentation, research ethics.
REPRODUCIBILITY STATEMENT
We are committed to conducting reproducible research. To ensure the integrity and transparency of work, employ multifaceted approach: First, meticulously compare the novelty of our research against existing literature. This involves a thorough examination of the current state of the field to identify gaps in knowledge and demonstrate the unique contributions of work. Second, we provide detailed derivations of our proposed approaches and theoretical results. By carefully outlining the mathematical underpinnings of methods, we enhance the understanding of work and facilitate its verification by others. Third, we conduct rigorous experiments using widely recognized benchmark datasets. This allows us evaluate the performance of our methods against established standards and provides a solid foundation for comparison with other approaches. Fourth, we meticulously report experimental details, including the specific datasets used, parameters chosen, and evaluation metrics employed. Finally, we make all experimental codes, accompanied by com- prehensive documentation, publicly available. This open-source approach empowers researchers to inspect our methods, verify our results, and build upon our work. By sharing our code, we foster collaboration, advance the field, and contribute to the overall reproducibility of scientific research.
REFERENCES
Michael Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: unifying framework for flows and diffusions. arXiv preprint SOXVv:2303308797 2023.
Michael Samuel Albergo and Eric Vanden Eijnee. Building normalizing flows with stochastic
interpolants. In The Eleventh International Conference Learning Representations, 2023. URL
http://wwwwe///oommt/eett net forum (o-rummid-1117geBbcRR
Brandon Anderson, Truong Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural
networks. Advances in neural information processing systems, 32, 2019.
Kevin Black, Michael Janner, Yilun Du, Ilya Kostrikov, and Sergey Levine. Training diffusion models with reinforcement learning. arXiv preprint w:Xiv:2305.13300 2023.
Nicholas M Boffi, Michael S Albergo, and Eric Vanden-Eijnden Flow map matching. arXiv preprint arXiv:2406.07507, 2024.
Nicholas M Boffi, Michael S Albergo, and Eric Vanden-Eijnden. How to build a consistency model: Learning flow maps via self-distillation. arXiv preprint urXiv:2505./8825 2025.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- Milne, et al. Jax: composable transformations of python+ numpy programs. 2018.
11

---

Published as a conference paper at ICLR 2026
Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. Generative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design. arXiv preprint arXhv:2402.0499) 2024.
Junsong Chen, Jincheng Yu, Chongjian Ge, Lewei Yao, Enze Xie, Yue Wu, Zhongdao Wang. James Kwok, Ping Luo, Huchuan Lu, al. Pixart-o: Fast training of diffusion transformer for photore- alistic text-to-image synthesis. arXiv preprint 1rXiv:2310.0044: 2023.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Xinlei Chen, Hao Fang, Tsung-Y Yi Lin, Ramakrishna Vedantam, Saurabh Gupta, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco captions: Data collection and evaluation server. arXiv preprint arXiv:1504.00325. 2015.
Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. Qlora: Efficient finetuning of quantized Ilms, 2023. URL https://axxiv. org/abs/2305. 14314, 2, 2023.
Kingma Diederik. Adam: A method for stochastic optimization. (No Title), 2014.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. Forty-first international conference on machine learning, 2024.
Marc Anton Finzi, Anudhyan Boral, Andrew Gordon Wilson, Fei Sha, and Leonardo Zepeda-Núñcz. User-defined event sampling and uncertainty quantification in diffusion models for physical dy- namical systems. In International Conference on Machine Learning, pp. 10136-10152. PMLR, 2023.
Kenji Fukumizu, Taiji Suzuki, Noboru Isobe, Kazusato Oko, and Masamori Koyama. Flow matching achieves almost minimax optimal convergence. arXiv preprint ar arXiv:2405.20879 2024.
Victor Garcia Satorras, Emiel Hoogeboom, Fabian Fuchs, Ingmar Posner, and Max Welling. E (n) equivariant normalizing flows. Advances in Neural Information Processing Systems, 34:4181- 4192, 2021.
Zhengyang Geng, Ashwini Pokle, and Zico Kolter. One-step diffusion distillation via deep equi- librium models. Advances in Neural Information Processing Systems, 36:41914 2023.
Zhengyang Geng, Ashwini Pokle, William Luo, Justin Lin, and J Zico Kolter. Consistency models made easy. arXiv preprint arXiv:2406.1455 2024.
Zhengyang Geng, Mingyang Deng, Xingjian Bai, Zico Kolter, and Kaiming He. Mean flows for one-step generative modeling. arXiv preprint arXiv:2505. 13447, 2025.
Ju He, Qihang Yu, Qihao Liu, and Liang- Chieh Chen. Flowtok: Flowing seamlessly across text and image tokens. arXiv preprini arXiv:2503. 10772, 2025.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge local nash equilibrium. In Advances in Neural Information Processing Systems, volume 30, 2017.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-683 68511 2020.
Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffu- sion for molecule generation in 3d. In International conference on machine learning, pp. 8867- 8887. PMLR, 2022.
Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al. Lora: Low-rank adaptation of large language models. ICLR, 1(2):3 2022.
12

---

Published as a conference paper at ICLR 2026
Yuhao Huang, Taos Transue, Shih-Hsin Wang. William M Feldman, Hong Zhang, and Bao Wang. Improving flow matching by aligning flow divergence. In Forry-sccond International Conference on Machine Learning, 2025. URL
FeZimu 66GG
Manuel Jahn, Robin Rombach, and Bjöm Ommer. High-resolution complex scene synthesis with transformers. arXiv preprint arXiv:2105.06443 2021.
Bowen Jing, Bonnie Berger, and Tommi Jaakkola. Alphafold meets flow matching for generating protein ensembles. arXiv preprint arXiv:2402.04845. 2024.
Andrej Karpathy and Li Fei-Fei. Deep semantic alignments for generating image descrip- tions. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 3128-310 1128-3137, 2015.
Diederik Kingma and Max Welling, drXiv:1312.6114,200 2013.
Auto-encoding variational bayes. arXiv preprint
Chieh-Hsin Lai, Yuhta Takida, Naoki Murata, Toshimitsu Uesaka, Yuki Mitsufuji, and Stefano Er- mon. Fp-diffusion: Improving score-based diffusion models by enforcing the underlying score fokker-planck equation. In International Conference on Machine Learning, pp. 18365-18398. PMLR, 2023.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference Learning Repre- sentations, 2023. URL
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. The Eleventh International Conference on Learning Repre- sentations, 2023a.
Xingchao Liu, Xiwen Zhang, Jianzhu Ma, Jian Peng, et al. Instaflow: One step is enough for high-quality liffusion-based text-to-image generation. In The Twelfth International Conference on Learning Representations, 2023b.
Cheng Lu and Yang Song. Simplifying, stabilizing and scaling continuous-timg consistency models. In The Thirteenth International Conference on Learning Representations, 2025. URL https:
//openzeview. net.
Cheng Lu, Kaiwen Zheng, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Maximum likelihood training for score-based based diffusion odes by high order denoising score matching. In International Conference on Machine Learning, pp. 14429-14460. PMLR, 2022.
Nanye Ma, Mark Goldstein, Michael S Albergo, Nicholas M Boffi, Eric Vanden- and Sain- ing Xie. Sit: Exploring flow and diffusion-basee generative models with scalable interpolant transformers. In European Conference on Computer Vision, pp. 23 40. Springer, 2024.
Chenlin Meng. Robin Rombach, Ruiqi Gao, Diederik Kingma, Stefano Ermon, Jonathan Ho, and Tim Salimans. On distillation of guided diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 14297-14306, 2023.
Amir Mosavi, Pinar Ozturk, and Kwok-wing wing Chau. Flood prediction using machine learning mod- els: Literature review. Water, 10(11):1536, 2018
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary Vito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems, volume 32, 2019.
Zeeshan Patel, James DeLoye, and Lance Mathias. Exploring diffusion and flow matching under generator matching. arXiv preprint 2rXiv:2412.10024 2024.
13

---

Published as a conference paper at ICLR 2026
Sarah E Perkins and Lisa V Alexander. On the measurement of heat waves. Journal of climate, 26 (0):4400--511,, 2013.
Adam Polyak, Amit Zohar, Andrew Brown, Andros Tjandra, Animesh Sinha, Ann Lee, Apoorv Vyas, Bowen Shi, Choh-Yao Ma, Ching-Yao Chuang, al. Movie gen: A cast of media founda- tion models. arXiv preprint arXiv:2410.1321 2024.
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific Data, 2014.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and pattern recognition, pp. 10684-10695, 2022.
Axel Sauer, Tero Karras, Samuli Laine, Andreas Geiger, and Timo Aila. Stylegan-t: Unlocking the power of gans for fast large-scale text-to-image synthesis. In International conference on machine learning, pp. 10105-30110 PMLR, 2023.
Johannes Schusterbauer, Ming Gui, Frank Fundel, and Björ Ommer. Diff2flow: Training flow matching models via diffusion model aligmment. In Proceedings of the Computer Vision and Puttern Recognition Conference, pp. 28347-28357, 2025.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learn- ing. pp. 2256-2265, pmir, 2015.
Jiaming Song. Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502, 2020.
Yang Song and Prafulla Dhariwal. Improved techniques for training consistency models. arXiv preprint arXiv:2310.1411 2023.
Yang Song, Conor Durkan, Jain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances in neural information processing systems, 34:1415- 1428, 2021.
Yang Song. Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv preprint arXiv:2303.01469, 2023.
Yuxuan Song, Jingjing Gong, Yanru Qu, Hao Zhou, Mingyue Zheng, Jingjing Liu, and Yeng Ma. Unified generative modeling of 3d molecules via bayesian flow networks. arXiv preprint arXiv:2403.15441, 2024.
Martin Wainwright, Michael Jordan, Graphical models, exponential families, and variational inference. Foundations and Trends@ in Machine Learning. 1(1-2):1-305, 2008.
Liang Wang, Nan Yang, Xiaolong Huang, Binxing Jiao, Linjun Yang. Daxin Jiang, Rangan Ma- jumder, and Furu Wei. Text embeddings by weakly-supervised contrastive pre-training. arXiv preprint arXiv:22/2.03533 2022.
Lemeng Wu, Chengyue Gong, Xingchao Liu, Mao Ye, and Qiang Liu. Diffusion-based molecule generation with informative prior bridges. Advances in neural information processing systems, 35:36533-36545, 2022.
Jinheng Xie. Weijia Mao, Zechen Bai, David Junhao Zhang, Weihao Wang, Kevin Qinghong Lin, Yuchao Gu, Zhijie Chen, Zhenheng Yang, and Mike Zheng Shou. Show-o: One single transformer to unify multimodal understanding and generation. arXiv preprint arXiv:2408. 12528, 2024.
Chen Xu, Xiuyuan Cheng, and Yao Xie. Local flow matching generative models. arXiv preprint 2:Xi::2410.02548, 2024a.
Minkai Xu, Alexander S Powers, Ron O Dror, Stefano Ermon, and Jure Leskovec. Geometrie latent diffusion models for 3d molecule generation. In International Conference on Machine Learning, pp. 38592-38610. PMLR, 2023.
14

---

Published as a conference paper at ICLR 2026
Yanwu Xu, Yang Zhao, Zhisheng Xiao, and Tingbo Hou. Ufogen: You forward once large scale text-to-image generation via diffusion gans. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 8196- 8206, 2024b.
Yufan Zhou, Ruiyi Zhang, Changyou Chen, Chunyuan Li, Chris Tensmeyer, Tong Yu, Jiuxiang Gu, Jinhui Xu, and Tong Sun. Towards language-free training for text-to-image generation. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 17907- 17917, 2022.
Zhijian Zhou, Junyi An, Zongkai Liu, Yunfei Shi, Xuan Zhang, Fenglei Cao, Chao Qu, and Yuan Qi. Guiding diffusion models with reinforcement learning for stable molecule generation. arXiv preprint 1rX1512508.16521, 2025.
15

---

Published as a conference paper at ICLR 2026
A TECHNICAL PROOFS
Theorem 4.1. The negative log-likelihood loss LNLL provides lower bound on the expected log- likelihood of the target distribution:
-A LNLL + C Ewige [log pe (Frgt)) H
H (Pegt) DKL (Pegt po),
(11)
where (Pogt) = - -Eeo [ogDug is the entropy of Pegt. DEL (Pege (pe) = Exar log Pes denotes
Po
the KL divergence between the target and the learned distributions, and A, B > 0 are constants.
proof of Theorem 4.1
We begin with the marginal likelihood:
log pe (Frgt)
log Ero (agg (20))]]
(13)
Euro [log po (Fegt [r0)]]
where the inequality follows from Jensen's 's inequality. Taking expectation over Figt gives
Eggge lo Po (Fege)
E Exg. log po (Fegt [zu]].
(14)
Now, by substituting the log-likelihood expression equation g
we obtain
Eggg Magt [log Po (Figt [mo)]
Eu Frgtt
(mo + úo (aro: 01) 1| c]
2 miu
E Bazg
|aFeet - (2r0 + ( 0)] +C
2(a
-0)2
2(amin o)² CNLL +C
(15)
Combining the inequalities, there exist constants A, C > O such that
A LILL + C S Eg Po Po
(16)
Finally, recall that
Eg [log po (Itg:)] = (pegt) - DKL (po),
(17)
which establishes the desired relation.
B EXPERIMENT SETUP
Flow map design. For a given pair = (mrg m1), we choose the conditional velocity field u. (x<x)2)
following (Albergo & Vanden-Fijnden, 2023). i.e.,
1(t)
7(t) (rr lm (1-1) (1-)) + (x) -ro):
(18)
where (t) 17(1 t) with (() 17, 7(1) 0, and 17 10-2 2,5 10 2, 10
Training Setup:
Loss Metric: Following Geng et al. 2025). we focus part (mean flow loss) of LRMFLOW 12 2 where denotes the regression error. In practice, we apply weight Ç. When m 0.5, this formulation becomes closely related to the Pseudo-Huber Huber loss introduced in Song & Dhariwal (202);; hence, adopt 0.5 for all experiments. The hyperparameters X and X₂ in LRMFloW are selected individually for each experiment, reported below.
expressed as A
m with
16

---

Published as a conference paper at ICLR 2026
Time sampling and condition: For the part IIII and II in LRMFIOW only need to sample zo that the time is always zero. Similar one sampling method used in Geng et al. (2025), we sample (t,r) such that p(t) 2t and
p(rlt) - lusrst + + ()-r -t)
t
so for given sampled i, we sample r from 44(0, t) with probability and set r = t probability q-
q is selected individually for each experiment, reported below. We use positional embedding for
(r, t), which are then combined and provided as the conditioning of the neural network. As used in Geng et al. (2025), is necessary for the network to directly condition (r, have ur, (:0)= = net(.41--
B.1 ABLATION STUDY
In this task, focus on how the value of Å₁ balances the Wasserstein (part in equation 12 and the likelihood (part Il in equation 12
B.1.1
GAUSSIAN MIXTURE
Table shows that the RMFlow has the best performance when X le 1, le 2.
A1
0
le-2
le-1
1e1
le2
I
25
24
SEA
Table 7: Ablation study the Gaussian mixture task with varying A1-
B.2 CHECKERBOARD
Table 8 shows that the RMFlow has the best performance when A le 1.
A1
le-2
le-1
lel
lc2
TV
0.238
0.201
0.173
0.222
0.289
KL
0.311
0.228
0.163
0.237
0.425
Table 8: Ablation study on the 2D Checkerborad synthetic task with varying A
B.3 QM9
Table shows that the RMFlow has the best performance when \ = 5e 2.
le-2
5e-2
1e-1
lel
Atomic Stab.
98.4
98.8
98.9
98.8
98.0
97.6
Mol Stab.
84.3
92.1
93.2
92.9
87.2
83.9
Table 9: Ablation study on the QM9 molecule generation task with varying A1-
B.3.1 CONFIGURATION
The Gaussian mixture and checkerboard experiments share the same configuration, differing only in the model's input and output dimensions. The configuration is summarized in Table 10
17

---

Published as a conference paper at ICLR 2026
Model
number of layers
hidden dim
256
activation
SiLU
Train
iteration batch size optimizer Ir schedule 1r A .22) ema decay precision À
1e5
256
Adam Diederik 2014
polynomial
le-4
(0.9, 0.95)
0.9995
fp32
[le-1, 1c-2]
Table 10: Training Setup and Backbone Configuration for MeanFlow and RMFlow on Synthetic Tasks
B.4 TEXT-TO-IMAGE
VAE Decoder: We use the VAE in stabilttyal/ss--ae--t--sss Rombach et al. 2022 and fine-tune the conv-in, up-blocks (0,-1), conv-out, quant-come conv, and parts of the VAE decoder to refine the decoded image quality on the tar- get dataset using the VAE decoder reconstruction loss Kingma & Welling (2033) with the learning rate le 5. 5.
post -uantuuantttoooo
Text Embedding Model: Following Section 5.4, optimize solely the MLP parameters that map pretrained e5 embeddings to the latent space under equation T2 Accordingly, tw,, is implemented as fixed e5 encoder plus a trainable MLP (i.e., PEFT setup).
U-Net Backbone and Training Setup We build the latent-space backbone for Mean Flow and RM- Flow by reusing selected U-Net blocks (~-80M parameters) from pretrained Stable Diffusion Rom- bach et al. (2022). with an additional time embedding for This is effectively transfer learning. Table shows the U-Net configuration and the training setup.
U-Net
block-out-channels down-block types
[320, 640, 1280. 1280] [CrossttnDownBleek2.. CressAttDowwnllek22.. CnessAttDoon Block22,, DownBlock2DD
layers-per-block pilentio-hheudddim
2
cross attention--iiie dim
768
Pre-train
CaMMM
epochs
500
batch size per GPU
16
optimizer Ir schedule
Adam polynomial
Warm up epoch
le-4 (0.9. 0.95)
Adam( 11. (aa) ema decay precision
0.9995
fp16
trainable param
480M
Post-train
CRMFIOW 500
epochs
batch size per GPU
16 Adam polynomial 10
optimizer Ir schedule up epoch
5e-5
Adam( 31. (2)) ema decay precision
(0.9, 0.95)
0.9995
bf16
trainable param A1
[5e-2, 1e-2)
Reg for di. Time sample
Aa plr #
le-4
0.25
Table 11: Training Setup and Backbone Configuration for MeanFlow and RMFlow on COCO Text-to- Image Dataset
18

---

Published as a conference paper at ICLR 2026
B.5
CONTEXT-TO-MOOE -OLLEULL
Context Embedding Model: We implement ou an MLP followed by single EGNN layer. The MLP projects the context vector into the data space, and the EGNN layer further refines these representations. The EGNN configuration matches that used in the Mean Flow/RMFlow backhone.
EGNN Backhone and Training Setup We use the same EGNN Backbone in (Hoogeboom et al. 2022) augmented with a time-embedding module for the additional scalar time variable r. Table I2 shows the training setup and configuration.
EGNN
number layers
acitivation hidden dim
SiLU 256
Pre-train
CeMFs
epochs
1500
batch size GPU
64
optimizer Ir schedule
Adam polynomial
Warm up epoch Ir
10
le-4
Adam (9,, (2)) ema decay precision
(0.9. 0.95)
0,9995
fp32
Post-train
loss epochs
CaMFlow
1500
batch size per GPU
64
optimizer
Adam
schedule
polynomial
Warm epoch
10
1c44
Adam((9) (2))
(0.9. 0.95)
ema decay
0.9995
precision
fp16
[1e-2, 5e-2]
Reg for
Ag plr # 1)
le-4
Time sample
0.5
Table 12: Training Setup and Backbone Configuration for MeanFlow and RMFlow on Context-to-Moleaul Generation QM9 Dataset
B.6 TIME SERIES: DYNAMIC SYSTEM
Trajectory dataset: we use the same dataset decribed in (Huang et al. 2025 Appendix B.1).
Models: (1) We implement the guidance embedding function @w as an MLP that maps the event- guidance vector together with the first three states æ (TT2), ((Ts) into the prior sample; (2) we adopt the UNet architecture from Finzi al. (023). adding an additional time embedding for r.
Training Setup: see Table 13
Train
iteration
batch size
500
optimizer
Adam
Ir
Ic-4
weight decay
0.995
X
Xa
le-4
Table 13: Training Setup for MeanFlow and RMFlow Dynamical System Forecasting Tasks
B.7 ADDITIONAL EXPERIMENTS
In this section, report the comparison of CLIP of our RMFlow, MeanFlow, and some models on the COCO dataset (2017 version dataset splitting, with 5000 images, following [iu et al. 2023b)). See Table 14
19

---

Published as a conference paper at ICLR 2026
Type
params
NFE
Teacher-free (or discrimissatr-frev
FID-Sk
Clip (1)
Diffusion v1.5 Rombach et
Dii
860M
>
0.315
InstaPlow | al (3024h)
OD3:
90KM
23.4
StyleGAN-T ST
PD-SD Meng
MexnFlow RMHoow
ODE Diff
620M 620M
38.5
0.273
27.9
0.291
Table 14: CLIP scores the generated images on the benchmark COC02017 dataset using different models.
B.8 POLICY GRADIENT WITH PHYSICAL FEEDBACK
We define reward r r(gsgt) quantify the molecule stability of the generated sample dgg: QM9 dataset, then define the policy gradient following Black et al. (2023):
VOLRI =
V log Pa ( (oo)r(Zrr)))
and the corresponding loss function
CRL == -E log po (Zigs 1
where
log Po aa)
2 min σ2) Figg (æo + ú0 (moi 0))1 + C,
We then perform reinforcement learning fine-tuning by augmenting the RMFlow objective 12
CRMF+RL (0,w) CRMFIOW (0, + (0)
where 77 is a hyperparameter.
20


Reviews:
[{"review_id": "wRM9FXT9xi", "reviewer": "Reviewer_cCGj", "summary": "The paper proposes RMFlow, a refinement of MeanFlow designed to improve single-step (1-NFE) generation quality across multimodal tasks such as text-to-image, molecule, and time-series generation.", "questions": "If this method were applied beyond text-to-image (T2I) generation to other applications, how would it perform? Would the existing approaches be directly applicable in the same way?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is clearly written and presents a coherent idea within the flow-matching framework.\n- The proposed formulation is lightweight  and the integration of noise injection is computationally efficient.\n- The experiments span diverse domains (synthetic, molecule, time-series, and COCO text-to-image), which demonstrates the model's general applicability.", "weaknesses": "- Performance for each downstream tasks: Although the paper claims near-SOTA performance, the actual COCO FID-30K (18.9) is still substantially higher than recent single-step diffusion models\n- Limitation or failure cases are not discussed in the manuscript.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pmYOv4jCbs", "reviewer": "Reviewer_8n7D", "summary": "This paper identifies a critical failure mode in 1-NFE (single-step) MeanFlow models, where they struggle to capture high-fidelity details in multimodal or structured data distributions. To address this, the authors propose RMFlow, a novel modification that introduces a noise-injection refinement step. This architectural change enables the use of a hybrid loss function that combines the Wasserstein distance objective of MeanFlow with a likelihood-maximization (KL divergence) objective. The authors demonstrate through extensive experiments on synthetic data, molecule generation, time series forecasting, and text-to-image synthesis that RMFlow significantly improves the quality of 1-NFE generation over the MeanFlow baseline, achieving competitive results for teacher-free single-step models.", "questions": "1) Regarding the text-to-image evaluation, could you please justify the decision to only report FID? Could you provide results using a text-alignment metric such as CLIP Score? This would substantially strengthen the claims of multimodal generation quality.\n\n3) Can you provide more insight into the sensitivity of the hyperparameter $\\lambda_{1}$ on the QM9 or COCO tasks? Does the optimal balance between the geometric ( $L_{CMFM}$ ) and likelihood ($\\mathcal{L}_{NLL}$) losses vary significantly across different data modalities?\n\n3) The noise-injection step uses a fixed noise level $\\sigma$. Have you experimented with making this a learned parameter or using a schedule? It seems plausible that the optimal amount of refinement noise might depend on the sample or context.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1) The paper is well-written, the motivations are clear and exemplified through failure cases of previous models, the contribution is well-theorized and explained. Experiments shows that the proposed approach is able to solve the failure cases and issues of previous approahes. \n\n2) The core idea of RMFlow is both simple and theoretically sound. Though the conceptual two-stage process (coarse transport + refinement), it is implemented in a single, efficient 1-NFE step. The key contribution is the rigorous connection (Theorem 4.1) between the noise-injection architecture and the ability to optimize for likelihood via an NLL loss term. This synergistic co-design of architecture and loss function is a significant strength.\n\n3) The authors validate their method across four different and challenging domains. The results are consistently positive, showing marked improvement over the MeanFlow baseline. The success in the context-to-molecule task, where RMFlow dramatically increases molecule stability, is a particularly powerful demonstration of the method's practical utility for structured data with hard constraints.", "weaknesses": "1) This is the most significant weakness of the paper. For the text-to-image task (Table 6), the sole reported metric is Fréchet Inception Distance (FID). As is now widely discussed in the community, FID is a flawed metric with several known issues:\n- It relies on an outdated InceptionV3 backbone trained on ImageNet, which is a poor feature extractor for the rich, diverse content produced by modern generative models.\n- It is known to correlate poorly with human perception of image quality.\n\nMost importantly for a text-to-image model, FID only measures the realism of the generated distribution and is completely agnostic to the text conditioning. The paper makes claims about \"multimodal generation\" but provides no quantitative evaluation of the text-image alignment. The lack of standard alignment metrics like CLIP Score or human preference scores (like HPSv2) is a major omission and makes it difficult to fully assess the model's capabilities on this task.\n\n2) The paper's performance depends on a key hyperparameter, $\\lambda_{1}$, which balances the Wasserstein and KL-divergence objectives. However, the ablation study for this hyperparameter (Table 7) is only performed on the simple 1D Gaussian mixture task. It is unclear how sensitive the model is to this trade-off on more complex, high-dimensional tasks like molecule or image generation, which would be crucial information for practitioners.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "auWFjOUiRJ", "reviewer": "Reviewer_T3eq", "summary": "This paper introduces RMFlow, a novel flow-based model designed for single function evaluation. The approach builds on Flow Matching, a generative modeling framework that maps two distributions by modeling the flow process using an ordinary differential equation (ODE). However, generating samples from the target distribution typically requires multiple function evaluations. To address this computational cost, several approaches have been proposed to reduce the number of function evaluations (NFEs). One such method, MeanFlow, optimizes an averaged velocity field that allows for single function evaluation, but it often yields suboptimal results. The authors hypothesize that the absence of KL divergence optimization may contribute to the reduced performance. To address this, they propose incorporating a noise injection step and an additional negative log-likelihood (NLL) loss during training, leading to their proposed model, RMFlow. The authors also provide a method for incorporating conditioning, which enables multimodality generation. Furthermore, the authors conducted an extensive empirical evaluation across multiple tasks, including synthetic datasets, context-to-molecule generation, and text-to-image generation.", "questions": "* Line 277 -  “During fine-tuning, we further strengthen training by integrating 1-NFE sampling with a policy-gradient objective that incorporates physical feedback on sample quality...”\nThis statement is a bit unclear. Could you clarify whether any additional reinforcement learning (RL)-based training was performed beyond directly optimizing equation (10)?\nIf so, did the other baselines, particularly the MeanFlow models, also undergo this additional training step?\n* Line 273: “...introduces additional gradient pathways...” -  Could the authors please elaborate on this point and provide more details or clarification?\n* Line 240: $\\varepsilon_1,varepsilon_2$ are not used in this section.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "* The paper is well written, the method is clear and easy to understand.\n* The results demonstrate strong empirical performance.\n* The authors present a clear and well-founded motivation for the design of their mode", "weaknesses": "* Comparison to other approaches: In the introduction, the authors mention several alternative methods for efficient flow modeling—such as Consistency Models (CM), distillation, and local Flow Matching (FM). However, it remains unclear whether any quantitative comparisons were conducted against these methods.\n* See questions", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)