GENOT: Entropic (Gromov) Wasserstein Flow Matching with Applications to Single-Cell Genomics

Single-cell genomics has significantly advanced our understanding of cellular behavior, catalyzing innovations in treatments and precision medicine. However, single-cell sequencing technologies are inherently destructive and can only measure a limited array of data modalities simultaneously. This limitation underscores the need for new methods capable of realigning cells. Optimal transport (OT) has emerged as a potent solution, but traditional discrete solvers are hampered by scalability, privacy, and out-of-sample estimation issues. These challenges have spurred the development of neural network-based solvers, known as neural OT solvers, that parameterize OT maps. Yet, these models often lack the flexibility needed for broader life science applications. To address these deficiencies, our approach learns stochastic maps (i.e. transport plans), allows for any cost function, relaxes mass conservation constraints and integrates quadratic solvers to tackle the complex challenges posed by the (Fused) Gromov-Wasserstein problem. Utilizing flow matching as a backbone, our method offers a flexible and effective framework. We demonstrate its versatility and robustness through applications in cell development studies, cellular drug response modeling, and cross-modality cell translation, illustrating significant potential for enhancing therapeutic strategies.

Paper

References (85)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer GBWj7/10 · confidence 3/52024-07-10

Summary

The paper proposes a multi-variable OT-based framework for solving Single-Cell related problems. The main idea behind the approach is to use the solutions for different variants of discrete entropy-regularize OT problems to learn a continuous parametric plan (measure) given by its conditional distributions, which approximates the true solution to the corresponding continuous problem. The authors propose utilizing conditional Flow Matching to parameterize the learned plan. The methodology is validated on various synthetic and real (single-cell) use cases. The obtained results testify the competitiveness of the framework.

Strengths

Despite the fact that the paper is full of specific technical details on single-cell stuff, the manuscript is clear and pleasant for reading. The method is clearly introduced. It seems that the authors adequately and fairly exhibit both the advantages of the method - its universality (coverage of different variants of OT formulations), simplicity (both for understanding and I guess for implementation); and limitations of the methodology - good limitation section + discussed reliance on discrete OT. Also, I would like to commend a good experimental section - a lot of interesting use cases are considered and interesting results (sometimes even negative) are presented. These experiment settings are worth considering regardless of the proposed method.

Weaknesses

* It seems the method could not (or hardly could) be adapted to high-dimensional (e.g., image-data) setups due to reliance on discrete EOT. At least, I expect that the learned plan in such setups will significantly deviate from the true EOT plan. Still, I understand that the proposed framework is targeted on single cell data, which is not that high-dimensional. * Reliance on Discrete OT is something I personally treat as a weakness. It introduces biases which hardly could be analyzed. Still, in practice, the minibatch OT seems to work well.

Questions

* The main of my questions is about the utilization of (conditional) flow matching for learning the parametric plan out of discrete minibatches. Why flow matching? As I understand, the proposed framework could be used on par with other generative techniques like GANs/normalizing flows/diffusion models. Why do you restrict yourself by flow matching? * I found it interesting that GENOT achieves (on average) the best results on EOT benchmark [3] among the competitors. Could the authors provide the code for the benchmark experiment? * Unbalanced Entropic OT work (to be mentioned): [1]. Also, a recent preprint you may find to be interesting: [2]. * Line 89: Do I understand correctly that $\mu$ and $\nu$ are not necessarily probability distributions, but could be positive measures with the same mass? * Why does in (UQEOT) one utilize $\varepsilon \text{KL}^{\otimes}$, while in (QEOT) - just $\varepsilon \text{KL}$ (without $\otimes$)? * Line 289: What is the “conditional mean” regime of the GENOT model? * Why does the unfused methods work badly on the single cell modalities translation problem with $\ell_2^2$ intra-costs (Figure 4, FOSCTTM metric is almost everywhere near 0.5)? Misprints: - (QEOT) formula: $Q_{c_X, c_Y}$ -> $D_{c_X, c_Y}$ - line 101: complexity if -> complexity is - line 351: Fig. 3 -> Fig. 4 [1] Pariset et. al., Unbalanced Diffusion Schrödinger Bridge, ICML’23 workshop. [2] Gazdieva et. al., Light Unbalanced Optimal Transport. [3] Gushchin et. al., Building the Bridge of Schrödinger: A Continuous Entropic Optimal Transport Benchmark ------------------------ POST REBUTTAL I thank the authors for the answers provided. I am satisfied with them. I have decided to raise my grade.

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Ok

