Learning diffusion at lightspeed

Weakness 1

We appreciate the suggestions to improve the exposition and agree with the reviewer's comment. Specifically, we now removed "phenomenal", "runs at lightspeed" and "few weeks old advancements" from the abstract and introduction, rephrasing so to be more factual in the exposition, focusing on the contributions. Thanks!

Weakness 2

We agree that the introduction blends with the related works rather abruptly. We thus prepared a revised version in which we smoothed it, gradually adding information for the reader less familiar with the literature.

For instance, we will add the paragraph "However, the JKO scheme entails an optimization problem in the probability space. Thus, the problem of finding the energy functional that minimizes a prediction error (w.r.t. observational data) takes the form of a computationally-challenging infinite-dimensional bilevel optimization problem, whereby the upper-level problem is the minimization of the prediction error and the lower-level problem is the JKO scheme." to better introduce the reader to what the bilevel optimization problem is about.

We also removed the focus from Figure 1 in the introduction.

Weakness 3

We thank the reviewer for pointing out this difference, which in fact is a strength of our method which we did not highlight appropriately and we hope we now did in the Remark 3.6 we prepared for the next revision of the paper. In a nutshell, it is true that $\mathrm{JKOnet^*}$ requires the construction of the optimal transport couplings beforehand. However, $\mathrm{JKOnet}$ constructs a new optimal transport plan at each iteration depending on the current estimate of the potential, whereas $\mathrm{JKOnet^*}$ needs to do so only once, at the beginning. We prepared a revision of the paper in which we added the time required to compute the optimal transport couplings in Section 4.1 ($0.03 \pm 0.01$s), and we provided an analysis of the comparison between Sinkhorn and plain linear programming for our scope in the newly added Figure 10 in Appendix C.2, which we now expanded (we report Figure 10 in the PDF attached to this rebuttal, and we will include the extended version of the Appendix C.2 for the next revision of the paper). We also better argued why the scaling to high-dimensions remains unaffected by the computation of the optimal transport couplings: the dimensionality affects only the construction of the cost matrix in the linear program, and otherwise the computational complexity of the linear program only relates to the number of particles. When the number of particles grows, one can apply the same batching applied in other methods to pre-compute the couplings. We add details regarding these practical considerations in the revised Appendix C.2. Finally, following the comments from reviewer vNRb, we introduced a real-world case study on learning and predicting molecular processes (see Figure 5 and the comparison table in the attached document). This way, we hope we strengthen our contributions not only by achieving state-of-the-art on a real-world application but doing so in under a minute of training (including the computation of the optimal transport couplings) compared to the hours required by the other methods.

Question 1

Thanks for pointing out the notation confusion. We defined $\mathrm{JKOnet^*_l}$ only later on, so now we introduced in the caption to clarify we refer to the linear parametrization.

Question 2

Thanks for pointing this out. We agree that this is an important point and we will discuss it in the revised version of the paper. In particular, the dimensionality affects only the construction of the cost matrix in the linear program associated with the optimal transport problem, and otherwise the complexity is only related to the number of particles. Specifically, the time required to construct the cost matrix scales linearly with the dimension, and in practice it is minimal and dwarfed by the actual solution of the linear program. When the number of particles grows, one can apply the same batching that is applied in other methods to pre-compute the couplings. We add details regarding these practical considerations in the revision of Appendix C.2. Please see also point 3 above.

Question 3

Thanks for observing the performance gains and pointing out this aspect. We prepared a revision of Section 4.1 in which we added a paragraph highlighting also the performance gains and discussing linear vs non-linear. In particular, the linear approximation has optimality guarantees, as long as the features are sufficiently rich. In high dimensions, however, the choice of feature is challenging and, thus, we recommend resorting to non-linear parametrizations.

Weakness 1

We thank the reviewer for the suggestion. We prepared a revision of the paper in which we added a reference to the appendix when referencing to content related to the appendix.

Weakness 2

Good catch, thank you! We prepared a revision of the paper in which we added the definition of $\rho$ below the Fokker-Plank equation.

Weakness 3

