Disentangling the Predictive Variance of Deep Ensembles through the Neural Tangent Kernel
Authors: Seijin Kobayashi
Pau Vilimelis Aceituno
Johannes von Oswald
Venue: NeurIPS 2022 Accept (2022)
Keywords: Deep Ensembles
Neural Tangent Kernel
Out-of-Distribution

Abstract: Identifying unfamiliar inputs, also known as out-of-distribution (OOD) detection, is a crucial property of any decision making process. A simple and empirically validated technique is based on deep ensembles where the variance of predictions over different neural networks acts as a substitute for input uncertainty. Nevertheless, a theoretical understanding of the inductive biases leading to the performance of deep ensemble's uncertainty estimation is missing. To improve our description of their behavior, we study deep ensembles with large layer widths operating in simplified linear training regimes, in which the functions trained with gradient descent can be described by the neural tangent kernel. We identify two sources of noise, each inducing a distinct inductive bias in the predictive variance at initialization. We further show theoretically and empirically that both noise sources affect the predictive variance of non-linear deep ensembles in toy models and realistic settings after training. Finally, we propose practical ways to eliminate part of these noise sources leading to significant changes and improved OOD detection in trained deep ensembles.

Disentangling the Predictive Variance of Deep Ensembles through the Neural Tangent Kernel
Seijin Kobayashi
Pau Vilimelis Aceituno Institute of Neuroinformatics
Department of Computer Science
ETH Zürich seijinkGethz.ch
University of Zürich & ETH Zürich pau@ini.. ethz. ch
Johannes von Oswald
Department of Computer Science
ETH Zürich voswaldj detddd.@etz.... ch
Abstract
Identifying unfamiliar inputs, also known as out-of-distribution (OOD) detection, is crucial property of any decision making process. A simple and empirically val- idated technique based on deep ensembles where the variance of predictions over different neural networks acts as a substitute for input uncertainty. Nevertheless, a theoretical understanding of the inductive biases leading to the performance of deep ensemble's uncertainty estimation is missing. To improve our description of their behavior, study deep ensembles with large layer widths operating in simplified linear training regimes, which the functions trained with gradient descent can be described by the neural tangent kernel. We identify two sources of noise, each inducing distinct inductive bias in the predictive variance at initialization. We further show theoretically and empirically that both noise sources affect the pre- dictive variance non-linear deep ensembles in toy models and realistic settings after training. Finally, we propose practical ways to eliminate part of these noise sources leading to significant changes and improved OOD detection in trained deep ensembles.
Introduction
Modern artificial intelligence uses intricate deep neural networks to process data, make predictions and take actions. One of the crucial steps toward allowing these agents to act in the real world is to incorporate a reliable mechanism for estimating uncertainty in particular when human lives at risk [1 2). Although the ongoing success of deep learning is remarkable, the increasing data, model and training algorithm complexity make thorough understanding of their inner workings increasingly difficult. This applies when trying to understand when and why a system is certain or uncertain about given output and is therefore the topic of numerous publications 10].
Principled mechanisms for uncertainty quantification would rely on Bayesian inference with an appropriate prior. This has led to the development of (approximate) Bayesian inference methods for deep neural networks (TIHTS). Simply aggregating an ensemble of models [16] and using disagreement of their predictions as substitute for uncertainty has gained popularity. However, the theoretical justification of deep ensembles remains a matter debate, see Wilson and Izmailov [17]. Although link between Bayesian inference and deep ensembles understanding of the widely adopted standard deep ensemble and it's predictive distribution still missing (20 21]. Note that even for principled Bayesian approaches there is valid theoretical or practical OOD guarantee without a proper definition of
be obtained, see [18] [9], an
data [22
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

One avenue to simplify the analyses of deep neural networks that gained lot of attention in
recent years is increase the layer width to infinity [23 24 or very large values [5 26].
In the former regime, an intriguing equivalence of infinitely wide deep networks at initialization and Gaussian processes allows for exact Bayesian inference and therefore principled uncertainty estimation. Although it is not possible to generally derive Bayesian posterior for trained infinite or finite layer width networks, the resulting model predictions can be expressed analytically by kernels. Given this favorable mathematical description, the question of how powerful and similar these models are compared to their arguably black-box counterparts arises, with e.g. moderate width, complex optimizers and training stochasticity [25] 32).
In this paper, we leverage this tractable description of trained neural networks and take first step towards understanding the predictive distribution of neural networks ensembles with large but finite width. Building on top of the various studies mentioned, we do so by studying the case where these networks can be described by kernel and study the effect of two distinct noise sources stemming from the network initialization: The noise in the functional initialization of the network and the initialization noise of the gradient, which affects the training and therefore the kernel. As we will show, these noise sources will affect the predictive distributions differently and influence the network's s generalization on in- and out-of-distribution data.
Our contributions are the following:
We provide first order approximation of the predictive variance of ensemble of linearly trained, finite- neural networks. We identify interpretable terms in the refined variance description, originating from 2 distinct noise sources, and further provide their analytical expression for single layer neural networks with ReI.U non-linearities
We show theoretically that under mild assumptions these refined variance terms survive nonlinear training for sufficiently large width, and therefore contribute to the predictive variance of non-linearly trained deep ensembles. Crucially, our result suggests that any finer description of the predictive variance of linearized ensemble can be erased by nonlinear training.
We conduct empirical studies validating our theoretical results, and investigate how the different variance terms influence generalization on in and out-of-distrbuttion We highlight the practical implications of our theory by proposing simple methods to isolate noise sources in realistic settings which can lead to improved OOD detection!
2 Neural network ensembles and their relations to kernels
Let fo f(, ho - Rh denote neural network parameterized by the weights @ E IR.. The weights consist of weight matrices and bias vectors {(Wi, bi) describing the following feed-forward computation beginning with the input data z
Tn Wi+1 ne + 11+1 with /'-1 ((241)). Vhu
(1)
Here h; is the dimension of the vector x' and o is a pointwise non-linearity such the softplus
log(1 [ + e*)) or Rectified Linear Unit i.c. max(0, æ) (ReLU) B3 We follow Jacot et al. 24] and use ow = /2 to control the standard deviation of the initialised weights W bf ~ NO0 1).
Given a set of N datapoints X = (x:)osieN E R Nxho and targets y = (y4 Josssn E NNbL, we consider regression problems with the goal of finding 0* which minimizes the mean squared error (MSE) loss (0)) = Σi [x,, 0) - ys 11.. For ease of notation, we denote by f(X, 0) E RNN- the vectorized evaluation of f on each datapoint and y E N- the target vector for the entire dataset. As the widths of the hidden layers grow towards infinity, the distribution of outputs at initialization fæx converges to multivariate gaussian distribution due to the Central Limit Theorem (23 The resulting function
then accurately be described as zero-mean Gaussian process, coined Neural Network Gaussian Process (NNGP), where the covariance of a pair of output neurons i, for data x and x! is given by the kernel
'Source code for all experiments: htt.th/t.wooeeee....oooooaaaaaa. predvar

---

K(x,, 2'), lim (r, (o) 06)]
(2)
h-000
with h = min(h1, 11-1). This equivalence can be used to analytically compute the Bayesian posterior of infinitely wide Bayesian neural networks [4].
On the other hand infinite width models trained via gradient descent (GD) can be described by the Neural Tangent Kernel (NTK). Given 0, the NTK O₂ of so is matrix in RN RR--zz with the (injr) given as the following dot product
(3) where we consider without loss of generality the output dimension of to be h₂ = 1 for ease of notation. Furthermore, we denote Og(X, == Vof(X, (X, 0) the matrix and Oz(x', X)) ) := Vaf(æ', (X), the vector form of the NTK while highlighting the dependen- cies on different datapoints.
(Vof(z,, 0), Vef(z,,
Lee et al. [25] showed that for sufficiently wide networks under common parametrizations, the gradient descent dynamics of the model with sufficiently small learning rate behaves closely to its linearly trained counterpart, i.c. its first-order Taylor expansion in parameter space. In this gradient flow regime, after training on the mean squared error converges, we can rewrite the predictions of the linearly trained models in the following closed-form:
flm(x) =f(x=f(xx,)) 0o) + Qo (I, X) f(x, 80))
(4)
where 2o, (x) == ON, (I, x), xx with Op the NTK at initialization, i.e. of f(, 0o).
The linearization error throughout training suP [.. f(() is further shown to decrease with the width of the network, bounded by O Note that one can also linearize the dynamics without increasing the width of neural network but by simply changing its output scaling (26..
When moving from finite the infinite width limit the training of a multilayer perceptron (MLP) can again be described with the NTK, which now converges to deterministic kernel O [24] result which extends to convolutional neural networks (27] and other common architectures 35
fully trained neural network model can then be expressed as
for (x) = (ar 0o) O (x. X) (2, -(xx0))))
(5)
where /((X),z)00) ~
(X,x))
2.1 Predictive distribution of linearly trained deep ensembles
In this Section, study in detail the predictive distribution of ensembles of linearly trained models, ie. the distribution of plin (x) given z over random initializations Og. In particular, for given data x,
interested in the mean [[f(z)] and variance ([f(=) of trained models over random initialization. The former is typically used for the prediction of deep ensemble, while the latter is used for estimating model epistemic uncertainty utilized c.g. for OOD detection exploration. To start, we describe the simpler case of the infinite width limit and a deterministic NTK, which allows us compute the mean and variance of the solutions found by training easily:
El (x, x)21
V[So -K(,, + Q ()0,
(x,
where we introduced (I, x)) (x,
For finite width linearly trained networks, the kernel is no longer deterministic, and its stochasticity influences the predictive distribution. Because there is probability mass assigned to the neighborhood of rare events where the NTK kernel matrix is not invertible, the expectation and variance over parameter initialization of the expression in equation diverges to infinity.
Fortunately, due to the convergence in probability of the empirical NTK to the infinite width counter- part (24). we know these singularities become rarer and ultimately vanish as the width increases to infinity. Intuitively, should therefore be able to assign meaningful, finite values to these undefined
3

---

quantities, which ignores these rare singularities. The delta method [37 in statistics formalizes this intuition, by using Taylor approximation smooth out the singularities before computing the mean or variance. When the probability mass of the empirical NTK is highly concentrated in small radius around the limiting NTK, the expression is roughly linear w.r. the NTK entries. Given this observation, we prove (see Appendix ( the following result, and justify that the obtained expression is informative of the empirical predictive mean and variance of deep ensembles. Rewriting equation 4 into
(x)) = -1,,
+ 1 (x. X) (ax, (())
where -(x,
and θ )( we state:
Proposition 2.1. For one hidden layer networks parametrized as in equation 7 given an input x and training data (X, when increasing the hidden layer width we have the following convergence in distribution over random initialization Og:
VM (a, X)
((2, dist. Z(x)
where Z(x) is the linear combination of Chi-Square distributions, such that V(Z(x)) = lim (x)
1-000
where
-V VTtp:
- -cg. (x() 2Co -2000p/0...000g...00000011111010
/ttp://00000...11100000000000000 0o)
We omit the dependence of 0 on the width h for notational simplicity. While the expectation or variance equation 4 for any finite width is undefined, their empirical mean and variance are with high probability indistinguishable from that of the above limiting distribution (see Lemma (..). Note that the above proposition assumes the noise in Oo to be decorrelated from (x, 0o), which can hold true under specific constructions of the network that are of practical interest as we will see in the following fc.f. Appendix A.3.2p.
Given Proposition 2.1 we now describe the approximate variance of flin (x) for L = which we extend to the general I. case using an informal argument (see A.2.2):
Proposition 2.2. Let f be neural network with identical width of all hidden layers, h1 = h2 = h11-11 h. We assume ||0,, - @1 Op())) Then,
(x)x
V'(x) + yeor (2or())))))))))))))))))(( + (Te)
where
= 00,,
your (x) 2E (x, (1))) (X, 210 (x,,,,)
xK(X)x) - =(X,
and yrea (x) O(h as well as K the expectation over initializations of the finite width counterpart of the NNGP kernel.
Several observations can be made: First, the above expression only involves the first and second moments of the empirical, finite width NTK, as well as the first moment of the NNGP kernel. These terms can be analytically computed in some settings. We provide in Appendix A.4.3 some of the moments for the special case of 1-hidden layer ReLU network, and show the analytical expression correspond to empirical findings.
Second, the decomposition demonstrates the interplay of distinct noise sources the predictive variance:

---

02-
04-
0.0
0.054
0.000
0.4
0.0
0.2
10084
0:386
0.040
0.108
04
0.0
0.2
00016
1.018
0.821 0.-
0.004
04
0.0
-0,001
0.002e
0.0002
0.006
0.0000
0.0000
04
0.000
0,4
0.000
044
0.000
04
0.0000
-0.25
0.00
0.25
-0.25
0.00
0.25
-0.25
0.00
0.25
-0.25
0.00 0.25
Figure 1: Empirical variance (=f of our kernel models flin, flin-a flin-c and flin-i on toy
regression problem, consisting regressing against a value of on the purple and -1 on the green datapoints. Left: [x)] as superposition of the three isolated components right of Center left: V[ [[ (z:] which correlates with the distance to the datapoints. Center right: V ((llite)ar)) (ar) which correlates with the angle to the datapoints. Right: V (xil)))))))) which depends on distance and angle to the datapoints.
ve is the variance associated to the expression in the first line of equation 7 the finite width counterpart of the predictive variance of the infinite width model (equation 6 as it assumes the NTK is deterministic. The variance stems entirely from the functional noise at initialization and converges to the infinite width predictive variance as the width increases.
Intuitively, it is
ye and yi stem from the second line of equation 7 v is first-order approximation of the predictive variance linearly trained network with pure kernel noise, without functional noise i.e. V xe V[Qo, (x, X))] On the other hand, ' depends the 2 noises, and can be identified as the predictive variance of a deep ensemble with a deterministic NTK θ and new functional prior ((x) = to, (z. (2' x) fxX, 6u). Intuitively, this new functional prior can be seen as a sata-specific inductive bias on the NTK formulation of the predictive variance (see Appendix A.3. for more details).
the interplay between
year is covariance term between the terms in equation 7 and also contains the correlation terms between Og, and f(z, 0o). In general, its analytical expression is challenging to obtain it requires the 4th moments of the finite width NNGP kernel fluctuation. Here, we provide its expression under the same simplifying assumption that the noise in O,, is decorrelated from f(x, fb). We therefore do not attempt to describe it in general, and focus in our empirical Section on the terms that are tractable and can be easily isolated for practical purposes.
Each of yc, V year decay in O(h-1), which, together with Va, provide a first-order approximation of the predictive variance of ((im(x)) Note that Va and Ve are of particular interest, as removing either the kernel or the functional noise at initialization will collapse the predictive variance of the trained ensemble to either one these terms.
2.2 Predictive distribution of standard deep ensemble of large width
An important question at this point is to which extent our amalysis for linearly trained models applies fully and non-linearly trained deep ensemble. Indeed, if the discrepancy between the predictive variance of a linearly trained ensemble and its non-linear counterpart is of a larger order of magnitude than the higher-order correction in the variance term. the latter can be 'erased" by training. Building on top of previous work, we show that, under the assumption of an empirically supported conjecture 38, , for one hidden layer networks trained on the Mean Squared Error (MSE) loss, this discrepancy is asymptotically dominated by the refined predictive variance terms of the linearly trained ensemble we described in Section
Proposition 2.3. Let f be neural network with identical width of all hidden layers, h h2 = hr,-1 = h, and such that the derivative of the non-linearity o' is bounded and Lipschitz continuous on RR. Let the training data (X, Y) contained in some compact set, such that the NTK of J on X is invertible. Let Je (resp. flin) be the model (resp. linearized model) trained on the MSE loss with gradient flow timestep with some learning rate. Assuming
Sup - 1F =
(8)
5

---

1/A
1/h
V 1,k
7 1/12
Figure 2: Empirical predictive variance W, VO V Ve of an ensemble of 100 models. Upper Left: Linearized feed forward networks various depths and widths trained on subset of MNIST. Upper Right: Linearized convolutional neural networks of various depth and widths trained on subset of CIFARIO. Lower Left (resp. Right): 1/h² scaling of yrea of an ensemble of MLPs (resp. CNN) of various depths and widths trained on subset of MNIST (resp. CIFARIO). Although the theoretical result on the scaling of 1/h of the variance terms influenced by the kernel noise as well as 1/h² of the residual holds only for depth
(line plots), the same scaling is observed for deeper networks as suggested by our informal result (L = 3 in dashed lines, in dashed-dotted lines). All plots plotted in the log-log scale.
Then, Vx, 48 0, 3C, H 4h H.
C 1 2
P[supllSt [x(1))12 < h
(9)
In particular, for one hidden layer networks, after training,
(10)
where V denotes the empirical variance with some fixed sample size.
The proof can be found in Appendix A.1.1 While only the bound @o Vi
has
been proven in previous works 255 many empirical studies including those in the present work (see Appendix Fig.
5) Table 3 have shown that the bound decreases faster practice, on the order of
O(h-') 25 38 Note that this result suggests the approximation provided in Proposition 2.2 is
good as gets for describing the predictive variance of non-linearly trained ensembles: the higher order terms would be of a smaller order of magnitude than the non-linear correction to the training, rendering any finer approximation pointless.
3 Disentangling deep ensemble variance in practice
The goal of this Section is to validate our theoretical findings in experiments. First, we aim to show qualitatively and quantitatively that the variance of linearly trained neural networks well approximated by the decomposition introduced in Proposition 2.2 To do so, we investigate ensembles of linearly trained models and analyze their behavior toy models and common computer vision classification datasets. We then extend our analyses to fully-trained non-linear deep neural networks optimized with (stochastic) gradient descent in parameter space. Here, we confirm empirically the strong influence of the variance description of linearly trained models in these less restrictive settings while being trained to very low training loss. Therefore we showcase the improved understanding of deep ensembles through their linearly trained counterpart and highlight the practical relevance of our study by observing significant OOD detection performance differences models when removing noise sources in various settings.

---

3.1 Disentangling noise sources in kernel models
To isolate the different terms in Proposition 2.2 we construct, from a given initialization 8g with the associated linearized model flin, three additional linearly trained models:
f (x) Q0, (z,
flima (x) = 00) + -0xx
(x) = ((;; + 2(x,2)) g(X 00))
where g(r, 0o) = On ( xXx 0o). Note that the predictive variance over random initialization of these functions corresponds to respectively Ve, Va, V as defined in Section 2.1
As one can see, we can simply remove the initialization noise from flin by subtracting the initial (noisy) function f(x,8u, do) before training resulting in a centered model flim Equivalently, we can remove noise that originates from the kernel by using the empirical average over kernels resulting in model flina Finally, we can isolate flin-i by the same averaging trick as in f but use as functional noise g 0b) which can be precomputed and added to flinee before training. Note that neglect the terms involving covariance terms and focus on the parts which are easy to isolate, for linearly trained as well as for standard models. This will later allow us to study practical ways to subtract important parts of the predictive distribution for neural networks leading for example to significant OOD detection performance differences. Now we explore the differences and similarities of these disentangled functions and their respective predictive distributions.
3.1.1 Visualizations star-shaped toy dataset
To qualitatively visualize the different terms, we construct two-way star-shaped regression problem on a 2d-plane depicted in Figure [ After training an ensemble we visualize its predictive variance on the input space. Our first goal is to visualize qualitative differences in the predictive variance of ensembles consisting of flin and the 3 disentangled models from above. We train a large ensemble of size 300 where each model is a one-layer ReLU neural network with hidden dimension 512 and I hidden layer. As suggested analytically for one hidden layer ReLU networks (see Appendix A.4.3). for example [[ff)((()]] depends on the angle of the datapoints while V [=]" depicts superposition of the 3 isolated variances. While the ReLU activation does not satisfy the Lipschitz- continuity assumption of Proposition 2.33] we use it to illustrate and validate our analytical description of the inductive biases induced by the different variance terms. We use the Softplus activation which behaved similarly to ReLU in the experiments in the next Section.
3.1.2 Disentangling linearly trained / kernel ensembles for MNIST and CIFAR10
Next, we move to a quantitative analysis of the asymptotic behavior of the various variance terms, as we increase the hidden layer size. In Figure 2 we analyze the predictive variance of the kernel models based on MLPs and Convolutional Neural Networks (CNN) for various depths and widths and on subsets of MNIST [39 and CIFARIO. As before, we construct binary classification task through a MSE loss with dataset size of N = 100 and confirm, shown in Figure 2 that VS, V decay by 1/h over all of our experiments. Crucially, we see that they contribute to the overall variance V even for relatively large widths. We further observe a decay in 1/h² of the residual term as predicted by Proposition 2.2] As in all of our experiments, the variance magnitude and therefore the influence on fii of the disentangled parts is highly architecture and dataset-dependent. Note that the small size of the datasets comes from the necessity to compute the inverse of the kernels for every ensemble member, see Appendix B for a additional analysis on larger datasets and scaling plots of yeoor
In Table M. we quantify the previously observed qualitative difference of the various predictive variances by evaluating their performance on ofd--is-iisiiiiionn detection tasks, where high predictive variance is used proxy for detecting duittf-distribution data. We focus our attention on analysing V[flin-(z)] and V[flii-*==)) as they are the variance terms containing purely the functional and kernel noise, respectively. As an evaluation metric, we follow numerous studies and compute the area under the receiver operating characteristics curve (AUROC, c.f. Appendix BB. We fit linearized ensemble on a larger subset of the standard 10-way classification MNIST and CIFAR 1 datasets using MSE loss. When training our ensembles on MNIST, we test and average the OOD detection performance on MNSH (FM) 40), E-MNIST (EM) [1 and K-MNIST (KM) 42 When training our ensembles on CIFARIO 10, compute the AUROC for SVHN 43
LSUN
44 TinyImageNet (TIN)
7

---

Table 1: Test set accuracy and AUROC for deep ensembles of size 10 of MLPs (h = 1024, - 3) and CNNs (h 256, L = 3) trained on a subset (N=1000) and on of full MNIST and CIFARIO. We indicate small standard deviations a < 0.005 obtained over ensemble size (E) with +00 In all experiments, the various disentangled models show significant differences behavior. All linearly trained models follow the gradient descent models behavior tightly. When optimizing with SGD, isolating initial noise sources still affect the ensemble behavior significantly and can lead to improved OOD detection as well as test set accuracy.
Model CNN, CIFARIO, N=1000.
E=10
MLP, MNIST, N=1000, E=30
Test (%)
SVHN
LSUN
iSUN
Test (%)
FM
EM
KM
flnn
36,432-90
-332#-006
8092.0004
.7832-000
91.53* 40
.962=-00 006
-922#0000
.982#-001
plue
37.2#.44
-677-000
-693..001
.674#.000
89.67 15
.935= 005
-881--003
-967#-002
fluna
30.90
-10+-006
.7642-003
.738--000
91.27# 06
-978=000 003
-922--601
.987#-000
flami
32.85=t.21
.591 000
.683
.660=t.000
91.60
.970=-000 004
-908--001
.983#-002
yed
39.70
516
0002
-789
.763=t-004
91.43
.971 .971
-244t-001
.986-t0001
fede
37.4
562
0004
.691 0000
.670st.002
89.67 06
.937
.884+.003
.968t..002
feda feds
30.53=1.15
.509
000
.734
90.73
.978 .978 0003
.922#.000
-877
3t.
.83 000
.656-t.00s 0000
.656
.638 .003
90.65 4. 35
.977 + 0003
.913* 913+..002
+987
Model
CNN, CIFARIO, N=50000, E=5
MLP, MNIST, N=50000, E=5
Test (%)
SVHN
LSUN
iSUN
Test (%)
FM
EM
KM
fedl
62.68=t-6
.557=-01
884#-00
864#-00
95.70=-12
.974#0005
-330
-99
frehe froda
57.03#.114
.554=-00
.791=-00
.781 #.00
94.43#-01
.224--016
-833.0006
-962#-004
58,83
.455 00
-844=00
-85
97.48=-13
.988=-002 0002
.943#.001
.995t-001
and CIFARIO0 (C100), see Appendix Table 4 for the variance magnitude and AUROC values for all datasets.
The results show significant differences in variance magnitude and AUROC values. While we do not claim competitive OOD performance, we aim to highlight the differences in behavior of the isolated functions developed above: we see for instance that for (MLP, MNIST, N=1000), flin-a generally performs better than flin OOD detection. Indeed, the overall worse performance of V[rlin- (x)) seems to be affecting that of ([ (z)] which contains both terms. On the other hand, we see that for the setup (CNN, CIFARIO, N=1000) (x)] is not well described by this interpolation argument, which highlights the influence of the other variance terms described in Proposition 2.2 Furthermore, the OOD detection capabilities of each function seem to be highly dependent on the particular data considered: Ensembles of flin-e are relatively good at identifying SVHN data OOD, while being poor at identifying LSUN and iSUN data. These observations highlight the particular inductive bias of each variance term for OOD detection different datasets.
We further report the test set generalization of the ensemble mean of different functions, highlighting the diversity in the predictive mean of these models as well. Note that for N >= 1000 we trained the ensembles in linear fashion with gradient flow (which coincides with the kernel expression) up until the MSE training error was smaller than 0.01.
3.2 Does the refined variance description generalize to standard gradient descent in practice?
In this Section, we start with empirical verification of Proposition 2.3 and show that the bound in equation 10] holds in practice. Given this verification, we then propose equivalent disentangled models as those previously defined but in the non-lincar setting, and 1) show significant differences in their predictive distribution but also 2) investigate which extent improvements in OOD detection translate from kernel / linearly to fully non-linearly trained models. We stress that we do not consider early stopped models and aim to connect the kernel with the gradient descent models faithfully.
3.2.1 Survival of the kernel noise after training
To validate Proposition 2.3 we first introduce f=)(() = f(x,0,), a model trained with standard gradient descent of steps i.e. Oe- = 0, - - Vof(X, f(X, 0.)). To empirically verify
8

---

Proposition 2.3
we introduce the following ratio
| f (x)) - -
R(f) = exp (Eew.
logl
(11)
[[V(x)) + (x(x)]
where the empirical variances are computed over random initialization, and the expectation over some data distribution which we choose to be the union of the test-set and the various OOD datasets. Given a datapoint 2;, the term inside the log measures the ratio between the discrepancy of the variance between the lincarized and non- non-linear ensemble, against the refined variance terms. R(f) then the geometric mean of this ratio over the whole dataset. Proposition 2.3 predicts that the ratio remains bounded the width increases, suggesting that the refined terms contribute to the final predictive variance of the non-linear model in a non negligible manner. We empirically verify this prediction for various depths in Fig. and Appendix Figure for functions trained subsets MNIST and CIFARIO. Note that for all our experiments we also empirically verify the assumption from Proposition 2.3 (see Appendix Figure 5 Table
3.2.2 Disentangling noise sources in gradient descent non-linear models
Motivated by the empirical verification of Propo- sition 2.3] we now aim isolate different noise sources in non-linear models trained with gra- dient descent. Starting from non-linear net- work fee, we follow the same strategy as before and silence the functional initialization noise by centering the network (referred as fal)(z)) by simply subtracting the function initialization. On the other hand, we remove the kernel noise with a simple trick: We first sample random weight 8i once, and use it as the weight initial- ization for all ensemble members. A function noise is added by first removing the function initialization from 0% and adding that of a sec- ond random network which is not trained. The resulting functions (referred as (z)) will induce and ensemble which will only differ in their functional initialization while having the same Jacobian
1. L=
L=4
S
2
1024
2048
3072
h
Figure 3: RJ)) of MLPs with multiple widths and depths (L E (2,3, 4}) trained on a subset (100) of MNIST As predicted, we observe R(f) bounded as we increase the width in support of our theoretical analysis.
yed-e (x) f(z 02) f(z,00),
yada (x) f(z 01) - f(z, 05) + f(x, 0o).
We furthermore introduce fad-1()) the non linear counterpart to flin=((x:) which we construct similarly to fzd)(()) but using g(z: 0, 06) O,, (x,,, (X',
(X', 0o) as the function initialization instead of fxx 0b) (see Section 2.1 and Appendix AA.31 for the justification). Unlike fada and fal:, constructing fauli requires the inversion of large matrices due the way is defined, challenging task for realistic settings. While its practical use thus limited, we introduce it to illustrate the correspondence of correspondence of the predictive variance of linearized vs non-linear deep ensemble.
Given these simple modifications of fad, we rerun the experiments conducted for the linearly trained models for moderate dataset sizes (N=1000). We observe close similarities the OOD detection capabilities as well as predictive variance between the introduced non-linearly trained ensembles and their linearly trained counterparts. We further train these models on the full MNIST dataset (N=50000) for which we show the same trend in Appendix Table We also include the ensemble' performance when trained on the full CIFAR dataset. Intriguingly, the relative performance of the ensemble is somewhat preserved in both settings between N=1000 and N=50000, even when training with SGD, promoting the use of quick, linear training on subset of data as a proxy for the OOD performance of fully trained deep ensemble.
Similar to the case of (MLP, MNIST, 1=1000/50000), we observe that fad ensemble performance is an interpolation of fad- and fad-a which interestingly performs often favorably, on different OOD data. To understand if the noise introduced by SGD impacts the predictive distribution of our disentangled ensembles, compared the behavior of fed and fnd in the lower data regime

---

of N = 1000. Intriguingly. we show in Appendix Table thtt no significant empirical difference between GD and SGD models can be observed and hypothesize that noise sources discussed in this study are more important in our approximately linear training regimes. To speed up experiments we used (S)GD with momentum (0.9) in all experiments of this subsection.
3.2.3 Removing noise of models possibly far away from the linear regime
Finally, we investigate the OOD performance of the previously introduced model variants f f'sdd- and figddd in more realistic settings To do so we train the commonly used WideResNet 28-10 [45] on CIFARIO with BatchNorm [46] Layers and cross-entropy (CE) loss with batchsize of 128, without data augmentation (see Table 3.2.3). These network and training algorithm choices are considered crucial to achieving sfate-of-the-arr and superior performance compared to their linearly trained counterparts. Strikingly. we notice that our model variants, which each isolate a different initial noise source, significantly affect the OOD capabilities of the final models when the training loss is virtually 0 as in all of our experiments. This indicates that the discussed noise sources influence the ensemble's 's predictive variance long throughout training. We provide similar results for CIFARI00 and FashionMNIST in Table B. and [. off of the Appendix B. We stress that we do not claim that our theoretical assumptions hold in this setup.
Table 2: Test set accuracy and AUROC for WRN 28- 10 ensembles of size 8 trained on CIFARIO on the cross entropy (CE) or MSE loss. While the models perform similarly on test set, a significant advantage of f in OOD detection is observed across most OOD datasets. Standard deviations computed over 5 seeds are indicated with +. In bold are values that outperform l with p < 0.2.
Model
Loss
Test (%)
C100
SVHN
LSUN
TIN
iSUN
fred
CE
0.900+t0.002 0.830 89.36+0.36
0.891
0.891#0.002 00.002
0.883+90000 00000
figde
CE
89.56
t9.56
0.831 +00003
0.8 0.899+00004
0.895 000003
0.862 +0.000
0,885 +00003
greda
CE
89.01 #0.32
0.827 +00000
0.894 000008
0.887#0.008
0.855=00.000
0.879++0.002
fied
MSE
77.94 20
0.739 +00004
0.863 0.863+0.0000 +0.00
0.823
0.822300.007 +0.007
40.95
0.813 +..00s
figie
MSE
77.88+0 +0.30
0.739 +00001
0.880+0.000
0.829 0.0
0.807 +00008
0.813 400.00
MSE
75.34 #0.18
0.707 0.0004
0.841 #:0011
0.784
0.0003
0.763 0.763 4:0.000
0.761 0:0,00
4 Conclusion
The generalization on in-and out-of-distribution data of deep neural network ensembles is poorly understood. This is particularly worrying since deep ensembles are widely used in practice when trying to asses if data is out-of-distribution. In this study, we try provide insights into the sources of noise stemming from initialization that influence the predictive distribution of trained deep ensembles. By focusing on large-width models we are able to characterize two distinct sources of noise and describe an analytical approximation of the predictive variance in some restricted settings. We then show theoretically and empirically how parts of this refined predictive variance description in the linear training regime survive and impact the predictive distribution of non-linearly trained deep ensembles. This allows us to extrapolate insights of the tractable linearly trained deep ensembles into the non-linear regime which can lead to improved by eliminating potentially unfavorable noise sources. Although our theoretical analysis relies on the closeness to linear gradient descent which has shown to result in less powerful models in practice, we hope that our surprising empirical success of noise disentanglement sparks further research into using the lens of linear gradient descent to understand the mysteries of deep learning.
detection of deep ensembles
Acknowledgments and Disclosure of Funding
Seijin Kobayashi was supported by the Swiss National Science Foundation (SNF) grant CR- $115_173721. Pau Vilimelis Aceituno was supported by the ETH Postdoctoral Fellowship program (007113). Johannes von Oswald was funded by the Swiss Data Science Center (J.v.O. P18-03). We thank Christian Henning, Frederik Benzing and Yassir Akram for helpful discussions. Seijin Kobayashi and Johannes von Oswald are grateful for Angelika Steger's 's and João Sacramento's 's overall support and guidance.
10

---

References
[1] Christian Leibig, Vaneeda Allken, Murat Seçkin Ayhan, Philipp Berens, and Siegfried Wahl. Leveraging uncertainty information from deep neural networks for disease detection. bioRxiv, 2017.
[2] Ashwin Carvalho. Predictive control under uncertainty for safe autonomous driving: Integrating data-driven forecasts with control design. In PhD Thesis, 2016.
[3] Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, D Sculley, Sebastian Nowozin, Joshua Dillon, Balaji Lakshminarayanan, and Jasper Snoek. Can you trust your model's uncertainty? Evaluating predictive uncertainty under dataset shift. In Advances in Neural Information Processing Systems 33, 2019.
[4] Shiyu Liang, Yixuan Li, and R. Srikant. Enhancing the reliability of out-of-distribution image detection in neural networks. In International Conference on Learning Representations, 2018.
[5] Eric Nalisnick, Akihiro Matsukawa, Yee Whye Teh, Dilan Gorur, and Balaji Lakshminarayanan. Do deep generative models know what they don't know? In International Conference on Learning Representations, 2019.
[6] Dan Hendrycks and Kevin Gimpel. A baseline for detecting misclassified and outtt--disttibuttin examples neural networks. International Conference on Learning Representations, 2017.
[7] Kimin Lee, Kibok Lee, Honglak Lee, and Jinwoo Shin. A simple unified framework for detecting out-of-distribution samples and adversarial attacks. In Advances in Neural Information Processing Systems 31. 2018.
[8] Jie Ren, Peter J. Liu, Emily Fertig, Jasper Snoek, Ryan Poplin, Mark Depristo, Joshua Dillon, and Balaji Lakshminarayanan. Likelihood ratios for distriiiiiiittttinn detection. In Advances in Neural Information Processing Systems 32. 2019.
[9] Johannes von Oswald, Seijin Kobayashi, Alexander Meulemans, Christian Henning, Benjamin Grewe, and João Sacramento. Neural networks with late-phase weights. In International Conference Learning Representations, 2021.
[10] Yeming Wen, Dustin Tran, and Jimmy Ba. BatchEnsemble: an alternative approach to efficient ensemble and lifelong learning. In International Conference on Learning Representations, 2020.
[11] David J. C. MacKay, A practical Bayesian framework for backpropagation networks. Neural Computation, 4(3):448 472, 1992.
[12] Max Welling and Yee Whye Teh. Bayesian learning via stochastic gradient Langevin dynamics. In International Conference on Machine Learning, 2011.
[13] Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference deep generative models. In Proceedings of the 31st International Conference on Machine Learning. PMLR, 2014.
[14] Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In Proceedings of the 32nd International Conference on Machine Learning. PMLR, 2015.
[15] Yarin Gal and Zoubin Ghahramani. Dropout bayesian approximation: Representing model uncertainty in deep learning. In Proceedings of The 33rd International Conference on Machine Learning. PMLR, 2016.
[16] Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In Advances in Neural Information Processing Systems 30. 2017.
[17] Andrew Gordon Wilson and Pavel Izmailov. Bayesian deep learning and a probabilistic perspective of generalization. arXiv, 2020.

---

[18] Bobby He, Balaji Lakshminarayanan, and Yee Whye Teh. Bayesian deep ensembles via the neural tangent kernel. In Advances Neural Information Processing Systems 34. 2020.
[19] Francesco D' Angelo and Vincent Fortuin. Repulsive deep ensembles are bayesian. In Advances in Neural Information Processing Systems, 2021.
[20] Seijin Kobayashi, Johannes Oswald, and Benjamin F Grewe. On the reversed bias- variance tradeoff in deep ensembles. ICML Workshop on Uncertainty and Robustness in Deep Learning, 2021.
[21] Frederik Benzing, Simon Schug, Robert Meier, Johannes von Oswald, Yassir Akram, Nicolas Zucchet, Laurence Aitchison, and Angelika Steger. Random initialisations performing above chance and how to find them. arXiv, 2022.
[22] Francesco D' Angelo and Christian Henning. Uncertainty out-of-distributiin detection requires suitable function space priors. arXiv, 2021.
[23] Jachoon Lee. Jascha
Jeffrey Pennington, Roman Novak, Sam Schoenholz, and Yasaman Bahri. Deep neural networks as gaussian processes. In International Conference on Learning Representations, 2018
[24] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in Neural Information Processing Systems 32. 2018.
[25] Jaehoon Lee, Lechao Xiao, Samuel S Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl-Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. Journal of Statistical Mechanics: Theory and Experiment, 2020.
[26] Lenaic Chizat, Edouard Oyallon, and Francis Bach, On lazy training in differentiable program- ming. In Advances in Neural Information Processing Systems 33. 2019.
[27] Sanjeev Arora, Simon S. Du, Wei Hu, Zhiyuan Li, Ruslan Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. In Advances in Neural Information Processing Systems 33, 2019.
[28] Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M. Roy, and Surya Ganguli. Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel. In Advances Neural Information Processing Systems, 2020.
[29] Jachoon Lee, Samuel Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite versus infinite neural networks: an empirical study. In Advances in Neural Information Processing Systems, 2020.
[30] Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. When do neural networks outperform kernel methods? In Advances in Neural Information Processing Systems,
2020.
[31] Dingli Yu, Ruosong Wang, Zhiyuan Li, Wei Hu, Ruslan Salakhutdinov, Sanjeev Arora, and Simon S. Du. Enhanced convolutional neural tangent kernels. arXiv, 2020.
[32] Mario Geiger, Stefano Spigler, Arthur Jacot, and Matthieu Wyart. Disentangling feature and lazy training in deep neural networks. Journal of Statistical Mechanics: Theory and Experiment, 2020.
[33] Richard Hans Robert Hahnloser, Rahul Sarpeshkar, Misha A. Mahowald, Rodney J. Douglas, and H. Sebastian Seung. Digital selection and analogue amplification coexist in a cortex-inspired silicon eircuit. Nature, 2000.
[34] Radford M. Neal. Bayesian Learning for Neural Networks. Springer Verlag, Berlin, Heidelberg, 1996.
[35] Greg Yang. Tensor programs II: Neural tangent kernel for any architecture. arXiv, 2020.
12

---

[36] Greg Yang and Etai Littwin. Tensor programs Ilb: Architectural universality of neural tangent kernel training dynamics. In Proceedings of the 38th International Conference on Machine Learning. PMLR, 2021.
[37] J. L. Doob. The Limiting Distributions of Certain Statistics. The Annals of Mathematical Statistics, 6:160 169, 1935.
[38] Mario Geiger, Arthur Jacot, Stefano Spigler, Franck Gabriel, Levent Sagun, Stéphane Ascoli, Giulio Biroli, Clément Hongler, and Matthieu Wyart. Scaling description of generalization with number of parameters in deep learning. Journal of Statistical Mechanics: Theory and Experiment, 2020.
[39] Yann LeCun, Corinna Cortes, and CJ Burges. MNIST handwritten digit database. ATT Labs. Available: http://yann.eccmm www../....o/ed///mniiiiii 2010.
[40] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms, 2017.
[41] Gregory Cohen, Saeed Afshar, Jonathan Tapson, and André van Schaik. EMNIST: Extending MNIST handwritten letters. arXiv, 2017.
[42] Tarin Clanuwat, Mikel Bober-Irizar, Asanobu Kitamoto, Alex Lamb, Kazuaki Yamamoto, and David Ha. Deep learning for classical japanese literature. arXiv, 2018.
[43] Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y. Ng. Reading digits in natural images with unsupervised feature learning. In NIPS Workshop on Deep Learning and Unsupervised Feature Learning, 2011.
[44] Fisher Yu, Yinda Zhang, Shuran Song, Ari Seff, and Jianxiong Xiao. LSUN: construction of large-scale image dataset using deep learning with humans in the loop. arXiv, 2015.
[45] Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. In Proceedings of the British Machine Vision Conference, 2016.
[46] Sergey loffe and Christian Szegedy. Batch normalization: accelerating deep network training by reducing internal covariate shift. In International Conference on Machine Learning, 2015.
13

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? Yes, our theoretical heavily relies on assumptions that have be met in practice. Arguably these assumptions are not met in
standard deep learning setups. [Yes]
(c) Did you discuss any potential negative societal impacts of your work? [Yes] Our study aims developing method and theory of improving out-of-disribution of deep neural networks. While advances in understanding deep ensembles and better OOD detection can improve AI systems in general and can have a lasting socictal impact, we do not anticipate any immediate direct, especially negative, societal impact originating from our work.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes] The full set of assumptions of all theoretical results is provided in the Appendix.
(b) Did you include complete proofs of all thcoretical results? [Yes] All proof's for our theoretical results are included in the Appendix.
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main exper- imental results (either in the supplemental material as a URL)? [Yes] We include code and instructions to reproduce main results as part of our submission and will include an URL linking to the code upon publication.
(b) Did you specify all the training details (e.g., data splits, byperparameters, how they were chosen)? [Yes] Complete list of hyperparameters provided in the SM.
(c) Did you report error bars (e.g., with respect the random seed after running experi- ments multiple times)? [Yes]
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, or cloud provider)? [Yes] Information about the used hardware can be found in the Appendix.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either in the supplemental material or as a URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you deseribe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
14


Reviews:
[{"review_id": "tF11q4vAphu", "reviewer": "Reviewer_CZjD", "summary": "This paper studies ensembles of linearly trained deep networks, for which a decomposition of the variance is proposed into two main components:\n - V_a results from the variance of the initial value f_theta_0\n - V_c results from the variance of the initial kernel Theta_0\nand other components that decay faster as the width of the network is grown.\n\nIt is also shown that this decomposition still hold when training non-linearly one hidden layer networks under stability assumption of the NTK during training. It is conjectured (and empirically somehow observed) that the difference in prediction between linear and non-linear training is small enough for the above decomposition to still be relevant in deeper networks.\n\nExperiments on toy datasets enable visualizing the different sources of variance. Experiments on more difficult (MNIST, CIFAR10) datasets show that conclusions vary across setups. A method is proposed to isolate sources of variance in actual setups.", "questions": " - it is not clear that the y scale in figure 2 and 5 is a log scale from just looking at the figure (and not mentioned in the text).\n - in figure 2, contrary to what is mentioned in the text, it looks that the decay of V_c is slightly slower than 1/h (thus suggesting that V_c does not decay in O(1/h)). It is difficult to see because of the log-log scale but it looks to be closer to 1/h^{3/4}\n - in figure 5 in appendix, I did not understand how many steps of GD were run for this experiment. I am guessing that it uses very few steps since the dataset if very small (N=100), and I am expecting the results to be different on actual setups. What do you think?", "limitations": "I did not find a discussion about potential negative societal impact, nor do I need that such paper should have one.\n\nThe limitations of the setups (linearly trained networks, small datasets) are discussed.", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "*Originality*: To the best of my knowledge (though I am not so familiar with recent literature), the idea of studying ensemble of ensembles of linearly trained networks is new and sound. The results linking infinite-width models to linearly-trained finite networks are new, in the line of works around the NTK formalism.\n\n*Quality*: The paper is well organised, the experiments and figures are relevant to the discussion. I had some comments about the experiments (see questions below)\n\n*Clarity*: The material is well presented. The theoretical results are clearly discussed, as well as the assumptions. I however was not sure what is the main message to get from the experiments on large networks/datasets. This is probably related to the latter, but I did not understand what are the \"practical ways to eliminate noisy sources\" claimed to be proposed in the abstract.\n\n*Significance:*: The study of ensembles of neural nets is of particular relevance as many recent papers use ensembles of deep nets. The decomposition between variance of the initial function and variance of the initial kernel is of interest for understanding the training dynamics of neural networks.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "kLLB4KTIfJN", "reviewer": "Reviewer_e4Rz", "summary": "This paper proposes to understand finite-width neural networks in terms of linearly trained finite width networks. The authors propose to decompose the variance of ensembles of linearly trained, finite width neural networks into different sources (of noise), and study empirically the implications of this decomposition for understanding actual finite-width networks trained with gradient descent and stochastic gradient descent. The work builds from [1][2] which proposes to approximate a sufficiently wide network $f_\\theta$ with its linearly trained counterpart $f_lin$, and in the infinite with limit with the fully trained network $f_\\inf$.\nThe paper decompose the variance of $f_\\inf$  and $f_lin$, and interprets different variance terms arising in ensembles of linear predictors in terms of kernel and functional noise (ln 140), which affects the predictive variance. The derivation relies on certain independence assumptions that do not hold in practice.\nThe paper additionally claims ways to eliminate sources of noise which improve OOD detection in ensembles. The results comparing the networks which only include disentangled noise sources show no consensus in how/if the proposed sources of noise affect the total variance. The claims of improving OOD detection after reducing sources of noise coming from the results don’t take into account that the results lie within the confidence interval (cf. Table 2).", "questions": "- what is it the implication of assuming the functional initialization noise is decorrelated from the $f(x)$ in proposition 2.1?\n- In line 140, why is it obvious that Va corresponds to the inf width model?\n- What are the implications of Fig 1, if we cannot draw any conclusions from the shape/location of the data points, as the authors mentioned this changes per model/dataset?\n- Could we add a Fig 2 which includes the baseline, total V? this would be useful to better understand the scale of different V. \n", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\nThis paper proposes to understand finite-width neural networks in terms of linearly trained finite width networks. The paper decompose the variance of ensembles of linearly trained, finite width neural networks into different sources (of noise). The work builds from  [1][2] which proposes to approximate a sufficiently wide network $f_\\theta$ with its linearly trained counterpart $f_lin$, and in the infinite with limit with the fully trained network $f_\\inf$. This is an interesting perspective for analyzing deep ensemblesl\n\nWeaknesses:\n- The paper decompose the variance of ensembles of $f_\\inf$  and $f_lin$.  The paper claims the variance for the ensemble of $f_lin$  includes two components which correspond to the functional initialization noise and the other to the kernel noise (ln 140), and that this noise affects the predictive variance. During the derivation it is assumed that the that the functional initialization noise is decorrelated from the $f(x)$ (ln 133). The authors mentioned that this is future work but this is a central assumption in this work.\n- The results comparing the networks which only include disentangled noise sources show no consensus in how/if the proposed sources of noise affect the total variance. Could we add a Fig 2 which includes the baseline, total V?\n- The claim of improving OOD detection after reducing sources of noise coming from the results in Table 2. In Table 2, we can see that the comparisons lie within each other's confidence intervals, so it is unclear why we can claim a model is better than another. \n- Some sections need additional clarification, i.e. the description of the variance components in paragraph ln 147 to 160 is not clear to me. Given that this paper is offering a new perspective on the variance decomposition a clear description of each component would be greatly beneficial. I had to come back to this twice. The same holds for the implications of the results in Fig 1. \n- The assumptions and implications of the propositions, in particular 2.1 and 2.3, could be clearer.  Proposition 2.3 can be true for a single network, is this correct? This proposition is in the deep ensemble section 2.2.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "aBIBHx46NQ", "reviewer": "Reviewer_rqPL", "summary": "This work is focused on decomposing the sources of variance in neural network predictions - the predictive variance of neural networks is an important quantity that allows for neural ensembles (and other approximate bayesian models) to detect OOD samples without seeing OOD data. However, this quantity is poorly understood - this work proposes to take a close, principled look at neural network predictive variance by studying _linear_ neural networks. Linear neural network predictions can be represented via a kernel expression, allowing a theoretical decoupling of model predictions and (noisy) gradient descent.\n\nThe work largely comprises of 2 main sections. The first section (section 2) contains the main theoretical results of the work, and begins by introducing linear neural networks and infinitely wide neural networks. The work then builds up a theoretical decomposition of variance terms in the predictive distribution of an ensemble of neural networks. To do this, the authors begin by analyzing an ensemble of infinitely wide neural networks, whose kernels are deterministic, demonstrating that it's predictive variance can explained by the _functional noise_ over different initializations. The authors then extend this analysis to _finite_ width linear networks, for one and more layers - under some assumptions, the authors show that the predictive variance of finite-width linear networks can be described largely by two noise sources: the _functional noise_ as described above, and the _kernel noise_ stemming from the variance over different kernels in the finite-width regime. Finally, the authors provide evidence that the discrepancy between a linearly trained model and it's non-linear, gradient descent counterpart, can be bounded on the order of it's layer-width, justifying the approximation of non-linear ensembles with linear models.\n\nThe second section (section 3) validates the authors hypothesis via a set of empirical experiments. The first set of experiments utilizes tractable linear models for which distinct components of predictive variance can be explicitly removed. Here the authors use a  toy dataset to provide an intuitive visualization of how each isolated component of the variance behaves. Next the authors study 2 real-world tasks, CIFAR-10 and MNIST, using a more complex linear model (both a CNN and MLP model), examining both ensemble performance and ensemble OOD detection rates. They find that different noise components have different levels of OOD detection - in some cases, the model with only kernel noise can achieve comparable or better performance in OOD detection than the full model. They then extend their empirical analysis to non-linear models - using a clever trick around model initialization, the authors derive two variants of gradient-descent models which isolate either functional initialization noise or initial kernel noise. They find that the isolated components of predictive variance in neural ensembles trained with (S)GD mirror those of the linearly trained models on OOD detection performance. Finally, the authors train fully modern neural networks, using residual connections, batch-norm, and cross-entropy loss on CIFAR-10, leveraging the same techniques to isolate each noise source. They show that, in this setting, isolating noise sources can have a significant effect on OOD performance - in particular, _centered_ models which remove initialization noise, outperform standard models.", "questions": "- The experiment in Figure 1 is a nice visualization, but the intuition it provides is not tied into later results - why is the distance or angle between datapoints a useful interpretation of functional or kernel noise? If it is not relevant to the next set of experiments due to data dependency then what is the toy experiment showing us?  ", "limitations": "yes", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The predictive variance of neural networks, and it's utility to OOD detection is of broad interest to the field. \n- Leveraging linear networks (and their kernels) to get a theoretical intuition about the different sources of variance in neural ensembles is a clever and novel approach to understanding the sources of variance, and their effects.\n- The theoretical analysis provided in this work is very interesting, and provides the foundation for future work aiming to understand predictive variance in neural models.\n- The final empirical result is relatively strong - namely, they show that a particular source of noise (kernel noise) is more effective at identifying OOD examples than standard neural models in modern neural architectures trained with SGD.\n\nWeaknesses:\n- The biggest weakness of this work lies in the empirical analysis and the takeaways from this analysis:\n\t- Overall, the benefits of each isolated noise source are not consistent. It is not clear what effect of each noise source has on OOD detection - in some cases isolated kernel noise outperforms isolated functional noise, and in other cases the opposite occurs.\n\t- The \"big takeaway\" of the role of each noise source in OOD detection largely seems to rely on the single result of the fully fleshed out neural network on CIFAR-10. Only is this model is there a clear story on the importance of isolated kernel noise.\n\t- In other words, the paper cannot empirically show us why the proposed decomposition is practically useful other than that it has _some effect_ on OOD detection.\n- Lastly, I think the structure of the paper could be improved:\n\t- For instance, much of the theoretical analysis lacks intuition around key quantities (e.g. the matrix Q is referenced all over the paper, and it would help the readability of the formulas to provide some intuition as to what this matrix represents).\n\t- Figures and tables are placed distant from their corresponding sections - additionally, sometime the corresponding figure or table is not referenced in the main text (e.g. Table 1 is not referenced in section 3.2.2).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "7KcNJS6LTi1", "reviewer": "Reviewer_VVkd", "summary": "This paper provides a first order approximation of the variance of a one-layer neural network. This approximation is composed of 5 terms:\n\n- $\\mathbb V^a$ is the finite-width equivalent of the predictive variance of an infinite width NTK, and depends on functional noise at initialization\n- $\\mathbb V^c$ is an approximation of the predictive variance of a linearly-trained network, and depends only on the kernel noise.\n- $\\mathbb V^i$ is the predictive variance of an ensemble with a NTK kernel and a data-specific prior\n- $\\mathbb V^{cor}$ and $\\mathbb V^{res}$ are two terms that are negligible in empirical experiments.\n\nBased on this decomposition, the authors define ensembles that depend only on one or another source of noise ($f^c$ and $f^a$), and show that isolating these sources of noise can improve OOD detection.", "questions": "- How would one extend the decomposition of Prop. 2.2 in the presence of additional sources of noise? Would an approach such as [1] be sufficient?\n- How are the AUROC values computed? Are they computed based on the (pointwise) predictive variances of the different ensembles?\n- Line 219: \"Interestingly, $\\mathbb V[f^{lin}(x)]$ depicts a superposition of the three isolated variances.\" Is this not a validation of the decomposition of Prop 2.2? Why is this surprising?\n- In cases where we have $\\mathbb V[f^{lin}] \\le \\mathbb V[f^{lin-c}] + \\mathbb V[f^{lin-a}]$, is this due to $\\mathbb V_i$? Does $\\mathbb V_i$ factor into defining favorable/unfavorable sources of noise?\n- Are there any difficulties in obtaining high-quality estimates of the different variance terms? Are the empirical estimates unbiased?\n\n\n[1] Understanding Double Descent Requires A Fine-Grained Bias-Variance Decomposition, Adlam & Pennington, 2020", "limitations": "The authors discuss different limitations of their work, mostly based around assumptions made about the behavior of NNs (e.g., 2.3). These assumptions are discussed in detail, evaluated empirically, and put in the general context of this work.\n\nIn my mind, the two major limitations of this work are the following:\n- The interplay between the two sources of noise (initialization & kernel noise) is not analyzed empirically, despite appearing with equal importance in the decomposition of Prop. 2.2. \n- I found the conclusions of the empirical analysis to be somewhat unclear. I believe the discussion could be improved by the authors stating early on which hypothesis they seek to verify, and how that hypothesis is or isn't validated by the experimental results. In the paper's current state, I find the conclusions to be somewhat unclear (which source of noise is conducive to OOD detection? Is a single source of noise sufficient in general?).", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strengths\n- Understanding the degree to which the variance of a deep ensemble is indicative of an ensemble's certainty in making a prediction is a crucial question as deep ensembles are deployed on real-world applications.\n- The authors push through a detailed decomposition of the variance of linearly trained neural networks, then generalize this decomposition under several assumptions to neural networks trained with gradient descent and SGD. \n- The authors analyze the terms in the proposed decomposition empirically across a variety of experimental settings (different training modalities, datasets, and dataset samples).\n\n# Weaknesses\n- I found the empirical section of this paper difficult to put into context with the previous decomposition (cf. \"Questions\" section below).\n- I am not sure I follow the author's claim that isolating unfavorable noise improves OOD detection.  This appears to be the case in certain experiments but not all (e.g., Table 1 for the CNN on Cifar-10), and it also seems like the \"unfavorable\" noise depends on the experiment setting (Table 1, linear models on CNN: $f^{lin-c}$ does better for SVHN but worse for LSUN).\n- I would've liked to see $\\mathbb V^i$ included in the empirical analysis, since it is the third major linear term in the approximation to the full variance. \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "6zbkwgEg1uZ", "reviewer": "Reviewer_GBTo", "summary": "This manuscript uses the concepts of Neural Network Gaussian Processes and Neural Tangent Kernels to analyze the sources of prediction variance from neural networks.  By making several reasonable assumptions, they can break the predictive variance into the superposition of several distinct terms.  Through this mathematical analysis and accompanying empirical studies, they demonstrate the practical consequences of this decomposition on several different neural networks.  Finally, they demonstrate that these decompositions can be used to build a greater intuitive understanding on out-of-distribution data.", "questions": "How exactly were OOD tasks performed?  Update: This has been largely clarified after the rebuttal.\n\nI do not fully understand the claims about the relationship between each variance term and OOD on data-specific performance.  Can you please elaborate on your claims of how the results highlight and explain these relationships?  Update: The response really highlights the challenges of interpretation.  Further claims and exploration are beyond the scope of this work, but I would encourage the authors to continue examining this question.", "limitations": "I believe that the authors have been forthright and fair in their descriptions of the limitations of their theoretical analysis.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "This manuscript clearly presents its theoretical arguments and accompanies them with empirical studies that confirm the findings.  Given the large focus on robustness to out-of-distribution data in recent years, this type of decomposition can help practitioners think about potential ways to understand and improve their models.  I believe that this decomposition is original, and has high quality and significance.\n\nThe final empirical studies, though, are underexplained and a bit confusing.  It is not clearly described enough what prediction tasks are being performed (simply “out-of-distribution” tasks) and what approach is being used to determine out-of-distribution samples.  This needs to be much more clearly described in a revised manuscript.  The claims about performance are not refined enough. They claim that “These observations highlight the particular inductive bias of each variance term for OOD detection.” However, it is not clearly described enough how these inductive biases relate to the particular predictive performance.  The analysis would be much enhanced with clearer statements about exactly how these inductive biases relate to performance.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept