[Paper Review] Radial Bayesian Neural Networks: Beyond Discrete Support In Large-Scale Bayesian Deep Learning

Sebastian Farquhar, Michael A. Osborne+1 more|arXiv (Cornell University)|Jul 1, 2019

Gaussian Processes and Bayesian Inference29 references21 citations

TL;DR

This paper introduces Radial Bayesian Neural Networks (Radial BNNs), a variational inference method that replaces the mean-field Gaussian posterior with a radial distribution in hyperspherical space, avoiding the 'soap-bubble' pathology of high-dimensional Gaussians. The method enables full support over weight space, reduces gradient variance, and outperforms MC dropout and deep ensembles on a large-scale medical imaging task without hyperparameter tuning.

ABSTRACT

We propose Radial Bayesian Neural Networks (BNNs): a variational approximate posterior for BNNs which scales well to large models while maintaining a distribution over weight-space with full support. Other scalable Bayesian deep learning methods, like MC dropout or deep ensembles, have discrete support-they assign zero probability to almost all of the weight-space. Unlike these discrete support methods, Radial BNNs' full support makes them suitable for use as a prior for sequential inference. In addition, they solve the conceptual challenges with the a priori implausibility of weight distributions with discrete support. The Radial BNN is motivated by avoiding a sampling problem in 'mean-field' variational inference (MFVI) caused by the so-called 'soap-bubble' pathology of multivariate Gaussians. We show that, unlike MFVI, Radial BNNs are robust to hyperparameters and can be efficiently applied to a challenging real-world medical application without needing ad-hoc tweaks and intensive tuning. In fact, in this setting Radial BNNs out-perform discrete-support methods like MC dropout. Lastly, by using Radial BNNs as a theoretically principled, robust alternative to MFVI we make significant strides in a Bayesian continual learning evaluation.

Motivation & Objective

• To address the pathological 'soap-bubble' effect in mean-field variational inference (MFVI), where samples from high-dimensional Gaussians are unrepresentative of the mean and cause high gradient variance.
• To develop a scalable Bayesian neural network method with full support over weight space, avoiding the discrete support of methods like MC dropout and deep ensembles.
• To create a robust, hyperparameter-insensitive inference method suitable for real-world, large-scale deep learning tasks such as medical image classification.
• To enable the use of BNN posteriors as data-dependent priors in continual learning, overcoming the limitations of zero-probability priors.

Proposed method

• Propose a radial approximate posterior distribution in a hyperspherical space for each network layer, transforming it into weight-space via a radial basis function.
• Define the posterior as a distribution on the unit sphere, with radial components drawn from a univariate distribution (e.g., Gamma or inverse-Gamma), ensuring samples cluster near the mean.
• Derive an analytic expression for the evidence lower bound (ELBO) that enables efficient training without explicit coordinate transformations.
• Use rejection sampling to truncate the Gaussian prior in MFVI as a baseline comparison, demonstrating that low-variance estimators improve performance despite bias.
• Implement the method with standard deep learning frameworks, maintaining training speed and ease of implementation comparable to MFVI.
• Apply the method to a 15M-parameter model on diabetic retinopathy detection with high-dimensional inputs (~230K dimensions).

Experimental results

Research questions

• RQ1 Can a radial posterior distribution mitigate the 'soap-bubble' pathology in high-dimensional mean-field variational inference?
• RQ2 Does a full-support posterior improve robustness to hyperparameter choices in large-scale Bayesian deep learning?
• RQ3 Can Radial BNNs outperform state-of-the-art methods like MC dropout and deep ensembles on real-world, high-dimensional medical imaging tasks?
• RQ4 Can Radial BNN posteriors serve as effective, data-dependent priors in continual learning settings?
• RQ5 Is the improved performance due to reduced gradient variance from more representative posterior samples?

Key findings

• Radial BNNs outperform MC dropout and deep ensembles on a real-world diabetic retinopathy classification task, achieving lower test RMSE and higher predictive log-likelihood.
• On the UCI regression benchmark, Radial BNNs achieve state-of-the-art performance, with a mean test RMSE of 0.64 ± 0.01 on the Wine dataset and 0.07 ± 0.00 on Kin8nm, outperforming MFVI and other baselines.
• The method is robust to hyperparameter choice, showing stable performance across different weight initialization variances and learning rates.
• Truncated MFVI (using rejection sampling) improves performance but remains inferior to Radial BNNs, indicating that the core issue in MFVI is high gradient variance.
• Radial BNNs maintain full support over weight space, enabling their use as priors in continual learning, where they prevent catastrophic forgetting more effectively than discrete-support methods.
• The analytic ELBO derivation enables training as efficiently as standard MFVI, with no additional computational overhead.