Flow Matching with Semidiscrete Couplings
Authors: Alireza Mousavi-Hosseini
Stephen Y. Zhang
Michal Klein
marco cuturi
Venue: ICLR 2026 Poster (2026)
Keywords: flow matching
optimal transport
semidiscrete optimal transport
Primary area: generative models
TL;DR: We use semidiscrete optimal transport to couple noise and data in flow-based generative models.

Abstract: Flow models parameterized as time-dependent velocity fields can generate data from noise by integrating an ODE.
These models are often trained using flow matching, i.e. by sampling random pairs of noise and target points $(x_0,x_1)$ and ensuring that the velocity field is aligned, on average, with $x_1-x_0$ when evaluated along a time-indexed segment linking $x_0$ to $x_1$. 
While these noise/data pairs are sampled independently by default, they can also be selected more carefully by matching batches of $n$ noise to $n$ target points using an optimal transport (OT) solver.
Although promising in theory, the OT flow matching (OT-FM) approach (Pooladian et al., 2023, Tong et al., 2024) is not widely used in practice. 
Zhang et al. (2025), pointed out recently that OT-FM truly starts paying off when the batch size $n$ grows significantly, which only a multi-GPU implementation of the Sinkhorn algorithm can handle.
Unfortunately, the pre-compute costs of running Sinkhorn can quickly balloon, requiring $O(n^2/\varepsilon^2)$ operations for every $n$ pairs used to fit the velocity field, where $\varepsilon$ is a regularization parameter that should be typically small to yield better results.
To fulfill the theoretical promises of OT-FM, we propose to move away from batch-OT and rely instead on a semidiscrete formulation that can leverage the fact that the target dataset distribution is usually of finite size $N$. The SD-OT problem is solved by estimating a dual potential vector of size $N$ using SGD; using that vector, freshly sampled noise vectors at train time can then be matched with data points at the cost of a maximum inner product search (MIPS) over the dataset.
Semidiscrete FM (SD-FM) removes the quadratic dependency on $n/\varepsilon$ that bottlenecks OT-FM. SD-FM beats both FM and OT-FM on all training metrics and inference budget constraints, across multiple datasets, on unconditional/conditional generation, or when using mean-flow models.

Published as a conference paper at ICLR 2026
FLOW MATCHING WITH SEMIDISCRETE COUPLINGS
Alireza Mousavi-Hosseini Apple, U. Toronto
Stephen Zhang" Apple, U. Melbourne
Michal Klein Apple
Marco Cuturi Apple
ABSTRACT
Flow models parameterized as time-dependent velocity fields can generate data from noise by integrating an ODE. These models are often trained using flow matching, i.e. by sampling random pairs of noise and target points (xo, X1) and ensuring that the velocity field is aligned, on average, with X1 Xo when evalu- ated along a time-indexed segment linking X. to X1. While these noise/data pairs are sampled independently by default, they can also be selected more carefully by matching batches of n noise to n target points using an optimal transport (OT) solver. Although promising in theory, the OT flow matching (OT-FM) approach (Pooladian et al., 2023; Tong et al., 2024) is not widely used in practice, Zhang et al. (2025) pointed out recently that OT-FM truly starts paying off when the batch size n grows significantly, which only a multi-Gi GPU implementation of the Sinkhorn algorithm can handle. Unfortunately, the pre-compute costs of rumning Sinkhorn can balloon, requiring (n2²/e) operations for every n pairs used to the velocity field, where is regularization parameter that should be typically small to yield better results. To fulfill the theoretical promises of OT-FM, we pro- pose to move away from batch-OT and rely instead on semidiscrete formulation that can leverage the fact that the target dataset is usually of finite size N. The SD- OT problem solved by estimating dual potential vector of size N using SGD; using that vector, freshly sampled noise vectors at train time can then be matched with data points at the cost of maximum inner product search (MIPS) over the dataset with cost ONN). Semidiscrete FM (SD-FM) removes the quadratic de- pendency on n/s that bottlenecks OT-FM. SD-FM beats both FM and OT-FM on all training metrics and inference budget constraints, across multiple datasets, on unconditional/conniiion generation, or when using mean-flow models.
INTRODUCTION
Flow-based generative models (Rezende & Mohamed, 2015; Dinh et al., 2016; Kingma & Dhariwal, 2018; Grathwohl et al., 2018) can gradually transform noise vectors into structured data. Flow models parameterized time-dependent velocity fields (Chen al., 2018) uy x) can be efficiently trained using flow matching (FM) (Lipman et al., 2023; Peluchetti, 2022; Albergo et al., 2023) to yield ofate-of-the-art performance (Zheng al., 2024; Esser et al., 2024). The FM training approach proposes to sidestep the need to differentiate through costly numerical ODE integration at training time by instead minimizing velocity fields locally through a regression loss. In a nutshell, the loss is formed by sampling pair of noise/data points (Xo, x1) and random time t E [0, 1] to evaluate vo(t, x;) 2 where X, == (1 -1-18 t)xo + t X1 is a barycenter in the [xo, x1] segment.
[x] Xg
Coupling Noise Data. A crucial ingredient in the FM methodology lies in choosing the coupling of noise/data used to sample pairs. Indeed, given source and target distribution, there are infinitely many probability paths that can interpolate between them; choosing a different coupling yields a different loss and hence a different flow (Albergo al., 2023; Liu, 2022). The default formulation of FM (Lipman et al., 2024) relies on the independent coupling, where noise and data are sampled independently. While simple and cheap, this approach is known to produce ODE paths with high curvature, which requires increased compute at inference time. In practice, the quality of samples is highly dependent on the number of function evaluations (NFEs) used integrate the ODE.
* Joint first authorship, work done while AMH and SZ. were interns Apple Correspondence: mousavi @ www.ttrooooddeeoo edu; syz@ syz@syz.id.... syz.id.au; michalk@@pple.coom cuturi @ apple.com

---

Published as a conference paper at ICLR 2026
I-FM
OT-FM
SEMIDISCRETE-FM
Noise N
Data
Xi...., Xp N
Precompute: fit g E R
Train: assign noise
Y1..... Nn ~ Data
95 J N RC M0
j max, x y, 91
~ N>>> y Data
X2 + Not
min I Jmnx x yy + 50,00
X ~ + Y3"
Figure 1: I-FM (left) assigns noise data purely at random. OT-FM (middle-left) samples batches of n noise and n data points and re-aligns them with an optimal matching permutation a². These matches are, however, inherently unstable, as n points do not reflect the whole noise distribution nor the dataset. Increasing drastically n can mitigate this issue (Zhang et al., 2025), but at significant cost. Our method, SD-FM (right), solves these issues in two steps: in a precompute phase, the semidiscrete OT problem (parameterized as a vector of size N, the dataset size) is solved using SGD. At FM train time, each newly sampled noise is assigned to a data point using a maximum inner product search, Laguerre cells (Mérigot, 201 1) being illustrated in the plot. Our figure uses no entropic regularization (c = 0) and neg-dot-produet cost for simplicity, see (4) for more generality.
Optimal Transport Flow Matching. Tong et al. (2024) and Pooladian et al. (2023) have proposed using an optimal transport (OT) coupling to select noise/data pairs motivated by the dynamic least- action principle of Benamou & Brenier. Both works advocate sampling noise and data points, compute an optimal n- nn-permutation (or n X n coupling matrix, from which pairs of indices are sam- pled from), as illustrated in the middle-left plot of Figure 1. While appealing in theory, OT-FM yields limited improvements for an additional precompute cost. In a follow-up work, Davtyan et al. (2025) have proposed keep in memory the pairings returned by multiple Hungarian solvers run along training iterations, but this approach is bottlenecked by a price of (n²)) at each batch. Zhang et al. (2025) hypothesize that the mitigated results of OT-FM are due to the choice of small n: Indeed, optimal matchings on small sample sizes are inherently unstable and cannot reliably approximate matches that would appear for far larger n due to the curse of dimensionality (Hütter & Rigollet, 2021; Chewi et al., 2024). Zhang et al. (2025) propose to instead use the Sinkhorn algorithm with significantly larger n (from 256 used originally to R 10%) using multi-GPU-node implementation, carefully ablating the role of regularization. They argue that these computations precompute costs, happening independently of FM training, and therefore can be cached beforchand. However, that cost grows O(n² /e2) (Dvurechensky et al., 2018; Lin et al., 2019) for every n pairs fed to FM training. Their finding that larger n / smaller e yield even better results hinders the adoption of OT-FM as currently implemented (see also our discussion in Table l).
The Semidiscrete Approach. Our work proposes an alternative route to OT-guided FM that does go to the costly process of optimally pairing batches of noises to data. We leverage the entropy regularized semidiscrete (SD) formulation of OT (Oliker & Prussner, 1989; Mérigot, 2011; Cuturi & Peyré, 2018; An et al., 2020), which studies OT from continuous (noise) to a discrete (data) distribution. SD-OT relies on a potential vector of size N of the target, fitted using SGD (Genevay al., 2016). At FM train time, freshly sampled noise is paired to point in the dataset in O(N) time, as summarized in Figure 1. Following the presentation of background material in Section 2,
We propose in Section 3 new convergence criterion and convergence analysis of SGD for regu- larized SD-OT that is applicable to both the E = O and E >> cases.
We introduce in Section 4 our semidiscrete FM (SD-FM) method, comparing it to FM and OT-FM in terms of memory and compute. We propose a generalization of the Tweedie formula (Robbins, 1956) for flows trained using SD-FM, with a correction ó: that vanishes when E R 0 or s & co. We use to sample from a geometric mixture of distributions using their velocity flows.
We illustrate in Section 5 in a varied set of (un)conditional generative tasks the advantages of SD- FM over OT-FM and FM. We show that SD-FM results in a better pairing of noise to data (as seen in better metrics) for a negligible computational overhead relative to the cost of FM training.
2

---

