PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficient and Scalable Bayesian Neural Nets with Rank-1 Factors

Michael W. Dusenberry, Ghassen Jerfel, Yeming Wen, Yi-An Ma, Jasper Snoek, Katherine Heller, Balaji Lakshminarayanan, Dustin Tran

arXiv:2005.07186v2cs.LGstat.ML

TL;DR

BNNs and deep ensembles provide useful uncertainty estimates but face underfitting and efficiency limitations. The paper addresses these issues with rank-1 BNNs and mixture variational posteriors, finding strong predictive and uncertainty performance across image and medical benchmarks. The approach remains subject to training and modeling limitations, especially for heavy-tailed variants and larger mixture sizes.

  • Problem

    BNNs and deep ensembles offer uncertainty benefits but respectively suffer from underfitting or substantial parameter and runtime costs.

  • Method

    The paper parameterizes distributions over rank-1 factors of weight matrices and trains rank-1 layers with mixture variational distributions.

  • Results

    Across ImageNet, CIFAR-10/100, and MIMIC-III, rank-1 BNNs achieve the best reported predictive and uncertainty metrics across in- and out-of-distribution data.

  • Takeaways & Limitations

    Rank-1 BNNs combine parameter efficiency and scalability with multimodal weight solutions while retaining strong accuracy, likelihood, and calibration performance.

  • Takeaways & Limitations

    Heavy-tailed Cauchy variants underperform in some settings, and larger mixtures can reduce per-component batch size under fixed-memory training.

Abstract

from arXiv · show

Bayesian neural networks (BNNs) demonstrate promising success in improving the robustness and uncertainty quantification of modern deep learning. However, they generally struggle with underfitting at scale and parameter efficiency. On the other hand, deep ensembles have emerged as alternatives for uncertainty quantification that, while outperforming BNNs on certain problems, also suffer from efficiency issues. It remains unclear how to combine the strengths of these two approaches and remediate their common issues. To tackle this challenge, we propose a rank-1 parameterization of BNNs, where each weight matrix involves only a distribution on a rank-1 subspace. We also revisit the use of mixture approximate posteriors to capture multiple modes, where unlike typical mixtures, this approach admits a significantly smaller memory increase (e.g., only a 0.4% increase for a ResNet-50 mixture of size 10). We perform a systematic empirical study on the choices of prior, variational posterior, and methods to improve training. For ResNet-50 on ImageNet, Wide ResNet 28-10 on CIFAR-10/100, and an RNN on MIMIC-III, rank-1 BNNs achieve state-of-the-art performance across log-likelihood, accuracy, and calibration on the test sets and out-of-distribution variants.

1. Introduction

BNNs offer robustness and uncertainty quantification but often underperform and impose substantial scaling costs. The paper proposes rank-1 BNNs and mixture posteriors to combine parameter efficiency with multimodal uncertainty modeling, achieving strong results across diverse benchmarks.

  • BNNs can improve robustness, uncertainty quantification, graceful failure, and generalization to out-of-distribution examples.
  • BNNs often underperform on accuracy and scale worse than simpler baselines, partly because effective parameterizations, priors, posteriors, and optimization strategies remain unclear.
  • Maintaining weight distributions is costly: mean-field inference doubles parameters, while ensembles or five MCMC samples require 5x the weights.
  • Rank-1 BNNs place each weight matrix’s distribution on a rank-1 subspace and use mixture posteriors to capture multimodal solutions efficiently.
  • For a size-10 ResNet-50 mixture, the rank-1 approach adds only 0.4% parameters instead of 900%, with Bayesian inference performed over thousands of dimensions.
  • Across ImageNet, CIFAR-10/100, and MIMIC-III, rank-1 BNNs outperform stated baselines on log-likelihood, accuracy, and calibration in-distribution and under distribution shift.

2. Background

The background introduces variational Bayesian inference and contrasts conventional deep ensembles with BatchEnsemble’s low-rank parameter sharing. BatchEnsemble represents member-specific weight changes through rank-1 vectors and vectorizes their computation across batches and operators.

  • Variational inference selects q(W) and minimizes its KL divergence to the posterior p(W | X, y).
  • The variational objective upper-bounds negative log-marginal likelihood and can be interpreted as approximate description length.
  • BNNs often underfit because prior and posterior choices, training stabilization, and sampling variance complicate optimization; even mean-field normals cost 2x parameters.
  • Deep ensembles aggregate predictions from multiple separately trained networks but require training and predicting with multiple parameter copies.
  • BatchEnsemble shares the original weight matrix while assigning each of K members trainable vectors r_k and s_k.
  • The rank-1 vectors r_k and s_k correspond to elementwise multiplication of input neurons and pre-activations, enabling efficient batch vectorization.
  • This vectorization extends beyond linear layers to convolution and recurrence.

3. Rank-1 Bayesian Neural Nets

