Efficient and Accurate Gradients for Neural SDEs

Patrick Kidger, James Foster, Xuechen Li, Terry Lyons

Research output: Chapter or section in a book/report/conference proceedingChapter in a published conference proceeding

Abstract

Neural SDEs combine many of the best qualities of both RNNs and SDEs: memory efficient training, high-capacity function approximation, and strong priors on model space. This makes them a natural choice for modelling many types of temporal dynamics. Training a Neural SDE (either as a VAE or as a GAN) requires backpropagating through an SDE solve. This may be done by solving a backwards-in-time SDEwhose solution is the desired parameter gradients. However, this has previously suffered from severe speed and accuracy issues, due to high computational cost and numerical truncation errors. Here, we overcome these issues through several technical innovations. First, we introduce the reversible Heun method. This is a new SDEsolver that is algebraically reversible: eliminating numerical gradient errors, and the first such solver of which we are aware. Moreover it requires half as many function evaluations as comparable solvers, giving up to a 1.98× speedup. Second, we introduce the Brownian Interval: a new, fast, memory efficient, and exact way of sampling and reconstructing Brownian motion. With this we obtain up to a 10.6× speed improvement over previous techniques, which in contrast are both approximate and relatively slow. Third, when specifically training Neural SDEs as GANs (Kidger et al. 2021), we demonstrate how SDE-GANs may be trained through careful weight clipping and choice of activation function. This reduces computational cost (giving up to a 1.87× speedup) and removes the numerical truncation errors associated with gradient penalty. Altogether, we outperform the state-of-the-art by substantial margins, with respect to training speed, and with respect to classification, prediction, and MMD test metrics. We have contributed implementations of all of our techniques to the torchsde library to help facilitate their adoption.
Original languageEnglish
Title of host publicationAdvances in Neural Information Processing systems 34
Subtitle of host publicationNeurIPS 2021
PublisherNeurIPS Proceedings
Publication statusPublished - 31 Dec 2021
Externally publishedYes
EventNeurIPS 2021: Conference on Neural Information Processing Systems - Virtual
Duration: 6 Dec 202112 Dec 2021
https://nips.cc/Conferences/2021

Conference

ConferenceNeurIPS 2021: Conference on Neural Information Processing Systems
Abbreviated titleNeurIPS 2021
Period6/12/2112/12/21
Internet address

Fingerprint

Dive into the research topics of 'Efficient and Accurate Gradients for Neural SDEs'. Together they form a unique fingerprint.

Cite this