We thank the reviewer for pointing out this difference, which in fact we believe to be a strength of our method which we did not highlight appropriately. To this end, we included a dedicated remark, Remark 3.6, for the next revision of the paper. In a nutshell, it is true that $\mathrm{JKOnet^*}$ requires the construction of the optimal transport couplings beforehand. However, $\mathrm{JKOnet}$ constructs a new optimal transport plan at each iteration depending on the current estimate of the potential, whereas $\mathrm{JKOnet^*}$ needs to do so only once, at the beginning. We prepared a revision of the paper in which we added the time required to compute the optimal transport couplings in Section 4.1 ($0.03 \pm 0.01$s), and we provided an analysis of the comparison between Sinkhorn and plain linear programming for our scope in the newly added Figure 10 in Appendix C.2, which we now expanded (we report Figure 10 in the PDF attached to this rebuttal, and we will include the extended version of the Appendix C.2 for the next revision of the paper). We also better argued why the scaling to high dimensions remains unaffected by the computation of the optimal transport couplings: the dimensionality affects only the construction of the cost matrix in the linear program, and otherwise the computational complexity of the linear program only relates to the number of particles. When the number of particles grows, one can apply the same batching applied in other methods to pre-compute the couplings. We add details regarding these practical considerations in the revised Appendix C.2. Finally, following the comments from reviewer vNRb, we introduced a real-world case study on learning and predicting molecular processes (see Figure 5 and the comparison table in the attached document). This way, we hope we strengthen our contributions not only by achieving state-of-the-art on a real-world application but doing so in under a minute of training (including the computation of the optimal transport couplings) compared to the hours required by the other methods.

Weakness 4

Good point, thank you for the comment. We have a dedicated section in Appendix B, which we now expanded to discuss the prediction schemes as well, and we have added an introduction to the problems in the experimental section. In particular, we added the data-generation equation, in which the role of the functionals $V(x)$ is apparent: $x_{t+1} = x_t - \tau \nabla V(x_t)$. We also added the name and equation of each functional in the figures where it was not listed (e.g., Figure 2).

Question 1

Thanks for the question. In our implementation, we solved the optimal transport problems via plain linear programming (using the POT python library). We prepared an updated version of Appendix C.2 that we will include in the revision of the paper that contains an analysis of the impact of different ways of solution algorithms on the final outcome in terms of computational time during pre-processing and Wasserstein error (see Figure 10 in the attached PDF). In particular, we conclude that, as long as the couplings are close to the correct one, the algorithm used to compute them does not impact the performance of $\mathrm{JKOnet^*}$. Since small regularizers slow down the Sinkhorn algorithm, we opted to directly solve the linear program (without regularization). In general, the solver choice can be considered an additional knob that researchers and practitioners can tune when deploying $\mathrm{JKOnet^*}$. Please see also point 3 above.

Question 2

Once an energy functional through $\mathrm{JKOnet^*}$ is learned, the equilibrium state of the system can be inferred in two ways. A simple approach consists of running sufficiently many iterations of the JKO scheme until an equilibrium is reached. Alternatively, the equilibrium state is well-known to be the minimum of the energy functional. Thus, the equilibrium state can be inferred by computing the probability distribution which minimizes the energy functionals, using tools from optimization in the probability space.

Weakness 1

We thank the reviewer for suggesting us one way to strengthen the presentation of our contributions.

We deployed our method to learn the diffusion dynamics of embryoid body single-cell RNA sequencing (scRNA-seq) data [1], a popular benchmark in the literature, and compared our results with nine other recent methods in the literature. We discuss the application and the results in the newly added Section 4.4, which will appear in future versions of the paper, and we report the related figure and table in the PDF attached to this response. We briefly summarize our experimental setting and results below.

Experimental setting (briefly): We follow the same data pre-processing as in [2,3]; in particular, we use the same processed artifacts of the embryoid data provided in their work, which contains the first 100 components of the principal components analysis (PCA) of the data and, following [2,3], we focus on the first five. We train on $60\%$ of the data at each time-step and test $\mathrm{JKOnet^*}$ to predict the evolution of the left-out data.

Results: $\mathrm{JKOnet^*}$ outperforms all existing methods in the literature (see the results table in the PDF attached). Importantly, our training takes less than a minute (including the computation of optimal transport plans), whereas the training time of all other existing methods takes hours.

We visualize the 2 principal components and the interpolations obtained with our method in Figure 5. Note that, for this application, we used a time-varying potential energy, of which we plot the level curves.

