Infinitely Deep Bayesian Neural Networks with Stochastic Differential Equations
Winnie Xu, Ricky T. Q. Chen, Xuechen Li, David Duvenaud
Abstract
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.
BibTeX
@InProceedings{pmlr-v151-xu22a,
title = { Infinitely Deep Bayesian Neural Networks with Stochastic Differential Equations },
author = {Xu, Winnie and Chen, Ricky T. Q. and Li, Xuechen and Duvenaud, David},
booktitle = {Proceedings of The 25th International Conference on Artificial Intelligence and Statistics},
pages = {721--738},
year = {2022},
editor = {Camps-Valls, Gustau and Ruiz, Francisco J. R. and Valera, Isabel},
volume = {151},
series = {Proceedings of Machine Learning Research},
month = {28--30 Mar},
publisher = {PMLR},
pdf = {https://proceedings.mlr.press/v151/xu22a/xu22a.pdf},
url = {https://proceedings.mlr.press/v151/xu22a.html},
abstract = { 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. }
}