PaperlayerView exact revision on arXiv

Source-linked AI summary

Bayesian Deep Ensembles via the Neural Tangent Kernel

Bobby He, Balaji Lakshminarayanan, Yee Whye Teh

arXiv:2007.05864v2stat.MLcs.LG

TL;DR

Exact neural-network posterior inference is computationally intensive, while deep ensembles lack a principled Bayesian justification despite strong uncertainty and out-of-distribution performance. The paper adds an untrainable random function to ensemble members, giving a posterior interpretation in the infinite-width limit; finite-width ensembles emulate analytic posteriors when available and often outperform standard ensembles out of distribution.

  • Problem

    Exact neural-network posterior inference is computationally intensive, and deep ensembles lack a principled Bayesian justification despite outperforming existing BNNs in uncertainty quantification and out-of-distribution robustness.

    Source: §1 · p. 1

  • Method

    They add an untrainable random function δ(x) to each NN's forward pass, yielding exact posterior samples at infinite width whose ensemble approximates the NTKGP posterior predictive.

    Source: §1 · p. 2, §3 · p. 4

  • Results

    Bayesian deep ensembles emulate the analytic posterior predictive when available and often outperform standard deep ensembles out of distribution for regression and classification.

    Source: §5 · p. 9

  • Takeaways & Limitations

    A posterior interpretation is compatible with empirical out-of-distribution gains over standard deep ensembles in both regression and classification.

    Source: §5 · p. 9

    The methods may perform worse than standard deep ensembles when confident predictions are not detrimental, and the analyses assume lazy learning without finite-width NTK corrections.

    Source: §5 · p. 9

Abstract

from arXiv · show

We explore the link between deep ensembles and Gaussian processes (GPs) through the lens of the Neural Tangent Kernel (NTK): a recent development in understanding the training dynamics of wide neural networks (NNs). Previous work has shown that even in the infinite width limit, when NNs become GPs, there is no GP posterior interpretation to a deep ensemble trained with squared error loss. We introduce a simple modification to standard deep ensembles training, through addition of a computationally-tractable, randomised and untrainable function to each ensemble member, that enables a posterior interpretation in the infinite width limit. When ensembled together, our trained NNs give an approximation to a posterior predictive distribution, and we prove that our Bayesian deep ensembles make more conservative predictions than standard deep ensembles in the infinite width limit. Finally, using finite width NNs we demonstrate that our Bayesian deep ensembles faithfully emulate the analytic posterior predictive when available, and can outperform standard deep ensembles in various out-of-distribution settings, for both regression and classification tasks.

1 Introduction

The introduction frames the lack of a Bayesian posterior interpretation for standard deep ensembles as an open question, despite their reported uncertainty and out-of-distribution performance. It proposes modifying training with an untrainable random function so independently trained members can approximate a posterior predictive in the infinite-width limit.

1 Introduction

  • Exact computation of the neural-network parameter posterior is intensive, motivating approximations including Laplace, MCMC, variational inference, and Monte-Carlo dropout.The posterior predictive marginalizes over parameters weighted by posterior probabilities and is optimal for predictive risk under a well-specified model.Source: §1 · p. 1
  • Standard deep ensembles are reported to outperform existing Bayesian neural networks in uncertainty quantification and out-of-distribution robustness, yet lack a principled Bayesian justification.Their relationship to Bayesian inference is identified as a key open question.Source: §1 · pp. 1–2
  • A training modification yields exact posterior samples for individual functions in the infinite-width limit, whose ensemble approximates a posterior predictive.The modification adds an untrainable function to each baselearner.Source: §1 · p. 2
  • Randomised priors can yield exact posterior samples under certain conditions, including linearity in parameters; wide neural networks local to initialization have been shown to satisfy that condition.The introduction uses this idea to model observation noise in regression and reviews the Neural Tangent Kernel.Source: §1 · p. 2

2 NTK Background

The section reviews how wide neural networks relate to Gaussian processes through the Neural Tangent Kernel and the linearised training regime. It explains why the Gaussian-process limit of a standard network trained with squared error does not generally provide a posterior interpretation.

