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 language | English |
---|---|
Title of host publication | Advances in Neural Information Processing systems 34 |
Subtitle of host publication | NeurIPS 2021 |
Publisher | NeurIPS Proceedings |
Publication status | Published - 31 Dec 2021 |
Externally published | Yes |
Event | NeurIPS 2021: Conference on Neural Information Processing Systems - Virtual Duration: 6 Dec 2021 → 12 Dec 2021 https://nips.cc/Conferences/2021 |
Conference
Conference | NeurIPS 2021: Conference on Neural Information Processing Systems |
---|---|
Abbreviated title | NeurIPS 2021 |
Period | 6/12/21 → 12/12/21 |
Internet address |