Authorsrebuttal2024-08-02

Additional Answer to Reviewer GBWj

> ***What is the “conditional mean” regime of the GENOT model?*** ➤ This refers to mapping a point $\mathbf{x}$ to the target domain via the conditional mean, i.e., by averaging multiple samples from the conditional distribution $\pi^\star_\varepsilon(\cdot|\mathbf{x})$ instead of taking a unique one $\mathbf{y}\sim\pi^\star_\varepsilon(\cdot|\mathbf{x})$. This is detailed in lines 134-138. While we would like to highlight that taking the conditional mean has to be taken with a grain of salt (as in most cases, this is not the quantity we are interested in when computing EOT plans), it can be beneficial in some single-cell applications with respect to certain metrics. For instance, in Figure 4 (translating modalities), we see that the FOSCTTM of GENOT-Q CM is lower (hence better) than that of GENOT-Q. Importantly, numerous single-cell applications explicitly model the barycentric projection (i.e., conditional mean), e.g., for translating modalities of cells ([Demetci+2022]), mapping cells to their spatial organization [Nitzan+2019], and aligning slides of spatial transcriptomics [Zeira+2023]. > ***Why does the unfused methods work badly on the single cell modalities translation problem with l2 intra-costs (Figure 4, FOSCTTM metric is almost everywhere near 0.5)?*** ➤ This is an interesting question. It seems that pairwise Euclidean distances are not sufficiently characteristic for cells. The suitability of the cost might depend a lot on the dataset. These observations are also observed in the original publication for using Gromov-Wasserstein for aligning modalities of cells (SCOT, [Demetci+2022], Fig. 6). Unfortunately, we are not aware of any method on how to choose the “best” cost a priori, but our experiments (as well as the experiments shown in Fig.1, pdf) show that the choice of the cost can have an impact on the performance. Note that the poor performance of using the squared Euclidean cost motivated our "extension" of SCOT, i.e. constructing a fused term which empirically helps. > ***Misprints.*** Thanks a lot for catching these typos. We adapted the relevant sections in the paper.

Reviewer GBWj2024-08-08

Thanks for the answers

I thank the authors for the provided answers. I have decided to raise the score (up to 7)

Authorsrebuttal2024-08-10

Many thanks for reading our rebuttal.

We are very grateful for your score increase, and we are happy to answer any further question on our work. The Authors

Reviewer NoBU6/10 · confidence 4/52024-07-12

Summary

The paper proposes a framework for realigning cells in single-cell genomics which lies in the field of neural Optimal Transport (OT) solvers. The method utilizes a generative flow-based model for computing entropic OT couplings and tackles several practical challenges, e.g., it allows for using arbitrary cost functions, learning stochastic transport plans, relaxing mass conservation constraint, and tackling challenges of the (Fused) Gromov-Wasserstein problem. Specifically, the method consists of training of the conditional flow matching model on top of estimated discrete entropic OT plans. The approach is tested in different tasks from the single-cell biology field.

Strengths

The paper proposes a flexible framework which allows for addressing different challenges of application of neural OT solvers in the single cell genomics field. It is shown that the approach provides good results in several tasks from the single cell genomics according to the metrics under consideration. The paper is well-written and structured.

Weaknesses