The paper parameterizes Bayesian neural networks with distributions over rank-1 factors rather than full weight matrices, enabling efficient stochastic modeling and scalable mixtures. Experiments examine priors, posterior mixtures, likelihood objectives, and ensemble diversity, finding that hierarchical priors and non-degenerate mixtures improve robustness and uncertainty-related performance.

  • Rank-1 parameterization: Rank-1 BNNs place distributions on factor vectors r and s while treating the shared weight matrix W as deterministic.For each weight matrix, W′ = W ◦ rs^T, reducing stochastic sampling from millions of weights to rank-1 factors.
  • Rank-1 parameterization: Rank-1 distributions reduce sampling variance because stochasticity comes from factor vectors rather than full weight matrices.The paper states that Equation 1 also permits sampling new r and s vectors for each example and supports arbitrary q(r) and q(s).
  • Rank-1 priors: Cauchy and Normal factor priors encourage sparse weight posteriors, while Cauchy has less mass near zero and heavier tails.The paper also relates factor priors to hierarchical weight priors and noncentered parameterizations.
  • Rank-1 priors: Priors over both input and output factors generalize better on corruptions than priors over only r or only s.The three settings perform comparably in accuracy on the test set, with differences mainly in NLL and ECE.
  • Mixture posteriors: 0.4% additional parameters support a ResNet-50 mixture posterior with K = 10, compared with a 900% increase for deep ensembles.On CIFAR-10 and CIFAR-10-C, performance improves as K increases from 1 to 8 but declines at K = 16.
  • Likelihood objectives: The average per-component log-likelihood consistently outperforms the log-mixture likelihood on the test set despite comparable training metrics.The log-mixture likelihood offers a tighter bound but couples gradients across mixture components, unlike independently trained deep ensembles.

4. Experiments

Across ImageNet, CIFAR, and MIMIC-III, rank-1 BNNs are evaluated against deterministic, Bayesian, ensemble, and related low-dimensional baselines on predictive performance, robustness, and calibration. They generally improve accuracy, likelihood, calibration, and distribution-shift robustness while using substantially fewer parameters than deep ensembles, though heavy-tailed variants show a recurrent limitation in some settings.

  • Experimental setup: Experiments cover ImageNet and ImageNet-C with ResNet-50, CIFAR-10/100 and corrupted variants with WRN-28-10, and MIMIC-III mortality prediction with recurrent models.Image and EHR tasks are compared using NLL, accuracy or AUC-PR, and ECE, with corrupted datasets probing robustness under distribution shift.
  • CIFAR-10 and CIFAR-10-C: 1.02 to 0.74 NLL, 3.7% accuracy, and 0.05 calibration improvements are reported for rank-1 BNNs on CIFAR-10-C relative to BatchEnsemble.Both Gaussian and Cauchy rank-1 BNNs also outperform similarly sized baselines on CIFAR-10 NLL, accuracy, and ECE.
  • Efficiency versus ensembles: Rank-1 BNNs outperform similarly sized deep ensembles on accuracy and outperform them under distribution shift on ECE, accuracy, and NLL.On WRN-28-10, they achieve robustness and uncertainty-calibration gains with 4 times fewer parameters, and at equal parameter counts outperform deep ensembles across all metrics.
  • CIFAR-100 and CIFAR-100-C: Cauchy rank-1 BNN mixtures significantly outperform BatchEnsemble and similarly sized deep ensembles across CIFAR-100 metrics, while Gaussian rank-1 BNNs lead baselines on NLL and ECE.The reported gains extend to CIFAR-100-C, supporting improved predictive performance, robustness, and uncertainty calibration without the expected accuracy compromise from heavy-tailed priors.
  • MIMIC-III mortality prediction: The rank-1 Bayesian RNN outperforms all MIMIC-III baselines across NLL, AUC-PR, and ECE using 25 evaluation samples and 25 random seeds.The method replaces variational LSTM and affine layers with a rank-1 counter-mixture model and adapts across architectures and data modalities.
  • Observed limitation: Cauchy rank-1 RNNs improve NLL and ECE but do not perform as well as Gaussian rank-1 RNNs on AUC-PR.The authors hypothesize that instability from heavy-tailed posterior samples contributes to underfitting in deep or recurrent architectures.

5. Related Work

Prior work has pursued efficient Bayesian inference and ensembling through low-dimensional or rank-1 representations, while scaling BNNs remains costly. Rank-1 BNNs build on these directions for broader architectural and distributional flexibility.

  • Rank-1 factors support exact weight sampling for convolutions and recurrence and arbitrary location-scale distributions.
  • Correlated weight noise and Flipout provide related variance-reduction approaches for sampling weights in settings beyond standard feedforward layers.
  • Monte Carlo Dropout and BatchEnsemble are parameter-efficient alternatives that approximate ensembling through stochastic or rank-1 perturbations.
  • Earlier ImageNet BNN studies improved selected metrics but used at least 2x or 9x the neural-network weights.
  • Low-dimensional Bayesian inference reduces the search space, but PCA-based projection can still require 5x the number of weights.

