Paper: A Wasserstein-2 Distance for Efficient Reconstruction of Stochastic Differential Equations (ICLR-2024)

Reviewer comment (Theory):
There are no generalization bounds.

Author response:
Q1 \& W1. Thank you for your question, and we appreciate the opportunity to
clarify our approach in addressing your question on ``generalization
bounds".

Our Theorem 2 provides a time-discretization error bound for the difference
between the Wasserstein
distance of solutions to two distinct SDEs given the same initial condition and the Wasserstein distance
of finite-time-point observations of the solutions to those two SDEs.

The
stochastic nature of the Stochastic Differential Equation (SDE) we aim
to reconstruct inherently complicates the straightforward extension of
"generalization bounds" defined for deterministic systems. Thus, in Examples
2, 3, and 4, we evaluate the performance of SDE reconstruction using
the metrics $|f-\hat{f}|$ and $|\sigma-\hat{\sigma}|$, representing
the errors in reconstructing the drift and diffusion terms of the SDE,
respectively. This choice aligns with our method's objective of
reconstructing general unknown SDEs from data.


It would be nice to obtain ``finite sample generalization
bounds". Even though a theoretical derivation of such bounds is quite nontrivial
challenging, we have conducted a sensitivity analysis in Examples 2,
3, and 4 to investigate how the number of samples impacts the
reconstruction of $f$ and $\sigma$ (refer to Figs. 2, 3, and 4). Our preliminary results
indicate that more samples lead to better accuracy and smaller $|f-\hat{f}|$ and $|\sigma-\hat{\sigma}|$.

Analyzing the error bound between the empirical Wasserstein distance
and the genuine Wasserstein distance is complicated and depends on
both the number of samples and the dimensionality of the SDE. This is
beyond the scope of our paper, as our focus was to introduce and
develop a novel and practical (see our examples and additional applications)
Wasserstein-distance-based SDE reconstruction method as well as
provide insight into its efficacy.


In theory, there is also a theorem for analyzing the errors between the
empirical Wasserstein distance with finite samples and the genuine
Wasserstein distance, see, \textit{e.g.}

Ref [1]. Fournier, Nicolas, and Arnaud Guillin, On the rate of
convergence in Wasserstein distance of the empirical measure,
$\textbf{Probability theory and related fields}$, 162(3-4): 707-738,
(2015).

Currently, our ongoing efforts involve both theoretical analysis of
our Wasserstein-distance SDE reconstruction method and practical
applications in reconstructing multiple SDEs with noisy dynamics in
biology. We have
cited Ref [1] in our Example 4 to clearly state that more analysis is
needed on the number of samples and dimensionality in our future work.

We hope that these clarifications effectively address your concerns
and we are open to further discussion or inquiries about our new
Wasserstein-distance-based loss function and SDE inference approach.

Reviewer comment (Theory):
A major concern is the tightness of equation (17). As a simple example, if $f,\sigma$ are both only functions of $t$ not $x$, then the finite projections will be jointly Gaussian, the WD between which is not likely to be just the sum of WD between marginals. Thus it is hard to assess how good the approximation is, even if it seems to work well in the provided examples.

Author response:
W2. We appreciate your thoughtful
consideration of Eq. (17) and your concern regarding its tightness.

Eq. (17) is an approximation to Eq. (16), and it is helpful to
evaluate how it approximates Eq. (16) in general cases.  One rationale
behind employing Eq. (17) in our examples is its advantageous
simplification for univariate SDEs. In such cases, Eq. (17) reduces to
Eq. (3), allowing for analytical calculations and a clear statistical
interpretation. Specifically, Eq. (17) aims to match the distribution
of reconstructed trajectories with the ground truth trajectories at
different time points.

In Example 4, we also directly minimized Eq. (16) (the blue line) as
the loss function. Our result shows that minimizing Eq. (16) performs
worse than minimizing Eq. (17). To provide further insight, we
incorporated a visualization of the loss functions concerning training
epochs in Example 4, demonstrating the superior numerical efficiency
of minimizing Eq. (17) relative to minimizing Eq. (16).

In theory, numerically calculating Eq. (16) could be less accurate
especially when the number of samples is limited, as the bound of the
discrepancy between the empirical Wasserstein distance and the genuine
Wasserstein distance deteriorates when the dimensionality increases.
In using Eq. (16), the dimensionality is $Nd$, where $N$ is the number
of timesteps and $d$ is the dimensionality of the SDE. See,
\textit{e.g.},

Ref [4]. Nicolas Fournier and Arnaud Guillin, On the rate of convergence in
Wasserstein distance of the empirical measure, $\textbf{Probability
  Theory and Related Fields}$, 162.3-4: 707-738, (2015).

