Designing a Conditional Prior Distribution for Flow-Based Generative Models
Authors: Noam Issachar
Mohammad Salama
Raanan Fattal
Sagie Benaim
Venue: Accepted by TMLR ()

Abstract: Flow-based generative models have recently shown impressive performance for conditional generation tasks, such as text-to-image generation. However, current methods transform a general unimodal noise distribution to a specific mode of the target data distribution. As such, every point in the initial source distribution can be mapped to every point in the target distribution, resulting in long average paths. To this end, in this work, we tap into a non-utilized property of conditional flow-based models: the ability to design a non-trivial prior distribution. Given an input condition, such as a text prompt, we first map it to a point lying in data space, representing an “average" data point with the minimal average distance to all data points of the same conditional mode (e.g., class). We then utilize the flow matching formulation to map samples from a parametric distribution centered around this point to the conditional target distribution. Experimentally, our method significantly improves training times and generation efficiency (FID, KID and CLIP alignment scores) compared to baselines, producing high quality samples using fewer sampling steps. Code is available at https://github.com/MoSalama98/conditional-prior-flow-matching.

Published in Transactions on Machine Learning Research (12/2025)
Designing a Conditional Prior Distribution for Flow-Based Generative Models
Noam Issachar' The Hehrew University of Jerusalem
nam..ssacharrmmalluuj.....
Mohammad Salama*
The Hebrew University of Jerusalem
Raanan Fattal
moneen.fttall@maiillij...coo
The Hebrew University of Jerusalem
Sagie Benaim
The Hebrew University of Jerusalem
Reviewed on OpenReview: https: // openreview. net/ forum? id= Teh9Bg4giF
Abstract
Flow-based generative models have recently shown impressive performance for conditional generation tasks, such as text-to-image generation. However, current methods transform general unimodal noise distribution a specific mode of the target data distribution. As such, every point in the initial source distribution can be mapped to every point in the target distribution, resulting in long average paths. To this end, in this work, we tap into non-utilized property of conditional flow-based models: the ability to design non-trivial prior distribution. Given an input condition, such as a text prompt, we first map it to a point lying data space, representing an "average' data point with the minimal average distance all data points of the same conditional mode (e.g., class). We then utilize the flow matching formulation to map samples from a parametric distribution centered around this point to the conditional target distribution. Experimentally, our method significantly improves training times and generation efficiency (FID, KID and CLIP alignment scores) compared to baselines, producing high quality samples using fewer sampling steps. Code is available at httts://github. com hww.:Moaaaaaa8////oaa wwt.:///cooooiiiiinnnii priinarll-------- floooorrr-ww---oooo-- f-wofllowaaaeee-- daloow-mtchingggg
1 Introduction
Conditional generative models are of significant importance for many scientific and industrial applications. Of these, the class of flow-based models and score-based diffusion models has recently shown a particularly impressive performance (27
S Although impressive, current methods suffer from long training and sampling times. To this end, in this work, we lap inlo a non-utilized property of conditional flow-based models: the ability to design a non-trivial prior distribution for conditional flow-based models based on the input: condition. In particular, for class-conditional generation and text-to-image generation, we propose robust method for constructing conditional flow-based generative model using informative condition-specife prior distribution fitted to the conditional modes (e.g,, classes) of the target distribution. By better modeling the prior distribution, we aim to improve the efficiency, both at training and at inference, of conditional generation via flow matching, thus achieving high quality results with fewer sampling steps.
Given an input variable (e.g., class or text prompt), current flow-based and score-basod diffusion models combinc the input condition with intermediate representations in learnable manner. However, crucially,
EEqual contribution

---

Published in Transactions on Machine Learning Research (12/2025)
these models are still trained to transform generic unimodal noise distribution to the different modes of the target data distribution. In some formulations, such as score-based diffusion (16 38 37), the use of a Gaussian source density is intrinsically connected to the process constructing the transformation. In others, such as flow matching (27 29 TD. a Gaussian source is not required, but is often chosen as a default for convenience. Consequently, in these seitings, the prior distribution bears little or no resemblance to the target, and hence every point the initial source distribution can be mapped to every point every mode in the target distribution, correspouding to given condition. This means that the average distance between pairs of source-target points is fairly large.
In the unconditional setting, recent works [3 H1,, show that starting from source (noise) data point that is close to the target data sample, during training, results in straighter probability flow lines, fewer sampling steps at test time, and faster training time. This is in comparison to the non-specific random pairing between the distributions typically used for training flow-based and score-based models. This suggests that finding a strategy to minimize the average distance between source and target points could result in a similar benefit. Our work aims to construct this by constructing a condition-specific source distribution by leveraging the input condition.
We, therefore, propose a novel paradigm for designing an informative condition--peeefic prior distribution for a flow-based conditional generative model. While in this work, we choose to work on flow matching, our approach can also be incorporated in other generative models, supporting arbitrary prior distributions. In the first step, we embed the input condition c to point z, lying in data space (which can be a latent one). For a discrete set of classes, this is done by averaging training samples corresponding to a given class in the data space. In the continuous case, such as text-to-image, we first choose a meaningful embodding for the input condition c (e.g. CLIP k44) Given training sample Ie and the corresponding ional embedding te, we train deterministic mapper function that projects ee to Te lying in data space. This results in an "average' data point of all samples x corresponding to the condition c. To enable stochastic mapping, we then map samples from a parametric conditional distribution centered Ze to the conditional target distribution (1(l))).
Thrget
Thrget
Flow Matching
Ours
Figure 1: Illustraiion of our approach. The LHS illustrates the flow matching paradigm, where every sample in the source Gaussian (shown as a circular point) can be mapped to every sample in the conditional target mode (shown as a cross point), where each class samples are shown in different color. In contrast, our method, shown on the RHS, constructs class-specific conditional distribution as source prior distribution. Each sample in the source distribution is, on average, closer to its corresponding sample in the larget mode.
While our approach can be implemented with any parametric conditional distribution, in our experiments, we chose to utilize a Gaussian Mixture Model (GMM). Specifically, the mean of each Gaussian is the "average' conditional data point. In the discrete condition case, each from the class-dependent training data, while for the continuous setting, it is fixed as a hyperparameter. Further, while the data space can be arbitrary, we choose it as the latent space of pro-trained variational autoencoder (VAE). These choices are derived from the following desirable properties: (i). One can easily sample from a GMM, (ii). Class conditional information can be directly represented by GMM, with each Gaussian corrosponding to conditional mode. (iii). We find empirically (sce Sec. 4.2) for real-world distributions (ImageNet
covariance is estimated directly
SS) and MS-COCO (26)), that the average distance between pairs of samples from the prior and data distributions (i.e. the transport cost) is much smaller than the unimodal Gaussian alternative (as in (27 (3)). Moreover, previous works (22 El 1b has shown that applying a GMM in a VAE space can be highly effective for clustering, suggesting that it can act as a suitable prior distribution. An illustration of our approach, for a simple setting consisting of eight Gaussians, each representing different class, is shown in Fig. I
To validate our approach, we first formulate flow matching from our conditional prior distribution (CPD) and show that our formulation results in low global truncation errors. Next, we consider a toy setting with
2

---

Published in Transactions on Machine Learning Research (12/2025)
known analytical target distribution and illustrate method's advantage in efficiency and quality. For real-world datasets, consider both the MS-COCO (text-to-image generation) and ImageNet-64 datasets (class conditioned generation). Compared to other flow-based (CondOT (27), BatchOT (33) and diffusion
(DDPM (TT) based models, our approach allows for faster training and sampling, as well as for a significantly
improved generated image quality and diversity, evaluated using FID and KID, and alignment the input. text, evaluated using CLIP score.
2 Related Work
Flow-based Models. Continuous Normalizing Flows (CNTs) ( emerged as a novel paradigm in generative modeling, offering a continuous- time extension to the discrete Normalizing Flows (NF) framework (24 (2). Recently, Flow Matching (27 22; T. has been introduced a simulation-free alternative for training CNFs. In scenarios involving conditional data (e.g., in text-to-image generation), conditioning is applied similarly to diffusion models, often through cross altention between the input condition and latent features. Typically, the source distribution remains unimodal, like a standard Gaussian (28). In contrast, our approach derives a prior distribution that is dependent on the input condition.
Conditioning in Flow-based Models. Conditional generation often adapts diffusion techniques, leveraging the connection between diffusson and probability flow ODEs (e.g.,
39). This typically involves vector field differences cetribbttoonddsiigniin
parameterizing ct) via Classifier Free Guidance inspired approaches [88
19), or cross- crossaattention, usually mapping from fixed Gaussian prior. Our adaptive, condition-specifie prior-it is orthogonal to these dynamics conditioning techniques. We leverage such standard mechanisms (specifically, cross-attention, as in our basclines) and apply our novel prior in addition. Informative Prior Design. Designing useful priors is studied in generalive models like VAEs (e.g., (10) and Normalizing Flows (e.g, (2I).. In score-based diffusion and flow matching, some works designed informative priors. For score-based diffusion, PriorGrad (25) addresses the well-known failure mode of score-based diffusion models to deal with extreme modes in the data, such as the voiced and unvoiced parts in waveform signals. This method whitens the data by normalizing mode-specific statistics, effectively simpler, quasi-isotropic distribution. This
flattening these modes that the diffusion process operates strategy reduces the burden on the network but also suppresses much of the original signal's 's structure. By contrast, our approach takes the opposite stance, instead of simplifying the data, we adapt the prior to better match the data distribution, thereby preserving its multimodal complexity while still enabling stable training. More recently, for flow matching, (33 [33 utilized dynamic optimal transport (OT) across mini-bateed batches to construct priors. Despite enabling efficient sampling, these OT-based methods face challenges with highly expensive training (quadratic OT computation) and reduced effectiveness for high-dimensional data, requiring exponentially larger batch sizes for performance gains. Our approach avoids these specific limitations by leveraging the conditioning variable to directly shape the prior distribution.
3 Preliminaries
We begin by introducing Continuous Normalizing Flow O in Sec. 3.1 and Flow Matching (27) in Sec. 3.2 This will motivate our approach, detailed in Sec. 4 which defines an informative conditional prior distribution on a conditional flow model.
3.1 Continuous Normalizing Flows
A probability density function over a manifold M is continuous non-negative function M + R, such that f p(x)dx= = 1. We set P to be the space such probability densitics on M. A probability path pe: [0, 1] + P is a curve in probability space connecting two densities po, Pi E P at endpoints = 0, 0,1-1. A flow di [0, 1] M + M is a time-dependent diffeomorphism defined to be the solution to the Ordinary Differential Equation (ODE):
d
v)(x)) = u (x)( Vu(x)= = r
(1)
dt
3

---

Published in Transactions on Machine Learning Research (12/2025)
subject to initial conditions where u, [0,1] x M - TM is time-dependent smooth vector field the collection of all tangent planes on the manifold TM (langent bundle). A flow t/ is said to generate a probability path Pe if it 'pushes' Po forward to P1 following the denendenn vector field uy. The path is denoted by:
dup
P: = P == Pol00 (x)) det
(x)
(2)
dx
where # is the standard push-forward operation. Previously, (6) proposed to model the flow ths implicilly by parameterizing the vector field ut, to produce pt, in a method called Continuous Normalizing Flows (CNF).
3.2 Flow Matching
Flow Matching (FM) (2r) is a simulation-free method for training CNFs that avoids likelihood computation during training, which can be expensive. It does so by fitting vector field of with parameters and regressing vector fields or that are known a priori to generate probability path Pr. P satisfying the boundary conditions:
(3) Note that Us is generally intractable. However, key insight of (27), is that this vector field can be constructed based on conditional vector fields (=(x(<1)) ) that generate conditional probability paths (x(r)))). The nush-forward of the conditional flow (x(()))), start at Pr and concentrate the density around x = 21 E M at t = 1. Marginalizing over the target distribution q recovers the unconditional probability path and unconditional vector field: pr(x) = J Pr (x)2)))((2 )dr, ue(ar) = JM (x(x)x) ( de,.
Po= = PI
px(2)
This vector field can be matched by a parameterized vector field vg using the (orm(0) objective:
B::= (1 x) (x|x1)]1
(4)
where || is a norm on TM. One particular choice of a conditional probability path Pe (x((x)) is to use the flow corresponding the optimal transport displacement interpolant (30) between Gaussian distributions. Specifically, in the context of the conditional probability path, (a( is the standard Gaussian, common convention in generative modeling, and P1 (2|²1) is a small Gaussian centered around 21. The conditional flow interpolating these distributions is given by 24 = (x(=r) = (1 t)zg + (xı, which results in the following conditional vector field (z([1] = - which is marginalized in Eq. 4 Substituting to uz(x/z1), one can also express the value of this vector field using a simpler expression: (xx = I1 To-
Flow matching (FM) was extended to conditional 20). In contrast the original FM, one first samples a input to the parametric vector field uo- The Conditional Generative Flow Matching (CGFM) objective Legtm (0) is:
Conditional Generation via Flow Matching.
generative modeling in several works 45
2
condition c. One then produces samples from Pa (x/c) by passing
(e (oo c, x) - (x(xx)11
(5)
In practice, e is incorporated by embedding it into some representation space and then using cross-altention between it and the features of Ug as in (55).
Flow Matching with Joint Distributions. While 127 considered the setting of independently sampled and 21, recently, [3 41 generalized the FM framework to an arbitrary joint distribution of p(zo,21) 21) in the unconditional generation setting, which satisfies the following marginal constraints: J p( 21 dez = q(xo), p(zo, 21 )dxo = ((x1). (38) modified the conditional probability path construction so at t- = 0, [+) = p( (1), where p( is the conditional distribution plao.a1), The Joint Conditional Flow Matching (JCFM) objective is:
((11)
Cjetm (0) 10 ue (x(x1)111
4 Method
Given a set {#1,: Hi of input samples and their corresponding conditioning states, our goal is to construct flow-matching model that samples from q(a)le) that start from our conditional prior distribution (CPD).
4

---

Published in Transactions on Machine Learning Research (12/2025)
4.1 Flow Matching from Conditional Prior Distribution
We generalize the framework of Sec. 3.2 to construction that uses an arbitrary conditional joint distribution of p( 21, c) which satisfy the marginal constraints:
/ 2120, ((21, / c) p(xo)
Then, building on flow matching, we propose to modify the conditional probability path so that at t == 0, we define:
po(20 c) = p(20/211.))
(7)
where x1 c) the conditional distribution p(vo.sse) Using this construction, we satisfy the boundary ((m1.))
condition of Eq. 3
po(x0) / (o(
(8)
= | p(zol=1, c)dxidc= p(xo)
(9)
The conditional probability path (x( does not need to be explicitly formulated. Instead, only its corresponding conditional vector field uz(2//11 c) needs to be defined such that points zo drawn from the conditional prior distribution po c), reach 21 at. I = 1, i.e., reach distribution pI (x/21, = 8(x x1). We thus purpose the Conditional Generation Joint FM Legitm objective:
10.. (t, x, ue (x c)12
(10)
where (20l Training only involves sampling from (3 21 e) and does not require explicitly defining the densities q(z0, and p,(r)211c). We note that this objective reduced to the CGFM objective Eq. when ( Z1. c)p(o).)
4.2 Conditional Prior Distribution
We now describe choice of condition-speeific prior distribution. When choosing a conditional prior distribution want to adhere to the following design principles: (i) Easy to sample: be efficiently sampled from. (ii) Well represents the target conditional modes. We design a condition-specific prior distribution based on a parametric Mizture Model where each mode of the mixture is correlated to specific conditional instribution p(a1lc). Specifically, we choose the prior distribution to be the following, easy to sample, Gaussian Misture Model (GMM):
Po GMM(N (pt. E.)
(11)
where E R is probability vector associated with the number of conditions (could be oo) and μt, M, are parameters determined by the conditional distribution g(<i/e;) statistics, i.c.
#4 = E|zi
Σ; = cov[2i
(12)
To sample from the marginal distribution p(zo x1, ci), we sample from the cluster Wi,,, E,) associated with the condition 6,-
Table 1: Average distances between samples from the prior and data dis- tributions (Le. transport cost) the ImageNet-64 and MS-COCO datascts across baselines.
Obtaining Lower Global Truncation Error. Our CPD fits GMM to the data distribution in favorable setting, where the association between samples and clusters is given. In this process, we fit dedicated Gaussian distribution to data points with the same condition. If the latter are close lo being unimodal, this approximation is expected to be tight, in terms the average distances between samples from the condition data mode and the fitted Gaussian. Tab. ! provides the average distances between pairs of samples from the prior and data distributions (i.e. the
ImageNet-64
MS-COCO
CondOT
640
630
BatchOT
632
604
Ours
570
510
5

---

Published in Transactions on Machine Learning Research (12/2025)
transport cost) of CondOT (2)). BatchOT (33) and our CPD over the ImageNet-64 ( and MS-COCO 26 datasets (note that both CondOT and BatchOT are applied in the conditinal setting, see Sec. 5.2.1 for details). As expected, BatchOT which minimizes this exact measure within mini-batches, obtains better scores than the naive pairing used in CondOT, while our CPD, which approximates the data using a GMM exploits the conditioning available in these datasets, and obtains considerably lower average distances.
As noted in (3), lower transport cost is generally associated with straighter flow trajectories, more efficient sampling and lower training time. We want to substantiate this claim from the viewpoint of cumulative errors in numerical integration. Sampling from fow-based models consists of solving time-dependent ODE of the form & = (xz), where d4 is the velocity field. This equation is solved by the following integral 2t = So (ds))dds where the initial condition zo is sampled from the prior distribution. Numerical integration over discrete time steps accumulate error at each step which is known the local truncation error Two which accumulates into what is known as the global truncation error en- This error is bounded by (40
lexl $
T (cLltn - -1)
(13)
hL
where h is the step size and L is the Lipschitz constant of ue. The distance between the endpoints of a path A = |ey - wol is given by I u, (da))ds which can be interpreted as the magnitude of the average velocity along the path z₁. Hence, the longer the path A, the larger the integrated flow vector field u, For example, if we scale a path uniformly by a factor C >> 1, i.e., tt + C(ze), we get, 4 C(xz)) = C(u,) in which case the Lipschitz constant L is also multiplied by C.
By shortening the distance between the prior and data distribution, errors which permits the use of coarser integration steps, which in turn yield smaller global errors. Thus, our construction allows for fewer integration steps during sampling. In Appendix A we also provide theoretical justification for why we expect the transition error to decrease as the prior distribution moves closer to the data.
CPD does, we lower the integration
4.2.1 Construction
for both the discrete case (e.g., class conditional generation) and continuous case (e.g., text conditional generation).
Next, we explain how we construct Po (Eq. 11
Discrete Condition. In the setup of discrete conditional gencration, we are given data (I1, G ) where there are finite set of conditions q.- We approximale the statistics of Eq. 12 using the training data statistios. That is, we compute the mean and covariance matrix of each class (potentially in some latent representation of pretrained auto-encoder). Since the classes at inference time are the same as in training, we use the same statistics at inference.
Continuous Condition. While in the discrete case we can directly approximate the statistics in Eq. from the training data, in the continuous case (e.g. text-conditional) we need to find those statistics also for conditions that were not seen during training. To this end, we first consider a joint space for training samples ( which represents the semantic distances between conditions c; and samples T1,- In the sctting where & is text, we choose a pretrained CLIP embedding. c; is then mapped to this representation space, and then mapped to the data space (could be a latent representation). using learned mapper Pa. Specifically, Py is trained to minimize the objective:
12
Cpeior (0) = Eg Eg (P(((() - 211.
(14)
where E is the pre-trained mapping to the joint condition-sample space (egg. CLIP). Pa can be seen as approximating E[zz|c], which is used as the mean for the condition specific Gaussian. At inference, where new conditions (c.g., texts) may appear, we first encode the condition e; to the joint representation space (c.g., CLIP) followed by Pg. This mapping provides us with the center b4 of each Gaussian. We also define Ey = of where os is hyper-parameter, ablated in Sec. 5.2.1

---

Published in Transactions on Machine Learning Research (12/2025)
4.3 Training and Inference
Given the prior po (either using the data statistics by training Po), for each condition c, we have its associated Gaussian parameters and Me. The map dr(01011 c) must be defined in order to minimize Eq. I0 above. This corresponds to the interpolating maps between this Gaussian at t = 0 and a small Gaussian around T1 at 1, defined by:
(x = ar c) + /4((1,, c);
(15)
(1), ((Omin)) I) (1 t)Ee/ and
(16)
() tx1 + (1 - t)fec
(17)
This results in the following target flow vector field
d
u (v(((1))1)))) c)) = dt
the (2|211 c) (Amiul 21/²) 21 Pe-
During inference we are given a condition and want to sample from ((xi(e). Similarly to the training, we sample zo ~ p(zole) and solve the ODE
d vx(x) = vo (1, (x), vo(x)) 3r
(18)
Training and implementation details are provided in Appendix
5
Experiments
0.0
0.2 t= 0.4 t- t 0.6 t= 0.7
0.8 t= 0.9
1.0
GT
CondOT
Ours
Figure 2: Trajectory illustration. A toy example illustrating the trajectory from the source to the target distribution for our method and conditional flow matching using optimal transport (CondOT).
NFE=2
NFE=3
NFE=4
NFE=6
NFE=8 NFE=10 NFE=15 NFE=400
GT
Prior
Samples
GT
CondOT
Train
Ours
Tesl
(b)
Figure 3: (a) NFE convergence illustration. A toy example illustrating convergence to the target distribution at different NFFs, for our method, compared to CondOT. (b) Generalization illustration. A toy example illustrating the generalization capabilities. LHS: Source prior and target samples for training classes RHS: As for LHS, but for test classes.
Overview. This section empirically evaluates our method across both synthetic and weal-world settings. Our goal to demonstrate that initializing the generative process from a structured conditional prior leads faster convergence, improved numerical efficiency, and higher-quality samples. We begin in Sec. 5.1 with controlled 2D toy experiments that illustrate the intuition behind approach and its advantages in

---

Published in Transactions on Machine Learning Research (12/2025)
Oun
CandoT
18.0
BatcnoT
17.0
50
DONN
17.6
172
0u5 Cordgr ВatcоoT
16.5
NFE
(a)
Ous
CondCT
16.75
ConSET
DORM
CLIP Score
DOFM
36.25
15.50
NFE
(b)
Figure 4: Numerical evaluation. (a) We compare method to class conditional flow matching using optimal transport paths (CondOT) 27 BatchOT (3), and DDPM (I7), on the ImageNet-64 dataset. We consider the FID score (LHS), KID score (Middle) and CLIP score (RHS). (b). As in (a) but for text-to-image generation on the MS-COCO dataset. As can be seen our mothod exhibit significant improvement per NFE, especially for low NFEs. For example, for NFEs, on ImageNel-64 and MS-COCO we get FID of 13.62 and FID of 18.05 respectively, while baselines do not surpass FID of 16.10 and FID of 28.32 respectively for the same NFEs. We consider up 40 NFE steps and note that DDPM converges to superior rosult given more steps.
convergence and generalization. We then turn to real-world image generation tasks in both class-conditioned (1mageNel-664) and lext-conditioned IMS-COCO) settings, where we quantitatively compare against existing flow- and diffusion-based baselines (Sec. 5.2.1p. Finally, analyze training efficiency, provide qualitative visualizations, and conduct ablation studies to isolate the contribution of key components of our method.
5.1 Toy Examples
We begin by considering the setting in which the prior distribution is mixture of isotropic Gaussians (GMM), where each Gaussian's mean represents the center of class (we set the standard deviation to 0.2). The target distribution consists of 2D squares with the same center as the Gaussian's mean in the source distribution and with a width and height of 0.2, representing large square. We compare our method to class-conditional flow matching (with OT paths), where each conditional sample can be generated from each Gaussian in the prior distribution.
In Fig. 2 we consider the trajectory from the prior to the target distribution. By starting from more informative conditional prior, our method converges more quickly and results in a better fitting of the target distribution. In Fig. 3 (aa, we consider the resulting samples for the different NFEs. NFE indicates the number of function evaluation used using a discrete Euler solver. Our method better aligns with the target distribution with fewer number of steps.
In Fig. 3bb, we consider the model's ability generalize to new classes not seen during training, akin to text-to-imant generation's setting. Training only subset of the classes, model exhibits generalization to new classes at test time, demonstrating that incorporating conditional prior does not impair the generalization behavior observed in vanilla Flow Matching.
In Fig. 5 we evaluate the method the case where classes are not uni-modal and there are intersec- tions the prior distribution, following data from VLines of the Datasaurus Dozen 122. We present generated samples from a model trained using CondOT (a) alongside samples from our model (b, c).

---

Published in Transactions on Machine Learning Research (12/2025)
The maximum mean discrepancy (MMD) computed on this data is 0.084 for CondOT, while we achieve an im- proved MMD of 0.072. In this set- ting, the class centers coincide, and the observed improvement arises from differences in the covariance of the condition-specifer priors, which pro- vide a better geometric alignment with the data distribution. In addition, we consider the case where the GMM con- sists of Gaussians for class A and 3 Gaussians for B (as shown in d and e respectively). These Gaussians can be found with per-class clustering method (applied training data) such as GMM-based clustcring. This results in improved performance compared to using single Gaussian per class (MMD of 0.067 vs 0.072).
(a)
(b)
(c)
(d)
(c)
Figure 5: Multi-modal classes. 1. A toy example illustrating multi- modal classes with intersections in the prior. Each color represents a class (class A or B), with samples as points and the prior distribution as contour lines. (a) shows a standard Gaussian prior (in black), while (b) and (c) show class-specific priors. While the mean each class falls on samples from the other class, our method results in an improved MMD score. (e) and (d) are in (b) and (c) respectively, but where separate GMM is used for each class.
5.2 Real World Setting
Datasets and Latent Representation Space.
For class-conditioned setting, we consider the ImageNet-64 dataset (8) while for text-to-image setting, we consider the 2017 split of the MS-COCO dataset (26), using standard train/validation//ee splits. We compute all our metrics on the ImageNet-64 validation set and the MS-COCO validation set. We perform flow matching in the latent representation of pre-trained variational auto-encoder
42
5.2.1 Quantitative Results
For a fair comparison, we evaluate our method in comparison to baselines using the same architecture, training scheme, and latent representation, as detailed above. We compare our method to standard class- conditioned or text-conditioned flow matching with OT paths (27) which we denote CondOT, where the source distribution is chosen to be standard Gaussian. We also consider BatchOT (3), which constructod prior distribution by utilizing the dynamic optimal transport (OT) formulation across mini-batches during training. Lastly, we consider Denoising Diffusion Probabilistic Models (DDPM) (16). We adapted the baselines DDPM, BatchOT/ CoondOT for conditional generation by incorporating condition into their shared U-Net architecture via standard mechanisms like cross-attention (detailed in Appendix E This architectural conditioning is orthogonal to our novel prior contribution. By using identical architectures and conditioning for all methods, our experimental validation fairly shows our prior enhances state-of-the-art conditional generation approaches, improving performance and convergence. To evaluate image quality, we consider the 3] and FID (14) scores. We also consider the CLIP score to evaluate the alignment. of generated images to the input text class, using the standard setting, as in
KID
Overall Performance. We evaluate the FID, KID and CLIP similarity metrics for various NFE values (as defined above), which is indicative of the sampling speed. In Fig. 4(a) and Fig, bb), we perform this evaluation for our method and the baseline methods, for ImageNet-644 (class conditioned generation) and for MS-COCO (text-to-image generation), respectively. As can be seen, our method obtains superior results across all scores for both ImageNet-6it and MS-COCO. For ImageNet-64, already, 15 NFEs our method achieves almost full convergence, whereas baseline methods achieve such convergence at much higher NFFs.. This is especially true for FID, where method converges at 15 NFEs, and bascline methods only achieve such performance at 30 NFEs. A similar behavior occurs for MS-COCO at 20 NFEs. We note that when considering NFEs for MS- COCC, we consider the pass in the mapper Pa to be marginal due to the small size of the mapper in relation to the velocity to, see Appendix B In Appendix B we also compare exclusively the diversity of generated samples in comparison to baselines, demonstrating improved performance.
9

---

Published in Transactions on Machine Learning Research (12/2025)
Training Convergence Speed. By starting from our conditional prior dis- tribution, training paths are on average shorter, and so our method should also converge more quickly at training. To
110
Ous CondBT
Car ConsoT
100
Sutudi
50-
13
evaluate this, Fig. G we consider the FID obtained at each epoch as well as the number of function evaluations (NFE) required for an adaptive solver to reach a pre-defined numerical tolerance, for a model trained on MS-COCO. Specifically, FID is computed using an Euler sampler with a constant number of function evalu- ations, NFE-20. As for the adaptive sampler, we use the dopri5 sampler with atol=rtol=1e-5 from the torchdiffeq library. Our method results in lower NFEs and superior FID, for every training epoch.
25
40
40
Epech
Epoch
Figure 6: Training time. For a text-conditional model trained on MS-COCO, we consider the NFE per training epoch. We compare our method with CondOT, BatchOT and DDPM. Note that DDPM had an FID value above 30 for all epochs not shown on the RHS.
Qualitative Results. In Appendix
NFE=3
NFE=5
NFE=8
NFE=10
NFE=15
NFE=20
NFE=100
we provide a visualization of our results for a model trained on MS- COCO, demonstrating both the sample corresponding the the text in the condi- tional source distribution, which resem- bles 'an average' image corresponding to the text, as well as samples corresponding to the text. We also provide a diverse set of images generated by our method, in comparison to flow matching. IN Fig. 7 we consider, for a model trained on MS- COCO and specific prompt, a visualization our results for different NFEs, illustrating the sample quality for varying numbers of sampling steps. As can be scen, method alrcady produces highly realistic samples at NFE=15.
C
Fig. g
Figure 7: Visualization of results for different NFEs. We consider a model trained on MS-COCO, and two different vali- dation prompts: Top: "There are yellow flowers inside a vase", Bottom: "A bowl full of oranges'.
Ablation Study. In the contimuous setting, as in MS-COCO, our method requires choosing the hyperparameter a, the standard deviation of each Gaussian. In Tab. 2 we report. the FID, KTD, and CLIP similarity values for different values of a. Our method results in best performance when 0.7, we believe that a relative large is necessary to allow a richer conditional prior due to the complex nature of the conditional image distribution. We also consider the case where our mapper Py takes as input a hag-of-words encoding instead of a CLIP encoding showing the importance of an expressive condition representation. As can be seen, performance drops significantly. As an additional ablation, we compare CondOT to our approach when mapping distributions directly to the original data space. When sampling with 15 NFFs, CondOT has FID score of 27.32 our 22.55 on ImageNet-64.
Table 2: Ablation study. Model per- formance for different values of a (the standard deviation) as a hyperparamo- ter for a model trained on MS-COCO. We also consider the case where our mapper Pe takes as input bag-of- words encoding instead of a CLIP.
FID
KTD4
CLIP
0.2
23.55
2.88
16.12
0.5
15.47
0.93
15.75
= 0.7
7.55
0.61
15.85
= 1.0
7.87
1.66
15.81
w/o CLIP
16.33
2.38
15.51
6 Conclusion
In this work, we introduce a novel initialization for flow-based generative models using condition-speeific priors, improving both training time and inference efficiency. Our method allows for significantly shorter probability paths, reducing the global truncation error. Our approach achieves improved performance on MS-COCO and ImageNet-64, surpassing basclines in FID, KID, and CLIP scores, particularly at lower NFEs. The flexibility of our method opens avenues for further exploration of other conditional initialization. While this work we assumed GMM structure of the prior distribution, different structures be explored.
10

---

Published in Transactions on Machine Learning Research (12/2025)
References
[L] M. S. Albergo and E. Vanden-Lijnden. Building normalizing flows with stochastic interpolants. ar preprint or "XXx:2209.15577 2022.
[2] L. Atanackovic, X. Zhang, B. Amos, M. Blanchette, J. Lee, Y. Bengio, A. Tong, and K. Neklyudov. Meta flow matching: Integrating vector fields on the wasserstein manifold, 2024.
[3] M. Binkkwski,, D. J. Sutherland, M. Arbel, and A. Gretton. Demystifying mmdd gans, 2021.
[4] L. A. Bull, P. Gardner, T. J. Rogers, E. Cross, N. Dervilis, and K. Worden. Probabilistic inference for structural health monitoring: New modes of learning from data. ASCE-ASME, Journal of Risk and Uncertainty in Engineering Systems, Part A: Civil Engineering, 7(1), Mar. 2021.
5
R. T. Q. Chen. torchdiffeg, 2018.
[6] R. T. Q. Chen, Y. Rubanova, J. Bettencourt, and D. Duvenaud. Neural ordinary differential equations, 2019.
[7] Q. Dao, H. Phung, B. Nguyen, and A. Tran. Flow matching in latent space, 2023.
[8] Deng, W. Dong, R. Socher, L.JJ. Li, K. Li, and L. Fei-Fei. Imagenet: largo-scale hicrarchical image database. In 2009 IEEE conference on computer vision and poltern recogriiiion, pages 248 255. Ieee, 2009.
(9] P. Dhariwal and A. Nichol. Diffusion models beat gans on image synthesis. Aduances in neural information processing systems, 34:8780-8794, 2021.
100 N. Dilokthanakul, han. Deep unsupervised clustering with gaussian mixture variational antoencoders. arXio preprint Xin: 16:::1111 02648, 2016.
Mediano, M. Garnelo, M. C. Lee, Salimbeni, K. Arulkumaran, and M. Shana-
[11] P. Esser, S. Kulal, A. Blattmann, R. Entezari, J. Müller, H. Saini, Y. Levi, D. Lorenz, A. Sauer, F. Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first International Conference on Machine Learning, 2024.
[12] C. Gillespic, S. Locke, R. Davies, and L. D'Agostino McGowan. datasauRus: Datasets from the Dalasaurus Dozen, 2025. R package version 0.1.9,
[13] J. Hessel, A. Holtzman, M. Forbes, R. L. Bras, and Y. Choi. Clipscore: reference-free evaluation metric for image captioning, 2022.
[14] M. Heusel, H. Ramsauer, T. Unterthiner, B. Nessler, and S. Hochreiter. Gans trained by two time-scale update rule converge to local nash equilibrium, 2018.
[15] G. E. Hinton and R. R. Salakhutdinov. Reducing the dimensionality of data with neural networks. science, (50057766):54-500 2006.
[16] J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-688 6851, 2020.
[17] J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. H. Larochelle, M. Ranzalo, R. Hadsell, M. Balcan, and H. Lin, editors, Aduances in Neural Information Processing Systems, volume 33, pages 6840 6851. Curran Associates, Inc., 2020.
[18] J. Ho and T. Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12500 2022.
[99 V. T. Hu, W. Zhang, M. Tang, P. Mettes, D. Zhao, and Snock. Latent space editing transformer- based flow matching. In Proceedings of the AAAI conference on artificial intelligence, volume 38, pages 2247-2255, 2024.
11

---

Published in Transactions on Machine Learning Research (12/2025)
[20] N. Isobe, M. Koyama, J. Zhang, K. Hayashi, and K. Fukumizu. Extended flow matching: method of conditional generation with generalized continuity equation, 2024.
[21] P. Izmailov, P. Kirichenko, M. Finzi, and A. G. Wilson. Semi-supervised learning with normalizing flows. In International conference on machine learning, pages 4615 4630. PMLR, 2020.
[22] Z. Jiang, Y. Zhong, H. Tan, B. Tang, and H. Zhou. Variational doep embedding: An unsupervised and generative approach to clustering, 2017.
[23] D. P. Kingma and Ba. Adam: A method for stochastic optimization; 2017.
[24] Kobyzev, S. J. Prince, and M. A. Brubaker. Normalizing flows: An introduction and review of current. methods. IEEE transactions on pattern analysis and machine intelligence, 43(11):3964 3979, 2020.
[25] S.-g. Lee, H. Kim, C. Shin, X. Tan, C. Liu, Q Meng, T. Qin, W. Chen, S. Yoon, and T.-Y. Liu. Priorgrad: Improving conditional denoising diffusion models with data-dependent adaptive prior. orXiv preprint or Xix::2106. 06406, 2021.
[66] T.-.. Lin, M. Maire, S. Belongie, J. Hays, P. Perona, D. Ramanan, P. Dollár, and C. L. Zitnick. Microsoft coco: Common objects in conlext. European conference on compuler vision, pages 740- 755, 2014.
[27] Y. Lipman, R. T. Chen, H. Ben-Hamu, M. Nickel, and M. Le. Flow matching for generative modeling. arxiv preprint 27X0022210.02747, 2022.
[28] A. H. Liu, M. Le, Vyas, B. Shi, A. Tjandra, and W.-N. Hsu. Generative pre-training for speech with flow matching. In The Twelfth International Conference on Learning Representations, 2024.
[29] X. Liu, C. Gong, and Q Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.0 030000,, 2022.
[30] R. J. McCann. A convexity principle for interacting gases. Advances in Mathematics, 128:153- 1997.
[31] M. F. Nacem, S. J. Oh, Y. Uh Y. Choi, and J. Yoo. Reliable fidelity and diversity metrics for generative models. In International conference machine learning, pages 7176 7185. PMLR, 2020.
[32] G. Papamakarios, E. Nalismick, D. J. Rezende, S. Mohamed, and B. Lakshminarayanan, Normalizing flows for probabilistic modeling and inference. Journal of Machine Learning Research, 22(57):1 64, 2021.
[33] A.-A Pooladian, H. Ben-Haam, C. Domingo-Earich, B. Amos, Y. Lipman, and R. T. Chen. Multisample flow matching: Straightening flows with minibatch couplings. arXiv preprint arXiv:2304. 14772, 2023.
[34] A. Radford, W. Kim, C. Hallacy, Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pages 8748 8763. PMLR, 2021.
[35] R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 10684-10695, 2022.
366 O. Ronneberger, P. Fischer, and T. Brox. U-net: Convolutional networks for biomedical image segmen- lation, 2015.
[37] J. Sohl-Dickstein, E. Weiss, N. Maheswaranathhnn, and S. Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learning, pages 2256 2265. PMLR, 2015.
[38] Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole. Score- generative modeling through stochastic differential equations. arXiv preprint Xiv:2001 1. 13456, 2020.
12

---

Published in Transactions on Machine Learning Research (12/2025)
[39] Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole. Score-based genera- live modeling through stochastic differential equations. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenRevvew.nee,, 2021.
[40] E. Süli and D. F. Mayers. An introduction to numerical analysis. Cambridge university press, 2003.
[41] A. Tong, N. Malkin, G. Huguet, Y. Zhang, J. Rector-Brooks, K. Fatras, G. Wolf, and Y. Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. arXiv preprint arXiv:2302.0044 2023.
[2] A. van den Oord, O. Vinyals, and K. Kavukcuoglu. Neural discrete representation learning, 2018.
[3] A. Vaswani, N. Shazeer, N. Parmar, Uszkoreit, L. Jones, A. N. Gomez, Kaiser, and 1. Polosukhin. Attention is all you need, 2023.
[44] P. von Platen, S. Patil, A. Lozhkov, Cuenca, N. Lambert, K. Rasul, M. Davaadorj, D. Nair, S. Paul, W. Berman, Y. Xu, S. Liu, and Wolf. Diffusers: State-of-the-arl diffusson models. https: diffusers,
//github.com//ugeen cwt.e/gggggggggaa////////o/a siffuruueere/eiffusrr.. 2022.
[45] Q- Zheng, M. Le, N. Shaul, Y. Lipman, A. Grover, and R. T. Q. Chen. Guided flows for generative modeling and decision making, 2023.
13

---

Published in Transactions on Machine Learning Research (12/2025)
Flow Matching
Ours
Figure 8: Visual comparison of randomly generated samples for prompts from the MS-COCO validation set. using method, in comparison to flow matching, for model trained on MS-COCO.
A Theoretical Justification for Lower Truncation Error
We follow our definition of the global truncation error in Eq. 13 and formalize why we expect it to decrease as the prior distribution moves closer to the data. This connection can be formalized under the reasonable assumption that the velocity fields u, learned via flow matching scale proportionally with the effective distance D between the source and target distributions. Empirically, this supported by faster convergence with our closer prior (o.g., Fig. Fig. Let u, Du', where v' velocity field with magnitude O(1). As argued in Sec. 4.2 the Lipschitz constant L of u, also scales proportionally: L DL'
Proposition 1. Assuming flow matching velocity fields u, have magnitude and Lipschitz constont scaling proportionally with D, then the global trancation error bound decreases D decreases.
Proof. The error is bounded by |enl <= maz (ef((t-to) - 1). Assuming maayyT also scales with D, and substituting L DL', the bound bocomes approximately D((ttt(t 1). The exponential dependence on D via shows that reducing D (closer prior) exponentially reduces the error bound, allowing larger steps h lower error.
14

---

Published in Transactions on Machine Learning Research (12/2025)
red and white plane in the sky
A light green kitchen some cabinets dish washer and sink
A black honda motorcycle with a dark burgundy seat
city street with multiple trees
Figure 9: A visualization of our results MS-COCO. We show, for four different text prompts: (a). The sample corresponding to the text in the conditional source distribution, which is used the center of Gaussian corresponding to the text prompt (LHS) (b). Six randomly generated samples from the learned larget distribution conditioned on the text prompt (RHS).
B Additional Quantitative Results
In Tab we present additional metrics (FID, KID, and CLIP-Scorc) for ImageNet-64 with 15 NFEs. We compare the performance of CondOT (27, BatchOT 33 and DDPM (T6). As shown, our model delivers significant improvements over the baselines.
Table 3: Numerical evaluation.
Quality of generated samples (FID, KTD), and conditional fidelity (CLIP- Score) for method in compari- son to baselines, for the ImageNet-64 dataset for 15 NFFs. We consider (CondOT) (27), BatchOT DDPM 1G.
As additional comparison, we note that while FID has a compo- nent addressing diversity, it measures both fidelity and diversity together. To this end, we considered the Recall and Coverage metrics from BI, which exclusively measure diversity (higher is better for both). On ImageNet-64, compared to CondOT and BatchOT, for 15 of 0.54/0.78 while CondOT obtains 0.03/0.0 /0.07 and BatchOT oblains 0.14/0.23. For 40 NFEs, obtain Recall/ Coverage of 0.55/0.79 while CondOT obtains 0.32/0.49 and BatchOT obtains 0.37 0.37/0.60. This demonstrates that our method improves the generated distribution's diversity, both low and high NFFs.
33 and
NFEs obtain
FID 4
KID4
CLIP1
DDPM
47.51
6.74
17.71
CondOT
16.16
1.96
18.02
BatchOT
16.10
1.43
17.72
Ours
13.62
0.83
18.05
C Visual Results
In Fig. we provide a visualization of our results for model trained on MS-COCO. We show, for four different text prompts: (a). The sample corresponding to the text in the conditional source distribution,
15

---

Published in Transactions on Machine Learning Research (12/2025)
Table 4: used for training each model
ImageNet-64
MS-COCO
Dropout
0.0
0.0
Effective Batch size
2048
128
GPUs
4
Epochs
100
50
Learning Rate
lc-4
1c-4
Learning Rate Scheduler
Constant
Constant
which is used as the center of Gaussian corresponding to the text prompt. (b). Six randomly generated samples from the learned target distribution conditioned on the text prompt. As can be seen, the conditional source distribution samples resemble 'an average' image corresponding to the text, while generated samples display diversity and realism.
In Fig. 7 we consider, for a model trained on MS-COCO and a specific prompt, visualization of our results for different NFEs, illustrating the sample quality for varying numbers of sampling steps. As can be seen, method already produces highly realistic samples at NFE-15.
In Fig. g we provide additional visual results for our method in comparison to standard flow matching for a model trained on MS-COCO.
D Additional Toy Examples
We construct two synthetic target distributions that share the same conditional prior to illustrate the behavior of our method when the prior statistics are identical. The first is "+" shape composed of intersocting horizontal and vertical line segments, while the second is obtained by rotating this distribution by 45" (Fig. 10p. Both distributions share an isotropic Gaussian conditional prior p(zolc) = (00 I), providing a controlled setting for analyzing equivalence and divergence regimes of our approach.
Figure 10: Toy distributions with identi- cal conditional priors. (Teft) A "4" shape formed by intersecting horizontal and vertical line segments. (Right) The same distribution rotated by 45% Both share an isotropic Gaus- sian conditional prior centered at the origin.
When both the prior and data distributions are isotropic, our method effectively reduces to standard Flow Matching (CondOT), producing comparable flow trajectories and sample quality. However, when the data variance deviates from one, the structured prior yields measurable advantages. In this case, we observe improved alignment between the prior and data distributions, with an MMD of 0.077 compared to 0.082 for the vanilla bascline (CondOT).
E Implementation Details
We report the hyper-parameters used in Table 4 All models were trained using the Adam optimizer with the following parameters: 31 0.9, 8g 0.999, weight decay 0.0, and 1e-8. All methods we trained (L.e Ours, CondOT, BatchOT, DDPM) using identical architectures, specifically, the standard Unet 36 architecture from the diffusers A4 library with the same number of parameters (872M) for the the same number of Epochs (see Table for details). For all methods and datasets, we utilize a pre-trained Auto-Eneen Encoder
P3
(42) and perform the flow/diffusion in its latent space.
In the case of text-to-image generation, we encode the text prompt using a pre-trained CLIP network and pass to the velocity up using the standard UNet condition mechanism. In the cass-conditional setting, we
16

---

Published in Transactions on Machine Learning Research (12/2025)
create the prompt 'an image of (class)' and use it for the same conditioning scheme in text conditional generation.
For the mapper Pe (Sec. 4.2 , we use network composed of a linear layer followed by two ResNet blocks, totaling 11M parameters. The training time required for Po is negligible compared to the flow network, accounting for only ~ 1% of the total training time.
When using an adaptive step size sampler, we use dopri5 with atol=rtol=1e-5 from the torchdiffeq 55) library.
Regarding the toy example Sec. 5.1 we use a 4 layer MLP with ReLU activation as the velocity to. In this setup, we incorporate the condition by using positional embedding (43) on the mean of each conditional mode and pass it to the velocity Ue by concatenating to its imput.
F Limitations and Future Work
Our proposed approach, utilizing a Gaussian Mixture Model (GMM) prior centered on mapped conditional means within a VAE latent space, demonstrates significant improvements for class-conditional and text-to- image generation lasks, proving effective in these complex, real-world settings.
Choice of Gaussian prior. In this work, intentionally model each conditional prior Gaussian distribution, which fully captures the first two moments (mean and covariance) of the conditional data. This design ensures analytical simplicily, stable training, and compulational efficiency, while already yielding substantial improvements in transport efficiency and generative quality. Nevertheless, richer or learnable priors could offer additional ficxibility. For example, mixture-basod or neural priors that capture higher-order moments may further reduce the transport distance or belter adapt to multi-modal condition distribu- tions. Exploring such priors-eithor learned jointly with the flow model or parameterized via amortized inference remains an important direction for future work.
Beyond image domains. While our approach proves effective for image generation, other conditional tasks involve fundamentally different output structures for which the current GMM-based prior may be suboptimal. For example, in conditional 3D human pose generation, data lie on a complex kinematic manifold, and an "average pose" used to center a Gaussian component might be physically invalid. More generally, structured modalities such as human motion, molecular graphs, or music exhibit strong geometrie, relational, or temporal dependencics that are not well captured by Gaussian components. Adapting our framework these setlings may require priors that are manifold-aware that explicitly model structural constraints such distributions defined over joint angles, graph topologies, temporal sequences. Developing such task- specific conditional priors constitutes promising direction for future work. Finally, investigating performance directly in the data space, without relying on a pretrained VAE encoder, also represents an interesting path for improving fidelity and interpretability.
17

---

Published in Transactions on Machine Learning Research (12/2025)
G
Training and Inference Algorithms
In this section, we provide concisc pscudocode for the main components of our method: (i) training the flow model from a conditional prior (CGJFM), (i)) inference (sampling) from the learned model, and (iit) training the prior mapper Pa for continuous conditioning (e.g., text).
Notation. 1 (2ol211, e) and (zolz1),) c) denote the interpolation and its time derivative, respectively, defined in Sec. (u, Σ) denoles Gaussian distribution with mean and covariance E. The loss Legifm is given in Eq.
Algorithm Training Flow Matching from a Conditional Prior (CGJFM)
Require: Data {21,c), conditional prior {Pe, Ech, flow network Vi
1: while not converged do
2: Sample (x1, c))(())11())) c), t~ 1), zo ~N
2r (xnla, Ut (ro
Legif(()) l|ua(t, x4, 1
O 0 VoCogiiu
end while
Algorithm Sampling from the Conditional Prior
Require: Condition e, trained flow ug prior (Hes Ee)) Sample Zg NNMe, Ee)
for / E [0, 1] do
ODE integration
④1 vo (t,
4: end for
5: return 3
Algorithm Training Pa for Continuous Conditions
Require: Data {x1,c), encoder E, network Pe
1: while not converged do
Sample c)-((rr))))
A. Pe(((c)
Cprior (0) = II - 21112
O 8 17 VoCpirio
6: end while
These pseudocode listings correspond directly to the procedures described in Sec. Algorithm trains the conditional flow network using our conditional prior distribution, Algorithm 2 performs sampling by integrating the learned ODE, and Algorithm trains Pa to map conditions (e.g., CLIP embeddings) to the mean of the condition-specific Gaussian prior.
18


Reviews:
[{"review_id": "wbf2jmZppR", "reviewer": "Reviewer_iUYo", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes to use a prior distribution for flow-based generative models, which enables the reduction of computational costs. The authors introduce two methods for designing the prior distribution using Gaussian mixture models, one for a class-labeled dataset and the other for text prompts. The experimental results demonstrate that the proposed method can improve training times and generation efficiency compared to baseline methods.\n\n[Strengths]\n- Introducing the prior distribution for flow-based generative models is well-motivated and interesting.\n- The proposed method experimentally works well to reduce the training and generation times and improve the generation quality.\n- The proposed method can be combined with various types of flow-based generative models and other techniques.\n\n[Weaknesses]\n- The proposed method requires additional computational cost for constructing the prior distribution, although it is not significant compared to the flow-network training cost as described in the Appendix. However, to introduce the prior distribution, the full training of the flow network is required.", "claims_and_evidence_yn": "Yes", "claims_explanation": "The motivation of the proposed method is clearly described. The use of the prior distribution for reducing the sampling steps of generation is convincing. The effectiveness of the proposed method is demonstrated through experiments on the toy dataset and the ImageNet and MS-COCO datasets.", "audience_yn": "Yes", "audience_explanation": "The generation cost of flow-based generative models is a central issue in the field. As the proposed method directly contributes to reducing this cost, it will likely interest researchers and developers working on generative models and their applications.", "requested_changes": "I think this paper is well-written. I do not have any critical issues or major concerns.\nThe following are some (minor) suggestions.\n\n- It might be better to provide the training and inference procedures of the proposed method as pseudocode, which would help readers understand the method in further detail.\n- It would be nice to release the code for the proposed method.\n- I cannot identify the exact architectures of $P_\\theta (E(c))$. Could you please explain the detailed implementation? Especially, in my understanding, $\\mathcal{P}_\\theta$ maps the encoded representation (continuous vector space) to the image space.", "broader_impact_concerns": "I do not have concerns about the ethical implications of this paper.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "YrjwUZsgb8", "reviewer": "Reviewer_2Pv1", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a novel method for improving the efficiency and performance of conditional flow-based generative models. The core idea is to replace the standard, generic, unimodal prior distribution (e.g., a standard Gaussian) with an informative, condition-specific prior.  The authors validate this approach on toy examples, ImageNet-64, and MS-COCO, demonstrating promixing improvements in FID, KID, and CLIP alignment scores compared to baselines, particularly at low sampling steps.", "claims_and_evidence_yn": "Yes", "claims_explanation": "- The experimental results are convincing. The method shows promising quantitative improvements over baselines (CondOT, BatchOT, DDPM) across multiple metrics (FID, KID, CLIP) and datasets (ImageNet-64, MS-COCO). Specifically, the model performs at a low number of function evaluations (NFE) for same quality, as shown in Figure 4. \n\n- The paper is well written, and the experiments are done thoroughly to support the claim. The authors also include ablations studies for in-depth analysis of the proposed approach.", "audience_yn": "Yes", "audience_explanation": "The central contribution—designing a conditional prior to shorten the transport path—is well-motivated, and intuitive. It addresses a clear limitation in existing conditional flow and diffusion models.", "requested_changes": "- As stated in Sec 5.1, the prior is modeled as an isotropic Gaussian. This feels like a significant simplification. The performance might be further improved by learning a more complex prior. The paper would be strengthened by a discussion of this design choice and its potential limitations.\n\n- The proposed method is only tested on ImageNet-64 dataset and MS-COCO dataset. It's better to evaluate on more complex real-world tasks such as video generation, 3D generation to validate the generalization ability.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "XEqpNkghcf", "reviewer": "Reviewer_cD6g", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "### Summary\nThis paper proposes a novel improvement of flow-based generative models to utilize condition-specific prior distributions. To be specific, the paper highlights that existing flow-based models employ a unimodal distribution to generate the target distribution, which typically involves a large flow distance from the source noise to the target point, resulting in more sampling steps and longer training times. To overcome this limitation, the authors focus on utilizing a conditional prior distribution in a conditional generative modeling method, specifically flow matching. Specifically, the authors describe how this concept can be implemented with a particular family of conditional prior distributions, namely the Gaussian Mixture Model. Using experiments on toy datasets and two real-world datasets, the authors demonstrate how the proposed method enhances overall training performance.\n### Strengths\n1. The concept of conditional prior distribution is well-motivated.\n2. Generally, the paper writing is clear and easy to understand.\n3. The experimental results support the improvement well.\n### Weaknesses\n1. I don’t find any major flaws or weaknesses in this paper. There are some questions or minor changes listed in the **Requested Changes** section.", "claims_and_evidence_yn": "Yes", "claims_explanation": "The main claim of the paper is that the use of conditional prior distributions can reduce both the number of samplings required during test time and overall training time, resulting in high-quality samples with more efficient training. The authors quantified the quality of the generated samples using three different metrics: FID, KID, and CLIP. They successfully demonstrated that the proposed method can converge faster than other existing methods without conditional prior distributions.", "audience_yn": "Yes", "audience_explanation": "To the best of my knowledge, the flow-based generative model is also one of the mainstream generative models in machine learning research. The underlying geometric nature of the normalizing flow maps each point in the prior distribution to some other point in the posterior distribution. Considering the flow of each point during sample generation, some readers may think that the topological/geometric mismatches between the prior and posterior distributions are cumbersome, and they would be interested in reducing these geometric gaps.", "requested_changes": "1. At the beginning of Section 5, please describe the main structure of the experimental section, with a summary of the goal of the experiments.\n2. I’m not quite convinced by the result in Figure 3(b). Do the authors have evidence (or experimental results) that the same result cannot be achieved unless conditional prior distributions are utilized?\n3. Figure 5 presents the case where the centers are the same, showing improvements in the maximum mean discrepancy. Do the authors claim that the improvement still comes from selecting class-specific priors with different covariances?\n4. I’m curious about the case when there are two target distributions whose conditional prior distributions are the same. For example, a uniform distribution over the union of horizontal and vertical line segments (+ shape) would have an isotropic Gaussian distribution as a conditional prior. Rotating this distribution by 45 degrees would result in the same conditional prior. Then, we can construct two target distributions with identical conditional prior distributions. Will the proposed method not have any improvement, because it will be the same as the existing methods? If there are still some improvements from the proposed method, why?", "broader_impact_concerns": "I don’t see a particular broader impact concern regarding this paper.", "additional_comments": "This paper is well-written in general. The use of conditional prior distributions is a novel contribution to flow-based generative models. The overall writing is clear and easy to understand. Also, the experimental results support the improvement well.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)