Infinitely Deep Bayesian Neural Networks with Stochastic Differential Equations
This work enables more flexible uncertainty quantification in deep learning models for researchers and practitioners, though it builds incrementally on existing Neural ODE frameworks.
The authors tackled scalable approximate inference in continuous-depth Bayesian neural networks by developing gradient-based stochastic variational inference for infinite-parameter settings, achieving competitive performance against discrete-depth alternatives while maintaining memory-efficient training.
We perform scalable approximate inference in continuous-depth Bayesian neural networks. In this model class, uncertainty about separate weights in each layer gives hidden units that follow a stochastic differential equation. We demonstrate gradient-based stochastic variational inference in this infinite-parameter setting, producing arbitrarily-flexible approximate posteriors. We also derive a novel gradient estimator that approaches zero variance as the approximate posterior over weights approaches the true posterior. This approach brings continuous-depth Bayesian neural nets to a competitive comparison against discrete-depth alternatives, while inheriting the memory-efficient training and tunable precision of Neural ODEs.