A comprehensive comparison between
Eq. (16) and the sum of Wasserstein distances between marginals
(Eq. 17) requires further discussion, especially when dealing with
multidimensional SDEs, which we outlined in Example 4. This intricate analysis goes beyond the scope
of our current manuscript, which primarily introduces the novel
WD-based SDE reconstruction method and provides preliminary insight
into its efficacy.

As applications, we have successfully applied Eq. (17) as a loss
function for reconstructing multiple multivariate SDEs.  Currently, we
are working on analyzing WD for multidimensional SDEs and exploring
whether there are more suitable loss functions than Eq. (17).

We hope that our responses can satisfactorily address your questions
and concerns and that you may better appreciate the new,
  practical, and efficient Wasserstein-distance-based SDE
  reconstruction approach that we developed.

Reviewer comment (Theory):
The major question is how tight the approximation (17) is, as it was extensively used in the examples, and is claimed to be an approximation of the original WD.

Author response:
Thank you for your valuable comments.

We wish to bring to your attention that both Eq. (16) and Eq. (17) are novel Wasserstein-distance-based loss functions for SDE reconstruction proposed by us, which provide two potential loss functions for reconstructing SDEs. Empirically, Eq. (17) performs better than Eq. (16). In our last response, we have also explained why numerically calculating Eq. (16) might not be so good due to the high dimensionality when coupling all time points. Furthermore, in our revised manuscript we also wrote in Example 4 that further discussion is needed to devise an optimal Wasserstein-distance-based loss function. 
Given the novelty and broad applications of our method in uncertainty quantification and general SDE reconstruction in biology, physics, and other related fields (as shown in https://drive.google.com/file/d/1HW5CBjx7g37DFx3s3A-sNDK0aoG6zqbT/view), we hope that our manuscript warrants your reevaluation.

Reviewer comment (Experiments):
I am not very familiar with this topic. There seems to be no major flaw in this paper, but I think it would be better if there are some more detailed comparisons with prior works and other losses, as the authors only provided some intuitive comparison with MSE and KL divergence losses in Appendix B.

Author response:
W1. Indeed in Appendix B we explained in theory why the MSE and KL divergence is not as good as the Wasserstein distance for the purpose of SDE reconstruction. Moreover, in Examples 1 and 2, we numerically compared our Wasserstein-distance based SDE reconstruction approach with MSE, maximal log likelihood, and Mean$^2$+Var approaches. Our Wasserstein distance loss gives the best reconstructed trajectories (Example 1) and smallest errors $f-\hat{f}$ and $\sigma-\hat{\sigma}$ (Example 2). In Examples 3 and 4, we compare our approach with recently proposed machine learning methods such as a Wasserstein Generative Adversarial Network (WGAN) approach and a Maximum Mean Discrepancy (MMD) loss. Once again, our Wasserstein distance method outperforms the WGAN and MMD approaches in both accuracy and efficiency.  We hope that these numerical experiments and comparisons with other methods convince you of the advantages of our new method.

Reviewer comment (Theory):
What kind of noise can the $W_2^2$ loss deal with?

Author response:
Q1. In principle, the $W_2^2$ loss can can
be calculated for a stochastic process obeying many kinds of noise.
For example, the Wasserstein distance between different
  solutions associated with different initial conditions of a general L\'evy process has been computed
, \textit{e.g.},.

Ref [1]. Jian Wang, $L^p$-Wasserstein distance for stochastic differential equations driven by L\'evy processes, $\textbf{Bernoulli}$, 1598-1616, (2016).

In our work, we minimize a Wasserstein loss to
reconstruct the SDE from data, and restrict ourselves to Brownian
noise processes.

Reviewer comment (Theory):
Can the authors provide some intuition about why $W_2^2$ loss is more robust to noise?

Author response:
Q2. The
robustness of the $W_2^2$ loss to noise can be attributed to its
nature as a distribution-based loss function, as exemplified by
Eq. (3). This type of loss is designed to quantify the difference
between the targeted distribution and the approximated distribution,
making it particularly well-suited for reconstructing the distribution
of trajectories. Compared to the MSE, $W_2^2$ captures the variation and compared to the Mean$^2$+Var , $W_2^2$ is able to capture more features such as multiple peaks, 
as illustrated in Example 1.

In Example 1, we also demonstrated the suitability of the $W_2^2$ loss for
reconstructing distributions. An additional advantage lies in the
flexibility of the $W_2^2$ loss. Unlike other commonly used loss
functions, such as Mean$^2$+Var or maximal log likelihood, the $W_2^2$
loss does not necessitate an assumption of a Gaussian
distribution. This characteristic enhances its robustness to various
types of noise, as it does not rely on specific distributional
assumptions.

Reviewer comment (Theory):
Is eq(10) purely of intellectual interest or there are some methods to estimate this expectation? It seems Theorem 2 didn't use this upper bound to construct an estimation of the $W_2^2$ loss.

Author response:
Q3. Eq. (10) is for Theorem 1. It is
a rather technical condition such that if satisfied, then
the $W_2^2$ distance between the solutions to two SDEs converges
uniformly to 0 as $\hat{f}\rightarrow f, \hat{\sigma}\rightarrow
\sigma$, as we explain on Page 4.

As a simple case, if $f, \sigma, \partial_x f, \partial_x \sigma$ are
uniformly bounded, then Eq. (10) is also uniformly bounded, and all
assumptions of Theorem 1 are satisfied.

Reviewer comment (Reproducibility):
Finally, I am curious what method did the authors use to minimize the $W_2^2$ loss?

Author response:
Q4. As described
on Page 6, training details are given in Appendix E. For all examples,
the AdamW stochastic optimization method is used for minimizing the
$W_2^2$ loss. The code implementation is available in the
supplementary files provided with the manuscript.

Reviewer comment (Theory):
Again with Remark 1, do the drifts $(f_i)$ and the diffusions components $(\sigma_i)$ satisfy the same assumptions in Theorem 1?

Author response:
If each $f_i, \sigma_i$ depend only on $X_i, t$ and satisfy the conditions in
Theorem 1, then Theorem 1 can be generalized to multivariate cases. We
are still working on extending Theorem 1 to interacting
multidimensional SDEs, which appears quite nontrivial.

Reviewer comment (Theory):
In Equations (16) and (17), the vectors $\boldsymbol{X}(t_i)$ and $\hat{\boldsymbol{X}(}t_i)$ are $d$-dimensional then the squared $(\boldsymbol{X}(t_i) -\hat{\boldsymbol{X}(}t_i))^2$ must be a squared-$L_2$ norm.

Author response:
We have revised
the $\ell^2$ vector norm notation $|\cdot|^2$ to $|\cdot|^2_2$ to avoid confusion.

Reviewer comment (Theory):
For the rotated $W_2$, I couldn't understand the intuition behind this application of rotations. I think there is an "embedding"-like of the origin components $X_1$ and $X_2$ of the SDE, hence it seems that we have changed the origin SDE, by changing its drift and diffusions components. So, the rotated $W_2$ is it an upper or a lower bound for the origin one?

Author response:
We
  appreciate your inquiry regarding the intuition behind the
  application of rotations in the rotated $W_2$ method. The rotated
  $W_2$ serves as a lower bound of the true $W_2$ distance. In higher
  dimensions, an analytical evaluation of the $W_2$ distances is
  challenging. Our approach approximates the $W_2$ distance in two
  ways: the decoupled $W_2$ and the rotated $W_2$.

  The primary method, which we refer to as the decoupled $W_2$
distance, involves projecting the 2D distribution in two orthogonal
directions and summing the $W_2$ distances between $1 \mathrm{D}$
projections of the true and reconstructed data. This method yields a
lower bound of the true $W_2$ distance and is exact (equal to the true
$W_2$) if the 1D projections $X_1$ and $X_2$ are mutually independent.

However, this approach may not effectively handle skewed (correlated)
distributions. We address this by applying a suitable rotation, as
illustrated in the figure https://drive.google.com/file/d/1wJdYAzZQhNEEOsQ-80Ne1GvWsuxxX0D-/view?usp=drive_link.

Practically, the
optimal rotation matrix is unknown, necessitating an average over
various rotations. Since each rotation is isometric, rotation 
does not change the true $W_2$ distance, and thus each decoupled
$W_2$ distance post-rotation remains a lower bound of the true $W_2$
distance. Therefore, the averaged $W_2$ distance after rotation also
constitutes a lower bound of the true $W_2$ distance.

We discuss the rotated $W_2$ distance in more detail in Appendix D of the revised manuscript 
as well as in the response to Comment 8 below.

Reviewer comment (Theory):
Does the rotated squared $W_2$ verify the same upper and lower bounds as in Theroems 1 and 2?

Author response:
The rotated squared $W_2$ is a lower bound
of the true $W_2$ distance and aims to approximate the true $W_2$
distance for the empirical distributions on the time partition
$\left(t_i\right)$.  Consequently, upper bounds for $W_2$ distances
associated with the time partition $\left(t_i\right)$ hold for the
rotated squared $W_2$ distance as well.  We have further clarified the
definition of the rotated squared $W_2$ loss in Appendix D of the
revised manuscript.

Reviewer comment (Writing):
In the definition of the rotated Wasserstein, does $m$ the number of rotations? If yes, so in Tables 6 and 7 $N_{rotate}$ should be replaced by $m$.

Author response:
Yes, $m$ is the number of rotations. We
have revised the notation in Tables 6 and 7 (now Fig. S1 and Fig. S2)
accordingly.

Reviewer comment (Evaluation):
According to Figure 4, I would prefer the reconstruction of $f$ and $\sigma$ using Eq. 16 and Eq. 17 than the rotated $W_2$. Indeed, according to their definitions on Page 17, they are acting on the empirical observations of the process in the time partition $(t_i)$. However, the rotated $W_2$ is acting on an embedding of the origin process, which might induce some loss of data information.

Author response:
To clarify, the rotation itself does not introduce additional
loss of information, as the rotation is an isometric linear
transformation. From an abstract perspective, the true
high-dimensional dynamics is not changed by a change of coordinates.
From another perspective, the original SDE of $\boldsymbol{X}_t$ and
the rotated SDE of $\boldsymbol{Y}_t \coloneqq
\boldsymbol{R}\boldsymbol{X}_t$, where $\boldsymbol{R}$ is a rotation
matrix, are related by
\begin{equation}
\mathrm{d} {\boldsymbol{Y}} = \boldsymbol{R}
\boldsymbol{f}\big(\boldsymbol{R}^{-1} (\boldsymbol{Y})\big)
\mathrm{d} t + \boldsymbol{R} \boldsymbol{\sigma} \big(\boldsymbol{R}^{-1}
(\boldsymbol{Y})\big)\mathrm{d} \boldsymbol{W}_t.
\end{equation}
 Once knowing the
rotated SDE, we can recover the original SDE by applying the inverse
transform $\boldsymbol{R}^{-1}$.
Despite this, averaging over different rotations and projecting in
different directions can lead to some loss of information. 

Lastly, we plug the decoupled and rotated $W_2$ distances into Eq. 17
as approximations to the true $W_2$ distance at each time point. 
While these methods provide analytical approximations to
the true $W_2$ distance, our manuscript also explores numerical
optimization procedures for a numerical approximation of the $W_2$
distance. The purpose is to compare the performance and computational 
costs of these methods. We discuss these
aspects in further detail in Example 4 of the revised manuscript.

Reviewer comment (Presentation):
Tables 6 and 7 could be more readable if they were represented as plots.

Author response:
We have replaced Tables 6 and 7 with Figs. S1
and S2 in the revised manuscript.

Reviewer comment (Theory):
Since the bound (7) is not elegant, perhaps it could be simplified appropriately to preserve the point about convergence.

Author response:
In Appendix A, we provide a detailed proof of Theorem 1, outlining each
intermediate step and the Cauchy inequalities employed to derive the
error bound. While the current form of Eq. (7) may not be the most
elegant, it effectively communicates how the Wasserstein distance
$W_2(\mu, \hat{\mu})$ is bounded by multiples of the differences in
the drift and diffusion terms, $f-\hat{f}$ and $\sigma-\hat{\sigma}$,
between two SDEs. Eq. (7) serves to illustrate the necessity of
minimizing $W_2(\mu, \hat{\mu})$ for SDE reconstruction, aiming to
reduce the errors $f-\hat{f}$ and $\sigma-\hat{\sigma}$ between the
ground truth and reconstructed SDEs.

We acknowledge the potential for further refinement by imposing
tighter inequalities in Appendix A to enhance the error bound.  What
is more interesting and nontrivial is to extend the error bound
Eq. (7) to solutions between two multidimensional SDEs.  We are
currently working on multidimensional SDEs but the analysis is
nontrivial, as described in Remark 1.

Reviewer comment (Reproducibility):
Some training set details for simulations seem missing like what was the step-size, time-horizon, number of samples etc ?

Author response:
At the beginning of Section 3, we have written that training
details for all examples such as the step-size, the number of samples,
the learning rate, etc, are given in Appendix E. The time-horizon is
given separately in each example (Eqs. (18), (20), (21), (23)).

Reviewer comment (Novelty):
There seem to be prior works which perhaps present a more comprehensive analysis of SDE wrt. optimal transport e.g., [1*], [2*], [3*], [4*]. It would be insightful if these are discussed in detailed in related work and perhaps appropriately compared with. I feel this is a major weakness of the work.

Author response:
We appreciate your providing these
references.  We were aware of some of these and had cited Reference
[2] from your list in our manuscript. However, these previous works do
not directly address the specific task of approximating an SDE by
another SDE or SDE reconstruction. Moreover, the references you
provided primarily focus on the Wasserstein distance between solutions
to the same SDE with different initial conditions or other specific
scenarios, such as stochastic contraction and numerical solutions;
they do not directly address the SDE reconstruction problems that our
proposed method aims to solve. We summarize them here:

 (1) Reference [1] analyzes the Wasserstein distance between
  solutions to the same SDE with different initial conditions, which
  is called stochastic contraction. It is not directly applicable to
  approximating an SDE by another SDE.

 (2) Reference [2] studies the Wasserstein distance between solutions
  to the same L\'evy process \& potential jumps with different initial
  conditions as well as coupling. It is not directly applicable to
  approximating an SDE by another SDE or SDE reconstruction problems
  either.

 (3) Reference [3] investigates the Wasserstein distance between the
  solution to an SDE and its Euler scheme approximation.  Yet, the results
  cannot be directly applied 
  for SDE reconstruction, which requires analyzing the Wasserstein distance
  between solutions to two SDEs.

 (4) Reference [4] studies the Wasserstein distance between the
  solution to an ergodic SDE and its numerical solution. It mainly
  discusses how to appropriately obtain the numerical solutions to
  SDEs of some specific form.  This paper does not mention how to
  reconstruction SDEs.  Yet, the Langevin equations it analyzes could
  be helpful if we wish to apply the Wasserstein distance to construct
  processes other then SDEs. We have cited this paper in our revised
  literature review.

In our Appendix F, we also carried out a numerical experiment showing
how the uncertainty in the initial condition affects the
reconstruction of an SDE. Of course, it would be helpful to carry out
more theoretical analysis on this, as was done in the references you
kindly provided to us.

As you and Reviewer 4 suggest, there seems to be no simple method
so far that exploits the Wasserstein distance for SDE
reconstruction. Thus, our method is not only new but it also appears
to be quite efficient in applications we explored. The most recent and
related work that directly tackles general SDE reconstruction or using
one SDE to approximate another unknown SDE are the WGAN-SDE method and
its extensions (Neural SDEs as infinite-dimensional GANs) and the MMD
method (Generative moment matching networks). We compared our
Wasserstein-2-distance-based SDE reconstruction method with two
methods in Example 3 and found our method was more
efficient and accurate. Furthermore, minimization of the Wasserstein-distance loss
function (Eqs. 16, 17) used in our method is much easier to implement than
the aforementioned WGAN and MMD methods.

Reviewer comment (Theory):
What do the measures \mu,\hat{\mu} represent? It seems they are measures corresponding to X(t) and \hat{X}(t) . If so, should not there be a \mu(t), \hat(\mu}(t) ? In which case, they cannot be used in the LHS of 5 as the RHS is integrated over time. I think the RHS is a generalization of optimal transport to SDEs, but the optimal objective need not be a wasserstein distance (LHS). Am I correct ?

Author response:
Q1. The measures $\mu, \hat{\mu}$ are two probability measures associated with
the entire infinite dimensional temporal trajectories $\boldsymbol{X}(t), t\in[0, T]$ and $\hat{\boldsymbol{X}}(t),  t\in[0, T]$
rather than specific probabilities for $\boldsymbol{X}(t), \hat{\boldsymbol{X}}(t)$ at a fixed time point $t$. Therefore, we would not
use the notation $\mu(t), \hat{\mu}(t)$.
The distance between two trajectories is thus defined as the temporal integration $\lVert\boldsymbol{X}\rVert\coloneqq
\big(\int_0^T \sum_{i=1}^d |X_i(t)|^2\text{d} t\big)^{\frac{1}{2}}$. In this context, we can define 
the Wasserstein-2 distance between the two measures $W_2(\mu, \hat{\mu})$ as Eq. (5).

Reviewer comment (Experiments):
Why are different sets of baselines used in the various simulation examples ?

Author response:
Q2. In Examples 1 and 2, we compare our method with several classical uncertainty quantification (UQ) methods. 
We find that the Wasserstein distance could indeed be better in reconstructing SDEs' trajectories and drift as well as diffusion terms than those classical UQ loss functions in statistics.
In Examples 3 and 4, for simplicity, we do not compare with those classical statistical UQ methods again. Instead, we focus on comparing with two most recent machine-learning-based SDE reconstruction methods (WGAN and MMD). We find that our method could also outperform these two methods in SDE reconstruction.