Weakness 2

We believe our approach might underperform in the following two cases.

First, while diffusion processes include many real-world phenomena, in some real-world applications it might be unknown if the particles are undergoing diffusion. If this is not the case, $\mathrm{JKOnet^*}$ might underperform. For instance, in the absence of noise and interaction energy, if the vector field is not the gradient of some potential energy function (e.g., it includes a solenoidal component $\nabla \times \psi$ for some function $\psi: \mathbb{R}^d \to \mathbb{R}^d$), we cannot expect $\mathrm{JKOnet^*}$ (as well as any other method learning a potential) to infer a reasonable potential. We prepared a dedicated section that we will include in the revision of the paper to discuss this failure mode.

Second, when learning both potential and interaction energy and noise level, there might be observability issues that prevent the distinction of the different components of the energy functional (e.g., a discrete-time population-level effect might be explained both by a potential energy and a noise level).

While we did not experience this issue in our experiments in Section 4.3, we are not aware of rigorous guarantees that ensure observability. In a dedicated section in Appendix G, we provide a small-scale analytical example illustrating this issue. As $\mathrm{JKOnet^*}$ is the only method capable of simultaneously learning all three energy components, we believe it can serve as the baseline to investigate this observability issue.

Weakness 3

We compared our method with others in our real-world application, discussed in 1) above and in Section 4.4 in the revised version of the paper.

Question 1, 2

Please refer to our answer to Weakness 1) above.

Question 3

In our experiments, we considered only vanilla parametrizations: two-layers MLP with 64 neurons in each layer, to compare directly with the other works in the literature. What is certainly required is a network that is expressive enough to approximate the energy functional of interest, so standard rationales apply for the choice of activation functions (we use softplus), dimension of the networks, etc. One of the limitations of our works is the data domain (we do not explore, e.g., images). We plan in future work to do so, and in that case more care will be needed to determine the most suitable architecture. Given that the same architecture and hyperparameters worked well across all the experiments (including the real-world experiment), we are confident to state that the learning algorithm itself is not particularly sensitive to the network architecture, which needs of course be chosen to be suitable to represent the energy in the application of choice. We also believe that exciting future work can be done in terms of understanding how different architectures can capture potential, internal, and interaction energies more efficiently: can e.g., transformers be used to learn an interaction energy more efficiently than a vanilla MLP?

[1] "Visualizing structure and transitions in high-dimensional biological data" by Kevin R Moon, David van Dijk, Zheng Wang, Scott Gigante, Daniel B Burkhardt, William S Chen, Kristina Yim, Antonia van den Elzen, Matthew J Hirn, Ronald R Coifman, et al. (2019)

[2] "Improving and generalizing flow-based generative models with minibatch optimal transport" by Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Kilian Fatras, Guy Wolf, and Yoshua Bengio. (2023)

[3] "Deep momentum multi-marginal schr"{o}dinger bridge" by Tianrong Chen, Guan-Horng Liu, Molei Tao, and Evangelos A Theodorou. (2023)

We believe the problem of score-matching in diffusion models to be fundamentally different from the one in our paper. In score-matching, one tries to "reverse" the time of a known diffusion process, e.g., to recover the uncorrupted state of a corrupted image. In our setting, instead, we use observational population data to learn the energy functional underlying an unknown diffusion process. Our goal is not to "reverse" the time of a given diffusion process to reconstruct its initial condition but rather to learn an unknown diffusion process to perform forward-in-time predictions.

From a technical perspective, our methodology heavily relies on optimal transport theory and tools from optimization in the probability space. Indeed, our loss can be interpreted as the "error" in satisfying a first-order optimality condition in the Wasserstein space. To the best of our knowledge, this approach has not previously appeared in the literature. For instance, the loss function in reference [1] suggested by the reviewer is constructed using tools from stochastic differential equations: Their loss function minimizes the errors in estimating the term $\nabla_x\log q_t(x)$ which appears when "reversing" the time of a (known) stochastic differential equation. For this reason, their loss cannot be reconciled to ours (in which the stochastic differential equation is unknown).

[1] "Soft diffusion: Score matching for general corruptions" by Daras, G., Delbracio, M., Talebi, H., Dimakis, A. G., & Milanfar, P. (2022).