Published as a conference paper at ICLR 2026
2 BACKGROUND AND RELATED WORK
Stochastic interpolants and flow matching. For two probability distributions μ. E P(RR),
let V(, C P(R" x IR") be the set of couplings between μ and v, i.e. all joint distribu- tions having μ as V as their first and second marginal, respectively. Consider souree and target pair of distributions No: in and E [(Po, a prescribed coupling between the two
For an interpolant function (1, y) + y) E Rd such that () = x and &p (x. bergo et al. (2023, Definition 2.1) prove that to a cou- pling π and interpolant function (p: corresponds bility path Pt E P(RR)) where A == P where # is the pushforward operator, and Pu Ho: PI = μ1. That interpolant can be written equivalently
y. Al-
dX, (X,,dt Xo ~ po- (1)
flow generating continuous proba-
probability flow ODE with a random initial condition (1), for a suitable time-parameterized velocity field u.- Access to ot then allows for samples Xo Po be transformed into samples Xi Jii through for the flow field v can be done by sampling (Xo, X1) π and E [0. 1]. setting X, == sp (Xo, X1) and minimizing (2).
evolution of (I). Training neural approximation
The linear interpolant (x. y) := (1 ty and the squared Euclidean distance for L are commonly used (Lippman et al., 2024, Eq. 2.3). 0
min Ex, L(X: Xo. (t). (2)
Albergo et al. (2023, Thm. 2.7) prove that the global minimum in (2) recovers ot from (1). Other interpolants or loss functions (Song& & Dhariwal, 2024; Kim et al., 2024) as well as penalizations facilitating one-step generation (Boffi et al., 2025, and references therein) have also been proposed.
Straighter flows... While the independent product coupling τΙ == po @ μi is the most widely used in practice (Lipman et al., 2024, $4.8.3), it may not be efficient in the sense that it induces curved flows. This means that adaptive ODE solvers consuming many neural function evaluations (NFEs) must be used to generate data at inference time. This is hardly surprising from OT perspective (Santambrogio, 2015, $4), because the independent coupling ny is known to incur a high transportation cost, where for any valid coupling æ E [[uo, ui) that cost is usually defined as C(T) = Ex.....)) | - 112 Concretely, although the flow v. learned from कy might validly link (4o μl, its kinetic energy fo [|vell/:(...d dt is high. Benamou & Brenier (2000) show how optimal transport arises naturally as the least-action dynamics that transports μo0 to fn- That map is exactly the Monge map, and the corresponding optimal flow paths are straight lines that would be trivial to integrate. To benefit from this insight, two families of works have emerged.
using Reflow. Liu (2022) proposes to straighten a pretrained flow model (Liu et al., 2022) using Reflow, a method that forms noise/generated-dinn pairs, used subsequently improve that model using FM. For Reflow to work, the pretrained flow model must be good enough to generate approx- imately the target distribution, and significant compute must be spent on generating each reflow step (through ODEs) large number of "virtual" data points. While successful (Liu et al., 2024), Reflow a costly recursive post-training procedure that uses FM training as an inner iteration.
using pairs sampled from OT couplings. Pooladian et al. (2023) and Tong et al. (2024) (see also discussion in Lipman et al. 2024, $4.9.2) propose OT-FM, a much lighter approach that simply replaces my by a coupling π that should, ideally, approximate the OT coupling x* arg mTn C(T). Compared to Reflow, which uses FM as an inner step, OT-FM adds pre-processing effort to FM training, to select better noise/data pairs, Because the ground-truth OT coupling "¹ is never known, however, they use OT matrices: sampling n 256 points from both No and μ1, they match them using an OT solver (e.g. the Hungarian algorithm, Kuhn 1955) and feed these pairs to the flow loss (2). Davtyan et al. (2025) proposed then FOOM-CFM, a hybrid approach that stores buffers of paired noise/data samples (using the Hungarian algorithm). Unfortunately, OT- FM yields only modest gains. Zhang et al. (2025) posit that this is due to the well-known statistical bias arising when approximating continuous OT couplings using samples (Hütter & Rigollet, 2021). To reduce the effect of this bias, they use significantly higher n handled with a multi-GPU, multi- node implementation of the Sinkhorn algorithm, carefully ablating the role of the regularization parameter. While encouraging, their results show better performance as n grows and Æ shrinks to O, an explosive cocktail of hyperparameters since the cost of running Sinkhorn (1964) grows as (n2/ Fundamentally, these approuches are bottlenecked by their reliance on the discrete view of OT problems- that is, they rely repeatedly on matching i.i.d. samples of noise and data.
3

---

Published as a conference paper at ICLR 2026
Conditional generation using OT. Chemseddine al. (2024); Kerrigan et al. (2024); Hosseini et al. (2025); Baptista et al. (2024) provided a noteworthy extension of OT-FM in the context of conditional generation, where each data point E R is paired with corresponding condition 2 E RP. The data is now an augmented random variable (X1, Z1) ~ jú E P((RáPP) with a marginal distribution i over conditions. The goal remains generate samples from noise in this setting, using the notion of condition-preserving triangular maps T (x, z) + (Tx(x), z) that couples conditional noise distribution 1i (0, @ F1,z 2 P((R(t) with augmented data in ū1. Kerrigan et al. (2024, Prop. 1) (see also Baptista et al. 2024, Theorem 2.4) state that such triangular maps can be reduced to transport of distributions supported on the product space Ed following exactly OT-FM principles, implemented as ODE trajectories using a conditional vector field X; (t, X, z). OT-FM then consists in sampling n noise and condition vectors, along with n data points and their known condition, which are then optimally paired using an augmented cost:
R. but conditioned on a vector z E IK of interest. OT-FM has a natural generalization
d(x, z), (x', zz) cx (x, x') + (cz(z,z'), 3>0.
Semidiscrete optimal transport. In the scenarios envisioned in this work, the noise measure M4 continuous, while the target measure μlr of data is finite. This setting fits the stochastic optimization approach outlined in (Genevay et al., 2016, $2). defined for two probability measures μu, v E P2(()). a cost function e R x Rd + R and entropic regularization e 2 0 as:
min EEI(P))
: [c(X, Y) + =KL(T @v)
(Pe)
This problem can be written in semidual form (Cuturi & Peyré, 2018), leveraging the fact that v is a finite measure, == = bydy, where b = (bi,..., bx) is a probability vector. For potential vector g (g). 191.....9hh E INN, we define the soft-c transform (Peyré & Cuturi, 2019, $5.3) as:
E log Σラ by exp
fg.e x E Rd
6
>>
E
R
(4)
- max [9; - c(x,y;)],
6 0.
One can solve (P₂) by solving the concave maximization semidual problem parameterized entirely by g, that Genevay et al. (2016, Alg. 2) proposed to solve using averaged SGD (our Algorithm I)
max F. (g) = E Ex-xxxaeeee L (X)] (b,g).
(Se)
BERN
For any 2 0, the values of (P and (S.) are equal. Moreover for E > 0, given a maximizing potential g" (gi..... ........)))) of (S.), ) one can recover a minimizer मं of (P.) that has density
dx. yz) = exp [gg.(()) + gj v,) by du(x).
(5)
3 IMPROVED OPTIMIZATION FOR SEMIDISCRETE OPTIMAL TRANSPORT
We adapt the approach of Genevay et al. (2016) to the much larger scales used in this work through two novel contributions: we propose an unbiased convergence criterion to monitor convergence, and theoretical analysis to inform the computation of SD-OT that works also for the case 0. Genevay et al. (2016) had not considered that case, which turns out to be the most promising when using SD-OT within the SD-FM methodology presented next in $4
Marginal Estimation for SD Couplings. The primal-dual relationship in (5) holds at optimality, but can also be used to create a couplingby plugging a vector E N into (5) to define a joint probability that only satisfies the first marginal constraint. The exponentiation in (5) can be rewritten using the softmax operator mapping a vector z = (21...., 2N) to a vector in the simplex A N.
bb (z)
exp(s,/e) E by exp(at/e]
& 8>0 0,
1[jEarg max, sr|b,
e 0.
(6)
== 1[kearg max, ze|ba
This helps to define coupling ne.g that arises from partial optimization of the dual, as
(x, == (x),(x)),d((x), where Se.erese.g(x) Sx.g(x) == (b, - c(x,)))))) E A By construction, the first marginal of ne.g is pc: Indeed, the measure Te.g (x..g(x,, yy) in (7) sums to du(x)
(7)
when integrated against the data for a fixed since summing w.r.t. index j is akin to summing weighted soft-max distribution. On the other end, Te.g would be an OT coupling if and only its
4

---

Published as a conference paper at ICLR 2026
1000==/sssxxxx// - E A (8)
second marginal, defined m(g) through an integral in (8), coincided with b, by analogy to the Sinkhorn algorithm (Peyré & Cuturi, 2019, $4.14). Note that Te solves (Pa) if g is a maximizer of (S.).
m(g) = [e, (TR." y;)1;
&EER
Convergence Criterion. To our knowledge, no convergence criterion has been considered in the stochastic SD-OT literature (An et al., 2020, Alg. (Genevay et al., 2016, Alg.2). Recall that is the solution to the dual problem and only if mgg) b, i.e. the v-marginal constraint satisfied.
A good criterion for convergence could be given by a distance between mgg) and b. A natural can- didate could be total variation, TV(
100
|| - bll1. as often used to track the conver- gence of the Sinkhorn algorithm. Unfortunately for continuous 1,, the expectation used in m(g) would need be replaced with an empirical mean, leading to biased estimate, because that expectation would be inside the norm. Reducing that bias would re- quire number of samples scaling linearly with N, making it prohibitively large. Fortunately, one can efficiently obtain unbiased estimator for the x² divergence, defined for two vectors P.q E AN as 2 (pllq) = (p;/g)) 1. That x² divergence can be formulated
(m [)
10-1
Adagrad
10-2
dooryy
10
10
of noise samples / N
Figure 2: SD-OT convergence: X2-
divergence vs. SGD optimization steps for ImageNet-32, averaged over 3 seeds. See Appendix for details.
expectation, highlighted the fact below, proved in A.1:
Fact x² (m() [] SK
(x
(9)
When is uniform, by Lemma 7, the above x² divergence is also equal to the rescaled squared norm of the semidual gradient. From (9), we propose in (10) an unbiased estimator using a batch i.i.d. samples X1,-.-a> XB y that be computed in D(NB B) time. Figure 2 shows that are able to effectively track X² as it decays towards zero along SGD updates.
(b(g)) [b)
BB-T Eje by
(E (8.. (x:)l;)² Ex(ta((x))33) -
(10)
Convergence Analysis. Studying the case 0 requires some additional regularity assumptions, that we use to state the convergence guarantee of SGD for SD-OT Theorem 2. Assumption 1. Suppose c(x,y) = xxy. y, 8 = min |ly, y > O, and admits density w.r.t. the Lebesgue measure. Additionally, suppose the maximum μ-surface area of any convex subset of Td is bounded by Cinax, i.e. Joa (xx)d5(xx) C for all convex A C Rdd S being the 1-dimensional Euclidean surface measure. When μ = -N),, Cma < /4 (Ball, 1993). Theorem 2. Suppose either > or Assumption
satisfied. Let LE = 1/s for > and Lo = C.maa /6 else. For any K E N, let π₁s = /A be a constant learning rate, where A = F F Let gs denote the SGD iterates with go = 0, and let ~ Unif (O.. 1}). Let g: be optimal dual solution Equation (S.). Then, for any & 2 and taking expectations w.r.t the randomness in noise samples and t,
and E(Ca( - C. x (1 2 1/2 (4
[|
ming by
If the number of iterations is not fixed priori, one can use the decaying schedule 7 which replaces 1/K with log K/K in the bounds above. The above statement shows how the second marginal of Te.g approximates the correct marginal v, and how its (entropic) transport cost approximates the optimal cost as K grows. Interestingly, both entropic E > 0 and unregularized 0 cases (pending Assumption 1) valid, supporting our use of g 0 throughout the paper.
ALek)))
5

---

Published as a conference paper at ICLR 2026
4 SEMIDISCRETE TRANSPORT FOR FLOW MATCHING
We discuss the comparative advantages of SD-FM over OT-FM and I-FM. We follow with a gen- eralized Tweedie identity for flows trained with regularized SD-FM that incorporates a correction term. That term vanishes when either E (I-FM) but also, more surprisingly, when e - 0.
SD-FM vs. OT-FM vs. I-FM. Building on Section 3, the OT problem is solved between Gaussian
noise Pi and the data measure μn supported on (x (x(")__ :....... x ). This framework can also accom- paired with p-dimensional condition, see Section 2. We follow Zhang et al. (2025, $3) and use the dot-product cost y)x,y) := -x"y y (augmenting it with temperature 3 for conditioning (3)) and rescale with the cost matrix standard deviation for a sampled reference batch. As illustrated in Figure 1, SSD-FM works in two steps, (1) solve SD>OTT between 14 and 141 before flow training. storing g TRN solving (S²); (2) in the FM train loop. pair a freshly sampled noise Xy to data x
modate the conditional generation setting, in which observations
where k ~ g (xo) E A
As recalled in (7), for E > O, this amounts to categorical sampling among N points, while for E = this reduces to an as- signment to the k*-th point in the dataset, k* = arg max; g + (Xo, xi (+)) (or sample uniformly among ties if they arise). The lat- ter operation is slightly faster than categor- ical sampling, as observed in Fig. 4. The differences between these three approaches 5, 6. In
Coupling
Memory
Preprocessing
FM Training
I-FM
NE
OT-FM
( + O(dn/e"))
FM
2NE
O(NEdn/E")
NEO
SD FM
N
Ndk
NE dN)
Table 1: Complexity of V/T//D-FM given N:
datasct size, d: data dimension, E: FM training epochs, nc Sinkhorn batch size, E: entropic regular- ization, θ: cost of computing FM loss' gradient for pair. K: SD SGD steps. OT-FM* is the cached vari- ant of OT-FM. θ dominates all costs, N n, e z O.
highlighted in Algorithms a nutshell, SD-FM uses a simple SGD pre- compute effort fit g, with a minor lookup effort to pair each noise with datapoint, rather than making potentially costly calls to Sinkhorn in OT-FM. That lookup is a MIPS procedure (Shrivastava & Li, 2014) when = O. solved exactly, incurring a Nd cost (noise x data dot-products) and max over N values. This might be sped-up with the use of approximate MIPS, left for future work.
Score Estimation and Guidance A key property that connects flow matching and diffusion mod- els is that learning the marginal velocity field under independent coupling and Gaussian noise is equivalent to estimating the marginal score along the probability path. Namely, let p denote the law of X, (1 11-t) x) + tx1, where Xo Pa = μ and X1 ~ P1 independently. Recall that,
V log (x) = Vx,
x, log puj (X, X1) X; x] (x)-x
(11)
1-t
Note that the above is Tweedie's formula since the flow matching velocity that transports Po to P1 satisfies (1 - (x)),xx) = E X, X, x]. While (11) relates score and velocity, the second equality above relies on the independence of Xo and X1; it may no longer hold under general couplings. Surprisingly. we find that for SD-OT couplings (11) still holds approximately for E z O. Proposition 3 (Generalized Tweedie's Formula). For any g E RNN, ler X X1 Te.g. let Pe denote the density of X, = (1 X) + LX1, 1. <. 1, and define ( F E| Xo Xx = x). Suppose Pa (O, Ia). Then, V log (x) 3x.(x) x+ (1 (1-1)8
holds for any & 2 0, where
(1-1)²
d. = -E x (o. X1) + c( X1) Xo. X, =x] xo 0, and du 0.
For small E and ignoring other constants, one can show that ||5_1 S e e-1/e/e, thus going to zero as E + O. Intuitively, this happens because X; becomes increasingly determined by Xo as + O. Note that ε + do recovers the independent coupling, and we obtain the same score as expected.
Correcting the Guidance. To see a direct application of Proposition 3, consider the case where we have two flow models (pı,))== and (p2:)l=o that are generated by velocity fields V1, and U2.r respectively. For simplicity, we consider the case where both flows start from po. Our goal is to sample from pf²) pi.eP2, / for some 7 2 O, where z(77 is the normalizing constant. One example is classifier-free guidance (Ho & Salimans, 2022) where P1 is a conditional model and p2 unconditional or, more generally, autoguidance (Karras al., 2024), where P2 is a weaker model compared to P1. While in practice we often directly combine the corresponding velocity fields, the
(7) and we need additional "correction" as laid out below. samples are not guaranteed be from Pt

---

Published as a conference paper at ICLR 2026
64
Curvature
0.04
32:
Coupiin
JFM
FID
16
SDFM
Sofeerr
Doprit
NEE NEE 1IEE-15--
Euker NFE
0.008
Faler NPE m:
10
10
10'
10a
10
10 Chi-Squared Error 8 (m((g)
10
10a
Chi-Squared Error (m (g)
b)
b)
Figure 3: Better SD Potential Estimation Better Curvature and FID: On ImgN32, convergence of dual potential SD-FM (s 0) curvature and FID; I-FM is shown lines. Note that curvatures of different solvers are computed on different trajectories, hence they are not comparable.
L.dd N== ~ Po. and define (x! (7: ) by the ODE Proposition 4 (Informal). Let (X) (xx
dx((x:) = (( (T)))) + (1 2) (7)1)) Jdtt E [r].
Furthermore, define the weights
w(ii) = --(7-19 (VL,((x( U2 (X(7)))) V log Pi,, (x( V log P2, (X, (x(r:)))dd
Draw I softmax(( (fw(n))) and let X, = x (7.)) Then, Law(X,) + ()) as + co.
This result can be seen as the flow matching variant of correctors used for diffusions (Bradley & Nakkiran, 2024; Skreta et al., 2025). Concretely, to sample from p², one must generate a number of i...d. samples using the combined velocity field, and draw the outcome according to the weights. However, calculating weights requires knowledge of the scores log Pt. Thanks to Proposition 3, either for the case of independent coupling (e = o0), or unregularized semidiscrete couplings (e 0), we can calculate them using the velocity model and obtain
( = > -1) /
1 [x.(((((T))) 22 22, xx/00))|1
EXPERIMENTS
5.1 IMAGENET PETFACE: UNCONDITIONAL AND CLASS-CONDIITS CONDITIONAL GENERATION.
We consider generation in raw pixel space for the ImageNet (ImgN) train fold (Deng et al., 2009), augmented by horizontal flipping (N=2.56 M images) in (32,32) d 3072, and (64,64) d = 12 288 resolutions, and the PetFace dataset (Shinoda & Shiohara, 2024) of animal faces (N= 1.27 M im- ages) in size (64,64). We measure generation quality using FID (Heusel et al., 2017) and flow cur- vature (Lee et al., 2023). I-FM, OT-FM and SD-FM training differ exclusively in the way noise-data pairs are sampled, all other aspects of FM training stay identical, using Pooladian et al.'s setup.
FID Potential Quality. We investigate whether spending precompute effort to get better poten- tial for SD-FM translates into better model performance on ImgN-32. Concretely, we record six different iterates of g along the SGD iterations in Algorithm 1, tracking their convergence criteria {²(m(g)] We train six SD-FM models using these six vectors to assign noise to data. We plot in Figure the FID and curvature across 4 ODE solvers. We observe that FID improves and curvature shrinks with lower x² confirming the main hypothesis of our paper that sharper SD-OT provides more useful coupling for FM, with gains that seem to saturate when X2 Rs 0.05. The runtime needed node of 8 H100 GPUs.
for SGD converge is of the order of 12 hours
FID vs. Pairing Time Cost. We consider the FIDs on ImgN returned by I-FM, replicate the setting in (Zhang et al., 2025) for OT-FM using E {0.01, 0.1) and E 12 2 219). and finally SD-FM using a fixed I budget for SGD with & E {0, 0.010 0.1). We show uncurated samples generated by I-FM and SD-FM in Figures and 11. We provide results (second line) in ImgN-32 where coupling computations happen in the = 500 principal components of the data, as proposed by (Zhang et al., 2025): this impacts all complexities presented in Table 1 by substituting k for d. We stress that PCA
7

---

Published as a conference paper at ICLR 2026
FID (Unconditiona)
FID
(Clas--conditoonal)
Fulerr
Fuler 8
Faler 16
DepriS
Fuler 4
Fuler 8
Fuler
Dopri5
37.90
9.10
12.95
7.39
ImageNet-64
SD-FM
45.62
23.75
15.02
8.42
26,04
12.06
7.32
3.63
(PCA 500)
0.01
45.68
23.91
15.10
8.45
26.15
12.26
7.51
0.1
45.90
24.05
15.35
8.58
25.92
12.16
7.42
3.64
I-MM
56.53
26.85
13.38
1.26
47.66
21.91
11.56
1.09
PetFace
SD-FM
==
20.54
12.77
7.50
1.26
19.10
12.10
7.13
1.05
(full d=12k)
0.01
20.58
12.91
7.67
1.23
19.32
12.10
7.06
0.1
20.96
12.96
7.59
1.18
19.46
12.20
7.18
1.04
Table 2: FID for ImageNet 64x64 and PetFace 64x64 unconditionallllas-cooodii generation.
only used to speed up noise/data pairings FM training is unchanged. For ImgN-64, the costs of running OTFM for large n is too large in full df 12288, we only report the PCA (A 500) setting. SD-FM improves FID compared to I-FM, specially for cheaper fixed-step Euler solvers, for far smaller time overhead compared to OT-FM. To put these overhead costs in perspective, display in a red-dashed line the average time @ (see Table 1) needed to compute the gradient of the FM loss one pair. In summary, for pairing overhead cost negligible w.r.f. O, SD-FM delivers dramatically better performance than I-FM, especially for small inference budgets.
ImageNet 32x32
Euler NFE 4
Euler NFE-8 8
Euler NFE-
DopriS
35
67.70
16
24.66
64
12.03
5.13
Full dim.
9.0
4.4
FID
4a
30
14
8.5
4.2
10
10
10
10
10
10
10
10
10
10
10
10
PCA. k= 1000
67.70
16
24.
12.03
5.13
9.0
4.4
FID
30
14
8.5
4.2
10
10
-5
10
10
10
10
55
10
10
3
ImageNet 64x64
PCA, k 500
79.95
26
37.90
16
19.68
9.0
50.0
FID
47.5
24
8.5
15
10"
10
10
10
Time per pair (s)
Time pair (s)
Time per pair (s)
Time pair (s)
8
0.01
0.1
SD-FM O
OT-FM, == 32768 A 131072
524288
I-FM
Figure 4: FID vs. time needed to form pair, when training I-FM, OT-FM (varying batch sizes n) and SD-FM. We use
(SD) only), 0.01, 0.1. Couplings are computed using full d or PCA
space. Red-dashed lines show the per-sample time needed compute the gradient of the loss for one pair. SD-FM yields significant improvements for negligible overhead.
Class-conditional ImgN and PetFace. We train class-conditional flow models using SD-FM, fol- lowing the principles presented in Section 4, taking the condition z to be a one-hot vector. ImgN and PetFace consist respectively of 1000 and 13 classes. As in the unconditional case and Figure 4, we show FID for ImgN as a function time-per per--air for generation in Figures 7 and 8. Compared to the unconditional setting, find that the gap between OT-FM and SD-FM is more pronounced, owing to the slower convergence of Sinkhom. At 64x64 resolution, the cost of matching points for OT-FM at large sizes becomes great (over 10 days) and is therefore not shown. See results in
8

---

Published as a conference paper at ICLR 2026
Table 2. Generated samples shown in Figures 10, 12, 14. Class-conditional results in Table 2 are consistent with our observations for ImgN, with SD-FM outperforming I-FM in all cases.
5.2
CONTINOUS-CONDDTIONA GENERATION: CELEBA SUPER-RESOLUTION
As an example application to continuous- valued conditions, we apply SD-FM to train flow models for conditional sampling in im- age super-resolution (SR) problems. We use the CelebA dataset (Liu et al., 2015) rescaled to 64x64 resolution, containing 162k train images. We downsize them to 16x16 (4x SR) or 8x8 (8x SR) and ap- ply Gaussian noise with standard deviation E {0.1, 0.2) to form condition and run SD-FM with continuous data and conditions described in Section 4. We compare to I-FM trained by sampling Xo ~ (0, La) @ jı,z and from the augmented data distribution of images x paired with their corrupted version z. We sample from the trained model, conditioned on corrupted images from CelebA validation set to obtain a high-resolution reconstruction.
Noise
0.1
Noise 0.2
LFM
SD.FM
LFFM
SD-FM
0
0.01
0.01
PSNR (1)
21.17
21.38
21.41
20.04
20.44
20.42
4x SR
SSIM (f)
0.67
0.69
0.69
0.61
0.63
0.63
PSNR (1)
17.52
17.87
17.94
16:50
17.35
17.34
8x SR
SSIM (1)
0.50
0.51
0.52
0.44
0.48
0.48
Table CelebA super-resolution results.
Since the ground-truth image is known, we rely the peak signal-to-noise ratio (PSNR) and structural similarity index measure (SSIM, Wang et al. (2004)) measure reconstruction accuracy. Results are summarized in Table and we show sam- pled reconstructions from the validation set in Figure 16.
0.35
0.175
Precision
0.30
0.150
Recall
SE-RM
0.25
0.125
5.3 GUIDANCE: IMGN32
0.20
0.100
For ImgN-32, Figure 5 shows the effect of r, where r is gener-
0
2
the CFG baseline without any correction and ating from the conditional model without guidance. Increasing improves generation quality (higher precision) but degrades diversity (lower recall). We use the precision and recall calcu- lation given by Kynkäänniemi et al. (2019).
Figure 5: Precision & recall on
ImgN-32 vs. of guidance sam- ples r. 7 2, NFE = 4, E 0.
5.4 ONE-STEP GENERATION WITH MEAN-FLOW LATENT SPACE IMGN256
Recently, consistency models (models that incorpo- rate trajectory consistency regularizers in their train- ing objective) have become popular choice for few- step generation with diffusions flows (Song et al. 2023; Song & Dhariwal, 2024). We demonstrate that the benefits of SD-OT couplings go beyond FM by testing them on the MeanFlow (MF) model of Geng al. (2025), noting that different consistency for- mulations can be unified and treated similarly (Boffi al., 2025; Sabour al., 2025). Figure compares SD-OT and independent couplings for training MF: we use the same setup Geng et al. (2025) train an unconditional DiT-XL/2 model for 300 epochs on ImgN-256x256 in latent space and M/2 for 40 epochs. Note that SDOT in latent space may be easier than pixel space owing to the Gaussianity promoted in auto-encoder latents, see 1ig.18.
I-MF XI.-2
SD-MF XL-2
I-MF
SD-MF M2
NEE
Figure 6: FID on ImgN-256 vs. NFE for
mean-flow
LIMITATIONS AND DISCUSSION
Some aspects of the SD-FM pipeline presented here leave room for improvement. Figure shows that computing the SD-OT potential to reasonable degree of precision impacts the performance of SD-FM. That SD-OT preprocessing cost occurs only once prior to FM training; only depends on the noise distribution and data, and is typically much cheaper than the FM training cost itself. However, if SD-FM were to be applied to datasets billions N of points, this could become a challenge (though still less than FM training). Because SD-OT is concave maximization prob- lem, one could leverage recipes such as batching, momentum or E-tempering. For z > 0, cat- egorical sampling over N- softmax vector would be costly, which is why highlighted

---

Published as a conference paper at ICLR 2026
the 0 case that can leverage MIPS and fast retrieval tools (both to train g and to pair fresh noise). For these reasons, and due to only very minor differences in performace, we would rec- ommend that users use SD-FM with 0. Finally, having in mind dataloader efficiency, one may want to flip the tables our sampling approach: Start from the j-th data point in memory. and match it with Gaussian noise restricted (Wu & Gardner, 2024) to be in its Laguerre cell
C = fx: x" (xx") x(()) + gu 9k 2 O, vk + 1}. While we have shown successful applications of the SD coupling preprocessing approach to extended FM methods such mean mean-floww guid- ance, SD couplings could be course used within even more complex approaches building upon FM, such Reflow. Our focus
ablating the impact of the SD coupling the indepen- dent coupling in the most widely used FM setup. Assessing how SD couplings interact with more advanced orthogonal methods is left for future work.
In concurrnt work, Kong et al. (2026) also demonstrate the utility of SD-OT for improving flow settings. comparison, we introduce
models, focusing on the unconditional and an efficient convergence criterion for using stochastic optimization to solve SD-OT, provide a theo- retical convergence analysis that handles both and & > 0, and generalize Tweedie's formula to SD-OT couplings, providing a score estimator that can be used in guidance techniques.
CONCLUSION
OT-FM can been used to guide the training of flow models by optimally pairing batches of n noise data points, with results only truly paying off for very large n. We proposed the more cost- efficient semidiscrete OT paradigm, building on two-step approach (fit potential vector as a one- off cost, assign noise when training through MIPS) that not bottlenecked by batch size n. Our experiments in various settings, both in unconditioned and conditioned scenarios, show much improved metries over I-FM and orders of magnitude cheaper compute compared OT-FM.
REPRODUCIBILITY STATEMENT
We explain all the necessary details to reproduce our results in Section and in Appendix C. We provide the code needed compute and sample from semidiscrete couplings an addition the existing open source OTT-JAX toolbox (Cuturi et al., 2022) available at ttttps: /tgg://tttttt github...... com/ ott ott-jax/oe /ott.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Michael Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: unifying framework for flows and diffusions. arXiv preprint UrXIv:2303.08797 2023.
Dongsheng An, Yang Guo, Na Lei, Zhongxuan Luo, Shing-Tung Yau, and Xianfeng Gu. Ae- ot: A new generative model based on extended semi-discrete optimal transport. International Conference on Learning Representations, 2020.
Keith Ball. The reverse isoperimetric problem for gaussian measure. Discrete & Computational Geometry, 10(1):111-42)) 1993.
Ricardo Baptista, Bamdad Hosseini, Nikola B Kovachki, and Youssef M Marzouk. Conditional sampling with monotone gans: From generative models to likelihood-frie inference. SIAM/ASA Journal on Uncertainty Quantification, 12(3):868 (10-80089000 2024.
Jean-David Benamou and Yann Brenier. A computational fluid mechanics solution the monge- kantorovich mass transfer problem. Numerische Mathematik, 84(3):375-393, 2000.
Nicholas Matthew Boffi, Michael Samuel Albergo, and Eric Vanden-Eijnden. Flow map matching with stochastic interpolants: mathematical framework for consistency models. Transactions on Machine Learning Research, 2025. ISSN 2835-8856.
Arwen Bradley and Preetum Nakkiran. Classifier-free guidance is predictor-corrector. arXiv preprint arXiv:2408.0900 2024.
Jannis Chemseddine, Paul Hagemann, Gabriele Steidl, and Christian Wald. Conditional wasserstein distances with applications in bayesian ot flow matching. arXiv preprint vxXhv:243.118005 2024.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Sinho Chewi, Jonathan Niles-Weed and Philippe Rigollet. Statistical optimal transport. arXiv preprint 1xXx:22407.18666 2024.
Marco Cuturi and Gabriel Peyré. Semidual regularized optimal transport. SIAM Review, 60(4), 2018.
Marco Cuturi, Laetitia Peng-Papaxanttos, Yingtao Tian, Charlotte Bunne, Geoff Davis, and Olivier Teboul. Optimal transport tools (ott): A jax toolbox for all things wasserstein. arXiv preprint arXiv:2201. 12324. 2022.
Aram Davtyan, Leello Tadesse Dadi, Volkan Cevher, and Paolo Favaro. Faster inference of flow- based generative models via improved data-noise coupling. In The Thirteenth International Con- ference on Learning Representations, 2025.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hier- archical image database. 2009 IEEE Conference on Computer Vision and Puttern Recognition, pp. 248-255, 2009. doi: 10.1 109/CVPR.20099 2000555220009500888888
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Real NVP: Real-valued volumelume pre- serving transformations for density estimation. arXiv preprint arXiv: 08803, 2016. Intro- duced coupling layers and established Real NVP architecture.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(61): 2011. URL http:://jmz: org/ hwtpg//appr/aaeeee vv2 duchilla. html.
Pavel Dvurechensky, Alexander Gasnikov, and Alexey Kroshnin. Computational optimal transport: Complexity by accelerated gradient descent is better than by sinkhorn's algorithm. In Jennifer Dy and Andreas Krause (eds.). Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 1367-1376. PMLR, 10-15 Jul 2018. URL https: WtLLhtttss//ppceeeeiiiie httppp/pceengg..... mlr. wwwppppppsaep...ppp..p v80 wwwp///awwreeeeccylaaaaaa html.
11

---

Published as a conference paper at ICLR 2026
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first international conference machine learning, 2024.
Aude Genevay, Marco Cuturi, Gabriel Peyré, and Francis Bach. Stochastic optimization for large- scale optimal transport. Advances in neural information processing systems, 29, 2016.
Zhengyang Geng, Mingyang Deng, Xingjian Bai, Zico Kolter, and Kaiming He. Mean flows for one-step generative modeling. arXiv preprint 2rXiv:2505.13444 2025.
Will Grathwohl, Ricky TQ Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. Ffjord: Free-form continuous dynamics for scalable reversible generative models. arXiv preprint arXiv: 2018.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Jonathan Ho and Tim Salimans. Classificr-free diffusion guidance. arXiv:2207.12598, 2022.
arXiv preprint
Bamdad Hosseini, Alexander W Hsu, and Amirhossein Taghvaci. Conditional optimal transport on function spaces. SIAM/ASA Journal on Uncertainty Quantification, 13(1):304-338 2025.
Jan-Christian Hütter and Philippe Rigollet. Minimax estimation of smooth optimal transport maps. The Annals of Statistics, 49(2), 2021.
Tero Karras, Miika Aittala, Tuomas Kynkäänniemi. Jaakko Lehtinen, Timo Aila, and Samuli Laine. Guiding a diffusion model with a bad version of itself. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Gavin Kerrigan, Giosue Migliorini, and Padhraic Smyth. Dynamic conditional optimal transport through simulation-free flows. Advances in Neural Information Processing Systems, 37:93602- 93642, 2024.
Beomsu Kim, Yu-Guan Hsieh, Michal Klein, Marco Cuturi, Jong Chul Ye, Bahjat Kawar, and James Thornton. Simple reflow: Improved techniques for fast flow models. arXiv preprint arXiv:2410.07815. 2024.
Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1x1 convolutions. Advances in Neural Information Processing Systems, 31, 2018. Advanced flow-based generative model with invertible convolutions.
Jun Kitagawa, Quentin Mérigot, and Boris Thibert. Convergence of a newton algorithm for semi- discrete optimal transport. Journal of the European Mathematical Society, 21(9):2033-2511, 2019.
Lingkai Kong, Molei Tao, Yang Liu, Bryan Wang, Jinmiao Fu, Chien-Chih Wang, and Huidong Liu. Alignflow: Improving flow-based generative models with semi-discrete optimal transport. ICLR, 2026.
Harold W Kuhn. The hungarian method for the assignment problem. Naval research logistics quarterly, 2(1-2):83-97, 1955.
Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Improved precision and recall metric for assessing generative models. Advances in neural information processing systems, 32, 2019.
Sangyun Lee, Beomsu Kim, and Jong Chul Ye. Minimizing trajectory curvature of ode-based gen- erative models. In International Conference on Machine Learning, pp. 18957-18973. PMLR, 2023.
12

---

Published as a conference paper at ICLR 2026
Tianyi Lin, Nhat Ho, and Michael Jordan. On efficient optimal transport: An analysis of greedy and accelerated mirror descent algorithms. In International conference on machine learning. pp- 3982-3991, PMLR, 2019.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Repre- sentations, 2023. URL https: // /ttpss///ppenne eeenreview, www.w//wweeeeeeeeeeeeeeeeeee wtt.:iwwwnee/oouu/eeeeeeoeoe wtw.rruumiiddPPPPRCC.C..
Yaron Lipman, Marton Havasi, Peter Holderrieth, Neta Shaul, Matt Le, Brian Karrer, Ricky TQ Chen, David Lopez-Paz, Heli Ben-Hamu, and Itai Gat. Flow matching guide and code. arXiv preprint arXiv:2412.06266 2024.
Qiang Liu. Rectified flow: A marginal preserving approach to optimal transport. arXiv preprint arXiv:2209. 14577, 2022.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow, 2022. URL httpp:://axxiy wtt./bs////.....///////////// 2000000033300
Xingchao Liu, Xiwen Zhang, Jianzhu Ma, Jian Peng, and qiang liu. Instaflow: One step is enough for high-quality diffusion-based text-to-imaae generation. In The Twelfth International Confer- ence on Learning Representations, 2024. URL wwtpt:tt/ppppnntteeeeeeee /tpp/rpeevviieeeeeeeeeeeeeeeeeee ew net/ forum? did= 1k4y2bbDqxX.
Ziwei Liu, Ping Luo. Xiaogang Wang. and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
Quentin Mérigot. A multiscale approach to optimal transport. In Computer graphics forum, vol- ume 30, pp. 1583-1592. Wiley Online Library, 2011.
Gaspard Monge. Mémoire sur la théorie des déblais et des remblais. Histoire de l'Académie Royale des Sciences, 1781.
V.I. Oliker and L.D. Prussner. On the numerical solution of the equation ....-(...) = f and its dis- cretizations, i. Numerische Mathematik, 54(3):271-2944 1989. URL http:://udml..rrgg
doc/133318.
Stefano Peluchetti. Non-denoising forward-time diffusions, 2022. openreview net/ forum? Ld=oVEIKuugffC
URL https://
Gabriel Peyré and Marco Cuturi. Computational optimal transport. Foundations and Trends@ in Machine Learning. 11, 2019.
Aram-Alexandre Pooladian, Heli Ben-Hammn Carles Domingo-Enrich Brandon Amos. Yaron Lip- man, and Ricky TQ Chen. Multisample flow matching: Straightening flows with minibatch cou- plings. In International Conference on Machine Learning, pp. 28 28100-28127. PMLR, 2023.
Danilo Jimenez Rezende and Shakir Mohamed. Variational inference with normalizing flows. In Proceedings of the 32nd International Conference on International Conference on Machine Learming- Volume 37, pp. 1530-1538, 2015.
Herbert Robbins. An empirical bayes approach to statistics. In Proceedings of the Third Berkeley Symposium on Mathematical Statisties and Probability, 1954-1955, vol. 1, pp. 157-163. Univer- sity of California Press, Berkeley and Los Angeles, 1956.
Amirmojtaba Sabour, Sanja Fidler, and Karsten Kreis. Align your flow: Scaling continuous-time flow map distillation. arXiv preprint arXiv:2506. 14603, 2025.
Filippo Santambrogio. Optimal transport for applied mathematicians. Springer, 2015.
Risa Shinoda and Kaede Shiohara. Petface: A large-scale dataset and benchmark for animal identi- fication. In European Conference Computer Vision, pp. 19-36. Springer, 2024.
13

---

Published as a conference paper at ICLR 2026
Anshumali Shrivastava and Ping Li. Asymmetric Ish (alsh) for sublinear time maximum inner product search (mips). In Ghahramani, M. Welling, C. Cortes, N. Lawrence, and K.Q. Weinberger (eds.). Advances in Neural Information Processing Systems, volume 27. Curran Associates, Inc., 2014. URL. https: httt:s::/prrcceedingg neurips cttpiisscccppppp_feeeeee htt/e/////l///// http:://w.......e......eeeeeeee.
Richard Sinkhorn. A relationship between arbitrary positive matrices and doubly stochastic matri- ces. Ann. Math. Statist., 35:876-888 879, 1964.
Marta Skreta, Tara Akhound-Sadegl Viktor Ohanesian, Roberto Bondesan, Alan Aspuru-Guzia Arnaud Doucet, Rob Brekelmans, Alexander Tong, and Kirill Neklyudov. Feynman-Kee correc- tors in diffusion: Annealing, guidance, and product of experts. In Forty-second International Conference on Machine Learning, 2025.
Yang Song and Prafulla Dhariwal.
Improved techniques for training consistency models. In The Twelfth International Conference Learning Representations, 2024. URL https:// id=WNzy 9bRDvG.
openreview www.reeiewnneteeeee
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. In Proceed- ings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Muchine Learning Research, pp. 3221 1-32252. PMLR, 23-29 Jul 2023.
Alexander Tong. Kilian Fatras, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research, 2024.
Zhou Wang, Alan C Bovik, Hamid R Sheikh, and Ecro Simoncelli. Image quality assessment: from error visibility to structural similarity. IEEE transactions on image processing, 13(4):600- 612. 2004.
Kaiwen Wu and Jacob R. Gardner. A fast, robust elliptical slice sampling method for truncated multivariate normal distributions. In NeurlPS 2024 Workshop Bayesian Decision-making and Uncertainty, 2024. URL https: /ttptppss//opeeeeeeeeee net/ forum? htte:idddbbbeppleeeeee.e
Stephen Zhang, Alireza Mousavi-Hosseini, Michal Klein, and Marco Cuturi. On fitting flow models with large Sinkhorn couplings. arXiv preprint arXiv:2506.055 055566 2025.
Zangwei Zheng, Xiangyu Peng. Tianji Yang. Chenhui Shen, Shenggui Li, Hongxin Liu, Yukun Zhou, Tianyi Li, and Yang You. Open-sora: Democratizing efficient video production for all. arXiv preprint arXiv:24 22.204444 2024.
14

---

X, Published as a conference paper at ICLR 2026
A OMITTED PROOFS
In this section, we provide missing derivations and proofs from the main text.
A.1 PROPERTIES OF SEMIDISCRETE COUPLINGS
We begin by stating quick proof of Equation (9)
Proof of Equation (9). By definition
m(g);
2 (m(()
[b)
Σ
b;
b;
j=l
(ʃ
Σ
(x.g(xx),
-1
by
j=l
N
Σ)
( (xx))xx(() - 1.
Next, we present the proof of the score formula for our semidiscrete couplings.
Proof of Proposition 3. We first consider the case = >> In this case, the conditional law X, X1 admits density. We use the general denoising score- score-matching identity
log Pe (x() = E log Pejz (X,) x,) X, = x],
(12)
which does not depend on the Xo. X1 coupling. For the sake of completeness, we present the derivation of the above identity:
Pxx(x)
px(x) / (x
V log
pr(x)
log Pijl (xx x21 Pe]1(X xx (dxx)
pc(x)
v log Piji (x X1 (11) x).
Now, using the definition X, = (1 Xu tX, and the change variables formula, we have
x -
Lx1
P x1) = (1- t)d Poj1
As result,
x Ix1
log Pxn( x1) = 1-t
V log Poj1
[x]
1 - t
Consequently,
V log (x) =
V log Poj1 (Xo
X1) X x
Let us define (X,) = by and g(x,) = 91 where is the index of the random vector X1, ie.
X, x Recall that for e > O. the coupling is given by
-(x
(X, X1)
((Xo)u(X)).
Ex.
p(x1)- X v(xi)
15

---

Published as a conference paper at ICLR 2026
Thus, for μr (O, we have
Vx, c(x, x( x()e s o(xoo.a)) v(x()
Vx, log Paj1 Pog (Xo (1) Xo
(Xo,
+-
1
Σ
Ex, )(x,)-c((x,,)) v(x1)
= c( X1) + - V cX, X1) Xol.
Plugging this back into (12), we obtain
E
Xu
X; x
V log px()
B
X1) E|V (Xo, X1) X, x)
t X;
c
E X
x + 8.
Moreover, since the optimal velocity field in flow matching is given by the expectation the condi- tional velocity field (Lipman et al., 2023), we have
X
X
x X X =
U/ (x) = E
X. =
t
Combining the two identities above finishes the proof for the s > 0 case.
g(x,) - c(xo, x1), where we can arbitrarily break ties, by e.g. picking the maximizer with the smallest index. Then, with probability over the draw of Xo. the samples Xo. X1 ~ Po.1 are given by Xi = Ti (Xo). As a result, we can write
For the E = case, we define Ti (xx) = arg max x Ex fim
X, = TXx) (1 X) + IT (Xo).
Recall that by Brenier's theorem, Tj can be written as a gradient of a convex function, therefore T. is strictly monotone and we write X, T. (X.) for all 1. Further, observe that Ti (Xx) is a piecewise constant function, thus its Jacobian O almost everywhere, i.e. we have (Xo) (1 I almost surely, where D denotes Jacobian.
For any x,, define Xy = T,- (xxx). We can now use the change of variables formula
Po(Xo) (x,)) = Idet DT, (x,()
Po(TT'(xx)) (x:))
t)
As result,
log (x,) = DT, (xx) Po = (DT, 'I log (x
1
V Vlog po(x))
1-t
Xo
1-t"
where we plugged in Pu = =-==(00. Ia) in the last step. Furthermore, since X, is X--measurable, we have
1y (x:) = X; X
x, T Tx,)
X; X0
X
X4
t
t
Combining the above equations completes the proof.
The following is the formal statement of Proposition 4, Proposition 5. Suppose V1,t and U2,t generate the probability paths PI.t and P2,z respectively. Let
( (x) (x),,(x) (1 :x),,,;xx,,
and define (x) as the solution of the ODE de! (x) w)f7((x) iee. w/80)(x (x) is the position of a particle initialized from Xo moving along the velocity v(1)), Define the weight
(x) with the initial condition
(x(xx = x(7-1) {" (v." (x)))))))))))))))))) ( (V - log (1/()) (x)))
(x))
16

---

Published as a conference paper at ICLR 2026
Then, for any test function Φ R - R,
/ ((x)) I (( (x))d(o(x)
I (px
Before presenting the proof of Proposition 5, we comment on its informal interpretation, Proposi- tion 4. Let (r) denote the stochastic empirical law of k i..d. random vectors (xx") from Then, if Pa sufficiently concentrated, we have almost surely for all test functions
f (. (x))d))
[eem(0) (16) (x))dpo(x S (px)
(13)
J (p (7) (x)
+ 0o. Define Pt (nm)) as the measure that satisfies
f ((x))d()))) ( (x) (0 (1) (x))dlp (r) (x)
Sewe(x)pph" (x)
for all test functions d. Then, sampling from Pi (nn) is exactly equivalent to drawing from (w1)) (x(
Di_ according to (unnormalized) weights (x.")) Finally, as - we have
f ((x)dp
(1))) (x) = S (e (016 (x))dpp (F) (x)
feunl))( (x))dpo(x) = ((x)dp{") (x)
I (x0 (r)
(xx(
as r 0o. Therefore, (1-0) converges almost surely, weakly to p This is the sense in which the convergence of Proposition 4 holds. We keep the argument informal since we do not study the exact concentration properties required for (13) to hold.
To prove Proposition 5, we recall this result about the solution of an advection-reaction PDE. Lemma 6. Suppose (P:)120 satisfies
O.p: = - (pres)) + pe(ace palaa-fooddn ordpp).
(14)
where Idd - Rdd and Rd + R. Let (x(x be the solution to the ODE dvv(x) with initial condition vu(x)) = x. Further define
(x(()(x)))
wx(x) ['' (x)(ss
(15)
J==
Then, for any L. 2 0 and any test function & Rdd IR,
f [(x)dx((x) = S ( (x))
(16)
S e (x)dpu(xx
Proof. We proceed by differentiating (16) against time to obtain (14). We have a. ʃ [(x)dp(x)) S wx)) (Vo(v.(x)) or (th (x))) + f (n)
I (x)d(x))))
f (() f femx(x)dpo(x)
(x"ddx))
/ (Vo(x), (x)) + / d(xx (a(xx)- fardp) dpp(x))
where the last equality holds by the change of variables formula. This is exactly the weak sense of (14).
As result, to prove Proposition 5, we only need to show that P (7) satisfies PDE of the form (14).
17

---

Published as a conference paper at ICLR 2026
Proof of Proposition 5. Let p
1-7 be the unnormalized measure, and recall 2
Pi,ef2,
(1) (x)dx the normalizing constant. We have
1-y 1- Peple
Pp2.
0.p =ypi,"' 1 P2, Øt
+ + - 700,,,, 0t
= - --pi, P2, - (01,)011,)) - u2) (22,)
V. (p) (2 (7))
+ 7(7 - -1)p/" (01, - 22,4 V Vog P1, vi
V log p2,1),
where recall vft)) 1, + (1 7) and the third identity follows from the product rule
V (mmm) = Vm v + mV v. Define
(x(x) = 7(7 - 1)( 02)(x), log 12,
Then,
(7)
0epi") (x) 2 0,p,
(1)
Pi 0,2,7
2
Note that 0,e0 is constant in x, and 0: I Pi (x)dxx = 0. Therefore.
0,2 0, / 0i (x) 2(-) / (x(x))dp(xx)
Therefore,
0ep(2) - (f
(m) +p/²)) (at / aydpt)
We have obtained our desired PDE, which concludes the proof.
A.2 PROPERTIES OF THE SEMIDUAL Loss
Since we are interested the semidiscrete case, optimal transport plan always admits density with respect to the product measure μ @ v, regardless of entropic regularization. Therefore, we restrict the optimization problem to those that admit a density, and abuse the notation by using π as the density of coupling with respect to μ @ v. Recall d (x)g Therefore, we can rewrite (T)) as
Se.g (x) bj
(T( = Σ/ ((xx + Σ (x)(()), log
du(x).
We first relate the dual objective Fe the distance between V/3 and Lemma 7. For any g R.N and æ 2 O, we have VF = b m(g) (we use subgradient for E 0). As result,
(m((g)
Σ J=1
(0g (g))2 and TV(m(g).u) -
In particular, for b = IN/N we have x2 (v (b) (g)
Recall the Pinsker inequality TV S /1KL. and the fact that KT. S log(1 + x2). Thus always upper bound total variation using the chi-squared divergence.
Proof of Lemma 7. Note that the Chi-Squared and total variation bounds follow immediately from 0, Fg(g) = by m(g),
(17)
since then
m(B))
O F. (g) 2=
Σ
b, -1 x=x((((g)) b),
b
bj
18

---

Published as a conference paper at ICLR 2026
and
N
IV
Σ
mm(g); byl = (mV(m b).
Thus, we turn to showing (17). For E > O,
2--0(x.11) bj
0g. (g))= = bj
bj (x.g(x))du((() by -m(g)).
Σx
96- 96-o
In fact, could show the above identity directly using duality and the envelope theorem, b m(g) = 0 is constraint of the minimization problem.
For similarly have
/
[ arg max; 91 c(), du(x) b, - su.g (x.g(x)
0g Folg)
Σ [k E arg max; 91 v().
which completes the proof.
The following allows us to turn an approximate stationary point (maximizer) of the dual into an approximate optimal transport plan. Lemma 8. For any RN we have (T S C: + +|g (g)
Proof. Note that since v is a discrete measure, any coupling with the correct marginals admits density with respect to & и. Hence, we can use the change of variables dr dx = sxx)dx(x)). Then, we can write the OT cost as
Σ / s(x), ((xx,5)) a log
(s(x)s - dx(xx) + bj
C
min s(x)20 [=(x)dd(x(==b>
Eg =(x)==1
When 8 O, can drop the constraint क 0 it is satisfied by entropy regularization. The Lagrangian is then given by
s(x)s
= Σ
s(x); ((xx,5)) + +elog
-e] du(xx) te
bj
==1
N
/ [(x) (1- (1-1- sxx) du(xx + (g. b- ((x))d(())
2=1
Maximizing over s yields s(x), = - 1)/s and leads to the following dual problem
max De(.g) == [ [(xx)d((x)+ (g, - f.9
Σ
exp [(x) + dΙ c(xx [bdd((x) + e,
1=1
This dual is maximized by
N
[] c(xx, bJ-
(xn) = - log Σ
exp
Note that at fa.. syg. coincides with Sa.g of (7). The semidual is defined as Fg(g) max, (e g), and given by
Fg(g) b) - / log Σ exp
gj c(x, y3) )0 dd(x)..
19

---

Published as a conference paper at ICLR 2026
for 0. Thus, we derived the semidual (S₂) for 0. Importantly, one verify
(e(g) (S.,e SR.E: 8) (Tx + (B, b-bb- (x)x(x)d(()))
(18)
We now turn to the case = The Kantorovich dual in this case is given by
max Do(f.g) = = f(x))x(xx + (g.
(.g: @egSc
when e 0, where f g E e means f(x) + - c(x,,,) for (Nov))ae (x, y,). This is achieved by
[k,(()) min c(x, vs) 99-
This yields the semidual
Fo(g) = (g. b). [min c(x, 93 9x]d
We once again derived the semidual (Sz). Using the definition of Ta.gz we have
Σ / [c(x,y5) - (x),
Folg) = (g. b) +
/=1
c(x, (xx,)))50, (())))))))()))))) +
b (B.b-1 f 200--So.gxxxe (x,x(x)du(x))
==1
Once again, we obtain
Fo(g) = Co(To.s) )+ + g. (000s//_(xxxxx
iee. we extend (18) to 0. By Lemma 7, for any 2 0 have
Vf.gg) =
se.gdj((xo)
The above formula directly implies the total variation bound. Therefore,
Fs (g) C. (Te.g) (g. (g).
By duality, additionally have F. (g) s C.. Combined with the Cauchy-Schward inequality, obtain Ce (TT.g) S C + Vc ||. completing the proof.
A.3 CONVERGENCE OF SGD
We begin by recalling the following classical result from convex optimization, and we include its proof for completeness.
Proposition 9. Suppose f RRN + RR is a stochastic convex function with subgradient Vf, and there exists a constant such that E[[V S G2 for all g E NR Define F RRN - via F(.) E.](,), where expectution is over the stochasticity of f. Let g" denote any minimizer of F, and F Fgg"). Let (- T-1 denote the iterates of SGD with the sequence of step size {nc TT- Define gT ΣT-d neBc n..
Then,
Σ.
T-1 T
Bl g01121 |go 1²1 + G
2 T
vt
1=0
additionally F is L-moothh, i.e. 2 Fl op - L. Then
T- 1E V (g)))
F gg
F
ΣT- m LG
T-1
2 21-11 1
Ei=0
11=0 π
20

---

Published as a conference paper at ICLR 2026
Proof. By the convexity of F, for any g have F 2 F(g) + VF(g) g). With this property, for any we can write
E g++1 g"" gil
llg: g¹11 - 2n
gi (g) g + 1 go
| gi
[gi g" 1 - 2m VF
(g)
g g" 1| 2mp (F' F(g:)) E ( || |
Taking expectation over gr, we obtain
Elllg+1 g" | < B g g01l - 2m B - F(g)) + G/.
By summing both sides from = 0 to - 1, we have
Bllga+1
='|l < go - g" +
n
-
- F(gi))
2
Σ
m
1=0
1=0
T-1
go g"11 + G
1=0
which proves the first inequality.
For the second inequality, we start with the smoothness lemma
L
F(g+1) S Fg.) + (f(g:) ( (g--1 - -g)) + マ 2118+11 -g:11².
Taking expectation conditioned on gr, we obtain
n f( gr)
E (((+1))) g:l - F(g:) vI(VF(g:)1
2
Taking another expectation, we obtain
LF
E| - B[F(g:) < {1
2
By rearranging the terms, summing over from 0 to T 1, and noticing that F S [(F(g))) we obtain
T-1
T-1
Σ [ F(go) F* +
LG
Σ
which completes the proof.
To apply Proposition 9, we need to estimate the smoothness constant of the dual, achieved by the following lemma.
Lemma 10. Consider the semidual Fg defined in (S.). For any cost c Rdd x Bd IR and any 8 > O, we huve || llop I 1/e. Furthermore, if y) -x y and μ admits density, then
for s we have || < 20mux/ 18, where ö minizs ||x1 ( x and Cuuxx is the maximum y-surface area of any convex set Rdd
Proof. By Lemma 7, we have
/ xo)ddx)))
VF.(g)
where we notice that for E = 0 we can write
s0.g (Xo): [[ + x y. gj + yy, Vil.
since μ admits density, and consequently, ties in the maximum occur with probability zero. Thus, for 0, by direct calculation obtain
(g) = Ss (x)s: (x) diag(ss
Using this calculation, one immediately obtains the cost function (the argument above holds for any other cost
22 < L. which in fact does not depend on
21

---

Published as a conference paper at ICLR 2026
For E 0, by Kitagawa et al. (2019, Theorem 1.3), we have
Folg) = W diag d),
where d; = -jjzi Wiz
1
Wa 0, and, Wi yy Yoll A(R.))A((R.))
u(xo)dS((x)) Vi #
where A(g, i)
{x gi + x y; 2 Dj + +x y, vj} and dS denotes the d - 1-dimensional surface
measure in a Euclidean space. Note that V < Fo is the Laplacian of weighted graph with weights We 2 0, therefore Fo()l S max;" Σ Ejft Wyy. Further,
W y1 vsll A(g.))A((B.)
1
u(x)dS(x) 1 f
v(x)dS(x)
Cuax
(A(g.))
8
which finishes the proof.
We are now in position to state the proof of Theorem 2.
Proof of Theorem 2. Let F - in Proposition 9. We begin by noting that for constant learning rate 7/ 1) L.K and Unif ... K 1}), interpret the gradient norm bound of Proposition 9 as
A G2
(g)
n 2
where we recall from Lemma 10 that V2 op S T Moreover, let
B Σ s 1=1
VF =
(x())))
B
denote the stochastic gradient where (x(1)) L.d.
μ. We can define the probability distribution
B
míg)= =
Σ
Seg (x) (1)
ENA
B
Then,
v g))
<
V P (g) = 2TV(b, m(g)) 2
As result, G 2 in Proposition 9, and obtain
L.A
E V 3//
(19)
K
Let Vmin == miny by for simplicity. Combined with Lemma 7, conclude that
[bb < - [IFF(()))
3
LA
22 (m(g)
Vmin
Umin
K
which concludes the first part the proof.
For the second part, Lemma 8 and the Cauchy-Schwartz inequality ineuulll obtain
ElC. )] S C + E[ll 1 (
From Proposition 9, we have for any fixed [K]
4A
E|lgg111 || || G2 l
T
As result, the above bound also holds for E(| g:
Combining this with (19) the proof.
B ALGORITHMS
We provide the details of all algorithms introduced in the paper in this section.
22

---

Published as a conference paper at ICLR 2026
We begin by detailing the optimization algorithm for solving SD-OT. me only depends on k, Algo- rithm I is SGD on the semidual loss (S²). For SGD, the typical choice of learning rate π 1/ Vk after an optional constant learning rate phase. By allowing T/e to depend on the gradients, Algo- rithm 1 becomes adaptive stochastic optimization et al., 201 1) to be a particularly effective choice of learning rate schedule, and provide details on the schedule of SGD and Adagrad in Appendix C.
the semidual loss. We found AdaGrad (Duchi
Algorithm SOLVESDOT
Input: data yI..... VN R b A", 2 0, threshold step sizes rj. batch size M. Output: Semidiscrete dual potential g.
g 0.g f 0, k
2: while X (m) > T do
Sample Nu..... XM lid
3:
g +
g + ?]
M
1 Eimr Re (xc)
5:
6:
= + FFI
++
8: end while
9:
return g
Once dual potentials are obtained from Algorithm 1, can use the following algorithm to pair noise with data X1, which can be used to train flow-based models.
Algorithm 2 ASSIGN
noise Xo, data x(" ...., x 2 0. Output: Coupled noise-data pair (xu, x() I: return (k) le (xx) E AN.
Input:
Next, to better see the contrast between I-FM and OT-FM/SD-FM, we recall the usual I-FM algo- rithm for training flow models.
Algorithm I-FM
Input: Noise no. Data μ1
Σx=1 L0)).
FM batch size B, flow model (e x), Epochs E.
1: for f L...., NE/B do
NEO
2: Draw (x ) ~ H4 &B
3: Draw (x[ ( ~ μ
4: 4- FMSTEP((O, (xx J (x B
1.)
x B
5: end for
We now present OT-FM and SD-FM. Steps specific to OT-FM are highlighted in blue, while those specific to SD-FM are highlighted in red.
23

---

Published as a conference paper at ICLR 2026
Algorithm 4 OT-FM
1: Input: Noise pa, Data μ11 N
1
FM batch size B, flow model (t x), Epochs
2: Regularization e > 0, OT batch-size n.
3: for é I,.... N/En do
NEO + NE/n O(dn' faadm//5 / NEO + O(NEn/E
4: Draw (xg ] (x=))===
5: Draw (x1 ) ) #1
6: P SINKHORN DEL., (xi 03 )= E A
O(dn2/e" x( k BP.., for every S O(n") (neglected) P O x D + x B
7: Sample from coupling matrix:
x
for k n...,///B
0 FMSTEP(O, (x0))
1-88111 (X) 1i=k-41)
(())))))))))))))
10: end for
11: end for
Algorithm OT-FM** Cached Sinkhorn computations.
1: Input: Noise no. Data μl
Y
Σド= (2)) FM batch size B, flow model (s((, x), Epochs E,
Regularization e OT batch-size n.
3: for é L,..., : E/n do (4)
> O(NE/E x di / ( EdN/E" - Memory: 2NV E
4: Store X₁ = (xg Jim1 ~ 140
> Memory; n (storing rng)
5: Draw (xl ( 1--1 ~ 41 @n
6: P SINKHORN((s) 1 (xf ))) E A
> O(dn2/E2) Memory:
7: Sample from coupling matrix: ke[1] +- ~ BP.,, for every i S n.
D O(dn?)
8: end for
9:
for 8 L,..., .....,NEE E/n do
NEO
10: Load (xp (4) )==1 Xx
11: Load paired (0) (ke(4) for i S n
12: for k = 0, O,.., n/
do
13:
0 +- FMSTEP(O, (x (++1)
1B411 (x) (000))+))))
θ x B
14:
end for
15:
end for
Algorithm SD-FM
Input: Noise pe, Data N En= 5 6xp,, FM batch size B, flow model (t x), Epochs E,
1: Regularization > O. Threshold
g SOLVESDOT((H). 7T) 3: for f L,..., NE/Bdo
NdK
Memory: N & -NEE+NNN)
4: Draw (xg (4) J== @
x( + ASSIGN(x)) g.5) for S B.
D d
0 FMSTEP(O, (x (x ) (x 1Z1))
7: end for
C EXPERIMENT DETAILS
C. SEMIDISCRETE OPTIMAL TRANSPORT
To solve for the semidiscrete dual potential g in practice, we implement Algorithm 1 using JAX and utilize built-in data parallelism to scale to handle large datasets over multiple GPUs and nodes. Since the problem (S.) concave (and strictly, for > 0) one can use variety of different optimizers and obtain a suitable convergence behavior. We use the dot-product cost and scale & with the standard deviation of the cost, as proposed by Zhang et al. (2025). the class-conditional setting, we choose 3 10 for ImageNet-32 and 3 2 x 10 for ImageNet-64 and PetFace.
In all cases, we opt to use AdaGrad with a constant learning rate for 200 000 iterations followed by inverse-square root decay for an additional 100 000 iterations and perform iterate averaging over the
24

---

Published as a conference paper at ICLR 2026
final 50 000 iterates. To choose the learning rate, notice that typically |lg/| is of order vN while ||V F+Fe(() (g)]l is of order 1/VN. Thus, expect T., Lipschitz constant of the gradient, be of the order 1/N, which is smaller than the conservative estimation Le used in Theorem 2. As Theorem suggests a learning rate that scales with /1/L (the typical scaling under bounded gradient norm), find vN be suitable heuristic for setting the learning rate. As N for our datasets is of order 10,, we use 103 as the initial learning rate to optimize the potential. Once computed, the optimal dual potential g" can be cached and reused. Optimal dual potentials for our ImageNet and PetFace experiments will be made public upon publication.
To estimate (m((g) b),, 213 that can further be sharded along multiple devices, and average (10) over these batches. For the ImageNet-32x32 dataset, this calculation takes less than 2 minutes on node H100 GPUs.
a total of 2 noise samples, which divide into batches of size
C.2 FLOW MODEL TRAINING
ImageNet (32x32)
ImageNet (64x64), PetFace
Channels
256
192
Depth
3
Channels multiple
1,2,2,2
1,2,3,4
Heads
4
4
Heads Channels
64
64
Atlention resolution
Dropout
0.0
0.1
Batch size / GPU
256
50
GPUs
4
16
Effective Batch size
1024
800
Epochs
350
575
Effective Iterations
138 000
957 000
Learning Rate
0.0001
0.0001
Learning Rate Scheduler
Polynomial Decay
Constant
Warmup Steps
200
Table 4: Hyperparameters used for flow model training, adapted from Pooladian al. (2023).
ImageNet In all cases, we use the same U-Net architecture and hyperparameter choices as de- scribed in Pooladian al. (2023, Appendix E) and reproduced in Table for completeness. For SD-FM we sample noise-data pairs using Algorithm 2. Flow matching models are trained on single 8 NVIDIA H100 node for 32x32 resolution and two 8 x NVIDIA H100 nodes for 64x64 resolution.
PetFace Images from PetFace are resized to 64x64 resolution. We use the same U-Net architecture and hyperparameter choices as used for ImageNet 64x64 resolution. Flow matching models are trained on two 8 x NVIDIA H100 nodes.
CelebA Images from CelebA are first rescaled to 64x64 resolution. From each image x, low- resolution, noisy images are constructed by first downscaling 16x16 (4x SR) or 8x8 (8x SR) and adding Gaussian noise with standard deviation o = 0.1 or 0.2, i.c z = downscale(x)) + 7), treated the condition paired with x, and solve conditional optimal transport with 3 25. Flow matching models are then trained to generate samples of 64x64 images x conditional on downscaled, noisy observations We parameterize the flow with a U-Net with the same hyperparameter choices as for 64x64 ImageNet, and conditions are handled by resizing the low-resolution 8x8 16x 16 image 64x64 and stacking together with the input to the U-Net. A batch size of 1 128 and a learning rate of 0.0002 are used for FM training for total of 500 000 iterations.
7 (N, o² I)... The sampled corrupted image
D ADDITIONAL EXPERIMENTAL RESULTS
25

---

Published as a conference paper at ICLR 2026
Full dimensional
Buler NFE-
Euler NFE-8
Euler NFE- 16
DopriS
27.5
31.78
11
12.18
6.95
3.0
3,44
6.25-
25.0
6.00
2.8
22.5
5.75 10
10
10
10
-5
10
10
10
10
Time per pair (s)
Time per pair (s)
Time per pair (s)
Time per pair (s)
PCA 1000
Euler NFE
Euler NFE-8 8
Euler NFE=16
Dopri5
27.5
31.78
11.0
12.18
6.95
3.0
3,411
10.5
25.0
2.8
22.5
10.0
6.00
10"
10
10
10
10
10
10
10
10"
10
10
10
Time per pair (s)
Time per pair
Time per pair
Time per pair (s)
0.01
0.1
SD-FM
OT-FM,
131072
I-FM
Figure 7: FID vs total training time for class-conditional ImageNet 32x32.
Full dimensional
Euler NFE
Euler NFE
Fuler NFE 16
DopriS
79.24
26
37.77
16.0
19.68
9910
50
FID
15.5
8.8
24
45
15.0
8.6
10
Time per pair (s)
Time per pair (8)
Time per pair (s)
Time per pair (8)
PCA 500
Buler NFE-
Euler NFE-8
Euler NFE= 16
DopriS
79.24
26
37.77
16.0
19.68
9.00
9.100
50.0
8.75
15.5
47.5
24
8.50
15.0
10
10
10"
Time per pair (s)
Time per pair (s)
Time per pair (s)
Time per pair (s)
SD-FM
OT-EM, A 32768 131072
524288
I-FM
Figure 8:
FID vs time-per-paii for ImageNet 64x64. We report üime-per-pan for training flow matching using SD-FM compared to I-FM and OT-FM with batch sizes 2 2,2 2
26

---

Published as a conference paper at ICLR 2026
Euler NFE=4
Euler NFE=8
Euler NFE-16
Dopri5
Figure 9: Images generated from unconditional models trained on ImageNet (32x32). (a) denotes independent coupling while (b) denotes semidiscrete OT coupling with e used for training the model.
27

---

Published as a conference paper at ICLR 2026
Euler NFE=4
Buler NFE=8
Fuler NFF=16
Dopris
Figure 10: Images generated from conditional models trained on ImageNet (32x32). (a) denotes independent coupling while (b) denotes semidiscrete OT coupling with s = 0 used for training the model. Rows selected classes: English setter, mountain bike, banana, plane, pizza, marmot, red wine, wall clock, cheeseburger, king penguin.
28

---

Published as a conference paper at ICLR 2026
Euler NFE=4
Euler NFE=8
Euler NFE-16
Dopri5
Figure 11: Images generated from unconditional models trained on ImageNet (64x64). (a) denotes independent coupling while (b) denotes semidiscrete OT coupling with e used for training the model.
29

---

Published as a conference paper at ICLR 2026
Euler NFE=4
Euler NFE=8
AB 
Euler NFE-16
Dopri5
Figure 12: Images generated from conditional models trained on ImageNet (64x64), (a) denotes independent coupling while (b) denotes semidiscrete OT coupling with s used for training the model. Classes the same as in Figure 10.
30

---

Published as a conference paper at ICLR 2026
Euler NFE=4
Euler NFE=8
Euler NFE=16
Dopri5
Figure 13: Images generated from unconditional models trained on PetFace (64x64).
31

---

Published as a conference paper at ICLR 2026
Euler NFE=4
Euler NFE=8
Euler NFE=16
Dopri5
Figure 14: Images generated from conditional models trained on PetFace (64x64), Rows correspond the following classes: cat, chimp, chinchilla, degus, dog, ferret, guineapig. hamster.
32

---

Published as a conference paper at ICLR 2026
a 0.10
0 0.20
21.5
20.5
E=0
PSNK
21.0
PSIN
E-000
E 0.1
IFM
20.5
100000
400000
500000
100000
300000
400000 500000
Ileratinn
lleration
4x SR
a- 0.10
0 0.20
21.5
0.0
1006
PSNRR
E 0.01
0.1
IFM
20.5
100000
200001
300000
400000
S00000
100000
20000
300000
400000
510000
Ileratinn
lleration
8x SR
a 0.10
a 0.20
0.54
0.68
0.62
E-00
SIM
==00
0.66
0.1
IFM
0.64
100000
200000
300000
400000
500000
160000
200000
300000
400000
$00000
lteration
Iteration
4x SR
O 0.20
0.50
0.46
SSIM
0.0
0.48
0.41
0.01
0.16
0.42
JFM
0.14
0.40
100000
200000
300000
100000
500000
100000
200000 300000
400000
500000
Iteratioon
Iteration
8x SR
Figure 15: Evolution of PSNR and SSIM over the course of training for 4x and 8x image super- resolution on CelebA.
33

---

Published as a conference paper at ICLR 2026
(a) 4x SR, SD-FM, 0,3 25
(b) 4x SR, I-FM
6
(c) 8x SR, SD-PM, 0,8 25
(d) SR, I-FM
Figure 16: Sample images for CelebA 4x and 8x super- reperrresolution with noise level 0.2, for models trained with SD-FM and I-FM. Columns correspond to the original clean image, downscaled and noisy image, respectively, followed by 8 samples from the flow model.
34

---

Published as a conference paper at ICLR 2026
E SENSITIVITY OF DUAL POTENTIAL TO
In 17, we provide fine-grained picture of the impact of when computing on the same ImgN64 dataset (in full dimensions).
SD-FM: ABLATIONS ON IMGN64
F.I DISTRIBUTION OF MIPS SAMPLING AS A FUNCTION OF
Figure 18 proposes to study the distribution of selected images in the dataset of N ImgN64 (full- dimension) as SD>FMM training occurs. This pictures provides a complementary view to Fig. 3 show- ing why the FID metric improve the X2 metric decreases.
F.2 SD-FM METRICS A FUNCTION OF T,E AND COST ABLATION
We simultaneously study in the following experiments three different factors that lie at the heart of SD-FM:
does it matter that the potential g. for certain is truly optimal, in the sense that it respects marginals and guarantees that every point in the dataset is properly sampled? In other words, does low t always correlate with better FID and curvature metrics?
does matter for SD-FM to work to narrow down on the optimal transport as defined using the "traditional" & cost cxx, y) - (x. y) Would results be different if one were to select different costs? To answer this we consider three other costs: the "anti-OT" cost c(T, y) = (Ir, v) that promotes, the contrary. the longest paths, and random corruptions of the (zr, y) = (a y). The holds the realizations of random (Rademacher) variable in 11, with probability 30% or 70% of taking the value 1. This value is sampled only once per experiment.
does it matter to have a deterministic association (s 0) or is it better to inject more randomness (8 > 0), having in mind that infinite e is equivalent to IFM.
While all of Figures 19, 20, 21, 22 provide wide wealth of information, we can confidently say, by looking at all of them, that apart from FID in the lowest NFE-4 case, which is slightly more nuanced, all metrics improve as & and r decrease (closer to "sharp" OT) and clearly highlight the crucial role of the & cost (i.e. not "any" OT plan would work, that with shortest paths results in better performance).
F.3 TRAINING SD-OT ON PIXEL SPACE VS. LATENT SPACE
See Figure 23 and its caption for a comparison between the evolution of the X2 convergence criterion when computing SD-OT on either ImgN32 in original pixel space and Img2256 in latent space.
35

---

Published as a conference paper at ICLR 2026
ImageNet-64, sensitivity of g to
0.01
E 0.1
100
- 0.1 or 0.01
-100
-200
-200
-100
0
100
g ===
Figure 17:
We compare the optimal dual potentials g for different levels of using scatter plots. Each scatter plot depicts 2.56M points, where 2D points ((9ž),, (gE))) are plotted for different * E', As can be shown, the potential vector g. varies smoothly, for each datapoint, with s. These computations are carried out on full-dimension 12, 288 on ImgN64.
Total MIPS assignment counts, 0.0
Total MIPS assignment counts, 0.01
2000
0.05
0.05
1000
1750
== 1.0
== 1.0
uniform
uniform
1500
counts
raw counts
1000
750
400
500
200
250
0
0.0
1.0
2.0
2.5
0.5
1.5
2.5 1e6
sorted sample
1e6
sorted sample
Figure 18: Effect of varying threshold 7 (in Alg. when computing g on the distribution of selected images in the dataset, for 0.0 on the left, E 0.01 right. As can be seen, loose threshold of (essentially no optimization for g) results in images that are never sampled, whereas a finer threshold of 0.05 equalizes the distribution of sampled images. These computations are carried out on PCA & 1000 of the data space originally in full-dimension = 12, 288, on ImgN64.
36

---

Published as a conference paper at ICLR 2026
Euler
Euler (8)
Eullr (16)
Dopri5
100
cost
40
-x..)
30
20
20
(Or x,y)
(0t0 x.
20
1.0
(x (anti-O1)
10
0
0
0.20
0.06
curvatura
0.02
0.0050
0.0025
0.01
0.02
0.05
0.0000
0.00
0.00
0.00
0.0
0.01
1.0
0.0
0.01
1.0
0.0
0.01
1.0
0.0
0.01
1.0
eps
eps
eps
eps
Figure 19: Low
0.05 regime (i.e. close to exact SD-OT computation).
Euler (4)
Euler (8)
Euler (16)
Dopri5
40
cost
30
(x (OT)
20
(05 xx
20
x,y)
10
(x (anti-OT)
10
0.20
0.0100
0.03
0.06
0.15
0.0075
0.04
0.10
0.01
0.05
0.0000
0.00
0.00
0.00
0.0
0.01
1.0
0.01
1.0
0.0
0.01
1.0
0.0
0.01
1.0
Figure 20: Middle 0.2 regime (i.e. exact SD-OT computation).
Euler (4)
Euler (8)
Euler (16)
Dopri5
40
cost
60
30
30
-X.,)
xx
40
20
20
-(ovo
(x.)) (anti-OT)
20
10
0
0.015
0.20
0.06
0.03
0.15
curvature
0.10
0.005
0.02
cps
ops
ops
ops
Figure 21: High 1.0 regime (no marginal preservation is guaranteed).
37

---

Published as a conference paper at ICLR 2026
Euler (4)
Euler (8)
Euler (16)
Dopri5
60
threshold
0.05
30
40
20
20
0.2
10
10
curvature
0.006
0.02
0.03
0.004
0.010
0.02
0.01
0.002
0.005
0.000
0.00
egs
eps
eps
eps
Figure 22: Original cost c(x,,) -xrr v), low regularization, ablating
Conwergenns 5or Ingh30 (s 3077) vt Soghi3S-aannttii 4340
0.10
Imageten226
0.00
0.75
0.00
13:033
Figure 23: compare the evolution of the X2 convergence criterion when computing SD-OT on either ImgN32 in original pixel space (d 3072) ImgN256 in latent space (d 4096). Here the total dataset size is the same (N = 2.56M). While one would expect SD-OT to converge faster for smaller dimensional space (i.e. ImgN32). this is not the case, because of the far more regular distribution of latents that designed to resemble a Gaussian.
38


Reviews:
[{"review_id": "dsoHapVWTt", "reviewer": "Reviewer_48uQ", "summary": "This paper proposes Semidiscrete Flow Matching to tackle the training computational cost of OT-FM. It studies a semidiscrete OT problem from a continuous noise to a discrete data distribution. This is solved in two stages: a one-time precomputation of a dual potential vector, followed by a MIPS lookup during training to pair noise with data. Theoretical convergence analysis and empirical results show SD-FM beats I-FM and matches OT-FM at a fraction of the computational cost.", "questions": "1. Figure 5 presents a nice analysis of diversity control through guidance samples. Could you provide more extensive analysis on how different parameters affect the sample diversity?\n2. When applied to datasets of billions of samples, what is the additional memory overhead of SD-FM compared to I-FM?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The core idea of reframing OT-FM as a continuous-to-dataset semidiscrete problem is novel and practically efficient. The associated convergence criterion is also a new contribution.\n2. The paper is well-written. The motivation is clear. Figures and tables effectively communicate the method and its advantages.", "weaknesses": "1. The paper alludes to fast approximate MIPS but does not investigate the trade-off between approximation error and final model quality, measured by FID. For large $N$, what is the impact of the approximation error from MIPS?\n2. The paper lacks validation on high-dimensional data. Current diffusion models often denoise in latent spaces, which have high data dimensions.", "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": "8LeYlXP9uE", "reviewer": "Reviewer_rkmY", "summary": "This paper proposes Semidiscrete Flow Matching (SD-FM), a novel approach that addresses the computational bottlenecks of Optimal Transport Flow Matching (OT-FM). By reformulating the noise-data pairing problem as a semidiscrete optimal transport problem from a continuous noise distribution to a discrete dataset, SD-FM fits a dual potential vector via SGD during a precomputation phase. This allows each newly sampled noise vector to be matched to a data point at training time through a simple maximum inner product search, effectively eliminating the quadratic dependency on batch size that plagues batch-OT methods. Extensive experiments demonstrate that SD-FM outperforms both standard FM and OT-FM across various datasets and training settings, achieving superior results in unconditional/conditional generation tasks under all tested inference budget constraints.", "questions": "See the \"Weakness\" section.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper introduces a fundamentally different approach to optimal transport in flow matching by leveraging the semidiscrete formulation, effectively circumventing the quadratic complexity that has limited practical adoption of OT-FM methods.\n- Beyond the algorithmic contribution, the work provides thorough theoretical analysis including convergence guarantees for both regularized and unregularized settings and extends the Tweedie formula to the semidiscrete case, demonstrating strong mathematical rigor.\n- The experimental design covers diverse scenarios including unconditional/conditional generation, super-resolution, and mean-flow models, with consistent improvements across different inference budgets, particularly notable in low-NFE regimes.\n- The decoupling of precomputation from training, combined with efficient MIPS-based matching, offers a viable path for scaling OT-based methods to larger datasets while maintaining performance benefits.", "weaknesses": "- While the 12-hour precomputation on 2.56M samples is reasonable, the paper lacks systematic analysis of how the method scales to modern large-scale datasets (e.g., >10M samples). The O(N) training cost, though better than quadratic, may still become prohibitive for massive datasets.\n- The comparison framework omits important contemporary approaches for improving flow straightness (e.g., Reflow, MinibatchOT variants), making it difficult to assess SD-FM's relative advantages in the broader landscape of efficient flow matching methods.\n- The paper provides limited insights into hyperparameter selection, particularly regarding the regularization parameter ε. Practical recommendations for choosing ε based on dataset characteristics (dimensionality, size, noise levels) would significantly enhance reproducibility and usability.\n- The experimental validation focuses primarily on lower-resolution datasets, leaving open questions about the method's effectiveness on high-resolution generation tasks (e.g., 256×256 and above) that are increasingly relevant in practical applications.", "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": "OX5IVa8KZU", "reviewer": "Reviewer_JxQX", "summary": "To overcome the computational cost but fulfill the theoretical promises of OT-FM, this paper leverages the entropy regularized semidiscrete OT (SD-OT) and proposes a SD-OT coupling. In the coupling stage, the SD-OT problem is solved by estimating the dual potential $g$ by SGD; then in the FM training stage, noise samples are pared to data samples by maximum inner product search (MIPS, when $\\epsilon = 0$) or SoftMax (when $\\epsilon > 0$), using the learned $g$. They also provide the convergence analysis of SGD for SD-OT. The experiments using ImageNet, PetFace and CelebA show that the proposed SD-OT coupling is better than I-FM and OT-FM in main metrics.", "questions": "1. Although $\\epsilon = 0$ is valid, how sensitive are results to $\\epsilon$? Also, what's the sensitivity for fitted $g$?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper is well-motivated by practical drawbacks of OT-FM, and the general coupling method SD-FM can potentially integerate into many diffusion-/ flow-based models. The convergence analysis justify the case for both $\\epsilon = 0$ and $\\epsilon > 0$. The validity of unregularized $\\epsilon = 0$ case avoids the choice of tuning parameter.", "weaknesses": "1. The paper mainly adapts SD-OT from Genevay et al., 2016 in FM setting, which may raise some novelty issues.\n2. The scalabililty of SD-OT precompucate for large $N$ remains a concern.\n3. In section 2 (Background and Related Work), are \"$\\ldots$\" in subtites typos?", "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": "eC2yuvCCKE", "reviewer": "Reviewer_WFDC", "summary": "This paper introduces Semidiscrete Flow Matching (SD-FM), a novel approach to define the noise-to-data coupling in flow matching. The core idea is to perform a offline pre-computation to solve a global, semidiscrete Optimal Transport (OT) problem, which yields a dual potential vector g*. This vector g* implicitly defines a partition of the noise space, where each region in the partition corresponds to a unique data point. During the online FM training phase, each sampled noise vector is coupled to its corresponding data point by performing a Maximum Inner Product Search (MIPS) to identify its region. This \"pre-compute then query\" strategy is designed to improve the efficiency and scalability of OT-guided flow matching. The authors demonstrate significant performance gains on some benchmarks.", "questions": "In table 2, the generation quality is better at $\\epsilon = 0$, while some theoretical results are only valid when $\\epsilon > 0$, e.g., Theorem 2, Proposition 3. Could you provide some explanation for this phenomenon?", "limitations": "", "rating": 2, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 1, "strengths": "- Solid Theoretical Analysis: The paper provides a convergence analysis for the core algorithm, with a systematic comparison with existing optimal transport-based methods. This lends to theoretical importance to the proposed method. \n\n- Clarity of Exposition: The paper is well-written and logically structured. Figure 1, in particular, is highly effective. It intuitively illustrates the core methodological differences between I-FM, OT-FM, and the proposed SD-FM, allowing readers to quickly grasp the paper's central contribution.", "weaknesses": "- Limited Experimental Scope: The method was only evaluated on two small-scale benchmarks. This is a significant omission, as the problem of computational inefficiency—which the paper claims to solve—is most critical on large-scale, high-dimensional datasets like ImageNet 256. The absence of such experiments makes it difficult to validate the method's scalability and practical utility.\n\n- Limited Baseline Comparisons: The set of baseline methods is narrowly restricted to only two types of flow-matching models. The paper lacks a comparison against current state-of-the-art (SOTA) generative models in the related domain. To properly contextualize its contributions, the method must be benchmarked against SOTA approaches in terms of both generation quality and efficiency.\n\n- Insufficient Analysis of the Key Hyperparameter ε: The vast majority of experiments are conducted with ε=0 (hard assignment). The parameter ε controls the stochasticity of the coupling (soft vs. hard assignment) and is highly likely to directly influence the diversity of the generated model. The paper lacks any investigation into how ε mediates the trade-off between generation quality (e.g., FID/precision) and diversity (e.g., recall).", "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)