PaperlayerView exact revision on arXiv

Source-linked AI summary

Amortized Reparametrization: Efficient and Scalable Variational Inference for Latent SDEs

Kevin Course, Prasanth B. Nair

arXiv:2312.10550v1cs.LGstat.ML

TL;DR

Training latent SDEs with adjoint methods requires solving initial value problems, creating computational burdens for continuous-time models. This paper avoids solving differential equations when approximating gradients, achieving similar performance with more than an order of magnitude fewer model evaluations during training.

  • Problem

    Adjoint methods for training continuous-time models require solving pairs of initial value problems and are poorly suited to parallelization.

    Source: §1 · p. 1

  • Method

    The method couples amortization with a reparametrization for expectations under Markov Gaussian processes to approximate latent-SDE training gradients without solving differential equations.

    Source: §1 · p. 2, §5 · p. 10

  • Results

    The approach uses more than an order of magnitude fewer model evaluations during training than adjoint methods while achieving similar RMSE on the validation set.

    Source: §4.1 · p. 8, §5 · p. 10

  • Takeaways & Limitations

    Gradient approximations have time and memory costs independent of data amount, time-series length, and model stiffness.

    Source: §5 · p. 10

    The method should be used to approximate the latent-state posterior only when it is approximately a Markov Gaussian process.

    Source: §3 · p. 6

Abstract

from arXiv · show

We consider the problem of inferring latent stochastic differential equations (SDEs) with a time and memory cost that scales independently with the amount of data, the total length of the time series, and the stiffness of the approximate differential equations. This is in stark contrast to typical methods for inferring latent differential equations which, despite their constant memory cost, have a time complexity that is heavily dependent on the stiffness of the approximate differential equation. We achieve this computational advancement by removing the need to solve differential equations when approximating gradients using a novel amortization strategy coupled with a recently derived reparametrization of expectations under linear SDEs. We show that, in practice, this allows us to achieve similar performance to methods based on adjoint sensitivities with more than an order of magnitude fewer evaluations of the model in training.

1 Introduction

The paper targets the computational burden of training continuous-time models and introduces a method for inferring latent SDEs without solving differential equations to approximate gradients. It reports comparable performance to adjoint-based methods with more than an order of magnitude fewer model evaluations in training.

1 Introduction

  • Continuous-time models can handle unevenly spaced observations and support hierarchical or adaptive prediction, while the proposed approach targets high-dimensional time-series tasks.The authors say its computational efficiency can support hyperparameter tuning and make state-of-the-art methods feasible on lower-performance hardware.Source: §1 · pp. 1–2
  • Adjoint-based training requires solving two initial value problems, while stiffness can make iteration costs explode and sequential solvers are poorly suited to parallel hardware.Adjoint methods also have stability issues for time-averaged gradients over long intervals in chaotic systems.Source: §1 · pp. 1–2
  • Regularizing dynamics and replacing ODE solvers with integration methods can speed training, but continuous-time models remain more expensive than discrete-time models..Source: §1 · p. 1
  • The proposed method approximates ELBO gradients without solving differential equations, with time and memory costs independent of data amount, time-series length, and stiffness.Its evaluation cost is a user-chosen constant, and model evaluations can run in parallel.Source: §1 · p. 2, Table 1 · p. 2
  • More than an order of magnitude fewer model evaluations achieved comparable performance to adjoint-based methods in training.The paper also reports avoiding adjoint-method instabilities for long time series generated by chaotic systems.Source: §1 · p. 2

2 Method

The method infers nonlinear latent SDEs by approximating their intractable latent-state posterior and optimizing a variational objective without solving a forward differential equation. It amortizes that approximation over short observation windows and uses reparameterization to form unbiased gradient estimates.

2.1 Problem description

  • The model treats observations as generated from a latent nonlinear SDE and seeks both data-aligned generative parameters and a parametric posterior over latent states.The posterior at a point in time is intractable; the variational posterior is conditioned on a small observation window.Source: §2.1 · pp. 2–3

2.2 Evidence lower bound

  • A Markov Gaussian Process defined by a linear SDE approximates the latent-state posterior; its Gaussian marginals support an ELBO that avoids a forward solver.The approximation is an intermediary for training; beyond the data window, the nonlinear SDE is used for forecasting.Source: §2.2 · pp. 3–4

