Natural gradient descent for improving variational inference based classification of radio galaxies

TL;DR

This paper investigates the use of natural gradient descent, specifically the iVON algorithm, to enhance variational inference in Bayesian neural networks for radio galaxy classification, leading to better uncertainty calibration and out-of-distribution detection.

Contribution

It demonstrates that natural gradient descent improves uncertainty calibration and out-of-distribution detection in BNNs for radio galaxy classification, highlighting the importance of optimizer choice.

Findings

iVON improves uncertainty calibration over previous methods.

Models trained with iVON can distinguish out-of-distribution data.

The cold posterior effect persists with iVON training.

Abstract

Bayesian neural networks (BNNs) are most commonly optimised with first-order optimisers such as stochastic gradient descent. However, when optimising for parameters of probabilistic models, incorporating second order information during optimisation can lead to a more direct path in the distribution space and faster convergence. In this work we examine whether using natural gradient descent can improve the performance of variational inference based classification of radio galaxies. We use the Improved Variational Online Newton (iVON) algorithm and compare its performance against a recent benchmark for BNNs for radio galaxy classification. We find that iVON results in better uncertainty calibration out of all the methods previously considered while providing similar predictive performance to the best performing inference methods such as Hamiltonian Monte Carlo and Bayes by Backprop based variational inference. Models trained with iVON can distinguish far out-of-distribution optical galaxy data, but they cannot reliably detect radio galaxy images from a telescope with different resolution and sensitivity. We find that the cold posterior effect persists in the models trained with iVON. Our results suggest that the choice of the optimiser can lead to qualitatively different solutions and future work using probabilistic neural network models should carefully consider the inductive biases being encoded through the optimisation process, in addition to the data, architecture and inference method.