2 NTK Background

  • The empirical NTK describes a network's functional gradient in terms of training loss and, at finite training time, remains random because initialization is random.For an MLP under NTK parameterisation, its infinite-width limit is deterministic and constant during gradient training.Source: §2 · p. 2
  • For wide networks, the Neural Tangent Kernel gives a deterministic limit that remains constant during gradient training and depends on architecture choices.Under NTK parameterisation, the empirical kernel converges to this limit as hidden-layer widths grow.Source: §2 · p. 2
  • The linearised regime expands the network function around its random initialization, and in the infinite-width limit its trained function matches the original network under the stated training conditions.The cited conditions include NTK parameterisation, gradient flow, and squared-error loss [23].Source: §2 · p. 2
  • The limiting trained function decomposes into a deterministic mean and a random component that captures predictive uncertainty from the initial function.Regularisation corresponding to observation or output noise can ensure invertibility and alleviate potential convergence issues from fast NTK eigenspectrum decay.Source: §2 · p. 3
  • In the infinite-width limit, the NNGP kernel K differs from the NTK Θ, so squared-error training does not generally yield a Bayesian posterior interpretation for a network or its ensemble.The cited condition for a posterior GP interpretation is that the training kernel and initialization kernel are multiples of one another [23].Source: §2 · p. 3

3 Proposed modification to obtain posterior samples in infinite width limit

Adding an untrainable random function to each network supplies the missing randomness needed for trained wide networks to sample the NTK-GP posterior. Ensembled NTKGP baselearners approximate its posterior predictive, with methods for observation noise, heteroscedasticity, and classification.

3 Proposed modification to obtain posterior samples in infinite width limit

  • The modification adds a random, untrainable function to the network forward pass, supplying randomness missing from the standard trained wide network and yielding an NTK-GP posterior sample in the infinite-width limit.The added function is constructed from an independently sampled parameter initialisation with the final-layer parameters set to zero; its randomness is independent of the standard network initialisation in the limit.Source: §3 · p. 4

3.1 Modelling observation noise

  • For Gaussian observation noise, the method extends training with randomised-prior parameter-space or function-space regularisation, correcting for the difference between the NNGP and NTK.The paper names these approaches NTKGP-param and NTKGP-fn, respectively.Source: §3.1 · p. 5

3.2 Comparison of predictive distributions in infinite width

  • In the infinite-width linearised regime, NTKGP predictive covariance is at least as large as standard deep-ensemble covariance for noiseless observations, and at least as large as randomised-prior covariance with observation noise.Predictive means match between deep ensembles and NTKGP when σ²=0, and between randomised-prior and NTKGP when σ²>0; benefits are not guaranteed under model misspecification.Source: §3.2 · p. 5

3.3 Modelling heteroscedasticity

  • For heteroscedastic regression, the method uses a Gaussian negative log-likelihood and trains an additional output head to predict input-dependent observation variance.The variance head is trained alongside the mean function.Source: §3.3 · pp. 5–6

3.4 NTKGP Ensemble Algorithms

  • The ensemble algorithm independently initialises baselearners and uses fixed, independently sampled target noise when modelling output noise, aiming to approximate the NTKGP posterior predictive.For NTKGP-param, the method anchors at initialisation because the linearised regime holds locally and this reduces parameter-sampling memory costs.Source: §3.4 · p. 6

3.5 Classification methodology

  • For classification, the method treats labels as one-hot regression targets, temperature-scales predictions on held-out validation data, and tunes target scaling to balance in- and out-of-distribution performance.The scaling factor is tuned using a small set of values to match the second moments of randomly initialised baselearners’ training-set logits.Source: §3.5 · p. 6

4 Experiments

Experiments compare Bayesian deep ensembles with analytic posterior predictions and standard ensemble methods across regression and classification tasks. Results show improved out-of-distribution behavior in several settings, with some task-dependent trade-offs.

Toy 1D regression task

  • NTKGP-param ensemble predictive variances match the analytic NTKGP posterior, while other methods make more confident predictions and are overconfident on out-of-distribution data in this example.The comparison uses a toy regression function and 95% predictive confidence intervals.Source: Figure 1 · p. 7, §4 · p. 7

Flight Delays

  • Bayesian deep ensembles perform slightly worse with little or no test shift but fail more gracefully as shift increases, and incur lower error across precision thresholds.The Flight Delays test sets begin at five different chronological points; the error contrast is magnified for more confident predictions.Source: §4 · p. 7, Figure 2 · p. 8