My major concern is that the provided approach is not guaranteed to learn the ground-truth plans. Most of the experimental details and theoretical results related to this topic are moved to Appendix which hampers the understanding of this question. Indeed, GENOT is based on the idea of distilling discrete entropic OT (EOT) solutions using the flow matching (which is close, but as the authors explain - different, to the idea from the work [1]). It raises a question: how well does such a kind of distillation of discrete plans approximate the ground-truth continuous ones (assuming that the approach works with batches of finite size)? Some explanations regarding the mini-batch biases are provided in lines 267-280. However, as far as I see it, the provided discussion (and references) do not cover the statistical properties of the discrete *unbalanced* linear or *quadratic plans*. The situation with quadratic setup is the most tricky here, since it is known that GW problem admits multiple solutions and, thus, the optimization problem may differ at each step of training depending on the calculated discrete entropic GW plans. From the empirical side, the authors address the ‘biasedness’ issue by testing their approach using a benchmark with analytically known EOT [2], U-EOT [3], and GW/Unbalanced GW [4] plans (all in Appendix E). According to these experiments, I tend to agree that the proposed approach provides meaningful results for the balanced (or unbalanced) linear case. However, the situation with the GW case is not clarified – according to the experiment description (Appendix E.3), the authors introduce some heuristics to “preserve the orientation of GW solution” and claim that in the case of biological tasks it is not needed. (Still, as a non-expert in biology, I can not assess the validity of this claim.) **In summary**, there are some evident issues with the implementation of the proposed approach in the Gromov-Wasserstein case. I appreciate that the authors mentioned this issue in the limitations section of their work (Appendix A). However, it is not obvious for me to what extent this is a serious drawback hampering the applicability of the approach to biological tasks. Besides, I think that as an important aspect of the implemented approach, this issue should be mentioned in the main body of the paper.

Questions

- In the related work section you cited several works on unbalanced EOT [5,6,7]. Why did you include only one of them [7] in comparison in Gaussians experiment? - In Table 1-3 you use abbreviations DBSM and CFM-SM which are not introduced in the text. Please introduce them in the relevant section. - Could you explain in more details why the issues of your implementation in the GW case will not arise in biological tasks? I see your explanation in lines 1345-1346, however, it seems that the issue might arise not only in the case of highly symmetric distributions

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors adequately addressed the limitations and potential negative societal impact of their work in a specific section of their work.

Authorsrebuttal2024-08-05

Additional Answer to Reviewer NoBU 1/2

