Simple ReFlow: Improved Techniques for Fast Flow Models
Authors: Beomsu Kim
Yu-Guan Hsieh
Michal Klein
marco cuturi
Jong Chul Ye
Bahjat Kawar
James Thornton
Venue: ICLR 2025 Poster (2025)
Keywords: flow matching
diffusion models
reflow
optimal transport
generative modeling
Primary area: generative models
TL;DR: Improvements to the ReFlow procedure for faster sampling of flow models by iterative straightening without perceptual losses.

Abstract: Diffusion and flow-matching models achieve remarkable generative performance but at the cost of many neural function evaluations (NFE), which slows inference and limits applicability to time-critical tasks. The ReFlow procedure can accelerate sampling by straightening generation trajectories. But it is an iterative procedure, typically requiring training on simulated data, and results in reduced sample quality. To mitigate sample deterioration, we examine the design space of ReFlow and highlight potential pitfalls in prior heuristic practices. We then propose seven improvements for training dynamics, learning and inference, which are verified with thorough ablation studies on CIFAR10 $32 \times 32$, AFHQv2 $64 \times 64$, and FFHQ $64 \times 64$. Combining all our techniques, we achieve state-of-the-art FID scores (without / with guidance, resp.) for fast generation via neural ODEs: $2.23$ / $1.98$ on CIFAR10, $2.30$ / $1.91$ on AFHQv2, $2.84$ / $2.67$ on FFHQ, and $3.49$ / $1.74$ on ImageNet-64, all with merely $9$ NFEs.