MNIST vs NotMNIST

  • Over 15% lower error at confidence τ=0.6 than all baselines, while MNIST in-distribution error increases around 0.2% for ensemble size 10.The comparison counts all NotMNIST predictions as incorrect; NTKGP ensembles at larger sizes closely match analytic NTKGP predictions on in- and out-of-distribution performance.Source: Figure 3 · p. 8, §4 · p. 8, Figure 4 · p. 9

CIFAR-10 vs SVHN

  • On CIFAR-10, NTKGP methods and RP-fn have the best in-distribution test error, but the out-of-distribution error-versus-confidence plot shows no clear difference.The ensembles are trained on CIFAR-10 and tested on CIFAR-10 and SVHN.Source: §4 · p. 8, Figure 5 · p. 9

5 Discussion

The discussion frames the work as a Bayesian interpretation of deep ensembles, while noting limitations and open directions. The appendix recap contrasts NTK and standard parameterisations.

5 Discussion

  • The paper presents a training modification that gives a posterior predictive interpretation to a wide ensemble, building on NTK results about standard deep ensembles.The discussion reports that Bayesian deep ensembles often outperform standard deep ensembles out of distribution for regression and classification.Source: §5 · p. 9
  • The methods may perform worse than standard deep ensembles when confident predictions are not detrimental, and the analysis is limited to the lazy-learning regime without finite-width NTK corrections.The authors say NTK hyperparameter tuning can alleviate the performance limitation.Source: §5 · p. 9
  • The authors identify NTK hyperparameter tuning for inductive biases, depth, architecture, and activation, as well as unsupervised settings, as future directions.Examples of unsupervised settings include active learning and reinforcement learning.Source: §5 · p. 9

A Recap of standard and NTK parameterisations

  • The two parameterisations give the same functional output distribution at initialisation but different scalings for parameter gradients in the backward pass.In the NTK parameterisation, weight and bias variances are hyperparameters of the infinite-width NTK.Source: A Recap of standard and NTK… · p. 14

B.1 Proof of Proposition 1

The proof material establishes an infinite-width Gaussian-process relationship for the modified initial function and gives covariance orderings for the compared methods.

B.1 Proof of Proposition 1

  • The proof considers joint finite-dimensional convergence of δ and f0 to independent Gaussian processes, using characteristic functions to establish the convergence.It uses the Lévy convergence theorem and Cramér–Wold device; the stated limiting relationship is for arbitrary input sets.Source: §B.1 · pp. 14–15

B.2 Proof of Proposition 2

  • For σ2=0, ΣNTKGP ⪰ ΣDE ⪰ ΣNNGP; for σ2>0, ΣNTKGP ⪰ ΣRP ⪰ ΣNNGP.These covariance orderings are the stated result of Proposition 2.Source: §B.2 · p. 15

C Alternative constructions of NTKGP baselearners

The section gives alternative NTKGP baselearner constructions, including a linearized option with an exact posterior-sample interpretation and a JVP-based function construction. The linearized option assumes a regime that holds throughout parameter space, making initialization irrelevant for strongly convex noisy regression.

C Alternative constructions of NTKGP baselearners

D Regularisation in the NTKGP and RP training procedures

The section relates parameter-space regularization in NTKGP and randomized-prior training to posterior sampling in Bayesian linear regression. It explains how the prior determines the regularizer and why the posterior interpretation depends on accounting for the distinction between the NNGP and NTK.

D Regularisation in the NTKGP and RP training procedures

  • Bayesian linear-regression posterior sampling requires prior-matched L2 regularization, with a Mahalanobis form for diagonal heteroscedastic priors.Source: D Regularisation in the NTKGP and RP… · p. 16
  • With parameters initialized at θ0, parameter-space regularization measures the Mahalanobis distance from θ0 rather than from the origin.Source: D Regularisation in the NTKGP and RP… · p. 16
  • Matching functional outputs across NTK and standard parameterizations holds in the linearized regime, but only NTKGP-trained baselearners gain a posterior interpretation because standard training does not account for the NNGP–NTK difference.The linearized regime holds only locally to the parameter initialization.Source: D Regularisation in the NTKGP and RP… · p. 16

