We derive reverse-mode (or adjoint) automatic differentiation for solutions of stochastic differential equations (SDEs), allowing time-efficient and constant-memory computation of pathwise gradients, a continuous-time analogue of the reparameterization trick.
 Specifically, we construct a backward SDE whose solution is the gradient and provide conditions under which numerical solutions converge.
 We also combine our stochastic adjoint approach with a stochastic variational inference scheme for continuous-time SDE models, allowing us to learn distributions over functions using stochastic gradient descent.
 Our latent SDE model achieves competitive performance compared to existing approaches on time series modeling.
 Deterministic dynamical systems can often be modeled by ordinary differential equations (ODEs). For training, a memory-efficient implementation of the adjoint sensitivity method (Chen et al., 2018) effectively computes gradients through ODE solutions with constant memory cost. Stochastic differential equations (SDEs) are a generalization of ODEs which incorporate instantaneous noise into their dynamics (Arnold, 1974; Øksendal, 2003) . They are a natural fit for modeling phenomena governed by small and unobserved interactions. In this paper, we generalize the adjoint method to dynamics defined by SDEs resulting in an approach which we call the stochastic adjoint sensitivity method. Building on theoretical advances by Kunita (2019), we derive a memory-efficient adjoint method whereby we simultaneously reconstruct the original trajectory and evaluate the gradients by solving a backward SDE (in the sense of Kunita (2019)) whose formulation we detail in Section 3. Computationally, in order to retrace the original trajectory during the backward pass, we need to reuse noise samples generated in the forward pass. In Section 4, we give an algorithm that allows arbitrarily-precise querying of a Brownian motion realization at any time point, while only storing a single random seed. Overall, this results in a constant-memory algorithm that approximates the gradient arbitrarily well as step size reduces by computing vector-Jacobian products a constant number of times per-iteration. See Table 2 for a comparison of our method against previous approaches in terms of asymptotic time and memory complexity. We incorporate SDEs into a stochastic variational inference framework, whereby we efficiently compute likelihood ratios and backpropagate through the evidence lower bound using our adjoint approach. This effectively generalizes existing model families such as latent ODEs (Rubanova et al., 2019) and deep Kalman filters (Krishnan et al., 2017) .
We present a constant memory gradient computation procedure through solutions of stochastic differential equations (SDEs) and apply the method for learning latent SDE models.
{"src": "scitldr-AIC_train", "paper_id": "HyeL9yh4KH"}