2.3 Amortization strategy

  • Amortization approximates the posterior over short, non-overlapping observation partitions, allowing multiple trajectories without increasing overall computational cost.The partition size should provide enough observations for the linear-SDE posterior assumption to be reasonable.Source: §2.3 · pp. 4–5

2.4 Reparametrization trick

  • Reparameterization combined with amortization yields an unbiased ELBO gradient approximation with O(R) time and memory cost.The stated scaling is independent of data amount, time-series length, and stiffness of the approximate differential equations.Source: §2.4 · pp. 5–6

3 Limitations & Related Work

The approximate posterior is assumed to be Gaussian-process-like over the observation window, while prediction samples from a nonlinear latent SDE. The approach also has a time-dependent-diffusion constraint and differs from latent ODE and weak-form methods.

3 Limitations & Related Work

  • The approximate posterior is assumed to be approximately MGP only over the observation window; prediction still samples from a nonlinear SDE and can have a complex distribution.The paper says the GP approximation is common in nonlinear state estimation.Source: §3 · pp. 6–7
  • The approach supports time-dependent diffusion, not state-dependent diffusion, though it could provide an initial guess for drift parameters when state-dependent diffusion is needed.Source: §3 · p. 7
  • Latent ODEs [1] [17] place uncertainty in initial conditions and rely on adjoints, whereas SDEs capture uncertainty accumulating over time.Source: §3 · p. 7
  • The approach replaces differential-equation solves with unbiased one-dimensional integral approximations, reducing training evaluations and avoiding adjoint instability for long time-averaged quantities in chaotic systems.Source: §3 · p. 7
  • Weak-form methods also turn differential-equation solving into integration, but learn only in data coordinates and use biased residual estimates that fail with widely spaced observations.They also require a carefully designed test space, unlike this approach.Source: §3 · p. 7

4 Numerical Studies

The numerical studies evaluate ARCTA on training cost, chaotic-system gradient stability, motion-capture forecasting, and video generation.

4.1 Orders of magnitude fewer function evaluations in training

  • More than one order of magnitude fewer model evaluations achieved similar validation RMSE to NODEs, even when tighter solver tolerances worsened NODE accuracy.The Lotka–Volterra comparison uses three NODE tolerances; Figure 2 also shows predictive test-set samples.Source: Figure 2 · p. 8, §4.1 · p. 8

4.2 Numerical instabilities of adjoints

  • As time-series duration increased, ARCTA remained numerically stable while adjoint methods produced prohibitively large gradients.The parameter-tuning experiment used data sampled at 200 Hz over intervals up to 100 seconds.Source: §4.2 · p. 9, Figure 3 · p. 9

4.3 Motion capture benchmark

  • ARCTA achieved competitive motion-capture performance, with test scores averaged across 10 models because only four independent test sequences were available.The table compares its performance with other state-of-the-art methods; the figure also shows selected-output predictive posteriors.Source: §4.3 · pp. 9–10, Figure 4 · p. 9

4.4 Neural SDEs from video

  • The method learned a generative model for a nonlinear pendulum video and produced a 30-second forecast from five initial frames.The experiment used 32 × 32 black-and-white frames collected at 15 Hz; Figure 5 overlays predictive-posterior samples on the data.Source: §4.4 · p. 10, Figure 5 · p. 10

5 Conclusions

The paper concludes that its approach reduces training cost while avoiding adjoint gradient instabilities in long chaotic time series, and demonstrates latent-SDE video generation.

5 Conclusions

  • Unbiased ELBO-gradient approximations achieve time and memory costs independent of data amount, time-series length, and latent-equation stiffness through amortization and Monte Carlo expectations.The method trades adaptive-solver numerical precision for Monte Carlo approximations and uses a change of variables for expectations under Markov Gaussian processes.Source: §5 · p. 10
  • Across test problems, the approach used more than one order of magnitude fewer function evaluations than adjoint methods while avoiding their numerical instabilities on long chaotic time series.The authors also demonstrated generative modeling of a simple video.Source: §5 · p. 10
  • Future work includes variance reduction over time expectations, generative modeling on more realistic problems, and applications to implicit densities.Source: §5 · p. 10

Appendices

The appendix re-derives an identity that rewrites expectations under linear SDEs using Gaussian marginal statistics, then notes a diagonal-case simplification.

