Title: Pairwise Optimal Transports for Training All-to-All Flow-Based Condition Transfer Model
Venue: NeurIPS 2025
Accepted: yes
Keywords: Flow Matching

Conditional Transfer
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n_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.\n\nQUESTIONS:\n1.\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.\n\nLIMITATIONS:\nAs the author said, the balance of dataset size and the hyperparameter $\\beta$ can pose a limitation.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 1, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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**\n\nQUESTIONS:\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- 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.\n\nLIMITATIONS:\nThe authors have addressed the limitations of their approach.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n# 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.\n\nQUESTIONS:\nProposition 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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 4, "ORIGINALITY": 4, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n## 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.\n\nQUESTIONS:\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   - 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.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}