SING: SDE Inference via Natural Gradients
Authors: Amber Hu
Henry Smith
Scott Linderman
Venue: NeurIPS 2025 poster ()
Keywords: latent SDEs
dynamical systems
variational inference
natural gradient
Primary area: probabilistic_methods
TL;DR: We propose SING, a method that leverages natural gradient variational inference for accurate and efficient inference and drift estimation in latent SDE models.

Abstract: Latent stochastic differential equation (SDE) models are important tools for the unsupervised discovery of dynamical systems from data, with applications ranging from engineering to neuroscience. In these complex domains, exact posterior inference of the latent state path is typically intractable, motivating the use of approximate methods such as variational inference (VI). However, existing VI methods for inference in latent SDEs often suffer from slow convergence and numerical instability. Here, we propose SDE Inference via Natural Gradients (SING), a method that leverages natural gradient VI to efficiently exploit the underlying geometry of the model and variational posterior. SING enables fast and reliable inference in latent SDE models by approximating intractable integrals and parallelizing computations in time. We provide theoretical guarantees that SING will approximately optimize the intractable, continuous-time objective of interest. Moreover, we demonstrate that better state inference enables more accurate estimation of nonlinear drift functions using, for example, Gaussian process SDE models. SING outperforms prior methods in state inference and drift estimation on a variety of datasets, including a challenging application to modeling neural dynamics in freely behaving animals. Altogether, our results illustrate the potential of SING as a tool for accurate inference in complex dynamical systems, especially those characterized by limited prior knowledge and non-conjugate structure.

SING: SDE Inference via Natural Gradients
Amber Hu Stanford University amberhu@stanford edu
Henry D. Smith Stanford University smi mmithhd@stann tanf ord. edu
Scott W. Linderman Stanford University swl 1@s @stanford...... edu
Abstract
Latent stochastic differential equation (SDE) models are important tools for the unsupervised discovery of dynamical systems from data, with applications ranging from engineering to neuroscience. In these complex domains, exact posterior inference of the latent state path is typically intractable, motivating the use of approximate methods such as variational inference (VI). However, existing VI methods for inference in latent SDEs often suffer from slow convergence and numerical instability. Here, we propose SDE Inference via Natural Gradients (SING), a method that leverages natural gradient to efficiently exploit the un- derlying geometry of the model and variational posterior. SING enables fast and reliable inference in latent SDE models by approximating intractable integrals and parallelizing computations in time. We provide theoretical guarantees that SING approximately optimizes the intractable, continuous-ting objective inter- est. Moreover, we demonstrate that better state inference enables more accurate estimation of nonlinear drift functions using, for example, Gaussian process SDE models. SING outperforms prior methods in state inference and drift estimation on a variety of datasets, including challenging application to modeling neural dynamics in freely behaving animals. Altogether, results illustrate the potential SING as a tool for accurate inference in complex dynamical systems, especially those characterized by limited prior knowledge and non-conjugate structure."
Introduction
Stochastic differential equations (SDEs) are powerful tools for modeling complex, time-varying systems in scientific domains such as physics, engineering, and neuroscience [1, 2]. In many real- world settings, do not observe the state of the system directly, but instead must infer it from noisy measurements. For example, in neuroscience, we often wish to uncover continuously evolving internal brain states underlying ditensiinenii neural recordings [3]. Latent SDE models address this problem by positing a continuous-time stochastic latent process that generates observations through a measurement model.
key challenge in latent SDE modeling is posterior inference over the latent state trajectory. Accurate inference enables principled understanding and prediction of the underlying state, but it is often computationally intractable to perform analytically due to non-conjugate model structure. Methods for approximating the intractable posterior include Markov chain Monte Carlo [4, 5]. approximate Bayesian smoothing [2. 6], and variational inference (VI) [7-11]. J. In particular, VI offers a flexible and scalable framework for approximate inference in latent SDE models. In their seminal work, Archambeau et al. [7, 8] formulate the VI problem for latent SDEs as constrained optimization over the family of linear, varying SDEs. While variants of this approach have been widely employed [9, 12-15), they often suffer from slow convergence and numerical instability. These shortcomings motivate the need for more efficient and robust VI methods in this setting.
"Equal contribution
'Code: http://ttttcccmmm///neeeaaaaaaaa
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