A Expectations under linear SDEs

  • Theorem 1 rewrites a bounded-functional expectation under a linear SDE as an expectation under a Gaussian with the solution’s marginal mean and covariance.The functional’s arguments are replaced by expressions involving B(t), the mean derivative, and the mean.Source: A Expectations under linear SDEs · p. 13
  • The proof uses the linear SDE’s Gaussian marginals and matrix Lyapunov equations to express A(t) in terms of S(t), then rearranges to obtain b(t).Source: A Expectations under linear SDEs · p. 13
  • When S(t) and L(t)ΣL(t)T are diagonal, Equation 17 simplifies.Source: A Expectations under linear SDEs · p. 13

B Evidence lower bound derivation

The reparametrized ELBO replaces differential-equation expectations with normal expectations, and amortization addresses its poor scaling with long trajectories. The framework also supports variational inference over generative parameters with little added computational cost.

B Evidence lower bound derivation

  • The reparametrized ELBO uses normal-distribution expectations, eliminating the differential-equation solver requirement, but initially stores the full latent-state estimate.Course and Nair applied the identity in (14) to derive the reparametrized ELBO.Source: B Evidence lower bound derivation · p. 14
  • Amortization yields an unbiased gradient approximation whose cost scales independently of data amount, time-series length, and differential-equation stiffness.It addresses the storage burden of the reparametrized ELBO, which otherwise requires storing the entire latent-state estimate.Source: B Evidence lower bound derivation · p. 14

C Evidence lower bound with priors on generative parameters

  • Stochastic variational inference over selected generative parameters adds only a marginal computational cost and leaves asymptotic complexity unchanged.The numerical studies place a log-normal prior on the diagonal diffusion term; this choice helped stabilize training in the examples considered.Source: C Evidence lower bound with priors on… · p. 14
  • The modified ELBO incorporates a KL-divergence between the approximate posterior and prior over generative parameters.The posterior over these parameters is reparameterized, and the resulting gradient estimate inherits the original approximation's properties.Source: C Evidence lower bound with priors on… · pp. 14–15

D Detailed description of recognition network

The recognition network maps observation windows to a reduced latent-state representation and uses a deep kernel to interpolate estimates over time. The appendices also specify a diffusion-parameter posterior, computing resources, and experiment-wide modeling choices.

D Detailed description of recognition network

  • The encoder reduces each observation batch to a latent state with h_i ∈ R^2d and can use varying amounts of data.The encoder is described as interpretable and cheap to compute when M is relatively small, since its evaluations can be parallelized; other encoders can also be used with Lemma 1.Source: D Detailed description of recognition… · pp. 15–16
  • Using a few neighboring frames helps resolve ambiguous latent states, such as a pendulum frame that does not reveal its swing direction.The approach approximates latent states using MGPs, so it cannot represent the multimodal posterior that a single ambiguous frame would require; the authors found K ∈ [1, 5] worked well in the problems considered.Source: D Detailed description of recognition… · p. 15
  • The deep-kernel encoder interpolates between observation encodings to estimate latent-state means and covariance over time.The authors report that without the deep kernel, validation accuracy was difficult to achieve on the datasets considered.Source: D Detailed description of recognition… · pp. 15–16

E Approximate posterior on diffusion term

  • The diffusion term is parameterized through C_θ(t)^−1 = diag(θ), with a log-normal approximate posterior described using variational parameters μ_i and σ_i.The loss requires the product C_θ(t)^−1, so the method parameterizes that product rather than L_θ alone.Source: E Approximate posterior on diffusion… · p. 16

F Computing resources

  • Experiments used an Ubuntu server with 24 CPU cores, 128GB RAM, and an NVIDIA GeForce RTX 4090 GPU.Most code was written in PyTorch; benchmarking used torchdiffeq, torchsde, and pytorch_lightning.Source: F Computing resources · p. 16

G Details on numerical studies

  • Across numerical studies, the model combines an encoder, deep kernel, decoder, and SDE drift and dispersion, with a log-normal dispersion prior and a linear KL schedule.The KL terms were increased from 0 to 1; the approximate posterior was assumed constant in time.Source: G Details on numerical studies · pp. 16–17

G.1 Orders of magnitude fewer function evaluations

This appendix section details the Lotka–Volterra predictive-model experiment, including its data generation, model configuration, and optimization settings.

