Efficient Training of Neural SDEs Using Stochastic Optimal Control

TL;DR

This paper introduces a control theory-inspired hierarchical approach for variational inference in neural SDEs, enabling more efficient training by decomposing control terms and leveraging optimal control for linear components.

Contribution

It proposes a novel decomposition of control terms in neural SDEs, combining stochastic optimal control with neural networks to improve training efficiency and convergence.

Findings

Faster convergence in training neural SDEs.

Reduced computational cost due to optimal linear control.

Maintained expressive power with neural network residuals.

Abstract

We present a hierarchical, control theory inspired method for variational inference (VI) for neural stochastic differential equations (SDEs). While VI for neural SDEs is a promising avenue for uncertainty-aware reasoning in time-series, it is computationally challenging due to the iterative nature of maximizing the ELBO. In this work, we propose to decompose the control term into linear and residual non-linear components and derive an optimal control term for linear SDEs, using stochastic optimal control. Modeling the non-linear component by a neural network, we show how to efficiently train neural SDEs without sacrificing their expressive power. Since the linear part of the control term is optimal and does not need to be learned, the training is initialized at a lower cost and we observe faster convergence.