Incorporating Neural Network Structure in the Bayesian Learning Rule
Abstract
We extend the Bayesian learning rule by explicitly incorporating the compositional structure of neural networks. The central observation is that backpropagation and exponential family variational inference share a common Lagrangian duality structure. Leveraging this connection, we develop a new Bayesian learning framework in which adjoint variables backpropagate layerwise sensitivity signals that converts into loss-site natural parameters. For Gaussian weight posteriors and Gaussian layer-state projections, the framework recovers deterministic back-propagation and Hessian backpropagation via delta-method approximations, while providing a unified perspective on localization-style prediction rules and moment matching methods. We also derive a new variant of IVON, clarify how its gradient and Hessian estimates differ from those of the standard version, and demonstrate how neural network structure can be incorporated into Bayesian learning in practice.