G.1 Orders of magnitude fewer function evaluations

  • The Lotka–Volterra dataset used noisy observations at 50Hz, with the first 50 seconds for training and the remaining 15 seconds for validation.The trajectory ran from time 0 to 65, with Gaussian observation noise of standard deviation 0.01.Source: §G.1 · p. 17
  • The NODE and proposed method used the same gradient-field and drift-function model, while the decoder was set to the identity.A Gaussian likelihood with constant standard deviation 0.01 was used.Source: §G.1 · p. 17
  • The experiment used R = 1, S = 10, and M = 256 in the nested Monte Carlo approximation, with a learning rate initialized at 10^−3.The KL schedule increased from 0 to 1 over 1000 iterations.Source: §G.1 · p. 17

G.2 Adjoint instabilities experiment

The Lorenz parameter-tuning experiment compares gradient behavior across time-series lengths. For intervals [0, 50] and [0, 100], the authors could not report average adjoint gradients because initial gradients were too large.

G.2 Adjoint instabilities experiment

  • The experiment evaluated parameter inference on five noisy datasets at each of four time intervals: [0, 1], [0, 10], [0, 50], and [0, 100].Observations were generated at 200Hz, with independent Gaussian noise of variance 1.Source: §G.2 · p. 17
  • The method’s reported gradients are average ELBO-gradient norms for σ, β, and ρ after 2000 iterations, while the adjoint gradient is reported at initialization.The initial parameter guesses were sampled around the true parameters, and five guesses were used for each time-series length.Source: §G.2 · pp. 17, 19
  • For intervals [0, 50] and [0, 100], initial adjoint gradients were too large to report averages over the optimization procedure.The adjoint comparison therefore reports the gradient at the starting iteration; error bars represent one standard deviation from the mean.Source: §G.2 · pp. 17, 19

G.3 MOCAP experiment

The MOCAP experiment reports its data, modeling assumptions, and training choices, while the video example describes a simulated pendulum and its architecture.

G.3 MOCAP experiment

  • The MOCAP experiment uses a preprocessed dataset from, a Gaussian likelihood, and log-normal priors on diffusion and likelihood variance.Its architecture is described as very similar to those used by [12].Source: §G.3 · p. 19
  • For MOCAP training, the setup uses batch size 512, nested Monte Carlo with R = 10 and S = 10, and a 200-iteration linear KL-weight schedule.The study compares Softplus [12] and tanh nonlinearities and reports the latter's better validation performance.Source: §G.3 · p. 19
  • Tanh nonlinearities provide better validation performance than Softplus, and results average test accuracy across 10 models trained from different random seeds.Models train for 100 epochs, with validation accuracy checked every 10 epochs.Source: §G.3 · p. 19
  • The report uses RMSE ± NEW ERROR to keep error units consistent with the original dataset, approximating RMSE from MSE with a first-order Taylor series.The passage contrasts this with prior reporting of MSE ± ERROR.Source: §G.3 · p. 19

G.4 Neural SDE from video

  • The neural-SDE video example simulates a nonlinear pendulum for 30 seconds, samples its state at 15 Hz, and uses the architecture shown in Figure 9.Figure 8 describes the MOCAP encoder, decoder, deep kernel, and drift function; its likelihood mean is denoted μ(t).Source: §G.4 · p. 19, Figure 8 · p. 20, Figure 9 · p. 21

H Numerical study on the effect of Monte-Carlo parameters

This study varies nested Monte Carlo sample counts for a four-dimensional predator-prey prediction problem. It finds reasonable validation accuracy at three S values, with increased total evaluations as S rises.

H Numerical study on the effect of Monte-Carlo parameters

  • The experiment studies validation RMSE convergence for prediction on a four-dimensional system of two predators and two prey with competitive predator dynamics.Figure 11 shows the encoder, deep kernel, and drift-function architectures used in the parameter study.Source: H Numerical study on the effect of… · p. 22, Figure 11 · p. 23
  • The system is simulated for 300 time units at 10 Hz with Gaussian noise of standard deviation 10^-2.The setup uses batch size 256, 1000 warmup iterations, and learning rate 10^-3.Source: H Numerical study on the effect of… · p. 22
  • S = 10, S = 50, and S = 100 each achieve reasonable validation accuracy; increasing S raises total function evaluations while parallel forward-model evaluations remain constant.Figure 10 plots validation RMSE against iteration and number of function evaluations.Source: Figure 10 · p. 22, H Numerical study on the effect of… · p. 22
Loading 2312.10550v1…