Stochastic Control for Bayesian Neural Network Training

Ludwig Winkler; César Ojeda; Manfred Opper

doi:10.3390/e24081097

Stochastic Control for Bayesian Neural Network Training

Entropy (Basel). 2022 Aug 9;24(8):1097. doi: 10.3390/e24081097.

Authors

Ludwig Winkler¹, César Ojeda², Manfred Opper^{2

3}

Affiliations

¹ Machine Learning Group, Technische Universität Berlin, 10623 Berlin, Germany.
² Artificial Intelligence Group, Technische Universität Berlin, 10623 Berlin, Germany.
³ Centre for Systems Modelling and Quantitative Biomedicine, University of Birmingham, Birmingham B15 2TT, UK.

Abstract

In this paper, we propose to leverage the Bayesian uncertainty information encoded in parameter distributions to inform the learning procedure for Bayesian models. We derive a first principle stochastic differential equation for the training dynamics of the mean and uncertainty parameter in the variational distributions. On the basis of the derived Bayesian stochastic differential equation, we apply the methodology of stochastic optimal control on the variational parameters to obtain individually controlled learning rates. We show that the resulting optimizer, StochControlSGD, is significantly more robust to large learning rates and can adaptively and individually control the learning rates of the variational parameters. The evolution of the control suggests separate and distinct dynamical behaviours in the training regimes for the mean and uncertainty parameters in Bayesian neural networks.

Keywords: Bayesian inference; Bayesian neural networks; learning.

Grants and funding

The research of M.O. was partially funded by the Deutsche Forschungsgemeinschaft (DFG)—Project-ID 318763901—SFB1294. The research of L.W. and C.O. was funded by the BIFOLD-Berlin Institute for the Foundations of Learning and Data (ref. 01IS18025A and ref. 01IS18037A).