Latent SDEs on Homogeneous Spaces

Summary

In this paper the authors develop machinery for performing variational inference on latent functions in models where observations are generated by a latent stochastic process. That is, the generative model is that a path in some latent space is generated according to a prior, and then we observe (noisy) values of this path at a subset of the points. In principle this paper deals with a general case where the latent space is a "homogeneous space" which is a manifold such that there is a Lie group such that for any pair of points in the manifold one can translate one of the points to the other using a single element of the Lie group. In practice, the paper focuses on the specific case of the manifold being an n-sphere (so that the Lie group is the rotations, SO(n)), and the generative model is that the initial point is sampled uniformly from the sphere and then evolves according to a (scaled) Brownian motion on the sphere. Once the authors develop their inference machinery, they apply their method to a number of regression, classification, interpolation, and extrapolation problems, showing that this latent SDE formulation can be useful in some settings.

Strengths

The proposed method is very elegant and the mathematics is nice. Moving to SDEs on a compact latent space allows for many nice features like being able to define uninformative priors and having nice correspondences between the initial distribution and the stationary distribution of the SDE. The simplicity and performance on the considered tasks is strong motivation for the usefulness of the construction, and having access to posteriors over entire latent trajectories allows for a number of interesting tasks like interpolation.

Weaknesses