Published as a conference paper at ICLR 2025
SIMPLE REFLOW: IMPROVED TECHNIQUES FOR FAST FLOW MODELS
Beomsu Kim Apple and KAIST
Yu-Guan Hsieh Apple
Michal Klein
Apple
Marco Cuturi Apple
Jong Chul Ye KAIST
Bahjat Kawar
James Thornton
Apple
Apple
ABSTRACT
Diffusion and flow-matching models achieve remarkable generative performance but at the cost of many sampling steps, this slows inference and limits applica- bility to time-critical tasks. The ReFlow procedure can accelerate sampling by straightening generation trajectories. However, ReFlow is an iterative procedure, typically requiring training on simulated data, and results in reduced sample qual- ity. To mitigate sample deterioration, we examine the design space of ReFlow and highlight potential pitfalls prior heuristic practices. We then propose seven im- provements for training dynamics, learning and inference, which are verified with thorough ablation studies on CIFARIO 32 x 32, AFHQv2 64 64, and FFHQ 64 64. Combining all techniques, (without with guidance, resp.) for fast generation via neural ODEs: 2.23 / 1.98 CIFAR1O, 2.30 1.91 AFHQv2, 2.84/ / 2.67 FFHQ, and 3.49 1.74 on ImageNet-64, all with merely 9 neural function evaluations.
achieve state-of-the- FID scores
INTRODUCTION
The diffusion model (DMs) paradigm the landscape of generative modelling of perceptual data, benefitting from scalability, stability and remarkable performance in diverse set of tasks ranging from unconditional generation (Dhariwal Nichol, 2021) to conditional generation such as image restoration (Chung et al., 2023), editing (Meng al., 2022), translation (Su et al., 2023), and text-to-image generation (Rombach et al., 2022). However, to generate samples, DMs require numerically integrating diffcrential equation using tens hundreds of neural function evaluations (NFEs) (Song al., 202 1 b;a). Naively reduc- ing the NFE increases discretization error, causing sample quality to worsen. This has sparked wide interest in accelerating diffusion sampling (Song al., 2021a; la; Lu et al., 2022: Zhang & Chen, 2023; Kim & Ye, 2023; Salimans & Ho, 2022; Song et al., 2023).
et al., 2015; Ho et al., 2020) has changed
closely related class of generative model sharing similar training and sampling procedures and enjoying similar performance to diffusion models. Indeed, FM and DMs coincide for particular choice of forward process (Kingma & Gao, 2023), and are also related to stochastic interpolants (Albergo & Vanden-Eijnden, 2023; Albergo al., 2023). Whereas diffusion models relate to entropically regularized transport (Bortoli al., 2021; Shi et al., 2023; Peluchetti, 2023), key property of flow matching models is their connection to non-regularized optimal transport, and hence deterministic, straight trajectories (Liu, 2022).
Flow matching models (FMs) (Lipman et al., 2023)
While there exist plethora of acceleration techniques, one promising, yet less explored avenue is ReFlow (Liu et al., 2022; Liu, 2022), also known as Iterative Markovian Fitting (IMF) (Shi et al., 2023). ReFlow straightens ODE trajectories through flow-mat fow-matching between marginal distributions coupled by a previously trained flow ODE, rather than using an independent coupling. Theoretically, with an infinite number of ReFlow updates, the resulting learned ODE should be straight, which enables perfect translation between the marginals with a single function evaluation (Liu et al., 2022).
In practice however, ReFlow results in drop in sample quality (Liu al.. 2022; 2024). To address this problem. recent works on sampling acceleration via ReFlow opt to use heuristic tricks such perceptual losses that only loosely adhere the underlying theory (Lee et al., 2024; Zhu al.,

---

Published as a conference paper at ICLR 2025
2024). Consequently, it is unclear whether the marginals are still preserved after ReFlow. This is problematic, as exact inversion and tractable likelihood calculation require access to a valid proba- bility flow ODE between the marginals. Moreover, these two functions are critical to downstream applications such as zero-shot classification (Li et al., 2023) etc.
The goal of this work is study and mitigate the performance drop after ReFlow without violating the theoretical setup. Although technical in nature, we call our method simple, as, similar to simple diffusion (Hoogeboom et al., 2023), it does not rely on latent-encoders, perceptual losses, or premet- rics, whose effect on the learned marginals is poorly understood. To this end, we first disentangle the components of ReFlow. Next, we examine the pitfalls of previous practices. Finally, we propose enhancements within the theoretical bounds, and verify them through rigorous ablation studies.
Our contributions are summarized follows.
We generalize and categorize the design choices of ReFlow (Section 3.1). We gener- alize the ReFlow training loss and categorize the design choices of ReFlow into three key groups: training dynamics, learning, and inference. Within each group. we discuss previ- ous practices, highlight their potential pitfalls, and propose improved techniques.
We analyze each improvement via ablations (Sections 3.2, 3.3, and 3.4). For each pro- posed improvement, we verify its effect on sample quality via extensive ablations on three datasets: CIFARIO 32 x 32 (Krizhevsky, 2009), FFHQ 64 x 64 (Karras et al., 2019), and AFHQv2 64 x 64 (Choi et al., 2020). We demonstrate that our techniques are robust, and they offer consistent gains in FID scores (Heusel et al., 2017)
all three datasets.
We achieve state-of-the-art results (Section 4), With all our improvements, we set state- of-the-art FIDs for fast generation via neural ODEs, without perceptual losses or premet- rics. Our best models achieve 2.23 FID on CIFAR 10, 2.30 FID on AFHQv2, 2.84 FID on FFHQ and 3.49 FID on ImageNet-64, all with merely 9 NFEs. In particular, our models outperform the latest fast neural ODEs such as curvature minimization (Lee et al., 2023) and minibatch OT flow matching (Pooladian et al., 2023). We are also able further en- hance the perceptual quality of samples via guidance, setting 1.98 FID on CIFAR10. 1.91 FID on AFHQv2, 2.67 FID on FFHQ, and 1.74 FID on ImageNet-64, also with 9 NFEs.
BACKGROUND
Let Io and Pi be two data distributions on Edd Rectified Flow (RF) (Liu et al., 2022: Liu, 2022) is an algorithm which learns straight ordinary differential equations (ODEs) between Po and P1 by iterating procedure called ReFlow. Below, we describe ReFlow, and explain how it can be applied to diffusion probability flow ODEs to learn fast generative flow models.
2.1 FLOW MATCHING AND REFLOW
Let us first define the flow matching (FM) loss and its equivalent formulation a denoising problem
(0: Q01)= = =)) (0) [<MSE(#; Io. v,(x,,,)) t)]]
(1)
= ((mm)a))~0 (0 -2 ( Do t
(2)
where vo Rdd 1) > Rdd is velocity parameterized by 8, EMSE (z. y) := || -
(1
t)æo + t. and Qui is a coupling, i.e., joint distribution, of Po and P1. and Do t) == xz two
a denoiser that is optimized to recover the original data wo given a corrupted observation æ, and time l as inputs. According to the FM theory, a velocity which minimizes Eq. (1) or denoiser which minimizes Eq. (2) can translate samples from P. to P1 i E {0, 1}. by solving the ODE
dx, = (z t) dt = t 'xx - Do t)) dt, from / to 1- (Lipman al., 2023). Let us call the denoiser which minimizes Eq. (2) w.r.t. the independent coupling Qo P. @ P1 as D/..
t E 10,
(3)
For n 2 1, ReFlow minimizes Eq. (2) with coupling induced by D to obtain D whose ODE
has lower transport cost. Specifically, observe that Eq. (3) with D induces coupling
cP' (z) )6(x0 - solve( 1 Dg. 1,))) dQi (wo, 21) = [dPPo (Po( )0(x solve(no. D 0, 1))
(4)
2

---

Published as a conference paper at ICLR 2025
to t with initial point æ. Concretely, given D2 ~ P, for i E {0, 1), we sample ø1-; ~ Q1-qi( (.)) by integrating Eq. (3) from t to 1 i starting from z,. The two expressions in Eq. (4) are equivalent, since an ODE defines bijective map between initial and terminal points.
where 8 is the Dirac delta, and solve(a; Dj to, 12) solves Eq. (3) from time t =
RF guarantees that if DØ+1 is minimizer of (O; Q))), the ODE with Dz+t converges perfectly straight ODE as n + oc. If an ODE has perfectly straight trajectories, it is possible to translate between the marginals with single Euler step, e.g., an Do(æ1, 1) when translating from / = 1 to (see Section 3 of Liu et al., 2022).
2.2 REFLOW WITH DIFFUSION PROBABILITY FLOW ODES
Given distribution Po on Rd and Gaussian perturbation kernel dQ2, (Woll0)) = N Boolyo. I2
DMs solve the denoising score matching (DSM) (Vincent, 2011) problems
DDSM (0) Es Eg Bv-Booo(-laun [EME (yo, Fo a))]
(5)
learn a denoiser F Rd (0. 00) + Rd Fg then defines probability flow ODE between Po and Pu * NO 8² I) ((, x² I) for a large Ō:
dy = - oy Fo σ)) do, a E (0, 00).
(6)
When F1 is standard normal, Eq. (3) with D3 and Eq. (6) are equivalent, as Eq. (3) with the change of variables (yo, o) == ( I- and Eq. (6) are identical (Lee et al., 2024). It follows that we can straighten diffusion probability flow ODE trajectories via ReFlow. Specifically, with the coupling
dQbi (Io, 21) = (211 (21)) solve( Fo, - 0)) d dPPo(zo)((e)) (1 - -)- solve( Fo. 0, 1-)
where t xz and solve(y, Fo, To. 1) solves Eq. (6) from σ do to with initial point y. we can minimize Eq. (1) to learn D and so on. Because optimizing Eq. (1) is often expensive, a typical procedure is to perform one ReFlow step with Eq. (7) to get D2 and distill Eq. (3) trajectories into a student model for one-step generation (Liu et al., 2022; Zhu et al., 2024; Liu et al., 2024).
3 IMPROVED TECHNIQUES FOR REFLOW
We now investigate the design space of ReFlow and propose improvements. Specifically, in Section
3.1, I, we generalize the FM loss and identify the components that constitute ReFlow. The components are organized into three groups training dynamics, learning. and inference. In Sections 3.2, 3.3, and 3.4, we investigate the pitfalls of previous practices and propose improved techniques in each group. To show that our improvements are robust, provide rigorous ablation studies CIFARIO 32 32, AFHQv2 64 x 64, and FFHQ 64 x 64. We find that ReFlow training and sampling are very different from those of DMs, and generally require distinct techniques for optimal performance.
3.1 THE DESIGN SPACE OF REFLOW
Generalizing weight and time distribution. Let the joint distribution of (æo, 2. f, t, be given by Co, 2i ~ dQoi. ~ T, and a; = (1 - l)a0 + (z), Then Eq. (2) can also be expressed as
LFM(A: Qo1) = (e Ex [w](t) (FM (;; Qo1, Zie, ()], where (F (f; Qo1, æe, 1) == Bee (CMSE( (o(a,, t))]
(8)
(9)
and wi(t) = 1-2 This shows that the FM loss is separable w.r.t. (æ,, 1), and the optimal denoising function is given by the posterior mean (Robbins, 1956): D P' may replace w(t) with general weight w(x,, t) and use general time distribution T
Eezo-Goje( [xo]] Hence,
LFM (0; C01) EarrEa,, [w(x), t) (F: Qo1, æx, f)]
(10)
This is minimized under the same condition, given that w(x, f 0 and T is supported on (0, 1).
3

---

Published as a conference paper at ICLR 2025
RF
RF++
Baseline
Simple ReFlow (Ours)
Train Dynamics (Sec. 3.2)
Weight (r(
1/12
1, 1ft
1/ sigg SglCCra(0: 20,
Time distribution dTl(f) xc
cosh(4(t 0.5))
cosh(4{t
0.5))
10
Loss function
FMSE, LPIPS
Pseudo-Hulee LPIPS
MMSE
Learning (Sec. 3.3)
Dy initialization with DM
dropout probability
0.15
Equal EDM
0.15 Backward
< 0.15 Forward, Projection
Sampling from Qoi
Backward
Backward
Inference (Sec. 3.4)
ODE Solver
Euler
Buler, Heun Unitorm
Heun Uniform
DPM-Solver Sigmoid x 20
Discretization of (0,
Uniform
(Liu ct al., 2022)
Reference
(Lee et al., 2024)
(Zhu et al., 2024)
Table I: Comparison of practices for optimizing the ReFlow loss Eq. (13) and solving the ODE Eq. (3). sg means stop gradient and HPF denotes high-pass filter. Baseline is the combination of most recent techniques which do not violate the flow matching theory.
Generalizing the loss function. We also consider using general loss functions in ReFlow, i.e.,
(F: Qo1, xt, t) (. [e(
(2t, t))]
(11)
for general é IRd x Rd IR. It is difficult to precisely characterize the class of that preserves the minimizers of Eg. (1), and popular losses such as LPIPS (Kendall et al., 2018) and pseudo-Huber (PH) (Song & Dhariwal, 2024) lack this guarantee. However, fMSE has been observed be sub- optimal compared to, e.g., LPIPS and PH for training fast models (Lee et al., 2024). To mitigate this trade-off between theoretical correctness and practicality, we consider a wider class of losses
l y) = [|)()) <ly))
(12)
for invertible linear maps Rd + Rd. This again ensures that the loss is minimized when Do outputs the posterior mean, and {ME special case of this loss with the identity map o I.
Generalized FM loss. Combining the two generalizations, we have our generalized FM loss
(G Qu1) FE-TEE (x< t) LGFM (6; Q01, Tt, t)], (11,,) (2o Do(a;, ())]
(13)
where LorM(0; Qo1, =
(14)
The following proposition ensures its theoretical correctness. Proof is deferred to Appendix E.I. Proposition Let (vr, l), (l) positive, and & be an invertible linear map. Then, 0 minimizes Eq. (13) if and only it minimizes Eq. (1).
3.1.1 TRAINING DYNAMICS, LEARNING, AND INFERENCE
We now observe that there are seven components that constitute ReFlow: time distribution T. train- ing dataset (empirical realization of Qu1), weight w f), loss function (g. denoiser De, and ODE solver and discretization schedule for solving Eq. (3). We categorize them into three groups below.
Training dynamics influence the path that the model takes towards the minimizers of Eq. (13) during training. Although the solution to which the model converges may change if dynamics changes, training dynamics do not impact the solution set itself. Weight function un [at, f), time distribution T, and loss function l belong here. Learning influence the solution set of Eq. (13) by constraining the hypothesis class or by changing the training dataset. Parameterization of Do and how we sample from Qoi belong here. Finally, inference influence generation or inversion of samples given a trained model. ODE solver and time discretization of the unit interval belong here.
In Tab. describe recent ReFlow practices within our framework. Bascline is the collection of most recent ReFlow techniques which do not violate FM theory (Lipman et al., 2023). We will build up improvements on this baseline setting in the subsequent sections.
3.2 IMPROVING TRAINING DYNAMICS
Evaluation protucol. To evaluate a training setting, we perform a single ReFlow step. Unless writ- ten otherwise, initialize ReFlow denoisers with pre-trained EDM (Karras et al., 2022) denoisers
4

---

Published as a conference paper at ICLR 2025
010 (00.) miitn, cl0:at..
DM
w((2,,
CIFARIO
AFHQv2
FFHQ
20
HeFFow
2.83
2.87
4.28
2.77
2.76
4.01
2.76
2.74
4.04
107
(o (.
2.78
1/Ez, [sg] Cara (e; Qoi ..
2.74
2.79
3.83
1/ sg Curee Qot. (..))
2.61
2.74
3.83
10P
1-
0.0
1122
0.4
0..6
0.8
1.0
Table 2:
Comparison of various wo(x;, (), combined with baseline T, and learning choices. Best num- bers are bolded, and second best are underlined.
Figure 1: Min., avg., max. relative losses after training on CIFARIO.
and optimize Eq. (13) with Qoi = Qbi Eq. (7) for 200k iterations. We sample 1M pairs from Qu by solving Eq. (6) from t= 1 to 0 with EDM models and use them for training. We measure the performance of an optimized model by computing the FID (Heusel et al., 2017) between 50k generated images and all available dataset images. Samples generated by solving Eq. (3) with the Heun solver (Ascher & Petzold, 1998) with 9 NFEs, and we use the sigmoid discretization instead of the baseline uniform discretization for reasons discussed in Appendix F.1. We report the minimum FID out of three random generation trials. See Appendix for a complete description.
3.2.1 Loss NORMALIZATION
Previous practice. There is little study on suitable loss weights for ReFlow training. For instance, Lee et al. (2024) use w(x,, t) 1, and such choices can be detrimental to training, the weighted loss w(x [ t) (G Qu1, 2t, t) without proper modulation by w w(x,, t) can have vastly different scales w.r.t. t, leading to slow and unstable model convergence. Typically, the loss vanishes t since at + vo. and to counteract this, previous works have suggested
1/1 w( (02 + 0.55)/(0.50)) 0 I-t
for ReFlow (Lee et al., 2024),
for DMs (Karras et al., 2022),
1/Berqq [ (0: Co1 Re, t)]]
for DMs (Karras et al., 2023b).
where sg[-] is graddgrrdntt..... The FM weight 1/t in Eq. (2) naturally emphasizes x 0 as well.
However, claim that such weights constant w.r.t. can be sub-optimal for ReFlow, as ReFlow loss scales can vary greatly w.r.t HG even for fixed For instance, the following proposition shows that, at initialization, relative loss for DM at be arbitrarily large. Proof is deferred to Appendix E.2.
is constant whereas relative loss for ReFlow can
Proposition 2. Assume output layer zero initialization for the DM denoiser and DM initialization for the ReFlow denoiser. Then maximum relative losses for DM and ReFlow at = 1 are
maxæ, (O; Po @ 11, 1), mink, (DM % 11 σι, 1 1)
maNe, LGFM (0; Cn 1 1)/ mille, CGFM (0; Ch #1, 1) ma |xo | ||/6 moll
resp., where no Ex,-P, [xo]]. and mill,,. maXæ, is taken w.r.t. æ; in the support of PP..
In fact, in Fig. 1, we observe that ReFlow loss varies greatly w.r.t. It after training as well. In contrast to DM training loss whose minimum and maximum values differ by a factor of at most 20, minimum ReFlow loss least 100 smaller than the maximum loss for all > 0.2.
Our improvement. Multi-task learning interpretation of loss normalization (Zhang et al., 2018; Karras et al., 2023b) (see Appendix G) motivates a simple improvement by using
w w(x,, 1/ sg| sg[CoFM(O: Qo1 xz, xt,,)]. t)].
(15)
Similar to Karras et al. (2023b), we keep track of the loss values during training with a small neural net that is optimized alongside Dg using the parameterization w(x,, t) exp( (x
Ablations. Tab. 2, we compare our weight with all aforementioned weights. As expected, the uniform weight w(x;, = 1 has the worst performance, as it is unable to account for vanishing loss as - 0. We get noticeable FID gain by using weights such as 1/t or which place a larger emphasis on O. Our weight, which accounts for loss variance w.r.t. both & and 2., yields the best FID across all three datasets. The gap between baselines and our weight is especially large on 10.
3.2.2 TIME DISTRIBUTION
5

---

Published as a conference paper at ICLR 2025
1=
1=1
(=0
(=1
Geoerated
Gesrnsted
t)
(b) FM with ÉMSE
FM with PH
Figure 3: Comparison of flow matching (FM) with ÉMSE and Pseudo Huber (PH) losses.
Previous practice. Liu al. (2022) and Zhu al. (2024) use a uniform distribution on (0, On the other hand, Lee al. (2024) notice better performance with time distribution whose density is proportional to shifted hy- perbolic cosine funetion, i.e., (t) x cosh(4(t 0.5)). which has peaks at = 0 and 1. The rationale behind using such a distribution is that, as Eq. (3) converges straight ODE with ReFlows, the denoiser needs to di- rectly predict data from noise at / As and vice versa at t = O, is beneficial emphasize those regions via T.
4006
LogNoruul
1 0.0
0.4
Figure 2: Time distribution densities.
Our improvement. The peak at t = of the baseline cosh time density compensates for vanishing loss as t 0, but as we normalize the loss with our weight, this peak now longer necessary. Thus, with density proportional to the increasing exponential, i.., dTT(I) a' for 1.
OT &
CIFARI0
AFHQv2
FFHQ
cosh ff 0.5))
2.61
2.74
3.83
loguormal
3.28
3.20
4,48
distribution
uniform (19)
2.65
2.70
3.85
10"
2.62
2.69
3.77
100'
3,93
Ablations. We compare the performance of the expo- Table 3: Comparison of various T, com- nential distribution with E {1, 10, 100}, where = 1 bined with our w(x,, t) and baseline corresponds to the uniform distribution. We also compare and learning choices.
with the lognormal distribution, which has been observed be effective for training DMs and CMs (Karras et al., 2022; 2023b; Song & Dhariwal, 2024) Fig. 2 displays time densities, and Tab. 3 shows training results. We first note that an emphasis on degradation with the lognormal distribution. We then observe that 10', which closely resembles the cosh distribution, but without the peak at O, has consistently good performance, while suffering from slight loss CIFARIO. Other choices such as 1' or 100' are either too flat sharp, yielding worse FID. Hence, we propose to take dT((t) x 10
is necessary, as evidenced by severe FID
3.2.3 LOSs FUNCTION
Previous practice. To accelerate convergence and mitigate sample quality degradation during Re- Flow, previous works have employed heuristic losses in Eq. (2) such as LPIPS and
(1/t) (||== v113 (ct) 1/2-c Pseudo-Hue Huberr (PH) (Lee et al., 2024).
((x, v) LPIPS x. y) + (1- -
LPIPS+PH (Lee et al., 2024).
While such losses perform better in practice than EMSE in terms of FID, they do not ensure Eq, (1) is minimized at optimality, and so lose the theoretical guarantees of FM. We demonstrate this below.
Fig. 3 compares FM with (MS and PH, where P1 is unit Gaussian and PP is a mixture of Gaussians. As shown in Fig. 3a, the two models learn distinct vector fields, so indeed induces different ODE trajectories. While the model trained with ÉMS translates between Po and IP, accurately, the model trained with PH generates incorrect distributions, e.g., the model density is not isotropic at 1, and modes are biased at et al. (2024)) to justify heuristic losses, we show that proper choice of the invertible linear map ø in Eq. (12) can still offer non-trivial performance gains while adhering to FM theory.
So, instead of relying empirical arguments (e.g., Section 4 in Lee
Our improvement. Previous works have observed high- fighffreyeenc features are crucial to diffusion- based modeling of image datasets (Kadkhodaie al., 2024; Zhang & Hooi, 2023; Yang al., 2023).

---

Published as a conference paper at ICLR 2025
Loss
CIFARIO
AFHQv2
FFHQ
CHEARIO
P10
ÉXISE
PH
2.62
2.50
2.69
2.71
3.77
3.25
0.0
FFBQ
ATTiQ22
PHI
LPIPS
2.81
2.65
4.02
3.1
- 10BPP
LPIPS-PH
2.63
2.72
3.79
(140. HPF
2.63
2.62
3.76
-0.1
2.58
2.55
3.69
2.7
HPP
3.20
2.79
4.43
01
0.1
10
100
1000
Table 4: Comparison of various l combined with our wxt,,) t) and T.
Figure 4: ablation.
Figure 5: CIFARIO training.
Moreover, since we initialize Do with pre-trained DM, we assert that the model already has good representation of low-frequency visual features. Hence, to accelerate the learning of highefrequency features, we propose calculating the difference denoiser output D3 Da (Ztz t) and clean data æo after
passing them through high-pass filter (HPF) using the linear map, i.e., use lo in Eq. (12) with
= I++ HPF
(16)
where À 0 controls the emphasis high-frequency features. The identity matrix in Eq. (16) is necessary to ensure that é is invertible, so Eq. (1) is minimized at optimality per Prop. 1. We also remark that using &s can be interpreted as preconditioning the gradient. Specifically, since
VD La Do t)) ot D. (2o. Do (æte f))),
(17)
using ls in place of ÉMSE is equivalent to scaling the original FM loss gradient along the eigen- vectors of o Φ by the corresponding eigenvalues. For instance, ls with Eq. (16) amplifies gradient magnitudes along high-frequency features by (A + 1), and leaves gradient magnitudes along low- frequeney features unchanged. This perspective provides another justification for using Ég. since using an appropriate preconditioning matrix can accelerate convergence (Kingma & Ba, 2015).
Ablations. Tab, 4 compares loss 4 HPF for X E {0.1, 10, 1000} with ÉMSE and the heuristic losses. If À is too small, f++A HPF has little improvement compared to ÊMSE, whereas if X is too large, & becomes nearly singular, leading to severe drop in the FID. Our loss with À 10 provides consistent improvement over ÊMSE doing even better than PH and LPIPS. Indeed, in Fig. 4 which visualizes FID change w.r.t. ÉMSE for various values of X, we observe performance across all datasets. Morever, CIFAR10 learning curves in Fig. 5 verify that 4 HPF enjoys fast convergence compared all other losses.
10 provides the optimal
3.3
IMPROVING LEARNING
3.3.1 MODEL. DROPOUT
Previous practice. Similar to simple diffusion (Hoogeboom et al., 2023), find dropout to be highly impactful. There is little study on the impact of dropout in denoiser UNets for ReFlow. Dropout rates in ReFlow denoiser UNets are usually set 0.15 (Liu et al., 2022; Zhu al., 2024), or equal to the dropout rates of DMs that are used to initialize ReFlow denoisers (Lee et al., 2024). For the EDM networks, dropout rates are 0.13 on CIFARIO, 0.25 on AFHQv2, and 0.05 on FFHQ.
Our improvement. We observe that learning a straight ODE is harder task than learning the the optimal DM denoiser only needs to
diffusion probability flow ODE. For instance, t predict the data mean Ex,,e,, for any input æ₁ ~ F1, but denoiser for perfectly straight ODE has to directly map P1 samples P. samples. This means we need larger Lipschitz constant for the ReFlow denoiser (Salmona et al., 2022) (see Appendix G for further discussion), so we use smaller dropout rates during ReFlow training favor of larger UNet capacity over regularization.
Ablations. To verify that smaller dropout rates are beneficial, we return to the baseline training setting (Tab. I), and run grid search over dropout probability p E [0, 0.15]. Fig. 6, which shows FID change w.r.t. baseline p 0.15, we find that smaller is always beneficial. In fact, optimal p are even smaller than those used to train EDM denoisers, despite using the same architecture written in row BSL+OP of Tab. 6. We
(Tab. 5). FIDs after applying optimal dropout to baseline also observe in row OYN+OP that optimal dropout rates can be combined with improved dynamics to further enhance performance without additional grid search over p.

---

Published as a conference paper at ICLR 2025
CIFARIO
AFHQQ2
FFHQ
AFDD
Baseline (BSL)
2.83
2.87
4.28
Dynamics (DYN)
-025
BSL + w(2; f)
2.61
2. 2.v0...
3.83vo.45
-0.0
BSL + w(22; 1) +T
2.62c0.21
2.09v0.18
3.77v0.51
CIFARJO
BSL 1r(22e
2.58c0.25
2.55
3.69
-0.75
APDQQ2
FIIQ
Learning (LRN)
BSL Optimal (OP)
(OP)
2.63.o.20
2.67 vo.20
3.60va.6s
0.0
0.05
0.06
0.09
0.12
0.36
BSI. + OP + Forward
2.57
2.63
3.60vo.68
Figure 6: Dropout p ablation.
BSL + OP Projected
2.57
2. 2.62
3.58vo.70
DYN & LRN
DYN + OP
2 2.43
2.53
3.17v1.11
CIFAR10
AFHQ²2
FFHQ
DYN + OP Forward
2.3 2.38 20.05
2.4470.43
3.14
DYN + OP Projected
2.38
2.47
3.13p1.1
RF Ours
0.15
0.15
0.35
0.09
0.00
0.03
Table 6: Summary of our training improvements. Sub-
scripts denote FID improvement w.rt. baseline. Evalu- ated with sigmoid discretization (see Append. F.I).
Table 5: Dropout p in each setting.
CIPAR1O
AFHIQ2
FFIQ
0.05
0.0
-0.05
-0.1
0.0
0.4
0.6
0.8
0.0
0.8
1.0
0.0
0.2
0.4
Figure 7: p ablation. Solid and dotted lines show results w/o and with improved dynamics, resp.
3.3.2 TRAINING COUPLING
Previous practice. common practice is to generate large number of pairs from Qbi by solving the diffusion probability flow ODE Eq. (6) backwards, i.e., from noise to data, and use the generated set as an empirical approximation of Qhi throughout training (Liu et al., 2022; Lee et al., 2024; Zhu al., 2024; Liu et al., 2024). However, the of generated 2o is only an approximation of the true marginal Fn. so naively training with generated data will accumulate error on the marginal at / = 0, as discussed by Alemohammad et al. (2024) (see Appendix G for further discussion).
Our improvement - forward pairs. To mitigate error accumulation at t = O, we incorporate pairs generated by solving the diffusion probability flow ODE forwards, starting from data, coined forward pairs. We assert forward pairs can be helpful, as æy are exactly data points.
To use forward pairs, we first invert the training sets for each dataset, which yields additional 50k pairs for CIFAR1O, 13.5k pairs for AFHQv2, and 70k pairs for FFHQ. Due to the small number of forward pairs, we use them in combination with backward pairs, and prevent forward pairs from being ignored due to the large number of backward pairs, we sample forward pairs with probability p and backward pairs with probability 1 at each step the optimization.
Our improvement projected pairs. We also propose projecting the coupling Q! to II P1), set of joint distributions with marginals PPo and P,, by solving the optimization problem
@+ = arg min W,, (T' Qh)) s.t. For E I1(
(18)
where W, is the p- Wasserstein distance (Villani, 2009), and using the projected coupling Qh in place of the original during training. Intuitively, this procedure can be understood fine-tuning the generated marginals to adhere to the true marginals without losing the coupling information in Q/1- We do not mix projected pairs with any other pairs. The full procedure is described in Appendix D.
Ablations. In Fig. 7, we see that it is always beneficial to use forward pairs, as long is p is not too high, e.g., P 0.5. Otherwise, the model starts overfitting to the forward pairs. Interestingly, on FFHQ, using forward pairs without improved training dynamics has no improvement in the FID, implying that improved dynamics may be necessary make the best out of the rich information contained in the forward pairs. In rows BSL+OP+Projeet and DYN+OP+Projected of Tab. 6, we observe that projected pairs also offer improvements in the FID score across all three datasets.
8

---

Published as a conference paper at ICLR 2025
CIFARIO
AFHQY2
FFHQ
Irlo
CIRARIC
Unif.
2.36
2.34vo.5a
2.97v1a1
3.2
APRIQV2
2,80
3. 3.01gn.74
6. 88
3.0
Ours
10
2.31
2.31
2.ST
2.5
20
2.25
2.30
2.84
30
2. 2.15pu.35
2.78cu.0
3.32mo.6
10
COM
Signovi
2.2
Table 7: Various discretizations applied to our best models and DPM-Solver with r 0.4.
16-5
0.0
1005
10.5
0.05
2.0
0
10
15
20
25
30
Figure 8: Truncation error.
Figure 9: K aablation.
3.4 IMPROVING INFERENCE
Previous practice. To generate data after ReFlow, previous works often use an uniform discretization {t; = i/N : 0. O..... NN) of [0, along with the Euler Heun integrate Eq. (3) from t = to 0 (Liu et al., 2022; 2024; Lee et al., 2024; Zhu et al., 2024).
Usilior
0.8
Signold
0.6
Our improvement. As ReFlow ODE converges we assert that high-curvature regions in ODE paths now occur near t {0, 13. While the previously proposed EDM schedule
straight ODE,
0.2
0.0
å 1/d = 0, === .. where o, (a 1/d + N (
6/N
=
0.25
0.5
0.75
1.0
for solving diffusion probability flow ODEs emphasizes E (0. Figure 10: Discretizations.
we note that it does not perform better than the uniform discretization, as shown in Tab. 7. Similar to Lin et al. (2024), we speculate that this is because iN < 1. Specifically, va( (t) # (2 , 1) since &N * 1, but the integration of the ODE is done with va(21, (x) in place of vs 1), leading to erroneous ODE trajectories.
Instead of tuning {amin, maax: d} to address this problem, propose simple sigmoid schedule
ft. (sig(k(} - 0.5)) sig( /)ig()) sig( 5)) : N,....,
(19)
with one parameter which controls the concentration of 4; at 1. {0,1} Here, sig is the sigmoid function. As ベ - 0, {tr] converges to the uniform discretization, and as x + cc, all t, with i < N/2 will converge to 0, and all t with i N/2 will converge to 1.
To solve the ODE Eq. (3), we consider DPM-Solver (Lu et al., 2022) with update rule
3t, f- Btst + (t: ( u ( s1+1) + (1 Σ/ (T tt+1))
(20)
where 8;-1 = + 4 and E (0, 1]. We recover the second order Heun update (the baseline solver) with 1, but we assert that we can obtain better performance by tuning r.
I-r
Ablations. In Tab. 7, we display results for solving Eq. (3) with various discretizations and DPM- Solver with r 0.4. First, row 10 shows that we can indeed mitigate the timestep mismatch problem in the EDM schedule. It also shows we can gain improvements by using r < 1 (in the baseline setting, we use the sigmoid schedule with x 10 and Heun). See Appendix F.2 for full ablation over Second, rows = 20, 30 tells us can get even better results by increasing sharpness, but too large K hurts performance.
To investigate the performance difference between discretizations, visualize local truncation error || in Fig. 8, where t given time-step t; and Die4 Q is defined as
== ( + (4 l++1 )o (Z), (2
4++1)
solve( Do, (41144)
We first note that the uniform distribution incurs large error near E 10, This highlights that we indeed must place more points near those in order to control discretization error. While the EDM schedule has less error at those regions, because xx # 1, the mismatch between the initial state æ and time ty does not ensure the ODE is solved properly. Finally, control the error at the extremes. While the error for our schedule increases near z 0.5, Fig. tells sacrifice accuracy at intermediate to prioritize perceptual quality by choosing large
that our schedule is able to

---

Published as a conference paper at ICLR 2025
CIFARIO
AFHQv2
FFHQ
ImageNet (cond.)
Method
Reference
NEE
FID
STN
NPE
FID
STN
NFE
FID
STN
NFE
FID
STN
DM ODE
EDM
35
1.97
14.19
79
1.96
28.41
79
2.39
27.15
79
2.30
26.76
(Karras al., 2022)
9
37.91
9
28.03
9
56.84
9
35.46
DPM-Solver
9
4.98
9
9.26
9
6.64
(Lu ct al., 2022)
AMED-Solver
9
2.63
9
4.24
9
5.60
(Zhou al., 2024)
FM ODE
MinCurv
9
8.76
5.87
9
13.63
10.45
9
10.44
10.49
(Lec 2023)
FM-OT
142
6.35
138
14:15
(Lipman et al., 2023)
OT-CFM
100
4.44
(Tong et al., 2023)
MOT:50
132
11.82
(Pooladian al., 2023)
HM
100
2.96
10.73
100
2.73
16.20
100
3.30
16.71
Baseline
MOT-512#
100
3.29
8.77
100
5.53
13.45
100
4.69
14.29
MOT-1024
100
3.18
8.59
100
5.53
13.45
100
4.84
14.07
MOT-4096*
100
3.16
8.34
100
6.18
12.68
100
4.92
13.47
ReFlow
110
3.36
(Liu et al., 2022)
Simple ReFlow*
2.23
2.30
3.30
2.87
2.72
Ours
Guidance
9
1.98
2.49
9
1.91
5.60
9
2.67
3.24
9
1.74
3.92
Table 8: Comparison of neural ODE methods. MOT-b is minibatch OT with minibatch size b, and MinCurv curvature minimizing flow. We report FID and straightness (STN); S(va) =
fo E [1( - zo) va(w,, 1)] dt. Star next method denotes our training results.
APPLICATIONS
Comparison to other fast flow methods. Our approach significantly outperforms other ODE ap- proaches, e.g.. minibatch-OT FM (Pooladian et al., 2023; Tong et al., 2023) and curvature mini- mization (Lee et al., 2023), see see Tab. 8. Where possible, we report straightness (Liu et al., 2022), which quantifies how an ODE trajectory deviates from straight line between its initial and terminal points (see Tab. 8 and Eq. (21) and further details in Appendix D). We attribute the inferior baseline performance due to bias in minibatch OT, and discuss this and other pitfalls in Appendix A.
Improving perceptual quality via guidance. DMs often use guid- ance such as classifier-free guidance (CFG) (Ho & Salimans, 2022) or autoguidance (AG) (Karras et al., 2024) to enhance the percep- tual quality of samples. As observed by Liu al. (2024), condi- tional ReFlow models can also be combined with CFG. While it
FID
BSL
DVN 1 LAN
unclear what effect guidance has the marginals of ReFlow models, we also apply AG/ / CFG to our best unconditional / condi- tional models, since perceptual quality may be interest for certain downstream tasks. We already achieve state-of-the-ar results for fast ODE-based generation, but we obtain even lower FID scores with guidance, shown in the last row of Tab. 8. See Appendix F.2 for a full ablation over guidance strength.
1.74
0.2
0.4
0.5
0.8
1.0
Figure 11: ImageNet-64 FID at NFEs without and with our improvements.
coas--ooditional ImageNet-64. We verify the scalability of our training dynamics (DYN), learn- ing (LRN), and inference (INF) choices on ReFlow with the class- cass-condiitonal ImageNet-64 EDM model. We use 8M backward, 4M forward pairs, and 0.2. Fig. 11 CFG scale w 0, i.e., guidance, confirms that our techniques are effective. DYN+LRN improves BSL FID from 4.27 to 3.91, and INF further improves the FID to 3.49, improving on prior state-of-the-art fast flow meth- ods. With CFG we = 0.4, our DYN+LRN+IN model achieves an even better FID) score 1.74. Also, our techniques consistently improve CFG FIDs, implying that they offer orthogonal benefits.
CONCLUSION
We decompose the design space of ReFlow into training dynamics, learning, and inference. Within each group, we examine prior practices and their potential pitfalls. We propose seven improved choices for loss weight, time distribution, loss function, model dropout, training data, ODE dis- cretization and solver. We verify the robustness of our techniques on CIFARIO, AFHQv2, and FFHQ, and their scalability on ImageNet-64. Our techniques yield SoTA results among fast neural ODE methods, without latent-encoders, perceptual losses, premetries. In terms of FID score, weight and dropout contributed most. However, in terms of novelty, we believe training data im- provements and the generalized loss function largest contributions.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work was supported by the National Research Foundation of Korea under Grant RS-2024- 00336454 and by the Institute of Information & communications Technology Planning & Evaluation (ITT)) grant funded by the Korea government (MSIT) (No.2019.0.0007,, Artificial Intelligence Graduate School Program (KAIST)).
REFERENCES
Michael S. Albergo and Eric Vanden-Eijnded Building Normalizing Flows with Stochastic Interpolants. In ICLR, 2023.
Michael Albergo, Nicholas M Boffi, and Eric Vanden-Lijaden. Stochastic interpolants: A unifying framework flows and diffusions. arXiv preprint urXiv:2303.05797 2023.
Sina Alemohammad, Josue Casco- Rodriguez. Lorenzo Luzi, Ahmed Imtiaz Humayun, Hossein Babaci. Daniel LeJeune. Ali Siahkoohi, and Richard Baraniuk. Self-Consuming Generative Models Go MAD. In ICLR, 2024.
Donald G. Anderson. Iterative procedures nonlinear integral equations. J. ACM, 1244)554,5500, October 1965. ISSN 0004-5411, doi: 321305.
10.11..1144/2222 10.1145/322296.221333 URL htt.:tpp//.....oo//11.11///////.
Uri M. Ascher and Linda R. Petzold. Computer Methods for Ordinary Differential Equations and Differen- tialAlgebraic Equations. Society for Industrial and Applied Mathematics, 1998.
Mare G Bellemare, Ivo Danihelka, Will Dabney, Shakir Mohamed, Balaji Lakshminarayanan, Stephan Hoyer, and Rémi Munos. The cramer distance as solution to biased wasserstein gradients. arXiv preprint arXiv:1705.1074 2017.
Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion Schrödinger Bridge with Applications Score-Based Generative Modeling. In NeurIPS, 2021.
Yunjey Choi, Youngjung Uh, Jaejun Yoo, and Jung-Woo Ha. StarGAN v2: Diverse Image Synthesis for Multiple Domains. In CVPR, 2020.
Hyungjin Chung, Jeongsol Kim, Michael T. Mccann. Mare L. Klasky, and Jong Chul Ye. Diffusion Posterior Sampling for General Noisy Inverse Problems. In ICLR, 2023.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In C.J. Burges, L. Bot- tou, M. Welling, Z. Ghahramani, and K.Q. Weinberger (eds.), Advances in Neural Information Processing Systems, volume 26. Curran Associates, Inc., 2013. URL https: /ptpccediocccssseoeoooppppc...ee courrps--.. cc/ papcr_filde
aat.a/20000000022222222222222312 Paaee22bbbbb3333-PPaapppp.. pdf.
Marco Cuturi, Laetitia Meng- Pang-Papaxanthos. Yingtao Tian, Charlotte Bunne, Geoff Davis, Olivier Teboul. Optimal transport tools (ott): A jax toolbox for all things wasserstein. arXiv preprint arXiv:2201.12324 2022.
Prafulla Dhariwal and Alex Nichol. Diffusion Models Beat GANs on Image Synthesis. In NeurlPS, 2021.
Tim Dockhorn, Arash Vahdat, and Karsten Kreis. GENIE: Higher-Order Denoising Diffusion Solvers. NeurlPS, 2022.
Jean Feydy, Thibault Séjourné, François-Xavier Vialard, Shun-ichi Amari, Alain Trouvé, Gabriel Peyré. Interpolating between Optimal Transport and MMD using Sinkhorn Divergences. In AISTATS, 2019.
Zhengyang Geng, Ashwini Pokle, William Luo, Justin Lin, and Zico Kolter. Consistency Models Made Easy. arXiv preprint arXiv:2406. 14548, 2024.
Ian J. Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative Adversarial Nets. In NewrlPS, 2014.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, and Bernhard Nessler. GANs Trained by a Two Time- Scale Update Rule Converge to Local Nash Equilibrium, In NeurlPS, 2017.
Jonathan Ho and Tim Salimans. Classifier-Free Difffusion Guidance. arXiv preprint arXiv:2207.1259. 2022.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising Diffusion Probabilistic Models. In NeurlPS, 2020.
11

---

Published as a conference paper at ICLR 2025
Emiel Hoogeboom. Jonathan Heek, and Tim Salimans. Simple diffusion: End-to-end diffusion for high lution images. ICML, 2023.
7ahra Kadkhodaie, Florentin Guth, Eero Simoncelli, and Stéphane Mallat. Generalization in diffusion models arises from ecometry-adaptiic harmonic representations. In ICLR, 2024.
Tero Karras, Samuli Laine, and Timo Aila. A Style-Bash Based Generator Architecture for Generative Adversarial Networks. In CVPR, 2019.
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating Design Space of Diffusion- -Based Generative Models. NeurlPS, 2022.
Tero Karras, Miika Aittala, Samuli Laine, Erik Härkönen, Janne Hellsten, Jaakko Lehtinen, and Timo Aila. Alias-Free Generative Adversarial Networks. In NeurIPS, 2023a.
Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine. Analyzing and Improving the Training Dynamics of Diffusion Models. In CVPR, 2023b.
Tero Karras, Miika Aittala, Tuomas Kynkäääniemi, Jaakko Lehtinen, Timo Aila, and Samuli Laine. Guiding a Diffusion Model with a Bad Version of Itself. arXiv preprint arXiv:2406.0253 2024.
Alex Kendall, Yarin Gal, and Roberto Cipolla. Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics. In CVPR, 2018.
Beomsu Kim and Jong Chul Ye. Denoising MCMC for Accelerating Diffusion-Based Generative Models. In ICMI,, 2023.
Dongjun Kim, Chich-Hsin Lai, Wei-Hsiang Liao, Naoki Murata, Yuhta Takida, Toshimitsu Uesaka, Yutong He, Yuki Mitsufuji, and Stefano Ermon. Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion. In ICLR, 2024a.
Sanghwan Kim, Hao Tang, and Fisher Yu. Distilling ODE Solvers of Diffusion Models into Smaller Steps. In CVPR, 2024b.
Diederik P. Kingma and Jimmm Lei Ba. Adam: A Method for Stochastic Optimization. In ICLR, 2015.
Diederik P Kingma and Ruiqi Gao. Understanding Diffusion Objectives the ELBO with Simple Data Aug- mentation. In NeurlPS, 2023.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009.
H. W. Kuhn. The hungarian method for the assignment problem. Naval Research Logistics Quarterly, 2(1-2): 83-97, 1955. doi: wwt.://.....g/....00//a//....a. 1002 3ttp:///avv..0000000
URL
wiley.
htt//doi/abs///
Sangyun Lee, Beomsu Kim, and Jong Chul Ye. Minimizing Trajectory Curvature of ODE-hased Generative Models. In ICML,, 2023.
Sangyun Lec, Zinan Lin, and Giulia Fanti. Improving the Training of Rectified Flows, arXiv preprint arXiv:2405.20322 203300 2024.
Alexander C. Li. Mihir Prabhudesai, Shivam Duggal, Ellis Brown, and Deepak Pathak. Your Diffusion Model is Secretly a Zero-Shot Classifier. In ICCV, 2023.
Shanchuan Lin, Bingchen Liu, Jiashi Li, and Xiao Yang. Common Diffusion Noise Schedules and Sample Steps are Flawed. In WACV, 2024.
Yaron Lipman, Ricky T. Q Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Lc. Flow Matching for Generative Modeling. In /CLR, 2023.
Qiang Liu. Rectified Flow: A Marginal Preserving Approach to Optimal Transport. MiXiv:2209.14577 2022.
arXiv preprint
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow. arXiv preprint arXiv:2209.0300 03003, 2022.
Xingehao Liu, Xiwen Zhang, Jianzhu Ma, Jian Peng, and Qiang Liu. InstaFlow: One Step is Enough for Ffigh-quaing Diffusion- Based Text-to-Imaae Generation. In ICI,, 2024.
12

---

Published as a conference paper at ICLR 2025
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps. NeurIPS, 2022.
Chenlin Meng, Yutong He, Yang Song, Jiaming Song, Jiajun Wu, Jun- Zhu, and Stefano Ermon. SDEdit: Guided Image Synthesis and Editing with Stochastic Differential Equations. In ICLR, 2022.
Stefano Peluchetti. Non-denoising reward-time diffusions. arXiv preprint arXiv:2312. 14589, 2021.
Stefano Peluchetti. Diffusion bridge mixture transports, schrüddinger bridge problems and generative modeling. Journal of Machine Learning Research, 24(374):1-51, 2023.
Gram-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo- Homingo-Enrich, Brandon Amos, Yaron Lipman, and Ricky T. Q Chen. Multisample Flow Matching: Straightening Flows with Minibatch Couplings. In ICML, 2023.
Herbert Robbins. An empirical Bayes approach to statistics. In Proceedings of the Third Berkeley Symposium on Mathematical Statistics and Probubility, 1956.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High-Resolution Image Synthesis with Latent Diffusion Models. In CVPR, 2022.
Tim Salimans and Jonathan Ho. Progressive Distillation for Fast Sampling of Diffusion Models. In ICLR, 2022.
Antoine Salmona, Valentin de Bortoli, Julie Delon, and Agnès Desolneux. Can Push-forward Generative Models Fit Multimodal Distributions? In NeurlPS, 2022.
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion Schrödinger Bridge Match- ing. In NeurlPS, 2023.
Richard Sinkhorn. A Relationship Between Arbitrary Positive Matrices and Doubly Stochastic Matrices. The Annals of Mathematical Statistics, 35(2):876 879, 1964. doi: 10.12144amm/117770339 URL https: //doi. org/ 10. 1214/ aoms / 11 77703597
Jascha Sock-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep Unsupervised Learn- ing using Nonequilibrium Thermodynamics. In ICML, 2015.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising Diffusion Implicit Models. In ICI.R, 2021a.
Yang Song and Prafulla Dhariwal. Improved Techniques for Training Consistency Models. In ICLR, 2024.
Yang Song, Jaseha Dock-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-Based Generative Modeling through Stochastic Differential Equations. ICLR, 2021b.
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency Models, ICML, 2023.
Josef Stoer and Roland Bulisch. Introduction Numerical Analysis, volume 12. Springer Scicncc+Business Media New York, 2002.
Xuan Su, Jiaming Song, Chenlin Meng, and Stefano Ermon. Dual Diffusion Implicit Bridges for Image-to- Image Translation. In ICLR, 2023.
Alexander Tong, Kilian Fatras, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector Rector-Brooks, Guy Wolf, and Yoshua Bengio. Improving and gencralizing flow- based generative models with minibatch optimal transport. arXiv preprint arXiv:2302.0000 00488 2023.
Cédric Villani. Optimal Transport: Old and New. Springer Berlin, Heidelberg. 2009.
Pascal Vincent, A Connection Between Score Matching and Denoising Autoencoders. Neural Computation, 23(7):666-77, 2011.
Xingyi Yang. Daquan Zhou, Jiashi Feng. and Xinchao Wang. Diffusion Probabilistic Model Made Slim. In CVPR, 2023.
Qinsheng Zhang and Yongxin Chen. Fast Sampling of Diffusion Models with Exponential Integrator. In ICLR, 2023.
Richard Zhang, Phillip Isola, Alexei A. Efros, Eli Shectman, and Oliver Wang. The Unreasonable Effectiveness of Deep Features Perceptual Metric. In CVPR, 2018.
13

---

Published as a conference paper at ICLR 2025
Yifan Zhang and Bryan Hooi. HiPA: Enabling One-Step Texa-to-Image Diffusion Models via High-Prequency- Promoting Adaptation. arXiv preprint arXiv:2311.18158 2023.
Zhenyu Zhou, Defang Chen, Can Wang, and Chun Chen. Fast ODE-based Sampling for Diffusion Models in Around 5 Steps. In CVPR, 2024.
Yuanzhi Zhu, Xingchao Liu, and Qiang Liu. SlimFlow: Training Smaller One-Step Diffusion Models with Rectified Flow. ECCV, 2024.
A FASTER ODES VIA COUPLINGS
A.1 FASTER SAMPLING VIA STRAIGHT PATHS
Generative ODEs with straight trajectories can be solved accurately with substantially fewer velocity evaluations than those with high-curvature trajectories (Stoer & Bulisch, 2002). In fact, probability flow ODEs with perfectly straight trajectories can translate one distribution to another with a single
Euler step. For an ODE with velocity Ug. can quantify its straightness (Liu et al., 2022):
S(v)) == So E (2) - xo) (2 [)]
(21)
where the expectation over ODE trajectories (æ, : I E [0, 1]} generated by ug. An ODE with zero straightness has linear trajectories, which means it can translate initial points terminal points with a single function evaluation.
One approach to encourage straight paths is to learn ODEs which minimize trajectory curvature (Lee et al., 2023) by parameterizing the coupling with neural network which takes as input some image and outputs a sample such the distribution of these samples is close Gaussian. Another approach based on connections to optimal transport.
A.2 CONNECTION TO OPTIMAL TRANSPORT
For any convex ground cost, the solution the dynamic optimal transport continuous support will be straight trajectories, see e.g. Liu (2022). addition, training a flow matching model on samples from an optimal transport coupling will preserve the coupling, providing the vector field of the flow matching model is conservative.
As aside, more generally performing bridge matching (Peluchetti, 2021; Shi et al., 2023) purropicall--eguuariiee
optimal coupling will preserve the coupling, though the trajectory will be given by an SDE, and hence no longer straight. In the limit as the entropic regularization term tends to zero, then this recovers the non-regularized optimal coupling with squared Euclidean ground cost (Shi et al., 2023; Peluchetti, 2023).
This motivates learning optimal transport (OT) coupling and then performing bridge / flow match- ing on this coupling. There are two dominant ways do this, either ReFlow (also known as Iterative Markovian fitting) (Liu et al., 2022; Liu, 2022; Lee et al., 2024; Shi et al., 2023) approximating the coupling using mini-batches (Tong et al., 2023; Pooladian et al., 2023).
A.2.1 MINI-BATCH OPTIMAL TRANSPORT FLOW MATCHING
We first make distinction between the loss batch size bess and coupling batch size beoupling where Dewupling 2 buss- The loss batch size, huss is the number of input pairs (mo, 21) used per training iteration in the flow matching loss Eq. (1), whereas the coupling batch size, bacupling. refers to the number of independently sampled pairs used as input into the mini-batch OT solver.
The procedure for obtaining mini-batch couplings is first independently sample becupling items Dcsppins from both marginal distributions ~ Pu and y/s P1. The next Drmening step is run mini-batch OT solver obtain a coupling matrix (pi.)) 1=1=1 such that Eig Pi.j = Pis I/beoplingg E, 1/0emppnn
denoted (ys))=miin (I:)
Dosspinn
The final step (ā,, Niy ~ P, which are then fed into loss Eq. (1) and corresponding standard gradient based optimization procedure.
sub-sample the coupling batch of size beoupling to obtain b aligned pairs
14

---

Published as a conference paper at ICLR 2025
Mini-batch bias. In stochastic gradient descent, for example, losses computed uniformly sam- pled batches are unbiased with respect to the measures the batches were sampled from. This is not true for mini-batch OT couplings with respect to the true OT coupling between marginals (Belle- mare et al., 2017). Indeed, marginal preservation within mini-batches may force points in each minibatchto be mapped together, such points may not be mapped, or have very low probability of being mapped, in the true OT coupling. Asymptotically the mini-batch couplings should converge the true OT coupling solvers the mini-batch size increases. Unfortunately, this is not practically feasible with discrete OT solvers for large datasets or indeed for measures with continuous support. The mini-batchs and true OT couplings would also be the same for infinite regularization, indeed the couplings would both be independent couplings and so not very informative.
Subsampling. Computing OT couplings for large batch sizes is not typically possible using the Hungarian algorithm (Kuhn, 1955) due to the cubic time complexity. However, entropic approxi- mations from Sinkhorn (Sinkhorn, 1964) is of only quadratic complexity and can be implemented on modern GPU-accelerators (Cuturi, 2013; Cuturi et al., 2022), hence enables fast computation of discrete entropic OT for batches in excess 100, 000 points.
Prior works (Tong et al., 2023; Pooladian et al., 2023) set bass cooupling and do not subsample. This is problematic as mini-batch OT is only justified as being close to optimal in the asymptotically large batch regime. Although for training the neural network via FM is limited by hardware memory and so it becomes infeasible Duoupling for large batch size. Prior works therefore small coupling batch size.
compute the coupling for large batch sizes, the optimization setup
to set bioss
Subsampling should still preserve marginal distributions. We observe in Tab. 8 that the straight- ness of the generative trajectories increases batch size grows, expected, however generative performance in terms of FID gets increasingly worse compared to regular flow matching. This is a surprising empirical result that warrants further investigation.
A.2.2 REFLOW AND ITERATIVE MARKOVIAN FITTING
ReFlow (Liu et al., 2022; Liu, 2022; Lee et al., 2024) and more generally Iterative Markovian Fitting (Shi et al., 2023) are procedures which iteratively refine the coupling between marginals. We shall focus on ReFlow for brevity. ReFlow first takes an independent coupling, then involves training a flow between samples from that coupling, known as a Markovian projection. Simulating from this trained flow is then used to define an updated coupling. This process is repeated between updating a flow and coupling until convergence. It has been shown that this process iteratively reduces the transport cost for any convex ground cost, and hence straightens the paths between coupling whilst retaining the correct marginals.
Note that ReFlow results in coupling which is slightly stronger than optimal transport. OT aims to minimize the transport cost for specific ground cost function, whereas ReFlow reduces transport cost for all convex costs. ReFlow can be limited to specific convex ground cost by ensuring the vector field takes a specific conservative form (Liu, 2022).
B FAST SAMPLING VIA HIGHER ORDER SOLVERS
One can use higher-order solvers which utilize higher order differentials of the ODE velocity to take large integration steps or reduce truncation error (Karras et al., 2022; Dockhorn et al., 2022; Lu et al., 2022; Zhang& Chen, 2023). While this approach is generally training- free, recent works (Zhou et al., 2024; Kim et al., 2024b) have incorporated trainable components which minimize truncation error further accelerate sampling.
C DISTILLATION AND CONSISTENCY MODELS
The goal of distillation within the field of diffusion models is typically to compress multiple steps along a probability flow ODE of a teacher diffusion model into a fewer steps steps of a student model. We refer to this discrete-time distillation (DTD). Representative methods are progressive distillation (Salimans & Ho, 2022), and consistency distillation (Song et al., 2023; Song & Dhariwal, 2024; Kim al., 2024a; Geng et al., 2024). While distillation and ReFlow similar in the aspect
15

---

Published as a conference paper at ICLR 2025
CIFARIO
AFHQv2
FFHQ
ImageNet-64
lterations
200k
200k
200k
500k
Minibatch Size
512
256
1024
Adam LR
2e-4
2e-4
2e-4
2e-1
Label dropout
0.1
EMA
0.9999
0,9999
0,9999
0.9999
Num. Backward Pairs
1M
1M
1M
8M
Num. Forward Pairs
50%:
13.5k
70k
4.M
Table 9: Training hyper-parameters.
that they train a new model using the outputs teacher diffusion model, we emphasize that they are, in fact, complementary approaches, and can benefit from one another. We discussion this point in more detail in the following section.
C.I REFLOW vs. DISTILLATION
We remark that faster ODEs have several practical benefits over discrete-time distillation alone. Since translation along ODE is bijective map, evaluation by integrating the ODE backwards starting from data.
achieve fast inversion and likelihood
Fast ODEs can be combined with discrete-time distillation. For instance, Lee et al. (2023); Liu et al. (2024); Zhu et al. (2024) have observed it is substantially easier to distill ODE models with straight trajectories. One may also any ODE solver with continuous time ODE, and there may be some benefit using adaptive solvers. Lee et al. (2023; 2024) also report combining RF models with higher-order solvers improves the trade-off between generation speed and quality.
D EXPERIMENT SETTINGS
D.1 TRAINING AND EVALUATION
To evaluate a training setting, we initialize ReFlow denoisers with pre-trained EDM (Karras et al., 2022) denoisers, and optimize Eq. (13) with Qoi Qh of Eq. (7). Specific optimization hyper- parameters are reported in Tab. 9. We sample backward and forward pairs from Qhi by solving Eq. (6) with EDM models and use them throughout training. Specifically, we use the EDM dis- cretization with the Heun solver (Ascher & Petzold, 1998). We use sampling budgets of 35 NFEs for CIFARIO and 79 NFEs for AFHQv2, FFHQ, and ImageNet. FIDs of backward training samples are reported in the first row of Tab. 8.
We measure the generative performance of the optimized model by computing the FID (Heusel et al., 2017) between 50k generated images and all available dataset images. Inception statistics are computed using the pre-trained Inception-v3 model (Karras al., 2023a). Samples are generated by solving Eq. (3) with the Heun solver with 9 NFEs, and we report the minimum FID score out of three random generation trials, as done by Karras et al. (2022). For reasons described Appendix F.1, we use the sigmoid discretization instead of the bascline uniform discretization.
D.1.1 BEST SETTINGS
Here, report seper-parameterss used produce results for our best models in Table 8.
CIFAR10. High-pass filter = 10, dropout probability 0.09, forward pairs with mixing ratio
p 0.4, sigmoid discretization with x = 20, DPM-solver 0.4, AutoGuidance scale w 0.6.
AFHQv2. High-pass filter À = 10, dropout probability 0.09, forward pairs with mixing ratio p =
0.4, sigmoid discretization with 6 20, DPM-solver 0.4, AutoGuidance scale 1.0.
FFHQ. High-pass filter = 10, dropout probabiliy 0.03, forward pairs with mixing ratio p = 0.2,
sigmoid discretization with 20, DPM-solver 0.4, AutoGuidance scale = 0.3.
ImageNet-64. High-pass filter = 10, dropout probability 0.05, forward pairs with mixing ratio
p ()2, sigmoid discretization with Æ. 20, DPM-solver 0.4, CFG scale = 0.4.
16

---

Published as a conference paper at ICLR 2025
D.2 FLOW MATCHING BASELINES
We strove to obtain competitive baselines for base and mini-batch OT flow matching methods, and indeed achieved superior performance comparable implementations from Tong et al. (2023) on the datasets considered.
Firstly, similar Karras al. (2022), formulate flow matching xo mean-prediction rather than using regression target Xo X1. We parameterize the mean- preddictcnn to be of form D2( t) = xt) + Cour (t) (cin(t)))) co (t)) where Fe is neural network:
XX),) Xoll
(22)
The scalar functions ca, Cskip Cout: Cin: A(t) are derived according to the reasoning of Karras et al. (2022) in Sec D.2.2. We set O0T = 0 for the independent coupling.
Throughout we use the similar setup as Karras et al. (2022) but with the flow matching loss and new preconditioning. In particular for AFHQv2. FFHQ and CIFARIO we use the SongNet from Song et al. (2021b) with corresponding hyperparameters from Karras et al. (2022) per dataset.
The time-sampling during training taken to be uniform and the Euler solver with 100 steps used for computing FID and straightness metrics, in order to be comparable to other reported baselines from Tong et al. (2023).
D.2.1 MINI-BATCH FLOW MATCHING
The mini-batee flow matching experiments use the same learning rate, networks, and training ob- jectives as base flow matching. The primary difference is in how the inputs, X,, Xo. are sampled.
We follow the procedure outlined in Sec. A.2.1 for sampling mini-batches, using Sinkhorn (Sinkhorn, 1964; Cuturi, 2013) as the mini-batch solver based on the OTT-JAX library (Cuturi al., 2022). Images were scaled [-1, 1] as is standard in diffusion models and flattened. The squared Euclidean ground cost was used.
The regularization parameter was set to = 2, qualitatively this provided reasonable trade-off between meaningful coupling visually and the time to compute using convergence threshold defaults from Cuturi al. (2022). The default regularization parameter from Cuturi et al. (2022) did not provide a visually meaningful coupling on large batches, and setting parameter less than e < 1 took over the maximum iteration threshold of 2, 000 iterations converge, and hence was not feasible for training.
Each Sinkhorn loop took approximately 100 200 Sinkhorn iterations without acceleration tech- niques, and wall-clock time up to roughly 0.8s for the largest coupling batch size 8192. We then ran acceleration techniques including Anderson acceleration (Anderson, 1965) with memory 2, epsilon decay starting from 10, and initializing potentials from prior batches to reduce runtime. This sped up the mini-batch process 0.4s per Sinkhorn loop, and convergence of Sinkhorn in approximately 20 30 Sinkhorn iterations.
We ablated the coupling batch size between 512, 1024, 4096, 8192. The loss batch size was kept constant at 512 for CIFAR10 and 256 for AFHQV2 and FFHQ
Scope for further improvements: Unfortunately, Cou[Xo, ær] o is not known for mini- batch couplings and hence as in the independent coupling we set oo, = 0 in the preconditioning computation. is possible that this sub-optimal and may be estimated better ways.
Although straightness improves with larger batch size and our implementation achieves better FID scores than prior baselines, mini-batch OT flow matching is still not well understood. is puzzling as to why performance in terms of FID gets worse compared to base flow matching. This is corrob- orated Table 5 of Tong et al. (2023) where FID for CIFAR1O is 3.74 with the mini-batel coupling and 3.64 with independent coupling, however notice significant discrepancy 2.98 FID for the independent coupling and 3.16 for the best mini-batch coupling. We leave further investigations to future work.
17

---

Published as a conference paper at ICLR 2025
D.2.2 PRECONDITIONING
In the interest of generality, derive EDM-style preconditioning (Karras al., 2022) for the more general case of bridge matching / stochastic interpolant (Peluchetti, 2023; 2021; Shi et al., 2023; Albergo et al., 2023) which recovers preconditioning for flow matching for y/s 0,
Let X, = Xo + B,XT + ThE where ~ I0) Consider prediction of form (o( = Cskip (t) + Cour (CC(())))) Co (t)) and X(-) weighted loss per Eq. (22).
The loss per Eq. (22) may be written:
X, Cou (X),,, t) Coun(t)
Cour(t) (Xo
Cakip! Cakip(t) [1)
(23)
Setting X(-). In order to uniformly weight the loss per time step, we set A(t) Cour Coun()) -2 similarly Karras al. (2022).
Setting Cin (l). We take the strategy of finding Gn such that Var[cia(() (1)
Let Var(Xo) og. Var [ar[Xx] of and Cov Xx, = rr
(arlci(())xp) cin(t (t) (aio0 + 3,00 20,3:00,7 + m2.
(24)
(1)(1) = [a700 + 86 + 20, + 7/
(25)
Setting Cakip and Cout- The prediction target of Do(I,, t) is Xg. hence the target of network F6 is Cour() (oun(t)-1 [oo- (t) We choose Cakip and Cout to ensure regression target has uniform variance Le. Var (Cour() [Xo- - Cskip (t) =1
Var (Cou((t)- -Xo - Cakip Cakip(t) (t) =
(26)
Cour (t) = Var [Xo Cakip (t)X,]
(27)
Caut (1) Var [(1 axccuiip (o)) Cskip Cskip(1) (1) + ]6)
(28)
Cour (t)² (1 attskip (t))²
(29)
23, 0g2 Cskip (1) (1) + Ctxip((t) 32 of + 7/ Cskip (t)²
(30)
(31)
Given the fixed relationship between Cskip and Cout. we choose Cakip to minimize Cout
d dessip
20x(1 ((Cssip(t))
(32)
28,0h, + can
(33)
+ 2c
(1)38 2c)xs(())
+ 2 erkip (1)
(34)
de dese
With first order condition
O. we obtain:
aoo + 3.00r
Cskip (t) ajo + 2,0 59 0 7
(35)
D.3 COUPLING PROJECTION
Recall that we propose projecting Qhi to II(Po,,,) at the end of each iteration
Q == (Q
(36)
and using Qh in place of Qh. However, the projection operation is well-defined only if there is suitable metric on the space under consideration (the space of distributions, in our case). An applicable metric the p- Wasserstein distance W,- Then, projection w.r.t. W, defined as
Q
arg min (T Qb1)
s.t. To Fo. r = F1.
(37)
Toi
18

---

Published as a conference paper at ICLR 2025
Furthermore, may parameterize
oT (To, x1) dl (20 or dFo( ((0
(38)
which means (with the first parameterization), only have enforce the marginal constraint
arg min (To1. Q11) s.t. To Po. df 01 dI 1o
(39)
Noting that
To = Po +
D(To, Po) =
(40)
for distances or divergences D, we can optimize
min D(T,, Po) + AW 2 s.t. dr = dr of d
(41)
Toi
for decreasing values of À and stop when oT Po) saturates. In practice, we solve
min (T Po) + ASKD Qo1) s.t. dT = of
(42)
Fo:
with gradient descent, where SKD stands for Sinkhorn Divergence (Feydy et al., 2019). We approx- imate Q01 mixture of diracs using the generated backward pairs, and approximate D using Generative Adversarial Network (Goodfellow et al., 2014). Since we do not know an appropriate value of X, we initialize À from large value, e.g., À 1000, decay it by factor of 0.1 every time FID saturates. If decaying À does not offer any more FID improvement, we terminate optimization, and use the optimized To as Q0- The full optimization procedure is described in Algorithm 2.
Algorithm 1 Coupling projection given
Inputs: Po. Qb Fo1. batch size B, discriminator De discriminator learning rate n, coupling learning rate 7, evaluate FID every NFID iterations, SKD, coefficient À
2: Initialize <- oolbbett <- To1, FIDbest < FTD Pa)
3: while training do
4: Sample { B
Po. (())))))) ) m= ~ Qb {(a +))) Toi
5: + + V
log (Tjj))) + 2 log(1 Dx(()))
6:
2 + 26
ag (log(1 D (x))) + ASKD, *SSDD 2. ) m= ((a)) xx) )
{
7:
itit
8:
if % NEID 0 then
9:
FID + (To, if FID curr FIDbest best then return oosbeets FIDbest
10:
11:
13: else
(To1, FIDbest) & (To1, FIDeurr)
end if end if
15:
end while
Algorithm 2 Coupling projection
1: Inputs: Pus Qb1- batch size B, discriminator Dy, discriminator learning rate 7). coupling learn- ing rate y, evaluate FID every NEID iterations, initial X. decay factor p (0, 1)
2: Initialize < 0, Fo <- Q Fo +- C FIDbest +- FI, Fo)
3: while training do
4: ToL, FIDearr < ALGI [Po, Qbi. Toi B, Da, n.Y, NFID,
5: if' FID eury FIDbest then
6: return FoL.best
7: else
8:
(A, Foi.estt FIDbest) & (p. Ton, FID
9:
end if
10:
end while
19

---

Published as a conference paper at ICLR 2025
E PROOFS
E PROOF OF PROPOSITION
Lemma 1. The following statements are equivalent.
(a) 0 minimizes (F; Qo1, De. t).
(b) 0 minimizes (G: Qu1 Le, t).
(c) Du(ærs EevonQu Co (-ba)) (aro].
Proof. We first observe that (writing Dg in place of Do(ær, 1) for brevity)
D, L0rm Qo1. Zt, 1) = "T (VVo(VV,, Craa(f; Qo1, 2t,t)}
(43)
and since @ is invertible, o @ is invertible as well, which implies
DD,CFFM Qo1, æt, =0 = VD. (0; Q01 or, t) = O.
(44)
Because both LGFM(0; Qoi, Rt, 1) and (FM(0: Qoi, æz, are strongly convex w.r.t. Do, this means 0 minimizes (0 Qoi, Lt, 1) iff 0 minimizes LFM(0: Qoi, ær, 1) iff
Da( t) Ez (11.) (20).
(45)
This establishes the equivalence of the three claims.
Lemma 2. Let μ be a o-finite measure. If 9 on a set A with p(A) SA f du > JA dd.
Proof. By linearity of integrals, we can assume g = 0. Since f > 0 on A, we may express
A= An = {x E A f(x) > 1/n].
(46)
Since p(A)= O, there is such that u(A..) 0. Otherwise, by subadditivity of measures.
((A) < E (*(A)) =0 which contradicts the assumption μ(A) > It follows that
(47)
f du 2 SA. f dp 2 SA, Adu u(An) 0.
(48)
This establishes the claim.
Proof of Proposition 1. Denote the measure of (t, xx) where t ~ unif 1) and æ Q as μ. (Assuming D3 can approximate a sufficiently large set of functions), define as the neural net parameter which satisfies
Dg. (xt, 1) = Exoo x2) [20]
(49)
for any (æt, () such that
LGFM(0; Qoi, xe,t) 2 (0" Qo1,
(50)
or equivalently,
(FM(0; Qo1, 2 LFxal0"; 20,, t)
(51)
for any (æy, f) and by Lemma
We now show that a minimizer of Eq. (13) minimizes Eq. (1). Suppose 0 minimizes Eq. (13), but there is set with positive measure, i.e., ((A) > O, such that
Do(æ,, # Extor ( (ax)) (2x)) [xo]
(52)
for all (ar, 1) A. By Lemma 1,
LCFM(0; Qo1, z1, Lor Qo1,
(53)
20

---

Published as a conference paper at ICLR 2025
for all (z), A, and since w( and dTT(t) positive by assumption, > dT(() w( t) LorM(f"; Qoi xt,
dT((t) w w( t) LGFM(0;
(54)
for all (ai, t) E A, so by Lemma 2,
F a ( (T)()) w(r, b) LorM(0; Qo1, ar, ((t, (A dT(t) w(x,, LGFM( Qoi Zz, t)]
(55)
(56)
where La( 1) if (2t, 1) E A and if not, and so LGFM(0; Q01) (( (dT(t) w( LGFM(O; Qo1 2t, t)] (((,,.))u (dT(()) (dTT))))(((((()))))))))))))))))) w(aee 1) LGrM(O"; Qo1; Ft 1)] (fr Qo1)
(57)
(58)
which contradicts the assumption that 0 minimizes Eq. (13). It follows that if 0 minimizes Eq. (13),
Da( t) = Bno~QQ (-1r=))) (())) [0o]
(59)
almost everywhere w.r.t p. it also minimizes
LPM Q01 xt, t)
(60)
almost everywhere w.r.t p by Lemma which implies 0 minimizes Eq. (1), The other direction can be proven in analogous manner.
E.2 PROOF OF PROPOSITION
Proof of Proposition 2. Let Q IP @ P1. If we assume zero initialization in output layer for Da
max LDM(0; 8 æ1, 1) min LDM 08 1)
(61)
max E Co11 ( (m) 1 D2 1)1l2/
min 10 - Dz(01, 1)1l2
(62)
21
min E Bg [|) 2a / max E 0g (1lm))
(63)
mn
max
x / min ~ro
æ 2 =1
(64)
On the other hand, if we pre-trained diffusion model to initialize Do,
Do(a1, no
(65)
such that
max CGFM(0; Q01 11, 1)/ min CGFM (0; Qoi x1, 1)
(66)
251
max E F (-m,, ||:00 14 / min E ( (11m))))
10.1
(67)
m1
po xo .max
min |00
3
(68)
because 21 + Zy ~ 0 (- [11 now bijective map between P, and P1 samples.
21

---

Published as a conference paper at ICLR 2025
FID
FID
AFTIQ22 FFRQ
AFTIQu2
FFHQ
3.0
2.5
2:29
2.5
2.67
2.25
2.0
2.0
1.98
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
1.0
Figure 12: DPM-Solver
Figure 13: AutoGuidance w
F ADDITIONAL EXPERIMENTS
F.1 LINEAR DISCRETIZATION LACKS DISCRIMINATIVE POWER
While all previous works use the uniform discretiza- tion to sample from ReFlow models, we use the sig- moid discretization to evaluate models in Sections 3.2 and 3.3. This is because, we found that the uni- form discretization lacks discrimination power, i.e, the ability to make the best of a given model, espe- cially small NFEs.
w( f)
Uniform
Sigmoid
2.88
2.87
2.89
2.76
1/22
2.93
2.74
(o + (.055/( 1/B CorM (0: Oo m.).
2.97
2.82
2.98
2.79
1/ CCppa Qo 2..))
2.95
2.74
To demonstrate this, in Tab. 10, we re-evaluate mod- els in Sec. 3.2.1 with the uniform discretization, and compare them with evaluation results with the sig- moid discretization with tion are better than the worst FID with the sigmoid discretization. Moreover, the model with our proposed weight, when evaluated with the uniform schedule, performs worse than the model with uniform weight.
Table 10: Uniform vs. sigmoid discretizations with Heun on AFHQv2.
10)
10. We observe that none of the FIDs with the uniform discretiza-
We speculate this happens because, as analyzed Sec. 3.4, large curvature regions for ReFlow ODEs occur near E {0, 1}. but the uniform discretization fails to account for them. So, the uniform discretization unable to accurately capture the differences ODE trajectories between different models. Due to these reasons, opt to the sigmoid discretization to distinguish training techniques that work from those that do not.
F.2 DPM-SOLVER AND GUIDANCE ABLATIONS
Recall that the DPM-Solver update for Eq. (3) is given
Et. f- Dtat1 + (4; - 4+ ( vo 8141) 1) (1- 전ه (x (4+)).
(69)
In Fig. 12, we show the FID for various values of r E (0, While we can get better FIDs than those in Tab. 7 by using tailored to individual datasets, we opt for simplicity and set r 0.4 as our improved choice, which still yields better FID than the Heun solver, i.e., using r
For conditional ReFlow models, classifier-free guidance (CFG) (Ho & Salimans, 2022) can be for- mulated as solving the ODE
dæ; {11 +w va( t,
ve 2)} dt
(70)
where Va (, c) is velocity conditioned on c, and vo (a I, Ø) is an unconditional velocity, and is guidance scale. Note that w O reduces the ODE standard class-conditional generation. In practice, we train conditional velocities with label dropout such that Va (a t, c) and vy (æ,, t, Ø) can be evaluated in parallel. by passing class labels to the former and null labels the latter.
For unconditional ReFlow models, AutoGuidance (Karras et al., 2024) can be formulated as solving da; = {(1 + w) my (0 () () w (2, 1)) dt where Gg is degraded version of vg. In practice,
(71)
ReFlow models trained with the baseline training configuration (see Tab. 1) for 10k iterations as 6g. While other choices of è may offer better FIDs, as AG is not the main topic of our paper, do perform extensive search.
22

---

Published as a conference paper at ICLR 2025
F.3 QUANTIFYING REFLOW BIAS REDUCTION
The extent of the bias introduced by ReFlow, and its reduction with our techniques can be calculated by comparing Tables 6, 7, 8. Specifically, since we measure the performance of our ReFlow models via Frechet Inception Distance (Wasserstein-2 distance between Gaussian approximations of true and model distributions in the feature space of the Inception network) (Heusel et al., 2017), we may use
CIFARIO
AFHQv2
FFHQ
DYN DYN+LRN
0.61
0.59
1.30
0.41
0.51
0.74
DYN+ IN+INE
0.26
0.34
0.45
DYN-LRN-INF-GD
0.01
-0.05
0.28
Table 11: Amount of bias introduced by Re- Flow under different settings. Negative bias means our model achieves a better FID than the diffusion model used to generate Q
(FID of Q++1 ) (FID of Q6)
proxy for the amount of bias introduced by Re- Flow. In our work, with n = 1, Q is the EDM model marginal and @ is our ReFlow model marginal. In Table 11, we summarize the extent of bias add improved training dynamics (DYN), improved learning (LRN), improved inference (INF), and guidance (GD) to the baseline (BSL) setting. With everything combined, techniques achieve significant reduction in bias.
23

---

Published as a conference paper at ICLR 2025
F.4 SAMPLE VISUALIZATION
(a) BSL. 2.83 FID with 9 NFEs
(b) DYN+LRN+INF, 2.23 FID with 9 NFEs
(c) DYNALNNNINTA 1.98 FID with 9 NFES
Figure 14: CIFAR10 samples with fixed random seeds
24

---

Published as a conference paper at ICLR 2025
(a) BSL. 2.87 FID with 9 NFEs
(b) DYN+LRN 2.30 FID with 9 NEEs
(c) DYNGLRNNINN+AGG 1.91 FID with 9 NFEs
Figure 15: AFHQv2 samples with fixed random seeds
25

---

Published as a conference paper at ICLR 2025
(a) BSL. 4.27 FID with 9 NFEs
(b) DYN+LRN+INE, 3.49 FID with 9 NFEs
(c) DYN+LRN+-NF+CFG, 1.74 FID with 9 NFEs
Figure 16: ImageNet-64 samples with fixed random seeds
26

---

Published as a conference paper at ICLR 2025
G FURTHER DISCUSSION OF OUR IMPROVEMENTS
Here, we provide further insight into why our techniques improve upon prior practice.
Loss normalization. In Section 3.1 show that the Generalized Flow Matching objective is collection of regression problems aggregated over (æt, 1), and at Section 3.2.1, we theoretically and numerically show that individual regression loss values during ReFlow can have vastly different scales with respect to (It, t). Multi-task learning interpretation of loss normalization (Zhang et al., 2018: Karras et al., 2023b) along with our observation motivates loss normalization with respect to both and and in Table 2, we demonstrate that loss normalization beats all other weights.
Time distribution. In Section 3.2.2, we explain that previous work (Lee et al., 2024) uses the cosh time distribution in order to oversample t near O or 1, where most of the learning happens. We also explain that we choose the increasing exponential time distribution, since our weight function already compensates for vanishing loss near 0.
Loss function. Section 3.2.3, we show that using & the loss function is equivalent to precon- ditioning the loss gradient, and it is well known that appropriate gradient preconditioning can accelerate model convergence (Kingma & Ba, 2015).
Dropout. In Section 3.3.1, we explain that we need models with larger Lipschitz constants if we wish to learn better ReFlow models. This is because ReFlow converges to straight ODE, and a straight ODE is ultimately fush-forward generative model, and (Salmona al., 2022) formally shows that a push-forward generative model needs to have a Lipschitz constant in order to map a unimodal distribution
multi-modal distribution accurately. For instance, Corollaries 6, 8 in (Salmona et al., 2022) show divergence or distance between data and model distributions is lower bounded by a decreasing function of Lipschitz constant of the push-forward model. This motivates us to increase effective model capacity by decreasing dropout probability.
Training data. Alemohammad et al. (2024) shows that recursively training generative models on data generated by itself reduces the quality and diversity of data. Alemohammad et al. (2024) also shows one can delay or prevent degradation by injecting real data into the training loop. Using forward pairs can be interpreted as an instance of injecting real data in the training loop, and using projected pairs can be interpreted as synthesizing new real data by solving the projection problem.
Discretization. In Figure 8, as evidenced by truncation error for the uniform discretization, the ODE after ReFlow has high curvature regions near and Truncation error for our sigmoid discretization schedule shows it is able to effectively control the error at the extremes of the interval.
Solver: We note that Karras et (2022) popularized Heun as an alternative to Euler, among the large set of solvers considered, based primarily on strong empirical performance. We argue that DPM-solver a generalization of Hcun (coinciding for r = 1), and we observe that setting r = 0.4 performs better than Heun's 's second order solver. can be tuned cheaply, especially since ReFlow models produce state-of-the- results with NFE- < 10.
27


Reviews:
[{"review_id": "JevZiNW4Vf", "reviewer": "Reviewer_7gZz", "summary": "The paper studies ways of improving the ReFlow procedure for diffusion and flow matching. They break the design space into training dynamics, learning and inference, and show isolated ablations w.r.t. the choices here (weighting, time distribution, loss function, initialization, dropout, sampling, ODE solvers, [0, 1] discretization). The main experimental results ablate these on CIFAR10, AFHQv2, and FFHQ.", "questions": "None", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I find the paper well-presented. The design space is clearly communicated and the ablations are convincing. These dimensions show how to squeeze the performance on these standard image generation settings.", "weaknesses": "I do not find many weaknesses in the paper, as it clearly communicates the goal and setting and explores it nicely:\n\n1. The experimental settings of CIFAR10, AFHQv2, and FFHQ are interesting, but only a small subset of all of the other potential uses and larger applications of flow matching and diffusion models. It would be interesting to see if the improvements continue holding in these.\n2. The paper is niche in the sense that it relies on wanting to do generative modeling with ReFlow, and discussing the relevant dimensions here.", "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": "P2i6EQny7s", "reviewer": "Reviewer_3r18", "summary": "This paper provides an empirical study of how to improve the FID of rectified flow based image generative models. They aggregate a number of design choices studied in the context of diffusion and reconsider them in the context of flow matching and rectification. These design choices focus on the various aspects of the loss function: the metric under which the regression is performed (e.g. MSE vs PIPs, vs generalized invertible mappings), the time weighting, and loss normalization. They then consider training choices, like how the training data is chosen based on the first deterministic coupling of the rectification (fixing a Gaussian base value for each data point, called forward pairs). Finally they run ablations over discretizations of the ODE.", "questions": "The reviewer is trying to understand the language of the following claim in 3.3.1:\n\n\"Our improvement. We observe that learning the OT ODE is a harder task than learning the diffusion\nprobability flow ODE. For instance, at t = 1, the optimal DM denoiser only needs to predict the data mean $\\mathbb E_{x_0 \\sim P_0} [ x_0] $ for any input $x_1 ∼ P_1$, but an optimal ReFlow denoiser has to directly map noise to image along the bijective OT map.\"\n\nShould the reviewer interpret this as being the claim that: The velocity field solving the probability flow is generically for $\\mathbb E[\\dot x_t | x_t = x]$, and for a deterministic coupling $P(x_0, x_1) = P_0 (x_0) \\delta (x - T(x_0)$, this conditional expectation collapses because only one $x_t$ passes through $x$? This would be true for any monge coupling, and not just the OT coupling.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper provides a number of useful empirical studies of what ultimately influences the performance of the flow matching models and when they are rectified to straighter paths. As the reviewer sees it, the paper aggregates many of the design space choices that arise from the diffusion model literature and uses them to study how these knobs affect flow matching. It is nice to have them all considered in one place. \n\nThe inclusion of a new high pass filter metric on the loss is nice, and could be well motivated in other domains where we explicitly know these frequency regimes are important, e.g. for generative models in the sciences.\n\nThe extent of the ablations is thorough and the reviewer acknowledges the size of this experimental undertaking!", "weaknesses": "Thanks to the authors for their hard work on this project.\n\n*Main remarks*:\n- While the paper provides a number of empirical insights for training, it would be nice if there was more of an analysis of why these tricks ultimately help. The reviewer stresses this a bit because a number of the proposed innovations are directly inspired from their application in other work, e.g. the improvement of loss normalization in Karras et al (2023b), and e.g. equations (8)-(12). \n\n- The reviewer is confused by the reference to OT solutions throughout, e.g. in Section 3.3.1, which suggests a potential misconception about the rectified flow procedure. In particular, as the reviewer understands it, one iteration does not give OT, but instead gives a monge map, in the sense that the coupling $P(x_0, x_1)  = P_0(x_0) \\delta(x - T(x_0))$, where $T$ is the transport endowed from learning the velocity field first with $x_0, x_1$ independently drawn in the loss. This is certainly not OT (brenier polar factorization theorem), but is a deterministic coupling. Can the authors address this? \n\n- The argument of proposition 2 is confusing, perhaps due to a typo. Is the first line supposed to specify the DM loss and not the GFM loss? Currently they are both statements about GFM. \n\n- These analyses are mostly limited to the image domain, and it's hard to understand how any of these \n\n- Can the authors address how large the bias is in rectification under the different procedures somehow? By this, the reviewer means that each iterate of rectification introduces a bias unless the flow was initially learned perfectly. The reviewer understands this to be difficult to measure, but it would be nice to understand how these ablations affect or relate to the extent of this bias. This might sit better in the questions section.\n\n\n\n\n*Small remarks*:\n-  the fundamental idea of of FM and rectified flow are equivalent to the method of stochastic interpolants, and should probably be cited accordingly. \n- There is in a sense a competing small set of remarks about the benefits of rectification vs distillation in the appendix section C, saying that there are benefits of rectification over distillation because rectification maintains an ODE-like structure. But this ignores the fact that the learning problem after a rectification is a distillation problem...and it too is now fitting to an imperfect teacher model. I'm not sure this remark really contributes much to the discussion.\n\n*Very small remarks*: A few typos to watch out for. \n- \"we  we\" top of page 5", "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": "Zcv583swMh", "reviewer": "Reviewer_Zbb5", "summary": "This paper proposes several improvements to training dynamics, learning, and inference for Rectified Flow models (e.g. initialization from the EDM model, loss reweighting, new loss function, refinements on noise-image pairs, etc.). Thorough ablation studies verify the effectiveness of the proposed improvements and the final experiment results are competitive among diffusion and flow-based models for few-steps sampling.", "questions": "I do not have other questions.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-written and easy to follow, with each step of the ablation study clearly described and the results effectively presented.\n2. The final experiment results after combining all improvements show lower FID and lower straightness compared with other diffusion-based and flow-based methods.", "weaknesses": "1. The projected pairs operation, which projects the coupling $Q_{01}^1$ to $\\Pi(P_0, P_1)$, complicates the training pipeline (GAN training and solving Sinkhorn) while yielding only marginal performance improvements, as indicated in Table 6 (with a maximum FID improvement of 0.02, which is not even statistically significant).\n\n2. The authors frequently use the notation 'OT ODE' and 'OT map' throughout the paper. However, as noted in references [1] and [2], the iterative rectified flow does not generally converge to the optimal transport mapping in dimensions higher than or equal to 2. Furthermore, the algorithm presented in this paper corresponds to the 2-rectified flow. Therefore, the use of the 'OT' notation is not appropriate.\n\n[1] Rectified Flow: A Marginal Preserving Approach to Optimal Transport. arxiv 2209.14577\n\n[2] Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow. arxiv 2209.03003", "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": "1frmlmSEIk", "reviewer": "Reviewer_VmAn", "summary": "This work addresses the performance degradation of ReFlow without violating its theoretical foundations. ReFlow aims to straighten the trajectories of ordinary differential equations (ODEs) by flow-matching between marginal distributions connected by a previously trained flow ODE. Theoretically, an infinite number of ReFlow updates should result in a straight trajectory that enables perfect translation between marginals with a single function evaluation. However, in practice, ReFlow suffers from a drop in sample quality, and previous attempts to mitigate this have relied on heuristic methods like perceptual losses, which deviate from the theoretical framework and may not preserve the underlying probability distributions.\n\nThe authors generalize the ReFlow training loss and categorize design choices into three key groups: training dynamics, learning, and inference. They analyze previous practices, identify potential pitfalls, and propose improved techniques within the theoretical bounds. Extensive ablation studies on datasets such as CIFAR10, FFHQ, and AFHQv2 demonstrate that their methods consistently improve sample quality and Fréchet Inception Distance (FID) scores. Ultimately, they achieve state-of-the-art results in fast neural ODE-based generation without relying on heuristic methods, offering a theoretically sound and practical approach to accelerating generative models.", "questions": "- As described in lines 100-111, the original ReFlow algorithm works by iteratively training $D^{n+1}$ using the coupling induced by $D^{n}$. Does Simple ReFlow proceed analogously? If so, how many iterations are used? It would be good to include this information in the paper.\n\n- The authors propose two improvements with regards to training data: forward pairs and projected pairs. Hence, there are three ways to obtain trajectories: backward pairs, forward pairs and projected pairs. How are the three approaches combined? For example, what is the approach used to obtain the results in Table 8?\n\n- The projected pairs approach is very interesting. The authors propose one way to solve the projection problem, which involves recasting the constraint on the marginals of $\\Gamma_{01}$ as a penalty term, which forces them to use GANs. It seems that solving a Schrodinger bridge problem between P_0 and P_1 with an appropriate regularization term would also achieve the final goal of the paper, but of course one then loses the connection with ReFlow. Are there other alternatives that may be simpler? Also, it would be good to provide a more detailed explanation in Subsec. D.3, e.g., more details in how problem 42 is solved.\n\n- What do the authors think is the improvement that is most responsible for the performance boost that is observed? As a reviewer, I think the training data improvements may be the most critical ones, but it would be nice to comment on that in the Conclusion section.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper does a very good job at ablating design choices to improve ReFlow. \n- The experimental results are state-of-the-art, by a good margin.\n- The paper can be very helpful for practitioners that face similar choices when training their models.", "weaknesses": "- Each of their contributions may be labeled as incremental, which means that the whole paper could also be labeled as incremental. Regardless, it is a very valuable paper to the community because it sheds light on several empirical choices one must make.\n\n- The statement “This again ensures that the loss is minimized when $D_{\\theta}$\noutputs the posterior mean” in line 184 is false. In fact, what holds is \n$D^*(x_t,t) = \\phi^{-1}(\\mathbb{E}_{x_0 \\sim Q(\\cdot|x_t)} [\\phi(x_0)])$.\nWhen $\\phi$ is far from the identity, $D^*(x_t,t)$ may be far from the desired point. This does not makes the empirical findings less valuable, but it does invalidate the theoretical claims that the authors make about the training loss based on $\\phi$ adhering to FM theory.", "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)