Stochastic Control for Bayesian Neural Network Training.

Area of Science:

• Machine Learning
• Bayesian Inference
• Stochastic Processes

Background:

• Bayesian models learn from data while quantifying uncertainty.
• Current training methods may lack robustness and adaptive learning rates.
• Variational distributions in Bayesian models encode uncertainty information.

Purpose of the Study:

• To develop a principled method for controlling learning rates in Bayesian models.
• To leverage Bayesian uncertainty information for improved training dynamics.
• To enhance the robustness and adaptivity of Bayesian model optimization.

Main Methods:

• Derivation of a stochastic differential equation for variational parameter training dynamics.
• Application of stochastic optimal control to variational parameters.
• Development of the StochControlSGD optimizer for adaptive learning rates.

Main Results:

• StochControlSGD demonstrates significant robustness to large learning rates.
• The optimizer adaptively and individually controls learning rates for mean and uncertainty parameters.
• Distinct dynamical behaviors observed for mean and uncertainty parameter training regimes.

Conclusions:

• Bayesian uncertainty information can effectively guide the learning procedure.
• StochControlSGD offers a novel approach to optimizing Bayesian neural networks.
• Adaptive control of learning rates improves training stability and performance.