Neural Linear Models with Functional Gaussian Process Priors
Abstract
Neural linear models (NLM) and Gaussian processes (GP) are both examples of Bayesian linear regression on rich feature spaces.
In contrast to the widespread use of nonparametric GPs for probabilistic nonlinear regression, NLMs remain an underused parametric alternative because standard type II maximum likelihood (ML) training leads to overconﬁdence outside of the data distribution.
Therefore, we propose to augment this training procedure through functional variational inference (fVI) proposed by Sun et al.
(2019), which is particularly well suited for NLMs due to their closed-form predictive distribution.
Additionally, we investigate whether an appropriate functional prior can guide parametric
NLMs to attain nonparametric GP performance, despite using fewer parameters.
Results show that functional priors can improve performance of NLM over ML training, and that the NLM performs on par with weight space BNNs in this setting.
1. Introduction
[Equation]
of data .
The models are trained with type II ML ( § ), functional ELBO ( ‡ ) and functional KL ( † ) objectives, and the NLMs use leaky ReLU activations.
Functional priors demonstrate that the NLM is expressive enough to learn diverse feature spaces and approximate a GP, but type II ML and fELBO training can lead to overﬁtting in the features resulting in reduced OOD uncertainty.
The neural linear model (NLM) (L´azaro-Gredilla and Figueiras-Vidal (2010); Ober and
Rasmussen (2019)) is a Bayesian neural network which combines deterministic neural network features with a distibution over the ﬁnal layer weights.
As such, it is a form of
Bayesian linear regression with a ﬁnite feature space, contrasting Gaussian processes (GP) with inﬁnite features spaces induced by kernel functions (Rasmussen and Williams (2005)).
These inﬁnite feature spaces make GPs rich probabilistic models for nonlinear regression, which are capable of expressing well-calibrated uncertainty quantiﬁcation (UQ) outside of the data distribution (OOD).
This UQ is crucial for tasks where sample eﬃciency and safety are critical, such as active learning, Bayesian optimization and model-based reinforcement learning.
Bayesian neural networks (BNN) have been widely criticized for their poor UQ (e.g. Foong et al.
; Osband et al.
(2018)), which impedes their performance in tasks for which Bayesian methods are preferred.
The neural linear model has several attractive qualities compared to alternative BNN approaches, such as Hamiltonian Monte Carlo (Neal (1995)), stochastic variational inference (Blundell et al.
Monte Carlo dropout (Gal and Ghahramani (2016)) and deep ensembles (Lakshminarayanan et al.
The predictive distribution is closed-form, rather than an implicit distribution, reducing the computational complexity of prediction.
As it is a Bayesian linear model, it also has fewer parameters for equivalent architectures compared to models with weight-space priors.
However, as this neural feature space is trained with type II maximum likelihood (ML), the features typically overﬁt to the data and struggle to express well-calibrated UQ when predicting OOD (L´azaro-Gredilla and Figueiras-Vidal, 2010).
Kernel feature spaces demonstrate less susceptibility to this overﬁtting, as fewer parameters are optimized.
Despite this weakness, as GPs struggle to scale as well as BNNs due to their nonparametric nature, we are motivated to enable NLMs that can match GP performance.
Therefore, we would wish to assess whether the ﬁnite neural feature space is even capable of matching or surpassing the performance of Gaussian process kernels.
To investigate this aspect of NLMs, we look at functional variational inference (fVI), which seeks to apply approximate inference over function distributions.
With the functional variational BNN (fBNN), Sun et al.
(2019) used a GP posterior as a functional prior for BNNs, which was then reﬁned on the data using the functional evidence lower bound objective (fELBO).
While this approach is essentially a form of empirical Bayes that requires two Bayesian models to be learned, we seek to use fVI to assess the representational power of NLMs in the context of GPs and other priors.
If NLMs are capable of matching or surpassing the performance of nonparametric GPs, then this motivates research into learning objectives for NLMs that encourages more kernel-like feature spaces.
2. Neural Linear Model
Intuitively, NLMs can be viewed as Bayesian linear regression in a projected feature space, where the projection is learned by a neural network.
Note that while the NLM can be easily deployed for multivariate regression, the following derivations (and later experiments) in this work focus on univariate targets for simplicity.
Let D =
( x i , y i ) } ni =1 be the observed data, x i ∈
R k , y i ∈ R , such that X ∈ R n × k and y ∈ R n .
Additionally, let φ ( · ; θ ): R k
→ R m be a feature space projection with parameters θ , φ i =
φ
( x i ; θ ) and Φ = (cid:2) φ (cid:62) 1 . . .
φ (cid:62) n (cid:3) ∈ R n × m , the matrix of vertically stacked row vectors.
It is common to add constant or linear terms to Φ to implicitly represent bias or identity terms.
However, for notational clarity, we ignore these terms and denote the projected feature space as R m .
A latent function f is modeled using Bayesian linear regression (Box and Tiao (1973);
Bishop (2006); Murphy (2012)) with weights β and additive, zero-mean, Gaussian noise (cid:15) with variance σ 2 , where
[Equation]
Placing a conjugate Gaussian prior N ( µ 0 , Λ − 1 0 ) over β results in a Gaussian posterior
N ( µ n , Λ − 1 n ) with an explicit Gaussian predictive distribution for query x ,
[Equation]
where µ n and Λ n are the mean vector and precision matrix of the posterior weight distribution.
The exact posteriors are detailed in Section A.
The observation noise σ 2 , prior weight parameters µ 0 and Λ 0 , and θ can be either set to constants or optimized jointly by maximizing the marginal likelihood.
With µ 0 = 0 , this model is equivalent to a Gaussian process with kernel k ( x , x (cid:48) ; θ ) =
φ ( x ; θ ) (cid:62) Λ − 1 0 φ ( x (cid:48) ; θ ) (Rasmussen and Williams (2005)).
This connection highlights that the NLM can be viewed as a distribution over functions p ( f ), in which f is restricted to the form of Equation (1).
3. Functional Variational Inference
Functional variational inference considers optimizing a parameteric posterior q ( f ) using the functional evidence lower bound objective (fELBO) through use of a functional KL divergence (fKL) and functional prior p ( f ),
[Equation]
Here, L ( D ) is a lower bound of the log-marginal likelihood log p ( D ), such that maximizing
L ( D ) translates to maximizing p ( D ).
In general, p ( f ) and q ( f ) can be arbitrary stochastic processes, i.e. distributions over functions, and the fKL between two stochastic processes (de G. Matthews et al.
(2016)) is a generalization of the regular KL divergence between two ﬁnite dimensional probability distributions.
However, the generalization to function space introduces obstacles in terms of tractability.
Roughly speaking, evaluation of the fKL requires integration in an inﬁnite-dimensional vector space, which is intractable as there is no inﬁnite-dimensional Lebesgue measure (Hunt (1992)).
Nonetheless, it is possible to approximate this quantity for Gaussian processes by averaging over a ﬁnite number of evaluations at s locations x j (Sun et al.
[Equation]
such that the quantity is concrete.
A key factor for a successful fKL approximation is the choice of { x j } .
Ideally, { x j } should be representative of the domain of f .
In general, it is impossible to entirely cover or reasonably represent the behavior of arbitrary distributions p ( f ) and q ( f ) over R k with any ﬁnite number of samples.
Therefore, it is necessary to select a bounded subset of the original domain as ‘domain of interest’, which can be challenging if the underlying data distribution is unknown.
4. Functional Priors in Practice
While the variational stochastic process q ( f ) can be modeled using any probabilistic function approximator, such as a NLM or any other Bayesian neural network, a beneﬁcial yet practicable choice for the prior stochastic process p ( f ) is less apparent.
White Gaussian noise (WGN), more speciﬁcally a Gaussian process with zero mean and isotropic covariance,
[Equation]
is a simple candidate which is often used by GPs (Rasmussen and Williams, 2005).
However, while GPs typically optimize the variance of this prior using type II ML, here it is ﬁxed.
To attain scale invariance across datasets, it is sensible to deﬁne this prior in whitened space.
Alternatively, it is possible to use a more informative prior, such as a Gaussian process posterior GP ( f | 0 , K ) whose kernel parameters were calibrated using the available training data (Sun et al.
; Shi et al.
The scalar k , vector k and matrix K are computed using the kernel, training data and function input x , so
[Equation]
Note that this choice of functional prior p ( f ) involves the training data D .
Therefore, it is not a pure prior in a Bayesian sense, but could be interpreted as a form of empirical Bayes (Morris (1983)), which may result in overﬁtting.
Lastly, it is also possible to use the fKL solely as the optimization objective, i.e. Equation (4) rather than Equation (3), so that q ( f ) approximates another stochastic process p ( f ), e.g. a Gaussian process posterior.
If p ( f ) adequately represents the observed data then q ( f ) will naturally ﬁt the data by minimizing the fKL while also inheriting other properties, such as extrapolation behavior and uncertainty quantiﬁcation.
In particular, this functional inference can transfer stochastic processes which are modeled by arbitrary probabilistic function approximators, e.g. GPs and NLMs, even if both models have considerably diﬀerent architectures.
This could enable model compression and distillation, which may be practical for nonparametric models like Gaussian processes.
The computational complexity of the training procedure for NLMs with the fELBO is of important practical concern and hence we discuss it here.
The prediction complexity is, regardless of the chosen training procedure, in O ( m 2 ) per prediction, assuming the feature-space dimension m to be the largest of all hidden-layer dimensions and neglecting the computation of the non-linearities of the network.
The total complexity during training of evaluating the fELBO (Equation (3)) is the sum of the complexities of evaluating the marginal likelihood and the fKL approximation (Equation (4)).
Under the previous assumptions, the complexity of the marginal likelihood objective for the NLM is in O ( nm 2 + m 3 ).
The complexity of computing the fKL approximation depends on the number of samples s and the types of involved stochastic processes.
Assuming that f p and f q are both Gaussian and already available, computation of the fKL term with s samples { x j } requires O ( sd 3 ) for d output dimensions, which simpliﬁes to O ( s ) for d = 1. Computing f q consists of evaluating the NLM’s predictive distribution, which is in O ( sm 2 ) for s samples. The complexity of computing f p depends on the chosen prior. The WGN prior is available in O (1) time, whereas the GP prior requires GP training once ( O ( n 3 )) and
GP prediction ( O ( sn 2 )) per objective evaluation.
[Table] [Figure] 5. Experimental Results This section discusses our nonlinear regression experiment, which aims to demonstrate the viability of combining NLMs with functional inference, leveraging the NLM’s explicit predictive distribution.
To this end, we replicated the experiments presented in Sun et al.
and evaluated the NLM on several UCI regression datasets.
The goal is to assess how the NLMs with functional priors compare to other Bayesian neural networks with weight space priors. In this setting, four diﬀerent optimization objectives, namely type II maximum likelihood with the conventional weight space prior ( § ), weight space ELBO ( ∗ ), functional ELBO ( ‡ ) and functional KL ( † ), were used. For the functional inference, we used white
Gaussian noise (WGN) and a Gaussian process (GP) as functional priors.
All models were implemented using PyTorch (Paszke et al.
For the GPs we additionally used GPyTorch (Gardner et al.
The NLM was implemented with µ 0 = 0 and a diagonal matrix with m distinct parameters to represent Λ 0 .
The observation noise parameter σ 2 , neural network weights θ and Λ 0 were learned via backpropagation.
To encode positive value constraints, σ 2 and Λ 0 were learned in log-space.
Following Sun et al.
(2019), NLMs used a single hidden layer with 50 units.
For optimization parameters, please refer to Appendix B.
Table 1 displays the results of our UCI regression benchmark experiments, comparing
NLMs to GPs and previous work.
In general, the NLM achieves competitive performance.
Although using the WGN prior results in underﬁtting for data with low observation noise (energy, yacht), it can match or surpass the predictive performance of a GP on data with high observation noise (boston, concrete, wine).
Applying fELBO inference with the GP prior, the NLM mostly outperforms the other Bayesian neural networks.
The fKL objective with the GP posterior has performance that does indeed match the GP across most datasets.
The NLM with fELBO training also generally outperforms the fBNN, apart from the boston dataset.
Given that the fBNN has approximately twice as many parameters, a more elaborate training procedure, and an implicit predictive distribution, this indicates that the NLM is a better function class for functional prior research.
6. Related Work
The study of functional priors for Bayesian inference is a burgeoning ﬁeld.
Our work builds oﬀ Sun et al.
(2019), who train a weight-space BNN with a GP posterior using a KL gradient estimator and mean-ﬁeld stochastic variational inference.
Shi et al.
(2019) in turn build oﬀ this, taking a mirror descent interpretation to allow for incremental training with minibatches and the Gaussian process prior.
As Gaussian processes are exact distributions over functions, sparse GPs may be viewed as approximate inference over functions (de G. Matthews et al.
, minimizing the fKL from its exact posterior via inducing points.
The noise contrastive prior (NCP) of Hafner et al.
(2019) is a similar idea where the training data is perturbed by random noise to serve as a ‘data prior’ for a BNN, in order to increase uncertainty estimation OOD.
While eﬀective empirically, the data prior is again akin to empirical Bayes as the prior is deﬁned by the data.
The practice of combining kernels in GPs has been translated to BNN architectures and activation functions, producing periodic and mixing phenomena in the network’s feature space for more expressive models (Pearce et al.
Variational implicit priors (Ma et al.
(2019)) use variational inference to deal with functional priors which you can only sample from, e.g. simulators, which provides the ﬂexibility of a broad range of complex processes to be adopted as priors.
7. Outlook
The goal of this work is to motivate further research into the NLM by demonstrating its potential in combination with functional variational inference.
Experiments have shown that this model is just as capable of representing a GP posterior as networks with weight priors, while being simpler and having closed-form predictive distributions.
Our results suggest that an appropriate functional prior could enable GP-like performance and uncertainty quantiﬁcation from NLMs, without requiring a GP posterior for training. Acknowledgments
Pascal Klink is funded by the DFG project PA3179/1-1 (ROBOLEAP).
Appendix A. Bayesian Last Layer Equations
In this section, we collectively state the equations for Bayesian linear regression (Box and
Tiao (1973); Bishop (2006); Murphy (2012)) with weights β and additive noise (cid:15) , where
[Equation]
Assuming known aleatoric uncertainty σ 2 (Bishop (2006))
[Equation]
a conjugate Gaussian prior over β results in a Gaussian posterior,
[Equation]
with an explicit Gaussian predictive distribution for output y and query x
[Equation]
The log-marginal likelihood can be written as
[Equation]
Appendix B. Optimization Parameters
Gradient-based optimization was performed using Adam (Kingma and Ba (2014)).
The GPs were trained until convergence for up to 3000 epochs at a learning rate of 1e-2.
The NLMs were trained at a learning rate of 1e-3, and the number of training epochs was determined by tracking the validation log-likelihood of 10% of the training data for up to 10000 epochs for ML training and up to 15000 epochs for all fVI variants.
The fKL was evaluated using 100 samples per epoch from a uniform distribution consistent with Sun et al. Appendix C. Experimental Results [Table] [Figure]