E Additional ensemble algorithms

The additional algorithms specify linearized and function-based NTKGP ensemble members, including a one-JVP randomized-prior construction. The section also describes ensemble aggregation as a Monte Carlo approximation to the NTK GP posterior predictive, with task-specific prediction averaging.

E Additional ensemble algorithms

  • NTKGP-Lin defines each member through a fixed Jacobian feature map applied to its trainable parameters.Source: E Additional ensemble algorithms · p. 17
  • NTKGP-fn uses one reweighted JVP to construct a randomized prior function that changes the initial function distribution from GP(0, K) to GP(0, 2Θ).The construction is chosen to reduce extra computational costs.Source: E Additional ensemble algorithms · p. 17
  • For NTKGP-fn, linearized network dynamics cannot reasonably be assumed throughout training because regularization is toward the origin rather than the initialized parameters.Source: E Additional ensemble algorithms · p. 17

F Aggregating predictions from ensemble members

  • Uniformly averaging member predictions gives a Monte Carlo approximation to the NTK GP posterior predictive; classification averages probabilities, while regression uses a moment-matched Gaussian.Source: F Aggregating predictions from… · p. 17

G Comparison of memory and computation costs for ensemble methods

The cost comparison finds negligible training-time computational overhead for NTKGP methods on fixed-size datasets when additive JVPs are stored in one pass. Extra parameter storage is required, but JVP memory is independent of network depth, and costs scale linearly with dataset size and parameter dimension.

G Comparison of memory and computation costs for ensemble methods

H Scaling for one-hot targets in classification

The classification experiments examine how the one-hot target scale κ affects ensemble comparisons. The authors vary how κ is set and report a trade-off between in-distribution performance and out-of-distribution detection.

H Scaling for one-hot targets in classification

  • For NTKGP classification, the authors scale the one-hot targets by κ; they tune κ² around Cζ₀ to match initial output second moments across C classes.An alternative is to scale the network output, which rescales its NTK; the implementation uses target scaling.Source: H Scaling for one-hot targets in… · p. 18
  • With κ fixed across methods to match the training-set NTK’s initial second moment, NTKGP methods slightly underperform in-distribution but outperform baselines on out-of-distribution detection.The authors describe the scale-dependent trade-off as applying to all ensemble methods, not only Bayesian ensembles.Source: Figure 6 · p. 19, H Scaling for one-hot targets in… · p. 19
  • Figure 6 holds κ constant across ensemble methods and matches it to the training-set NTK’s second moment at initialization.Its error bars represent five independent runs.Source: Figure 6 · p. 19, H Scaling for one-hot targets in… · p. 19

I Experimental Details & additional plots

The appendix gives experiment settings and additional plots across a toy regression example, flight-delay regression, and two image-classification comparisons. The plots report ensemble-size effects on the toy example and entropy distributions for CIFAR-10 and SVHN.

I.2 Flight Delays

  • For Flight Delays, baselearners are four-layer MLPs trained for 10 epochs, with inputs and targets standardized and early stopping based on validation loss.The validation set contains 50,000 examples sampled from the 700,000-example training set.Source: §I.2 · p. 19

I.1 Toy 1d example

  • At small ensemble sizes, NTKGP-param’s analytic mean-predictor approximation degrades relative to RP-param, while predictive uncertainties remain well matched.The authors attribute the mean degradation to more untrainable noise in initialized NTKGP baselearners and leave separate mean and uncertainty baselearners as a possible solution.Source: §I.1 · p. 19, Figure 7 · p. 20

I.3 MNIST vs. NotMNIST

  • MNIST and NotMNIST experiments use two-hidden-layer MLPs and temperature scaling; analytic NTKGP predictions are approximated with Monte Carlo samples.The experiments use half of each test set for resource reasons, and analytic NTKGP target scale κ is not tuned.Source: §I.3 · pp. 19–20

I.4 CIFAR-10 vs SVHN

  • On CIFAR-10 and SVHN, Figure 8 shows a less noticeable difference between ensemble methods than in the MNIST–NotMNIST comparison.The histograms cover in-distribution CIFAR-10 test data and out-of-distribution SVHN test data.Source: Figure 8 · p. 21, §I.4 · p. 21

Related papers

Loading 2007.05864v2…