Functional Priors for Bayesian Neural Networks through Wasserstein Distance Minimization to Gaussian Processes
Abstract
The Bayesian treatment of neural networks dictates that a prior distribution is considered over the weight and bias parameters of the network.
The non-linear nature of the model implies that any distribution of the parameters has an unpredictable eﬀect on the distribution of the function output.
Gaussian processes oﬀer a rigorous framework to deﬁne prior distributions over the space of functions.
Our proposal is to impose such functional priors on well-established architectures of neural networks by means of minimising the Wasserstein distance between samples of stochastic processes.
Early experimental results demonstrate the potential of functional priors for Bayesian neural networks
1. Introduction
The concept of prior distribution in Bayesian inference oﬀers a mathematical formulation that allows us to describe the family of solutions that we consider acceptable, before having seen any data.
While there are cases for which picking a good prior is easy or intuitive given the context (O’Hagan, 1991; Rasmussen and Ghahramani, 2002; Srinivas et al., 2010;
Cockayne et al., 2019; Briol et al., 2019), for nonlinear parametric models with thousands (or millions) parameters, like deep neural networks ( dnn s ) and convolutional neural networks ( cnn s ), this choice is not straightforward.
As these models are nowadays accepted as de facto standard in machine learning (LeCun et al., 2015), the community has been actively proposing ways to reason about the uncertainty of their predictions, with the Bayesian machinery being at the core of many contributions (Graves, 2011; Chen et al., 2014; Gal and Ghahramani, 2016; Liu and Wang, 2016).
Despite many advances in the ﬁeld (Kendall and Gal, 2017; Rossi et al., 2019; Osawa et al., 2019; Rossi et al., 2020), it is reported that in some cases the predictive posteriors are not competitive to non-Bayesian alternatives, making these models—and Bayesian deep learning, in general—less than ideal solutions for a number of cases, if no remedies are taken (Wenzel et al., 2020).
In this work, we focus our discussion on the prior distribution of Bayesian neural networks ( bnn s ).
For such models, the common practice is to deﬁne a prior distribution on the network weights and biases, often a Gaussian distribution.
A prior on the parameters induces a prior on the functions generated by the model, which also depends on the network architecture.
However, due to the nonlinear nature of the model, the eﬀect of this prior on the functional output is not obvious to characterize and control. [Figure]
Consider the example in Fig.
1, where we show the functions generated by sampling the weights of bnn s with tanh activation from a Gaussian prior N (0 , 1).
We see that as depth increases the samples tend to form straight horizontal lines with large marginal variance, which is a pathology known in the literature (Duvenaud et al., 2014).
Also, it has been recently shown that for ReLU activation, the prior distribution of unit outputs become more heavy-tailed as the network depth increases (Vladimirova et al., 2019).
We stress that a ﬁxed Gaussian prior on the parameters is not always problematic, but it can be, especially for deeper architectures.
Nonetheless, this kind of generative priors on the functions is very diﬀerent from shallow Bayesian models, such as Gaussian Processes ( gp s ), where the selection of an appropriate prior typically reﬂects certain attributes that we expect from the generated functions.
A gp deﬁnes a distribution of functions that is characterized by a mean and a kernel function κ .
The gp prior speciﬁcation can be more interpretable than the one induced by the prior over the weights of a bnn , in the sense that the kernel eﬀectively governs the properties of possible functions, such as shape, variability and smoothness.
Contributions
We seek to tune bnn s so that their (functional) prior distributions exhibit interpretable properties, similar to shallow gp s .
We consider the Wasserstein distance between the distribution of functions induced by bnn prior over the network parameters, and a target gp prior.
We propose an algorithm that optimizes bnn hyper-parameters based on the distance between the bnn and a desired gp .
To the best of our knowledge, our approach is novel and it leads to sensible improvements over classical choices for the prior.
Related works
It is well-known that the prior over functions induced by a bnn converges to a gp as it width becomes large (Neal, 1996; Matthews et al., 2018).
This motivates a line of research that attempts to map gp priors to bnn priors (Flam-Shepherd et al., 2017; Pearce et al., 2019).
Closest to our work is that of Flam-Shepherd et al.
(2017), which propose to minimise the Kullback-Leibler ( kl ) divergence between the bnn prior and a target gp prior.
As there is no analytical form for this kl due the entropy term, the Authors propose to ignore this term and employ an early stopping scheme as it merely acts as a regularization term.
We compare with the kl -based approach in the following sections.
2. Using the Wasserstein distance to map GP prior in Bayesian NN
The equivalence between function-space view and weight-space view of linear models, like
Bayesian linear regression and gp s (Rasmussen and Williams, 2006), is a straightforward application of Gaussian identities, but it allows us to seamlessly switch point of view accordingly to which characteristics of the model we are willing to observe or impose.
We would like to leverage this equivalence also for bnn s but the nonlinear nature of such models makes it analytically intractable (or impossible, for non-invertible activation functions).
We argue that for bnn s —and Bayesian deep learning models, in general—starting from a prior over the weights is not ideal, given the impossibility of interpreting its eﬀect on the family of functions that the model can represent.
We therefore rely on an optimization-based procedure to impose functional priors on bnn s using the Wasserstein distance as a similarity metric between such distributions.
In the general case, given two probability measures π ( x ) and ν ( y ) deﬁned on a separable metric space X (e.g. R d ), the 1–Wasserstein distance is deﬁned as follows
[Equation]
where D ( x , y ) is a proper distance metric between two points x and y (e.g. Euclidean norm distance D ( x , y )
=
k x − y k )
and Γ( π, ν ) is the set of functionals of all possible joint densities γ whose marginals are π and ν .
With the exception of few cases where the solution is available analytically (e.g. π and ν being Gaussians), solving Eq. 1 directly or via optimization is intractable.
Fortunately, the Wasserstein distance in Eq.
1 admits the following dual form (Kantorovich, 1942, 1948),
[Equation]
[Equation]
where φ is a 1-Lipschitz continuous function deﬁned on X → R .
This is eﬀectively a functional maximization of φ on the diﬀerence between two expectations of φ , evaluated under π and ν , constrained to its Lipschitz continuity.
2.1. Wasserstein distance optimization for mapping priors
Assume to have a bnn with weights w on which we have a prior distribution p ( w ; ψ ), where ψ is its set of parameters (e.g. for a Gaussian prior, ψ = { µ, σ 2 } ).
This prior on the weights induces a corresponding prior distribution over functions p nn
[Equation] where p ( f | w ) is deterministically deﬁned by the network architecture.
Our target gp prior is p gp ( f ) =
N ( f | 0 , K ), where K is the covariance matrix obtained by computing the kernel function κ for each pair of { x i , x j } in the training set.
We aim at matching these two stochastic processes at a ﬁnite number of measurement points X M def =
[ x 1 , ..., x M ] > sampled from a distribution q ( x ).
To achieve that, we propose a sample-based approach using the 1-Wasserstein distance (Eq. 2) as the objective
[Equation]
where f M are the function values at X M and φ is 1-Lipschitz function.
Following recent literature (Goodfellow et al., 2014), we parameterize the Lipschitz function φ by a neural network with parameters θ 1 .
We alternate between n Lipschitz steps of maximising L with respect to the Lipschitz function’s parameters θ and one step of minimizing L with respect to the prior’s parameters ψ .
Notice that the objective is fully sample-based.
As a result, it is not necessary to know the closed-form of the marginal density p nn ( f ; ψ ).
Lipschitz constraint.
Based on the fact that a diﬀerentiable function is 1-Lipschitz if and only if the norm of its gradient is at most one everywhere, Gulrajani et al.
(2017) propose a soft constrain on the gradient norm of the output of the Lipschitz function φ θ with respect to its input.
The loss of the Lipschitz function is then augmented by a regularization term
[Equation]
Here ˆ f is the distribution of ˆ f = ε f nn + (1 − ε )
f gp for ε ∼ U [0 , 1] and f nn ∼ p nn , f gp ∼ p gp being the sample functions from bnn and gp priors, respectively; λ is the penalty coeﬃcient.
Choice of the measurement set.
We are using ﬁnite measurement sets to have a practical and well-deﬁned optimization.
For low-dimensional problems, one can simply use a regular grid or apply uniform sampling in the input domain.
For high-dimensional problems, one can sample from the training set or an augmented version of it, where noise is injected into the data.
We follow a combination of the two approaches: we use the training inputs (or a subset thereof) to which we include points that are randomly sampled (uniformly) in the input domain.
A pseudocode of the procedure is formally presented in Algorithm 1. Algorithm
Requires : B , mini-batch size; q ( x ), sampling distribution for measurement set; n Lipschitz , number of iterations of Lipschitz function per prior iteration; while ψ has not converged do draw X M from q ( x ) //
Sample measurement set ; for t = 1 , ..., n Lipschitz do draw GP functions { f ( i ) gp } Bi =1 ∼ p gp ( f ; κ ) at X M ;
3. Experimental validation
Before presenting the experimental results on the classic UCI benchmark for regression tasks, we report the details on the choice of prior distributions and we illustrate the advantages oﬀered by the use of the Wasserstein distance compared to alternative approaches.
3.1. Parametrization of BNN Prior
First, we observe that Algorithm 1 is completely independent of the choice of the form of the prior on the weights and biases of bnn s .
The question that now remains is “What form of
[Figure]
prior should we place on w ?” Since the family of functions the bnn can represent depends on both the architecture of the network and the activation function, we argue that there is not a unique choice.
We consider a layer-wise factorization with two independent Gaussian distributions for weights and biases.
The parameters to adjust are ψ = { σ 2 l w , σ 2 l b } Ll =1 , where σ 2 l w is the prior variance shared across all weights in layer l , and σ 2 l b is the respective variance for the bias parameters.
For any weight and bias entries w, b ∈ w , the prior is:
[Equation]
and
[Equation]
In the next section, we refer to this parameterization as the gp induced Gaussian ( gp i g ) prior.
Although this simple approach retains the Gaussian prior on the parameters, in many cases it is shown to suﬃciently capture the target functional priors.
3.2. KL divergence or Wasserstein distance?
The kl divergence is popular choice as an optimization criterion capturing similarity between distributions.
It was used in the same vein as in this work by Flam-Shepherd et al.
(2017), where they consider the KL divergence between samples of a bnn and a gp .
This requires an empirical estimation of the entropy, which is a challenging task for high-dimensional distributions (Delattre and Fournier, 2017).
These issues were also reported by FlamShepherd et al.
(2017), where they propose an early stopping scheme to what is essentially an optimization of the cross-entropy term.
In our experiments, we have found that a scheme based on Wasserstein distance converges more consistently without the need for additional heuristics.
We demonstrate the convergence properties of our scheme against Flam-Shepherd et al.
(2017) in Fig.
We monitor the evolution of Wasserstein distance from the target gp prior and performance metrics for the
Boston UCI dataset (test negative loglikelihood ( nll ) and root mean square error ( rmse )).
The KL-based scheme provides improvements for the ﬁrst few iterations, and then early stopping is required, as prescribed by Flam-Shepherd et al.
(2017) (dashed red line).
Our approach, instead, improves over the iterations and does not need early stopping to achieve converge to the desired prior, and it oﬀers consistently better performance.
3.3. Regression experiments on the UCI benchmark
We evaluate our proposed method on a classic UCI benchmark for regression.
We consider two ﬁxed priors: (1) the Fixed Gaussian prior ( fg ), N (0 , 1); (2) the Fixed Hierarchical
R M S E ( ← )
Boston 4 . 85 .
6 Concrete 0 .
5 Energy 0 .0 .0 .
Kin8nm 3 .
Power 3 . 84 .
Protein 0 .0 .0 .
Wine 2 . 3 .
N LL ( ← ) 3 .3 .3 .
− 1 .− 1 .2 .2 .2 .2 .2 .2 .0 .
1 .1st
3rd 5th
Rank ( ↓ ) 1st
3rd 5th
FG
FG + TS
GPiG (Ours)
FH
Deep Ensembles prior ( fh ) where the prior variance for each layer is sampled from an Inverse-Gamma distribution, Γ − 1 (1 , 1) (Chen et al., 2014); and our approach, (3) the gp induced Gaussian ( gp i g ) prior.
We also compare bnn s against Deep Ensembles (Lakshminarayanan et al., 2017), a state-of-the-art approach for uncertainty estimation in deep learning (Ashukha et al., 2020; Ovadia et al., 2019).
Furthermore, we additionally compare to the “cold” posterior (Wenzel et al., 2020) that uses a Fixed Gaussian prior and temperature scaling ( fg+ts ).
We use stochastic gradient Hamiltonian Monte Carlo ( sghmc )
(Chen et al., 2014) for posterior inference with scale adaptation (Springenberg et al., 2016), where the hyperparameters of sghmc are adjusted during a burn-in phase.
We sample four independent chains; for each chain, the number of collected samples after thinning is 30 except for the Protein dataset for which we collect 60.
We use a step size of 0 .
01 and a momentum coeﬃcient of 0 .
For fh , we resample the prior variances using a Gibb step everyiterations.
Each dataset is randomly split into 90%/10% train/test sets.
This splitting process is repeated 10 times except for the Protein dataset which uses 5 splits.
We use a 2 hidden-layer bnn with tanh activation function, containing 100 hidden units for smaller datasets and 200 hidden units for the Protein dataset.
We map a target hierarchical- gp prior with an RBF kernel to gp i g by using our proposed Wasserstein scheme.
We place a prior LogNormal(log √ 2 D, 1) on the length-scales and a prior LogNormal(0 .
1 , 1) on the variance ( D is the number of input dimensions).
Fig.
3 illustrates the average test nll and rmse .
For most datasets, our optimized priors provide the best results and outperform Deep Ensembles in terms of both rmse and nll .
We notice that tempering the posterior only shows small improvements on fg .
Instead with gp i g , the predictive performance of the true posterior improves considerably.
4. Concluding remarks and future work
We have analyzed the eﬀect of placing a prior on the output of a bnn rather than on its parameters.
To this goal, we borrowed the ﬂexibility of the Wasserstein distance as similarity metric between functions sampled from gp s to functions sampled from bnn s .
We tested this contribution on a classic regression benchmark and we showed sensible improvements over classic choices and state-of-the-art inference methods.
On these premises, our ongoing work focuses on deeper models, such as cnn s , and on tasks where having good uncertainty estimate is key for obtaining good results, such as active learning and Bayesian optimization.