Pairwise Optimal Transports for Training All-to-All Flow-Based Condition Transfer Model
Authors: Kotaro Ikeda
Masanori Koyama
Jinzhe Zhang
Kohei Hayashi
Kenji Fukumizu
Venue: NeurIPS 2025 poster ()
Keywords: Flow Matching
Conditional Transfer
Primary area: deep_learning
TL;DR: We present a scheme of learning condition transfer achieving pairwise optimal transport, from a generic dataset of from {(x, c)} where a sample corresponding to a specific c is unique

Abstract: In this paper, we propose a flow-based method for learning all-to-all transfer maps among conditional distributions that approximates pairwise optimal transport. The proposed method addresses the challenge of handling the case of continuous conditions, which often involve a large set of conditions with sparse empirical observations per condition.  We introduce a novel cost function that enables simultaneous learning of optimal transports for all pairs of conditional distributions. Our method is supported by a theoretical guarantee that, in the limit, it converges to the pairwise optimal transports among infinite pairs of conditional distributions. The learned transport maps are subsequently used to couple data points in conditional flow matching. We demonstrate the effectiveness of this method on synthetic and benchmark datasets, as well as on chemical datasets in which continuous physical properties are defined as conditions.

Pairwise Optimal Transports for Training All-to-All Flow-Based Condition Transfer Model
Kotaro Ikeda
Masanori Koyama The University of Tokyo Preferred Networks, Inc.
The University of Tokyo
Preferred Networks, Inc.*
kotaro- -ikdro-ikedaßg ecc. uecc.u ac. jp.
Jinzhe Zhang Preferred Networks, Inc.
Kohei Hayashi The University Tokyo Preferred Networks, Inc.
Kenji Fukumizu
The Institute of Statistical Mathematics Preferred Networks, Inc.
Abstract
In this paper, propose flow-based method for learning all-to-all transfer maps among conditional distributions that approximates pairwise optimal transport. The proposed method addresses the challenge of handling the case of continuous conditions, which often involve a large set of conditions with sparse empirical observations per condition. We introduce novel cost function that enables simul- taneous learning of optimal transports for all pairs of conditional distributions. Our method is supported by theoretical guarantee that, in the limit, it converges to the pairwise optimal transports among infinite pairs of conditional distributions. The learned transport maps are subsequently used couple data points in conditional flow matching. We demonstrate the effectiveness of this method on synthetic and benchmark datasets, as well as on chemical datasets in which continuous physical properties are defined as conditions. The code for this project can be found at https://gitaae www/kotttaateeeeecc honn/Kottuuuurii wroom A2A-
Introduction
Recent advances in generative modeling have been largely driven by the theory of dynamical systems, which describes the transport of one probability measure to another. Methods such as diffusion models [40] 18 and flow matching (FM, [31 32 2 achieve this transport by leveraging stochastic or ordinary differential equations to map an uninformative source distribution to target distribution(s). Incorporating conditional distributions is critical aspect of these dynamical generative models, as this enables the generation of outputs with specific desired attributes. Considerable research has been dedicated designing mechanisms to condition these models. For example, in image and video generation, text prompts are commonly used to guide the output toward specified content 12 [7]. Similarly, in molecular design, conditioning on physical properties enables the generation of molecules with target characteristics [46 25).
This paper focuses condition transfer (Fig. T (a)), essential task in conditional generative modeling whose goal is to transport an arbitrary conditional distribution to another. In applications, it is often used to modify the specific attributes conditions of given instance while preserving its other features. Applications span various domains; for example, in computer vision, image style transfer has been an active area of research
[16] 48 whereas chemistry, modifying molecules to
achieve desired physical properties is vital for exploring new materials and drugs 23 25
"This work was partially done while Kotaro was intern in Preferred Networks.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

(a)
Condition transfer problem
(b)
Grouped data
Non-grouped data
*g Image datssel istels
Hg Miceaae datase with plypaca propariee
Paro
Pas
Data
Data, 2
Pox
Pa
Data A
Deto
c(1) (4)
(3)
(3)
Condition
c(+)
Conditin C
Condition
Applicabillty
Mullmariginal /I/EFM EFM
A2A-FM (Ours)
Figure 1: (a) The task is to transport Egre Peane to generate Starg P Prang for arbitrary (Carc: Ctarg) pair, where P. denotes the conditional distribution. Red and blue arrows respectively represent the case of (Carc: Ctarg) = (c()), c( and (Care. Crarg) = (c(3), cl4)). (b) Left: Grouped data is the type of dataset that can be grouped into subsets D,to of large size, whose members are i.id. samples from Pao. Many condition transfer methods including Multimariginal Stochastic Interpolants (SI) [ and Extended Flow Matching (EFM) [1] leverage this data format. Right: In non-grouped data, a sample corresponding to a given condition can be unique. Proposed method, A2A-FM, can learn condition transfer both cases in the form of pairwise optimal transport. (see Section for comparision with related works).
Flow-based generative models have been explored for condition transfer. Among others, Albergo et al. [2] proposed Stochastic Interpolants (SI), a flow-based method that interpolates between two distributions. Tong et al. [41] introduced OT-CFM, which uses optimal transport 43 to couple the data points in minibatches for conditional FM, and applied it to tasks such as image style transfer and
single-cell expression data. Meanwhile, recent methods, such Multimarginal SI 3 and Extended
Flow Matching (EFM, (21) have extended FM approaches to support (lllto-al(multiwaay transports between multiple conditional distributions.
challenge in condition transfer arises particularly when dealing with continuous condition variables. This scenario encompasses critical scientific applications, such modifying molecules based on continuous physical properties. Such continuous conditions can pose the difficulty that each condition may be associated with only a single observation æ,, resulting in limited information for each condi- tional distribution. Moreover, the situation may involve an infinite number of source-target condition pairs; applying method designed for two distributions all the pairs would be computationally infeasible. Most existing approaches cannot resolve these difficulties because they assume what we call grouped data, a type of dataset in which each observed condition is associated with a sufficiently large number of samples (Fig. [ (b)). To the best of our knowledge, no method that can learn the all-to-all transfer model for condition transfer scalably on general datasets of observation--oondiioon pairs, including non-grouped ones.
This paper proposes All-to-All Flow-based Transfer Model (A2A-FM), novel method that solves condition transfer tasks on general datasets (even in non non-grouped data setting) by simultancously learning a family of flows that approximates the optimal transport (OT) between any pair of conditional distributions. Inspired by the technique of [ 27), we develop a novel cost function with a theoretical guarantee that empirical couplings converge the pairwise OT in the infinite sample limit.
The main contributions this paper are as follows.
A2A-FM, an Fb-based method, is proposed to learn pairwise optimal transport maps across all conditional distributions from general datasets, including both grouped data and non-grouped data, irrespective of whether the conditional variables are continuous.
We introduce a novel cost function for coupling samples and prove that the coupling achieves the pairwise optimal transports in the infinite sample limit.
A2A-FM is applied chemical problem of altering the target attributes of molecule while preserving the other structure, and it demonstrates competitive and efficient performance.

---

2 Preliminaries
Suppose that have a joint distribution P on the space X' x C, where X' C IR and C C Rd denote the data space and the condition space, respectively. Let {Palc E C} denote the family of conditional distributions parameterized by c. Our method belongs the family of Flow Matching (FM) frameworks [1 32 2 Therefore, we first review FM and conditional optimal transport in the context of conditional generation, before turning to our condition-transfer task. Throughout the paper, we use the notation
Q to denote that follows the distribution Q, and =(() to denote the time derivative of a smooth curve srl). The set of integers (1....,N) N) denoted by [1:N].
2.1 Flow Matching (FM)
FM was originally developed to learn a transport from a fixed, possibly uninformative source distribution, such as the normal distribution, to a target distribution, for which samples {2ا0) are available for training. Below, review standard FM method, but adapted to learn the conditional distributions. The source distribution Po is transported to the target P. via an ODE
i - (t), xo(0) ~ Po.
so that h.., (x,()))))) follows Pe. where (=)))) 1:((() is the flow defined by the ODE. FM learns the vector field ( c) via neural networks with the parameter set 0 using the loss function
[(0) (l) t|c) (t)l
(2)
where vha 10, X' is random path such that (0) and the (1) follow Po and Pc, respectively.
2.2
Optimal Transport (OT) and OT-based FM methods
popular method for constructing the random paths (l) is minibatch optimal transport (OT-CFM, which aims to approximate the ODE of optimal transport.
41 36)
Specifically, OT-CFM uses linear path (t) t (1) + (1 ())) The key to this approach lies in how the start point (0) ~ Pa and the end point (b((1) ~ P. are coupled. OT-CFM utilizes a coupling based the optimal transport map between these two distributions. The theoretical basis for this is the optimal transport problem. Given two distributions Q1, Q2 on X, the Kantorovich formulation of the OT problem is follows:
inf IEr (2) / ||0 x2ll22 (11,
(3)
where r(Q1, Q2) is the set of all joint distributions of (I1, x2) whose marginals are Q1 and Q2 re- spectively. The square root of the infimum is known as the 2- 2-Wasserstein distance, (2(Q1, Q2) [44
When Q1 satisfies certain conditions (e.g., is absolutely continuous), this problem can be reduced to the Monge formulation, which seeks transport map T X X
inf T: x / ||= ([=
where the infimum taken over all maps T' that satisfy the push-forward condition Q2 = T#Q
Denoting the optimal transport map that solves Eq. 4 for Q: = Po and Q2 = P. as To-te. the ideal random path for OT-CFM is defined by the deterministic mapping (1) (0), which gives
(t) = t tooo + (0).
(5)
In practice, however, this map Ta- is unknown. Therefore, OT-CFM approximates this optimal coupling the minibatch level. The sampling process consists of two steps: (i) drawing independent batches (f=))) (9) ), (x²2)) from Pe and Pa, respectively, and (ii) using a finite-sample OT algorithm, such as [10], to obtain coupled pair (zg. xx)) between these batches. This pair (1), (c))) =
(zg, ze) is then used to construct the linear path (t).
3

---

2.3 Conditional Optimal Transport (COT) for conditional generation
proposed an algorithm and theory of conditional optimal trans- port(COT) that minimizes conditional Wasserstein distance. Using dataset ((x(), clt)),, their methods learn a transport from the conditional source Po.e to the conditional target distributions Pa on the joint space X' x C. When using empirical OT in such a situation, key challenge is that the condition-we tiitii optimal coupling from conditional source to conditional target does not necessarily converge the minimizer of conditional Wasserstein distance (see Example 9, D... To solve this problem, several works proposed a coupling algorithm on X' x C that provably converges to the optimal transport for each c ranges over GN the symmetric group of [1 N..
Some recent work 28
20
20 9 Their objective to minimize the following cost, where π
N
Σ
|
(1)
x(1) 1
Bl(c) C 13.
(6)
Here, the two batches Bg {(2 (0) (1) ) and B1 = {(xt) () . are drawn from the source-condition joint distribution Pe and data distribution P. respectively. Chemseddine et al. [9] have shown that, if we regard the optimal coupling T3 as the joint distribution H = E. (af) .4 xxt 2500
Cg
-Sa 2300 on (X' X C) x (X x C), then with increasing 3 and sample
size N (up to subsequence), the distribution I13 converges to a joint distribution supported on f(x, .) E (X CC) x (X' x C)} that achieves the conditional 2- Wasserstein distance
E.J [P3
(7)
where Po.e and Pa denote the conditional distributions given for the source Pe and target P. respectively. Because 7 is the average 2-Wasserstein distance of the conditional distributions given the above result implies that the algorithm with the cost (6) constructs a desired coupling that realizes the optimal transport per condition for a large sample limit.
3 Purpose and Method
In this section, we first restate the goal of A2A-FM and describe its applicability to general dataset. We then introduce the training objective and provide theoretical support for the algorithm.
The task of condition transfer for the family of conditional distributions [Palc E C} can be cast as the problem learning the (c1, C2) -parameterized transports Tap from Pe, Paz- A2A-FM is method to learn pairwise optimal transports {Ta--ey C1 02 C) that minimize the transport cost, More formally, writing the transport Tor by its induced joint distribution Pr(-lc,, 62) on X' x X with marginals Pe and Po (Kantorovich formulation). A2A-FM's goal is to learn the maps that simultaneously minimize the pairwise transport cost
Pairwise optimal transports:
|m|
x2|
22 a2/e1, 02) for all (cI, (2))
(8)
x22
Pairwise transport cost ( is the transport cost between the conditional distributions Pe, and Poz.
Applicability to general dataset: Most importantly, the scope of A2A-FM covers any type of dataset D = (xx(0) e(t) drawn i.i.d from the joint distribution P, regardless of whether the data is grouped or non-grouped, While some recent methods 21 2 they require the more restrictive grouped data to obtain the information about P. and to compute its coupling to other distributions in the dataset. The significance of A2A-FM lies in the capability of learning the pairwise OT also from non-grouped data (Fig. (b)), which arises naturally with continuously valued conditions in many scientific applications.
also aim to learn all-to-all condition transports,
3.1 Training objectives and Procedure
Aligned with the OT-CFM 41]. we propose a flow-based method through minibatch OT to achieve pairwise optimal transport. To learn T. Tez we use (c), 02) parameettreiee ODE like (T]
te (t) = v(Ter,o (o (t), (2),, where Fer,cs (0) ~ Pe.: Fer,o2 (1) ~ Pay.
(9)

---

Algorithm Training of A2A-FM
Input: (i) Dataset of sample-condition pairs == ((x(() ,,)))
). where each x:11) E X sampled
from Pata- (ii) A parametric model of vector field by X' [0, x C x C - X. (iii) The
scalar parameter B. (iv) An algorithm OPTC for optimal coupling.
Return: The parameter 0 of
1: for each iteration do
# Step 1:
Compute the coupling
2: Subsample batches B1 {(x ( ) B2 = ((xg , 22 from D.
3: Minimize 11 about =3 over N indices by OPTC
# Step 2: Update Up
4: Sample 4 unif[0, 1], E [1 N].
5: Update 0 by (o)( with
L(O) Ei |1 (t.) t.lef C2 ) (t:)ll where v.(() (1 t)x{
+ txg
6: end for
B1
82
B
82
1
1
X2
7
0
1
0
2
2
3
3
b
3
C
5
5
C
2
2
5
-4
6
7
b
(8)
Figure 2: (a) Batches Bi, Bg drawn independently from P on X' x C, where = (a, ct. (b)
Couplings between B1 and B2 by (1I). With large 8, the cost favors such that (c() -6 (c)()) ((0)
The model for the vector field e X x 1] x C x C - R& is trained with random path weyy whose endpoints are drawn from the coupled points in two minibatches, as explained below. The training objective for u is given by
E Eereereeee (t), 02) de, (t))
(10)
The crux of A2A-FM how we construct the coupling that defines ee.. We begin by sampling two independently chosen batches B, {=11 (i) )) B2
(1x) C2 ) of the same size
N from the dataset D. The objective function the coupling
N
Σ
|
(0)
=11) 12 + (1c4 - 12 11 12). x2
(11)
where GN Letting FH be the minimizer of define the path the the with
(cI, (2) (e +0 for i [1 M]by by
(t) (1- (1-t)x + tx +3(0)
To transport sample Te, Pay sample in Pe, using the trained v, follow the same procedure as in standard FM and solve the ODE ( forward from t = 0 to = with initial condition Tories (0) Teg- The algorithmic training procedure summarized in Algorithm
3.2 Intuition behind the objective
The objective function (TT controls the continuity of Pa with respect to in the transfer maps through balance of the hyperparameter and the sample size [D] = N. When there is only one sample Te
5

---

per each it is impossible to learn P. separately, let alone the optimal transport between Pe, and Pag. However, if small enough, data points Ie Pa and Tez ~ Pay can be coupled for almost any 01 and C2, so that the information between Pej and Paz is shared. Conversely, as B + 0o, (11
would force ( to be close to ck (1) requiring that for each pair of (c), e2), the sample #1 with condition of near C₁ is coupled with (2) with condition c near C2. This would allow the samples with similar c values to be transported collectively.
We illustrate our way of coupling when C is finite (Fig. where the functionality of [I] is more
intuitive. In that case, for - oc, the second term in 11 would diverge to infinity unless it is zero However, for [C] < (x, the sample size N increases, there will be many nontrivial permutations
with 6 Ck (I= 1,, which makes the second term zero. Therefore, the minimum of (II)
(i)
(() for all E [] NJ, 1,2]. We is attained by permutation in G { E GN | dh Ck
can see the function 6 more clearly by partitioning [1 N]] into the subsets closed within C;
[1 N] (la),c)ecc (c1,e)) where (()) {i [1 [1:N] C/ Ck. (k 1, 2) For example, in Fig. 2 the subset 6 consists of permutations that act separately on {1, 8). (3, 5)4())((4) and {6}. Consequently, permutation in G is decomposed into permutations within ((().og)) and Go is the the product of symmetric groups Uee).ez) The minimization
TT reduced the sum of
EiEU |z| - =10)
og))
where each summand is the standard cost of transporting data from class in B1 to class in B2. Therefore, the optimal x3 achieves the optimal coupling among {=] (i) i E U(er.e))) {x2 (+)
i ). In Fig. 2 we see U(a.b) = {1, 8,8, and the optimal z would optimally transport (1) _x (x()) (7) _ (8)) in Class a to {--
(7) _.(1) in Class b. As N increases, each ((c),ea) contains larger number of indices, that z better approximates the OT from Pe, to Po, for every (c1, 02).
Interestingly, although we have discussed the case of finite C in this section, this argument extends to the general case of continuous We claborate on this fact in the next subsection.
3.3 Theoretical guarantee
We show that the coupling given by IT converges to pairwise OT. We state the result informally here and defer the formal result to Appendix
Proposition 3.1 (Informal). Let IT's be the joint distribution on (X' x C) (X' x C) defined by the coupling & that minimizes
11 that is
IT; - D=
((x) 9
.6
).2
-300 +300
.t2
Then, for any sequence 3y. + 0o, there exists an increasing sequence of the sample size N₃, such that I13, converges to II for which (., C1, C2), the corresponding conditional distribution of (I1, x2) given (c1,02), is a joint distribution on x X that achieves below for almost every (c1,02):
f |||| - 221 2e1--21laait
62) W?((Po, Poz).
(12)
We will show experimentally in Section 5 that, with the objective function 11 A2A-FM learns approximate pairwise optimal transport even on the generic dataset drawn from joint distribution P.
Related Works
other works such as 41 32] used flow models to transfer between two distributions with different attributes (e.g.. images of smiling faces vs. non-smiling). Kapušniak et al. [24] also incorporated the geometry of the data manifold to the flow. Although one can apply these methods to many pairs of conditions individually, such strategy is computationally inefficient for the all-to-all condition transfer task.
In addition to the direct application of optimal transport transfer tasks ||I 35]
In this regard, Albergo al. 3 presented Multimarginal SI, which uses the principle of generalized geodesic [5). In this framework, produces Te, Ta for arbitrary C1 C2 lincar ensemble of

---

(a) Grouped data (Fig. 3 (a)
(b) Non-grouped data (Fig. (b)
Method
MSE from pairwise OT
Method
MSE from pairwise Of
Ours (Tz) Generalized geodesic
(5.81 + 2.22) X 10
A2A-FM (Ours)
(1.51 = 0.17) x 10-2
Partial diffusion
(6.77 = 0.14) x 10-2
(1.03 + 0.04) x 100
Multimarginal SI
(4.90 + 0.28) 10-2
Table 1: The discrepancy from pairwise OT for synthetic data. The error in the table shows the standard deviation over 10 evaluation runs.
the optimal transports Te)) from arbitrary barycentric source Po to Pam. By design, this approach requires grouped data for the learning of Tg_pote) Isobe et al. [21] take another approach using a matrix field for the system of continuity equations, from which an arbitrary To, -e, can be derived. For the learning of the matrix field, grouped data is again required. In another related note, [6] recently proposed method that can transport any source distribution to target distribution by encoding the source population itself, and conditioning the vector field with the encoded population. This approach again requires grouped data for the learning of population embedding.
Meanwhile, Manupriya et al. 33 uses COT 9 28 to learn the OT map from Pe to Qe for each e when given the jointly sampled dataset of (z,, c) P and (y, c) ~ Q. They handle non-grouped data, but they cannot create arbitrary Te- Although the proof technique in COT is an inspiration for our proof, COT pursues conditional generation as opposed to the fundamentally different task of all-to-all condition transfer. In this regard, we emphasize that the objective (TT) is fundamentally different from the one used in conditional OT (6). Note that, directly applying the method of similar nature as COT by a simple replacement of a single instance e with (C1, C2) would fail to learn a map between 01 F 02; if we learn COT for v(-le1, 02) via (6) with (c1, 02) in place of (cg.ct), the model
will only learn transfers between a pair of Cj and C2 that are close to each other.
5 Experiments
In this section, we present a series of experiments to validate our theoretical claims in Section 3 well as the effectiveness of our model on real-world dataset. The methods we use for comparison include (1) partial diffusion (25), a method that adds limited noise to source sample and then denoises using classifer-free guidance to obtain a target-conditioned sample while preserving smilarity to the input (2) Multimarginal SI (see Section 4 and (3) an application of OT-CFM in which we use the sample dataset itself as the source Pg. In non-grouped settings, we generally choose 1/dde, where de is the dimension of C (see Appendix A.3 A.3 for theoritical backgrounds for this rate).
5.1 Synthetic Data (grouped and non-grouped data)
We demonstrate with synthetic data that both the coupling ग used to make the supervisory paths and the trained vector fields v in A2A-FM approximate the pairwise optimal transport.
We compared the supervisory vector fields die, ca of both A2A-FM and generalized geodesics against the pairwise OT on a synthetic dataset of three conditional distributions. each having the distribution of two-component Gaussian mixtures (Fig. B (a)). Note that A2A-FM more closely resembles the numerical approximation of the true pairwise OT [1]. Denoting by (x1, -x,(
Grouped data:
(vée, eg (0), ea, (1)) the coupling of Pa,, Pop by a given method, we also
validated this result quantitatively with Ez,,c, [I]TT, -cz (21) e1 -e2 (z1) [[]] (Table | (a)), where
co is the ground-truth pairwise optimal transport map that minimizes
Non-grouped data: We compared the transfer map of A2A-FM against ground-truth pairwise OT, partial diffusion, and Multimarginal SI on non- con-grouped data. The dataset consisted of samples from a 2D polar coordinate quadrant (r E [1, 2], was uniform along r. This represents non- non-grouped data. For Multimarginal SI, we discretized 6 into K bins because can only handle grouped data. See also Appendix B We visualize the learnt transport in Fig. 3 (b). Partial diffusion produced nearly random couplings. Multimarginal SI failed to generate the target distribution's 's marginals because of discretization. Quantitative evaluations in Table (b) also confirm that A2A-FM approximates pairwise OT more accurately than the rivals.
[0. 222)), where was a continuous condition and Po
7

---

Pairwise OT
Ours
(b)
Pairwise OT
Ours
Partial diffusion Multimarginal Si
Generalized geodesic
(
Pnjectory
initial points
target points generated points
Figure 3: (a) Results for grouped data. The sample size was 103 in each of conditions {clo), c(1), c( and 10*, (b) Results for non-gru non-grouped data. The gray points in the background
show samples from the training dataset. The presented pairwise OT is numerical approximation by
13.
The red lines in the right column shows the bins for training Multimarginal SI.
1.0
(Oun)
COATI-LDM
(K-10)
0.8
AZAFM (Ours)
0.8
PD+CFG (T=300)
OT-CFM
Success ratio
(T=1000)
0.6
Pri
0.2
0.2
0.2
0.4
0 0.8
1.0
Pr(eram
cx]
0.0
25
100
500
1000
50000
Figure 5: Sampling Efficiency Curve for LogP-TPSA benchmark. See Appendix BB. and Table:]oor notation. K is the number of discretization bins.
Maximum oracle calls
Figure 4: Sampling Efficiency of A2A-FM and the partial diffusion model of (26].
Nearby Sampling for Molecular Optimization
Molecule design is multi-constraint task: a lead must bind strongly and selectively while remaining drug-liks and nontoxic. Random candidates rarely meet every criterion, so chemists traditionally resort to scaffold hopping-ggnerattie generatinn many close analogs of existing molecules in the hope that some retain favorable properties while mitigating undesirable ones. The machine learning approach frames this as nearby sampling, in which the goal to sample molecule with target property within a structural neighborhood of reference molecule.
QED experiment: For a Quantitative Estimate of Drug-likeness (QED) optimization task 23 the goal is to transfer a molecule with QED < 0.8 to QED > 0.9 while maintaining the original structure measured by Tanimoto similarity, whose threshold is set to 2 0.4. A2A-FM was trained on a 500K ZINC22 subset, where X represented the latent space of molecular representations (26] and 3-32 was the space of QED embeddings.
as much as possible. The similarity the original
Following the same protocol as previous methods [26] 19 38 we evaluated the al algorithm by the
success rate of discovering a molecule of desired property (QED, similarity) within the prescribed
Method
AUC Values
Method
Success (%)
PD+CFG (T=500)
0.027
PD+CFG (T=300)
0.450
DESMILES 34
76.9
PD+CFG (T=1000)
0.060
QMO 19
92.8
SI (K=10)
0.583
MolMIM 38
94.6
OT-CFM
0.819
COATI-LDM 25
95.6
A2A-FM (Ours)
0.990
A2A-FM (Ours)
97.5
Table 3: AUC Values of LogP-TPSA task.
Table 2: Nearby sampling success rate.
PD+CFG is the Partial diffusion method with Clas- sifier Free Guidance with T noise steps.

---

number of oracle sample calls made from the initial molecule (Fig. 6 See Table for the comparison
against SOTA [26]. and see Appendix [B.2 for details and generated samples.
The results in Table 2 show that A2A-FM surpasses
Sampled molecules
SOTA on default oracle calls (50000). Fig. 4 also shows that we achieve much higher success rates across all maximum oracle calls. We note that, un- like QED which is computationally inexpensive to validate, other properties like AAG and vertical ionization potential 47 require costly calculations like DFT and FEP. and biological assays. For such cases, making many oracle calls can be in- feasible, so achieving high sampling efficiency is a feat of significant scientific interest.
QED 0.93
Inilial molecule
Similaity = 0.62
45
29 50
Similarity - 0.76
QED 0.77
QED 0.74
Similarity 0.25
LogP& TTPS experiment: To evaluate the sampling efficiency in all-to-all condition transfer task, we con- ducted the nearby sampling similar the QED exper- iment, except that we chose random 1,024 molecules from ZINC22 as the initial molecules, and aimed changing their two other properties (LogP and TPSA) to a randomly selected pair of target values embed- ded in C = TR For evaluation, we measured Tanimoto similarity &Tani and normalized condition error Cere for each transferred sample, and plotted Pr(cerr S a) against Pr(sTani S a) to visualize how much we need to trade off 8Tani to increase the probability of sampling molecules within the desired Gerr threshold. A2A-FM substantially outperformed rivals AUC (Fig. 5 and Table 3p. See also Appendix B.3
QED 0.94
Figure 6: QED optimization task. An initial molecule is marked success if there exists a molecule with Tanimoto similarity 2 0.4 and QED > 0.9 among the molecules sampled by transferring the original. The size of the sample set is called the maximum oracle calls. See Table [ for the ratio of initial molecules marked success in the dataset of Jin et al. [23]
5.3 Additional experiments
To apply methods like B 6 requiring grouped data (Fig. (b), left), data D = {xz, } is needed to be partitioned into K bins of M samples with similar c's, constrained by |D|= = M x K. A choice of large K (small M) impairs reliable P. estimation and, for Multimarginal SI, leads to O(K²) optimization costs. Mcanwhile, small K (large M) coarsens the partitions, reducing precision (see Fig. 9 Appendix C). Rising K up to K = [D] is computationally infeasible for Multimarginal SI since costs O 2(22 for pairwise optimization. In contrast, the computational cost of A2A-FM depends only on DD. which scales identically to OT-CFM (different only in its cost function) and is independent of K for fixed DD), as shown in Table 6 Appendix C
Computational cost:
The choice of 3: As in the case of (9) the selection of the parameter 3 is crucial to the performance. For grouped datasets, 3 can be set sufficiently large (see Section 3.2). However, the same method does not directly apply to the general case due to the trade-off; a large 33 may fail to accurately approximate pairwise OT, while small 3 can lead to reduced precision in the terminals (Fig, Appendix C). Although the difficulty of this trade-off is a potential limitation of A2A-FM, we empirically observed that the choice of B = N1/(22d) inspired by 14] was effective in achieving good balance. We validated this heuristic and assessed the method's 's robustness to 3 on synthetic non-grouped datasets (Fig. 7 Appendix C). The results confirm that the optimal 3 aligns with our heuristic and demonstrate the method's 's stability across range of /3 values spanning order of magnitude. In Appendix A.3 we also dicuss a theretical necessary condition on the rate of 3 that the convergence in 3.1 holds.
Discussion
An advantage of using pairwise OTs: Section demonstrated that A2A-FM performs competitively in transporting Pc, to Po, in real-world settings. The efficacy of A2A-FM in condition transfer task may be partly explained by its connection to the function representation theorem (30] This theorem states that if a random variable X has a feature C = gx(X) with ge deterministic, then there exists an independent feature 7 (x) such that (C, Z) can generate X (i.e., X decomposes into C
P(- ((X])(X) c1) Po (-) to Zegz may leverage the and 2). When transporting sample Ze

---

representation 2e, = f for an invertible function f. and perform transport by mapping f(z c1) to Iry = C2) while keeping the z-component fixed. This representation closely related to Nomain-adversarr training [1]] which effectively seeks in this expression for condition (domain)- invariant inference, and is used in [5 for transporting between conditional distributions. is also known that, in the case of one-dimensional x, the transport of type ((2,c1) H f(2, agrees with the OT with quadratic cost [42] Sec 2.2]. Also, (1,( f32 e2)]l [lf(z1. - f(21, for all (zi, et) and (z2, c2), then the OT trivially favors plan that does not alter Such situation may arise when z takes discrete values such that any modification to z results in a larger change a than modifying c. This is also in Fig. B (a), where each cluster is transported to another cluster.
Cycle consistency: It is reasonable to require the consistency property Te-->e Te id for transport maps. In A2A-FM, this can be ensured by enforcing the antisymmetry condition Vaycs .0g = -og.ey in the model of vector fields. More generally, in some transfer applications, one might prefer the cyclic consistency property: Teg Teg-ve
Ta - Tep-e -ta- For instance, some prior works (21 [ build this cycle consistenct directly into their models. However, A2A-FM may not necessarily enforce this property. because OT does not generally satisfy it. On the other hand, adopting this antisymmetry in QED experiment enhanced the success rate from 94.65 4..6% to 97.5% suggesting some justification for this regularization. While the full effect of this anti-symmetry restriction is a matter of future work, we emphasize that A2A-FM provides a scalable method for non-grouped data, whereas conventional methods with cycle consistency [21] are inapplicable or computationally prohibitive for such data classes (see Appendix 5.3 for details on computational cost).
A2A-FM on large grouped data: Many text-labeled image datasets can be regarded as grouped data, as they often provide sufficient data for each independent textual condition. For example, while a sample "yellow dog' may be rare, there can be many 'dog' samples and 'yellow" samples. As demonstrated in Appendix C A2A-FM scales effectively to such image datasets. For grouped data, when category-specific classifiers can be trained effectively, classifier-dpennentt methods are viable options for condition transfer [49 37.. However, we note that A2A-FM is designed for more general situations in which such classifiersregressorr can unavailable or unreliable.
7 Conclusion
We proposed A2A-FM, an FM-based condition transfer method that can learn pairwise OT from general data type, including the domain of continuous conditions. To achieve this purpose, we introduced an objective function that realizes pairwise optimal transport in the infinite sample limit. The balance of dataset size and the hyperparameter 3 can pose limitation, but we provided stable heuristic whose theory can be subject of possible future research. We applied A2A-FM to a chemical application of modifying the target attribute of molecule, demonstrating wfte-of-theeart performance.
8 Acknowledgements
This work was partially conducted during KI's summer internship at Preferred Networks. KF is partially supported by JST CREST JPMJCR2015 and JSPS Grant-in-Aid for Transformative Research Areas (A) 22H05106.
References
[1] R. Abel, L. Wang, E. D. Harder, B. Berne, and R. A. Friesner. Advancing drug discovery through enhanced free energy calculations. Accounts of chemical research, 50(7)1625--162. 2017.
[2] M. S. Albergo, M. Goldstein, N. M. Boffi, R. Ranganath, and E. Eijnden Stochastic interpolants with data-dependent couplings. arXiv preprint arXiv:2310.03725, 2023.
[3] M. S. Albergo, N. M. Boffi, M. Lindsey, and E. Vanden-Eijndenn Multimarginal generative modeling with stochastic interpolants. In The Twelfth International Conference on Learning Representations, 2024. URL
10

---

[4] J. Alfonso, R. Baptista, A. Bhakta, N. Gal, A. Hou, V. Lyubimova, D. Pocklington, J. Sajonz, G. Trigila, and R. Tsai. A generative flow model for conditional sampling via optimal transport. In NewrlPS 2023 Workshop Optimal Transport and Machine Learning, Dec. 2023.
[5] L. Ambrosio, N. Gigli, and G. Savare. Gradient flows: In metric spaces and in the space of probability measures. Lectures in Mathematics. ETH Zürich. Birkhauser Verlag AG. Basel, Switzerland, 2 edition, Dec. 2008
[6] L. Atanackovic, X. Zhang, Amos, M. Blanchette, L. Lee, Y. Bengio, A. Tong, and K. Neklyudov. Meta flow matching: Integrating vector fields on the wasserstein manifold. In ICML 2024 Workshop on feeometry-groundee Representation Learning and Generative Modeling, 2024. URL
[7] G. Carlier, A. Galichon, and F. Santambrogio. From knothe's transport to brenier's map and a continuation method for optimal transport. SIAM Journal on Mathematical Analysis, 41(6): 2010.
[8] Chemseddine, P. Hagemann, C. Wald, and G. Steidl. Conditional Wasserstein distances with applications in Bayesian OT flow matching, 2024.
[9] J. Chemseddine, P. Hagemann, G. Steidl, and C. Wald. Conditional Wasserstein distances with applications in Bayesian OT flow matching. Journal of Machine Learning Research, 26(141); 1-47, 2025.
[10] M. Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In
C. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K. Weinberger, editors, Ad- vances in Neural Information Processing Systems, volume 26. Curran Associates, Inc., 2013. URL
http://prrrcceennnneeneeeurreppp neurips. cc/ wwtppppr.pp..../.p http:/ppe_ffleessaaeeeeeeeee pttll///ppper/23///ieee/ 47710067775477771177775555666663 .pff
[11] B. B. Damodaran, B. Kellenberger, R. Flamary, D. Tuia, and N. Courty. Deepjdot: Deep joint distribution optimal transport for unsupervised domain adaptation. In V. Ferrari, M. Hebert, C. Sminchisescu, and Y. Weiss, editors, Computer Vision - ECCV 2018, pages 467-483, Cham, 2018. Springer International Publishing. ISBN 978-3-030-01225-00
[12] P. Dhariwal and A. Q. Nichol. Diffusion models beat GANs image synthesis. In A. Beygelz- imer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
[13] R. Flamary, N. Courty, A. Gramfort, M. Z. Alaya, Boisbunon, S. Chambon, L. Chapel, A. Corenflos, K. Fatras, N. Fournier, L.. Gautheron, N. T. Gayraud H. Janati, A. Rakotomamonjy. I. Redko, A. Rolet, A. Schutz, Seguy, D. J. Sutherland, R. Tavenard, A. Tong, and T. Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78): 2021. URL html
http://jmlr. wwwggpapers//22220-444
[14] N. Fournier and A. Guillin. On the rate of convergence wasserstein distance of the empirical measure. Probability Theory and Related Fields, 162:707-738, 2015.
[15] Y. Ganin, E. Ustinova, H. Ajakan, P. Germain, H. Larochelle, F. Laviolette, M. March, and V. Lempitsky. Domain-adversarin training of neural networks. Jourmal of Machine Learning Research, 17(59):1-355 2016. URL http:://jmx: wrg. httpg/aapers/v1/115... 239 .ttll
[16] L. Gatys, A. Ecker, and M. Bethge. A neural algorithm of artistic style. J. Vis., 16(12):326. Sept. 2016.
[17] J. Ho and T. Salimans. Classifier-free diffusion guidance. In NewrIPS 202/ Workshop on Deep Generative Models and Downstream Applications, 2021. URL httpss:/ net/ forumPid=qu8AKsfE0I
[18] Ho, A. Jain, and Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, volume 33, pages 6840-6851. Curran Associates,
Inc., 2020. URL htt.:tt::/ppoceediniiaee www.wps...e.s....ooc.oo.o cttppp.cccppppr.....p htt//paper____ llappr_fflleee wwtp.s//papae///////////a///

---

[19] S. C. Hoffman, V. Chenthamarakshan, K. Wadhawan, P.-Y, Chen, and P. Das. Optimizing molecules using efficient queries from property evaluations. Nature Machine Intelligence, 4(1): 21-31, 2022.
[20] B. Hosseini, A. W. Hsu, and A. Taghvaei. Conditional optimal transport on function spaces. SIAM/ASA Journal on Uncertainty Quantification, 13(1):304-338 2025.
[21] N. Isobe, M. Koyama, Zhang, K. Hayashi, and K. Fukumizu. Extended flow matching: a method of conditional generation with generalized continuity equation. arXiv preprint WrXiv:2402.18839 2024.
[22] Y. Jiang, Z. Huang, X. Pan, C. C. and Z. Liu. Talk-to-edit: Fine-grained facial editing via dialog. In Proceedings of International Conference on Computer Vision (ICCV), 2021.
[23] W. Jin, K. Yang, R. Barzilay, and T. Jaakkola. Learning multimodal graph-to-graph translation for molecule optimization. In International Conference on Learning Representations, 2019.
[24] K. Kapušniak, P. Potaptchik, T. Reu, L. Zhang, A. Tong, M. M. Bronstein, J. Bose, and F. Di Giovanni. Metric flow matching for smooth interpolations on the data manifold. In The Tiihtt-eighth Annual Conference on Neural Information Processing Systems, Nov. 2024.
[25] B. Kaufman, E. C. Williams, R. Pederson, C. Underkoffler. Z. Panjwani, M. Wang- Henderson, N. Mardirossian, M. H. Katcher, Z. Strater, J.-M. Grandjean, B. Lee, and Parkhill. Latent diffusion for conditional generation of molecules. bioRxiv, 2024. doi: 10. 10100024.000 22. 609169. URL https:://wwwt biorxiv 609169
08. 22
[26] B. Kaufman, E. C. Williams, C. Underkoffler, R. Pederson, N. Mardirossian, I. Watson, and J. Parkhill. Coati: Multimodal contrastive pretraining for representing and traversing chemical space. Journal of Chemical Information and Modeling, 64(4):1145-1157 2024.
[27] G. Kerrigan, G. Migliorini, and P. Smyth. Functional flow matching. In S. Dasgupta, S. Mandt, and Y. Li, editors, Proceedings of The 27th International Conference on Artificial Intelligence and Statistics, volume 238 of Proceedings of Machine Learning Research, pages 1934-3942 PMLR, 02-04 May 2024. URL httttp:://proceedinng barggggs.lrr.... peress/v238/kerrte
html
[28] G. Kerrigan, G. Migliorini, and P. Smyth. Dynamic conditional optimal transport through simulation- free flows. In The Thirty-eighth Annual Conference on Neural Information Process- ing Systems, 2024. URL
net/forummice tit.idddttttuntttt.
[29] T.SS. Lee, B. K. Allen, T. J. Giese, Z. Guo, P. Li, C. Lin, T. D. McGee Jr, D. A. Pearlman, B. K. Radak, Y. Tao, et al. Alchemical binding free energy calculations in amber20: Advances and best practices for drug discovery. Journal of Chemical Information and Modeling, 60(11): 2020.
[30] C. Li and A. E. Gamal. Strong functional representation lemma and applications to coding theorems. In 2017 IEEE International Symposium on Information Theory (ISIT), pages 589-593, 2017. doi: 10. (009///TT.217.8800666.
[31] Y. Lipman, R. T. Q. Chen, H. Ben-Hamu, M. Nickel, and M. Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations, 2023. URL
[32] X. Liu, C. Gong, and Q. Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. The Eleventh International Conference on Learning Representations, 2023.
[33] P. Manupriya, R. K. Das, S. Biswas, and S. N Jagarlapudi. Consistent optimal transport with empirical conditional measures. In S. Dasgupta, S. Mandt, and Y. Li, editors, Proceedings of 27th International Conference on Artificial Intelligence und Statistics, volume 238 of Proceedings of Machine Learning Research, pages 3646-3654. PMIR, R, 02-04 May 2024. URL. https://proceedingg mlr, prwss/ httpas//222//aampia2a22... html
12

---

[34] P. Maragakis, H. Nisonoff, B. Cole, and D. E. Shaw. A deep-learning view of chemical space designed to facilitate drug discovery. Journal of Chemical Information and Modeling, 60(10): 4487-4496, 2020.
[35] Y. Mrouch.
Wasserstein style transfer. In International Conference on Artificial Intelligence and Statistics, pages 842-852. PMLR, June 2020. UtLhtttt//prooeeeingg mlr. press/ v 108/mroueh20a. html ISSN: 2640-3498.
[36] A.-A. Pooladian, H. Ben-Hamu, C. Domingo-Enrich, B. Amos, Y. Lipman, and R. T. Q- Chen. Multisample flow matching: Straightening flows with minibatch couplings. In A. Krause, E. Brunskill, K. Cho, Engelhardt, S. Sabato, and J. Scarlett, editors, Pro- ceedings of the 40th International Conference on Machine Learning, volume 202 of Pro- ceedings of Machine Learning Research, pages 28100-28127, PMLR, 23-29 Jul 2023. URL htttps://proceedingg. mlr httppp/eellaaaa...ppsp/////////p
html
[37] K. Preechakul, N. Chatthee, S. Wizadwongsa, and S. Suwajanakom, Diffusion autoencoders: Toward a meaningful and decodable representation. In Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pages 10619-10629, 2022.
[38] D. Reidenbach, M. Livne, R. K. Ilango, M. Gill, and J. Israeli. Improving small molecule generation using mutual information machine. arXiv preprint arXiv:2208. 09016, 2022.
[39] R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer. High-resolution image synthesis with latent diffusion models. In 2022 IEEF/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 10674-10685, 2022. doi: 10. 10.1100922R222888222222.2.
[40] J. Song. C. Meng, and S. Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations, 2021. [ttp:atts:///peerreeee http:eerrriiie..ee/oorrrr.r id=StigiarCHLP
[41] A. Tong. N. Malkin, G. Huguet, Y. Zhang. J. Y. Bengio. Improving and generalizing flow-based generative models with minibatch op- timal transport. Transactions on Machine Learning Research, 2024. ISSN 2835-8856. URL
K. Fatras, G. Wolf, and
https:://penreevie. wwwee//ooumiidieeeeeeeeee
[42] C. Villani. Topics in Optimal Transport. American Mathematical Society, 2000.
[43] C. Villani. Optimal Transport: Old and New. Grundlehren der mathematischen Wissenschaften. Springer, Berlin, Germany, Dec. 2009.
[44] C. Villani. Optimal Transport, old and new. Springer, 2009.
[45] L. Wang. Y. Deng, J. L. Knight, Y. Wu, B. Kim, W. Sherman, J. C. Shelley, T. Lin, and R. Abel. Modeling local structural rearrangements using fep/rest: application to relative binding affinity predictions of cdk2 inhibitors. Journal of chemical theory and computation, 9(2):1282-129) 2013.
[46] M. Xu, L. Yu, Y. Song, C. Shi, S. Ermon, and J. Tang. GeoDiff: A geometric diffusion model for molecular conformation generation. In International Conference on Learning Representations, 2022.
[47] C.-G. Zhan, J. A. Nichols, and D. A. Dixon. lonization potential, electron affinity, electronega- tivity, hardness, and electron excitation energy: molecular properties from density functional theory orbital energies. The Journal of Physical Chemistry A, (11())):4444411999 2003.
[48] Y. Zhang, N. Huang, F. Tang, H. Huang, C. Ma, W. Dong, and C. Xu. Inversion- style transfer with diffusion models. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 0014610116, 2023.
[49] M. Zhao, F. Bao, C. Li, and J. Zhu. EGSDE: Unpaired image-to-image translation via energy- guided stochastic differential equations. In Advances in Neural Information Processing Systems, Oct. 2022.
[50] F. Zhu, F. A. Bourguet, W. F. Bennett, E. Y. Lau, K. T. Arrildt, B. W. Segelke, A. T. Zemla, T. A. Desautels, and M. Faissol. Large-scale application of free energy perturbation calculations for antibody design. Scientific Reports, 12(1):12489, 2022.
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
we provide a theoretical statement to support our claim and provide its proof in the Appendix. We also provide empirical support to our claim, and large scientific dataset to support the scalability of our method.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discussed about the potential limitation of the proposed method in Sec. 5.3 and also refered to it in the conclusion of the paper (Sec.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and prool's
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
14

---

Answer: [Yes]
Justification: Our theoretical result (Prop. 3.1 (informal) Thm. A.1 (formal)) stated with full set of assumptions and its proof is in Appendix A Prop. 3.1 in the main paragraph.
We also provide intuitive explanations
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof's provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide training and data preprocessing code in the supplementary ma- terial, offering full reproducibility. We also elaborate on the details of the experiments in Appendx
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
15

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: Code and instructions for training and data-preprocesssing is provided in the supplementary material. Datasets used in this paper (ZINC22, CelebAHQ-Dialog) be downloaded via internet.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ /twpps.ccppuulleeee wwtpcccppppccccceeeoec
/wweeeeeeeooeemmmeeceeee commmsssions aessonooiicy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: Full details of the experimental settings are provided as code the supplemen- tary material while important points are also refered in Appendix B
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
For the synthetic experiments we included standard deviations calculated through multipule runs of evaluations. For chemistry experiments, did not include error bars since this was impossible because of the computational cost (more than 24 hours for evaluation run).
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
16

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative eπror rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We reported the number and types of GPUs used in the experiments in Appendix B
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental rums well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: Authors read the code of ethics and throughly confirmed that the paper does not contain viotates it.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [NA]
17

---

Our reaserch is rather foundational one and its societial impact is less since it does not include any deployments or tied to particular applications.
Justification:
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has socictal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our paper does not poses such risks since it does not provide any trained models or preprocessed data.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but encourage authors to take this into account and make best. faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All the models and datasets used in this project is cited in this paper and the reference to its URLs are refered in the README file included with the code.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
18

---

The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g.. CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, pwwpsssipppsseeeee thcode com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: The code included in the supplementary material has instruction documents inside. New datasets models were not introduced.
Guidelines:
The answer NA means that the paper does not release assets.
Researchers should communicate the details the dataset/code/moodl as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: This paper does not include any experiments using crowed sourcing nor human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
supplemental material is fine, but if the main contribu-
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
19

---

This paper does not include any experiments using crowed sourcing nor human subjects.
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
Declaration LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The proposed method does not involve LLMs as any important, original. or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM for what should or should not be described.
20

---

Technical Appendices and Supplementary Material
A Theories
A.1 Notations
In this section, we provide the notations that we will be using in the ensuing mathematical formulations and statements.
C Rd= Space of observations.
C C Rde; Space of conditions.
P(A): the set of all distributions on measure space A.
Wi.... the projection onto .1..... igth component. For example, wi AxB + A, we C.
Ax B - B, and 22, A
M2) == (1I P(Ax
#Il - μ4 and #2 = (12) for 41, E P(A). QAIB=b E P(A): Regular conditional distribution of Q E PP(A x B) where A, B are respectively the random variables on A, B with distribution Q.
P & Q: For probability distributions on X and Q on V. P & Q denotes the product probability defined by (P @ (A x B) (B))Q(B). This gives independent marginals.
A.2 Convergence to the pairwise optimal transport
Let P2(X)) denote the space of probabilities with second moments; i.e.,
P2(x)= = 3 PEP(X) / |=]ldd((x) < 200 The 2-Wasserstein distance W2(P, for P.Q E P2(X) defined by inf / W (P.Q) (PT(P xxx
|x
y/1 dll(z y).
In the sequel, P2(X)) is considered be metric space with distance W2.
A.2.1 Conditional Wasserstein distance
As preliminaries, we first review the conditional Wasserstein distance, which was introduced in [9] Let Z be subset of IR"" which is used set of conditions and let P E P(x' x 2). For
4 (28).
example, in the setting of conditional generations, Z C.
To relate the transport plans for conditional distributions and the plans for joint distributions, we introduce the 4-plans as [9] Let v be probability P(Z). For E P(Z), we define the class of joint distributions with the same marginal on Z:
P(X'' z) (P P(X 2) (#2#P) = v).
For two probabilities P.Q E E(X' x 2; v), define (P, by
(P. = (II E P((x 2 +2 A#v),
(13)
where A Zx 2, > (2, 2). the diagonal map and 22,4 the projection map to the second and fourth component.
For P,Q x(X'xZ;)) with finite p-th moments, the conditional Wasserstein distance W (P,((
is defined as follows 9 28 4:
/ [[(1,) =1) - (x2, 22) 22)l|P
W p.u (P.Q) =
inf
(14)
(ner)(P.))
It is known [Prop. 1.9 that the conditional p- Wassersstiin distance is in fact the average p- Wasserstein distances of conditional distributions:
Proposition A.1. Let (X1, Z) and (X2, Z) be random variables on X x Z with distributions pl and 22 E P(X 2;v), which both have finite p-th moments. Then,
(P. p2)p
Wp(P P Pdu(s)
(15)
21

---

A.2.2 Convergence of A2A-FM
Let Di = (a1 =1 x C). Recall that the proposed A2A-FM uses the optimal plan or permutation distribution P P2 (X
and D. = (x 1 c2 سر be two independent copies of i.idd samples with
N which achieves the minimum:
N
min TEON N
Σ a
1T 22
Bllc p + - 1p
(16)
We formulate " No as an optimal transport plan that gives an empirical estimator of the conditional Wasserstein distance. For this purpose, we introduce augmented data
D ((1))) of .
and
(0) 91 D2 : (1x2
(3) ) ,
(17)
For notational simplicity, we use 200 (c) (1) o and thus, Di ((x1) ) and D2
{(a) )))
First, it is easy to see that the objective function 16 of the coupling is equal to
min REGN
Σ da (( (
20) (a)
(r(1)
((x())))
where the cost function dø on X (C C) is given by
dz((21, 21), (I2, 22)) = ||= - x22l + B||21 - -22l1
(18)
This means that the minimizer of (16) gives the optimal transport plan between the augmented datasets Dj and D2 with the cost d3. We express this optimal plan by I atomic distribution in P2((X' C x C)= x (X xC x C)).
Second, consider the following random variables and distributions on X' x C x C:
(X1, x == 112,4 ② P):
(19)
(X2, C2, Q² w32, P).
For Q', the variables X1 and C1 are coupled with the joint distribution P. while C2 is independent; for Q2, on the other hand, X2 and C2 are coupled with distribution P. and Ci is independent. The datasets D, and D2 are obviously i.i.d. samples with distributiion Q* and Q², respectively. Note that the distribution of Z == (C1, C2) is equal to v := (P @ PP) E P(C x C), and common to Q' and Q². An important fact is that the conditional distributions given Z satisfy
Q'2, 2=(er = Qx Pxi|ci=a
2x ==(e))0) = 2x =Pxs|c==er
(20)
where we use the independence between (X1, C1) and C2 for Q' and similar relation for Q?.
Let @N (a = 1,2) be the empirical esitribution of Da. and Uy be that of (0)) = (cf") .02 (i)) that is,
2% 1
(af2 20) (a 1,2) and DN N Σ 8.1 E Pz(C x C). (20)
Note that, since z-component in Di and D2 are identical, the coupling 11 is a plan in T(Q), (x:
With the above preparation, it is not difficult to derive the next result, which shows that, as regu- larization 3 and sample size N go infinity, the above empirical optimal plan IT converges, up to subsequences, to an optimal plan that gives optimal transport plans for all pairs of conditional probabilities PxiC=e, and Px|C=or
22

---

Theorem A.2. Suppose that X' and C are compact subsets of R and IR',, respectively, and that Bk is an increasing sequence of positive numbers. Let T be the optimal transport plan in ( for W2.9, (2 ②x), defined as above. Then, the following results hold.
(i) There is subsequence (N&) so that I converges in P2((X 2) to T3( Q²) that is an optimal plan for (2,' Q2), where 2 = C x C.
(ii) If the optimal plan 11 is identified as a 3-plan (Q',, Q?), it satisfies
JocCJexx |21 -
- 221l2dlT (I1, 32 2 P (c)))d=
- -f exc (P(xXI :x ((a) (21)
where 1a II is the disintegration of the 3-plan II. (iii) Consequently, ee is the optimal plan to give W2 (2 Px|c= for (w2#P) @
(wit P)-almost every (1, 02.
Proof. Let Z = C C. Since Q' and Q² have the second moment, the empirical distributions Ø and 2 converge to Q' and Q², respectively, in W2(X' x 2). Then, from Prop. 12 of Chemseddine et al. (9]. for each k there is subsequence (N&) such that IT; converges in W2 to an optimal plan II (t''', Q²) for W2, (Q', Q²) ks + This proves (i).
For (ii), it follows from Proposition A. that
[ (xxx2) [[ =1) (w22 22) dII (1 21, T2, = W2(Q (x,1z==)
(22)
Since 21 = 22 almost surely for II, by writing as an element in ( Q?), the left hand side of
J Xxxx2 121 x2l 2 x2,2) /
||
x2 (d)
XxX
From 20 and the fact / = (w2# P) & (#22# P), the right hand side
22 is equal to that of (21).
Since it holds generally that
|| 1|2 2 02 (f1 > WW (xx0====)xx((=o)))
(23)
Xxx
21
implies that the equality
23 must hold for (w2 @
every c1, C2-
A.3 Theoritical support on the choice of |
In this subsection, we derive necessary condition on the rate of 3 that the convergence of Theorem A.2 holds. Noting that the dimension on the condition space C C is 2d,, we see that the necessary condition below is 3 = ((V1/de). The choice 3 = (1/(2d,) used our experiments satisfies this condition, although it may not be the maximal of the necessary condition.
In the sequel, we discuss the general condition space C dimension d. Let X C RT and C c Rd and C' be random vector taking values in C with distribution Pc. Let (X, C), ~ P and (Y,C) Q be random vectors taking values in A' x C, where P, Q E P2(x' x C) have the same marginal distribution Pc, and X and Y are bounded: || |||| < M almost surely.
Suppose that we have i.i.d. samples (II, (In, with distribution and (yi, ci),..., (Mn, C.) with distribution Q such that and are conditionally independent given & for each i. Consider the following empirical conditional OT problem 20 28
(n) arg min Ex(o), Fa(o) == .EG
1 n
Σ
1 | No(o 112 + 8 72
||a c (24)
B
23

---

where B > 0 is constant. The optimal transport plan on (X' C) x (X x C) corresponding to
0. 0 (n) denoted by ) i.e., 0g (m)
E (( (o)co..
Proposition 12 in [8 shows that, for any increasing sequence 3r 00 (k E N), there is subsequence
(nn) converges w.r.t W2((X' x C) x (X x C)) to the OT plan
(n=) of such that, as + sc, Br
ao E (P, Q) for the following conditional OT problem:
/ || ||x- y/l
min (PPg))
Here, Q)) is the set 4-plans, defined by
2.4 (2(P,Q) = {a E((,)Q
# a == AxPc)
where A C C
Cx C, c (c,
We wish to prove 3g = ( as k - co. co. For this purpose, we make the following assumptions about the probabilities.
2/d
The probability P. has a density function f with respect to the Lebesgue measure such that there is constant By that satisfies
Assumption 1:
f(c) s Bg < 00
for a.e. e E C.
Assumption 2:
Ec I (lk-yyldPxiccell v1(ddxxc(()ddd11100) # Ec [Wa(Pxic Qr c)3].
The non-equality of Assumption 2 is not restrictive. fact, by the definition of W2, generally holds
Ec(W3 Qxic)l s Ec [/le- -Ill= ul c(x)dQric(v)
The assumption requires that, for almost all the conditionally independent ditribution Px|c & Qric does not give W2 (2((xi Qyic). In general, the optimal plan with the cost || 1ır-w1: y| is attained by independent distributions only if the supports of the distributions are orthogonal. In particular, if one of the distributions has density function with respect to the Lebesgue measure, this is not possible.
We have the following proposition.
Proposition A.3. Let Bk - oc. be an increasing sequence, and be subsequence of N such that the optimal plan a& for (24) converges the conditional OT plan ao w.r.t W2 x (X x C)). Under Assumptions and have
(nw)
Bx = 00 2/d
(25)
as h: oo.
Proof. We will derive a contradiction assuming that 25 does not hold. In that case, there is a subsequence (k') of N such that B -2/d sequence B% and ny satisfy
+ oo as k W.l.o.g.. assume that the original
BETk 2fd
(k xo).
(26)
Take 7 = (B 2/0)-1 which converges to as - co, and define the radius r%k > 0 by
2/d 1/2
0r 1/2 n
1/d
Tk == n4 k
(27)
for which
Burk ///2 - 1/d
nx
7/
1/2
(k: + 00)
(28)
holds. Also, we have n = 7
4/2
+ 0, which implies from Proposition A.4 that
PP
min
min ||e. c; < mx) +
(A: + 00).
331
24

---

This means that the probability of the event
F = [|c. eyll > T3 for all 1 si<j S m
tends to for Is + co. Hereafter, we consider a random evnet on Eg. For simplicity, us write of == 0
(me) and define
J& (1 sis no o1)()) # 1)
and my :- [Ja]- It follows that
B% T2 1=1
2 la - (0) 11x Thk ieJ.
Σ
lla - co))011
8k mik
Tk-
On the other hand, since of minimizes Prg, we have
1
1 7k
BR
ll - caz(1)1 S Fnz (ok) S Fny (id) =
I1 = Op()).
Thk
From the above two inequalities, we have Bx mark Op(1), and thus it holds from 28 that
mi = Op((3Arr) = (1) (k: + oo)-
(29)
T
Next, the first term of Fng (at) is
f |/
y| d/
(nx)
Thk Σ 1| yo
1
||=; yo2(o)
12k iEJu
ig
Σ
20
vill
Σ
1|,,
11112
Σ
|
Moz(s)11
(30)
17k
Th
Thk
lEJE
seJ
The second and third terms in the last line are upper bounded by
Σ 4M2
Thk
JEJx
which converges to zero in probability as ma/ns (1). As result, the first term in the last line converges to E - 2 with X LLLYC' C' in probability. Because the left hand side of (30) converges
f || v da,, implies
E| - 211 E. (PxI:
which contradicts Assumption
Proposition A.4. Let X Xn be i.i.d. valued random variables with density f on S C IRd.
Assume that
f(x) < Bu < xc for all x E S. For > 0, let == Vol(B(O,1)) be the volume of the unit ball in Rdd Then
min 1, (z) By va vare
P
min min |X X]|5rr
1sisn
Sn
3
Proof. The proof is standard and omit it.
B Experimental Details
All of the model training done using internal NVIDIA V100 GPU cluster.
25

---

Property
Configuration
Data space X'
Latent space [26] (R2122
Condition space C Training dataset
QED values embedded in 12 same as [25]
500K subset of ZINC22
Initial points z evaluation
800 reference molecules provided in 23
Initial conditions at evaluation
Embedded QED of T1
linspace (0.84, 0.95, 10) (maximum oracle calls 50000)
Target conditions 02 evaluation
[0.9, 0.91, 0.92, 0.93, 0.94 (otherwise)
Table 4: Training and evaluation configuration for the QED experiment
B.1 Synthetic datasets
In the synthetic experiments shown in Fig. 3
(a), we used a dataset sampled from a Gaussian mixture:
1 (M -1. -11.0..11))))))) (=]0==000-))]) + -i
6
+N
(31)
In the synthetic experiments shown in Fig. 3 (b), we used the dataset consisting of 10" samples and used a version of MLP with residual connection except that, instead of the layer of form x - (fxx), we used the layer that outputs the concatenation x + [ar, ofxx), incrementally increasing the intermediate dimensions until the final output. We used this architecture with width=128, depth=8 for all of the methods. For A2A-FM, we used 3 10, and we selected this parameter by searching 3 = 0.01, 0.1, 1, 10, 100 (see Appendix C for details). The batch size for all the methods were 1 x 10 In SI, we discretized the conditional space into 5 equally divided partitions. Also in partial diffusion, we used the classifier free guidance method with weight 0.3, used timesteps of T" 1000 and reversed the diffusion process for 300 steps. For non-grouped data, the model was trained on one NVIDIA V100 GPU. In both grouped and non-grouped data experiments, we calculated the metric Eg, P .c2 Il| 2, (z1)) TOT -cg (x1)))l empirically using 100 i.i.d. samples from the uniform distributions on the support of ci. C2. For error analysis, we ran this evaluation 10 times and reported the mean and standard deviations of the multiple runs in Table T
B.2 QED experiment
For the task shown in Section 5.2, we trained A2A-FM on 500K subsampled ZINC22 dataset, and adopted grid search algorithm similar (25) (see the pseudocode in Algorithm 2 to evaluate the methods to be compared. We used four NVIDIA V100 GPUs for training. As mentioned in the main manuscript, this search done for each initial molecule by making multiple attempts(Orach) Calls) of the transfer, and the initial molecule is marked success if the method of interest succeeds in discovering the molecule satisfying both the Tanimoto similarity and condition range requirement. The initial molecule is marked fail if the method fails to find such a transferred molecule within fixed number of oracle calls. We denote the set of initial molecules by M, in the pseudocode of 2 and we chose this set to be the same set used in (25). In the grid search approach, we adopted a boosting strategy, which is to scale the estimated velocity field v with a certain parameter b and calculate the ODE using Uboosted = b v. Using this strategy, we enabled our flow-based model to generate various molecules similar to the diffusion models with large guidance weights. For the grid search configurations (B, C, N) in Algorithm 2 we used B linspace(0.5s (0.8, 2.5, 20), N linspace(10 10 -3 3, 20), C = linspace(O..m (0.84, 0.95, 10) when MAX ORACLE CALLS = 50000. For MAX ORACLE CALLS {25, 100, 500, 1000}. we used B [1, 4,5), N [0]. C
[0.9, 0.91, 0.92, 0.93, 0.94]. We give samples of generated molecules in Fig. 10, We summarize the setting of the QED experiment in Table
In modeling the vector fields v, we used formulation inspired from Isobe et 21 order to reduce the computational cost given by
v(Ver,o2 (4), Me,, 02) = o(Va.o (1)c(())o2) - - (1), c())()))),
(32)
<(t) = 02 + * (1-t).
26

---

Algorithm 2 Grid search for QED conditioned molecular transfer
Input:
Set of initial molecules Mo {2, (QEN)}
Boosting weights B
Noise intensities N
Target conditions C
Trained velocity field v
Number of maximum oracle calls
Function for grid searching
1:
function GRID_ SRID_SEARCH(2o.CC co)
2: 1-1
3: while n & MAA__RACLE CRACCE_CALLES do
4: for b B do
5: for E C do
6:
for E sample ~ (0. 1), ån xn + 31 < ODESolver(t v(- (1) 71 n n+1
7:
8:
9:
10:
if i is not decodable then continue
11:
12:
end if
13:
if Tanimoto similarity(Fo. z1) 2 0.4 and QED(x1) 2 0.9 then
14:
Return: SUCCESS
15:
end if end for
16:
17:
end for
18:
end for
end while
20:
Return: FAIL
21: end function
# Loop for all initial points
mum success &-
23: for (Io, co) E Mo do
if GRID SEARCH(Z0, cy) SUCCESS then
25:
num sum__sccess &- mum_succoss + end if
27:
end for
28: Return: nur_ssuccess Mol
and used the UNet architecture proposed in Kaufman al. [25 in which replaced the convolution with dense layers. The network parameters were the same as the ones used in the QED nearby sampling benchmark of Kaufman al. [25]. This formulation ensures that transport will be con- ducted when 2 C1 and avoids learning trivial paths. In the training procedure, we first normalized the condition space with the empirical cumulative density functions that the empirical condition values would become uniformly distributed, and then embedded them using the TimeEmbedding layer to obtain its 32 dimensional representation; this is the same treatment done Kaufman al. (25). We used the latent representation provided in Kaufman et al. (25 and trained A2A-FM on a 500K subset of ZINC22 with batch 102c=1024, 3 = (batch (eatch_size) 1/2d, (1.2419) 1/2, where d. 32 is the dimension of the conditional space C. For this experiment, we found the choice of (batch_size)) 2/2e) to perform well, which is the reciprocal to the speed of the convergence of pair of empirical distributions [4]
B.3 Logp& TPSA experiment
In the evaluation of methods on LogP-TPSA benchmark (Section 5.2 we used the sampling efficiency curve plotted between (1) normalized discrepancy of LogP and TPSA between generated molecules and the target condition Cer [0, 1] and (2) Tanimoto similarity of Morgan fingerprints between
27

---

Property
Configuration
Data space X' Condition space C Training dataset Initial points æ1 evaluation Initial conditions at evaluation Target conditions 0g at evaluation Table 5: Training and evaluation configuration for the LogP&TPSA experiment
Latent space of [26] (P
LogP and TPSA values embedded in 12222
3.7M subset of ZINC22
Random samples from ZINC22
Embedded LogP and TPSA of I1
meshgrid 3, 4]. [10. 45, 80, 115, 150))
the initial molecule and the generated molecule STani E [0, 1]. Since the decoder of Kaufman et al. [25 does not always succeed in mapping the latent expression to valid molecule, we made 10 attempts for each target condition by adding different perturbations to the initial latent vector. Detailed procedure for calculating these two metrics for each initial molecule is in Algorithm [3] In our experiments, we used a 3. 3.7M subset from the ZINC22 dataset and created a validation split of initial molecules containing 1024 molecules. We used four NVIDIA V100 GPUs for training. For B in A2A-FM, we used [ = (1.25) 1/2. For Multimarginall SI, we created 10 clusters in the C space using the K K-means algorithm and treated the class labels as discrete conditions. For OT-CFM in this evaluation, we used the full training dataset as Po and trained a flow to conditional distributions using the regular OT-CFM framework for conditional generations. Other configurations, including batch size and model architecture, were common among competitive methods and were the same as the QED benchmark. We illustrate samples of generated molecules in Figures [ to [5 the training and evaluation configurations in Table
We summarize
The sampling efficiency curve in Fig. iss described mathematically as,
y = G(F-1(x))
(33)
F F(x) := Pr[STani S xz),
(34)
G(x) == Pr[c S x))
(35)
where Closs STani is the output of Algorithm 3 Pr[A < x²] is the ratio of members a E A that satisfies
2t, and Cerr == Closs/Cmas is the normalized Gloes by the normalization factor Guax = 7 so that Cper E [0, 1]. This way. the curve describes how much one has to trade off the tanimoto similarity threshold (similar to the size of the neighborhood) in order sample the molecules with Cery under certain error threshold.
C Additional experiments
The robustness of A2A-FM different 3: To further analyze the robustness of A2A- FM against the hyperparameter 3 in non-grouped settings, we trained our model with / = 0.1, 1,5, 5, 10, 20, 50, 100 and calculated the MSE from the ground-truth pairwise OT using the same procedure as Table (b). The result of this experiment is illustrated in Fig. 7 g From this experiment, we found out that the hyperparameter is least robust over range of order of magnitude (3 [1, 10). The selection of 3 used in the experiments (3 (batch_sizze 1/2ds) is within an order of magnitude of this range since we used batch size = 103, de = in the non-grouped synthetic data experiment.
Computational complexity of Multimarginal SI and A2A-FM: In the synthetic experimental setup depicted in Fig. 3 (b), one might claim that methods that are designed for grouped data, such as Multimarginal SI, could be applied to the setting of general (x, e) datasets by binning the dataset. However, this adaptation of methods like Multimarginal SI is challenging in practice. Firstly, as illustrated in Fig. [9] when a small number of bins are used, Multimarginal SI suffers from low transfer accuracy because the partitioning of C is just simply too coarse. On the other hand, increasing the number of bins not only leads to a substantial increase in computational cost compared to A2A-FM, but also results in diminished number of data samples per bin, consequently degrading its accuracy of OT estimation (see also (6) and Fig. g
CelebA-Dialoo HQ 256 Dataset: To demonstrate the applicability of A2A-FM to digh-oame grouped data, we trained A2A-FM on the 256 x 256 downscaled version of CelebA-Dialog HQ datasct [22] which contains 200K high-quality facial images with 6-level annotations of attributes:
28

---

Ours
Partial diffusion
SIK=))
0.06
0.05
MSE fram OT
0.04
0.03
0.02
0.00
10
Figure 7: Quantitative results from the additional experiment on the robustness hyperparameter The results of partial diffusion and multimariginal SI (SI) in the graph are taken from Table (b).
0
4
beta=0.01
beta=0.1
beta=0.5
beta=0.01
beta=0.1
beta=0.5
beta-1 1
beta=5
beta=10
beta=1
beta=5
beta=10
beta 20
beta=50
beta=100
beta 20
beta
beta 100
target points
generated points
initial points
Figure 8: Samples produced in the additional experiment on the robustness of hyperparameter 3.
Method
MSE from OT
Training time (min)
SI (K=2)
1.25 10-1
14.8
SI (K=5)
4.90 x 10-2
32.2
SI (K=10)
2.75 x 10-2
185.0
parital diffusion
6.77 x 10-2
18.1
A2A-FM (Ours)
1.51 x 10-2
29.1
Table Quantitative results from the additional experiment regarding the change in the number of bins K in Multimarginal SI. Although the performance measured by the MSE from OT improves as K increases, the training time will increase exchange. Results of partial diffusion and A2A-FM is also reported for comparison. The training of the shown models were done using the hardware with NVIDIA V100 GPU.
29

---

Algorithm 3 Grid search for LogP-TPSA conditioned all-to-all molecular transfer
Input:
Set of initial molecules Mo {(a,, (TPgP))CTSAA)
Noise intensity 0.1
Target conditions C meshgrid (0, 1, 4]. [10, 45, 80, 115, 150))
Trained velocity field
Transfer algorithm Trans(e,
Zi. Eo. co, co,c1) c1)
Condition calculation function Cond(-) X + C
Function for transferring with several attempts
1: function TRANSFER(2D. Co, C1)
2: с closs= < 00. .0......)). 0), similarity + [
3: for n E {1, 10} do
4: sample 2 (), 2g t- 20 E2 Trans(1, Ho, Co, 61)
5:
6: 21 not decodable or Sanimoto_similrityy(ee = then continue
7:
8: end
9: [oss[[]] MAE(c1, Cond(x1))
10: similarity/z + Tanimoto sitpootooooiiiittttyarr x1)
11: end for
12: Return: (min(c_loss) similarity/arrgmi(cc
13: end function
# Function to calculate the mean of all-to-all transfer
14: function GETMEAN(22,, Co)
15: c_ c_loss_meas mean <- 0, similarity_ment mean <-
16: for 6j C do
17:
c_loss', similarity' + TRANSFER(20, Co. C1)
c_loss_mean
c_loss_mean mean c_loss', similarity_mes similarity mean
similarity"
end for
20:
Return: c_loss mean/(C] similarity Siclarity_mean/CC
end function
# Loop for all initial points
Classy STani f- [0, .......)
23: for i E {1, .1..... [ol] do
Fo: Co Mo
25: Closs [i], STani + GETMEEEEMNn co)
26: end for
27: Return: Closs:
8Tani
Bangs, Eyeglasses, Beard, Smiling, Age. In this experiment we used only Beard and Smiling labels and treated these two labels as swo-dimensional condition in [0 5 We used the latent expressions of Rombach al. 39 and trained vector field using the UNet architecture with slight modification for multidimensional conditional inputs. This was done that we can adopt the independent condition embedding layers each dimensions the condition space and add them to the time embedding as the regular condition embeddings. We display examples of all-to-all
transfer results of A2A-FM in Figures I6
17
18
The results of this experiment support our claim that A2A-FM is scalable since the final dimentionality 12, 288. We shall note that many text-labeled image datasets
of the latent space was 64 64 are technically grouped databecause computer vision datasets oftentimes contain sufficient data for each independent textual condition, such as 'dog' and 'brown'. In such case, a classifier for each category may be trained, and condition transfer methods with classifiers like Zhao et al. [4),, Preechakul et al. [37] may become a viable option practice. We stress that A2A-FM also designed to be able to handle cases in which such good classffien/reegresooss cannot be trained. For example, on the aforementioned dataset like 7INC22 with continuous condition, difficult to train a high-pett pererraeform regressor/Classifier that can be used for classifier-dependent methods. For
30

---

K=2 (00:14:52)
K=5 (00:32:12)
K= 10(03:05:00)
4
K=2 (00:14:52)
K=5 (00:32:12)
K= 1000:0000000
π
-
initial points
target points
generated points
Figure 9: Samples from the additional experiment of multimarginal stochastic interpolants (SI) with different number of bins K. The red lines show the boundary of bins. The subtitle of each figure reads as [K # (bins) (trainingtime)) This experiment was done using the same hardware with NVIDIA V100 GPU for all models.
example, in Kaufman et al. (26), classifier guidance method for conditional generation does neither perform as well as classifier free guidance methods its flow-matching counterpart.
31

---

O
B
Initial molecules
qed: 0.77
qid: 0.78
ced: 0.77
qed: 0.78
Loto
Transfered molecules
0.31 by
simlarity
similariy
similaity
Initial molecules
qed: 0.70
qed: 0.79
oed: 0.78
qed: 0.78
qed: 0.73
Transfered molecules
ged. 0.91 similrity 0.41
cRd: 0.01 similantt 0.44
ged: 0.92 simliarty 0.40
qedi 0.91 similrity 0.42
oRd: 0.90 similaait 0.64
The first and third row the initial samples and the second and fourth row are the successful sampled molecules.
Figure 10: Samples of the nearby sampling in Section 5.2
Instaa Matecu TPSA Leg7-11.37
Thrget Logh 2.00
Target
Terget TPSA
Terget TPSA Log?
Terpet TPSAI Logh-4.03 -.00
Tsninot Ry40
fantinat latojo
AZA PA tourtt
Teeenste Senda by-2
The Santari
Tunmnote Samile k0
Tex moto San
Laghts
bretets Srrirrg-1.3
beemrtt Shre ameertpee
Logh- 1.22
Actast brvivees Lw
122
Acal
brerves Seriebee 2.1
Actuii 1944-177 Logeu 1.07
Actal Te56 Light Shilaroyyo
Actusi Testuce latimate 1a
Artal Teiteed 24 Birimots 2rriargeee..
Amal 195d- Lreh- 12t beirret
Actuni TDEA- logh- Accal TSSA- 185 lege-260
Actast TRSK-66 Logh-q.cc
Actal TRiA- Logh-3.07
Acsal TISA- Logi- 112
(D<CrarT13000
Figure 11: All-to-all transfer examples of experiment in Section 5.2
32

---

TPStwee Loghoe333
TPS 150 1,00
TPSA Log5
TPS Loghed
TPS Logh-3.06
TPSh 10 LegPed.05
Actast Tanmon 10
Tast Siibbrry ao
latinot lafoy-0
Actual 0.9
Acual 124 losp- Tenmrdo
Actast 72 2:46
Amall 3.34 "irtmtte 2nriurteee:17
Amal 196c Tantrrte ortierteee..
Tananote Samila Ayw253
Tarmnto latywo
Actuet PSSAES 30
Actast TISANDS Lo 33 Senilarr
Actest T7SA Log h0.33
Actall TISANES Loghta.s
Acuast 1ISA Logh=1.0
9 10)
Artast T95A-100 100 tigtunee et
Acbat TiSG-ds Lagite
Actaa TRAR- Looted of
Acual Teckebe Loottd
Areet TiSAk::: Lopt-2
1D 101
Attuat 175A4122 tughwo.s
Atual 7544224 ingf-115
1758 Loghtdus
Aqual TTSA Loghta.s
Adual TTSa Logto4
Actuii inge-russ Tanmnot Strila
TeMtun antrres orlgen
Tatmmdo simlaryyy
lattmate labyyo
Figure 12: All-to-all transfer examples experiment in Section 5.2
33

---

TPSAwd Log
TPSA 1,00
TPSA Log5
TPSAw Loghed
TPS Logh-3.00
TPSh LegPod.05
Actaal loght
Laghtd
Tanmon Simill Tanmmss Sinil
Artall
2,34
Actal TPSA Logh--22.1 22.7 Tenmoto Sentaroyeo
Actust Logh-4.22 Tarmntt
Amall 'irtmtt irriurtee..22
Amal lantrrtz atgeet.
Tanmnote Samilarly Samiardye2..r
Acball TSSAS Log
Actatt TISAkgg Logheq
Actust T7SANBT Loghegu4
Acual: TISARD Lightgat
Acuat TISA Lightdu2
9 10)
Artast 1954-1122 toge-.21
Atbat T956-99 Laghtd so
Actaa TeAk-1s Losted
Acall Texeeb Loot
Amal irsa-
10 100
Actuat 175A-112 LightLdt Actut TTSAed Lag
1754e Laghteuss
Aqual TTSA Loghtd.zz
Aqual: 1TSa-11 Loght2.6
Actuii TeMooo LnghoL Tarmmoo dentarry
Amal lntrres arteet
Tammntt Ry11
fattmate labyyoo
"
Figure 13: All-to-all transfer examples experiment in Section 5.2
34

---

TPStwee Log
TPSA 150 1,00
TPSA Log5
TP5 Loghed
TPS Logh=3.0
TPSh 10 Leghod.05
Actasl Tanmon Simlll Ro
Tast Sinlarr
Artull: Tananote Samila tywd
Actal Lagh-22 2. Tenmoto Sestariyeo
Actust 10 3.48 Tarentt latywo
Amall "iremate amriurgeee.3
3.13
Amal lantrrtz et
Acball TSA 40
Acbatt TISA Light Nesberyel
Actuat T7SA455 Log 49
Acual: TISAASS LightAAt
Acuat TISA Ligh-
9 10)
Acbal PRAA Lagte1
Anall TIA--21 inp-10L
Atbal THA Lostet1
Acaal TeLK-R2 Loote
Amet Tisa Lret-dit
10 100
Actuat TTSA-1D4 LigheLtt
Actur TTSke:: Log
175865s Laghedde 20 Aqual TTSA Loghtd.s
Actust T7Sams toghei.
Actuii TIM THA-- 1opk-121 Tanmmoo Aimlaryye
Amal lettrees
Tammntt Rp-014
lntinate Sre labyyoo
Figure 14: All-to-all transfer examples experiment in Section 5.2
35

---

TPS Log
TPSA 150 1,00
TPSA Log5
TPS Loghed
TPS Loghe3.06
TPSh 10 Leghed.05
Actaal Tanman Simill 13
Tast Sinllrry aettaroyo
Artull Logi-
Lagh-2.01 2.11 Tenmoto Sentaroyeo
Actual
Amall "irtmtte ammlutgeee.6
3.23 Amal 54 lantrrte et
Tananote Samilarly
Tarento 2a
Ackael Log DO
Actast TISA Log to
Actest T7SANIS Logiea.20
Actall TISAAES Lighto.d4
Acual: TISA-13 Logh-4.2
9 10)
Artast 1954-132 312 tnghoree
Accal TINA- 105 inst-not tot
Actual TeA Lopted
Actal Teskwet Looted
Amael TinAk-d Logh-
10 1001
Attuat 17SA-101 Lughwo.2
Acturt TTSA hap
1754 Loghtd
Aqual TTSAeet Loghtdat
Aqual: 1SSa-12 Loght
Artuni Tiikk Tanmmtt Serda
Tatimdo Rellaryyn
Figure 15: All-to-all transfer examples of experiment Section 5.2
36

---

Beard
original
Smiling
Figure 16: Transfer examples of CelebA-Dialog HQ dataset. A2A-FM is scalable to large scale all-to-all transfer tasks.
37

---

Beard
original
Smiling
Figure 17: Transfer examples of CelebA-Dialog HQ dataset. A2A-FM is scalable to large scale all-to-all transfer tasks.
38

---

Smiling with Beard 0
original
Beard with Smiling = 0
Figure 18: Additional transfer examples of CelebA-Dialog HQ dataset. Here, we fixed the Beard to changing Smiling attribute and Smiling 0 at changing the Beard attribute.
39


Reviews:
[{"review_id": "kmW8NqjYVW", "reviewer": "Reviewer_26Cp", "summary": "This paper seeks to simultaneously solve a collection of optimal transport problems defined between families of conditional distributions. The authors propose A2A-FM, a method based on flow matching, which uses couplings constructed with a particular cost function such that in theory one recovers the OT plans between the corresponding conditionals. The authors show the convergence of the empirical solution to the OT problem with this cost to the true pairwise OT solution. Experiments are conducted on some synthetic datasets and tasks based on molecular optimization.", "questions": "- I found the motivation for the proposed cost (Equation 9) very unclear. While the COT cost is somewhat intuitive (Eqn 3) since it enforces that the condition for a given datapoint is close to the condition of the datapoint it is paired to, this intuition is lost to me in Eqn. 9. It seems that this arises from taking the $i$th datapoint in the first batch $(x_1^i, c_1^i)$ and then appending the value of $c_2^i$, i.e., the conditioning variable for the $i$th datapoint in the second batch. \n   - Doesn't this mean that the coupling is sensitive to the ordering of the batches? \n   - Is it possible to write out the corresponding cost function in terms of $(x_1, c_1, x_2, c_2)$? Note that as written Equation 9 does not have this form as each term depends on the whole batch and so it is not a proper cost function i.e. depending only on two samples.", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths:\n- The paper works on an interesting and important problem that existing generative models often overlook\n- The synthetic experiments are a clear demonstration of the method and that A2A-FM can recover pairwise OT plans\n- The method seems to obtain strong performance on the molercular optimization experiments (although I am not an expert on this domain)\n\n## Weaknesses:\n- I found the motivation for Eqn 9 quite unclear (see questions section below)\n- The experimental settings could be more compelling/varied, particularly for the real-world experiments. While the authors do demonstrate some results on images in the appendix, this is only in the group-conditioning case, and doesn't really demonstrate the key appeal of the method (i.e., the ability to handle continuous conditioning variables). \n\n\n## Minor\n- There are some problems with references throughout.\n   - L178: Reference to [32] seems incorrect; I don't see where COT is used in this reference\n   -  The references in Section 2.2 are not correct. The authors cite [7] as developing the cost in Eq 3 but this is actually due to [CGS 2008] and also appears in the works [4] and [HHT 2023] which appeared before [7].  \n\n\n\n[CGS08]: From Knothe’s transport to Brenier’s map and a continuation method for optimal transport. Guillaume Carlier, Alfred Galichon, Filippo Santambrogio, 2008. \n\n[HHT 2023] Conditional optimal transport on function spaces. Bamdad Hosseini, Alexander W. Hsu, Amirhossein Taghvaei, 2023.", "was_revised": true, "final_justification": "I think this is paper is pretty solid in terms of its methodological contributions. The proposed technique is novel, interesting, and a non-trivial extension of existing techniques. The experiments are a clear demonstration of the fact that the methodology works as claimed. I still think that a more diverse set of experiments would strengthen the paper, but I don't think it is a necessary condition for publication of this work.", "rebuttal": "Thank you very much for constructive comments and suggestions, we would like to respond to each one of them below: \n\n\n> I found the motivation for the proposed cost (Equation 9) very unclear. While the COT cost is somewhat intuitive (Eqn 3) since it enforces that the condition for a given datapoint is close to the condition of the datapoint it is paired to, this intuition is lost to me in Eqn. 9. It seems that this arises from taking the $i$th datapoint in the first batch $(x_1^i, c_1^i)$ and then appending the value of $c_2^i$, i.e., the conditioning variable for the $i$th datapoint in the second batch.\n\nThank you for the insightful question. You are right to point out that the motivation for \"appending the value of $c_2^i$ \" is not immediately obvious when compared to the standard COT cost.  In addition to our explanation of intuition in 3.2, let us further clarify on this point in comparison with COT:\n\nIn standard COT (Eq. 3), the task is Conditional Generation: \"Generate an object with property c.\" The task is fully defined by a single condition c. Therefore, it's natural to enforce that the paired source and target points share this single condition c.\nIn our A2A-FM (Eq. 9), the task is All-to-All Transfer: \"Transform an object with property $c_1$ into an object with property $c_2$.\" To define this transformation, a single condition is not enough. We need a pair of conditions $z = (c_{\\rm src}, c_{\\rm targ})$, which corresponds to \"source condition\" and \"target condition\" respectively. \n\nGiven that our task is defined by a condition pair $z$, we use dataset where each data point is associated with such a pair. We start with two independent batches, $B_1 =\\{(x_1^{(i)},c_1^{(i)})\\} $ and $B_2 =\\{(x_2^{(i)},c_2^{(i)})\\}$.  For each $i$, this will give rise to pairs $(x_1^{(i)},c_1^{(i)})$, and this is to be coupled with $(x_2^{\\pi(i)},c_2^{\\pi(i)})$ for some permutation $\\pi$.  Here, (Eq. 9) enforces that the $\\pi$  does not greatly alter $(c_1^{(i)}, c_2^{(i)})$, so that  $x_1^{(i)} \\to x_2^{\\pi(i)}$ is approximately mapping from $P( \\cdot | c_1^{(i)})$ to $P( \\cdot | c_2^{(i)})$.  This way, the triplet $(x_1^{(i)},c_1^{(i)},c_2^{(i)})$ may be read as (source variable, source condition, target condition) where the target condition can be \"approximate\"--- $(x_1^{(i)},c_1^{(i)}) $ will be coupled to a variable whose condition is approximately $c_1^{(i)}$.   \n\n> Doesn't this mean that the coupling is sensitive to the ordering of the batches?\n\nAlthough we used the order of batches to randomly append conditions by assuming the batches to be randomly shuffled in the implementation, the core of this trick is that the randomly created tuple $(c_1^{(i)}, c_2^{(i)})$ is independently sampled. As mentioned in the Appendix (Equation (18)), the convergence guarantee relies on this independency. If for some reason the dataset cannot be shuffled globally, this independence can be maintained by shuffling within each minibatch before creating the condition pairs.\n\n>  Is it possible to write out the corresponding cost function in terms of $(x_1, c_1, x_2, c_2)$? Note that as written Equation 9 does not have this form as each term depends on the whole batch and so it is not a proper cost function i.e. depending only on two samples. \n\nYou are correct in your observation. Equation 9 is not a cost function between two individual samples, but rather the global objective function for the batch-wise optimal transport problem. However, this structure is not unique to our method. It is a standard formulation for any method that utilizes minibatch Optimal Transport to find a coupling, including OT-CFM and COT.\n\nMoreover, parallel to Monge formulation of OT problems as well as COT, it is possible to write down the cost function using joint distributions as\n\n$$ \\mathbb{E}\\left[ \\|X_1-X_2\\|^2 + \\beta( \\|C_1 - C'_1\\|^2 + \\|C_2 - C'_2\\|^2)\\right] $$\nwhere,  random variables $(X_1, (C_1, C_2), X_2, (C'_1, C'_2))$ follows a distribution $\\Pi_N \\in \\Gamma(\\hat{Q}^1_N, \\hat{Q}^2_N, \\hat{\\nu}_N)$ defined in the section A.2.2 in the Appendix.\n\n> The experimental settings could be more compelling/varied, particularly for the real-world experiments. While the authors do demonstrate some results on images in the appendix, this is only in the group-conditioning case, and doesn't really demonstrate the key appeal of the method (i.e., the ability to handle continuous conditioning variables).\n\nSorry for not being able to provide more real-world experiments in group-conditioning cases. We concentrated on scientific application scenarios where condition transfers in continuous domains are more important and well discussed. Moreover the purpose of the additional experiment on image generation in the appendix section is to prove the scalability of the method against high dimensional datasets and we believe that this purpose was achieved by this set of experiments. \n\n> L178: Reference to [32] seems incorrect; I don't see where COT is used in this reference\n\nThank you for kindly pointing out the mistake in citations. Reference to [32] should be:\n Manupriya, Piyushi, et al. \"Consistent optimal transport with empirical conditional measures.\" International Conference on Artificial Intelligence and Statistics. PMLR, 2024. We will fix this point in the camera-ready version.\n\n> The references in Section 2.2 are not correct. The authors cite [7] as developing the cost in Eq 3 but this is actually due to [CGS 2008] and also appears in the works [4] and [HHT 2023] which appeared before [7].\n\nThank you for mentioning the proper origins of Eq. 3. We will mention to the proposed articles in the revision."}, {"review_id": "jakQ8utWHT", "reviewer": "Reviewer_uaVd", "summary": "The paper proposes A2A-FM, a flow-based generative modeling framework for learning all-to-all condition transfer maps via pairwise optimal transport. Unlike existing methods that assume grouped datasets with many samples per condition, A2A-FM is designed to handle more general settings, including non-grouped data and continuous condition spaces where each condition may appear only once. The key contribution is a novel coupling cost function that jointly considers data and condition similarity, enabling the learned couplings to converge to the true pairwise optimal transports in the large-sample limit. Integrated within the flow matching framework, these couplings supervise vector fields that transport samples across condition pairs. Empirical evaluations on synthetic and molecular datasets demonstrate that A2A-FM achieves more accurate and efficient condition transfer than prior methods, particularly under the challenging non-grouped setting.", "questions": "Proposition 3.1 claims convergence to pairwise OT as \\( N \\to \\infty \\) and \\( \\varpi \\to \\infty \\). However, this is proven under the assumption of joint sampling from \\( P(x, c) \\).  \n- How does this convergence guarantee hold when the support of conditional distributions \\( P_c \\) is non-overlapping or degenerate (e.g., single-point support)?\n- More specifically, in the continuous setting, how does the method ensure absolute continuity of \\( P_{c_1} \\) and \\( P_{c_2} \\), which is a necessary condition for the existence of OT maps?\n\nThe coupling cost mixes sample and condition distances, controlled by \\( \\varpi \\), but the paper offers only a heuristic choice \\( \\varpi = N^{1/(2d_c)} \\).  \n- Can you derive this scaling from a principled bias-variance decomposition of the coupling error?  \n- Does this choice minimize some bound on generalization error? If not, what does it optimize?\n\nThe model approximates transport maps \\( T_{c_1 \\to c_2} \\) via training a vector field \\( v(x, t | c_1, c_2) \\).  \n- Do you assume Lipschitz continuity of the transport maps w.r.t. both \\( x \\) and \\( (c_1, c_2) \\)? If so, is this assumption realistic in your non-grouped datasets?", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "# Strengths\n\n1) A2A-FM is one of the first methods to handle non-grouped datasets, where each condition might have only one associated sample. This significantly broadens the applicability of conditional generative models to scientific domains with sparse or continuous condition spaces.\n2) The paper proves that the proposed coupling objective converges to pairwise optimal transport maps in the infinite sample limit, providing a solid mathematical foundation for the method.\n3) A2A-FM outperforms prior baselines such as OT-CFM, Multimarginal SI, and partial diffusion on both synthetic and real-world tasks, including challenging molecule optimization benchmarks, demonstrating both accuracy and sampling efficiency.\n\n# Weaknesses\n\n1) Although more efficient than fully multimarginal approaches, the minibatch optimal transport coupling still involves solving a linear assignment problem per iteration, which may become a bottleneck for large batches\n\n2) The success of the method heavily depends on the choice of the coupling regularization parameter. While a heuristic is provided, tuning remains nontrivial, especially for real-world datasets with unknown condition geometry.", "was_revised": true, "final_justification": "I thank the authors for their thoughtful and detailed rebuttal. After carefully considering their responses and discussions with the other reviewers and the area chair, I am increasing my score to 5. My reasoning is as follows:\n\nResolved Issues:\n\n* The authors provided satisfactory clarifications to all of my initial concerns, including minibatch OT couplings and regularizations, Lipschitz continuity and bias-variance decomposition.\n\nRemaining Concerns:\n\n* No major concerns remain after the rebuttal. Some minor clarifications could still be made in the final version for improved clarity.\n\nWeighting:\n\nSince all of my core questions were convincingly addressed, and the overall contribution is sound and relevant to the community, I have given significant weight to the improved clarity and completeness offered during the rebuttal.\n\nOverall, I believe this paper makes a meaningful contribution and merits acceptance.", "rebuttal": "Thank you very much for the constructive comments, we would like to resolve each one of them below: 　\n\n\n> Although more efficient than fully multimarginal approaches, the minibatch optimal transport coupling still involves solving a linear assignment problem per iteration, which may become a bottleneck for large batches.\n\nWe agree that this could be a bottleneck if the algorithm required very large minibatch to operate.  Interestingly, however,  even in our chemical case of non-grouped dataset of size 500K, we were able to achieve competitive results using a batch size of only 1024, for which the coupling computational cost is negligible.  Although the degree of computational burden may vary depending on the sparsity of the joint distribution, our method does not impose a greater computational cost than any other minibatch based FM algorithms.  We also note that it is possible to cache the optimal couplings before training process or use more efficient OT calculation algorithms (e.g. Sinkhorn) and accelerate the training procedure to mitigate potential bottlenecks in  large scale scenarios.\n\n\n> The success of the method heavily depends on the choice of the coupling regularization parameter. While a heuristic is provided, tuning remains nontrivial, especially for real-world datasets with unknown condition geometry.\n\nThank you for recognizing the heuristics.  It is true that, when the condition manifold is embedded in very high dimensions, it is difficult to obtain the true intrinsic dimension of the condition. However, when the conditions to control are explainable as variables in $R^{d_c}$ and if all elements in  $R^{d_c}$  are valid conditions, this will not be a problem in the real world settings since we can make use of the provided heuristics.\n\nThe difficulty may indeed arise if $P_c$ defined on  some Euclidian space $R^{d}$ which has a very low dimensional support, but such a situation also poses serious challenges in terms of out of distribution generalization for conditions that are possibly not observed.\n\nMoreover, provided that our method is robust against $\\beta$ up to an order of magnitude (see Fig. 7 in Appendix) and also that the value of $\\beta = N^{1/(2d_c)}$ does not largely change against $d_c$, we think that tuning this parameter is not so difficult in such cases with unknown geometries.\n\n> How does this convergence guarantee hold when the support of conditional distributions ( P_c ) is non-overlapping or degenerate (e.g., single-point support)?\nMore specifically, in the continuous setting, how does the method ensure absolute continuity of ( P_{c_1} ) and ( P_{c_2} ), which is a necessary condition for the existence of OT maps?\n\nIndeed, as is required in the well known Benamou-Brenier formula[Benamou 2000], the absolute continuity of $P(\\cdot | c_k)$ with respect to the common Lebesgue measure is necessary for the existence of the vector field corresponging to the OT maps.\nHowever, the results regarding our convergence proof are based on the Wasserstein Distance and OT plan in the form of Kantorovich formulation, for which such assumptions are not required.\nAlso, we note that we do not directly use the population-level OT maps in our method. Rather, we construct optimal couplings (i.e., OT plans) between conditional distributions just for the construction of Flow Matching objective. Therefore, the absolute continuity with respect to the Lebesgue measure, which is required for defining continuous OT maps or vector fields in the Benamou–Brenier formulation, is not necessary in our case.\n\n\n[Benamou2000] Benamou, Jean-David, and Yann Brenier. \"A computational fluid mechanics solution to the Monge-Kantorovich mass transfer problem.\" Numerische Mathematik 84.3 (2000): 375-393.\n\n> Can you derive this scaling from a principled bias-variance decomposition of the coupling error?\n> Does this choice minimize some bound on generalization error? If not, what does it optimize?\n\nThank you very much for the suggestion.  While this is not completely within the scope of this work,  it is an interesting direction for future studies.  Using the notation of the Appendix A.2.2,  we assume the bias-variance decomposition you mentioned can be written as\n\n$$\\mathbb{E}[(W^2_{2,\\beta}(\\hat{Q}^1, \\hat{Q}^2)-W^2_{2,\\nu}(Q^1, Q^2))^2] = (\\mathbb{E}[W^2_{2,\\beta}(\\hat{Q}^1,\\hat{Q}^2)]-W^2_{2,\\nu}(Q^1, Q^2))^2 + \\operatorname{Var}(W^2_{2,\\beta}(\\hat{Q}^1,\\hat{Q}^2))$$\n\nwhere, $Q^a$ denotes the augmented data distributions (eq.(17)) and $\\hat{Q}^a$ denotes the empirical distribution of $Q^a$. $W_{2,\\beta}$ denotes the 2-Wasserstein distance with the $\\beta$ -weighted cost function (eq. (16)) and $W_{2,\\nu}$ stands for conditional 2-Wasserstein distance (eq. (12), (13)). From Theorem A.2, the left hand side converges to $0$ with an increasing sequence of $\\beta_k, N_k$. However, obtaining convergence rates of the two terms in the LHS is non trivial and needs further discussion. Possibly, some kind of justification according to the convergence rates of empirical Wasserstein distances [Fournier2023] may be possible. We would be happy if you can provide any ideas in this direction.\n\nHowever, for theoretical understanding behind this heuristics, we can show that this selection of $\\beta$ meets the necessary condition for the convergence in Prop. 3.1. See our reply to T9Aq. Therefore, although it needs further discussion for finding an optimal value of $\\beta$, we can say that our bound has some justification based on necessary condition. \n\n[Fournier2023] Fournier, Nicolas. \"Convergence of the empirical measure in expected wasserstein distance: non-asymptotic explicit bounds in ℝd.\" ESAIM: Probability and Statistics 27 (2023): 749-775.\n\n\n\n>  Do you assume Lipschitz continuity of the transport maps w.r.t. both ( x ) and ( (c_1, c_2) )? If so, is this assumption realistic in your non-grouped datasets?\n\nWe do not assume Lipshitz continuity of the transport maps with respect to $( x )$ and $( (c_1, c_2) )$ because our theory does not impose explicit assumption on the model with which to approximate the family of pairwise OTs. Indeed, the existence of OT for each pair $(c_1, c_2)$ does not require  such an assumption.  At the same time, Lipschitz continuity is related to the uniqueness of global solution, as mentioned in Prop 8 of Chemseddine et al."}, {"review_id": "9jL7WmQ7kB", "reviewer": "Reviewer_CVUQ", "summary": "The paper proposes a novel cost function which allows for learning optimal transport (OT) maps between multiple pairs of conditional distributions simultaneously. Then the obtained OT maps are used to couple data points in conditional flow matching analogously to OT-CFM paper (Tong et al., 2023). The resulting model (A2A-FM) can be used to learn translation between conditional distributions with continuous conditions. The authors provide theoretical justifications of the model’s performance and test it on toy examples and on the chemical datasets.", "questions": "- Could you please elaborate on the principle of generalized geodesics given in lines 169-170 and its connection to the method of multimarginal Stochastic Interpolants (Albergo et al., 2024? \n- Why section 5.1 misses comparison with (Kapusnyak et al., 2024)?\n- Why is the hyperparameter $\\beta$ not robust to values bigger than $10$ in experiment from section 5.1 in non-grouped settings?\n\n**References.**\n\nTong, Alexander, et al. \"Improving and generalizing flow-based generative models with minibatch optimal transport.\" Transactions on Machine Learning Research, 2023.\n\nK. Kapusniak, P. Potaptchik, T. Reu, L. Zhang, A. Tong, M. M. Bronstein, J. Bose, and F. Di Giovanni. Metric flow matching for smooth interpolations on the data manifold. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, Nov. 2024.\n\nB. Kaufman, E. C. Williams, C. Underkoffler, R. Pederson, N. Mardirossian, I. Watson, and J. Parkhill. Coati: Multimodal contrastive pretraining for representing and traversing chemical space. Journal of Chemical Information and Modeling, 64(4):1145–1157, 2024.\n\nL. Ambrosio, N. Gigli, and G. Savare. Gradient flows: In metric spaces and in the space of probability measures. Lectures in Mathematics. ETH Zürich. Birkhauser Verlag AG, Basel, Switzerland, 2 edition, Dec. 2008.\n\nM. S. Albergo, N. M. Boffi, M. Lindsey, and E. Vanden-Eijnden. Multimarginal generative modeling with stochastic interpolants. In The Twelfth International Conference on Learning Representations, 2024.", "limitations": "The authors have addressed the limitations of their approach.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths.** The proposed type of cost function (which is used to find the pairwise OT mappings for the family of conditional distributions) allows for dealing with datasets with any type of conditions including continuous ones. The method is supported by the theoretical result showing that it provides good approximations of the ground-truth OT mappings between any pair from the chosen family of conditional distributions (in the limit w.r.t. to the hyper-parameter $\\beta$ and number of training samples $N_k$). The method provides good results in the real-world chemical experiment.\n\n**Weaknesses.** Some parts of the paper are hard to understand. ‘Preliminary’ section is missing the overview of definitions from optimal transport theory. This seems to be necessary since later parts of the paper exploit the definitions of OT mappings and OT problem in Kantorovich formulation without defining them previously (lines 110-112). Figure 1 uses captions referring to the methods (Multimarginal SI/EFM) which are not yet defined. The principle of generalized geodesic (Ambrosio et al., 2008) and method of partial diffusion (Kaufman et al., 2024) which are included in comparison on synthetic data (section 5.1, Table 1a,b) are not well explained also.\n\nI am also wondering why section 5.1 misses comparison with (Kapusnyak et al., 2024) included in 'Related work' section. Comparison with this method is needed to understand the position of the proposed approach w.r.t. other methods. The computational inefficiency of (Kapusnyak et al., 2024) stated in line 166 seems to be not a crucial issue for toy experiments.\n\nMy other concern is related to the provided theoretical result (Proposition 3.1), specifically, its conditions.  According to this result, the proposed approach allows for learning good approximations of all-to-all OT mappings (between any pair of conditional distributions) in the limit w.r.t. to the number of training samples and hyper-parameter $\\beta$, i.e., when $\\beta$ is sufficiently large. At the same time, according to the provided experimental details in Appendix C, the authors showed that the parameter $\\beta$ is robust up to value $10$ in experiments of section 5.1 in non-grouped settings. It is not evident for me why the hyperparameter is not robust to bigger values, i.e., why the method fails to correctly align the distributions in such cases?\n\n**Overall**, I think that the paper provides interesting cost functions allowing for finding the coupling between multiple pairs of conditional distributions at once. However, the flow-based model by its own is not that novel, it represents an instantiation of the well-known OT-CFM model (Tong et al., 2023) with just another way of finding the coupling. From the practical side,  the approach gives good results in toy experiments as well as the one from the chemistry field, but contains visible artifacts in the experiment with faces in Appendix, see Figure 16. Since I am not an expert in chemistry, it raises questions regarding the overall significance of the proposed approach.\n\n*Typos:*\n- line 28 - **in** chemistry\n- line 45-46 - **there are** no methods\n- line 65 - missing point\n- before line 145 - double **the**", "was_revised": true, "final_justification": "The authors have addressed most of my concerns during the rebuttal phase, thus, I changed my score to positive.\n\nI did not assign higher score since I remain slightly concerned by the experiments - I can not judge the experiment on chemical datasets as I am not an expert in this field, while other experiments are either sunthetic or given only to show the scalability of your approach, i.e. do not demostrate its superior performance.  I remain slightly concerned by the experiments - I can not judge the experiment on chemical datasets as I am not an expert in this field, while other experiments are either sunthetic or given only to show the scalability of your approach, i.e. do not demostrate its superior performance.\n\nOverall, I took into account that other reviewers find the chemical experiment meaningful, and assign slightly positive score, i.e., 4.", "rebuttal": "> ‘Preliminary’ section is missing the overview of definitions from optimal transport theory.\n\nWe are sorry that we could not include the introduction to optimal transport theory in the main manuscript of the paper due to page limitation.  \nIn the revision we will explain the Kantorovich formulation and provide references. \n\n> Figure 1 uses captions referring to the methods (Multimarginal SI/EFM) which are not yet defined. \n\nWe are sorry for the confusion, we will provide citations in the captions and also provide reference pointers to the paragraphs explaining these methods.  \n\n> The principle of generalized geodesic (Ambrosio et al., 2008) and method of partial diffusion (Kaufman et al., 2024) which are included in comparison on synthetic data (section 5.1, Table 1a,b) are not well explained also.\n\nWe are sorry for the lack of explanation. We will explain each one of them below\n\n- Partial diffusion \n\nPartial Diffusion is a method that is presented in [22],  and it operates the condition-transfer from $c$ to $c'$ by (1) adding a limited noise to the embedded representation of the initial molecule of condition $c$ then (2) denoise the perturbed representation with the classifier free guidance for class $c'$. We conducted partial diffusion with different noise steps, please see Figure 5 as well as [22]. In the chemistry experiments, we used the same architectural and model settings used in [22], please see Appendix C of [22] for details.  \n\n- Multimarginal Stochastic interpolants and Generalized Geodesics \n\n As described in section 4, Generalized Geodesics [5], which is the base of Multimarginal stochastic interpolants (SI) [2], is a framework of interpolating the optimal transport. Given the dataset of $P_{c_i}$ and the source distribution $P_\\emptyset$, it first learns $T_{\\emptyset \\to c_i}$ empirically with grouped dataset, and generates $P_c$ by first writing $c = \\sum \\alpha_i c_i$  and using the transport $\\sum \\alpha_i T_{\\emptyset \\to c_i}$.  SI [2] further learns the flow for this tranport, optimizes the interpolation path by optimizing the path in the $C$ space by incrementally transforming $c$ with this interpolation. We used the same architectural design as our method to design the vector flow of [2]. \n\n> I am also wondering why section 5.1 misses comparison with (Kapusnyak et al., 2024) included in 'Related work' section.\n\nWe are sorry for the confusion. Although Kapusnyak et al. (2024) is correctly cited as Flow matching research that incorporates data manifold their work focuses on unconditional smooth interpolations over data manifold, and does not involve conditioning on external variables or modeling conditional distributions. In contrast, our problem setting is specifically aimed at all-to-all transport between conditional distributions—that is, between  for arbitrary pair of conditions. We acknowledge that our inclusion of this work in the Related Work section may have caused confusion and we will revise that section to more clearly distinguish between unconditional and conditional transport settings.\n\nWe did not include the method of Kapusnyak et al for comparison in the submission, since the method designed for unconditional interpolation do not naturally generalize to the current setting. However, it may be possible to adopt Kapusnyak et al. (2024) in conditional generation settings by extending the neural network to accept an additional input for conditions. We experimented this method with this strategy on the LogP-TPSA benchmark mentioned in our paper and obtained:\n```\nAUC value: 0.179\n```\nwhich is a low value compared to other similar methods like OT-CFM. We would like to report that this may be partially due to the instability of the geometry training procedure of Kapusnyak et al., (2024). The time derivative calculation in geodesic path learning led to unstable learning when adopted to high dimensional data like the molecular latents. We need more explorations in the future to adapt the method to the chemical generation scenarios.\n\n> However, the flow-based model by its own is not that novel, it represents an instantiation of the well-known OT-CFM model (Tong et al., 2023) with just another way of finding the coupling. \n\nWe would like to respectfully emphasize that our method addresses a substantially different and more challenging task than prior OT-CFM-based approaches. Specifically, as noted by reviewer **uaVd**, *“A2A-FM is one of the first methods to handle non-grouped datasets, where each condition might have only one associated sample. This significantly broadens the applicability of conditional generative models to scientific domains with sparse or continuous condition spaces.”* Unlike previous methods, which often rely on grouped data on finite conditions, our model is designed to operate effectively under non-grouped data, making it applicable to a wider range of real-world scientific settings.\n\nFurthermore, while our method builds upon the OT-CFM framework, many recent advances in flow matching also derive from OT-CFM through alternative coupling mechanisms. For example, [24], published last year, is algorithmically equivalent to OT-CFM, differing only in that it employs a parameterized coupling between joint distributions. Similarly, Riemannian Flow Matching [Chen2023], another widely cited work, modifies CFM primarily by replacing linear interpolation with geodesics.  Although these works, including ours, share a common foundation, developing algorithms and theory capable of accurately learning the target vector field. In our case, we propose a principled and novel approach to simultaneously learn pairwise optimal transport, supported by a rigorous theoretical framework.\n\nWe respectfully ask the reviewer to reconsider the novelty and significance of our contribution in terms of the broader scope of applicability and the difficulty of the problem our method addresses.\n\n[Chen2023] Chen, Ricky TQ, and Yaron Lipman. \"Flow Matching on General Geometries.\" The Twelfth International Conference on Learning Representations.\n\n>  From the practical side, the approach gives good results in toy experiments as well as the one from the chemistry field, but contains visible artifacts in the experiment with faces in Appendix, see Figure 16. Since I am not an expert in chemistry, it raises questions regarding the overall significance of the proposed approach.\n\nThank you for this valuable feedback. We understand your concern that the artifacts in the face experiment could cast doubt on the overall significance of our approach, and we would like to clarify our reasoning. The experiment in the Appendix served a specific and distinct purpose: to verify that our method is scalable to high-dimensional data, a common hurdle in many scientific fields. The model's ability to process this data and produce structured—though not visually perfect—outputs confirms its scalability. The artifacts themselves are a direct consequence of not using any vision-specific architectures, which lies outside the scope of our contribution. Our core results are in the chemistry domain, where the method yields compelling outcomes. Therefore, we respectfully argue that the scalability is well support by the experiment. \n\n>  Could you please elaborate on the principle of generalized geodesics given in lines 169-170 and its connection to the method of multimarginal Stochastic Interpolants (Albergo et al., 2024)?\n\nPlease see our comment above. \n\n> Why section 5.1 misses comparison with (Kapusnyak et al., 2024)?\n\nPlease see our comment above. \n\n> Why is the hyperparameter $\\beta$ not robust to values bigger than $10$ in experiment from section 5.1 in non-grouped settings? It is not evident for me why the hyperparameter is not robust to bigger values, i.e., why the method fails to correctly align the distributions in such cases?\n\nGiven *a fixed number of samples $N$*,  indefinitely making the $\\beta$ value large in (9) would make the model ignore the $x$ terms and force the choice of $\\pi$ for which $c_k^{i} = c_k^{\\pi(i)}$ for $k=1,2$.  In the case of *unique $x$ sample per each $c$*, this can cause the permutation $\\pi$ to exclusively choose the identity map. As a result, the model would learn the couplings with just a single pair of samples and this can lead to larger errors in OT approximations. When the $\\beta$ is of appropriate size with respect to $N$, the learning based on equation (9) would keep the balance of minimizing the transport cost between the set of samples with approximately $c_1$ condition and the set of samples with approximately $c_2$ condition.\n\nPlease also see our reply to T9Aq where we derived an inequality that shows for $\\beta \\gg N^{1/d_c}$, Prop. 3.1 will not converge. Therefore, we can also theoretically state that $\\beta \\ll N^{1/d_c}$ is necessary."}, {"review_id": "sl1u4DBDdY", "reviewer": "Reviewer_T9Aq", "summary": "This paper introduces A2A-FM, a flow-based method for learning all-to-all transfer maps between continuous conditional distributions via pairwise optimal transport. The authors propose a novel coupling cost function that integrates data and condition terms, enabling simultaneous learning of OT maps for all condition pairs using minibatch OT-coupled flow matching.", "questions": "1.\tWhile the heuristic $\\beta$ demonstrates empirical effectiveness, it lacks a theoretical justification. It would be beneficial to understand the sensitivity of performance to $\\beta$'s specification, particularly in high-dimensional settings (large d). Are there theoretical bounds that can guide the selection of $\\beta$?\n\n2.\tIn scenarios with highly sparse data (e.g., a large number of unique c values), under what conditions does A2A-FM's approximation of pairwise Optimal Transport (OT) break down? Are there theoretical or empirical guidelines for the minimum number of samples required per condition to ensure reliable approximation?\"\n\n3.\tPlease clarify the application of minibatches to the molecular dataset. What is the strategy for creating minibatches from molecular data?\n\n4.\tThe discussion of antisymmetry lacks depth. The phrase 'suggesting some justification for this regularization' is unclear. Please provide a more detailed explanation of the rationale behind this regularization and the evidence supporting its effectiveness.", "limitations": "As the author said, the balance of dataset size and the hyperparameter $\\beta$ can pose a limitation.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 1, "significance": 3, "originality": 3, "strengths_and_weaknesses": "_Strengths_:\n\n1.\tTackles the challenge of non-grouped data with continuous conditions (e.g., physical properties in molecules), where existing methods (e.g., Multimarginal SI) fail.\n\n2.\tAuthors provide convergence analysis proving the approximation of pairwise OT maps..\n\n_Weakness_:\n\n1.\tUnclear molecular implementation. Minibatch OT requires fixed-dimensional data, but molecular structures vary in atom count. The paper does not clarify how this dimensional inconsistency was resolved.\n\n2.\tWhile antisymmetry regularization boosts performance, as the author mentioned in Sec. 6, its mechanism and impact lack theoretical or empirical exploration.\n\n3.\tIn addition to the synthetic data, the manuscript's drug design experiment requires further development. It lacks essential elements, such as a clear introduction to baseline models and a definition of the evaluation metrics used (e.g., AUC). Furthermore, the analysis appears overclaimed, as key properties like $\\Delta\\Delta G$ and vertical ionization potential were not experimentally tested. The performance of the model on these properties, or others requiring computationally expensive calculations, remains unknown.", "was_revised": true, "final_justification": "My concerns about the minibatch, antisymmetry, and scenarios on sparse data are addressed. My remaining concern is the clarity of the paper. I would lean to accept a little.", "rebuttal": "Thank you very much for the comments and suggestions, we would like to respond to your concerns below.\n\n> Unclear molecular implementation... the paper does not clarify how this dimensional inconsistency was resolved.\n\nWe are sorry we were not clear enough.  As we mention in section 5.2, we conducted our algorithm on the latent space used in Kaufmann et al, which provides a Molecular autoencoder based on [23]. The explicit construct of the encoder is given in Section 3.1 and Appendix B of [22], which encodes the SMILES representation of each molecule to 512 dimensional real valued vector.  \nThe minibatch therefore consists of vectors of dimension 512.  \nWe used the pretrained encoder provided by [22], which was achieving the SOTA result prior to our research. \n\n> While antisymmetry regularization boosts performance, as the author mentioned in Sec. 6, its mechanism and impact lack theoretical or empirical exploration.\n\nWhile antisymmetry regularization is not the main focus of this work,  this mechanism is inspired from the solutions of Extended Flow Matching and Stochastic interpolant [2, 18] which both state that , in the general formulation of multimarginal flow, the flow ODE from condition $c_1$ to condition $c_2$ takes the form of \n$$\\frac{d}{dt} x_{c_1 \\to c_2}(t)  = M(x(t), c(t)) [ \\frac{d}{dt}  c(t)  ]    $$   \nwhere $c(0) = c_1,  c(1) = c_2$,  $M$ is a matrix of dimension $d_x \\times d_c$.  The theoretical work [18] in particular shows that Dirichlet energy, which is the generalization of the transport cost, is minimized by the flow above with $c(t) = (1-t) c_1 + t c_2$, in which case the flow becomes $M(x(t), c(t))(c_2 - c_1)$.   \nWhile [2,18]  were not able to provide a scalable algorithm in the case of non-grouped data, these two works provide extensive theories behind this flow design.  We will elaborate this point in the revision. \n\n>  ... experiment.. lacks essential elements, such as a clear introduction to baseline models and a definition of the evaluation metrics used (e.g., AUC). \n\nWe are sorry that we did not clearly explain them in the main part of the manuscript. To reiterate, our baselines include partial diffusion, stochastic interpolant, and OT-CFM.   We will make sure to include them in the revision. \n\n- Partial diffusion \n\nPartial Diffusion is presented in [22] as well, and it operates the condition transfer from $c$ to $c'$ by  adding a limited noise to the embedded representation of the initial molecule of condition $c$ and apply the classifier guidance with class $c'$ for denoising. We conducted experiments with different noise steps, please see Figure 5.  We used the same architectural and model settings used in [22], please see Appendix C of [22].  \n\n- Multi-marginal stochastic interpolant \n\nMulti-marginal stochastic interpolant (SI)[2] is, as described in the related works section, is a framework of constructing a multi-marginal optimal transport based on the interpolation of optimal transport called Generalized Geodesics [5]. As described in section 4,  given the dataset of $P_{c_i}$ and the source distribution $P_\\emptyset$, Generalized Geodesics first learns $T_{\\emptyset \\to c_i}$ empirically with grouped dataset, and generates $P_c$ by first writing $c = \\sum \\alpha_i c_i$  and using the transport $\\sum \\alpha_i T_{\\emptyset \\to c_i}$. SI  approximates this with the flow and  further optimizes the interpolation path by optimizing the path in the $\\mathcal{C}$ space by incrementally transforming $c$ with this interpolation. We used the same architectural design as our method for the flow. \n\n- OT-CFM \n\nThis is a classic method of [38]. As described in section, we conduct conditional transfer with this method by treating the unconditional distribution as the source distribution. \n\n\n- AUC of figure 5. \n\nWe provide this information in the Appendix B.3.  The curves of Fig 5 are the plots between \"(1) normalized discrepancy of LogP and TPSA between generated molecules and the target condition cerr ↔[0,1] and (2) Tanimoto similarity of Morgan fingerprints between\nthe initial molecule and the generated molecule sTani ↔[0,1].\"  Because it is more generally difficult to alter just the desired property while maintaining the similarity to the original molecule, (1) and (2) are in tradeoff relation. To compare the wellness of this tradoff relation, we used \"Area Under the Curve (AUC) of the plot of (1) against (2).  We will include this in the main manuscript as well. \n\n\n>  the analysis appears overclaimed, as key properties like $\\Delta\\Delta G$ and vertical ionization potential were not experimentally tested. \n\nWe appreciate reviewer's  mention of these properties. At the same time, *validating* these quantities requires expensive quantum chemistry simulations such as DFT or FEP, whose computational cost on the benchmark dataset of considered size can be infeasbile in academic settings.  For this reason, even many prior *pure chemistry* works (e.g., COATI [23], MolMIM [35]) evaluate on low-cost surrogates like QED or LogP, and we followed the same practice to ensure comparability.  As the general machine learning research, we are providing the results comparable to the state of the art *chemical academic research*. \nIn the revision, we make sure that our chemistry experiments are not stand-alone *industrial level* results. \n\n\n> heuristic $\\beta$ demonstrates empirical effectiveness, it lacks a theoretical justification.\n\nFirst we note that, theoretically providing optimal values for $\\beta$ is a difficult task. For example, although the application and the coupling cost function differs, Kerrigan et al [24] that was accepted last year also uses the similar parameter, and do not provide an effective heuristics nor a theoretical intuition on how to decide the parameter. We believe that we are the first one to provide this type of heuristics based on the convergence speed of Wasserstein convergence rate and this heuristics itself should be regarded as a contribution.\n\nAlthough creating precise bounds to obtain optimal values for $\\beta$ requires discussions on the convergence rates in Prop.3.1, which is far beyond the scope of this paper, it is possible to observe that the provided heuristics $\\beta = N^{1/{2d_c}}$ will fulfill necessary conditions for the convergence. In Prop. 3.1, we stated that the cost function:\n\n$$C(N, \\beta) = \\min_{\\pi \\in S_N} \\Delta_{\\pi}(X) + \\beta\\Delta_{\\pi} (C)$$\n\nwhere,\n$$\\Delta_{\\pi} (X) = \\frac{1}{N} \\sum_{i=1}^N \\left[\\|x_1^{(i)} - x_2^{(\\pi(i))}\\|^2\\right]$$\n$$\\Delta_{\\pi}(C) =  \\frac{1}{N} \\sum_{i=1}^N \\left[ \\|c_1^{(i)} - c_1^{(\\pi(i))}\\|^2 + \\|c_2^{(i)} - c_2^{(\\pi(i))}\\|^2\\right]$$\n$$ S_N: \\text{symmetric group of [1:N]}$$\n\nconverges to a finite value under a subsequence $\\beta_k, N_k$. On the other hand, from an argument similar to [Yukich2006, section 4.3], with high probability there exists a positive value $K$ such that $ KN^{-1/d_c}  \\leq \\min_{\\pi \\in S_N, \\pi \\neq \\operatorname{id}} \\|c_1^{(i)}-c_1^{(\\pi(i))}\\|^2 +\\|c_2^{(i)}-c_2^{(\\pi(i))}\\|^2  $ holds, as this order corresponds to the square of probabilistically high of minimum distance between the nearest neighbors when $N$ instances of $(c_1, c_2)$ are sampled from $R^{d_c + d_c}$.\n Since the minimizer $\\pi$ of $C(N,\\beta)$ is not identity maps in regular situations of continuously valued $c$s in the presence of  $\\Delta_{\\pi} (X)$, this bound is likely to hold for the minimizer of $C(N, \\beta)$. At the same time,  when we substitute $\\beta = N^{1/{2d_c}}$, we obtain\n\n$$ C(N, \\beta) \\geq \\min_{\\pi \\in S_N} \\Delta_{\\pi} (X) + K \\underbrace{\\beta N^{-1/d_c}}_{N^{-1/(2d_c)}}.$$\n\nTherefore, by choosing this order of $\\beta$, we can not only prevent the cost from exploding in the limit of $N\\to \\infty$, \nbut also require $\\Delta_{\\pi}(C)$ part to approach to $0$ without literally realizing $\\pi = id$, thereby facilitating $\\Delta_{\\pi} (X)$ to learn the coupling between $P(\\cdot | c_1)$ and $P(\\cdot | c_2)$ for each instance of $(c_1, c_2)$.  \nAlso,if  $\\beta \\gg N^{1/{d_c}}$, the cost itself may not converge in the limit of $N\\to \\infty$.\nThe investigation of $\\beta$ is an interesting future research question, and  can also help advance the application of Kerrigan et al [24] for which the heuristic has not been established. \n\n\n[Yukich2006] Yukich, Joseph E. Probability theory of classical Euclidean optimization problems. Springer, 2006.\n>  In scenarios with highly sparse data (e.g., a large number of unique c values),  under what conditions does A2A-FM's approximation of pairwise Optimal Transport (OT) break down?\n\nWe first note that our method is compatible with dataset in which all c values to appear are unique.  However, just as is the case for all generative models, approximation would break down when the joint data is sparse with respect to the distribution. \nThat is, when the set observations of $c$ does not approximate $p_c$, then the model may not capture the continuity pf $p(\\cdot | c)$ with respect to $c$.  \n\n>  What is the strategy for creating minibatches from molecular data?\n\nAfter encoding the molecules in the SMILES representation into 512 dimensional vectors with the encoder mentioned above, we sampled the minibatches from randomly shuffled samples, using the standard PyTorch implementation.  \n\n>  The discussion of antisymmetry lacks depth.\n\nPlease see our comment above."}]

Author rebuttal:

Meta-review: The paper extends conditional OT works by considering the case where the conditioned variable could be different in the source and target. Also, unlike most works, grouped data is not assumed. The paper proposes a flow based objective and methodology to solve the corresponding problem. Both theoretically and simulations-wise, the correctness of the methodology is proven. Simulations on a molecular opt benchamrk show that the method may potentially scale to real-world applications.

Overall, the reviews are postive and there is an agreement that the problem is new and the proposed methodology is interesting. The paper seems to reduce a gap in literature in the sense that the previous methods either focus on conditioned variable being the same (COT) or assume grouped data. However, with simple changes in transport-plan parametrization to include two conditioned variables in methods like [1*], may also work. Some discussion or comparison with such simple extensions of existing methods may improve the work.

It seems most of the reviewers initially found the motivation for the basic eqn (9) and other details hard to follow. Even in final comments some reviewers have stressed on clarity. I request the authors to please elaborate on such aspects in their next revision.

[1*] https://proceedings.mlr.press/v238/manupriya24a/manupriya24a.pdf
Decision: Accept (poster)