In this work, we introduce SING: SDE Inference via Natural Gradients, method that leverages natural gradient variational inference (NGVI) for latent SDE models. SING exploits the underlying geometry of the prior process and variational posterior, which leads to fast and stable latent state inference, even in challenging models with highly nonlinear dynamics and/or non-conjugate observation models. Our contributions are as follows:
(i) We derive natural gradient updates for latent SDE models and introduce methods for computing them efficiently;
(ii) We show how SING can be parallelized over time, enabling scalability to long sequences;
(iii) We provide theoretical guarantees that SING approximately optimizes the continuous-ting objective of interest; and
(iv) We demonstrate that improved inference with SING enables more accurate drift estimation, including in models with Gaussian process drifts.
2 Background and problem statement
2.1 Latent stochastic differential equation models
We consider the class of models in which the latent states evolve according a stochastic differential equation (SDE) and give rise to conditionally independent observations. 2 E R takes the form
The prior the latent states
p(æ)= dt((() - ((=(()) + Σ dwo(),
x00) ~ ((,,
0 0 Tmax
where f RDDD - RR is the drift function describing the system dynamics, Σ is a time-homogeneous a standard D D-dimennsional Brownian motion. and V
noise covariance matrix, and (())oss).T.... are the mean and covariance of the initial state. For our exposition, we will model f as deterministic drift function with parameters 0. We will subsequently expand our scope Section 3.5 to the case of a Gaussian process prior on f. For simplicity of presentation, we assume f is time-hoc mome-homgeeneous: however, our framework can readily extended to the inhomogeneouus case as well.
While the latent trajectory (a=()) is unobserved, we observe data = (t,, (t,)) where t; E [0, m are observation times and (t+) E IR V are corresponding observations modeled as E (ty( x] = g (Cx(tt, + d).
Here, g(-) is a pre-determined inverse link function and {C, d} are learnable hyperparameters defining an affine mapping from latent to observed space. Of particular interest is the setting D < N, where the latent space is much lower-dimensional than the observations. We provide a schematic of the generative model in Figure 1A.
For the latent SDE model, we aim to solve two problems: (i) inferring a posterior over the latent trajectories conditional on the data D and (ii) learning the prior and output hyperparameters, collectively denoted by θ = (0, d.
2.2 Natural gradient variational inference for exponential family models
SING leverages natural gradient variational inference (NGVI) to perform approximate posterior inference in latent SDE models. Before presenting our method, we first provide relevant background on NGVI for exponential families. In this section, we will consider the general VI problem in which we have a variational posterior q(= (z)))) with parameters 17 E TR.. The goal of VI to maximize lower bound L(n) to the marginal log likelihood [16, 17].
One way to maximize L(n) is via natural gradient ascent [18] with respect to n, which adjusts the standard gradient update to account for the geometry of the distribution space. In practice, this corresponds to the update rule
1 1 (n n1) (++1) arg min 17 V (())
F(n T
())
(3)
p
DDu(((<)) ((=1))
- > (++1) n7(
(4)

---

Generative model
Natural gradient VI H1R
C
Variational natural parameters
Hioooall g(Cr0))
dosnaaaion din
h
gradient
ft
Jo
L
h;
Lo
J1
literddinn
1-
precondition inverse Fisher
hr L--1
Jr
Sttent
dxt() = ((e (WW
parameserr
Figure I: A schematic figure describing SING. A: In the generative model, low-dimensional latent SDE gives rise to noisy, conditionally independent observations at timestamps {t; B: SING leverages NGVI perform fast and reliable approximate posterior inference in latent SDE models. NGVI exploits the geometry of the model by effectively preconditioning updates by an inverse Fisher information matrix, often leading faster convergence than vanilla gradient ascent. C: On discretized time grid т, the variational posterior is a multivariate Gaussian distribution with a block tridiagonal precision matrix.
where m7()) denotes the variational parameter at iteration j. p is the step size, and F(n) Eggas/n) V,, log ((=nn) log 7(2(m) is the Fisher information matrix. The key idea of natu- ral gradient methods [17-21] is to perform an update that both increases the ELBO and remains close the current variational distribution (in KL divergence). This interpretation relies upon second-order Taylor expansion of the KL divergence, as denoted in eq. (3); we provide a derivation in Appendix B.I. See Figure IB for a geometric interpretation of natural gradient ascent.
However, the inverse Fisher information in the update eq. (4) is intractable to compute when r is high-dimensional. This is common in models with temporal structure, including the latent SDE model presented in Section 2.1. The key insight behind NGVI [22-27] is that eq. (4) can be computed efficiently when ((2(n) constitutes minimal exponential family:
((2|n) 1xz)) exp (n,, (())) - A(r)), A(n) log (f
(fb())exp M[z))
((n, T(z))) dz
(5)
Here, 17 are called the natural parameters of the exponential family, T22) are the sufficient statistics, μ Eg(=im) [=))] are the mean parameters, and A(n) the log normalizer. The exponential family density eq. (5) is defined for all v for which A(n)
We refer the reader to Appendix B for an
overview of exponential families. Let u(()) := in [T())] ]] be the bijective mapping from natural mean parameters. In Appendix B, we show that the Fisher information is equal to the Jacobian of the mapping (u()) i.e., F(n) = (n)((n). As result, eq. (4) can be rewritten as
n7 (++1)
()) + pV ())
This reparametrization avoids need to compute or invert large Fisher information matrix, allowing for tractable updates even for high-dimensional parameter spaces.
2.3 Variational inference in latent SDE models
In the latent SDE model in Section 1, one show that the true posterior p(=x(D) is time-varying SDE with the same diffusion coefficient as the prior:
p( D) (t = f(x(tt, + dw) (0)) PD. 0 osts Tmax-
The drift of the posterior SDE is the solution linear partial differential equation (PDE) [28-31]. In the case where the prior (1) is a linear SDE and the observation model eq. (2) is linear and Gaussian, both ff and the marginal density of the posterior (()(t) can be computed via Kalman- Bucy smoother [32]. However, in non-conjugate settings, neither quantity admits closed-form, and solving for them numerically using PDE solvers can be computationally expensive.
We propose VI method for approximate inference the latent trajectories (#=()) As in prior works [7, 8], we posit a variational family ((x) of linear varevngrying SDEs,
q(x) dxx(t) (t(t)æ(t) + b(t)) + ddw x(0) (m 0 5 t < Tiax. (7)
3

---

where the variational parameters are (A(t))os and (b(t))osisTi In other words, we approxi- mate the posterior drift f with a linear function fz(x((),t) == A(t)x(t) + b(t) at each time.
In the setting where we have multiple independent trials & sampled from the generative model (ax*9) (1)) (1) {y (2) (4,))== the true posterior factorizes across trials. Hence, we choose our variational approximation to also factorize across trials. In other words, we approximate the posterior for each trial p(x(() (D) with a linear, time-varying SDE as in eq. (7), where the variational parameters allowed to differ across trials. In practice, find that this variational family is highly flexible (see Section 5) and is capable of accurately inferring latents sampled according to SDEs with highly nonlinear (e.g., bifurcating) dynamics.
Our choice of variational family (7) yields the following evidence lower bound (ELBO) to the marginal log likelihood:
Lcom (g, θ) = F.g [lo ply) æx θ)] DRI. (g(ar) p(z)
(8)
Our goal is maximize Leon )) with respect to q. given O fixed. However, this objective is typically intractable to optimize directly. In practice, common approach is to instead perform inference over a finite-dimensional distribution (( where æ; := xxT,) and T (T) S
[0, m,a], To 0, TT = max is set of time points that includes observation times {t:
When the prior SDE p(æ)) is nonlinear, the transition distribution p(æo:r) does not have closed-form expression. Therefore, we approximate p with its Euler-Maruyama discretization, which we call p. Since both p and y have Gaussian transition densities, this leads to a tractable approximation the continuous time ELBO Leaat (q. θ),
Lappoox (1, θ :- F,y
x. θ)] DRr. (1(#0))) |l p(ao:r ()).
(9)
For derivations of the ELBOs Leout and Lapprox. see Appendix C.
3 SING: SDE Inference via Natural Gradients
3.1 NGVI for latent SDE models
We propose SING, a method that uses NGVI to optimize the ELBO Lappox in the latent SDE model. SING can be applied to any combination of prior drift and likelihood model and enjoys fast convergence by exploiting the geometry of the natural parameter space. SING builds upon recent work from Verma et al. [33]; however, it has several notable differences that we detail in Appendix I.
We begin by recognizing that since eq. (7) is Gaussian and Markovian, its marginalized distribution ((xor)) can be written as a minimal exponential family,
(( = exp
1
T Σ i=0
J.C: σ
Σ 1=0
T-1
x L.E;
Σ h x. (n[=:-A(m)
T A(n)
(10)
2=0
In eq. (10), the natural parameters are = [1: J.
1= { L ) and the corresponding
sufficient statisties are T(xor)) = {z, u. 1T o. x1 T-
Since q(xo:)) is Markovian,
its full precision matrix is block tridiagonal, as depicted in Figure IC.
To apply the NGVI update eq. (6) to Lapprox: must first consider g( in terms of its mean parameters, which are simply the expectations of T(#or) under We denote these together as y == == [M,,, 24.2 (14,3) From here, eq. (6) simplifies to the following updates on the natural parameters, which decompose over time steps O,...,T:
(h, (115-1) 1(6+1) =(1 (())) ,("') + PV (1)) Be [log Mogr(ys)en (z.;) )6; log p( ar. (0,10;11) 1)] 1(++1)= = -11-)) + P
mi. PVp.. B [log ((
(11)
In eq. (11), q0)) := q( o(xuurly ( is the variational posterior at iteration j and 5, is an indicator denoting whether there is an observation at time T,. We provide a derivation in Appendix D.
(i))
As written, the updates in eq. (11) are challenging to compute because they involve (i) approximating intractable expectations and (ii) converting from natural to mean parameters between iterations. In the subsequent sections, we propose computational methods to address each of these problems.

---

3.2 ELBO approximation guarantee
Prior to detailing how we compute the SING updates eq. (11), we provide theoretical justification that maximizing the ELBO Luppeox will approximately maximize the true, continuous-timg ELBO of interest Centt Moreover, on subset of the variational parameter space, the approximation error tends to zero uniformly as the size of the grid tends to zero.
Theorem 1: Under mild regularity conditions (see Appendix E), there exists constant C such that for any mesh size At max A;, A, T T; of the time grid, |Ccent - Cuppox S C(At) 1/2
Proof sketch. One compares the continuous- and discrete time log-densities and applies error bounds on the Euler-Maruyama approximation. The full statement and proof can be found in Appendix E.
To the best of our knowledge, our work is the first in the topic of variational inference for latent SDE models to provide an explicit rate of convergence for the discrete time ELBO to its continuous counterpart.
3.3 Computing expectations of prior transition densities
The updates in eq. (11) rely on computing expectations of the form E,y [gllog ((0 difficult in general because the prior drift f(x(t)) can be nonlinear and the expectation is taken over the pairwise distribution g 2441), which has total dimension R2D In the following proposition, show using Stein's 's lemma that integration dimensionality to RP. The proof can be found in Appendix F.
[z:)]. This is
rewrite the expectation over just q(x)), which reduces the
Proposition 1: The term log p ( [a:)]] can equivalently be written in terms of the mean parameters p4g M4+1 and the expectations Egg [f(x)) Bg [f(x;) Σ - f=x)(, and E, [Jy (x;)), where J (a.,) denotes the Jacobian of f-) at 2c;-
The key implication of Proposition 1 is that the difficult expectations over the nonlinear prior transition terms can be reduced in dimensionality from IR22 to RD. This makes these terms in SING's updates significantly more tractable. In practice, there are many methods to compute these expectations, including Gauss-Hermite quadrature don-dimensional systems or Monte Carlo in higher-dimensional systems. Empirically, observe that even with single Monte Carlo sample, SING performs accurate approximate inference over the latent trajectories in up to 50-dimensional latent spaces (Section 5.1).
3.4 Parallelizing SING
In many real-world settings, sequential inference on long time series can be computationally chal- lenging. In this section, we show how SING can be fully parallelized, yielding efficient inference that scales logarithmically with the number of time steps.
The computational bottleneck of Section 3.1 is converting from natural parameters v to mean parameters μu for the updates in eq. (11). By the identity A()) μ., we can perform this conversion by first computing A(7)) and then applying autodifferentiation with respect to 17. However, existing algorithms for this problem take ODD)T) time (see Appendix G. for further details). While D is typically small in our setting, the linear scaling in T may be prohibitive for long sequences.
Our key insight is that the log normalizer A(n) of eq. (10) be computed in parallel time by using associative scans [34]. Given a binary associative operator and a sequence (a1, Q2,..., a,), the an)) in logarithmic time
associative scan computes the prefix sum (ai, (a) d (a) 02 using a divide-anddconquer approach. By recognizing the log normalizer as
A(n) = log / exp (-)
2o Jox + 1o Φo II exp (-2x) Jax x 1;-18:11 11
dxwe (12)
a 1.0
we can apply the associative scan to compute A(π). We define the unnormalized Gaussian potentials in eq. (12) as the sequence elements 0;-1,i and marginalize out variable at time via the binary associative operator Gz,j aj.k = ot.jaj, == a; To marginalize out the variable Tr. we also define the dummy element 0r = 1. As result, the log of the full scan (a-1.0 ag,1 e0T,T41 evaluates exactly to A(n). This algorithm has time complexity O log T), resulting in significant
5

---

speedups as sequence length increases (Figure 5). In Appendix G.2, we derive the analytical expression for the Gaussian marginalization operator and analyze our algorithm's 's time complexity.
3.5 SING-GP: Drift estimation with Gaussian process priors
So far, we have assumed that the drift function f (a)(t) @) is deterministic with learnable hyperpa- rameters O. However, in practice, we often have limited prior knowledge about the functional form of the drift but would still like to learn or infer its structure. A natural approach for this problem is to place Gaussian process (GP) [35] prior over the drift,
fa(-) Fd (o((-)))(-.))
dord=1.... D.
where f()) = (f........... fD())]] (::) is a kernel function, and are kernel hyperparameters. This results in a well-studied model called the GP-SDE [13, 14], which is appealing for its ability to encode rich prior structure in dynamics and provide posterior uncertainty estimates. However, the practical utility of GP-SDEs has been limited by the inference methods used to fit them [7, 8), which we show in Section 5.2 often exhibit slow convergence and numerical instability.
Here, we extend SING to enable accurate and reliable inference in GP-SDE models, resulting in a new method called SING-GP. While we defer technical details to Appendix K, we outline the key ideas here. Following Duncker et al. [13], we introduce an augmented variational posterior of the form q f) and use sparse GP techniques with inducing points [36] to perform approximate inference over q(f). This leads to convenient closed-form updates q(f). In Section 5.2, we illustrate that SING-GP performs competitively to flexible parametric drift functions in low-data regimes, while also directly providing posterior uncertainty estimates on the drift. In Section 5.4, we use SING-GP on a challenging application of modeling latent neural dynamics in freely moving animals.
Related work
VI for latent SDE models A number of VI methods for latent SDEs posit larger variational
families than eq. (7), consisting of non-Gaussian diffusion processes [10, 11, 37-41]. Ryder et al.
[10] parameterize the variational family the collection of diffusion processes with neural network drift and difffusion coefficients. Hereafter, we refer to the probability distributions belonging to this variational family as neural stochastic differential equations, 'neural-SDEs'. Li et al. [11] also consider a neural-SDE variational family and introduce a stochastic adjoint method to optimize the continuous- time ELBO. While memory efficient, the stochastic adjoint method requires sampling from the variational posterior at each gradient step. Most recently, Bartosh et al. [41] introduce a simulation-free variational inference algorithm, 'SDE Matching', for neural-SDE variational families. In Appendix L..2, we compare SING to SDE Matching on a Lorenz attractor benchmark from [1 I]. We find that, although SING learns Gaussian process approximation to the posterior, it outperforms SDE Matching on recovery of the ground truth latents and achieves competitive, albeit slightly worse, performance on drift recovery.
In settings where statistical and computational efficiency are critical, many works have studied Gaussian diffusion process variational families of the form eq. (7). In their seminal work, Archambeau et al. [7. 8] propose an algorithm 'VI for Diffusion Processes' (VDP), which optimizes the continuous- time ELBO in eq. (8) via Lagrange multipliers by incorporating the evolution of the marginal mean and covariance of (()) as constraints. A detailed overview of VDP is provided in Appendix H. Duncker al. [13] use sparse GP approximations to extend VDP to GP-SDE models, method we refer to as 'VDP-GP'. Related to VDP, Course and Nair [9, 12) express the continuous-time ELBOO entirely in terms of the marginal mean and covariance and then approximate these functions using a neural network encoder architecture.
Most relevant to our work, Verma et al. [33] propose applying NGVI to the discrete-time ELBO from eq. (9) in the case of a deterministic drift. Unlike SING, Verma et al. [33] do not address how to tractably compute Lapprox. and they do not suggest how efficiently convert from natural to mean parameters. Moreover, no theoretical guarantees are given as to how well Lappoox approximates Line- A complete comparison is given in Appendix I.
Connection to GP regression problems Several works have applied NGVI for efficient inference in GP regression problems by exploiting the fact that many GPs admit equivalent representations as

---

True iatents Budynanics
Observations
Inferred latents
Inference convergence
cbs indor
Gausion abc
())r
inferred
RMSE treets
1.0
VDP
0.5
Adam
-3
10
0.0
time
time
Iteration
H
8
truc
2-
Polsson counts
RMSE of Istents
SING
VDP
obc
CHGS
-3
Time
time
200
Figure 2: (Top row) We apply SING to LDS with Gaussian observations. A: True latents are sampled from LDS characterized by a stable spiral. B: Observations are 10-dimensional Gaussian variables. C: True vs. inferred latents on an example trial, with 95% posterior credible intervals. D: Comparison between SING and several baselines over iterations. (Bottom row) We apply SING to simulated place cell activity, where both the prior and observation models are nonlinear. E: True latents sampled from Van der Pol oscillator and represent trajectories through 2D space. Tuning curves modeled by radial basis functions represent expected firing rates at each location in latent space. F: Observations Poisson spike counts for neurons. G, H: See C, D.
linear time-invariant SDEs [2, 42]. For instance, Chang al. [43] and Dowling al. [44] consider GP regression problems with non-Gaussian likelihoods. To avoid cubic scaling the number of data points, they reformulate the GP linear SDE and then apply NGVI to perform approximate inference in this non-conjugate model. Hamelijnck et al. [45] combine this approach with sparse GP approximations to perform inference in ppatio-tempont setting. While these works focus linear SDE representations of GPs in regression problems, SING targets estimation of both the latent process and drift and applies to the broader class of nonlinear SDE priors.
Parallelizing sequential models Associative scans have been successfully applied to parallelize inherently sequential computations in areas such as classical Bayesian inference and deep sequence modeling [45-50]. Särkkä and FerenneFernnndez [46] and Hassan al. [47] derived associative scans for parallelizing Bayesian filtering and smoothing algorithms, and Hamelijnck et al. [45] extended these algorithms to inference non-conjugate GP regression. We note that our associative scan in Section 3.4 is reminiscent of the one in Hassan al. [47], which operates potentials for inference in the hidden Markov model. However, SING differs in that it leverages these techniques to compute the log normalizer in Gaussian Markov chain with observations, which enables efficient parameter conversion in its natural gradicnt-bart based updates.
Experiments
5.1 Inference on synthetic data
We begin by applying SING perform latent state inference in three synthetic datasets. Throughout this section, we assume that all hyperparameters θ are fixed to their true values during model fitting, allowing directly assess the quality of latent state inference of SING compared to several baseline methods. We provide full experimental details for this section in Appendix I..I.
Linear dynamics, Gaussian observations First, simulate 30 trials of 2D latent SDE with linear drift exhibiting a counter-clockwe ure 2A). Then, generate IOD Gaussian observations conditional on the latents (Figure 2B). We fit SING for 20 iterations and find that it is able to accurately recover the true latent states (Figure 2C).
stable spiral using Euler-Maruyama discretization (Fig-
We evaluate the inference quality of SING using root mean-squared еттоr (RMSE) between the true and inferred latent states over iterations (Figure 2D). We compare SING to three baselines: (1) Kalman smoothing (KS) [51], (2) VDP [7), and (3) direct optimization of Lappoox (1) using Adam [52]. Since this is a linear Gaussian model, exact posterior inference of g(*0:T) is tractable and can be obtained via KS. In Figure 2D, we find that SING recovers the true posterior in single iteration,

---

True latents dynamics
GP dhift post. mean & vanance
10 2.0
20.0
100omn PMEE
Anumic RMMM
dynmmes HMM
10-3
21
0.0015
DC15 de
SING
Iteration
---OPP
GP
Rersion
Noural
Palynomial basis
Figure 3: A comparison of SING and VDP for drift estimation. A: True latent trajectories evolve according to the 2D Duffing equation. B: Posterior mean (arrows) and variance (shading) correspond- ing GP prior drift with RBF kernel. SING-GP places high posterior uncertainty in regions unseen by the (true) latent trajectories. C, D: Latent RMSE and dynamics RMSE for SING and VDP across three classes of prior drift (GP, neural-SDE, polynomial basis). SING consistently outperforms VDP for hyperparameter learning, across all three drift classes. E: Dynamics RMSE for SING and VDP as function of the grid size At for the neural network drift.
consistent with the known result that NGVI performs one-step exact inference in conjugate models (see derivation in Appendix J). In contrast, VDP takes several iterations to converge in this setting; we provide theoretical justification for this finding in Appendix Finally, while Adam uses momentum and adaptive step sizes, it is outperformed by both SING and VDP in this experiment.
Place cell model Next, we simulate a dataset inspired by place cells in the hippocampus [53]. Place cells fire selectively when an animal occupies specific locations space, and are often modeled using -ependnt tuning curves. We simulate 30 trials of latent states from a Van der Pol oscillator to represent animal trajectories through 2D space. We then construct tuning curves for N 8 neurons using radial basis functions with centers placed along the latent trajectories (Figure 2E). Finally, we simulate spike counts for neuron n in time bin T; as yo (T:) ~ (I( (Figure 2F).
D = 50
40
30
10
-20
-10
This setup is particularly challenging due the non-conjugacy between the nonlinear dynamics and Poisson observations. Nev-
latents RMSE
D
ertheless, SING accurately infers the underlying latent states (Fig-
Monte Carlo
quadrature
ure 2G). We compare SING to several baselines, including VDP
3
0,1419
0.1417
and Adam on the natural parameters as in the previous experiment
5
0.1870
0.1869
(Figure 2H). Since KS does not apply in this non-conjugate setting,
10
0.2726
X
we instead use the discrete-time conditional moments Gaussian
20
0.3978
x
smoother (CMGS) [54] as a third baseline, which assumes linear the dynamics and observation models.
50
0.6778
x
Gaussian approximations
Altogether, we find that SING outperforms all baselines in both convergence speed and final accuracy. In contrast, VDP requires careful tuning of its learning rate for stability and CMGS relies on linear approximations that incur error in highly nonlinear models like this one. These results highlight SING's 's robustness in settings with nonlinear dynamics and observation models, while existing methods struggle.
Figure 4: (Top) Recovered
latent trajectories in the first dimensions of a 50-dimensional embedded Lorenz attractor. (Bottom) Comparison of latents RMSE between Monte Carlo- and quadrature-based approxi- mations of expectations. Monte Carlo results are averaged over 5 random seeds, with negligible standard errors (omitted).
Embedded Lorenz attractor We next demonstrate that SING can perform accurate inference in high-dimensional latent dynam- ical systems. To do so, we embed the Lorenz attractor into latent spaces of dimension D = 3, 5, 10, 20, 50 by sampling the first three coordinates according an SDE with Lorenz drift and the remaining dimensions according to random walk. For each D, we simulate 30 trials and generate 100-dimensional Gaussian observations. We fit SING for 1000 iterations using a single Monte Carlo sample per expectation, while the quadrature baseline uses 6 nodes per latent dimension (Section 3.3). Full experimental details can be found in Appendix 1.3.3.
Figure 4 (top) shows that SING accurately recovers the underlying attractor even in 50-dimensional latent space. The table in Figure 4 (bottom) further shows that Monte Carlo-based expectations match the accuracy of quadrature for low latent dimensions (D 5), while remaining tractable D
8

---

D
=
D
Wall eboc-time 50
-- paralee, trial sequential, trie
104
ver paralee.. 100 trais
100
sequental, 100 trias
64
4096
4096
512
4096
Number of timesteps
Figure 5: Runtime comparisons between parallelized SING and its sequential counterpart.
This demonstrates that Monte Carlo enables scalable and memory-effictt -efircceff inference and does not sacrifice accuracy in the low-dimensional regime where both methods are feasible.
increases.
5.2 Drift estimation on synthetic data
Next, we demonstrate that the fast, stable latent state inference enabled by SING results in improve- ments to drift estimation and parameter learning. We will seek to learn the drift parameters & together with the output parameters {C, d}. All experimental details can be found in Appendix L.3.
Our synthetic data example consists of 2D latent nonlinear dynamical system evolving according to the Duffing equation, which has two stable fixed points at (+v2, 0) in addition to an unstable fixed point at the origin (Figure 3A). We sample four trials starting near the rightmost fixed point with 10D Gaussian observations at 30% of the grid points. This example is challenging insofar as observations are sparse and the true latent trajectories do not cover the entire latent space.
We consider three classes of prior drift functions for our experiments: (i) Gaussian process drift with radial basis function (RBF) kernel (Section 3.5): (ii) neural network drift; and (iii) linear combination of polynomial basis functions to order three. We perform 50 variational EM iterations, alternating between performing inference and learning the hyperparameters O During the learning step, we optimize the drift parameters using Adam on the ELBO. The updates for the output parameters (C, d} can be performed in closed form (see Appendix L..).
Our results demonstrate that SING facilitates accurate and efficient parameter learning for diverse classes of prior drift. Although the GP prior constitutes nonparametric family of drift functions, performs comparably to flexible parametric drift classes, both in terms of latents RMSE (Figure 3C) and RMSE between the learned and true dynamics (Figure 3D). Unlike the parametric drift classes, though, SING-GP quantifies the uncertainty in the inferred dynamics about the leftmost fixed point (Figure 3B). Whereas SING recovers the true latents and dynamics in fewer than 25 iterations, VDP results in worse performance on both metrics in 50 vEM iterations. In Figure 3E, we observe that for both SING and VDP, decreasing At results in more faithful approximation the continuous-tmme ELBO Leont and hence more accurate recovery of the dynamics.
5.3 Runtime comparisons
Here, demonstrate the computational speedups enabled by using associative to parallelize the natural to mean parameter conversion SING (Section 3.4). Recall that while the standard sequential method of performing this conversion takes O(D²T), our parallelized version takes O O log T). In Figure 5, we compare the wall clock time of these two approaches for inference in a LDS with Gaussian observations. For different values of D and T, fit SING for 20 iterations over 5 randomly sampled datasets from this model and report the average runtime on an NVIDIA A100 GPU.
On single trials, parallel SING achieves nearly constant runtime scaling with sequence length, and roughly 100x faster than sequential SING for T 4096. This improvement holds across all tested values of D. Even for batches of 100 trials, where parallel resource constraints become more significant, parallel SING maintains favorable scaling compared to the sequential baseline. We provide experiment details and additional results over different batch sizes in Appendix L.4.
5.4 Application to modeling neural dynamics during aggression
Finally, we apply SING-GP to the challenging task of modeling latent neural dynamics of aggressive behavior in freely moving mice. Prior work identified approximate line attractors governing neural dynamics in this setting [55, 56). More recently, Hu et al. [14] fit GP-SDE models (see Section 3.5) with smooth, piecewise linear kernel uncover such dynamics. However, they relied on VDP-GP for inference, which demonstrated can be slow to converge and may do to suboptimal solutions.

---

A
Reconstr ucted activity
Inferred latents and dynamics
Trouus nO
AF
56
400
way mnd
time 15
300
D
time
41
0.1
4t 0.06
100
0.7
opeced
peScivve
50
0.5
4.6
0.3
0.06
00
30
15
30 15
5o000nds swm.boodddddddddooooooooodd forward
Figure 6: Results on modeling neural dynamics during aggression. A: Experimental setup, adapted from Vinograd et al. [55]. The data consists of calcium imaging from mouse interacting with two intruders. We model intruder effects using step function inputs. B: Reconstructed neural activity from SING-GP model fit. C: Inferred neural latents and dynamics from SING-GP, with high- probability slow point regions shown in purple. Comparison between SING-GP and VDP-GP on final expected log-likklinne liellikelioodd and convergence speed for different discretizations, Error bars are + 2SE over initializations of @ E: Comparison of forward simulation R² metric (see Appendix L.5 for details) assess the goodness of fit for latents and dynamics.
We revisit these analyses using a dataset from Vinograd et al. [55], which consists of calcium imaging of ventromedial hypothalamus neurons from a mouse interacting with two intruders (Figure 6A). We model this data using GP-SDEs that incorporate external inputs to capture intruder effects (see Appendix K.5 for more details), and we choose the same kernel from Hu et al. [14] to encode interpretable structure into the drift. We fit these models both using SING-GP and VDP-GP for 100 variational EM iterations. Full experimental details can be found in Appendix L.5.
Our SING-GP model captures variance in observed neural activity (Figure 6B) and uncovers low- dimensional representations consistent with prior findings (Figure 6C). Moreover, the GP-SDE enables posterior inference over slow points; we use this verify that SING-GP indeed finds an approximate line attractor (Figure 6C, purple). In Figure 6D, we assess the effect of discretization in each method on final expected log-likelihood and convergence speed. SING-GP is robust At in both metrics, while VDP-GP's performance suffers for large At values. This highlights a key advantage of the SING framework: robustness At enables accurate, memory-effiiien inference at coarse discretizations. Finally, we evaluate the quality of model fit via a forward simulation metric (Figure 6E). While SING-GP maintains high predictive R²2 for trajectories simulated up to 15 seconds forward, VDP-GP's predictive performance degrades more rapidly and is again sensitive to the choice of AL.
Discussion
We introduced SING, a method for approximate posterior inference in latent SDE models that har- nesses the exponential family structure of the variational posterior to achieve rapid, stable convergence. By deriving tractable NGVI updates and parallelizing the natural-to mean parameter conversion via associative scans, SING scales to long time series and handles nonlinear dynamics along with complex observation models with ease. Although we optimize a discrete-time ELBO, our theoretical bound (Theorem 1) ensures that SING closely approximates the continuous-time objective. We show via extensive experiments, including an application to real neural data, that SING outperforms prior methods in inference accuracy, drift estimation, and robustness to discretization.
However, SING makes simplifying assumptions that suggest several directions for future work. First, rather than fixing discretization T of [0, Tmaa] priori, one could consider adaptively learning locations at which to discretize inference in order to better approximate rapidly changing dynamics. Second, future work should explore incorporating learnable locations of inducing points in SING-GP to allow for more flexible and scale-aware GP priors on the drift. Nevertheless, our results already demonstrate SING's effectiveness in range of challenging practical settings. Altogether, SING is positioned to be reliable and broadly applicable tool for facilitating scientific discovery complex dynamical systems.
10

---

Acknowledgments and Disclosure of Funding
This work was supported by grants from the NIH BRAIN Initiative (U19NSI 13201, RO1NS131987. ROINSI 13119, & RFIMH133778), the NSF/NIH CRCNS Program (ROINS 130789), and the Sloan, Simons, and McKnight Foundations. Henry Smith supported by the NSF Graduate Research Fel- lowship (DGE-21
and the Lnight-Hennessy graduate fellowship. We thank Aditi Jha, Etaash Katiyar, Hannah Lee, and other members of the Linderman Lab for helpful feedback throughout this project. We also thank Adi Nair, Amit Vinograd, and David Anderson for their assistance with our analysis of the hypothalamic data [55]. The authors have no competing interests declare.

---

References
[1] Bernt Oksendal. Stochastic Differential Equations: An Introduction with Applications. Springer Science & Business Media, 2013.
[2] Simo Särkkä and Amn Solin. Applied Stochastic Differential Equations. Cambridge University Press, 2019.
[3] Saurabh Vyas, Matthew D Golub, David Sussillo, and Krishna V Shenoy. Computation through neural population dynamics. Annual Review of Neuroscience, 2020.
[4] Umberto Picchini. Inference for SDE models via approximate Bayesian computation. Journal of Computational and Graphical Statistics, 2014.
[5] Gareth O Roberts and Osnat Stramer. On inference for partially observed nonlinear diffusion models using the Metropolis- Hastinggs algorithm. Biometrika, 2001.
[6] Yu Ryan Yue, Daniel Simpson, Finn Lindgren, and Hävard Rue. Bayesian adaptive smoothing splines using stochastic differential equations. Bayesian Analysis, 2014.
[7] Cedric Archambeau, Dan Cornorrd, Manfred Opper, and John Shawe-Tayllo, Gaussian process approximations of stochastic differential equations. In Gaussian Processes in Practice. PMLR, 2007.
[8] Cédric Archambeau, Manfred Opper, Yuan Shen, Dan Cornford, and John Shawe-Taylee Variational inference for diffusion processes. Advances in Neural Information Processing Systems, 2007.
[9] Kevin Course and Prasanth Nair. Amortized reparametrization: efficient and scalable variational inference for latent SDEs. Advances in Neural Information Processing Systems, 2023.
[10] Tom Ryder, Andrew Golightly, A Stephen McGough, and Dennis Prangle. Black-box variational inference for stochastic differential equations. In International Conference Machine Learning. PMLR, 2018
[11] Xuechen Li. fing-Kam Leonard Wong, Ricky TQ Chen, and David Duvenaud. Scalable gradi- ents for stochastic differential equations. In International Conference on Artificial Intelligence and Statistics. PMLR, 2020.
[12] Kevin Course and Prasanth B Nair. State estimation of a physical system with unknown governing equations. Nature, 2023.
[13] Lea Duncker, Gergo Bohner, Julien Boussard, and Maneesh Sahani. Learning interpretable continuous-time models of latent stochastic dynamical systems. In International Conference on Machine Learning. PMLR, 2019.
[14] Amber Hu, David Zoltowski, Aditya Nair, David Anderson, Lea Duncker, and Scott Linderman. Modeling latent neural dynamics with Gaussian process switching linear dynamical systems. Advances in Neural Information Processing Systems, 2024.
[15] Lukas Köhs, Bastian Alt, and Heinz Koeppl. Variational inference for continuous -imeeuuuuuu--uueeeee-ee switching dynamical systems. Advances in Neural Information Processing Systems, 2021.
[16] Michael Jordan, Zoubin Ghahramani, Tommi S Jaakkola, and Lawrence K Saul. An introduc- tion to variational methods for graphical models. Machine Learning, 1999.
[17] David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians. Journal of the American Statistical Association, 2017.
[18] Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural Computation, 1998.
[19] Yann Ollivier, Ludovic Arnold, Anne Auger, and Nikolaus Hansen. coammation-geomeeti optimization algorithms: unifying picture via invariance principles. Journal of Machine Learning Research, 2017.
12

---

[20] James Martens, New insights and perspectives on the natural gradient method. Journal of Machine Learning Research, 2020.
[21] Mohammad Emtiyaz Khan and Hävard Rue. The Bayesian learning rule. Journal of Machine Learning Research, 2023.
[22] Matthew D Hoffman, David M Blei, Chong Wang, and John Paisley. Stochastic variational inference. Journal of Machine Learning Research, 2013.
[23] Mohammad Khan and Wu Lin. Conjugate-computationn variational inference: Converting variational inference in non-conjugate models to inferences in conjugate models. In Artificial Intelligence and Statistics. PMLR, 2017.
[24] Mohammad Emtiyaz Khan and Didrik Nielsen. Fast yet simple natural-gradien descent for variational inference in complex models. In International Symposium Information Theory and Its Applications (ISITA). IEEE, 2018.
[25] Hugh Salimbeni, Stefanos Eleftheriadis, and James Hensman. Natural gradients in practice: Non-conjugate variational inference in Gaussian process models. In International Conference on Artificial Intelligence and Staristics. PMLR, 2018.
[26] Wu Lin, Mohammad Emtiyaz Khan, and Mark Schmidt. Fast and simple natural-gradient variational inference with mixture of exponential-fane family approximations. In International Conference on Machine Learning. PMLR, 2019.
[27] Matt Jones, Peter Chang, and Kevin Murphy. Bayesian online natural gradient (BONG). Advances in Neural Information Processing Systems, 2024.
[28] Cedric Archambeau and Manfred Opper. Approximate inference for continuous-time Markov processes. Bayesian Time Series Models, 2011.
[29] Harold J Kushner. On the differential equations satisfied by conditional probability densities of Markov processes, with applications. Journal of the Society for Industrial and Applied Mathematics, Series A: Control, 1964.
[30] Ruslan Leont' evich Stratonovich. Conditional Markov processes. In Non-linear transformations of stochastic processes. Elsevier, 1965.
[31] Etienne Pardoux. Équations du filtrage non linéaire de la prédiction et du lissage. Stochastics, 1982.
[32] Rudolph E Kalman and Richard S Bucy. New results in linear filtering and prediction theory. Journal of Basic Engineering, 1961.
[33] Prakhar Verma, Vincent Adam, and Arno Solin. Variational gaussian process diffusion processes. In International Conference on Artificial Intelligence and Statistics. PMLR, 2024.
[34] Guy E Blelloch. Prefix sums and their applications. Technical report, School of Computer Science, Carnegie Mellon University, 1990.
[35] Christopher KI Williams and Carl Edward Rasmussen. Gaussian processes machine learning. MIT Press, 2006.
[36] Michalis Titsias. Variational learning of inducing variables in sparse gaussian processes. In Artificial Intelligence and Statistics. PMLR, 2009.
[37] Christian Wildner and Heinz Koeppl. Moment- variational inference for stochastic differential equations. In International Conference on Artificial Intelligence and Staristics. PMLR, 2021.
[38] Ali Hasan, Joao M Pereira, Sina Farsiu, and Vahid Tarokh. Identifying latent stochastic differential equations. Transactions on Signal Processing, 2021.
[39] Benjie Wang, Joel Jennings, and Wenbo Gong. Neural structure learning with stochastic differential equations. In International Conference on Learning Representationss 2023.
13

---

[40] Sebastian Zeng, Florian Graf, and Roland Kwitt. Latent SDEs on homogeneous spaces. Advances in Neural Information Processing Systems, 2023.
[41] Grigory Bartosh, Dmitry Vetrov, and Christian A Naesseth. SDE Matching: scalable and simulation-free training of latent stochastic differential equations. In International Conference on Machine Learning. PMLR, 2025.
[42] Jouni Hartikainen and Simo Särkkä. Kalman filtering and smoothing solutions to temporal Gaussian process regression models. In International Workshop on Machine Learning for Signal Processing. IEFE, 2010.
[43] Paul E Chang, William Wilkinson, Mohammad Emtiyaz Khan, and Arno Solin. Fast variational learning in state-space Gaussian process models. In Workshop on Machine Learning for Signal Processing (MLPS). IEEE, 2020.
[44] Matthew Dowling, Yuan Zhao, and Il Memming Park. Linear time GPs for inferring latent trajectories from neural spike trains. International Conference on Muchine Learning. PMLR, 2023.
[45] Oliver Hamelijnck, William Wilkinson, Niki Loppi, Amo Solin, and Theodoros Damoulas. spatio-tempordl variational gaussian processes. Advances in Neural Information Processing Systems, 2021.
[46] Simo Särkkä and Ángel F Farefa-Fermández. Temporal parallelization of Bayesian smoothers. Transactions on Automatic Control, 2020.
[47] Syeda Sakira Hassan, Simo Särkkä, and Ángel F GarreiaFFeránnder Temporal parallelization of inference hidden markov models. IEFE Transactions on Signal Processing, 2021.
[48] Jimmy TH Smith, Andrew Warrington, and Scott W Linderman. Simplified state space layers for sequence modeling. In International Conference on Learning Representations, 2023.
[49] Albert Gu and Tri Dao. Mamba: linear-time sequence modeling with selective state spaces. In First Conference on Language Modeling, 2023.
[50] Xavier Gonzalez, Andrew Warrington, Jimmy Smith, and Scott Linderman. Towards scalable and stable parallelization of nonlinear RNNs. Advances in Neural Information Processing Systems, 2024.
[51] Simo Särkkä and Lennart Svensson. Bayesian filtering and smoothing. Cambridge University Press, 2023.
[52] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Interna- tional Conference on Learning Representations, 2015.
[53] Edvard Moser, Emilio Kropff, and May-Britt Moser.
Place cells, grid cells, and the brain's spatial representation system. Annual Review of Neuroscience, 2008
[54] Scott W Linderman, Peter Chang, Giles Harper-Donnelly, Aleyna Kara, Xinglong Li, Gerardo Duran-Martin, and Kevin Murphy. Dynamax: A Python package for probabilistic state space modeling with JAX. Journal of Open Source Software, 2025.
[55] Amit Vinograd, Aditya Nair, Joseph H Kim, Scott W Linderman, and David J Anderson. Causal evidence of line attractor encoding an affective state. Nature, 2024.
[56] Aditya Nair, Tomomi Karigo, Bin Yang, Surya Ganguli, Mark J Schnitzer, Scott W Linderman, David J Anderson, and Ann Kennedy. An approximate line attractor in the hypothalamus encodes an aggressive state. Cell, 2023.
[57] Erich Leo Lehmann and Joseph P Romano. Testing Statistical Hypotheses. Springer, 1986.
[58] Martin J Wainwright, Michael Jordan, et al. Graphical models, exponential families, and variational inference. Foundations and Trends in Machine Learning, 2008.
14

---

[59] Giuseppe Peano. Intégration par séries des équations différentielles linéaires. Mathematische Annalen, 1888.
[60] Vincent Adam, Paul Chang, Mohammad Emtiyaz Khan, and Arno Solin. Dual parameterization of sparse variational Gaussian processes. Advances in Neural Information Processing Systems, 2021.
15

---

Appendix
Appendix A: Technical assumptions
Appendix B: Exponential families and natural gradients
B.I: Overview of exponential families
B.2: Natural gradients and exponential families
Appendix C: ELBO derivation
Appendix D: Derivation of SING updates
Appendix E: Proofs for continuous-time approximation
Appendix F: Computing expectations prior transition densities
Appendix G: Parallelizing SING
G.1: Sequential computation of log normalizer
G.2: Parallelized computation of log normalizer
Appendix H: Background on Archambeau et al., 2007 (VDP)
Appendix Comparison of SING to Verma et al., 2024
I.1: Overview of Verma et al., 2024
I.2: Differences between SING and Verma et al., 2024
Appendix Verma and SING the conjugate setting
Appendix K: SING-GP: Drift estimation with Gaussian process priors
K.1: GP generative model
K.2: Sparse GP approximations for posterior drift inference
K.3: Inference and learning using SING-GP
K.4: Posterior inference of the drift at unseen points
K.5: Incorporating inputs
Appendix L: Experimental details and additional results
L.1: Inference on synthetic data
* L.1.1: Linear dynamics, Gaussian observations
* L.1.2: Place cell model
L.1.3: Embedded Lorenz attractor
L.2: Comparison neural-SDr variational posterior
L.3: Drift estimation on synthetic data
L.4: Runtime comparisons
L.5: Application to modeling neural dynamics during aggression
16

---

A Technical assumptions
Throughout our paper, we consider probability measures defined on the measurable space (C m mmux] _ where C(0, Tma) P is the Banach space of continuous functions endowed with the supremum norm.
In order to ensure the SDEs eq. (1) and eq. (7) admit unique, strong solutions, we will assume f and J satisfy the standard Lipschitz and linear growth conditions:
Assumption (Drift Regularity): There exist constants C₁, C2 2 0 such that for all æ, E TR and t, 8 [0, Tmax]:
[f(x,t) - fly s)ll & Ci +t- s(), [lf( < (0(++1zz)])
and the same bounds hold for fq-
In addition, we will assume all SDEs are sufficiently regular such that Girsanov's 's Theorem (Ok- sendal [1]. Theorem 8.6.4) can be applied. See, for example, Novikov's criterion. For the sake of completeness, we restate this result below:
Theorem 2 (Girsanov's Theorem for SDEs): Suppose the SDEs
V1 dxx(t) by (==(t, + dw)) 0 Sst T
(13)
V/ : dIr( (z)())))))))))) + by (1) 1) + (l))(( 0515 max
(14)
have the same initial law v1 (x)) (0)) and admit unique, strong solutions. Then 1/2 is absolutely continuous with respect to and the Radon-Nikodym density is
Bram (11ll = exp (% Tun E 1/2 bzx)) t), dw (t))
Tim
E 1/2 (t))
2 Jo
where (w (1))(()) (o)SsTT is U1 -y-rowwnian motion. In particular, for any bounded functional + defined on C10,, max]
Ex,(((())) Ev, (0(x) (v) [1]1
In particular, Theorem 2 implies that the KL divergence between eq. (1) and eq. (7) is given by
((xx(0) (0))(0) Zr... (v11aa)
DKL (q ll p) = Byg log
('im
Tim
DRi. (g(=((0) || pl (()))
(EI((((())))) ()1=))i)
Here, and throughout the appendix, we adopt the notation || = = MM for M R2xx symmetric, positive semi-definite.
B Exponential families and natural gradients
B.1 Overview exponential families
Consider the collection of probability densities indexed by IRP
((z|n) = h(2) exp((n, T(z)) - A(n)), A(n) = log (f ((()))
exp((n, [((z))xp(nn
T(z)))
with respect to v, which we consider to be either (i) p-dimensional volume (i.e., Lebesgue measure) (ii) counting tuples of p integers (i.e., counting measure). q(2(m) is well-defined probability density whenever A(π) < co. The set E {n E TRP A(n) < 20) } is the natural parameter space of the exponential family, and v are its natural parameters. We assume E is an open set, in which case the exponential family is said to be "regular". We call l(x)) the base measure of the
Exponential families be detined for arbitrary measure but these the two most common instances.
17

---

exponential family; T(z) R² are its sufficient statistics; and A(n) is its log normalizer or log partition function. Corresponding to the sufficient statistics are the mean parameters μ IR of the exponential family, which are simply the expectations of the sufficient statistics under the exponential family, μr := [T(()
As an example, consider the p-dimensional multivariate normal distribution (V ( A) with unknown mean & and unknown, positive-definite covariance A. Then we can write this family in exponential family form as follows:
1/2 (( A) (22)--/21A1 -1/2
exp 1-211=-111
1
=2 p/2 exp
tr (A (A-AA112)) -2= )) + A A- flog A +
Let us identify vec(M), M E R2 represents the paeepp=a pi-dimensional vector representation of the matrix M. The natural parameter space is E = (vec(M) M E _pxP symmetric, positive-defimit) x R The sufficient statistics of the exponential family are T(z) = (- vec(=z ], z) and the mean parameters It remains to compute the log normalizer A(n). In particular,
(M, 77)) = (vec(A
- A A)) as well as h(z) = (2T) p/2 where
are μ = (- ] vec(A + EE )
log |A| = - |A 1 log 17
< A ----- - n7! -m1
which implies A(n) (n (m) - log |m l). Here we slightly abuse notation by treating m interchangeably as a matrix and a vector.
Our definition of the exponential family ((z()) leads us to two important results regarding the log normalizer.
Lemma 1:
The log normalizer A(n) is infinitely differentiable on the natural parameter space E. Moreover, the log normalizer satisfies
V (n)((n) = E, [x])]= v2,(() = Cov,(T(2)).
(15)
Proof. A proof of the first statement be found in Lehmann and Romano [57, Theorem 2.7.1]. For the second, differentiate the log normalizer twice
V vn)(()) ,n log
(T())((2) exp (n. ((z))u(dz) J h12))) expl (dz))
T(z)h(z) exp( (n. T()) A(n))u(dz))
Ex))(2)
V v2A(1) / v, (T) A(m))
E,, (, [())
As an immediate consequence of the identities in eq. (15), we can give equivalent expression for the Fisher information matrix of the exponential family ((2(n).
A2))(n) be the Jacobian of the mapping from natural to mean parameters. Then the following are equivalent:
Lemma 2:
Consider ((xn)) constituting an exponential family. And let de(n) dn
V²A(()) Coy [(x)] = F(n).
where F(n) denotes the Fisher information of the exponential family.
18

---

Proof. The first identity is exactly the second statement in Lemma 1. For the second, recall the identity
F(n) (o [V, log ((2
for the Fisher information. Plugging in the density of the exponential family, we get
F(n) Cov (( V T(z)) = Cova(=lm) [T()]]
Define the set of mean parameters attainable by some probability density with respect
M: { R μ Bx(=)) [=)),, q probability density wrt 2).
Note that, in this definition, 7(z) does not necessarily belong to the exponential family ((z(n). Proposition demonstrates that V, An)) maps from the interior of the natural parameter space to the mean parameter space M. In fact, this mapping is surjective onto relint(MM), the relative interior of M (see Wainwright et al. [58, Theorem 3.31). Recall that for a convex set C,
relini(C)) = {o E C for every d' f c, there exists d E C, A E (0, such that c= Ac' + (1- x)d].
A question of considerable interest is when, in addition, Vn)4(n) will be one-to-one. In other words, when can each mean parameter μl be identified with unique member of the exponential family.
An exponential family is said to be minimal if its sufficient statistics do not satisfy any linear constraints. If p-dimensional exponential family is not minimal, dimensional exponential family by rewriting one sufficient statistic as a linear combination of the others. In the multivariate normal family example, the exponential family as written is not minimal, since zz" I is symmetric, and hence there are duplicate statistics along the off-diagonal. To reparameterize this exponential family in minimal form, one can define the symmetric vectorization operator MV) := (M1,, /2M21, M22, V2M3i, V2M22, M33..... ..., Mpp) E IR((++1)/2 for M E RPxP symmetric and then rewrite the multivariate normal density as
be reduced
(p 1)-
7(2 A) P/2 exp (- Tvee((222 T)
1)+ T -- flog A + A 2
- A--E)
The natural parameters of the minimal exponential family are = (171 772 = (svec A ) A -1 Note that in minimal exponential family form, the natural parameter space has dimension p(p 1)/2 + whereas in the non-minimal form it has dimension p²
It is classical result (see Wainwright et al. [58, Proposition 3.2]) that when g(z i)) constitutes a minimal exponential family, then A(n)(m is also one-to-one. Moreover, the mean parameter space M full-dimensional, meaning in) relint(M Observe that since V,,A(n) is then both surjective and one-to-one, the inverse mapping [VA]- int(M - 3 is well-defined. As consequence of the Inverse Function Theorem, we have that for minimal exponential families, the inverse mapping (V (p)) has Jacobian (Covg(=)m(ue)) [T)))] F(n(11)-
Lastly, we give an expression for the KL. divergence an exponential family.
Lemma 3: Suppose ((z|m)) and ((2171) belong to the same p-dimensional exponential family with natural parameters 7), and 772, respectively. Then
DKL (7(2lm) 7(2 ((2))))) = (m - 72. 11) A(m) A(72).
Proof. Write out the difference between the log densities of 7(2 m7) and ((1112). Then recognize that the contribution of the base measure h(z) cancels and that E ( T(z) = #1.
Recall that in Section 2.2, we motivated natural gradient VI by claiming
DKL (π(*)m)) || ((()(12)) ~ 2 (m2 T F(()) - -
19

---

We are now equipped to justify this claim by Taylor expanding A(m²) about 7h and invoking Lemma 3:
A(n)) A(m) (m2 7h. 11) +
2 0/ 11)
1
F((m)(m - m) + O(1 mll²)
m1)
(m 12, (4) A(m) + A(72)
112 F(m)((z m) O( m 11)
=0x m) lI ((z/m²))
In the first line, we used the result V,,((m) from Lemma I as well as the result An))))) F(m)
from Lemma 2.
B.2 Natural gradients and exponential families
Next, we discuss why the natural gradient VI update (4) simplifies to a gradient update of the ELBO ((n) with respect to the mean parameters μ., rather than the natural parameters 17, when the variational posterior constitutes minimal exponential family.
Recall eq. (4) that the natural gradient VI update is
0++1) = " + ((
From Lemma 2, we have that in minimal exponential family the Jacobian of the change-of-variabe from natural mean parameters is the Fisher information matrix F(n). This implies, by the Inverse Function Theorem, that the inverse Fisher information matrix is the Jacobian of the change-of-variable from mean natural parameters. This allows rewrite the NGVI update as
n ((+1) 17()) + (- (V,)))-V,,,, (())))
dm((1) (1(12) v,,((11)) du
= n7 +
7 + (V
In the final line, we used the chain rule. The equivalence between these two updates boils down to the fact that, in minimal exponential families, one can change between natural and mean parameters by multiplying by the Fisher information matrix.
Prior to concluding this section, we provide some helpful natural gradient identities that will enable us to compute our SING natural gradient VI updates.
Lemma 4: Suppose q(2 (m) and 7(2(m2) belong the same p-dimensional exponential family with natural parameters 17, and 172, respectively. Then
V V,, DrI. (7(=lm) |l g(2 (m)) = m - 7)2
V my Bg(=))) log ((z mn) 7/
Proof. Differentiating the expression for the KL divergence Lemma 3, we obtain V pn DRa. (q(=)m) g(2)m²))
Vn, ((m) - 172 H1) A(m) + A(n2))
dm(w)) (11) l1 + (m 72) - (dr(u)) (11)
d
A(m)
dp
djp
dn
7)1 722
as claimed. As for the gradient of the negative entropy of exponential family, we have by a nearly identical argument
V Mp Eg( m log logg(z/m))
Vi ((n, H1) A(m))
(dT(H)((1)) pe m
10--1174 (dnlte) (11)
d
A(mi dy dn
djp
7h
20

---

An important point is that while the natural gradient of the KL divergence and entropy admit simple, closed forms an exponential family, the same not true for the gradient with respect to the natural parameters. Differentiating the expression for the KI. divergence from Lemma 3, we obtain
V DKL (7(z/m) (2)) = v (on 72, M1) A(m) + A(72))
= P1 + (du(m) (m)) (m - -72) pr
dn
F(m 72).
Notably, this expression requires multiplication by the Fisher information matrix, which, as pointed out in Section 3.1, is intractable in high-dimensional exponential families.
C ELBO derivation
We first provide derivation of the continuous-time ELBO, feomt. Using the notation from Theorem
2, we have
logply/() = -log plyla, (
( log / p(y)2r, 0) 2 (gll)) p( (0))
((x Q (x))de
(1))
p( (()(0))
2 log p(y)ar θ 3T. me (g
((x)ddx
((=(0))
E, loggll θ)] DEL (1(x) p(x ())
1 B2 Ey x θ)] DRI. (1(x(()) p()(0))) 2 //1 Tum
//mm (2(l(((()))))) (t ())
=Dx((((===(00)) p(ar)=(0))e))
Equality (i) follows from Girsanov's Theorem and inequality (ii) follows from Jensen's inequality. Here, the notation p(z[O) )d is adopted for readability. Formally, p( (eee@)d denotes integration with respect to the probability measure p( ()) over paths C(O, mmxx]
Note that if the parameters of p(x(0)) = (v, V) are learned, then during the learning step one will set p(x(0)) (xx())), and the second term of the ELBO disappears.
Next, for the discrete-time ELBO, replacing the path measures q(x) and p(x) with g( and p(ZorT), respectively, and following an identical argument yields the lower bound on the marginal log likelihood
Ey log q
log 2 Egg plyl θ)]
p O)
Dia (g( ) p( (e))
Substituting p( (0) for p(Ze:r ()) in the denominator of the second term yields Laparox-
D Derivation of SING updates
In this section we derive the SING updates eq. (11) the natural gradient ascent update eq. (4) performed the approximate ELBO Lapprox.
Prior to computing the gradient, we first define m; E,,x, and S = Cov Cov, the marginal mean and covariance of æ; under g as well as S,+1, = Cov, [z,+1. 2y T to be the pairwise covariance of 21+1 and æ; under q. We shall adopt this notation throughout the remainder of the appendix. By our definition of the mean parameters from Section 3.1, we have |4,1 mi, 24, S; + m.m," p4.3 S m
to be
Now, to compute the NGVI update step eq. (4), we first invoke the exponential family identity for the negative entropy from Appendix B.2, Lemma 4: V JullE log q( 17. Consequently, eq. (4)
21

---

simplifies to
n n = -n(5) + PV (v
= (1 1-))n()) + vV By ()p(p) θ)]
(1 p-) + ,V, (Bylogg + Bollo O)
B. loo po 0)] + Ey plyr 0)ōrl
From here, we consider the gradient with respect to M1.1. P42, Ma, separately. Looking at our expression for E. [[oog((=++1 [, θ)) Appendix F, Proposition we see that E, ((llgg((a,,11) [ng Y)]
16,2, depends only on (m, mi+1, S.. S, 1,+,, and hence only on (14,1, 14,,1, M4+12 (4,3). Likewise, BEg [log [.., θ) depends only on the marginal parameters of q at time Ta (m, S.).
and hence only ( (4,2). Eg (log po }) depends only on (mo, So). and hence only (Mo.1 (0.2). Separating out the contributions p1.1 21,.. 34,3, we recover the local updates eq. (11).
In order to compute the gradients with respect to cal -farr-bbee- for each (us m4+1, S(+1) ) + (m, m4+1, S, S,,1, S.t1.s) in order to first compute gradients with respect (m,, m1+1: S,, S S(+1.)) and subsequently multiply the result by the Jacobian of this map. Note that the Jacobian is of dimension D22 + 3D) D + 3D), which importantly does not scale with the size of the grid T.
(, 22, (4.3), we perform a lo- 0...., 1. Specifically, we use the bijective map
E
Proofs for continuous-time approximation
Theorem (Precise Statement): Assume f is globally Lipschitz, jointly in (æ,
Il (ax 1) - f( s)ll S (11 yll + (1-s)] sl)
in addition to
sup 0StSTTn
||A(t))
sup 0.e
[[((0]| M, IA(I) A())l. (9(1))(((() b(+)] S MI s/. 0 SSS Time.
Suppose further that At) is differentiable on 0 < t < T
Then the error between the continuous and approximate ELBO satisfies
Coom (g, θ - Lapprox (g. θ) & (22) C(E, Toax M,, 11-1//11,, |
where At max?"" -1 A; the mesh of the grid r. The constant CEE, |==(0) TTTax M, D, Ly) does not depend on (A(t))oSISTI.. or (b(t))oseST.... except through M.
Proof of Theorem 1. For ease of notation, we will drop the dependence of f on O. And without loss of generality, will assume 11/ I. This is because the ELBO is invariant under invertible transformations and applying the change-of- variable efficient to identity and changes the prior drift to Σ 1-122 o f 1/2 The variational posterior drift also transforms in the same way. The transformed prior drift has Lipschitz constant ||/ |
-- 1/2 reduces the diffusion co-
Recall that defined the difference between subsequent grid points be A, 7++1 тÜ and let us define (=( continuous and approximate ELBO
(,), u) = fy (x(u). u) f (œ(u), u). Then we decompose the error between the
Tun E, Ef1f((x,t) - (2,,t))1
A T-1
C.eoat (q. +) - Lapprox (g. O)
<
Σ
Eg
(1
?)
2
2 1=0
A T-1
T
(x [x;)
+
Ey (1)
log
II
(16)
2
p p(0,11 xx))
1=0
0
22

---

We will proceed by bounding each of these two terms individually. To do so, we first state and prove the following three:
Lemma 5:
1g (B,, [[=)] + Daeax + ex2 + m)Tma == (o(EE]|2e00)1 Trax, M, D)
sup OstgTTn
Proof. By our assumptions on (A(u))osuTTh. and (ousST)) the SDE admits unique strong solution that satisfies So E, [z]( co. By Ito's Lemma,
d()2 ((x Su (x((1) 0) + ((r((), dwo()) Ddl.
Taking the expectation of both sides and recognizing that the Ito integral a bona fide martingale,
(,)1==(111 = Eg]1x00)) S (z (2) (q(()) w)
D)du
s < / (I + M)) + DJdu
{ + DD ++ Mc+ / (2M + ( (F.
Grönwall's inequality yields the desired inequality,
sup OSiSTan
By]la(t)]1 (E, [z)) DT + cxp(xp(( 1)T max )
Lemma 6: For all 0 S S T maxx we have the following bound
Eg [1]1=x(1) <<)) (E, Tiax M, D) (l 8)
for C1 (1,.y11000 Tmax M, D) not depending on 8, I.
Proof. We start by rewriting the squared norm difference between xtt) and #(s)
[|x(t) x(s)]|² - -11." dwo + J. fo (z(() widu
[ drw f." A(u)(z(w) =(s))du f. (IA(u)- A(s))a(ss) (b(u) b(s)))du
(t - ())((()()))
(1l.. dew (x)
+ +1/." A(n)(=(u)- - =(s))du
11. ((A(w)- A(s))a(s) (b(u) b(s)))
+t (t- s))((((s)((s)) + +60)))]
where we use the bound (E az)² < LO az). We take the expectation of both sides to obtain
Egg (x) - xs))112
u])2 +E, E, 1 A(u)(z(()) x(s))du
s D(t 3) + (t-
sup OSusTTma
21001++02 b(s)))duu )du
+ Eg f/'' ((A(u) -As))22, + (b(u)
b(s)}
23

---

By applying Sauchy-Schward to each of the final two terms, we get
gg | A(w)(zr()) - (m()))d
(1 6) L By (g(((()(a(w))) - (B
S M () L. Bellx(()) (b))]|2d
as well as
B2 f. ((A(u) - A(5))=(()+
+ (((()) bb())) ( - () L. Eel(((()) A()))(()) (b(u) d(s))
2((11 (e S. E,(l(((u)) ( [[b(u) ((s)l?)du
2 - M² ! (B,]|x(>))1 + 1)du
2t s) 8)* M
E, (II++
sup oSsuSTim Tmx
Combining above bounds and invoking Grönwall's incquality, we get
Ey11((t) x+))|| < 4 D(t - 8) + (t - $)²
sup OSuSTme
(xll w]l
+ 2(t ²2²
Eg |x(2|1 + +1 exp (t - s)²}.
sup 0
Recognizing the leading term as (t s) and invoking the bound on supogusTi.. Ex](+(u)]1 from Lemma 5 we obtain the result
[g)1=(() 2s)) < C1 (E, [=(1) T M, (t- s).
Lemma 7: Under g, the mean 7 and covariance S,,41 of g(x++1 satisfy
m = 24 A.f((E + (A.)
x +r:²) (2)
S S I (A.)
2 (A(T) + A A(T:) (A,) R:
where [|x !|r?
(2)
I1 | S C2(M) are independent 0 TT.
Proof. Since q the (strong) solution the linear SDE eq. (7), then g(x1+1|2;) is Gaussian with mean and covariance
m,,41 (T++1,7:)) +
TT A, + u)b(u)du
S,,41 =
(T(TT + A. + u
D. Ti + u) du,
st,,s):= : 1+ E _
A(ui) A(ux)du ... duk, 05 StS Tmax
(17)
where the series form of the state transition matrix is known as the Peano-Bakee series [59]. Truncation of the Peano-Bake series at its second term yields
V (t, :6:s)==0 = [
A(u)du - (t,
24

---

where
[ (t, (t,
Σ
A(w1) A(uu))du duk
A(u)) A((ux)du duk
(1-8)(M MM*K
k!
k=
exp((1 - -1 (1 ( s)² s)²
exp( s)).
In inequality we used the fact that the series converges absolutely and for equality (ii) we used Jh J.u juk-1 du1 ... dux = (t 3)* /kl as well as our assumption | I M. Moreover, by the Value Theorem,
[" A(u)du = (L (1- + (1,8) (I) S M(1 - -))²
(18)
where we used our assumption that A() is M-Lipschitz on [0, max and hence its derivative is bounded in operator norm by M. Plugging this into our expression eq. (17) for the mean and covariance matrix.
M x; + (z, T.) (A.)) (r.")
(1) xz 12:+1 (2))
S A (
(A(T)) + ((T) + (A,) R.
where 11 II R < (2(M) are independent of the index i.
As immediate consequence of Lemmas 5 and 6, have for any 0 S 058 T Tiax.
(xll)((=((1,.)) +)]) $ 3E, { [|A(t) =ss)] + ||((t) A(s) ([s)||[[((5|| + (lb(t) b(s)
s $3ME,, .s))] (t (t- s) {
1
sup ogusTmm.
E,(la(u))1?
I CG(E,lI (0) T M, D) (t s)).
We also have that for any 0 < t < Tmmx
(1)) <2 2 M { 1+ sup
B. [a OSuSTia
+ sup oSuSTam
I CE, (=00)
Tmxx, M, D).
For the second inequality we invoke the linear growth condition on f and Lemma 5. Equipped with these results, we are finally prepared to handle the two terms in the decomposition eq. (16). For the first term,
Time
[.llf((())))))))
- (1) Σ A By (1 mll
2
0
(2), - -E (2) (zr( ui)l du
Eng
Eg { (11, Ti)l 11 ()
(x)(), u) (
-
du
1=0
25

---

u) dll
Till
2v 20 + =(u)ll2 + 2Lz
u)
A 3/2 C3 (By]|zol1 Tuax, M, M,
i=
< < (By] xo 11 max M, Lt).
Altogether, we have proven that the first term eq. (16) (((A))122)
For the second term, we know that the strong solution q to the linear SDE eq, (7) has Gaussian transition densities g(x(++1/2;) By the KL formula for two Gaussian distributions
Egg log g( (x:) p(a=+1 (e;)
1 E.
(( - D D log A; log Saitl
(æ; + A, (as, T.))1
2
A.
A.
where 7,,+1 and 5 1,1+1 defined as in Lemma 7. By Lemma 7, the mean term is
Eg m - (a: + A.f( T:))
24
2 Bg [ Ti) - f(x ( M,, M, D. Ly)
2 Eg |/(,, T.) + (E, M, D,
In the second line, we expand the square and invoke the bound on supogigTm.. E,][2(t)]1 from Lemma 5, Also by Lemma 7, the covariance term satisties
A tr ((..)))) = D + (L(((()) A(T:)
log |S.141l D log A; + A (())) A +
Altogether, have shown
A Eg (7 Ti) - E.g log ((0++1 23)))
A,, Trux: M, D, I.))
2
((141)0;)
Summing across time steps yields
T-1
T q ( (xx))
Eg (11(2, Eg log
< (E,)) Tmaxx M, D, Ly).
p (x1+11<;;)
F Computing expectations of prior transition densities
Here, we provide proof of Proposition from Section 3.3, which says that (gy(a) (log p(æ; 1 [x; O)] can be written as an expectation with respect q(æx)) only. To do this, first prove the following statement, which is consequence of Stein's 's lemma.
Lemma 8: Suppose
Øi+1
æ;
N
mi+1
m:
S
S
S
S.11
where x, x:+1 RR,. Then, for
any f : RD - RD [f ( (m m
E
[J (x:)) ST where (x( is the Jacobian
of f evaluated at æ.,
Proof. First, for notational simplicity, we denote
t = C; Li+1
g = f(
[(x++1)]
f(x))
m
mi+1
m:
S
S
ST
S
26

---

Next, we apply Stein's lemma to the Gaussian random vector ø and function g()). This yields the identity,
(() [ m) [] ( [T)( S,
(19)
where 99 (z) is the Jacobian of evaluated at a:: (g y (m))
y (T4441)
Reading off the bottom right block of matrices in eq. (19) gives,
(m:)) m;+1) [] [m:)) S S/
Using this lemma, we can now prove Proposition 1.
The term E, E,,log((a; [ θ)] can equivalently be written in terms of the mean parameters (Ju,, 14++)) and the expectations E. [x;; IR,y [f(x;]() T [)] and Eg J (x:)), thereby reducing the dimensionality of integration from R2D to R.P
Proposition 1:
Proof. Expanding the transition expectation term, have Eg Eyllog((zz4) æ, O)] = Egg logllo//aas (ax+11 I; + (z) A,E)]
D
(2T 24,
Eg || - x; - A, 0)
2
We can rewrite the expectation in the above line as Bgg || - 2; - (x; [0)]
= -E,l1 1 | 3/+ + (2
20 - - -- f(x;(0)
+2A,2, Σ Σー f(x |
-tr(E mi+ 1 tr(
m.m A7
( 1-
- -2tr(E (S + +mmm.1)) 2A, [ Σー - O))
24 Eng
| Σ f(x; [)
Now, applying Lemma the second to last term, we have
Eg - f(x/ tr\ [f, ( m Es J (z:)) s
Applying Stein's lemma to the last term, we have E,g . Σ f (z. [
tr - ( ()) + ((.)
Thus, we have shown that Egg log x [z;; θ)] can be written in terms of pi, M441, and the expectations Eg [f(x Eg f )) Σ f(z; θ) and Eg (z:)]
G Parallelizing SING
G.1 Sequential computation of log normalizer
Here, we describe the naive sequential approach of converting from natural parameters 17 to mean parameters μ of g( The idea is to marginalize one variable time eq. (12), starting from To and ending at æT.
We'll first show how marginalize out øy and æ1, and then use this to derive general recursive update for this sequential approach. To marginalize out Zy, we focus on the following integral which contains all terms depending Zp.
exp
Tn Jnxn + (hn- - LJj x x0 Too
dxp
(2m)
|Jol exp
j ho -otno ho
exp h (T)
2 x
f(m] x
2
: Zo
27

---

where Zo is the contribution to the log normalizer from marginalizing out 2o, and
J|m = Lo h Lo ho.
Next, when marginalizing out ø1, we should integrate the potential exp (h (P) from the previous step multiplied by the rest of the terms in A(n) that depend on at.. This gives the following integral,
21 _ 1
exp (- a Lं Ji 26 11 h2 21 dx 1 21
exp
h
(D)
21
2-
2
1 T
exp
2 1
(Ji +
(h1
+ h1
-b [l)) w22)
dx
=
=1/
(2T) $
[1 1
exp
1 t J
h0
exp
h22
(p)
022
2 22
1
J2 x2
2
:21
where Z1 is the contribution to the log normalizer from marginalizing out æ1, and
n = =- LT, h = [
Generalizing these recursions leads to the following sequential algorithm:
1. Initialize J (() = 0,, 10 (p)
0, and log Z = O.
2. For i ..... T:
(a) Update J.(0) J; J. and h (6)
h: + n(20)
(b) Update log + log Z + log Z..
(c) Update [(m L.J (c) L and h h &- -L,J(()-11 (n
3. Return A(n) = log Z.
Time complexity In each iteration i, this sequential algorithm has time complexity O(D3) which comes from the matrix inversion and multiplication in - L Since each iteration relies on values computed in the previous iteration, this algorithm must be implemented sequentially, leading to total time complexity of O(D*)).
G.2 Parallelized computation of log normalizer
Here, we present our new method of converting from natural to mean parameters in parallel time using associative scans, leading to substantial computational speedups on modern parallel hardware. As discussed in Section 3.4, we accomplish this by defining elements in the scan as unnormalized Gaussian potentials oj,j and combine them using the binary associative operator,
a ajK dr
(20)
Note that eq. (20) is indeed a binary associative operator, as it satisfies the identity:
(ai.j aj.k) ag. = /
(J (.. an.
= I ai.j (ʃ (ossannda dxz
as.j
(aj.k a...).
In the following calculation, we show how to compute the binary associative operator between two potentials analytically. We will start with computing the operator between two consecutive potentials,
28

---

and use this calculation to generalize to computing any pair of potentials with shared variable to be marginalized.
Let 1,-1, and 01,++1 be two consecutive potentials. Then,
0,-1, 01,441
1
T Jaa - 10c.-xx1140 10-1L;-10; + h 2t
(-2dit1Jet1aan - x. L x:41 + 1 dæ
1 (-2
exp
Ri
exp
2
exp
2141J1412141
T
1 f exp
1
Jaw: + (ls; -1-10:111 L 20411) "'ee" xss
do:
2
(2
8 JJl
exp
1 xi-1
0i+1
L (p)
((
(
J
Øi+1)
xi-1
1i+1
xi-1
hog
(p))
Zaand
01-1,1+1
where Ziocal the contribution to the log normalizer from this operation, and
J( L; LL J, Li-1,
J Ji+1 L:J," L,
h (p) J. hi; L (P) L.J; Li-1-
h hit1 L,JØ hi:
From the calculation above, observe that the element digy, where and j are not necessarily consecutive, takes the general form
1 I
L J2
h1
0i. Ziocal exp
L
2;
h2
Using this form, we now derive the analytical result of the binary associative operator between any two Gaussian potentials with a shared marginalization variable.
Oj jjj.k
Ziocal exp
F2
J1
L
a
a
h1
2
L
J2
h22
Ziocal exp
Øjj
J
L
T
(hj
dx,
L
J2
h2
Ztocal Zoczl exp
220 23 JI
28,
h
T 2
J2æk + xx h2
2
/ exp
1 3
T
(J2
+ J. Jay
+ 2..
(- [(,)
[P) Tx + h2 + 12)
dx.
:=)(()
( exp
1
xi
T
(p)
L(e)
x;
c;
(p)
2
(i)
[(P)
+
h
0j,k
where
{(P) J1 - L -()))
J - -.
(P) h L 1(0)-1
h(p) h, (
(p) (
(1 n (()-)1()))))))))) 2
(p) Local gooca (2T)) [ exp local
In practice, instead of computing (n)) we instead compute log ( Since each binary operation marginalizes out one variable, the final result of the associative scan,
0-1,0 00.1 0r,++1
will be the result of accumulating all log 7 terms, and will hence be equal to the desired log normalizer A(n).
29

---

Time complexity The time complexity of the associative scan, assuming T processors, is O(C'log T) where C represents the cost of matrix operations within each application of the op- erator [34, 48). By the derivation above, C' = ((D*) which comes from the matrix inversion and multiplication in L j(0)) LL and analogous terms. Therefore, the parallel time complexity is O(D T).
H Background on Archambeau et al., 2007 (VDP)
In this section, describe the variational diffusion process (VDP) smoothing algorithm proposed by Archambeau et al. [7]. The namesake for this algorithm is adopted from Verma et al. [33].
Rather than first approximating the continuous-time ELBO Leeat. eq. (8) with Lapprox. eq. (9) and then maximizing with respect to the variational parameters, the VDP algorithm directly maximizes feom- In particular, the VDP algorithm involves iteratively solving (continuous-time implemented numerically, though, these ODEs are solved using a discrete-time scheme, such as Euler's s method. Therefore, like Verma et al. [33] and SING, VDP only approximately maximizes the true, continuous-time ELBO.
ODEs. When
Archambeau et al. [7] rewrite the ELBO Lcomt in terms of (A(1), b())os)TT.. well as (m,)), SesesTT. the marginal mean and covariance of g(x) at each time S S Tmax These parameters related by the Fokker-Planck equations for linear SDEs
dm((t) A(t)m(t) b(t)
(21)
dt
dS(() (())(()) 5(1)((1) E.
(22)
dt
To ensure that, at the optimum, equations eq. (21) and eq. (22) are satisfied, VDP introduces the Lagrange multipliers (A(t))oSeT.... and (1())osTTTT.. and incorporate eq. (21) and eq. (22) into the objective
Loout Cooar((A ((A(t), ((1))),,) :) 1,"0a Tan 2(t) 0) 0)+ 0
(dm(t)
A(t)m((t) b(t) dt ( dt
T
(W()T
ds()) - -000) xx))uL dl.
dt
Performing an integration by parts yields
Ccam Camn((A)) ((A(1), (()) S(1))ases..) +) 9)-1 ('um
1)-((mm X(() (A(I)m(l) + b(t)) dt
ffrme L (V(() ( (A(())())))))))))))))))))))))))) ( ((() + Σ} di f. Tino
(t)(t)) T 5(0) }
dA(t) T
( di
m(1) dt
dl
(T(()(Tmm))) S(( ir (0)
S(0)) (()) A(Tma)) "mm(Tmu (Tmxx) (())
m0))
(23)
The authors proceed by performing coordinate ascent on this objective with respect to (A(t), b(t), mt), S(1), ((), (1))oss) However, when taking the (variational) derivatives with respect to the functions A()), b(t), m()), and S(t), the dependence of A()) and b(t) on mt) and St)), and vice versa, is ignored.
Taking the derivative of eq. (23) with respect to m(t) and S(t) yields the updates for the Lagrange multipliers
dv (tt) (A(t)) (T() ds(() (-DKL (q || p)), dt
L (Tmax) 0
(24)
dA(t)
d
-Att X(t) + dm(l) ( (-DiL. (q p)), A(Tmax) =
(25)
dt.
30

---

subject to the jump conditions at observation times {ti)
((,++) (4,) dS(4)
Bg (ogp(y(4))((((.))
(26)
d
X(1)) = X(1)) + dm((1)
By logg(u (4))((()))(((44)) ((4)).
(27)
Taking the derivative of eq. (23) with respect to A(t) and b(t) yields the updates
dd(xx
A(1) (((a(())
(L)
(28)
dx
btt = (q(((t) [ [)] A())m(t) - Ex(t).
(29)
Finally, for p( (v, V). taking the derivative of eq. (23) with respect to v and V yields
m(0)
V\(0)
(30)
(()) (2) + V-1)-1
(31)
Altogether, the VDP variational inference algorithm alternates between the following steps, until convergence
1. Solve for (m()), (()))SISTT... according to eq. (21) and eq. (22).
2. Update (A(t), ()) (t)) starting
from A(Tmax) = and V (Tmax) = O, according to eq. (25) and eq. (24) and taking into account the jump conditions eq. (27) and eq. (26) at observation times (4))--1
3. Update (A(t), ((t)ss..TT... according eq. (28) and eq. (29).
4. Update (0) and (0) according to eq. (30) and eq. (31).
For stability, Archambeau et al. [7, 8] propose replacing the update for (A(t), b(1)) at iteration l with the "soft" updates
( (t) A + w ( (1)
(32)
b (CtD (1) = 6 (1) + w (1) (00) (1)),
where (t) and (t) are defined as in eq. (28) and eq. (29). In experiments, observe that choosing w to be small (e.g., 10-3) is often necessary to ensure convergence of VDP.
Unlike Verma et al. [33] and SING, VDP does not benefit from convergence guarantees. Indeed in Appendix we show that even the setting where the prior SDE is linear and the observation model is Gaussian, VDP will not recover the true posterior in one step.
I Comparison of SING to Verma et al., 2024
I.1 Overview of Verma et al., 2024
In this section, we outline Verma et al. [33], the work upon which SING is built, and discuss key differences between the two variational EM algorithms.
Like SING, Verma et [33] maximizes the approximate ELBO Cupprox from eq. (9) with respect the variational parameters (A(t), b(t))(()ste..). well as the model parameters O To accomplish this, they perform a change-of-measue
pl Do:r θ PL (x(or)))
p(Zorr 2) PL (x
where Pr. (For) = (( belongs to the same exponential family as the variational posterior and has natural parameters 7). (h~_ JL , Suppose q(FoTT|)) is initialized 17(0) m.. Then the corresponding NGVT update steps can be written as
Æ0+1) (1 - p)(({)) p pV Egtn (oggrty 1:: + A,V + (V(
j(+1) = (1 p)j{()) + pv nr, E, [log (:) + A. (2:, V.
(33)
i (1- (,))) + DVV,, A. (x [i+1]]
h. (0+1) = 1,(++1) + n²" J(5+1) (65+1) j L (++1) = i (041) 1²
31

---

Here, (z, 2++1) = log (f( ()/pL(2:41 S is transition probabilities from I; I1+1 under p and pr.- The NGVI updates that appear in Verma et al. [33] are written in this unusual form to separate the contribution of pt, which belongs to the same family as q, from the other terms that appear in the ELBOO.
is the log ratio of the
The authors propose alternating between (i) performing numerous NGVI updates (33) and (ii) updat- ing Pt. at iteration j according to a statistical linearization step. Specifically, statistical linearization of the prior drift computes the best linear approximation to the nonlinear prior drift under the current variational posterior according to
A by = arg min R. ((mm(1651))1
A az + by - faz (
(xf(x:(0) ())
(34)
A E (w, (0)
b..h E (f()) A i- O.....
dx
By the log normalizer computation discussed Appendix G. one can use (A,, by. 1 to obtain new set of natural parameters 1/w.new (h new J new L Lew. One then updates the natural parameters by setting
h h,6)) + (h!
- new), J(H t j(, (J+ -J.--J J. 1 1 + (L)
(35)
and lastly s7 <- new. The updates eq. (35) ensure that updating the parameters ny' of pt does not change the parameters of the
Importantly, the statistical linearization step appearing in Verma al. [33] does not matter for inference, since is apparent that by "undoing" the change-of- measure, one can rewrite the NGVI update (33) to exactly match the SING update (11). In fact, we point out that repeatedly performing statistical linearization, as is suggested by Verma et al. [33], is unnecessary: one will obtain the same o7t by performing only a single statistical linearization step at the end of inference.
The reason for introducing Pt that the authors contend that it aids hyperparameter learning, insofar it allows one to couple the E and M-steps via nL. which from eq. (34), evidently depends on O In other words, during the M-step Verma et al. [33] solve
arg max C(n) + 7. O), n = (h. J..
e
This idea is adopted from Adam et al. [60]. One can imagine there are many functions 0 ; R/ = that couple the E and M-steps by writing = (())= + (1) (())). The authors argue that defining & in eq. (34) is desirable insofar as optimizing 7) (()) with respect to n for O fixed is less
dependent on θ than for alternate 0. One issue with this argument is that, unlike Verma et al. [33], we are also optimizing the parameters of the likelihood p(y /rr, θ), and is not evident that statistical
linearization of the prior eq. (34) best couples the E and M-steps.
1.2 Differences between SING and Verma et al., 2024
SING aims to simplify and improve upon the algorithm presented by Verma al. [33]. First, Verma et al. [33] does not address how to compute the term
1 E.g log p(are la; )) = (u(-b) 1
Eg (x,, x1+1)) =
- 11
x, + bj. - f
A;
PL (r( 1zxs)
Fi+1 - x; (A/ x + b (zx (A, + b/:)
A
Σ-1
which, a priori, requires approximating expectations in 2D dimensions. In Appendix F, we explicitly discuss how to compute analogous quantities for SING using D-dimensional expectations, which makes our proposed inference algorithm both faster and more memory-efficient. Moreover, as we previously discussed, changing the base measure from p PL is unnecessary for inference, and so we remove it. We do not explore coupling the E and M-steps with SING.
Additionally, the inference algorithm proposed by Verma et al. [33] requires, for each NGVI step, sequentially converting between natural and mean parameters, which has linear time complexity and can be computationally expensive when the number of steps is large. In Appendix G.2, propose
32

---

parallelized algorithm for computing the log normalizer of the variational family, which leads to significant speedup in the runtime for inference (Figure 5).
Using the sparse variational GP framework from Titsias [36], we propose an extension of SING, SING-GP, to the Bayesian model that includes Gaussian process prior the latent SDE drift. Previous works that perform approximate inference in GP-SDE models, including Duncker et al. [13] and Hu al. [14], perform inference on the latent variables (a=(t)o<<sTT... using VDP, and hence inherently slow to converge and unstable for large step sizes A; (Figure 2 and Figure 3).
J VDP and SING in the conjugate setting
Proposition 2: Suppose f(z :|0)= A (t) (t) batior (t) and p(y((t))(((4.)) Q) -N(2x((() d, R). Then, when replacing p with in upppoxx SING performs exact inference on T when 1. Specifically, after one full natural gradient step, g( (1) = p p (D A).
Proof. First, we point out that we are able to replace ÿ with p in Lupprox since f linear.
From here, we write
p(Zorr )p( (o:r)
T
1
Σ
T
J. xi Σ
2 (1)(())
Cx((k) dl
wi
2 1=0
2
T
T-1
T Σ (n + 8.C R (y(7)) dd)" æ2
exp (
- 0 (J. + &.C R C'a=
Σ <=0
x
L
:
where 8; is the indicator denoting whether an observation occurs at т, It is straightfor- ward to see that this coincides with the posterior over paths p(æ) p(æ D), marginalized to grid written p()oor((). Hence, p(za:r(D) is Gaussian with natural parameters 17tue
(no + (y(TT) - d), -J.rio + S R oo Je-0 (-
If we write Luporox with p substituted for p as
p(an.r (y
Luppoox (1. + = Eg log
(x1:T)
then the NGVI update eq. (11) with p 1 becomes
77 1)) 17 (0) + Cappox (17
(1) n
(0)
(17 tree (0) = mic
+
In equality (i), we used the identity from Appendix B.2, Lemma 4 for the natural gradient of the KI. divergence between two distributions belonging to the same exponential family.
Consider the same setting as in Proposition 2. Then VDP algorithm described in Appendix H does not recover p(z] in one step.
Proposition 3:
Proof. For simplicity of presentation, we will assume Σ = I.
Suppose we have (m)),, S(t))oStTT.. the marginal means of covariances of the variational posterior q. as well as (1()), (()))ser..... that specify the drift of the variational posterior. After the first step of the Archambeau et al. [7] VDP algorithm, these consistent.
In order to analyze VDP in this setting, first compute the (variational) derivative of each term in the ELBO Loomt with respect to mt) and S(t). To this end,
B pdor (t)x(() +
b (t))
b(0)))
2
=F,,
(A A(t)) ( (f)))=()) + ( (
(AP A(t)) (bor"(t) A(t))((6"""()) b(6))
const
2
1.
tr (A) - -A(t)) ( (A A(t)) (()))((((++()m((()) m(t)m(())
2
m(1)" (ATT
(A (t) A(1)) (b) b(t)) const
33

---

well as
Bg log. (olt)) + +d
E
2 - -
const
1 (ti) C
( C + C R y(t)) xxti) T C R d + const
- g
2
tr (C" R C(S(() ml(.)(())))) + m(4)) TC C R (y((,) - d) + const
where const represents a generic constant with respect to (m(t), St)
From Appendix H, these imply the following updates for the Lagrange multipliers
dA(t) -Att) ((t) (A prior (t) - ( (Aprio(tt ((t))
(Aprior
dt
(Aprior()) (t) - At)) prior (t) b(t))
(1) A() T 4 (t) (t) 2 (A A(t)) (A prior (t) A()))
dt
with the jump conditions at times (4:
A(t)) X(t,) + + R (y(t,) - -d) Cm(t)
M (1tt) = ((4) 2 C
CTT-C. R.
WLOG, assume t, < max for S terminal condition is the same as the jump condition at time m. One can check these match the updates appearing on Archambeau et al. [7], page 7. The last step in the VDP algorithm is to update the variational parameters according to
and set ((Tmax) = 0. L (Tmax) 0 (otherwise, the
(1) A prior (1) (L)
b(t) A Aprior (t)m((t Borio (t) A(t)m(t) A(t) (Aprior (1) A(1) bprior A(t)
Next, we determine the analytical form of the posterior drift in this setting. In particular, the Kushner- Stratonovich-Partee equations [29-31] give closed-form expression for the SDE satisfied by the posterior path measure
q dtx [Aprior x; bprior + V log log dt dw(().
where p (t)) represents the conditional probability of observations yt:)). 2 conditioned on xt))= x. Observe also that p(v22læx(t) will have jumps occurring at observation times.
In general, cannot compute p(y=+/=<()) in closed-form, and hence do have an analytical expression for the posterior drift. In the linear case, though, we can compute this explicitly. In order to do so, we make couple of observations. First, for 4:11 < I. S 44, we see that by the Markov property
(()))(((() x ((x(()))=(())
II (ott A=n
p(u((t))<<((a))
Since each of these factors Gaussian, (x)))((()) will again be Gaussian. Let us write
p(yz:/a(t)) = exp { [ (())x(()) + g(t) x(t) + =dt)). 4- <t
Second, (xx 1) = p(yzel (1)) satisfies the Backward Kolmogorov Equation (BKE)
(2, t) + (A (t)x(() + paio (t)) T (xx + 2 0t
A (xx t) = 0
34

---

in between observation times 44-1 < <ts &.. Plugging
(xx, () exp
2" G(t)æ + g(t) x x =(t)
into the BKE, we get
(0 ()(A (t)æ 6utn (t) (g(() G())x) + v(x, 1/ {((()) G(())))( G()))) G())) (n
1 d = --v(2, 1)
--20 dt
G(t) ae g() J x + dt cft)
d
Simplifying,
(Antiar (f)a: + b (1)) (G(() 9(()) + (((())=-(00)00 +(()))++(((()))))
1 d G(()x di 00(1) x + d c(1) 2 dt
dt
(tt) = - (Aprior (t)) G(() Gtt) A ptior (t) + G(t) G(())
dt
dt g() = (AAPin (t)) T (t) + G(t) prior (t) + G(t) g(t)
where used the facts that 2x A (t)) G(t)a= x (Aprio (t)) G(t)x+ ())) T()()) A x
G(t)
and that if mmTmax) is symmetric, then (G(1)) will be symmetric.
In other words, we have shown that p(y (()) satisfies the Riccati equation
di G(t) = - ((A prir (t)) G(t) - A.A
(
prior (t) + G(t) G(() G(()
dt g()) = (AA prior (D)) g() + G(t) T paiir + G(())g()) g())
jump conditions
G(t,) = G(t)) C C
g(t+) = C (R)))
and initial conditions G(Tmax) = 0, (TTmxx = 0. The posterior drift is then given by
ptlor (t)x((t))))))) + botior (t) + V logp(y>;lz(t) = (Aprior - G(t))=x(t) + (borior + g()).
Now, though, we can directly compare (t) to G(() and (A) - A(t) + A(t) to g(). In particular,
d(2nT (t)) A(t)) (2 (t) (2 (t) A(t) (A prior (t) A(t)) (t) A(t)), dt
(Arrio
24 (t.) (1) C C which is not equivalent to the Riccati equation satisfied by (G(t))oSiSTt>>
Suppose, for example, (A(t), ((t))((()))ss)) are initialized to (Aprior (t), butiar (t))ost Then
d(22 (AA mior (24 - (1) (1), = 24d (
dt
K SING-GP: Drift estimation with Gaussian process priors
Here, we present the GP-SDE generative model and full inference algorithm using SING-GP from Section 3.5. Our presentation of the GP-SDE model primarily follows that of Duncker et al. [13] and Hu et al. [14].
35

---

K.1
GP-SDE generative model
The GP-SDE model consists of the latent SDE model from Section 2.1 combined with Gaussian process priors on the output dimensions of the drift function f(((t)). More formally, we have
dx((1) = d()((() + E)dw((), Ely((4))= ((Cx((,)) d)
where
fa(.) Hd (P((),o(:.-),, d d= d=.....,;
Ke is the kernel covariance function, and 0 are kernel hyperparam eters. As in Section 2. we denote all model hyperparameters as @ = {0, C, d]. Throughout this section, assume Σ I without loss of generality. In addition to primary goal of performing inference over at)), we would also like to perform inference over f()). In Appendix K.2, we outline approach from Duncker et al. [13] that sparse GP approximations for inference over f(-).
Here, f(-) = [
K.2
Sparse GP approximations for posterior drift inference
Following Duncker et al. [13], we augment the generative model from Appendix K. with sparse inducing points [36], which enable tractable approximate inference of f(-). This technique allows the computational complexity of inference of f(-) evaluated at batch of T. locations be reduced from (LL)) to (LL²), where M < L is the number of inducing points. Let us denote the inducing locations { m=1 C RR and corresponding inducing values as (ua)d RR- After introducing these variables, the full generative model from Appendix K. becomes
x, x. f.u () ppyl ax O)p( f) II
P(fa ua O)p(ua (). d=1
Treating un pseudo-observations,, we assume the following prior on inducing points which derives from the GP kernel:
p( ua θ) =N(uu K.x)
where K: is the Gram matrix corresponding to the batch of points {zm JM conditional Gaussian distributions, this implies that fa(-) un is another GP.
m=1- By properties of
fa(-) ua (P( (-), Kalaaa (-:)))
where
P us (x:) = kuz Bu.kK- K=2 faluus (x, x') = Ko (x, x²)
kis K-2 Kgn-
Above, kaz = [N #1), ................................ No (x, zM)) and hizn TT.
To perform VI the GP-SDE model, we follow Duncker et al. [13] and define an augmented variational family of the form
(x f, =0= q(x) II p(fa ua O)q(ua).
For the variational family on inducing points, we choose
q(u.) = Nuu m Sal)
so that m and S:d have convenient closed-form updates, describe in the next section.
36

---

We use this augmented posterior to derive a continuous-tims time ELBO analogous to Coont in eq. (8) of the main text. By Jensen's inequality,
log p(y ))
- -log / p(y æ. O)p(a) f)p(f u, O)p(u O)dadffu
2 f q(2, f. u log
ply æ, p)p( f) u, p)p( )
dxd du
(x, f, u)
= =ʃ q( f. u) log log
ply x pp f) T p(u
) drdfdu
((x)) I g(ua)
D DEL (g(ua) |l p(us ())
θe Ea DxL (g(x)) I| p( f))l
Σ d=1
=
Lam.aap((())) θ)
where
(f) II p(fa ud. ()q(ua)duaa
(36)
Since C (()), θ) is intractable compute directly, discuss in Section 2.3, introduce an upproximation to the continuous-tim ELBO that is analogous to Loont in eq. (9). This objective is
Capprox (P (1( q(u), O):-
Eg(as [oggly x; + - [DKI. (g((0):) || ((xorr f, (])
Σ DKL (q(ua) || p(ua θ))
(37)
In eq. (37), we use definitions from Section 2.3: ((For)) is the finite-dimensional variational posterior on time grid and p(æo.r) ) is the corresponding prior process the same time grid discretized using Euler-Maruyama. We note that Lappox has two main differences from Lapprox: (1) the second term now involves expectation with respect q(f), which comes from the Gaussian process prior the drift, and (2) Lapprox-Ge has an additional third term regularizing the posterior on inducing points towards its GP-derived prior.
K.3 Inference and learning using SING-GP
In GP-SDE model, the two inference goals are (1) inferring posterior over the latent states 7 ((ror) ) and (2) inferring a posterior over the drift function (i.c., dynamics) q(f). In SING-GP, we propose to SING as described in Section 3 perform latent state inference, and derive novel closed-form updates for q( ((u,) based on Lapprox-GP. In the next section, we will show how Gaussian conjugacy allows for straightforward inference of q(f) given q(ua).
To update q(ua), we would like maximize Lapprox-GP (q( q(u), ①) with respect the varia- tional parameters m Sdde given the current estimates of q( and O By direct differentiation of Lapprox Gp (g(x::r) ), quu), ②) with respect m and Side one can derive the following analytical updates for mă* and Sddt
sd K25
K
At
E Bygn (Kzz, Mezss]
K.:
1=0
T-1
Mi+t
m.
T-1
dkzz
S - S.
m
+
By(a))
At
(38)
In the above equation, m E MDD contains m in each column. The required kernel expectations in eq. (38) have analytical expressions for, e.g., the commonly used radial basis function (RBF) kernel, and otherwise can be approximated via Gaussian quadrature for specialized kemness. Moreover,
37

---

observe that since the dependence on T only appears in sums, this update can be computed in parallel across time steps.
Finally, the learning goal the GP-SDE is to optimize the kernel hyperparameters O. Prior work using GP-SDEs with structured GP kernels [14] found that alternating between updating ((u) and optimizing 0 was prone to getting stuck in local maxima of the ELBO. Hence, we adopt the kernel hyperpararmeter learning approach proposed by Hu et [14], which uses the update
0
arg m18x
m1x Eapprox (( ((u))))))()))))) 00] v(u)
The inner maximization can be performed via the closed-form updates as discussed above. We perform the outer maximization, which depends both (()) and Lappoxx-CP (qgr) , ( 0), via Adam optimization. In practice, we find that this kernel hyperparameter learning method performs well on both synthetic and real datasets.
K.4 Posterior inference of the drift at unseen points
The sparse approximation (() allows to compute the posterior distribution f f(x*) at any unseen latent space location x" E PRP Recall that
(l = II plSa M(S und O)q(ua)dus. d=1
Since quua)) Gaussian and p(fī [ uhd ②) is conditionally Gaussian, then ((f") is also a closed-form Gaussian distribution. This takes the form,
(f2) = (fa ka K_ m =o(x",xx") ) - ka== Ker + ka K 22 S. K k
K.5
Incorporating inputs
For our application of SING-GP to real neural data in Section 5.4, we incorporate external inputs that may affect the underlying latent trajectories. Following Hu et [14], we model this as
dxx(t) = (f(x()) + Bu(t))dt + dur)((
where vtt) E / is time-varying, known input signal and B E Dx is a linear transformation of the input into the latent space. We denote 0; == v(T.) where T; time point in the grid We derive the optimal update B" that maximizes Capprox-Gi as
T-1 Σ
(
(AL
T-1
B
(m:+1 m. - q(e Do
>
im
L Experiment details and additional results
For all experiments, we fit our models using either an NVIDIA A100 GPU or NVIDIA H100 GPU.
L.1 Inference synthetic data
Here, we provide additional experimental details to supplement Section 5.1. For this section only, we will abuse notation and use (t) denote individual coordinates of xxt) E R for 0 t < Tmax- This not to be confused with time steps of the grid т, as introduced in Section 2.3.
L.1.1 Linear dynamics, Gaussian observations
Details on simulated data We generate 30 independent trials from SDE with linear dynamics
f =
0.997
cos 0
sin sin 0 cos 0
-1)
At
with x x/250. To simulate these trials, we use an Malle-Maruyama discretization with A/ 0.001 and Tmax = 1. Initial conditions are sampled uniformly in [-2, 2]2. We then simulate observations ys ~ (CCx, + d, R) every time point т,, where the entries of C, d are sampled as i.id. standard normal variables and we set R 0.35I.
38

---

Details on model fitting For our main result in Figure 2 (top row), we fit SING with 20 iterations and step size p with hyperparameters + fixed to their true generative values. As expected. convergence in this fully conjugate setting occurs after only 1 iteration. We use analogous settings for VDP: 20 variational EM iterations, each consisting of single forward/backward pass to update variational parameters. For our Kalman smoothing baseline, we use the Dynamax package [54]. For Adam-based optimization, we sweep over initial learning rates {10-4.5 --5100+ 10-+, 10 5. 1003, 10- fit for 20 iterations, and choose the best fit according to the final ELBO value, which in this case corresponded to learning rate 5 10-3
Computing latents RMSE follows:
The latents RMSE metric from Figure 2D and H are computed as
1/2
latents RMSE T+1 Σ
B
a. true
[1)]
Since (3x,)) Nar, m, St), this can be computed analytically using the identity
x 1
tr(S.)
m: x
L.1.2 Place cell model
Details on simulated data We generate 30 independent trials from an SDE whose drift is charac- terized by a Van der Pol oscillator.
1 (0)(1) (1) x02(1)
((x(()) 7/ 4 (1)
f2(x(()) = 27 (1)
/
with 7 10 and μ 2. To simulate these trials, we use an Maler-Maruyama discretization with At = 0.001 and Tmax = 2. Initial conditions are sampled uniformly in 3|3.
For our observations, we use Poisson count model inspired by spiking activity in place cells. To model the tuning curves (i.e., expected firing rates) for this model, we use radial basis functions. For N neurons, these take the form
(- -
Tm (x: = a cxp
2 | - c,112 + au. n n=
where [ca h= C R are centers placed along the latent trajectories, & controls the width of the curve, and a, 0g control the peak and background firing rates respectively. For experiments, set f = 0.5, a = 2.5, and ag = 0.25. We then generate Poisson spike counts for each time bin т; according to
(T, ~ (xoi,,
In this likelihood model, the expected log- likeliieehood term IE, [logg(yyl()) is not available in closed form, so we approximate it using Herus- quadrature.
Details on model fitting For our main result in Figure 2 (bottom row), we fit SING for 500 iterations. We use the following step size schedule inspired by [25]: p log-linearly increased from 10-3 to 10-1-5 1.5 for the first 10 iterations, and then kept at 10 1.5 for the rest of the iterations. We fit VDP for 500 iterations and use the "soft" updates in eq. (32) ensure numerical stability. After manually tuning the learning rate, we find that VDP performs best with w = 0.05. For the conditional moments Gaussian smoother (CMGS), we use an implementation from Dynamax [54]. For Adam- based optimization, we sweep over initial learning rates (10-1, 5 10 5 10-3 10-2), 2), fit for 500 iterations, and choose the best fit according to the final ELBOO value, which in this case corresponded to learning rate 10
10-1.5
39

---

L.1.3 Embedded Lorenz attractor
Details simulated data We generate 30 independent trials from a D-dimensional latent SDE whose first 3 dimensions evolve according to a Lorenz attractor drift, and the remaining dimensions random walk. These drifts are characterized by
fi(x(t)) a(x2(() (t))
f2(x(()) = 01 (1) xz(()) - x2(())
(39)
fa(x((1) æ (1)=2 (t)
fa( = O,
where (a, 8) (10, 28, 8/3). To simulate these trials, we use an Suler-Maruyami discretization with At 0.001 and Tmux 1. Initial conditions are sampled from independent standard normal distributions. We then simulate 100-dimensional observations yi ~ (Cx; + d. R), where the entries of C and are sampled i.idd from (0, 0.01) and we set 0.05I.
Details model fitting For results in Figure different methods of approximating transition expectations: (1) Monte Carlo and (2) Gauss-Hermite quadrature. For each method, fit SING for 1000 iterations, keeping output parameters and prior drift parameters fixed to their true values. The step size is log-linearly increased from 10 3 to 10-2 for the first 10 iterations, and then kept at 10-2 for the of the iterations. For the Monte Carlo approach, we use a single Monte Carlo sample per expectation. For the quadrature approach, we use quadrature nodes per latent dimension for = 3 and 5; beyond D 5, quadrature results in out-of met-or- errors on a NVIDIA A100 GPU.
fit SING for D 3, 5, 10, 20, using two
L.2 Comparison to neural-SDE variational posterior
Background on SDE Matching SDE Matching, proposed by Bartosh et al. [41], is an algorithm for performing joint (approximate) inference and learning in the latent SDE model. The authors approximate the intractable posterior distribution with a neural stochastic differential equation
((x)) (Ir f(x(()))d /1/)
x (0) Nmm,, Sa), 0 < S Tmax,
where á are the parameters of the variational posterior and f(ll) is neural network drift function. Unlike SING, SDE Matching also applics to models with state-dependent diffusion coefficient. Compared to prior works that perform VI with neural-SDE variational families, SDE Matching is advantageous since it does not require drawing samples from the variational posterior.
Also dissimilar from SING, SDE Matching is an amortized VI algorithm, meaning that the varia- tional approximation is "shared" across samples (t)) In particular, the posterior drift f(x (t))ø) is modeled also as function of the observations (f) (t.))
Details on simulated data We compare SING to SDE Matching [41] on the three-dimensional stochastic Lorenz attractor benchmark dataset as described Li al. [11]. particular, we generate latents according to the three-dimensional SDE on 0 < t S Tmax. Tmax = 1 with drift function
(z((tt) = a<z2(t) (t))
f2(x(t)) (1) xz(()) - x2(
(40)
fs(x(t)) = (1(()=2(t) -Bzz(t)
for (a, p, B) (10, 28, 8/3). diffusion coefficient 0.15, and initial condition (0) I). We sample the SDE according to the Malur-Maruyama discretization with At = 0.00025. We observe evenly spaced grid of size 0.025, where C
(t:)) according (yy(t.) Car(t,) + d. (0.01)²T) and d are also chosen as in et al. [11]. We sample 1024 total trials.
Details on learning Following Bartosh al. [41], we model the prior p(x) as a neural-SDE whose neural network drift has a single softplus activation. However, unlike Li et al. [11], Bartosh et al. [41] who consider four-dimensional latent space, we consider a tiree-dimensional latent space so that can directly assess recovery of the ground truth Lorenz attractor dynamics.
40

---

For SDE Matching, we use the authors' publicly available codebase, which includes code for evaluation on the same benchmark; we do not further tune any hyperparameters. For SING, we perform 10 iterations (10 E-steps, 10 M-steps each), and for SDE Matching we perform 101 iterations; these choices ensure that the number of gradient updates is the same for both algorithms.
To perform the E-step of the vEM algorithm with SING, we define the time grid T using evenly spaced points with At 0.0025. We increase the NGVI step size p on a log-linear scale for 100 iterations from 10 - to 10 and hold it constant for subsequent iterations. To perform the M-step, we update both the output parameters and the drift parameters using Adam. We hold the learning rate for the drift parameters constant at 10 33 And for learning the output parameters, we decay the learning rate from 10- 3 10-7 7 for 500 iterations and then hold it constant. To compute expectations of the prior drift under the variational posterior, we use single Monte Carlo sample.
Latents RMSE
Dynamics RMSE
# Samples
SING
SDE Matching
SING
SDE Matching
10
.50 = .05
= .07
.25 004
.49 + .05
30
.52 .05
.78 .03
.45 .38
.25 - .04
50
.49 + .05
.70
31 .06
.23 .03
100
.50 .04
:63
:19
.19 .01
300
.46 + .05
.0
+5
.17 .01
500
.48 .05
.25
.17
1024
.58 - .08
.57
.56 .44
.17 .01
Comparison of SING and SDE Matching, A comparison of latents and dynamics RMSE on the stochastic Lorenz attractor dataset for SING and SDE Matching. Errors represent 2SE computed over 10 random initializations of the prior drift and output parameters and random seeds for the algorithms.
Figure 7:
Evaluation metrics We compare the performance SING and SDE Matching using two metries: the latents and normalized dynamics root mean-squared error. We compute the latents RMSE exactly as in Appendix 1.1.1. And for the normalized dynamics RMSE, we 1024 draw independent samples (t)) from the ground -ound-truth stochastic Lorenz attractor and estimate
( (inon
[[(x(()) f(((()))))) f(a()))) [5(x(())12
1/2
-(a())))(((()))))) [[(()(()))
1/2
E
trials
where == { is the Muler-Maruyamu sampling grid. Note that the samples on which the normalized dynamics RMSE is computed are different from those used to perform joint inference and learning (with either SING or SDE Matching).
Results Our experimental results, summarized Figure 7, provide compelling evidence that the structured Gaussian process variational approximation, combined with the fast and numerically stable SING inference algorithm, leads to more accurate recovery of the ground truth latents compared to SDE Matching, In particular, SDE Matching requires approximately 500 samples to match the latents RMSE that SING achieves with only 10 trials. This is likely because at each iteration SDE Matching updates the posterior using only a single time in [0, ] whereas SING updates the variational posterior on the entire grid T.
However, SDE Matching outperforms SING in recovering the ground truth Lorenz attractor dynamics. In our experiments, we observe that while performing vEM with SING can achieve dynamics RMSE comparable to that of SDE Matching, the dynamics RMSE varies greatly depending on the random initialization of the prior neural network and output parameters. Improving learning with SING remains future research direction.
41

---

L.3
Drift estimation on synthetic data
Details on simulated data
The Duffing equation is a second-order ODE described by
d'v
du d
ov + Bu 0, 0 s S T max.
d12
4t we can convert the Duffing equation into system of first-order ODEs. Moreover, we add an independent Brownian motion with noise a 0.2 each coordinate. Altogether, this yields the two-dimensional SDE with drift
By setting σ1
22
du
fi(x(()) = xt( fz(x(()) = az1 (t) - 3z1 (t)( -100--0000 (t)
(41)
is straightforward to verify that the Duffing equation has three fixed points when a/3 > 0: one at the origin and two at /a/3 0). We generate four trajectories from the Duffing SDE eq. (41) with parameters (a, 8, 7) = (2, 1, 0.1) and initial condition x00) (va/3 0.1, 0.1), i.e., near the rightmost fixed point. We sample the solution to the SDE using the Euler- Marlyama discretization with Timx = 15 and At = 0.015. All 10 time steps are taken to define the grid T. A visualization of two of the sampled trajectories is provided in Figure 3A.
From the 10 total time steps, we sample 3 102 at random at which to include Gaussian observations + d, R). The true emissions parameters C IR Nx and d E R. are chosen by sampling i..dd standard normal random variables. R is fixed to a diagonal matrix with entries 0.1.
Details on learning We simultaneously perform inference and learning in the latent SDE model according to the variational EM (vEM) algorithm. For each variational EM step, we take 10 E-steps (according to either VDP or SING) as well as 50 M-steps. In each M-step, we update both the drift parameters O as well as the output parameters {C, d]. We update the drift parameters by performing Adam on the approximate ELBO Laparox. and we use a closed-form update for the output parameters (see next section). In the case of the GP prior drift, the only learnable parameters are the length-scale and output scale of the RBF kernel. The updates to the posterior over inducing points p( (u) are performed during the E-step.
prirr
For the GP prior drift we use 12 x 12 grid of induc- ing points on |-6, 2 To choose the size of this grid, we looked at the first two principal component scores of the observations (tt(t, (/=r For the neural-SDE, we use a neural network with two hidden layers of size 64 and ReLU activations. And for the polyno- mial basis drift, we consider all polynomials up to order three {1, I1, 22, 22 xं ......J. Note that by the form of the Duffing equation eq. (41), this polynomial drift function can, in theory, perfectly represent the true drift.
A
B
Neural-SDE drift
Polynomial basis drift
-2
-2
X1
-2
Figure 8: Prior drift functions learned via vEM with SING and Adam. A: Learned neu- ral network drift. B: Learned polynomial ba- sis drift.
We initialize the output parameters C' and d by taking the first two principal components of the observations and the mean of the observations, respectively. The GP length scales and output scales are initialized to one; the neural network weight matrices are initialized by i.i.d. draws from a truncated normal distribution and the biases are initialized to zero; the coefficients on the polynomial basis are initialized by i..dd draws from (0, (0.1)²).
When using SING during the E-step, we increase p on a log-linear scale from 10 3 1 during the first ten iterations, and we then fix it at 10-1 for the remaining 40 vEM iterations. We found this procedure of starting with P small and subsequently increasing it on a log- -ineer scale during the first k iterations to result in effective inference and learning [25]. As for VDP, we choose the soft update parameter w to be as large as possible while maintaining the stability of the algorithm. For the GP prior and neural network drift classes, this selection procedure yielded w = 0.05. For the polynomial basis drift, we were not able to obtain convergence with any reasonable choice of w (see Appendix H for discussion of the parameter w). For this reason, it is omitted from Figure 3.
42

---

In Figure 8, we visualize the learned dynamics corresponding to the neural network and polynomial basis prior drift functions. The posterior mean and variance of the GP prior drift are visualized in Figure 3B. When performing inference with SING, both parameteric drift functions, in addition to the GP posterior mean, match the drift eq. (41) of the Duffing equation (Figure 3D). We measure the discrepancy between the learned and true drift by computing the dynamics root mean-squared error
1/2
dynamics RMSE
Σ w(b))lf(o) f ((( VEV
(42)
Here, we take a fine grid of cells, each with equal D-dimensional volume, covering the latent space; denote the collection of their centers by V. Then, compute the proportion of true latent trajectories that fall into the cell with center v, written w(v), such that Evev w(v) = Mathematically, eq. (42) approximates
1/2
trials
Tum
(1()))))
trials Tmax
j=1
where trials represents the number of trials and (x (t))sstsT. is the true latent trajectory corre- sponding to the jth trial.
Closed-form updates for output parameters Next, we derive closed- form updates for the output parameters C E NxD d
R.N. R E RNN in the Gaussian observation model
ply x θ) II Cx((t) + R),
where assumed be diagonal. These same update equations appear Hu al. [14]. Although not directly mentioned in the main text, we can also learn the entries of the emissions covariance matrix R.
Writing out the expected log-likelihood of ply (a)) under g(ar), we have
blllgg v 2)] log
R
Ey
Cx((.) dl
+ const
1m1
=-2 log
R
nd R d
y(tt)) R 'Cm(()( + + d) - 1 -uitt) R vt())
2
2 C
R-'C (((t)) + mlt m(t:mm(t)) d R (m((tt))) + const,
where const represents generic constant with respect (C, d, R). Differentiating with respect to C, d, and R yields the update equations
c
Σ ((1)) + m(()mm(())))
(m(())) - d' d
(43)
d' n Σ (v(()) c'm(t.)
(44)
R Σ (v,))))² (2, m(t.) ((C)) m(t.))² + (C7) S(NC)) (W)))))(()) -d))
where Ry is the jth diagonal entry of R and Cj is the jth row of C.
43

---

By defining y = n -i=E
ii= v(4)) as well as m = =n
- -n'E== m((t)) and plugging eq. (44) into
eq. (43), we obtain
c
Σ ((L) + (mi(t)) - m)(((t))) - -m)") =mmmm
Σ (v((s) (m(((())) m)) mm
- (u(tt)) - (m((t) -m)
((m(t)) m)(m(ti) - -m)" + s(ti))
From the equation for C is straightforward to solve for d' and Ri 1 isis
Discretization experiments In Figure 3E, we investigate the effect of the discretization size At on the performance of the vEM algorithm with SING and VDP. Recall that our true latent trajectories are sampled with Euler-Maruyama step size 0.015, and so we would not expect Lapprox to better approximate Lcont for smaller discretizations т. In the case of VDP, though, we demonstrate that taking At smaller than 0.015 leads to improvements in the stability of the algorithm. Moreover, for large Át, performing vEM with VDP is highly unstable and may diverge depending on the neural network initialization. For this reason, we exclude the dynamics RMSE corresponding to At > 0.15.
In order to coarsen i.c.. downsample the original grid т, fill in grid points where observations were not observed with zeros. Then, we consider each jth observation in r. that observation is take the nearest observation, should there
zero (i.e., there is not observation at that grid point), consecutive points in the coarsened grid. exist an observation in the original grid that lies between This procedure increases At j At. In order to make the original grid finer, i.e., upsample, add j zeros between the existing observations, again treating missing observations as zeros. This procedure decreases Al A + 1). Note that coarsening the grid may lead to fewer observations, but making it finer will neither increase nor decrease the number of observations.
L.4 Runtime comparisons
Experiment details We perform runtime comparisons simulated datasets consisting of 2D linear dynamical system (LDS) with Gaussian observations. Each point in Figures 5 and 9 represents the average wall clock time (in seconds) over 5 runs of each method randomly sampled dynamical systems. All runtime comparisons were performed on an NVIDIA A100 GPU.
We compare two methods: (1) 'parallel' SING, which uses associative scans as described in Sec- tion 3.4 to perform natural-to-men to-mnann parameter conversion, and (2) 'sequential' SING, which uses the standard sequential algorithm to perform this conversion, as we detail in Appendix G.1.
For each run, randomly generate the data follows. We construct linear drift function of the form
f(x(t)) Ax +
where A _ is constrained have negative real eigenvalues for stability and b E R has small non-zero entries. Using Buler-Maruyann discretization with At = 0.001, we sample for T time steps from this 2D linear SDE. Then, for each time step Tú. we sample 10D Gaussian observations,
(T,)) (((Cx(T..) d. R)
where and d are randomly sampled with i.i.d. standard normal entries and = 0.35I. For each run, we fit the corresponding method (either 'parallel' or 'sequential') for 20 iterations, keeping hyperparameters é fixed.
Additional results on different batch sizes In Figure 9, we provide additional results from runtime experiments varying batch sizes. Here, show wall clock time for batch sizes of 1, 20, 50, and 100 for each method. We find that for a single trial, parallel SING exhibits nearly constant scaling with sequence length, while sequential SING scales linearly in time, expected.
For larger batch sizes, our implementations of parallel SING and sequential SING both parallelize over the batch dimension. Thus, it is expected that parallel resources single GPU may be
44

---

parallel
Sequential
D 1, trial
= 2, trial
D = 3. 1.3.00 trial
10a
10a
D = 1, 20 trials
D 2, 20 trials
D 3, 20 trials
10
Wall clock time b9)
10a
D 1, 50 trials
D 5, trials
D 3, 3,50 trials
10
10a
1, 100 trials
D 2, 100 trials
D - 3, 100 trials
10a
10a
512
4096
512
4096
512
4096
timesteps
Figure 9: Runtime comparisons between parallel SING (using associative scans) and its sequential counterpart on NVIDIA A100 GPU. Parallel SING scales favorably with sequence length compared to sequential SING on different latent dimensions and batch sizes up to 100 trials.
constrained for models that operate on both large batch sizes and long sequences. Figure we show that despite these constraints, parallel SING still scales favorably with respect to sequence length compared its sequential counterpart for batch size up to 100 trials. We also note that our implementation of SING supports mini-batching smaller batch sizes in order to maintain this favorable scaling behavior, even on datasets with many (i.e., hundreds or thousands of) sequences.
L,5 Application to modeling neural dynamics during aggression
Dataset specifications We apply our method to publicly available dataset from Vinograd et al. [55], which can found at: https:// //dp::://dandaae aar archive. org/ http/ddnisee///eeeeeeeeeaaaesesd sonddiist/00100/ /001037. We use the dataset labeled as 'sub-MOL' which consists of neural activity from one trial of mouse exhibiting aggressive behavior towards two consccutive intruders. This dataset consists of calcium imaging from ventromedial hypothalamus neurons. In total, 56 neurons
imaged over 420.1 seconds 10 Hz, resulting in data matrix with = 4201 and N = 56. Before fitting models, we z-score the data such that each neuron's 's activity has mean 0 and standard deviation 1.
GP-SDE model details We fit GP-SDE models with kernel function chosen to be the "smoothly switching linear" kernel from Hu et [14]. This kernel takes the form
x(x, x²) = Σ (( TM c,) o Tú ( T
j=1
(5 (n,a'))
pat
(mm
where
exp w 6(x)/T)
(
1 Σ == exp (w) 6(x)/T)
45

---

A
B
gPSLDS likely of dynamics
gpSLDS inferred posterior variance
12
12
10
-5
11 -5
-9
X1
X1
11
Figure 10: Additional results for the same SING-GP model fit as in Figure 6C. A: Inferred latent trajectory, dynamics, and learned input effect using SING-GP. The dynamics arrows are colored by most likely linear dynamical state, as learned by the smoothly switching linear kernel from Hu et al. [14]. The orange arrow represent the direction of learned inputs (i.e. columns of B) capturing the effect of the intruders' entrances in the mouse's environment. As expected, these arrows explain the fast change in latent trajectories during intruder entrances. B: Same inferred dynanics as panel A, but the latent space is colored by inferred posterior variance of q(f) from SING-GP model fit. The model expresses low uncertainty in regions with data and high uncertainty in regions without data.
Intuitively, this kernel encodes prior structure in the dynamics such that it "switches" between different linear dynamical systems based on location latent space, but maintains smooth dynamics between different linear systems. This is enforced by the linear kernel Kinn
()) (az, æ²) and partition kernel "opu (xx x²) for each linear regime j. Above, J is the total number of discrete states that the function switches between; we choose J 4 and D to be consistent with prior work [56, 55, 14]. We
also choose d(x) = [1 21 x2] The kernel hyperparameters are = {M, o {ej, w T. which we optimize via Adam [52] in variational EM. In particular, for learning we use the approach proposed by Hu et al. [14], which optimizes hyperparameters jointly with the closed-form posterior drift update. Finally, model observations
yt.) ~ NC(Ca(t) + d, R)
for every observation time f;. Unless otherwise specified, we fit models whose discretization matches the imaging rate; i.c., At = 0.1.
Model fitting details To fit GP-SDE models using both SING-GP and VDP-GP, we chose inducing points to uniformly spaced 8 x grid in [-8, x [-4, 12]. These bounds were chosen by examining two-dimensional latent trajectories obtained by principal component analysis (PCA) on the neural data matrix. Similarly, we initialize C' and d using PCA and keep them fixed during model fitting. We initialize kernel hyperparameters randomly from i.i.d. standard Gaussian distributions.
For both SING-GP and VDP-GP, we perform variational EM for 30 iterations. In each E-step we perform 15 updates (natural gradient steps for SING-GP and forward/backward passes for VDP- GP). In each M-step we perform 50 Adam iterations with initial learning rate 1e-4 to learn kernel hyperparameters.
For SING-GP, we again learning rates for the first 10 iterations to be log-linearly spaced between 10-1 and 10 We then keep the 10 learning rate for the rest of fitting. For VDP-GP, use "soft updates" (eq. (32)) and manually tune w to the largest value which still maintains numerical stability; in this case w = 0.005.
suggested learning rate schedule from [25]. In particular, set the
Since the smoothly-switchin repeat our experiments for different random initializations of O In Figure 6D-E, we report average metrics and 2-standard error bars across these 5 initializations.
linear kernel from Hu et al. [14] has several hyperparameters,
Computing posterior probability of slow points Figure 6C, we use the probabilistic formulation of the drift in the GP-SDE to identify gigh-probability regions with slow points. To do this,
46

---

compute a "posterior probability of slow point" metric for fine grid of locations in latent space. For a particular R we define this as:
D II
F( is a slow point)
P (P < d=1
where e > is chosen to be a small threshold. Intuitively, this probability large if the inferred drift at x" is close to 0 in the posterior. Recall from Appendix K.4 that ((a((=")) is closed-form Gaussian distribution, so this quantity can be easily computed.
Experiment details for Figure 6D In Figure 6D, compare SING-GP and VDP-GP on two metrics as a function of discretization A, ranging from 0.06 to 0.1. We fit both methods using 100 variational EM iterations.
The first metric for model comparison is final expected log-likelihood; precisely, this is E, So, æ)] at iteration 100. We choose this metric to be able to directly compare the goodness- of-fit for each model; in contrast, the ELBOs are not directly comparable due to their slightly different formulations of the term DEL. (g(x) || p| f))
The second metric for model comparison is the number of iterations until convergence. Here, we say that the model has converged iteration j if ()) @) L(g(i-1), @) 200 nats, where [(-) denotes each model's respective ELBO function. We note that other choices of convergence threshold yielded similar trends the plot in Figure 6D.
Forward simulation experiment Here, we provide details on the forward simulation experiment in Figure 6E. To compute the predictive R² metric in Figure 6E, we carry out the following steps.
(a) For a given model fit and discretization At, select 20 uniformly spaced time points tn along the inferred latent trajectory, well as their corresponding marginal means, m E,, [=](()]]-
(b) For each to and m(fo), use the posterior mean of inferred dynamics q(f) to simulate forward the latent state for seconds using discretization At, where = 1,2...., 15.
(c) Forward simulation will yield predicted latent state, #t(0o + s). We compute the predictive R² for each choice as,
R2 1 Eeo (to + 5) yto + s))
E, y + s) 0
where itu + = C# (tu 5) + d, ÿ is the mean of all observations in the trial, and the sums are taken over all 20 initial times tg-
Additional results In Figure 10, present additional experimental results to supplement the SING-GP inferred flow field and line attractor from the main text, Figure 6C. These results provide new angles of analysis of the SING-GP model fit. Recall from above that the smoothly switching linear kernel from Hu et al. [14] encourages smoothly switching linear dynamics to be inferred.
In Figure 10A, flow field arrows are colored by most likely linear state. That is, each set of colored arrows represents a different linear dynamical system. Encoding this structure into the kernel allows for convenient downstream analysis of these inferred dynamics, in particular using established properties of linear systems. Moreover, the orange arrows in Figure 10 represent the columns of the learned input effect matrix B Dx2, We find that these input directions correspond to the direction of the latent trajectory at both timepoints at which an intruder enters.
In Figure 10B, we plot the inferred posterior variance of the same SING-GP model. We find that model uncertainty low in the regions traversed by the inferred latent trajectory and high otherwise. In fact, this posterior variance is exactly the quantity that we use to compute posterior probability of slow points in Figure 6C. This demonstrates how the probabilistic formulation of the GP-SDE, which we can now perform reliable inference and learning in using SING-GP, allows us to better understand and interpret inferred flow fields.
47

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We propose SING, a practical method for inference latent SDE models, We describe our proposed method in Section 3, including our main theoretical result, and provide empirical results for our method in Section
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: We clearly state all assumptions of our model and, in particular, the variational approximation to the posterior in Section 2.3 and Section 3.1. Furthermore, we discuss
limitations and suggest directions for future work in Section 3.3 and Section 6.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged
the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
48

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: We provide the full set of assumptions and the proof for our main theoretical result (Theorem 1) in Appendix E.
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide detailed descriptions of our experiments in Section 5 and Ap- pendix L.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access hosted model (e.g.. in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
49

---

5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: We provide details to reproduce experiments in the supplemental material. We also provide link to our open-source codebase containing the implementation of the method.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nipg cttp//ccpcccc///////////c/.c/cct cc/
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: // //npps.cc/ppppppppp htc//pbllc/gguudeeeo
Sodeeuumissiia wwwisssssssiisoooioo for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide all experimental details in the main text (Section 5) and the appendix (Appendix L).
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Yes, we provide thorough quantitative comparisons with error bars for the experiments in Section 5.3 and Section 5.4.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
50

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
Question:
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification: Yes, we provide information on the hardware used for all experiments and perform runtime experiments in Section 5.3.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that NeurlPS Code of Ethics https : /ttp///eerrrp.eeeee hc/pppbliccthhicsssiedelll..
didn't make into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the
Answer: [Yes]
Justification: The authors have reviewed and conducted research that conforms to the code ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [Yes]
Justification: We discuss the broad applicability of our method throughout the paper, and focus on particular application neuroscience in Section 5.4. We do not expect negative
socictal consequences.
51

---

Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper does not involve data or models that have a high risk for misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All code used in this paper owned by the authors. All datasets are either generated by the authors or properly credited in Section 5.4.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
52

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswithcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
Our open-source codebase linked in the paper is well documented.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
This paper does involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
53

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
The development of our method did not involve the use of L.MSs in an original way.
Justification:
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy (https: /ttp::///euuipppeee cct.CoCeeeeeecccs/222325. for what should or should not be described.
54


Reviews:
[{"review_id": "ILSC74Oxpr", "reviewer": "Reviewer_ma3P", "summary": "The paper introduces a new method for inferring (latent) stochastic differential equations based on variational inference of natural parameters of probability densities represented in terms of families of natural parameters.\n\n\nThe work builds on the natural-gradient formulation of variational inference and extends existing literature in two directions. First, it provides a derivation for the closed-form natural-gradient updates for a latent SDEs discretised with the Euler–Maruyama scheme, exploiting the Gaussian structure of the discretised state marginals. Second, it introduces a parallelisation scheme for converting natural-to-mean parameters\nvia associative scans that reduce the computational cost from $\\mathcal{𝒪}(D^3 T)$ to $\\mathcal{𝒪}(D^3 log T)$ enabling thereby learning of long sequences.\n\nThe authors prove that the evidence lower bound (ELBO) computed on a finite grid converges to its continuous-time counterpart linking thereby the discrete-time and continuous-time objectives.\n\nThe authors validate their approach empirically on two synthetic datasets: one created with a latent linear dynamical system with Gaussian observations, and one with latent dynamics from a Van der Pol oscillator coupled with Poisson observations. In both settings, the proposed algorithm converges more quickly and recovers the latent dynamics more accurately than existing variational and filtering baselines. The authors further compare the approach against the variational diffusion process of Archambeau et al.. They further demonstrate their method on calcium imaging data for on modeling neural dynamics during aggression. \n\n \n\nWhile natural-gradient variational inference for SDEs is not entirely new, see for example  Verma et al. (2024), the present work differentiates itself through the parallelisation acceleration and the convergence analysis, both detailed in the appendix.\n\nI find the paper well fit for the NeurIPS audience.", "questions": "- How would the method scale computationally with increasing latent dimension and/or increasing observation dimension (number of observed neurons)?\n\n- For the SING-GP, how do you select inducing points? How sensitive is the performance of the method to the number and the method of selection of inducing points?\n\n- Is there a reason you do not compare the synthetic datasets with the VDP-GP?", "limitations": "the authors discuss the limitations of their approach in discussion and in the appendix.", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths:**\n\n- The paper tackles a challenging problem, that of inference of latent stochastic systems, where the posterior inference is intractable for general systems.\n- The authors provide rigorous derivations.\n- They provide empirical evidence on both conjugate and non-conjugate models, as well as on synthetic and experimental datasets. \n\n\n**Weaknesses:**\n\n- The authors demonstrate the results relying only on measurements up to 10 neurons. While impressive I wonder how the performance and computational demands of the method would scale if the considered neurons were on the order of hundreds.\n- Variational inference with natural gradients has been introduced in previous work.", "was_revised": true, "final_justification": "All reviewers were quite positive to accept this paper. The concerns raised during the review were minor technical questions that the authors addressed either by providing additional information or by providing additional numerical results.", "rebuttal": "We thank the reviewer for taking the time to read our submission and for their thoughtful response! We are especially glad that they found our work to be rigorous and a good fit for the NeurIPS audience.\n\n## Weaknesses\n>  I wonder how the performance and computational demands of the method would scale if the considered neurons were on the order of hundreds.\n\nAs with most variational inference methods, the computational scaling with respect to observation dimension primarily depends on the structure of the likelihood $p(y | x)$, which is used to compute $\\mathbb{E}[\\log p(y | x)]$ in the ELBO. For our work, we consider a GLM likelihood (e.g., Gaussian, Poisson) that factorizes over the observation dimension. In this case, the complexity of the expected log-likelihood scales linearly with the observation dimension and can be computed in parallel across it.\n\nWe note that SING could also be applied in settings with more complex or non-factorized likelihoods, though such cases may introduce additional computational overhead. This tradeoff, however, is not specific to SING; it arises in virtually all variational inference frameworks from the expected log-likelihood in the ELBO.\n\nWe conduct additional experiments that test the scalability of SING with respect to the observation dimension on synthetic data. We first simulate ground-truth latent trajectories from the SDE $dx(t) = f(x(t)) dt + dw(t)$ where $f(x)$ is a Lorenz attractor with parameters $(\\rho, \\sigma, \\beta) = (28, 10, 8/3)$. We then simulate observations $y_n(t) \\sim \\mathcal{N}(c_n^\\top x(t) + d_n, (\\sqrt{0.05})^2)$ for $n = 1, \\dots, N$, where $N$ is the observation dimension (i.e. number of neurons). Here, the entries of $c_n$ and $d_n$ are sampled *i.i.d.* from $\\mathcal{N}(0, 0.1^2)$. We fit SING using a prior SDE in the same class as the generative model and learn the drift parameters. We report the RMSE between the true and inferred latents, as well as the RMSE between the true and learned dynamics at the locations of the true latents. We also report runtime in seconds on an A100 NVIDIA GPU. Each entry reports the mean and 95% confidence interval obtained from $5$ random initializations of hyperparameters.\n\n| obs dim ($N$) | latents RMSE         | dynamics RMSE        | runtime (s)  |\n|---------|----------------------|-----------------------|--------------------|\n| 10      | 0.2637 ± 0.0004      | 0.3245 ± 0.0674      | 57.02 ± 0.15     |\n| 25      | 0.2058 ± 0.0001      | 0.1523 ± 0.0025      | 57.08 ± 0.20     |\n| 50      | 0.1643 ± 0.0001      | 0.1726 ± 0.0012      | 58.05 ± 0.37     |\n| 100    | 0.1416 ± 0.0001     | 0.1315 ± 0.0028     | 58.33 ± 0.40     |\n| 200    | 0.1211 ± 0.0001     | 0.1446 ± 0.0016     | 58.22 ± 0.21     |\n\nWe observe that both latents RMSE and dynamics RMSE generally improve with greater observation dimensions, likely because seeing more data improves performance. Furthermore, the runtime remains effectively constant. This is due to computing the expected Gaussian likelihood in parallel over observation dimension. While extremely high-dimensional observations (e.g., thousands of neurons) could eventually strain parallel resources, we stress that standard GLM likelihoods that factorize over observation dimension remain highly scalable in practice.\n\n> Variational inference with natural gradients has been introduced in previous work.\n\nWe agree with the reviewer that natural gradient variational inference has been introduced in prior work, and we certainly build upon it in our work here. Our contribution lies in deriving closed-form natural gradient updates specifically for latent SDEs, enabling parallelization of these updates for scalability to long sequences, providing theoretical bounds on the ELBO approximation error, and extending this approach to inference in models with nonparametric (Gaussian process) drift. \n\n## Questions\n> How would the method scale computationally with increasing latent dimension?\n\nThe bottleneck with respect to latent dimension lies in computing expectations of functions of the drift as described in Section 3.3. For the experiments performed in our submission, we used Gaussian quadrature, which scales poorly with latent dimension (i.e. the number of quadrature nodes necessary to achieve fixed approximation error increases exponentially). \n\nMoreover, we have conducted additional experiments using Monte Carlo to approximate these expectations, which we find enables efficient scaling to higher-dimensional latent spaces while maintaining accurate inference.\n\nFor our additional experiments, we apply SING with Monte Carlo to systems with latent dimension $D = 3, 5, 10, 20, 50$. For all of these settings we simulate ground-truth latents from an “embedded Lorenz attractor” whose first 3 dimensions evolve according to an SDE $dx(t) = f(x(t)) dt + dw(t)$ where $f(x)$ is a Lorenz attractor with parameters $(\\rho, \\sigma, \\beta) = (28, 10, 8/3)$, and whose other dimensions $d > 3$ evolve according to a random walk with identity diffusion. We simulate observations $y_n(t) \\sim \\mathcal{N}(c_n^\\top x(t) + d_n, (\\sqrt{0.05})^2)$ for $n = 1, \\dots, N$, where here we fix $N = 100$. As before, the entries of $c_n$ and $d_n$ are sampled i.i.d from $\\mathcal{N}(0, 0.01)$. We report latents RMSE, dynamics RMSE, and total runtime (including compilation time) on an A100 NVIDIA GPU, along with their 95% confidence intervals obtained from $5$ different initial seeds for Monte Carlo sampling. All results are obtained by using a single sample per expectation evaluation.\n\n| latent dim ($D$) | latents RMSE        | dynamics RMSE       | runtime (s)        |\n|------------------|---------------------|----------------------|--------------------|\n| 3                | 0.1416 ± 0.0001     | 0.1241 ± 0.0107      | 58.05 ± 0.16       |\n| 5                | 0.1867 ± 0.0001     | 0.1413 ± 0.0055      | 60.11 ± 0.36       |\n| 10               | 0.2721 ± 0.0001     | 0.3102 ± 0.0059      | 89.20 ± 0.09       |\n| 20               | 0.3971 ± 0.0001     | 0.2043 ± 0.0048      | 136.28 ± 0.18      |\n| 50               | 0.6761 ± 0.0001     | 0.2193 ± 0.0059      | 561.44 ± 0.63      |\n\nGaussian quadrature (with $6$ quadrature nodes per dimension) achieves latents and dynamics RMSE $0.1415$ and $0.1220$ in $D = 3$ latent dimensions and $0.1866$ and $0.1337$ in $D = 5$ latent dimensions.\n\nWe observe that SING applied with Monte Carlo recovers the true latents and dynamics with comparable accuracy to Gaussian quadrature in $3$ and $5$ latent dimensions. Although latents RMSE continues to increase with $D$ beyond $5$ latent dimensions, this is expected for two reasons: (1) the error is computed over higher-dimensional vectors, and (2) accurate inference becomes more challenging as latent dimension increases. Additionally, we find that the runtime of joint inference and learning of SING remains tractable in up to 50 latent dimensions. We will include a description of this Monte Carlo extension along with implementation details in our revised paper.\n\n> For the SING-GP, how do you select inducing points? How sensitive is the performance of the method to the number and the method of selection of inducing points?\n\nIn practice, we typically select inducing points by (1) running PCA (for Gaussian observations) or a discrete-time state-space model (for non-Gaussian observations) to estimate the scale of the latents, (2) choose a grid of inducing points that covers this region of latent space, and (3) initialize model output parameters to those obtained by the dimensionality reduction technique from step 1. \n\nIn general, SING-GP tends to improve in dynamics accuracy with more inducing points, as this increases the flexibility of the variational posterior over the dynamics. However, this comes at the cost of increased computational complexity, since each update of the GP posterior costs $\\mathcal{O}(TM^2)$ where $M$ is the number of inducing points and $T$ is the length of the time grid on which inference is performed. We have not yet systematically explored sensitivity to the number or selection method of inducing points for SING-GP, but we agree that this is an important direction. \n\nSeveral works in the GPLVM literature (e.g., [1], [2]) adopt a similar initialization strategy but subsequently treat the inducing point locations as learnable parameters, often optimized jointly with other variational parameters. These approaches could potentially be incorporated into SING-GP; for instance, it would be possible to learn inducing locations by direct optimization of the ELBO. It may also be possible to incorporate adaptive inducing points selection to SING-GP, as in [3]. We will add a discussion of these future directions to the paper.\n\n> Is there a reason you do not compare the synthetic datasets with the VDP-GP?\n\nWe thank the reviewer for their comment and apologize for the confusion our plot may have caused. In Figure 2, we do indeed compare against VDP-GP (purple dashed line). SING-GP outperforms VDP-GP in terms of recovering the ground truth latents (‘latents RMSE’) as well as the ground truth Duffing Oscillator dynamics (‘dynamics RMSE’).\n\n**Thank you again for your positive comments and thoughtful response!** If you have any further questions, we would be happy to answer them.\n\n– The SING authors\n\nReferences:\n[1] Lalchand et al. Generalised GPLVM with stochastic variational inference. AISTATS,  2022.\n[2] Jensen et al. Manifold GPLVMs for discovering non-Euclidean latent structure in neural data. NeurIPS, 2020.\n[3] Uhrenholt et al. Probabilistic selection of inducing points in sparse Gaussian processes. Uncertainty in Artificial Intelligence, 2021."}, {"review_id": "IOJ9ntAc4j", "reviewer": "Reviewer_bEC3", "summary": "The authors propose to use natural gradient VI for inference in latent SDEs, proposing an algorithm SING where a variational family of linear SDEs is prescribed over the latent dynamics, and the parameters of the drift and link function are learnable. The authors provide several numerical exploits to speed up inference with parallelization and provide illustrative and quantitative examples of their method in simulated and real data.", "questions": "- I am concerned with the linear family for the posterior drift, since this would preclude dynamics that for instance exhibit multiple modes. Similarly, it seems to me that if the prior SDE allow for bistable behavior (e.g. bifurcations in the trajectories), using a linear variational family cannot model this. \n\n- On the other hand, the prior family on the drift could admit bifurcations for some values of $\\theta$. How does this reconcile with the linear variational family for the latent SDE? In particular, should one expect any connection between (1) the latent linear SDE learned by NGVI and (2) the learned vector field?\n\n- The authors mention several related works that use more general families of processes for inference of trajectories. Presumably these are more computationally intensive or less accurate overall. However, are there some scenarios (such as bifurcations) where more flexible model classes may help? Some discussion of this would be valuable I think. \n\n- Can anything be said in the case where the diffusion coefficient has some dependency on the state $x$? For instance, SDEs of the Cox-Ingersoll-Ross type. Such systems may have applications in some biophysics domains, e.g. [1]\n\n- As I understand it, SING infers a (variational) posterior over latent trajectories, but MAP estimate for the parameters $\\Theta$. Is extension to the fully Bayesian setting straightforward, or would it require new techniques?\n\n- For the question on the dimensionality of the latent dynamics, I think it would be valuable (and potentially not too much work) for some empirical results comparing the use of Gaussian quadrature to Monte Carlo for estimation of expectations. It seems to me that the most glaring limitation in the present work is that all examples are 2D; at least in the synthetic examples it would be good to see some results for moderate dimensions (2-5 dims). \n\n\n[1] Frishman A, Ronceray P. Learning force fields from stochastic trajectories. Physical Review X. 2020 Apr 1;10(2):021009.", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\nThe paper is well written and compelling. It provides a pedagogical introduction to Bayesian inference problems for SDEs, natural gradient VI, etc. \n\nThe formulation of SING is rigorous and the use of a linear SDE variational family clearly leads to a statistically efficient and elegant model. This is evident from the mathematical derivations as well as the numerical experiments. \n\nThe extension to SING-GP is a nice addition and in my view it significantly broadens the relevance of the work, since in many applications the underlying class of drift functions would not be known.\n\nThe numerical results presented are polished and clearly illustrate the favourable behavior of the proposed algorithm. \n\nWeaknesses:\n\nIn all the examples shown, the latent dynamics are two-dimensional. The authors mention that this is a limitation of their method. I understand that this arises from the need to compute expectations in dimension $d$. The authors mention they use Gaussian quadrature, but it would be good to have an idea of what would happen if Monte Carlo estimation were used instead. The estimates would be unbiased, so some commentary on whether this stochastic approach is likely to work would be nice. \nIt is furthermore unclear from the related work section whether the baseline methods would also have the same scalability issues to higher dimensional problems.", "was_revised": true, "final_justification": "The authors' rebuttal helped clarify a key misunderstanding I had about the work, namely that the linear SDE parameters are learned for each trajectory rather than across trajectories. I am also satisfied with the additional results on scaling up to higher latent dimensions, the results look very promising and would aid the paper in making its point. I therefore keep original high score.", "rebuttal": "We thank the reviewer for taking the time to read our submission and for their thoughtful response! We especially appreciate the comments highlighting our work as “a pedagogical introduction” to the topic, as well as the rigor and applicability of our approach.\n\n## Weaknesses\n> The authors mention they use Gaussian quadrature, but it would be good to have an idea of what would happen if Monte Carlo estimation were used instead.\n\nWe thank the reviewer for their suggestion and have conducted additional experiments demonstrating that our method scales well with increasing latent dimension by using Monte Carlo.\n\nWe apply SING with Monte Carlo to systems with latent dimension $D = 3, 5, 10, 20, 50$. For all of these settings we simulate ground-truth latents from an “embedded Lorenz attractor” whose first 3 dimensions evolve according to an SDE $dx(t) = f(x(t)) dt + dw(t)$ where $f(x)$ is a Lorenz attractor with parameters $(\\rho, \\sigma, \\beta) = (28, 10, 8/3)$, and whose other dimensions $d > 3$ evolve according to a random walk with identity diffusion. Then, we simulate observations $y_n(t) \\sim \\mathcal{N}(c_n^\\top x(t) + d_n, (\\sqrt{0.05})^2)$ for $n = 1, \\dots, N$, where $N$ is the observation dimension. Here, the entries of $c_n$ and $d_n$ are sampled i.i.d from $\\mathcal{N}(0, 0.1^2)$. We fit SING using a prior SDE in the same class as the generative model and learn the drift parameters. We report latents RMSE, dynamics RMSE, and total runtime (including compilation time) on an A100 NVIDIA GPU, along with their 95% confidence intervals obtained from 5 different initial seeds for Monte Carlo sampling. All results are obtained by using a single sample per expectation evaluation.\n\n| latent dim ($D$) | latents RMSE | dynamics RMSE|runtime (s)|\n|------------|---------------------|----------------------|--------------------|\n| 3 | 0.1416 ± 0.0001| 0.1241 ± 0.0107     | 58.05 ± 0.16     |\n| 5| 0.1867 ± 0.0001| 0.1413 ± 0.0055     | 60.11 ± 0.36   |\n|10| 0.2721 ± 0.0001| 0.3102 ± 0.0059     | 89.20 ± 0.09    |\n| 20| 0.3971 ± 0.0001| 0.2043 ± 0.0048     | 136.28 ± 0.18     |\n| 50| 0.6761 ± 0.0001| 0.2193 ± 0.0059     | 561.44 ± 0.63     |\n\nGaussian quadrature (with $6$ quadrature nodes per dimension) achieves latents and dynamics RMSE $0.1415$ and $0.1220$ in $D = 3$ latent dimensions and $0.1866$ and $0.1337$ in $D = 5$ latent dimensions.\n\nWe observe that SING applied with Monte Carlo recovers the true latents and dynamics with comparable accuracy to Gaussian quadrature in $3$ and $5$ latent dimensions. Although latents RMSE continues to increase with $D$ beyond $5$ latent dimensions, this is expected for two reasons: (1) the error is computed over higher-dimensional vectors, and (2) accurate inference becomes more challenging as latent dimension increases. Additionally, we find that the runtime of joint inference and learning of SING remains tractable in up to 50 latent dimensions.\n\n> Would the baseline methods from related work would have the same scalability issues to higher-dimensional problems.\n\nSeveral related works on latent SDE inference already use sampling-based approaches to estimate intractable expectations, which allows them to scale to higher-dimensional settings (e.g., [1], [2]). We agree that Monte Carlo estimation is a natural and effective choice in this context, and in our additional results above we show that this can be seamlessly integrated into the SING variational inference algorithm. As a result, we expect SING (with Monte Carlo) to exhibit similar scaling behavior to these related methods. We will revise our manuscript to include this Monte Carlo extension and its implementation details.\n\n## Questions\n\n> Choice of linear posterior drift and whether it allows for inference in bifurcating systems. How does linear variational family reconcile with the prior SDE, which could admit bifurcations?\n\nThe reviewer is correct in noting that the prior SDE need not be linear and indeed could exhibit bifurcations, while the variational posterior is chosen to be a time-dependent linear SDE, which necessarily corresponds to a Gaussian process. This mismatch is reconciled by the fact that the prior SDE is shared across trials, whereas the variational posterior is not. \n\nSpecifically, under the generative model the latent trajectories are *i.i.d.* draws $x^{(i)} \\sim p(x)$ from the prior SDE. Because the observation likelihood also factorizes across trials, then the joint distribution of $N$ trials can be represented as $ \\prod_{i=1}^N p(x^{(i)}) p(y^{(i)} | x^{(i)})$. This implies that the true (intractable) posterior distribution also factorizes across trajectories. Hence, we construct our variational approximation as $q = \\prod_{i=1}^N q(x^{(i)})$, where each $q(x^{(i)})$ is represented by a linear, time-varying SDE with parameters $A^{(i)}(t)$ and $b^{(i)}(t)$.  And so, although each $q(x^{(i)})$ is linear, the collection $(q(x^{(i)}))_{i=1}^N$ can represent multiple modes induced by a nonlinear/bistable system.\n\n> Should one expect any connection between (1) the latent linear SDE learned by NGVI and (2) the learned vector field?\n\nThe latent linear time-varying SDE from NGVI is a trial-specific linear approximation of the posterior drift; the learned vector field (prior drift) is global and shared. The ELBO term $\\sum_{i=1}^N \\text{KL}(q(x^{(i)}) || p(x^{(i)} |  \\Theta))$ encourages the learned vector field to explain the dynamics of the inferred variational posterior for each trial. \n\n> However, are there some scenarios (such as bifurcations) where more flexible model classes may help? \n\nFor the reasons we detailed in our previous response, we find that SING performs comparably to methods with more flexible methods variational families, even on examples where the underlying latent dynamics are bifurcating (e.g., the Lorenz attractor).\n\nWe compare SING to SDE Matching [2] on the stochastic Lorenz attractor dataset [3] with 3D neural-SDE prior and Gaussian observation model. SDE Matching performs amortized variational inference by using a normalizing flow to model the posterior marginals. When trained on $1024$ trials, i.e. latent trajectories, with an equal number of gradient steps, SING recovers the underlying latents more accurately than SDE Matching (RMSE $0.440$ versus $0.549$). The neural network prior recovers the Lorenz attractor dynamics slightly worse than with SDE Matching (root mean square relative error $0.297$ versus $0.164$). In the data-limited regime ($< 30$ trials), SING outperforms SDE Matching on both metrics.\n\n> Can anything be said in the case where the diffusion coefficient has some dependency on the state $x$?\n\nWe agree with the reviewer that latent-SDE models with state-dependent diffusion coefficients are of practical relevance, and we detail how SING can be applied in this setting. To “whiten” the noise of the state-dependent prior process, one can apply the Lamperti transform $h$, $z(t) = h(x(t))$ (see e.g., [4] Section 3.6). We propose applying inference, with SING, and learning in the transformed latent space $z$. One can then apply Itô’s rule to recover the variational posterior in the original space $x$, which may not be a Gaussian process. \n\nAs to the reviewer’s example of the Cox-Ingersoll-Ross equation, $dx(t) = a(b - x(t)) dt + \\sigma \\sqrt{x(t)}dw(t)$, the Lamperti transform is $h(x) = \\frac{2}{\\sigma} \\sqrt{x}$, which takes nonnegative values. And so the linear, time-varying SDE variational approximation made by SING is not appropriate to approximate the posterior of $z(t) = h(x(t))$.\n\nWe validate our proposed approach on the one-dimensional geometric Brownian motion $dx(t) = \\mu x(t) dt + \\sigma x(t) dw(t)$ on $[0, 2]$, which has Lamperti transform $h(x) = \\frac{1}{\\sigma}\\log(x)$ taking values in $\\mathbb{R}$. We sample $1024$ latent trials with $\\mu = 1$ and $\\sigma = 1$, and observations at every $\\Delta t = 0.02$ that are equal to the latents perturbed by Gaussian noise with standard deviation $0.1$. We choose our prior to be a geometric Brownian motion with unknown $\\mu$, initialized at $0$. We fix the output parameters $C, d$ as well as $\\sigma$ and aim to (i) perform inference with SING and (ii) recover $\\mu$. We demonstrate that SING recovers the latents in the true latent space ($x$) with MSE $(0.084)^2$ as well as $\\hat{\\mu} \\approx 1.08$.\n\n> Is extension to posterior inference over $\\theta$ straightforward?\n\nYes, extension of SING to perform posterior inference over parameters $\\Theta$ is straightforward.\n\nOne could put a prior $p(\\Theta)$ on model parameters, and use a variational family that factorizes over parameters and latents, $q = q(\\Theta) q(x)$. Then, one could alternate between inferring $q(x)$ using SING and inferring an approximate posterior over parameters using gradient descent (or natural gradient descent, if possible) over $\\Theta$. \n\nWe also note that, if the goal is to get an approximate posterior over the drift, one could use SING-GP to solve this problem. SING-GP models the drift function using Gaussian processes, and it allows for uncertainty quantification over the dynamics directly in function space as opposed to in model parameter space. This approach could be a better alternative than Bayesian inference over drift parameters $\\Theta$ if the true drift of the system is unknown or if the parameter space $\\Theta$ is high-dimensional. \n\n**Thank you again for your positive comments and thoughtful response!** If you have any further questions, we would be happy to answer them.\n\n– The SING authors\n\nReferences:\n[1] Course and Nair. Amortized reparametrization: efficient and scalable variational inference for latent SDEs. NeurIPS, 2023. [2] Bartosh et al. SDE Matching: Scalable and simulation-free training of latent stochastic differential equations. ICML, 2025. [3] Li et al. Scalable gradients for stochastic differential equations. AISTATS, 2020. [4] Pavliotis. Stochastic processes and applications. Texts in Applied Mathematics, 2014."}, {"review_id": "Wn9pJuM8Hu", "reviewer": "Reviewer_Q7Cf", "summary": "The authors consider a method for estimating latent stochastic differential equations (SDEs) through a natural gradient estimator. The natural gradient provides gradient updates that adhere to the geometry of the KL divergence. The model is assumed to have observations of a known link function applied to an affine transformation of a latent SDE, and the objective is to recover the affine transformation as well as the posterior latent distribution. The authors consider a family of linear SDEs for the variational distribution, which provides Gaussian time marginals and is amenable for use with natural gradient variational inference. The authors test the proposed method on a number of different experiments and illustrate the performance of the natural gradient estimator versus baselines.", "questions": "How limiting is the variational distribution being a linear SDE? For example, since we approximate objects of the form $g(Cx+d)$, this would imply a limitation on the volatility, for example, of the process of $y$.\n\nIs computing the Jacobian in the ELBO expensive? Can the authors describe some more details on this. \n\nIs there a way to estimate the parameters of $g$ using this method? \n\nFor other SDEs (for example, GBM) where the volatility is a function of state, are there other ways to extend the method to these cases? The GBM case has log normal marginal distributions, so I assume something similar could work by reformulating in terms of the parameters of a log normal, but is this generally extensible? \n\nMany of the experiments suggest that the true drift can be recovered but I was wondering if a more concrete theoretical statement could be made regarding this? Along the lines of the previous questions, how does the constrained variational distribution affect this?", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths:**\n\nThe authors provide a bound on the ELBO under a discrete setting to the continuous time setting that has the same order of convergence as the Euler Maruyama method. \n\nThe method empirically displays good performance on a variety of tasks. A number of comparisons are particularly impressive, for example the Adam vs NGVI plots are very informative.\n\nThe derivation of this formulation is neat and a slightly different take than many other papers along the lines of latent SDEs.\n\n**Weaknesses:** \n\nThe authors mention a variety of related work but do not consider comparisons to them. It would be interesting to see how these methods compare in both computation speed and memory usage. \n\nIt would be useful to get more context as to where using the natural gradient is more appropriate versus using standard gradient descent with an ELBO with a more flexible variational distribution. \n\nThe variational distribution being a linear SDE to be a major limitation especially since the volatility needs to be constrained quite heavily. On that note, the main techniques involve making the variational distribution Guassian and applying NGVI is not entirely original.\n\nThis field has a lot of related work and has been well studied from a variety of different angles, so it would be good to get more context on the contributions made in the paper compared with prior work. While this work focuses on the natural gradient variational inference under a linear SDE variational distribution for discovering SDE parameters, others have considered adjacent strategies. Just as an example of some related work missing: \n\nLatent SDEs on Homogeneous Spaces, NeurIPS 2023\n\nNeural Structure Learning with Stochastic Differential Equations, ICLR 2024\n\nIdentifying Latent Stochastic Differential Equations, IEEE Transactions on Signal Processing 2021\n\nLearning interpretable dynamics of stochastic complex systems from experimental data, Nature Communications 2024\n\nComputing high-dimensional invariant distributions from noisy data, Journal of Computational Physics 2023", "was_revised": true, "final_justification": "After reading the response, I believe the method provides useful insights into the some topics regarding estimating SDEs and would be an interesting addition to the conference. I am still not certain about the novelty of the method since the main technical results are largely combinations of existing work. Additionally, requiring an SDE to satisfy the Lamperti transform can be quite difficult, I think for a general SDE estimation method this is something that should be studied in greater detail.", "rebuttal": "We thank the reviewer for their positive and thoughtful response! We especially appreciate the reviewer’s comments on the derivation of our ELBO bound and the strong empirical performance of SING across tasks.\n\n## Weaknesses\n> Comparisons to related work, in particular computation speed and memory usage\n\nWe have conducted additional experiments comparing SING to SDE Matching [1], a recent work that uses amortized normalizing flows to model posterior marginals, which are not necessarily Gaussian. This is similar to [2], which still assumes Gaussian posterior marginals but also applies an amortized encoder.\n\nWe compare the methods on a stochastic Lorenz benchmark [3], where 3D latents evolve on $t \\in [0, 1]$ with $(\\rho, \\sigma, \\beta) = (28, 10, 8/3)$ and diffusion $0.15$. Observations are generated every $\\Delta t = 0.025$ as $y(t) = \\mathcal{N}(Cx(t) + d, 0.01^2)$. For SDE Matching, we used the authors’ codebase. We ran SING for $10^3$ vEM iterations (10 E-steps, 10 M-steps each) and SDE Matching for $10^4$; the number of gradient updates is the same. Both use a neural-SDE prior with softplus activation. We evaluate methods using the RMSE between the inferred and ground truth latents (‘latents RMSE’) and between the learned prior drift and ground truth dynamics (‘dynamics RMSE’).\n\n**Latents RMSE**\n| n_trials | SING  | SDE Matching |\n|----------|-------|--------------|\n| 10 | 0.443 | 1.173|\n| 30 | 0.443 | 0.807|\n| 50 | 0.442 | 0.711|\n| 100 | 0.440 | 0.719|\n| 300 | 0.440 | 0.637|\n| 500 | 0.440 | 0.567|\n| 1024 | 0.440 | 0.549|\n\n**Dynamics RMSE**\n| n_trials | SING  | SDE Matching |\n|----------|-------|--------------|\n| 10 | 0.272 | 0.551|\n| 30 | 0.265 | 0.329|\n| 50 | 0.300 | 0.232|\n| 100 | 0.265 | 0.239|\n| 300 | 0.330 | 0.179|\n| 500 | 0.270 | 0.161|\n| 1024 | 0.297 | 0.164|\n\nSING outperforms SDE Matching in data-limited settings ($< 30$ trials) and recovers latents more accurately in up to $1024$ trials. We hypothesize that this is because SING explicitly models the entire latent trajectory during inference, whereas SDE Matching performs updates using the inferred posterior at a single time point. However, SING does not effectively refine its prior estimate of Lorenz dynamics with more trials; we will investigate improvements to learning with SING following the rebuttal period.\n\nSDE Matching is preferable in runtime and memory, as it updates the amortized posterior at individual time points. SING’s complexity depends on $T$, the length of the time grid at which inference is performed. Still, SING parallelizes its updates across $T$, yielding $\\mathcal{O}(D^3 \\log T)$ complexity. \n\n> More context on natural vs. standard gradient with a more flexible variational distribution\n\nSING’s variational posterior factorizes across trials $i$ as $q = \\prod_{i=1}^N q(x^{(i)})$, where each $q(x^{(i)})$ is represented by a linear time-varying SDE with parameters $A^{(i)}(t), b^{(i)}(t)$. Consequently, if the latent trajectories can be well-explained by per-trial locally linear dynamics, then using a linear time-varying SDE as the variational posterior is reasonable. In this case, SING exploits the exponential family structure of the problem to achieve faster convergence than standard gradient descent or even Adam, as we showed in our experiments. Moreover, in data-limited regimes, using a linear SDE could be more statistically efficient relative to more flexible variational posteriors. That said, for latent SDEs exhibiting sharp nonlinearities, more flexible variational distributions may yield better estimates, provided enough data is available.\n\n> Linear SDE is a major limitation, especially since the volatility needs to be constrained\n\nAs mentioned above, the variational posterior factorizes across trials $i$. Thus, as long as each trajectory $x^{(i)}$ is well-approximated by time-varying linear dynamics, we would expect our variational approximation to be accurate. This claim is supported by our experiments showing that SING performs accurate inference in the Duffing oscillator (Fig 3) and stochastic Lorenz attractor.\n\nAs to volatility, we outline our approach to allowing for state-dependent diffusion coefficients below.\n\n> Main techniques are not entirely original\n\nWe agree with the reviewer that natural gradient methods have been well-studied. The goal of our work is to develop a robust and scalable natural gradient method for latent SDEs, along with theoretical guarantees on its approximation error to the continuous-time setting. Our contribution lies in deriving closed-form natural gradient updates specifically for latent SDEs, enabling parallelization of these updates for scalability to long sequences, providing theoretical bounds on the ELBO approximation error, and extending this approach to inference in models with nonparametric (Gaussian process) drift.\n\n> Additional related work\n\nWe thank the reviewer for bringing these to our attention and will include them in our revised paper.\n\n## Questions\n\n> Limitation on the volatility of the process of $y$\n\nWhile the variational posterior of $x(t)$ is assumed to be a linear SDE, its learned time-dependent linear coefficients $A(t)$ allow the marginal distribution of $y(t)$ to have variance that varies in time. This is because the marginal variance of $x(t)$ evolves according to the ODE $dS(t) = A(t)S(t) + S(t)A(t)^\\top + \\Sigma$.\n\n> Is computing the Jacobian in the ELBO expensive?\n\nThe ELBO does not require computing a Jacobian. The author may be referring to the naive NGVI update (eq. 4), which requires computing the inverse Fisher information matrix, i.e. the Jacobian of the mapping from natural to mean parameters. Naively, this would have complexity $O(D^3 T^3)$. However, we do not actually compute the inverse Fisher explicitly; due to the exponential family structure of our variational posterior, the NGVI update can be equivalently written as eq. (6). We show that SING can compute these updates in $O(D^3 \\log T)$ time, enabling scalability to long time series.\n\n> Is there a way to estimate the parameters of $g$?\n\n$g$ is a pre-determined inverse link function, but the parameters $C$ and $d$ of the affine mapping can be learned.\n\n> For other SDEs (for example, GBM) where the volatility is a function of state, are there other ways to extend the method to these cases?\n\nWe thank the reviewer for their insight regarding how SING can be applied to latent-SDE models with state-dependent diffusion. The example of the GBM can indeed be generalized to a larger class of state-dependent diffusion coefficients via the Lamperti transform (e.g., [4] Section 3.6). For a 1D process $dx(t) = f(x(t)) dt + \\sigma(x(t)) dw(t)$, the Lamperti transform is defined as\n$$\nh(x) = \\int_{x_0}^x 1/\\sigma(s) ds \n$$\nfor $x_0$ an arbitrary fixed point in the state space. Applying Itô’s rule to the map $z(t) = h(x(t))$ and process $x(t)$ yields a diffusion process with identity noise and modified drift. The observation model also transforms accordingly through $h^{-1}$. \n\nOnce we have transformed the prior and likelihood models, we can leverage SING to perform inference in the transformed space $z$. One can then apply Itô’s rule to $h^{-1}(z(t))$ to recover the approximate posterior in the original space $x$. Because the continuous-time ELBO $\\mathcal{L}_{\\text{cont}}$ remains invariant under reparameterization (e.g., [5] Lemma E.2.1), the optimization problem in the transformed space is equivalent to that in the original space.\n\nWe note that a Gaussian process variational approximation in the transformed space $z$ is not appropriate when the Lamperti transform takes values in only a subset of $\\mathbb{R}$. For example, for the Cox-Ingerson-Ross equation $dx(t) = a(b - x(t)) dt + \\sigma \\sqrt{x(t)}dw(t)$, the Lamperti transform $h(x) = \\frac{2}{\\sigma} \\sqrt{x}$ takes nonnegative values.\n\nWe conducted additional experiments validating our proposed approach on latents that evolve according to the 1D GBM $dx(t) = \\mu x(t) dt + \\sigma x(t) dw(t)$ on $[0, 2]$ with $\\mu = 1$ and $\\sigma = 1$. We have observations at every $\\Delta t = 0.02$ that are equal to the latents perturbed by Gaussian noise with SD $0.1$, and we use $1024$ trials. We choose our prior to be a GBM with unknown $\\mu$, initialized at $0$. The Lamperti transform is $h(x) = \\frac{1}{\\sigma}\\log(x)$. We fix the output parameters $C, d$ as well as $\\sigma$ and aim to perform inference and learn $\\mu$. We find that SING recovers the latents in the true latent space ($x$) with MSE $(0.084)^2$ as well as $\\hat{\\mu} \\approx 1.08$.\n\n> More concrete theoretical statement on true drift recovery? How does the constrained variational distribution affect this?\n\nWe agree that a more concrete theoretical analysis on drift recovery would be valuable, particularly with respect to the choice of variational family. While our work focuses primarily on empirical drift recovery rather than formal guarantees, we note that the choice of a linear time-varying SDE can still capture nonlinearities in the latent processes (as above) and encourage learning of a nonlinear drift. \n\n**Thank you again for your positive comments and thoughtful response!** If you have any further questions, we would be happy to answer them.\n\n– The SING authors\n\nReferences:\n[1] Bartosh et al. SDE Matching: Scalable and simulation-free training of latent stochastic differential equations. ICML, 2025.\n[2] Course and Nair. Amortized reparametrization: efficient and scalable variational inference for latent SDEs. NeurIPS, 2023.\n[3] Li et al. Scalable gradients for stochastic differential equations. AISTATS, 2020.\n[4] Pavliotis. Stochastic processes and applications. Texts in Applied Mathematics, 2014.\n[5] Dupuis and Ellis. A weak convergence approach to the theory of large deviations. John Wiley and Sons, 1997."}, {"review_id": "IMXS82YhNU", "reviewer": "Reviewer_c8GD", "summary": "This paper introduces SDE Inference via Natural Gradients (SING) as a method that leverages natural gradient VI to improve inference in latent SDE models. The paper also describes a method for parallelisation, further improving efficiency and extends the approach to GP-SDEs. The paper provides extensive theoretical support and proofs as well as solid experimental support via validation on simulated and real-world data.", "questions": "In related work, the authors should also describe alternative approaches to latent SDE estimation based on flow matching. See e.g. https://arxiv.org/abs/2502.02472. As these methods significantly improve on VI it would be important to contrast against these. Ideally the authors show competitive performance with these alternative approaches that have been introduced to address the limitations of VI.", "limitations": "The assumption that the diffusion term is known can be acknowledged as a limitation if this cannot be addressed in the current setup.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "This paper provides a solid contribution to the field by introducing SING as a way to enable fast and reliable inference in latent SDE models. The claims are supported by solid proofs. The parallel approach yields convincing improvements. The experimental validation clearly demonstrate its benefits compared to methods like VDP. The paper is very well written in general.\n\nThere are a couple of points that warrant attention. \n\nThe diffusion term Sigma is assumed given. However, it can also be considered a model parameter. Can the theory be extended to incorporate fitting of Sigma? If yes, I suggest to incorporate this. If no, please explain why and explain why this is not an issue when fitting models on real-world data where the process noise is unknown.\n\nThe authors state that the method is closest to that of Verma et al. However, comparisons are with VDP. While the authors clearly describe the differences in the appendix I would have expected an explicit comparison in terms of empirical results. Please add these comparisons to the plots or show this using a targeted comparison in the appendix or make more explicit why such a comparison is irrelevant.", "was_revised": true, "final_justification": "I would like to thank the authors for their response and additional experiments. This paper is an important contribution to the community.", "rebuttal": "We thank the reviewer for taking the time to read our submission and for their thoughtful response! We are especially glad that the reviewer appreciated the theoretical and empirical contributions of our method, and found our paper to be well-written.\n\n## Weaknesses\n\n> Assumptions on the diffusion term $\\Sigma$\n\nWe thank the reviewer for their comment and point out that for the purpose of inference, SING can handle arbitrary fixed, positive-definite diffusion coefficient $\\Sigma^{1/2}$. Moreover, when learning the output parameters $C, d$ together with the prior drift $f$, the diffusion coefficient should be fixed to $\\Sigma^{1/2} = I$ to ensure identifiability of the model.\n\nIndeed, suppose our prior over latents is $p(x): dx(t) = f(x(t)) dt + \\Sigma^{1/2} dw(t)$. Then by applying the transformation $z(t) = \\Sigma^{-1/2}x(t)$, the generative model becomes $p(z): dz(t) = \\Sigma^{-1/2} f(\\Sigma^{1/2} z(t)) dt + dw(t)$, $\\mathbb{E}[y(t)] = g(C\\Sigma^{1/2} z(t) + d)$. And the posterior process becomes $q(z(t)): dz(t) = \\Sigma^{-1/2}( A(t) \\Sigma^{1/2} z(t) + b(t) )dt +dw(t)$, which is again a linear, time-varying SDE. Moreover, it is a standard result (see e.g., [1] Lemma E.2.1) that the continuous-time ELBO $\\mathcal{L}_{\\text{cont}}$ remains invariant under invertible transformations. As a consequence, for any prior process with non-white noise $\\Sigma^{1/2}$, one can always perform an appropriate transformation to whiten the noise, perform approximate inference with SING, and then apply the inverse transformation to recover the inferred latents in the original space.\n\nIf model class remains the same under the mapping $f \\to \\Sigma^{-1/2} f(\\Sigma^{1/2})$ (e.g., neural networks, linear maps, linear combination of polynomial basis functions) and $C, d$ are learned together with the drift $f$, then the prior noise $\\Sigma$ is non-identifiable in the model. For this reason, we fix $\\Sigma = I$ when performing learning.\n\nLastly, we point out that for certain state-dependent diffusion coefficients, it remains possible to perform both learning and inference with SING. This is done using a similar approach to that outlined above: whitening the noise of the prior process and variational posterior, applying inference and learning with SING, and then performing the inverse transformation. Precisely, we transform the prior SDE using the Lamperti transform (see e.g., [2] Section 3.6). \n\n> Comparison to Verma et al., 2024\n\nWe appreciate the reviewer’s suggestion and agree that an empirical comparison with Verma et al., 2024 would be ideal and have explored this direction. However, we encountered persistent issues running their publicly released code. We carefully followed the installation instructions provided in their repository and additionally attempted several manual interventions, such as modifying version constraints and testing across different environments. Despite these efforts, we were unable to get the code running due to unresolved dependency conflicts. We have alerted the authors to inform them of this issue.\n\nIn light of this, we provide a more explicit comparison between SING and Verma et al., 2024 which we will also include in the appendix. Verma et al., 2024 propose an inference algorithm that repeatedly factorizes the prior on latents as $p(x) = p_L(x) v(x)$ where $p_L$ is a linear SDE and $v(x) = p(x) / p_L(x)$. During inference, they alternate between this factorization, performing natural gradient descent to update the variational posterior and applying a statistical linearization step to determine a new $p_L$. In contrast, SING avoids these repeated linearization and factorization steps by leveraging the fact that NGVI is parametrization invariant; that is, the particular factorization of the generative model does not affect the resulting updates. As a result, SING achieves equivalent inference while reducing computational cost. \n\nMoreover, SING introduces a novel theoretical contribution not present in Verma et al., 2024: a bound on the approximation error between the discretized ELBO optimized in practice and the true continuous-time ELBO. This result provides theoretical justification for applying SING to continuous-time latent SDE models, bridging an important gap between practice and theory that was not addressed in prior work.\n\nWe also acknowledge that Verma et al., 2024 propose a modified M-step that couples inference and learning by updating model parameters and variational parameters simultaneously. While this approach may reduce the number of M-step iterations required for convergence, it increases the computational cost per iteration. More importantly, the practical benefits of this approach are difficult to assess without empirical comparison, which is not feasible at this time due to the technical limitations described above.\n\nFinally, we would like to emphasize that SING introduces several additional practical and computational improvements which we believe will make our method more readily usable and scalable in real-world applications. These include simpler but equivalent natural gradient updates during inference, parallelizable updates which scale logarithmically with respect to sequence length, and tractable computations for expectations in the ELBO. We have developed and plan to release a modular and efficient codebase that implements SING and its extensions, including support for Gaussian process drift models. \n\n## Questions\n\n> Comparison to SDE Matching\n\nWe thank the reviewer for their suggestion to compare SING to SDE Matching [3]. SDE Matching performs amortized variational inference by using normalizing flows to model the posterior marginals. These marginals are not necessarily Gaussian. This is similar to another related work [4], which still assumes Gaussian posterior marginals but also applies an amortized posterior encoder.\n\nWe compare the methods on the stochastic Lorenz attractor benchmark dataset, in which 3D latents evolve according to Lorenz attractor dynamics on $t \\in [0, 1]$ with parameters $(\\rho, \\sigma, \\beta) = (28, 10, 8/3)$ and diffusion coefficient $0.15$. Observations are made every $\\Delta t = 0.025$ are generated as $y(t) = \\mathcal{N}(Cx(t) + d, 0.01^2)$. The full experimental setup is detailed in [5]. For SDE Matching, we used the authors’ codebase. We compare SING with $10^3$ vEM iterations, each consisting of 10 SING E-steps and 10 M-steps, to SDE Matching with $10^4$ iterations; the number of gradient updates is the same for both methods. SING updates are performed with a single Monte Carlo sample. We model the prior for both methods with a 3D neural-SDE, whose drift is a neural network with softplus activation.\n\nWe compare the methods using two metrics: the RMSE between the inferred and ground truth latent (‘latents RMSE’) and the root mean squared relative error between the learned prior drift and ground truth Lorenz attractor dynamics (‘dynamics RMSE’).\n\n**Latents RMSE**\n| n_trials | SING  | SDE Matching |\n|----------|-------|--------------|\n| 10       | 0.443 | 1.173        |\n| 30       | 0.443 | 0.807        |\n| 50       | 0.442 | 0.711        |\n| 100      | 0.440 | 0.719       |\n| 300      | 0.440 | 0.637       |\n| 500      | 0.440 | 0.567       |\n| 1024     | 0.440 | 0.549       |\n\n**Dynamics RMSE**\n| n_trials | SING  | SDE Matching |\n|----------|-------|--------------|\n| 10       | 0.272 | 0.551        |\n| 30       | 0.265 | 0.329        |\n| 50       | 0.300 | 0.232        |\n| 100      | 0.265 | 0.239       |\n| 300      | 0.330 | 0.179       |\n| 500      | 0.270 | 0.161       |\n| 1024     | 0.297 | 0.164      |\n\nWe observe that SING outperforms SDE Matching on both metrics in data-limited settings ($< 30$ trials). Moreover, even up to $1024$ trials, SING recovers the latents more accurately than SDE Matching. We hypothesize that this is because SING explicitly models the entire latent trajectory during inference, whereas SDE Matching performs updates using the inferred posterior at a single time point. However, SING is not able to effectively refine its prior estimate of Lorenz dynamics as the number of trials increases; we will investigate improvements to learning with SING following the rebuttal period.\n\n**Thank you again for your positive comments and thoughtful response!** If you have any further questions, we would be happy to answer them.\n\n– The SING authors\n\n[1] Dupuis and Ellis. A weak convergence approach to the theory of large deviations. John Wiley and Sons, 1997.\n[2] Pavliotis. Stochastic processes and applications. Texts in Applied Mathematics, 2014.\n[3] Bartosh et al. SDE Matching: Scalable and simulation-free training of latent stochastic differential equations. ICML, 2025.\n[4] Course and Nair. Amortized reparametrization: efficient and scalable variational inference for latent SDEs. NeurIPS, 2023.\n[5] Li et al. Scalable gradients for stochastic differential equations. AISTATS, 2020."}]

Author rebuttal:

Meta-review: Authors propose a novel approach for variational inference in latent SDEs that effectively addresses the common issues of slow convergence and numerical instability. In recent years, there have been many approaches to this problem. While natural-gradient variational inference for SDEs is not a new concept, the authors successfully differentiate their work through this parallelization acceleration and a detailed convergence analysis (in the appendix). The reviewers were positive about the paper and indicated that it is a good fit for the conference. The authors also honestly acknowledged one limitation: that SING does not effectively refine its prior estimate of Lorenz dynamics with more trials. This work represents a signficant and orignal contribution relevant for both the machine learning and statistical neuroscience communities.
Decision: Accept (poster)