- The main weakness for me is the presentation of the material. I will try to provide more specific feedback in the following points, but overall there was a large disconnect between the setup in equations (1) and (2) and then the actual applications. Specifically, (2) is listed as the objective function, but then it's not clear what the objective function would be for the regression and classification tasks. The classification task presumably has a different generative model than the one presented between equations (1) and (2), and it's not clear where the class label would go, and what data is used where. What parts of the models are learnable? - The introduction is a bit confusing. For example, it is unclear while reading the introduction what is path-valued and what is not. Line 105 says everything is path-valued but then lines 93-94 say that we only have a finite number of discrete observations of the path. - I found the condition on $h$ in line 103 confusing. Is this at all a restriction? It seems like equation (1) automatically makes $X_t | Z_t$ normally distributed with mean $\mu_t = h(Z_t)$ and variance $\mathbf{R}$. - I defer to the authors on how they want to present it, but I was confused by the presentation of the ELBO around (2) because in general the KL divergence between distributions over functions determined by SDEs will be infinite unless the diffusion terms are the same. The authors address this thoroughly and clearly later in the manuscript, but sweep this (to me) important point under the rug in the introduction. - Equation (12) and Figure 2 make it clear that the inference network takes in the $\mathbf{x}$'s and then learns the hyperparameters of the initial distribution and the SDE. Is the inference network necessary because of the sample sizes considered here? Wouldn't it be possible to just directly optimize the set of $\mathbf{K}_i^\phi$'s that specify the SDE (and the parameters of the initialization) for each set of observations? Some comment on the distinction of what is being used for amortization (and why) vs. modeling flexibility would be useful. - The notation around equation (7) is very confusing to me. It appears that the time points $t_k$ represent the time that has passed since time $0$, but then it feels like $Z_{t_j}$ should be $G(t_j)G^{-1}(t_{j-1})Z_{t_{j-1}}$. Perhaps relatedly, it's not obvious to me what the connection is between $G(t_j)$ and the SDE in (4). In particular it feels like to solve the SDE for the interval $[t_{j-1}, t_j]$ you would need a different initial condition (the distribution over $G(t_{j-1})$ obtained by solving the SDE up to that point. Apologies if I'm being slowing here. - The example of extrapolating the rotated MNIST was a bit confusing. Doesn't the extrapolation rely exclusively on the Chebyshev polynomials behaving well after the end of the interval over which they are trained? Relatedly, how is the end time point chosen during the VI optimization and does it matter? The KL term seems sensitive to how long of an interval is considered. E.g., if we consider data sampled on [0, 1], but then compute the KL on paths from [0, T] with T>>1, we would want the drift terms to eventually relax back to matching the prior. - Lines 371-372 are confusing: what is meant by saying that only the initial time point is observed? The prior is driftless, so it feels like if one just observes the initial state it should be hard to get directional rotations using the latent SDE model. Typos: - "on a various time series interpolation" --> "on various time series interpolation" - "the paradigm of Parameterizing the vector fields" --> "the paradigm of parameterizing the vector fields" - Line 171: "in context of" --> "in the context of" - Line 186: "are show in" --> "are shown in" - Line 189: "This allows to select" --> "This allows selecting" - I believe the equation prior to line 192 is missing a square on the Frobenius norm term since the vector norm on the lefthand side is squared - Line 265: "with label switches need" --> "with label switches needing" - Line 358: "Upon receiving first" --> "Upon receiving the first" - I believe equation (35) has an erroneous $\mathbf{z}$ on the lefthand side (the righthand side is matrix valued)

Questions

The actual generative model is quite simple, and very little of it is learned. In particular, the prior distribution has only a single learnable parameter. Presumably most of the modeling flexibility comes from learning the mapping from the latent space to the observable data (i.e., $p(\mathbf{x}(t) | \mathbf{z}(t))$). Is there a reason for not learning a drift term for the prior? Would one run into identifiability issues with more complex priors? How difficult would it be to extend the implementation and framework to other homogeneous spaces? Are there any applications where such a generalization would be obviously useful?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

1 poor

Contribution

3 good

Limitations

The authors have adequately addressed potential social implications.

Summary

This paper deals with the problem of variational inference for sequential data, i.e. time series, using latent SDEs. The idea of this class of methods is to assume the stochastic process that is observed is related to a generative SDE in a latent space, whose parameters need to be inferred from data in a Bayesian fashion. The novelty of this paper is to consider SDE priors that live on homogeneous spaces, and in particular on the sphere (acted upon transitively by matrix multiplication by $SO(n)$). Many works have taken the path of having either a spherical or sequential latent space recently, but have never considered latent SDEs on the sphere as priors for variational inference of the posterior process. This is made possible by designing an uniformative sequential prior process on the sphere (with a tractable KL divergence) and a parametric class of posterior SDEs that provably lead to solutions on the sphere, and using an associated geometric Euler Maruyama scheme for the numerical integration of the SDE. The model is then optimized using an ELBO loss that is adapted to sequential variables and the SDE at hand, using the Power Spherical law as a reparametrization-friendly prior distribution on initial condition, and furthermore fitting the drift of the latent SDE. The method is then tested on a number of datasets and machine learning tasks, for which having both a sequential and spherical prior may or may not be a natural choice. Results indicate that the proposed method is at least comparable to other SOTA approaches on those problems on the tasks adressed.

Strengths

Strengths : - The theory of ODEs on Lie Groups and the associated geometric integration schemes is nicely put to use, and is potentially applicable to any homogeneous space. It would be quite nice to try and extend all this work to hyperbolic spaces, which have been used a lot in ML recently, and happen to also be homogeneous spaces. - I found the paper well written and enjoyable to read. - The example of the sphere is compelling and leads to the specification of stochastic processes defined through SDEs on the sphere both for the prior and posterior. The tools introduced lead to a new prior distribution for stochastic processes on the sphere, and tractable posterior distributions shown to be usable in a variational inference context. - An extensive experimental study is carried out on different standard datasets. The proposed framework is applicable to a number of different tasks involving time series : classification (of each time step or of the global sequence), interpolation, regression. The method proves competitive with respect to SOTA approaches for each task.

Weaknesses

Weaknesses : - As often in works involving spherical, or more generally hyperbolic latent space, the motivation of using them is not always crystal clear for all applications. For data involving some kind of periodicity, as most of those tested here (pendulum, rotating MNIST), this makes sense, but for others it is not so obvious why one would want to use spherical latent spaces. Would this method work for e.g. time series prediction of chaotic data, e.g. the Lorenz system ? However this is a minor remark. Once the need for a spherical latent space has been established, the proposed work extends successfully variational inference for sequential data to this setting. - I find it a bit disappointing not to find an application of the proposed method to uncertainty quantification, or at least one where the capacity to sample from the posterior process is not put to use. To me, this is the main interest of working with latent SDEs instead of ODEs : one gets a sequential generative model capable of sampling, computing expectations, covariance or other statistics. This could be showcased e.g. through a simple experiment of forecasting time series, showing predictions and confidence intervals.

Questions

- Are the error bars reported in the tables computed though multiple training sessions or simply several samples of the posterior, for a given trained model ? - On a related note, I am not sure how the model is adapted for tasks such as per-frame classification. I get that the purpose of the h function is to potentially map to another space, e.g. using a softmax for classification problems, and I imagine the second term of the ELBO is the one that needs to be adapted depending on the model. It could be useful to detail one example (e.g. classification) and explaining what eq. (2) becomes in that case. Could h be learned jointly with the posterior instead of being assumed to be known ? - In figure 1, I would have been interested in seeing a few trajectories of posterior samples that do not have a constant label, to see if the proposed intuition for the trajectories still holds in that case. - How does the neural net used to fit the parameters of the initial condition given by the spherical power law constrain them to be positive for the concentration parameter and on the sphere for the location parameter ? - Note that Figure 2 does not display properly on all the different pdf readers I tried. For one of them, I cannot see the boxes and arrows, but just the text, which makes the figure hard to read.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

- One of the main limitations of the work is the lack of motivation of using SDEs/generative models instead of simply ODEs/deterministic models, if the sampling capacity of the model is never really put to use, while in many domains involving time series uncertainty quantification is a key scientific issue (see above for detailed comments on this point). - Another minor limitation is the limitation of a geometric Euler Maruyama solver for the SDE, while more accurate solvers (of Runge Kutta Munthe Kaas type) exist. Those are not used so as not to end up with a too much computational burden, but I wonder how much of a problem this would be. More generally, I also would like to have an idea of the required computational time/complexity of the proposed model, and how much it would depend on the used solver.

Summary

The authors are interested in learning neural SDE models. Instead of parameterising arbitrary latent SDEs, the authors restrict their attention to homogeneous spaces, and in particular the unit sphere, in order that they can leverage the transitive group (the Lie group) to construct an SDE in the space in terms of an SDE in the Lie group, whose logarithm is a linear SDE, which leads to convenient solutions. The major advantage of this is using a straightforward discretise-then-optimise approach without an explosion of computational cost. They perform a thorough evaluation against other neural ODE/SDE methods, and show competitive performance despite the more restrictive form of the latent SDE. Post-discussion: The major concern has been addressed. The authors have pledged to extend their evaluation of the efficiency of their method to other experiments in the paper, which will be sufficient to address that part.

Strengths

1. (major) To my knowledge the approach is novel and markedly different to other latent SDE methods, but is applicable to many of the same tasks, so is highly relevant. 2. (major) The evaluation is generally very thorough (though see weaknesses), with a wide range of related methods and different tasks evaluated. 3. (major) The code is available, which improves reproducibility, and is based on a widely used framework (pytorch) which improves possible impact. 4. (minor) The paper is quite clearly written, including good discussion of the empirical findings.

Weaknesses

1. (major) The method is claimed to be efficient for learning, but no evidence is provided to support this claim -- the experiments as they are demonstrate that the method can produce a test performance which is competitive with apparently more flexible approaches, and that the relative performances may depend somewhat on the task. For the reader to judge how far this method is more efficient, a quantitative evaluation of the time for learning for different models would help.

Questions

typos: * lines 25, 58: Paramaterizing -> parameterizing * line 178: extra $≤$ * line 180: arithmetics -> arithmetic

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

Limmitations are well discussed explicitly in the paper.

Summary

The paper provides an affirmative answer to a very natural and intriguing question: Can we simplify the underlying latent model describing the dynamics of a temporal process so that it can overcome the computational and technical challenges with neural ODEs/SDEs while accurately modeling the real-world phenomenon? The authors propose SDE models that arise from the action of a Lie group on a homogeneous space instead of an arbitrary SDE. They also demonstrate at-par performance of their approach on benchmark regression, classification and interpolation problems when compared to the existing methods.

Strengths

This work is really interesting as it opens a new direction of research that could reduce the model's degrees of freedom and yet achieve (or nearly) the sota. It may encourage machine learners to leverage recent developments in the SDE literature and further simplify learning a time series phenomenon. The paper is well written.

Weaknesses

1. The paper claims in the introduction that their approach significantly reduced computing efforts. “ However, .......... computing gradients.” However, there is no comparison with other approaches I could find in the numerical section. Also, a mathematical discussion on the reason behind computational gain is absent. 2. I think the paper would benefit from adding more discussion on the statement below in the main paper. “Numerical solutions to such an SDE are computed with a simple one-step geometric Euler-Maruyama scheme for which the “discretize-then-optimize” strategy of backpropagating gradients during learning is not a limiting factor."

Questions

1. Why do we need such a specific SDE form for G_t in the display (4)? 2. Why do we not need a reparametrization trick for the proposed approach? I found no discussion on this in the paper (except a standard comment in line 111). Minor. You followed [47] (in the paper)or [48](in the appendix) for preprocessing the human activity dataset. Which one did you follow?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The paper includes a discussion on the limitations and possible societal impact.