> ***The situation with quadratic setup is the most tricky here, since it is known that GW problem admits multiple solutions and, thus, the optimization problem may differ at each step of training depending on the calculated discrete entropic GW plans.*** ➤ We agree that the quadratic setup is the most tricky one, and we thank the reviewer for explicitly mention this issue again. As mentioned in the general response, we set out to investigate this further (Alg. 1, pdf), and results can be found in Fig. 2, pdf. We consider two datasets: 30-dim. Gaussians (20 dim. quadratic + 10 dim. fused term): For both the source ($\mathcal{\mu}$) and the target ($\nu$) distribution we sample mean vectors from $\mathcal{U}([-1, 1])$ and construct the covariance matrices by sampling each entry of their square roots from $\mathcal{U}([0, 3])$. 30-dim. single-cell data (the one used for modality translation, but no train/test split): 20-dimensional PCA of the normalized ATAC data, 20-dimensional PCA projection of the gene expression data, and 10-dimensional PCA of VAE space (shared space). To assess the stability of the solution of discrete EOT, we fix one sample $X_0 \sim \mu$ from the source distribution, and compute the variance of $\hat{\pi}(\cdot|X_0)[10:]$, i.e. of the incomparable space $\tilde{\mathcal{Y}}$ only, which is obtained from a discrete EOT solver with different samples in source and target. We do this for $\varepsilon=0.01$ and $\varepsilon=0.0001$ with the Sq. Euclidean cost. We consider: - outer coupling: this serves as orientation. - Gromov-Wasserstein (GW): We use $\tilde{\mathcal{X}}$ and $\tilde{\mathcal{Y}}$, i.e., we use the last 20 dimensions of both the source and the target distribution. - Fused Gromov-Wasserstein (FGW): We use $\Omega \times \tilde{\mathcal{X}}$ and $\Omega \times \tilde{\mathcal{Y}}$, $\alpha=\frac{1}{11}$. - Gromov-Wasserstein with geodesic cost (Geodesic): We use $\tilde{\mathcal{X}}$ and $\tilde{\mathcal{Y}}$ with the geodesic cost. - Gromov-Wasserstein with initialisation (with init): We use $\tilde{\mathcal{X}}$ and $\tilde{\mathcal{Y}}$. We use the initialisation scheme as discussed in paper (Appendix A, E). > ***However, the situation with the GW case is not clarified – according to the experiment description (Appendix E.3), the authors introduce some heuristics to “preserve the orientation of GW solution”*** ➤ We visualise what we mean by "preservation of the orientation" of the GW solution in Fig. 1 in the provided pdf, where we transport a 1-dimensional Gaussian to a 1-dimensional Gaussian. The LHS of Fig. 1 shows the result obtained without initialisation, and the learnt coupling is clearly a mixture of solutions, while on the RHS of Fig. 1, we visualise that GENOT-Q learns a coupling which appears to be no mixture of solutions. While this is only an example in 1 dimension, the results in Fig. 2, pdf, highlight that using the proposed initialisation scheme stabilises the orientation of the GW solution also in high dimensions, for both the Gaussian and the single-cell experiments. > ***[...] in the case of biological tasks it is not needed. (Still, as a non-expert in biology, I can not assess the validity of this claim.)*** ➤ We thank the reviewer for pointing this out, apologize for this bold statement, and would like to correct ourselves and weaken the statement to: "In the case of single-cell data, choosing the right cost might help to overcome the need for the initialisation scheme to a certain degree." as can be seen from Fig. 2, pdf, where the geodesic cost seems to help make the solution more unique with $\varepsilon=0.0001$, as well as the results in Fig. 4, manuscript, LHS. Moreover, we will add the sentence "Additionally, most single-cell applications are a FGW problem (as opposed to a GW problem), and empirically, the fused term helps to make the set of solutions of the EOT coupling less diverse, as can be seen from Fig. 5, FIg. 19, and [provided pdf, Fig. 2].". Indeed, the only pure GW problem in common single-cell applications we are aware of is SCOT [Demetci+22]. As we realised the performance of both SCOT and GENOT-Q is poor, we introduced a novel way to include a fused term into the task of modality translation by constructing a joint space by creating a joint conditional VAE embedding from gene activity (source space) and gene expression (target space). > ***However, it is not obvious for me to what extent this is a serious drawback hampering the applicability of the approach to biological tasks.*** ➤ We hope to have clarified this with the additional experiments and the answers above.

Authorsrebuttal2024-08-05

Additional Answer to Reviewer NoBU 2/2

> ***Besides, I think that as an important aspect of the implemented approach, this issue should be mentioned in the main body of the paper.*** ➤ We fully agree with the reviewer, and will definitely include the discussion mentioned above in the main body of the paper, highlighting this limitation. Moreover, we will include the results on the stability of quadratic EOT couplings as outlined in the provided pdf document. > ***In the related work section you cited several works on unbalanced EOT [5,6,7]. Why did you include only one of them [7] in comparison in Gaussians experiment?*** ➤ We assume the reviewer is alluding to lines 56-59: [Eyring+2024], [Lübeck+2023], [Yang+2019] proposed a way to incorporate unbalancedness into deterministic linear OT maps, while unbalanced formulations for entropic OT in both the linear and the quadratic case have not been explored.” In fact, [Eyring+2024], [Lübeck+2023] are unbalanced Monge Maps (i.e. deterministic) estimators, and hence do not solve EOT. [Yang+2019] suggest a solver for both deterministic and entropic OT couplings, but their implementation only allows for the deterministic OT coupling. We now clarify this, and also include references to other preprints (thanks to reviewer GBWj) by adapting the statement in lines 56-59 to: "[Eyring+2024], [Lübeck+2023], [Yang+2019] proposed a way to incorporate unbalancedness into deterministic linear OT maps, while [Yang+2019] extended their method to the entropic setting (but do not provide code), while recent preprints [Pariset+2023], [Gazdieva+2024] also suggested methods to approximate unbalanced linear EOT solutions. To the best of our knowledge, unbalanced neural formulations in quadratic OT have not yet been explored." > ***In Table 1-3 you use abbreviations DBSM and CFM-SM which are not introduced in the text. Please introduce them in the relevant section.*** ➤ We thank the reviewer for pointing us to this. We will add all references to Table 1. > ***Could you explain in more details why the issues of your implementation in the GW case will not arise in biological tasks? I see your explanation in lines 1345-1346, however, it seems that the issue might arise not only in the case of highly symmetric distributions*** ➤ As mentioned above, this has two main reasons: First, empirically, the geodesic cost seems to overcome this limitation to a certain extent (see superior performance in Fig. 4, paper, and lower conditional variance in Fig.2, pdf). Second, most single-cell applications can be extended to a fused problem. Yet, we apologize for our formulation that this issue only arises in highly symmetric contributions, and will correct the statement to "Empirically, both changing the cost and including a fused term can help to account for this limitation. If none of these approaches work, GENOT can be run with its quadratic initialisation scheme, which empirically helps to make the set of EOT solutions less diverse, at the cost of loosing simulation-free training."

