An Empirical Analysis of the Advantages of Finite- v.s. Infinite-Width Bayesian Neural Networks
Abstract
Comparing Bayesian neural networks (BNNs) with different widths is challenging because, as the width increases, multiple model properties change simultaneously, and, inference in the finite-width case is intractable.
In this work, we empirically compare finite- and infinite-width BNNs and provide quantitative and qualitative explanations for their performance difference.
We find that when model is misspecified, increasing width can hurt BNN performance.
In these cases, we provide evidence that finite-width BNNs generalize better partially due to the properties of their frequency spectrum that allows them to adapt under model mismatch
Introduction
Recent works in Bayesian deep learning note a counter-intuitive phenomenon—that larger architectures (specifically larger width) can hurt model performance.
For example, Lee et al.
observe that on image classification tasks, as the width of the Bayesian CNN increases, classification accuracy decreases; Pleiss and Cunningham make similar observations on regression tasks.
Despite better performance of smaller NNs, posterior inference in finite-width BNNs is challenging [Coker et al., 2022, Foong et al., 2020], whereas, in the infinite-width limit, BNNs converge to a well-behaved model, the Neural Network Gaussian Processes (NNGPs)
(e.g., [Lee et al., 2017, Matthews et al., 2018, Neal, 2012]).
Furthermore, infinite-width BNNs easily admit theoretical analysis while finitewidth BNNs are difficult to analyze due to the non-Gaussian processes [Noci et al., 2021] they define in function space.
Thus, to choose between finite- versus infinite-width BNNs, we need to know: when and why does solving the harder inference problem present a better trade-off for improved performance?
In this paper, we provide empirical explanations for the counter-intuitive phenomenon that larger width hurts BNN performance, by comparing the inductive biases and frequency spectra of finite- and infinite-width BNNs.
We study one-hidden-layer BNNs with erf and ReLU activations, whose infinite-width counterparts are referred to as limiting NNGPs in the following.
Our contributions are 1 :
We show that when there is model mismatch between the limiting NNGP and the data generating model, finite-width BNNs can generalize better than infinite-width BNNs; 2.
To study the above phenomenon, we quantitatively show that finite-width BNNs have different inductive biases from NNGPs (finite-width BNNs generate more diverse datasets than NNGPs); 3.
We also qualitatively compare the inductive biases of finite- and infinite-width BNNs by comparing their frequency spectra from spectral analysis.
We demonstrate that, compared to limiting NNGPs, the prior predictive distributions of finite-width BNNs define heavier-tailed coefficient distributions over the frequency spectrum (they place more mass on large spectral coefficients).
Additionally, after inference, the posterior predictive distributions of finite-width BNNs define frequency spectra that are more similar to those defined by the data generating model; 4.
We hypothesize that because finite-width BNNs define frequency spectra with heavy-tailed spectral coefficient distributions, they are more flexible in modeling the data using different frequencies.
To test this, we design an inference procedure to filter out high-frequency components from the functional prior of BNNs.
We see that the generalization performance of BNNs decreases as more high-frequency components are removed from the function, supporting our hypothesis.
Additionally, we reconstruct datasets by removing the high-frequency components.
We observe that finite-width BNNs outperform the limiting NNGPs with their spectral coefficients of highfrequency components more concentrated at 0 , which resembles the true data generating process.
In this work, we consider both erf and ReLU nonlinearity, since in these cases we obtain closed-form kernels for the limiting NNGPs, which allows for exact inference and theoretical analysis.
Additionally, both activation functions are well studied in the literature (e.g., [Williams and Rasmussen, 2006,
Cho and Saul, 2011]) and are suitable for modeling various types of functions [Lee et al., 2020].
Related Work
A large body of work studies when and why finite-width (non-Bayesian) NNs succeed.
For example, Ghorbani et al.
and Refinetti et al. empirically find that finite-width NNs outperform kernel methods when the data contains low-dimensional structures.
Aitchison [2020] shows that finite deep linear networks yield stochastic kernels, which are more flexible in learning representations from the data.
However, it is not known whether these insights apply to BNNs.
In the context of BNNs, Pleiss and Cunningham study the performance degeneration of infinite-width
BNNs by studying deep GPs, which the authors consider to be a more generalized BNN model.
The authors prove that as the width of a deep GP increases, the posterior collapses from a mixture function of data-dependent bases to that of data-independent bases.
Thus, infinitely wide deep GPs are less capable of adapting to the data.
Similar insights have been made in Aitchison [2021].
We study BNNs from a spectral perspective and find that finite-width BNNs are more flexible in that they can use functions of different frequencies.
Analogously to standard NNs [Rahaman et al., 2019], existing works observe that, unlike their infinite counterparts, finite-width BNNs tend to be heavy-tailed in function space (e.g., [Noci et al., 2021,
Pleiss and Cunningham, 2021]).
We examine how and why this heavy-tailedness affects the prediction tasks under model mismatch.
Finally, a recent line of work studies the spectral properties and the reproducing kernel Hilbert space of infinite-width NNs (e.g., [Chen and Xu, 2020, Bietti and Bach, 2020]).
Another body of work provides theoretical and empirical analyses of how the power spectrum of standard, wide NNs changes during training (e.g., [Rahaman et al., 2019, Cao et al., 2019, Yang and Salman, 2019]).
However, both bodies of work do not consider the finite-width, Bayesian settings, nor the relationship between the power spectrum and generalization performance.
Background We consider a regression task with dataset D
= { x i , y i } ni
=1 , x i ∈ R d in , y i ∈ R .
For an
NN with one hidden-layer of width H , the output is given by f NN ( x i )
= [-
-]1 √ H w 1 ϕ ( w 0 x
i + b 0 )
+ b 1 , with w 0 ∈ R H
×
d in , b 0 ∈ R H × 1 , w 1 ∈ R 1 × H , b 1 ∈ R , and ϕ is a nonlinear activation function.
We denote neural network weights as W and biases as b .
For a Bayesian Neural Network, we specify prior distributions p ( W ) ∼ N (0 , σ 2 W I ) , p ( b ) ∼ N (0 , σ 2 b I ) .
The posterior distribution of the weights is given by: p ( W , b |D ) ∝ p ( D| W , b ) p
( W ) p ( b ) , where p ( D| W , b ) = (cid:81)
i p ( y i | x i , W , b ) , p ( y i | x i , W , b ) =
N ( y i |
f NN ( x i ) , σ 2 ϵ ) and σ 2 ϵ is the variance of the output noise.
Given a set of new inputs D ∗ = { x ∗ i } mi =1 , the posterior predictive distribution is given by p ( f ∗ | x ∗ , D ) =
(cid:90)
W , b p ( f ∗ | x ∗ , W , b , D ) p ( W , b |D )
d W d b , where f ∗ =
f NN ( x ∗ ) .
We use NUTS [Hoffman and Gelman, 2014] for sampling (Appendix C).
With erf or ReLU activations, a BNN will converge to a GP with Arcsin or Arccos kernel, respectively (Appendix A)
[Williams and
Rasmussen, 2006, Cho and Saul, 2011].
In practice, Arccos kernel has eigenvalues that decay slower than those of Arcsin kernel and is thus more suitable for learning rough functions [Lee et al., 2020].
Generalization Performance of Finite- and Infinite-Width BNNs
We compare the regression performance of finite- and infinite-width BNNs under model mismatch— we train BNNs whose limiting NNGPs are different from the data generating model.
Thus, the performance of the limiting NNGP is not optimal, leaving room for improvement.
We want to see if, under this model mismatch, finite-width BNNs perform differently from infinite-width BNNs.
Datasets We generate data using GPs with three kernels: RBF, Arcsin and Arccos (GP-Arccos generates rougher functions that require a larger model capacity to learn).
For GP-RBF, the lengthscale l is chosen from { 0 . 5 , 1 . 0 , 1 . 5 , 2 . 0 } .
For GP-Arcsin/Arccos, the hyperparameters σ 2 W , σ 2 b in
Equation 6, 7, are chosen from { 0 . 5 , 1 . 0 , 1 . 5 , 2 . 0 } .
For each kernel setting, we draw 200 functions, each as a univariate dataset (details in Appendix B).
We then fit each dataset using BNNs with erf and
ReLU nonlinearity whose prior variances σ 2 W , σ 2 b are chosen from { 0 . 5 , 1 . 0 , 1 . 5 , 2 . 0 } and hidden widths H vary from 2 to 200 .
For data drawn from a GP-RBF, there is always model-mismatch between the limiting NNGP and the data generating model.
For data drawn from a GP-Arcsin/Arccos, when the hyperparameters are equivalent to the prior variances of BNNs with erf/ ReLU nonlinearity, there is no model mismatch.
Evaluation Metrics To quantify the performance of finite-width BNNs, we calculate the expected difference of test negative loglikelihood (NLL) and test mean squared error (MSE).
Specifically, for each kernel setting with datasets {D ( s ) } Ss =1 , we calculate
[Equation]
When there is no model mismatch, both terms will converge to 0 as hidden width H goes to infinity.
Results When there is a model mismatch between the data generating model and the limiting
NNGP, finite-width BNNs can outperform NNGPs.
Previous work observed similar phenomena on real datasets (e.g.,Pleiss and Cunningham [2021], Aitchison [2021]), which inevitably have model mismatch.
In Figure 1a, 1b, we see that finite-width BNNs have better generalization performance (the black curves are below the red dashed lines) and converge to the limiting NNGP when the hidden width increases (the black curves approach the red dashed lines).
We see similar trends on large real datasets (Figure 9).
Additionally, a larger model mismatch increases the difference in the generalization performance (Figure 7, 8).
[Figure]
Inductive Biases of Finite- and Infinite-Width BNNs
Why do finite-width BNNs outperform infinite-width BNNs even though they have a smaller model capacity?
To answer this question, we first empirically show that finite-width BNNs have different inductive biases from infinite-width BNNs (finite- and infinite-width BNNs are two distinct model classes); thus, one may generalize better than the other given different situations.
3.1 Quantitatively Comparing Inductive Biases Using the Data Likelihood
There are different ways to quantify inductive biases, one of which is to evaluate how much the model assumptions overlap with the true data generating process.
In the Bayesian literature, log marginal likelihood (LML) is often used to quantify this overlap.
The LML of generating dataset D under model M can be computed by marginalizing out the prior over the function values f ,
[Equation]
In our setup, the dataset D is generated from a GP and f is modeled by either a BNN or an NNGP.
Equation 3 quantifies the inductive bias (overlap between the model assumptions and the data generating process) by answering the question “ how likely is the dataset to be generated by the model ?”
The literature argues that a model has a good inductive bias when it can achieve high LML for various datasets (a model is too complex when achieving modest LML for many datasets and too simple with high LML for limited datasets (e.g., [Mackay, 1992])).
Additionally, higher LML does not necessarily imply better generalization performance (e.g., [Lotfi et al., 2022]).
Despite many good properties of LML, Equation 3 is hard to compute (the Monte Carlo estimator exhibits large bias and variance given the large parameter space of BNNs).
In this work, we use an alternative way to quantify the model class overlap (we note that this is not a proxy for LML): we still compute Equation 3 but now the datasets D are drawn from finite-width
BNNs or the limiting NNGPs while the model M is the data generating GP.
We refer to this metric as log data likelihood (LDL) to avoid confusion with LML.
A model has a desirable inductive bias when it can generate diverse datasets with high LDL under the data generating model.
Given a data generating GP with prior GP ( m ( x ) , K ) , LDL can be efficiently computed as:
[Equation]
Σ − 1 f ( y − m ( x ))
− 1 2 log | Σ f | − n 2 log 2 π, where Σ f = K + σ 2 ϵ I.
Results Finite- and infinite-width BNNs have distinct inductive biases .
In Figures 2a and 2b, we see that BNNs with larger widths generate less diverse datasets: the CDF of larger BNNs is higher in the high LDL region (i.e., they generate limited datasets with high LDL).
[Figure]
3.2 Qualitatively Comparing Inductive Biases in Function Space
While the LDL allows us to quantitatively compare the inductive biases of two Bayesian model classes, we also want to qualitatively describe the differences between them.
For the latter, we want to compare the differences between the prior and posterior predictive distributions of finite- and infinite-width BNNs in function space.
Previous works prove that finite-width BNNs yield heavy-tailed priors in function space [Noci et al., 2021, Zavatone-Veth and Pehlevan, 2021].
However, they do not discuss how heavy-tailed priors affect the generalization performance.
[Figure]
One common approach to quantifying heavy-tailedness is to compute the higher-order moments of the distribution; however, this is not easy because the function space is large and requires a prohibitively large number of samples.
We instead analyze the samples drawn from the prior and posterior predictive distributions using discrete cosine transform (DCT).
DCT is a linear transformation that expresses a function as a weighted sum of cosines of different frequencies.
Formally, given a function f ⊺ =
[ f 0 , . .
, f N − 1 ] ⊺ , DCT is a linear and invertible function,
T DCT : R N −→ R N , that transforms f as:
[Equation]
T DCT [Equation] where a ⊺ =
[ a 0 , . .
. , a N − 1 ] ⊺ are the DCT coefficients representing the weights on cosines of different frequencies in ascending order ( a 0 is the weight on the lowest frequency and a N − 1 , on the highest).
We qualitatively compare the DCT coefficient distributions of finite- and infinite-width BNNs.
Results Figure 3a shows that comparing to the NNGPs, the functional priors of finite-width BNNs define a spectrum with more mass on large coefficient values across different frequencies (the 99% percentile of DCT coefficients defined by finite-width BNNs is larger).
Figure 3b shows similar trends after inference.
Additionally, the frequency spectrum defined by finite-width BNNs behaves more similar to that of the data generating model , which also puts more mass on the larger coefficients (the black dashed line is higher).
3.3 Connecting the Inductive Bias to Performance Difference Through Frequency Filtering
In the previous section, we show that the frequency spectrum defined by prior predictive distributions of finite-width BNN places more mass on large coefficient values than that of the limiting NNGP, which makes them more flexible in using functions of different frequencies.
In this section, we conduct two sets of experiments to confirm this hypothesis, using “low-pass filtering”.
[Figure]
We chose thresholds empirically ( t ≤ 0 .
Low-pass Filtering Given a function f , we first compute its DCT coefficients a via Equation 5, and reconstruct that function via f = T ⊤ DCT a .
We then remove high-frequency components from the function by setting the corresponding DCT coefficients to 0 :
˜ f
= T ⊤ DCT I ( k ) N T DCT f : =
Cf , where k = tN, t ∈ [0 , 1] , and I ( k ) N is an N × N identity matrix with the last k diagonal elements set to 0 .
A larger t removes more high-frequency components, which makes the reconstructed function ˜ f smoother (Figures 14 and 15).
t = 0 recovers the original function.
Low-pass filtered BNNs We compare the generalization performance of a finite-width BNN with that of a “low-pass filtered” BNN—a BNN in which we filter out high-frequency components from the prior predictive distribution (i.e., we “smooth” functions drawn from the functional prior).
We note that our approach modifies the function-space prior of BNNs directly, as opposed to post-hoc modifying the posterior predictive distributions, which can be coupled with any inference method (details in Appendix D).
We find that the model performance decreases now that finite-width BNNs are constrained to the low-frequency regime.
Results Removing high-frequency components hurts the generalization performance of finitewidth BNNs.
In Figure 4, we show that with more high-frequency components removed ( t increases), the generalization performance decreases ( ∆ MSE increases).
Given different BNN architectures, removing the same proportion of the high frequency components has different effects (when t = 0 .
93 , low-pass filtered BNNs with ReLU outperform the limiting NNGP while low-pass filtered BNNs with erf do not).
Additional results are in Appendix E.3.
Low-pass filtered Datasets To provide additional evidence that finite-width BNNs are more flexible, we construct datasets with high-frequency components removed.
We expect that finite-width BNNs can adapt to this model-mismatch by putting low mass on high-frequency components.
We apply low-pass filtering to data generating GPs and compare the DCT spectrum defined by the posterior predictive distributions of finite- and infinite-width BNNs.
Results Finite-width BNNs can adapt to the model-mismatch we designed .
Figure 5 shows that for finite-width BNNs, the distributions of DCT coefficients of the high-frequency components put more mass around 0 (i.e., the percentiles of DCT coefficients are lower as the solid line is mostly below the dashed line), which resembles the data generating process ( a i = 0 for i > k ).
Discussion
In this paper, we quantitatively and qualitatively compare the behaviors of finite- and infinite-width
BNNs.
We empirically show that finite-width BNNs define a frequency spectrum with more mass on large coefficient values, which allows them to adapt to the data faster under model mismatch.
Our experiments suggest that there are advantages to using non-Gaussian models, such as finite-width
BNNs or non-Gaussian approximation of the true posterior, even though inference may be challenging.
Future work is needed to determine the optimal width given the dataset size.
Acknowledgements
This material is based upon work supported by the National Science Foundation under Grant No.
IIS-2007076.
Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of the National Science Foundation.
A Analytical Form for NNGP Kernels
Define the following terms,
[Equation]
For a one hidden-layer BNN with erf nonlinearity, the BNN will converge to a GP with Arcsin kernel defined as,
[Equation]
where θ = sin − 1 (cid:16) [Equation]
[Equation]
For a one hidden-layer BNN with ReLU nonlinearity, the BNN will converge to a GP with Arccos kernel defined as,
[Equation]
where
[Equation]
The complete derivation can be found in [Williams and Rasmussen, 2006, Cho and Saul, 2011].
B Datasets
Synthetic Datasets We generate data using GPs with three kernels, RBF, Arcsin and Arccos (GP-Arccos generates rougher functions that require larger model capacity to learn).
For GP-RBF, the lengthscale l is chosen from { 0 . 5 , 1 . 0 , 1 . 5 , 2 . 0 } .
For GP-Arcsin/Arccos, the hyperparameters σ 2 W , σ 2 b in Equation 6, 7, are chosen from { 0 . 5 , 1 . 0 , 1 . 5 , 2 . 0 } .
For each kernel setting, we drawfunctions, each as a univariate dataset.
For each dataset, we uniformly sample 20 training points from [ − 3 , − 1] (cid:83) [1 , 3] aggregate with x = 0 to facilitate learning ( N train = 21 ) and N test = 100 test point from [ − 3 , 3] .
We add Gaussian noise with σ ϵ = 0 .
1 to the training data.
Real Datasets We test 3 real datasets: UCI boston, UCI Energy (the output dimension are fit individually), Sic97 [Shah et al., 2014].
We split the data into train, test, valid with the ratio 80 − 10 − 10 .
We perform cross validation on NNGPs and choose the GP hyper-parameter with the highest validation MSE.
C Models and Inference
We fit each dataset using BNNs with erf and ReLU nonlinearity whose prior variances σ 2 W , σ 2 b are chosen from { 0 . 5 , 1 . 0 , 1 . 5 , 2 . 0 } and hidden widths H vary from 2 to 200 .
We use NUTS implemented in NumPyro [Phan et al., 2019] to obtain posterior samples from p ( W , b |D ) .
We used 4 parallel chains, each with 10 , 000 samples during burn-in, 40 , 000 MCMC samples and a thinning of 5 .
D Low-pass filtered BNNs
We modify the original BNN model,
[Equation]
[Equation]
so that the f NN is filtered to include only its low-frequency components.
We denote this low-pass filtered function by ˜ f NN , and describe a general algorithm for computing ˜ f NN and incorporating it into the model, allowing it to be paired with any inference algorithm (e.g., NUTS, which requires ˜ f NN to be differentiable).
For simplicity, we assume the inputs are 1-dimensional ( d in = 1 ).
Step 1:
From weight-space to function space.
To filter out high-frequency components of a function f NN , we cannot simply operate on the corresponding neural network weights W , b .
Instead, we decompose the function (evaluated on a grid of points) into a weighted sum of cosines (using the
DCT) and then remove the highest frequency terms in the sum.
Specifically, we define a grid of N grid uniformly spaced inputs x grid n ∈ R d in , denoted jointly by X grid ∈ R N grid × d in .
We perform the DCT on f NN ( · ; W , b ) evaluated on this grid.
We further ensure that our training and test sets lie on this grid by mapping them to their nearest neighbors on the grid (a fine grid will only negligibly modify the training and test sets).
Step 2: Performing the DCT.
We compute the type-II orthonormal DCT, a k ( W , b ) ∈ R N grid , of f NN ( X grid ; W , b ) as follows:
[Equation]
for k = 0 , . . .
, N grid − 1 .
Step 3: Low-pass Filtering.
We select a filtering threshold t ∈ [0 , 1] , where t = 0 recovers the original f NN , and as t increases, frequencies are filtered out from highest to lowest.
We then zero out all DCT coefficients a k ( W , b ) for which k ≥ (1 − t ) N grid , and construct the low-pass filtered function using the type-II inverse orthonormal discrete cosine transformation:
[Equation]
for n = 0 , . . .
, N grid − 1 .
Step 4: Inference.
Lastly, we incorporate the low-pass filtered function into the BNN model by replacing f NN with
˜ f NN in the model in Equation 8 as follows:
[Equation]
where 0 x is a vector of 0s with a 1 corresponding to the row of X grid closest to x .
Since the low-pass filter is a linear operation, it is differentiable and enables pairing with any inference method available in standard probabilistic programming languages.
Computational Efficiency.
We can vectorize the above algorithm as follows:
Let T ∈ R N grid × N grid denote the matrix corresponding to the type-II orthonormal discrete cosine transform on a grid of points X grid ∈ R N grid :
[Equation]
So, a = T DCT f and T ⊤ DCT is the inverse transform.
Let I N grid ,t denote an identity matrix of size N grid with ( I N grid ,t ) kk = 0 for k ≥ (1 − t ) N grid .
Let
˜ f NN
( X grid ; W , b ) =
T ⊤ DCT
I N grid ,t T DCT · f NN ( X grid ; W , b ) .
For a batch of inputs X ∈ R N × N grid , we can write the likelihood as
[Equation]
where 0 X = ( 0 ⊺ x 1 , . .
. , 0 ⊺ x N ) ⊺ ∈ R N × N grid is a matrix of 0s and a single 1 in each row that functions to select the elements of ˜ f NN ( X grid ; W , b ) corresponding to the training inputs x 1 , . . .
, x N .
E Additional Results E.1 Generalization Performance of Finite and Infinite Width BNNs
We show additional results comparing the generalization performance of finite and infinite width
BNNs.
Under all model mismatch (the limiting NNGP is not equivalent to the data generating model) we design, finite width BNNs can outperform infinite width BNNs.
Additionally, we observe that larger model mismatch results in larger performance gap (Figure 7, 8).
We also provide empirical results on the real dataset (Figure 9) and we see that finite width BNNs perform better on large datasets (Figure 9b, 9c, 9d).
[Figure]
[Figure] [Figure]
[Figure]
E.2 Log Marginal Likelihood (LML)
We provide additional results of the CDF of LML.
Similar to Figure 2a, 2b, we see that under different types of data generating models, BNNs with smaller hidden widths generate more diverse datasets than the limiting NNGP (the CDF of LML over datasets drawn from the data generating GP has more mass on the tail).
E.3 Low-pass filtered BNNs
Removing high-frequency components from the BNN hurts its ability to adapt to model mismatch.
We fix the data-generating process to be that of a GP.
Across varying widths, we then compare the ∆ MSE (Equation 2) of a BNN with that of low-pass filtered BNNs with different amounts of filtering, t ∈ [0 , 1] .
As expected, as the width increases, ∆ MSE of the BNN converges to 0 .
Moreover, under no model mismatch (i.e., when the limiting NNGP model matches the data generating process, as in Figure 11a), the BNN always under-performs the limiting NNGP.
In contrast, under model mismatch (remainder of Figure 11, as well as Figures 12 and 13), small-width BNNs always outperform the NNGP (i.e., ∆ MSE < 0 ).
As we hypothesized, removing the high-frequency components from prior functions hinders the BNN’s ability to adapt to model mismatch.
[Figure] [Figure]
A qualitative comparison of a BNN vs. a Low-pass filtered BNN.
Figures 14 and 15 show how a BNN ( t = 0 ) differs from a low-pass filtered BNN ( t > 0 ), both in its prior and in its posterior predictive distributions.
Specifically, as t increases, the functions being drawn become more smooth.
[Figure] [Figure]