Differential Bayesian Neural Nets

TL;DR

This paper introduces a Bayesian extension of Neural ODEs using stochastic differential equations and Bayesian neural networks, enabling uncertainty quantification while maintaining model expressiveness.

Contribution

It develops a Bayesian Neural ODE framework with SDEs and SGLD inference, improving stability and model fit over non-Bayesian Neural ODEs.

Findings

Enhanced stability on synthetic time series tasks

Better model fit on UCI regression benchmarks

Effective uncertainty quantification in predictions

Abstract

Neural Ordinary Differential Equations (N-ODEs) are a powerful building block for learning systems, which extend residual networks to a continuous-time dynamical system. We propose a Bayesian version of N-ODEs that enables well-calibrated quantification of prediction uncertainty, while maintaining the expressive power of their deterministic counterpart. We assign Bayesian Neural Nets (BNNs) to both the drift and the diffusion terms of a Stochastic Differential Equation (SDE) that models the flow of the activation map in time. We infer the posterior on the BNN weights using a straightforward adaptation of Stochastic Gradient Langevin Dynamics (SGLD). We illustrate significantly improved stability on two synthetic time series prediction tasks and report better model fit on UCI regression benchmarks with our method when compared to its non-Bayesian counterpart.