6. Discussion

Across ImageNet, CIFAR-10, CIFAR-100, and MIMIC-III, rank-1 BNNs combine mixture variational distributions with low-dimensional Bayesian inference and achieve the best reported predictive and uncertainty metrics in- and out-of-distribution.

  • Rank-1 BNNs place priors over rank-1 factors and train with mixture variational distributions, making Bayesian inference parameter-efficient and scalable.
  • Across four datasets, rank-1 BNNs achieve the best results on predictive and uncertainty metrics for both in-distribution and out-of-distribution data.

A. Variance Structure of the Rank-1 Perturbations

The variance-structure analysis compares full-rank weight perturbations with rank-1 perturbations around a local optimum. Under multiplicative covariance structure, the rank-1 representation can encode the same broad class of score-function fluctuations efficiently.

  • The analysis compares local score-function fluctuations induced by full-rank weight matrices and rank-1 representations around a local optimum.
  • The theoretical setup assumes a multiplicative covariance structure for full-rank perturbations and maps it to perturbations of the rank-1 factors.
  • Theorem 1 establishes correspondence between covariance in weight perturbations and covariance in the rank-1 factors.
  • Because the covariance matrix can be any symmetric positive semi-definite matrix, rank-1 perturbations can encode a wide range of fluctuations in the weights.
  • The construction is especially suited to multiplicative noise, which motivates the rank-1 parameterization’s variance representation.
  • The proof develops local gradient, Hessian, and variance expressions in the full-rank weight space and the low-rank vector space.

B. Additional Experimental Details and Hyperparameters

The experiments evaluate mixture Gaussian and Cauchy rank-1 BNNs across image and EHR tasks, with implementation choices covering initialization, priors, training schedules, architectures, and baselines. Ablations favor deterministic standard-deviation initialization and real-valued factor priors in key settings.

  • Experiments use mixtures of Gaussian and Cauchy priors and variational posteriors, averaging results over 10 image-task runs and 25 EHR-task runs.
  • The standard configuration fixes four mixture components and varies prior standard deviation and variational-posterior mean initialization.
  • Rank-1 factors are applied to all layers except normalization and MIMIC-III embedding layers, while biases remain non-Bayesian.
  • Training uses linear KL annealing over two-thirds of the epochs, with task-specific training budgets for CIFAR, ImageNet, and MIMIC-III.
  • Deterministic initialization sets variational standard deviations near zero so early optimization resembles a deterministic network and helps prevent underfitting.
  • Accuracy and calibration decrease as initialized dropout rate increases, while NLL remains roughly unchanged; deterministic initialization is recommended for its accuracy advantage.
  • Real-valued r and s priors match positive-valued priors on test accuracy but outperform them on test NLL and ECE.

C.3. Number of Evaluation Samples

Using multiple weight samples per mixture component during evaluation improves metric performance, especially on corrupted data, even though training uses one sample.

  • More evaluation samples improve metric performance, with significant gains across all corrupted metrics.The model was trained using only a single weight sample per mixture component and example.

D. Additional Discussion and Future Directions

The discussion highlights gains from heavy-tailed priors and multiple evaluation samples, while identifying scaling, training stability, and higher-rank extensions as future directions.

  • Future Directions: Future work includes scaling to larger ImageNet models, applying rank-1 parameterization to MCMC, using independent chains, and extending to higher-rank factors.Prior work suggests diminishing returns that practically stop at ranks 3 or 5.
  • Additional Discussion: Heavy-tailed priors on low-dimensional subspaces improve robustness and uncertainty calibration while maintaining or improving accuracy.The authors associate this with heavier tails covering more points in loss-landscape valleys.
  • Additional Discussion: Additional evaluation samples improve corrupted metrics, while Cauchy rank-1 BNNs with four samples outperform Gaussian variants on all metrics except ECE.Training still uses a single weight sample per mixture component and example.
  • Additional Discussion: Minimizing the upper bound in Eq. 13 is easier and can improve generalization, while Eq. 12 and Eq. 14 are equivalent for classification predictive distributions.The equivalence does not generally hold when predictive parameters are marginalized before computing negative log likelihood.

F.1. CIFAR-10-C Results

The CIFAR-10-C results report median performance across corruption types and increasing skew-intensity settings.

  • Figure 10 reports median performance across corruption types for increasing skew-intensity settings on CIFAR-10-C.

F.2. CIFAR-100-C Results

The results cover CIFAR-100-C and ImageNet-C, reporting median performance across corruption types and increasing skew-intensity settings.

  • Figure 11 reports median performance across corruption types for increasing skew-intensity settings on CIFAR-100-C.
  • Figure 12 reports median performance across corruption types for increasing skew-intensity settings on ImageNet-C.
Loading 2005.07186v2…