Reviewer GAfL6/10 · confidence 3/52024-07-12

Summary

The authors present Generative Entropic Neural Optimal Transport (GENOT), a flexible method for learning entropic couplings for linear and quadratic entropic optimal transport (EOT), unbalanced OT, and OT across incomparable spaces. At its core, their method uses conditional flow matching (CFM) to solve for the OT couplings. Through extensive empirical experiments, the authors show the use-case of their method for learning couplings of cell trajectories, predicting single-cell response to perturbations, and translating between single-cell data modalities.

Strengths

The authors present promising and comprehensive set of empirical results that support the applicability of their method for OT related problems in single-cell biology. They demonstrate a novel use case of learning OT couplings between different data modalities in singel-cell biology, i.e. from RNA-seq to ATAC-seq. Their proposed method is also supported with sound theoretical background and justification. In general, the paper is well motivated and addresses important and challenging problems in single-cell biology.

Weaknesses

Overall, this work has all the elements that constructs a complete and thorough application-based contribution to the field of neural optimal transport and single-cell biology. However, there remain several items that I feel hinder the understanding and impact of this work: - It seems to me that a key novelty of this work is the use of quadratic OT for learning couplings between incomparable spaces -- i.e. from ATAC-seq to RNA-seq for the single-cell case considered in this work. The authors consider a large quantity of experiments for both linear and quadratic OT over various settings, with many of the experiment being reported in the appendix. This makes some of the results and discussion hard to follow, and overall diluting the presentation of the key items of this paper. Given the large quantity of experiments and dense content of this paper, maybe it would be beneficial to shift more focus on these aforementioned experiments and move other to the appendix. This may help the reader better understand some of the key contributions of this work. - A large amount of items that are placed in the appendix tare referenced in the main text and seem to be important. This in general disrupts the overall fluidity of the paper. For instance, propositions in section 3 are listed in the appendix. Moreover, some experiments discussed in the main text only reference figures that are in the appendix. For example, experiments in section 5.1 "GENOT-L on simulated data" and "U-GENOT-L predicts single-cell responses to perturbations". In general, there seem to be a large proportion of experimental results that are placed in the appendix and referenced in section 5, which seem to be given equal wait to the experiment included in the main text. This makes it difficult to decipher the significance/importance of the respective results. As mentioned in my previous comment, possibly focussing the paper slightly may be beneficial to the overall presentation of this work.

Questions

- For the experiment on single-cell response to perturbations, it appears there are no comparisons to baseline methods? Is the reason for this because existing baseline methods are deterministic? Is there intuition on why the stochastic approach is helpful (or necessary) for this task?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors discuss limitations and broader impacts in the appendix.

Reviewer GAfL2024-08-08

I thank the authors for their detailed response and answers to my questions. I will keep my positive score.

Authorsrebuttal2024-08-10

Thanks for the response to our rebuttal

We appreciate the positive feedback and are more than happy to address any further questions. The Authors

Reviewer NoBU2024-08-13

I thank the authors for providing the answers to my concerns and raise my score to 6.

Authorsrebuttal2024-08-13

Thanks for the response to our rebuttal

We are glad we could address all concerns and are more than happy to address any further questions you may have about our work. The Authors

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC