F UNCTIONAL V ARIATIONAL B AYESIAN N EURAL N ETWORKS
Anonymous authors
Paper under double-blind review
A BSTRACT
Variational Bayesian neural networks (BNN) perform variational inference over weights, but it is difﬁcult to specify meaningful priors and approximate posteriors in a high-dimensional weight space.
We introduce functional variational Bayesian neural networks (fBNNs), which maximize an Evidence Lower BOund (ELBO) deﬁned directly on stochastic processes, i.e. distributions over functions.
We prove that the KL divergence between stochastic processes is equal to the supremum of marginal KL divergences over all ﬁnite sets of inputs.
Based on this, we introduce a practical training objective which approximates the functional ELBO using ﬁnite measurement sets and the spectral Stein gradient estimator.
With fBNNs, we can specify priors which entail rich structures, including Gaussian processes and implicit stochastic processes.
Empirically, we ﬁnd that fBNNs extrapolate well using various structured priors, provide reliable uncertainty estimates, and can scale to large datasets.
I NTRODUCTION
Bayesian neural networks (BNNs) (Hinton & Van Camp, 1993; Neal, 1995) have the potential to combine the scalability, ﬂexibility, and predictive performance of neural networks with principled
Bayesian uncertainty modeling.
However, the practical effectiveness of BNNs is limited by our ability to specify meaningful prior distributions and by the intractability of posterior inference.
Choosing a meaningful prior distribution over network weights is difﬁcult because the weights have a complicated relationship to the function computed by the network.
Stochastic variational inference is appealing because the update rules resemble ordinary backprop (Graves, 2011; Blundell et al., 2015), but ﬁtting accurate posterior distributions is difﬁcult due to strong and complicated posterior dependencies (Louizos & Welling, 2016;
Sun et al., 2017; Zhang et al., 2017; Shi et al., 2017).
In a classic result, Neal (1995) showed that under certain assumptions, as the width of a shallow BNN was increased, the limiting distribution is a Gaussian process (GP).
Lee et al.
(2017) recently extended this result to deep BNNs.
However, the relationship of ﬁnite BNNs to GPs is unclear, and practical variational BNN approximations fail to match the predictions of the corresponding GP.
Furthermore, because the aforementioned analyses related speciﬁc BNN architectures to speciﬁc GP kernels, it’s not clear how one would design BNN architectures to match a given kernel structure.
Given the rich variety of structural assumptions that GP kernels can represent (Rasmussen & Williams, 2006; Lloyd et al., 2014; Sun et al., 2018), there remains a signiﬁcant gap in expressive power between BNNs and
GPs (not to mention stochastic processes more broadly).
In this paper, we propose to perform variational inference directly on the distribution of functions.
Speciﬁcally, we introduce functional variational BNNs (fBNNs), where a BNN is trained to produce a distribution of functions with small KL divergence to the true posterior over functions.
We prove that the KL divergence between stochastic processes 1 can be expressed as the supremum of marginal KL divergences at ﬁnite sets of points.
Based on this, we present a novel objective functional ELBO (fELBO).
Then we introduce a GAN-like minimax formulation and a VAE-like maximization formulation for functional variational inference.
To approximate the marginal KL divergence gradients, we adopt the recently proposed spectral Stein gradient estimator (SSGE) (Shi et al., 2018).
[Figure]
Our fBNNs make it possible to specify stochastic process priors which encode richly structured dependencies between function values.
This includes stochastic processes with explicit densities, such as Gaussian Processes which can model various structures like linearity, smoothness and periodicity (Lloyd et al., 2014; Sun et al., 2018).
In addition, we can also use stochastic processes with implicit densities, such as distributions over piecewise linear or piecewise constant functions.
Furthermore, in contrast with GP, fBNNs efﬁciently yield explicit samples of the function.
This enables fBNNs to be used in settings that require explicit minimization of sampled functions, such as
Thompson sampling (Thompson, 1933) or predictive entropy search (Hernández-Lobato et al., 2014;
Wang & Jegelka, 2017).
One desideratum of Bayesian models is that they behave gracefully as their capacity is increased (Rasmussen & Ghahramani, 2001).
Unfortunately, ordinary BNNs don’t meet this basic requirement: unless the asymptotic regime is chosen very carefully (e.g. Neal (1995)), BNN priors may have undesirable behavior as more units or layers are added.
Furthermore, larger BNNs entail more difﬁcult posterior inference and larger description length for the posterior, causing degeneracy for large networks, as shown in Figure 1.
Our fBNNs do not have this problem, since the prior is deﬁned directly over the space of functions, and therefore the BNN can be made arbitrarily large without changing the functional variational inference problem.
Hence, their predictions behave well as the capacity is increased, as shown in Figure 1.
Empirically, we demonstrate fBNNs generate reasonable extrapolations for both explicit periodic priors and implicit piecewise priors.
Furthermore, we show fBNNs outperform competing approaches on both small scale and large scale regression datasets.
Finally, we manifest fBNNs output reliable uncertainty estimations by obtaining state of the art performance on the contextual bandits benchmarks (Riquelme et al., 2018).
B ACKGROUND 2. V ARIATIONAL I NFERENCE FOR B AYESIAN N EURAL N ETWORKS
Given a dataset D = { ( x i , y i ) } ni =1 , a Bayesian neural network (BNN) is deﬁned in terms of a prior p ( w ) on the weights, as well as the likelihood p ( D| w ) .
Variational Bayesian methods (Hinton &
Van Camp, 1993; Graves, 2011; Blundell et al., 2015) attempt to ﬁt an approximate posterior q ( w ) to maximize the evidence lower bound (ELBO):
[Equation]
The most commonly used variational BNN training method is Bayes By Backprop (BBB) (Blundell et al., 2015), which uses a fully factorized Gaussian approximation to the posterior, i.e. q ( w ) =
N ( w ; µ , diag( σ 2 )) .
Using the reparameterization trick (Kingma & Welling, 2013), the gradients of
ELBO towards µ, σ can be computed by backpropagation, and then be used for updates.
Most commonly, the prior p ( w ) is chosen for computational convenience; common choices include independent Gaussian or Gaussian mixture distributions.
Other weight priors, including log-uniform priors (Kingma et al., 2015; Louizos et al., 2017) and horseshoe priors (Ghosh et al., 2018;
Louizos et al., 2017), have been proposed for particular purposes such as model compression and model selection.
However, the relationships of weight space priors to the functions computed by networks are difﬁcult to characterize.
2. S TOCHASTIC P ROCESSES
A stochastic process is deﬁned as a random function F : X → Y .
For any ﬁnite subset x 1: n , F can compute its marginal joint distribution over the function values F ( x 1: n ) .
For example, the Gaussian
Process has its marginal distributions as multivariate Gaussians, and the Student-t Process (Shah et al., 2014) has its marginal distribution as multivariate Student-t distributions.
However, directly representing a stochastic process is difﬁcult because X is typically inﬁnite.
Kolmogorov Extension Theorem (Øksendal, 2003) provides an alternative using ﬁnite dimensional marginal distributions.
Speciﬁcally, for a collection of joint distributions ρ x 1: n , Kolmogorov Extension Theorem states that, ρ deﬁnes a random process F such that ρ x 1: n is the marginal distribution of
F ( x 1: n ) if ρ satisﬁes the following two conditions:
Exchangeability For any permutation π of { 1 , · · · , n } , ρ
π ( x 1: n ) ( π ( y 1: n ))
=
ρ x 1: n ( y 1: n ) .
Consistency For any 1 ≤ m ≤ n , ρ x 1: m ( y 1: m ) =
(cid:82) ρ x 1: n ( y 1: n ) dy
m
+1: n .
Bayesian neural networks satisfy both exchangeablility and consistency, therefore it deﬁnes a stochastic process.
2. S PECTRAL S TEIN G RADIENT E STIMATOR (SSGE)
When applying Bayesian methods to modern probabilistic models, especially those with neural networks as components (e.g., BNNs and deep generative models), it is often the case that we have to deal with intractable densities.
Examples include the marginal distribution of a non-conjugate model (e.g., the output distribution of a BNN), and neural samplers such as Generative Adversarial
Networks (Goodfellow et al., 2014).
A shared property of these distributions is that they are deﬁned by a easy sampling process, despite the density is intractable.
This kind of distribution is called implicit distributions (Huszár, 2017).
Spectral Stein Gradient Estimator (SSGE) (Shi et al., 2018) is a recently proposed method for estimating the log density derivative function of an implicit distribution, only requiring samples from the distribution.
Speciﬁcally, given a continuous differentiable density q ( x ) , and a positive deﬁnite kernel k ( x , x (cid:48) ) in the Stein class (Liu et al., 2016) of q , they show
[Equation]
(2) where { ψ j } j ≥ 1 is a series of eigenfunctions of k described by the Mercer’s theorem: k
( x , x (cid:48) )
= (cid:80) j µ j ψ j ( x ) ψ j ( x (cid:48) ) .
Then the Nyström method (Baker, 1997; Williams & Seeger, 2001) is used to approximate the eigenfunctions ψ j ( x ) and their derivatives.
The ﬁnal estimator is given by truncating the sum in eq.
(2) and replacing the expectation by Monte Carlo estimates.
F UNCTIONAL V ARIATIONAL B AYESIAN N EURAL N ETWORKS 3. F UNCTIONAL E VIDENCE L OWER B OUND ( fELBO )
Given dataset D =
( X D , y D ) , variational inference in weight space is based on the equivalence between maximizing the ELBO and minimizing the KL divergence from the true posterior.
Now we introduce functional variational inference, where the functional ELBO (fELBO) is deﬁned with respect to a variational posterior over functions ( f : X → Y ) rather than over weights.
[Equation]
Speciﬁcally, we assume a stochastic process prior p ( f ) over functions.
This could be a Gaussian
Process, but one could also use stochastic processes without closed-form marginal densities, such as distributions over piecewise linear functions.
We consider a variational posterior q φ ( f ) ∈ Q deﬁned in terms of a neural network with stochastic weights and/or stochastic inputs.
Speciﬁcally, we sample a function from q by sampling a random noise ξ and reparameterize f ( x ) =
g φ ( x , ξ ) .
For example, the standard Bayesian neural networks with Gaussian weight distributions allow this reparameterization trick (Kingma & Welling, 2013).
Note that because a single vector ξ is shared among all input locations, it corresponds to randomness in the function , rather than observation variance; hence, the sampling of ξ corresponds to epistemic, rather than aleatoric, uncertainty (Depeweg et al., 2017).
In theorem 1, we prove that the functional KL divergence is equal to the supremum of marginal KL divergences over all ﬁnite sets of input locations X
= x 1: n , which we term measurement points .
Intuitively, this follows from the Kolmogorov Extension Theorem in section 2.2 which guarantees that we only need ﬁnite marginal distributions to deﬁne the stochastic process.
A full proof is given in Appendix A.
Theorem 1 (Functional KL Divergence) .
For two stochastic processes P, Q , the KL divergence is the supremum of marginal KL divergences over all ﬁnite subset of inputs x 1: n :
[Equation]
fELBO Using this characterization of the functional KL divergence, we rewrite the fELBO:
[Equation]
Hence, maximizing L ( q ) can be viewed as a two-player game analogous to a GAN (Goodfellow et al., 2014): the generator is maximizing L X ( q ) with respect to q , and the discriminator is minimizing
L X ( q ) with respect to X .
Interestingly, we can show that if measurement points contains all training examples, L ( q ) is a lower bound of the log marginal likelihood log p ( D ) .
A proof is in Appendix B.3.
Theorem 2 (Lower Bound) .
If X contains all training input locations X D , then
[Equation]
3. F UNCTIONAL V ARIATIONAL I NFERENCE
Now we present practical algorithms for functional variational inference.
Adversarial Functional Variational Inference Analogously to GANs 2 , L ( q ) is likely to be inﬁnite if the discriminator is completely unrestricted, because most stochastic process priors assign measure zero to the set of functions representable by a neural network (Arjovsky & Bottou, 2017).
However, by limiting the capacity of the discriminator, we obtain a meaningful minimax objective which forces the generator to produce samples which resemble the true posterior.
Most obviously, we can restrict the discriminator to choose measurement points of size N m .
Then we have the adversarial functional variational inference:
[Equation] [Equation]
To solve this minimax problem, we can adopt the similar inner and outer loop optimization like did in
GANs (Goodfellow et al., 2014).
In the inner loop, we minimize L X ( q ) with respect to X ; in the outer loop, we maximize L X ( q ) with respect to the generator q .
Sampling-Based Functional Variational Inference Adversarial functional variational inference plays a minimax game, which is difﬁcult to optimize and takes longer time for convergence (Goodfellow et al., 2014).
Here we present sampling-based functional variational inference as an alternative, which only jointly optimizes a same objective like VAE (Kingma & Welling, 2013).
Speciﬁcally, we replace the minimization in eq.
(6) with a sampling distribution c , and then maximize the expected
L X ( q ) under this distribution.
However, without the minimization, the objective will be inclined to overﬁt on training locations because of the log likelihood term.
We therefore let the measurement points X include random subset
X D s of training data and denote X M =
X \ X D s .
Then the sampling-based functional variational inference becomes,
[Equation]
where X M are M points independently drawn from c .
Denote f M , f D as the function values for
X M , X D , respectively.
By theorem 2, if X contains all of X D ,
[Equation]
Maximizing L X M , X D ( q ) is equivalent to minimizing the KL divergence from the true posterior on points X M , X D .
Therefore, another interpretation of sampling-based functional variational inference is that we want better posterior approximation on more "interesting" locations weighted by c .
KL DIVERGENCE G RADIENTS
Computing fELBO requires to compute the likelihood term and the KL divergence.
Although the likelihood term is tractable, the KL divergence term remains intractable because we don’t have the explicit formula for variational posterior q φ ( f X | X ) .
Note ∇ φ KL[ q φ ( f X ) (cid:107) p ( f X )] is given by
[Equation]
It is easy to check that the ﬁrst term in eq.
(9) is zero (Roeder et al., 2017).
Besides that, using parametric variational posterior, the gradients ∇ φ f X can be computed easily by backpropagation.
All we left are the log-density derivatives ∇ f log q ( f X ) and ∇ f log p ( f X ) .
As introduced in Sec 2.3,
SSGE offers us an principled way to estimate the score function at both in-distribution samples and out-of-distribution samples.
Therefore, we can apply SSGE to estimate the two derivatives above.
Additionally, for priors with explicit densities like Gaussian Processes (Rasmussen, 2004) 3 , Student-t process (Shah et al., 2014) and Wishart Process (Dawid, 1981), ∇ f log p ( f X ) is analytic.
3. T HE A LGORITHM
Throughout this paper, we mainly investigate sampling-based functional variational inference.
Now we present the whole algorithm for fBNNs in Algorithm 1.
Because the log likelihood term has unbiased mini-batch estimations, we estimate L X ( q ) with mini-batch D s for fast computation.
Speciﬁcally, we are optimizing (cid:88)
[Equation]
which is also a proper lower bound of log p ( D s ) .
Require: Dataset D , sampling distribution c , variational posterior g ( · ) , prior p (explicit or implicit).
1: while φ not converged do 2:
X M ∼ c ; D S ⊂ D
(cid:46) sample measurement points 3: f i =
g
( { X M , X D S } , ξ i ; φ ) , i = 1 · · · k.
(cid:46) sample k function values 4: g 1 = 1 k (cid:80)
i (cid:80) ( x,y ) ∇ φ log p ( y | f i ( x ))
(cid:46) compute log likelihood gradients 5: g 2 = SSGE( p, f 1: k ) (cid:46) estimate KL gradients 6: φ ← Optimizer ( φ, g 1 − g 2 ) (cid:46) update the parameters 7: end while
3. C ONSISTENCY
We can enumerate all ﬁnite subsets of for neither adversarial nor sampling-based method.
However, we show desirable consistency results under assumption of Gaussian Processes.
Theorem 3.
When p and q are both Gaussian Processes, q ’s family is expressive enough and measurement points cover all training locations, | X | ≥ |D| +2 and supp( c )
=
X , then the following statements are equivalent: • Adversarial functional variational inference attains its optimum. • Sampling-based functional variational inference attains its optimum.
The variational posterior process and the true posterior process are identical.
R ELATED WORK
Bayesian neural networks Variational inference was ﬁrst applied to neural networks by Peterson (1987) and Hinton & Van Camp (1993).
More recently, Graves (2011) proposed a practical method for variational inference with fully factorized Gaussian posteriors which used a simple (but biased) gradient estimator.
Improving on that work, Blundell et al.
(2015) proposed an unbiased gradient estimator using the reparameterization trick of Kingma & Welling (2013).
There has also been much work (Louizos & Welling, 2016; Sun et al., 2017;
Zhang et al., 2017) on modeling the correlations between weights using more complex Gaussian variational posteriors.
Some non-Gaussian variational posteriors have been proposed, such as Multiplicative Normalizing Flows (Louizos & Welling, 2017) and implicit distributions (Shi et al., 2017).
Neural networks with dropout were also interpreted as
BNNs (Gal & Ghahramani, 2016; Gal et al., 2017).
However, all these methods place priors over the network parameters.
Often, identical independent Gaussian priors are placed over the weights for convenience and the prior assumptions are not well-understood.
Structured Priors There have been other recent attempts to train BNNs with structured priors.
Flam-Shepherd et al.
(2017) trained a BNN prior to mimic a GP prior, but their method still required variational inference in weight space.
Noise Contrastive Priors (Hafner et al., 2018) are somewhat similar in spirit to our work in that they use a random noise prior in the function space.
However, instead of doing proper Bayesian inference, the prior is incorporated by adding a regularization term to the weight-space ELBO, and is not rich enough to encourage extrapolation and pattern discovery.
Neural Processes (Garnelo et al., 2018) tries to model any conditional distribution given arbitary data points, which struggles to scale to even medium dimensions like 10.
Variational Implicit Process (Ma et al., 2018) speciﬁes BNN priors and approximately ﬁt a Gaussian Process for the posterior, thus the mis-speciﬁed prior problem also applies here.
Scalable Gaussian Process Gaussian processes are difﬁcult to apply exactly to large datasets since the computational requirements scale as O ( N 3 ) time, and as O ( N 2 ) memory, where N is the number of training cases.
Multiple approaches have been proposed to reduce the computational complexity.
However, sparse GP methods (Lázaro-Gredilla et al., 2010; Snelson & Ghahramani, 2006; Titsias, 2009; Hensman et al., 2013; 2015) still suffer for very big dataset, while random feature method (Rahimi & Recht, 2008; Le et al., 2013) and KISS-GP (Wilson & Nickisch, 2015) are only applicable to speciﬁc kernels.
Furthermore, these methods don’t solve the problem of obtaining explicit samples of the function.
E XPERIMENTS
In this section, we present some experiments testing fBNNs’ ability to extrapolate using a variety of structural motifs.
We also evaluate their predictive performance on regression benchmarks and their ability to guide exploration in the contextual bandit setting.
In all of our experiments, the variational posterior is represented as a stochastic neural network with independent Gaussian distributions over the weights, i.e. q ( w ) =
N ( w ; µ , diag( σ 2 )) .
4 We always
., in principle,
4      (a) BBB    4      (b) FBNN    4      (c) GP used the ReLU activation function unless otherwise speciﬁed.
Denote xmin , xmax as dimensionwise minimal and maximal input locations, xd = xmax − xmin .
We set the sampling distribution c =
U [xmin − xd / 2 , xmax
+ xd / 2] unless speciﬁc notice.
For experiments where we used GP priors, we ﬁrst ﬁt the GP hyperparameters to maximize the marginal likelihood on subsets of the training examples, and then ﬁxed those hyperparameters to obtain the prior for the fBNNs.
5. E XTRAPOLATION U SING S TRUCTURED P RIORS
Making sensible predictions outside the range of the observed data requires exploiting the underlying structure.
In this section, we consider some toy datasets where fBNNs are able to use structured priors to make sensible extrapolations.
5.1. L EARNING P ERIODIC S TRUCTURES
Gaussian processes can model periodic structure using a periodic kernel plus a RBF kernel:
[Equation] [Equation]
) where p determines the period.
In this experiment, we consider 20 inputs randomly sampled from the interval [ − 2 , − 0 .
5] ∪ [0 . 5 , 2] , and targets y which are noisy observations of a periodic function: y = 2 ∗ sin(4 x ) + (cid:15) with (cid:15) ∼ N (0 , 0 . 04) .
We use a BNN with ﬁve hidden layers, each withunits.
The inputs and targets were normalized to have zero mean and unit variance.
For all methods, the observation noise variance was set to the true value.
We compared our method with Bayes By
Backprop (BBB) (Blundell et al., 2015) (with a spherical Gaussian prior on w ) and Gaussian Processs.
We used the trained GP as the prior of our fBNN.
In each iteration, measurement points included all training examples, plus 40 points randomly sampled from [ − 5 , 5] .
We used a training budget of 20,000 iterations, and annealed the weight of the KL term from 0 to 1 linearly.
As shown in Fig.
2, BBB fails to ﬁt the training data, let alone recover the periodic pattern (as we would expect, since its prior does not encode periodic structure).
The GP (as usual) ﬁt the data well and extrapolated the period structure beyond the range of the data; we can view the GP predictions as a sort of gold standard.
The fBNN was able to extrapolate similarly to the GP, although it predicted slightly smaller uncertainty.
Interestingly, despite the ReLU activation functions, the fBNN made smooth predictions.
5.1. I MPLICIT P RIORS
Since the fBNN is based on implicit variational inference, the priors need not have convenient expressions for the marginal densities, i.e. we need not limit ourselves to GPs.
In this section, we consider two implicit priors: a distribution over piecewise constant functions, and a distribution over piecewise linear functions.
Appendix D.2 provides the detailed explanation on the sampling process.
Some samples from these priors are shown in Figure 3.
In each run of the experiment, we ﬁrst sampled a random function from the prior, and then sampled 20 locations from [0 , 0 . 2] and 20 locations from [0 . 8 , 1] , giving a training set of 40 data points.
The standard deviation of observation noise is 0.02.
In each iteration, measurement points included all training examples, plus 40 points randomly sampled from [0 , 1] .
We used a fully connected network
C o n s t a n t 0.
L i n e a r 0.
with 2 hidden layers of 100 units, and tanh activations.
(For this experiment, we used tanh rather than
ReLU because ReLU would have been unfairly well-suited to the non-smooth priors.)
The network was trained for 20,000 iterations.
We show the predictive samples and variances of three different runs in Figure 3.
As shown in Figure 3, the fBNN captured the basic structures for both piecewise constant priors and piecewise linear priors, although the posterior samples did not seem to capture the full diversity of possible explanations of the data.
Interestingly, even though the tanh activation function encourages smoothness, the network learned to generate functions with sharp transitions.
5. P REDICTIVE P ERFORMANCE
To compare with previous work on predictive performance, we next evaluated fBNN in regression experiments with publicly available datasets.
5.2. S MALL S CALE D ATASETS
Following previous work (Zhang et al., 2017), we ﬁrst experiment with standard regression benchmark datasets from the UCI collection (Asuncion & Newman, 2007).
In particular, we only used the datasets with less than 2000 data points so that we can ﬁt GP hyperparameters using marginal likelihood.
The datasets were randomly split into training and validation sets, comprising 90% and 10% of the data respectively.
This splitting process was repeated 10 times to reduce randomness.
[Table]
For all datasets
[Table]
we used networks with one hidden layer of 50 hidden units.
For all experiments, we ﬁrst ﬁt GP hyper-parameters using marginal likelihood with a budget of 10,000 iterations.
In each experiment, the fBNN was trained for 2,000 epochs.
And in each iteration, measurement points included 20 training examples, plus 5 points randomly sampled from the data range.
We compared our method with Bayes By Backprop (BBB) (Blundell et al., 2015) and Noisy Natural
Gradient (NNG) (Zhang et al., 2017).
In accordance with (Zhang et al., 2017), we report the standard metrics including root mean square error (RMSE) and validation log-likelihood.
The results are summarized in Table 1.
On most datasets, our fBNNs outperformed both BBB and NNG, sometimes by a signiﬁcant margin.
[Table] [Table]
L ARGE S CALE D ATASETS
Observe that fBNNs are naturally scalable to large datasets because they access the data only through the expected log-likelihood term
[Table]
which can be estimated stochastically.
In this section
[Table]
we verify this experimentally.
We compared fBNNs, BBB and BP with large scale UCI datasets, including Naval ,
Protein Structures , Video Transcoding (Memory, Time) , GPU kernel performance and YearPrediction .
We randomly split the datasets into 80% training, 10% validation
[Table]
and 10% test.
We used the validating set to select the hyperparameters and perform early stopping.
Both methods were trained for 80,000 iterations.
5 We used 1 hidden layer with 100 hidden units for all datasets.
For the prior of the fBNNs, we use a GP with Neural Kernel Network (NKN) kernels as used in Sun et al.
We note that the GP hyperparameters are ﬁtted using minibatches of size 1000 with 10000 iterations.
In each iteration, measurement sets consist of 500 training samples and 5 or 50 points from the sampling distribution c , tuned by validation performacne.
We ran each experiment 5 times, and report the mean and standard deviation in Table 2.
5. C ONTEXTUAL B ANDITS
[Table]
One of the most important applications of uncertainty modeling is to guide exploration in settings such as bandits, Bayesian optimization (BO)
[Table]
and reinforcement learning.
In this section, we evaluate fBNNs on a recently introduced contextual bandits benchmark (Riquelme et al., 2018).
In contextual bandits problems, the agent tries to select the action with highest reward given some input context.
Because the agent learns about the model gradually, it should balance between exploration and exploitation to maximize the cumulative reward.
Thomposon sampling (Thompson, 1933) has shown to be extremly promising by keeping a posterior of the model and selecting the action in proportion to the posterior distribution.
Therefore, the key problem of Thompson sampling is to learn a reliable
Bayesian uncertainty model.
We compared our fBNNs with the algorithms benchmarked in (Riquelme et al., 2018).
We used the multi-task GP of Riquelme et al.
as the prior for the fBNNs.
We run the experiments for all algorithms and tasks using the default settings open sourced by Riquelme et al.
For fBNNs, we also kept the same setting, including batchsize, training epochs and training frequency.
Measurement sets consisted of training batches, combined with 10 points sampled from c .
We ran each experiment 10 times and report the mean and standard derivation in Table 3.
Similarly to
Riquelme et al.
(2018), we also report the mean rank and mean regret.
As shown in Table 3, fBNNs outperformed other methods by a wide margin.
Additionally, the fBNNs maintained consistent performance with various network sizes.
By comparison, BBB suffered signiﬁcant performance degeneration when the hidden size was increased from 50 to 500.
This is consistent with our hypothesis that functional variational inference can gracefully handle networks with high capacity.
We also show a BO experiment in App C.2, which also needs reliable uncertainty.
C ONCLUSIONS
In this paper we investigate the variational inference between stochastic processes, i.e., distributions of functions.
We prove the KL divergence between stochastic processes equals to the supremum of
KL divergence for marginal distribution over all ﬁnite input subsets.
Then we present functional variational Bayesian neural networks (fBNNs).
We show our fBNNs enable to specify priors entailing rich structures, including both explicit stochastic process like GP and implicit stochastic process like piecewise constant functions.
Empirically, we demonstrates fBNNs extrapolate well over various structures, estimates reliable uncertainties and scales to large datasets.
A S TOCHASTIC P ROCESSES A. B ASIC M EASURE T HEORY N OTATIONS
π -system A π -system is a collection of subsets that is closed under ﬁnitely many intersections.
σ -algebra generated by cylinder sets Suppose
[Equation]
is a set of real-valued functions.
Let B ( R ) denote the Borel subsets of R .
A cylinder subset of X is a ﬁnitely restricted set deﬁned as
[Equation]
Each
[Equation]
is a π -system that generates a σ -algebra Σ t 1 , ··· ,t n .
Then the family of subsets
[Equation]
is an algebra that generates the cylinder σ -algebra for X .
Deﬁnition 1 (KL divergence (Gray, 2011)) .
Given two probability measures q and p on measurable space (Ω , B ) , the KL divergence of q with respect to p is deﬁned as
[Equation]
(12) where the supremum is over all ﬁnite measurable partitions Q of Ω .
A. P USH F ORWARD AND K OLOMOGROV E XTENSION T HEOREM
Deﬁnition 2 (Pushforward measure) .
Given probability spaces ( X, Σ X , µ ) and ( Y, Σ Y , ν ) , we say that measure ν is a pushforward of µ if ν ( A ) =
µ ( T − 1 ( A )) for a measurable T : X → Y and any
A ∈ Σ Y , denoted by ν = µ ◦ T − 1 .
Deﬁnition 3 (Canonical projection map) .
Let T be an arbitrary index set, and { (Ω t , F t ) } t ∈ T be some collection of measurable spaces.
For each subset J ⊂ I ⊂ T , deﬁne Ω J = (cid:81)
t ∈ J Ω t .
We call π I
→ J the canonical projection map from I to J if
[Equation]
Theorem 4 (Kolmogorov extension theorem) .
Let T be an arbitrary index set.
(Ω , F ) is a standard measurable space, whose product space on T is (Ω T , F T ) .
For each ﬁnite subset I ⊂ T , suppose we have a probability measure µ I on Ω I that satisﬁes the following compatibility relation: For each ﬁnite subset J ⊂ I ⊂ T , we have
[Equation]
Then there exists a unique probability measure µ on Ω T such that ∀ ﬁnite subset I ⊂ T ,
[Equation]
In Gaussian processes, µ is a Borel measure on a separable Hilbert space ( H , B ( H )) .
µ I is a marginal
Gaussian measure at any ﬁnite set (i.e., I ) of input positions (Mallasto & Feragen, 2017).
A. P ROOF
We ﬁrst introduce several lemmas for proving the theorem.
Lemma 5.
[Equation]
is a π -system.
And σ ( G X ) generates the same cylinder σ -algebra of X as F X .
Deﬁnition 4.
Let H be a set of functions, deﬁne its T -set τ
( H ) as
[Equation]
Here × represents product-algebra.
Intuitively, τ ( H ) represents all timesteps that restricts the functions in H .
Or we can say if some t doesn’t belong to τ ( H ) , then function values of t can be any number in real line R .
Lemma 6.
For any set G ∈ σ ( G X ) , its T -set τ ( G ) is countable.
Proof.
We build the sigma-algebra σ ( G X ) through an iterative process.
We begin by H
=
G X , in each iteration, let
[Equation]
Recursively applying this extension, we can generate σ ( G X ) =
H ∞ .
Because τ ( H c ) =
τ ( H ) , and countable union of countable sets is still countable.
We know the two updates above keep the countablility of its T -set.
Therefore, ∀ G ∈ σ ( G X ) , τ ( G ) is countable.
Theorem 7.
For two stochastic processes P, M on functional measurable space (Ω T , F T ) , the KL divergence between P and M ,
[Equation]
Where the supremum is over all ﬁnite subsets, and P T f , M T f represent the marginal distribution of
P, M on T f , respectively.
Proof. Sketches
The cylinder σ -algebra deﬁnes a σ -algebra σ ( G X ) on the function space.
By Lemma 5. 2.
Each set in the σ ( G X ) is only related to countable number of timesteps.
By Lemma 6. 3.
Any ﬁnite measurable partitions are related to countable number of timesteps, thus correspond to the marginal distribution over a countable subset by an inclusion map.
Therefore the KL between stochastic processes reduces to the KL over marginal distributions of countable subsets. 4.
KL over marginal distributions of countable subsets can be represented as supremum of KL over marginal distributions over its ﬁnite subsets.
By Deﬁnition 1,
[Equation]
where the sup is over all ﬁnite measurable partitions of the function space, denoted by Q Ω T :
[Equation]
By Lemma 6, τ ( Q ( i )Ω T ) is countable.
Therefore, T c :=
[Equation]
Consider the canonical projection mapping π T → T c , which induces a partition on Ω T c , denoted by
Q Ω T c :
[Equation]
The pushforward measure deﬁned by this mapping is
[Equation]
Then we have
[Equation]
Denote F ( T c ) as the collection of all ﬁnite subsets of T c .
For any ﬁnite set T j ∈ F ( T c ) , we denote P T j as the pushforward of P on ﬁnite set T j .
According to Kolmogorov Extension Theorem (Theorem 4), P T j is the marginal distribution of P at T j .
Because T c is countable, we have
[Equation]
Note for any ﬁnite subset T j , we can build a ﬁnite measurable partition Q Ω T such that its T c equals to T j .
Let Ω 1 , Ω 2 be any two-element partition of Ω .
In fact, we can let Q Ω T be a partition with 2 | T j | subsets, that each subset is a cylinder subset with timesteps in T j and for each timestep, the function values belong to Ω 1 or Ω 2 .
Therefore, we have shown the set { T j | T j ∈ F ( T c ) , T c } contains all ﬁnite subsets.
On the other hand, every element in { T j | T j ∈ F ( T c ) , T c } is a ﬁnite subset, we know { T j | T j ∈
F ( T c ) , T c } equals to the collection of all ﬁnite subsets.
Therefore we have proven the theorem.
B A DDITIONAL T HEORECTICAL R ESULTS B. KL D IVERGENCE BETWEEN C ONDITIONAL S TOCHASTIC P ROCESSES
In this section
[Equation]
we give an example of computing the KL divergence between two conditional stochastic processes.
Consider two datasets D 1 , D 2 , the KL divergence between two conditional stochastic processes is
[Equation]
Therefore, the KL divergence between these two stochastic processes equals to the KL divergence between marginal distributions on the observed dataset D 1 ∪D 2 .
In terms of D 2 = ∅ , p ( f |D 2 ) = p ( f ) , this shows that the KL divergence between posterior process and prior process are the KL divergence on the marginal distributions of observed locations.
This also justiﬁed our usage of M measurement points in the adversarial functional VI and samplingbased functional VI of section 3.
B. O PTIMIZE GP H YPERPARAMETERS USING M INI -B ATCH
Consider multivariate Gaussian distribution,
[Equation]
Its derivative towards K θ is − 12 K − 1 θ ( K θ − yy (cid:62) )
K − 1 θ , which attains zero only at K (cid:63) =
yy (cid:62) .
Therefore we have
[Equation]
Gaussian Process optimization maximizes data likelihood log p ( y | 0 , K θ ) .
According to eq.
(24), its maximum is obtained at K (cid:63) =
yy (cid:62) .
In contrast, mini-batch GP optimization maximimzes expected subset likelihood E S (cid:2) log p ( y S | 0 , K Sθ ) (cid:3) , with S being a subset.
According to eq.
(24), log p ( y S | 0 , K Sθ ) ≤ log p ( y S | 0 ,
y S y (cid:62) S ) .
Therefore, its maximum is also obtained at K (cid:63) =
yy (cid:62) .
Based on the results above, for expressive enough kernels, mini-batch optimized GP reaches the same maximum as full-batch optimized GP, which justiﬁes using mini-batch MLE for hyperparameter optimization. B. P ROOF FOR E VIDENCE L OWER B OUND This section provides proof for theorem 2. Let X M =
X
\ X D be measurement points which aren’t in the training data.
Proof. [Equation] [Equation] B. C ONSISTENCY FOR G AUSSIAN P ROCESS
This section provides proof for consistency in theorem 3.
Proof.
By the assumption that both q ( D ) and p ( f |D ) are Gaussian Processes, [Equation]
[Equation]
where m and k denote the mean and covariance functions, respectively.
In this theorem, we also assume the measurement points cover all training locations, based on theorem 2, we have
[Equation] Therefore, when the variational posterior process reaches its optimum, it must have
KL[ q ( f X ) (cid:107) p ( f X |D )][-= 
-]
= 0 at any X
= [ x 1 , . . .
, x N m ] (cid:62) ∈ X N m , which is
[Equation]
(30) where m ( X ) =
[ m ( x 1 ) , . .
, m ( x N m )] (cid:62) , and K ij = k
( x i , x j ) .
Given N m ≥ 2 , then for ∀ 1 ≤ i < j ≤ N m , we have m p ( x i )
=
m q ( x i ) , and k p ( x i , x j ) =
k q ( x i , x j ) .
Remember that in sampling-based functional variational inference, supp( c ) = X , M ≥ 2 , thus X covers all ( x , x (cid:48) ) pairs, thus we have
[Equation]
[Figure]
Remember that in adversarial functional variational inference, X is chosen to minimize L X ( q ) .
If it reaches optimum log p ( D ) , then for all ( x , x (cid:48) ) pairs, the same mean and covariance also holds. Thus again we have
[Equation]
Because GPs are uniquely determined by their mean and covariance functions, we arrive at the conclusion.
C A DDITIONAL E XPERIMENTS C. T IME - SERIES E XTRAPOLATION
Besides the toy experiments, we would like to examine the extrapolation behavior of our method on real-world datasets.
Here we consider a classic time-series prediction problem concerning the concentration of CO 2 in the atmosphere at the Mauna Loa Observatory, Hawaii (Rasmussen &
Williams, 2006).
The training data is given from 1958 to 2003 (with some missing values).
Our goal is to model the prediction for an equally long period after 2003 (2004-2048).
In Figure 4 we draw the prediction results given by BBB, fBNN, and GP.
We used the same BNN architecture for BBB and fBNN: an ReLU network with 2 hidden layers, each with 100 units, and the input is a normalized year number augmented by its sin transformation, whose period is set to be one year.
This special design allows both BBB and fBNN to ﬁt the periodic structure more easily.
Both models are trained for 30k iterations by the Adam optimizer, with learning rate 0 .
01 and batch size 20.
For fBNN the prior is the same as the GP experiment, whose kernel is a combination of RBF, RBF ×
PER (period set to one year), and RQ kernels, as suggested in Rasmussen & Williams (2006).
Measurement points include 20 training samples and 10 points sampled from U [1958 , 2048] , and we jointly train the prior
GP hyperparameters with fBNN.
In Figure 4 we could see that the performance of fBNN closely match the exact prediction by GP.
Both of them give visually good extrapolation results that successfully model the long-term trend, local variations, and periodic structures.
In constrast, weight-space prior and inference (BBB) neither captures the right periodic structure, nor does it give meaningful uncertainty estimates.
C. B AYESIAN O PTIMIZATION
In this section, we adopt Bayesian Optimization to explore the advantage of coherent posteriors.
Speciﬁcally, we use Max Value Entropy Search (MES)
(Wang & Jegelka, 2017), which trys to maximize the information gain about the minimum value y (cid:63) ,
[Equation]
Where φ and Ψ are probability density function and cumulative density function of a standard normal distribution, respectively.
The y (cid:63) is randomly sampled from the posterior of function minimas and γ y (cid:63) ( x ) = µ t ( x ) − y (cid:63) σ t ( x
[Equation]
With a probabilistic model, we can compute or estimate the mean µ t ( x ) and standard deviation σ t ( x ) .
However, to compute the MES acquisition funtion, samples y (cid:63) of function minimas are required as [Figure] well, which leads to bigger difﬁculties.
Typically when we model the data with a GP, we can get the posterior on a speciﬁc set of points but we don’t have access to the extremes of the underlying function.
In comparison, if the function posterior is represented parameterically, we can perform gradient decent easily and optimize for approximating function minimas.
We use 3-dim functions sampled from some Gaussian Process prior for bayesian optimization.
Concretely, we experiment with samples from RBF, ArcCosine and Matern12 kernels.
We compare three parametric approaches: fBNN, BBB and Random Feature (Rahimi & Recht, 2008).
For fBNN, we use the true kernel as functional priors.
In contrast, ArcCosine and Matern12 kernels do not have simple explicit random feature expressions, therefore we use RBF random features for all three kernels.
When looking for minimas, we sample 10 optimals y (cid:63) .
For each minima, we perform gradient descent along the parametric function posterior with 30 different starting points.
We usedimensions for random feature.
We use network with 5 × 100 for fBNN.
For BBB, we select the network within 1 × 100 , 3 × 100 .
Because of the similar issue in Figure 1, using larger networks won’t help for BBB.
We use batch size 30 for both fBNN and BBB.
The measurement points contain 30 training points and 30 points uniformly sampled from the known input domain of functions.
We train fBNN and BBB for 20000 iterations and anneal the coefﬁcient of log likelihood term linearly from 0 to 1 for the ﬁrst 10000 iterations.
The results with 10 runs are shown in Figure 5.
As seen from Figure 5, fBNN and Random feature outperform BBB by a large margin on all three functions.
We also observe fBNN performs slightly worse than random feature in terms of RBF priors.
Because random feature method is exacly a GP with RBF kernel asympotically, it sets a high standard for the parametric approaches.
In contrast, fBNN outperforms random feature for both
ArcCosine and Matern12 functions.
This is because of the big discrepency between such kernels and
RBF random features.
Because fBNN uses true kernels, it models the function structures better.
This experiment hightlights a key advantage of fBNN, that fBNN can learn parametric function posteriors for various priors.
C. T EXTURE E XTRAPOLATION
[Figure]
In this section we perform a texture extrapolation experiment.
With a texture image, texutre extrapolation tries to recover the missing region based on the observed region in some image.
We use a image of 224 × 224 , and the central 60 × 80 region is removed from the image.
Treating each pixel as a data point, this forms 45376 training examples andtest examples.
Because the training size is too big to be afforded by a standard Gaussian Process.
Previous GP methods (Wilson et al., 2014; Sun et al., 2018) investigate the additive structure of such patterns and take advantage of the Kronecker structure of the kernel matrix.
However, such kronecker structure is only applicable to limited structures.
This experiment adopts fBNN with a GAN-like generator to perform the recovery.
We use a GP prior with the kernel used in (Sun et al., 2018).
We ﬁrstly train the prior hyperparameters with GP regression with batch size 300.
For the variational network, we use a GAN-like generator to generate the whole image starting from a 40 dimensional U [0 , 1] noise.
The network arranges as 40-(64*7*7)-(32*28*28)-(16*112*112)-(8*224*224)-(1*224*224).
We use ﬁlter size 5 and leaky
ReLU activation.
Measurement points include 800 training points and 200 points sampled from the missing region (FBNN w. missing).
We use learning rate 0 .
003 to train fBNN for 40000 iterations and anneal the kl coefﬁcient from 0 to 1 linearly for the ﬁrst 20000 iterations.
We also evalute a setting where measurement points only contain training regions (FBNN w.o. missing).
The results are shown in Figure 6.
As shown in Figure 6, fBNN correctly recovers the underlying structure and extrapolates to the missing region.
However, if measurement points do not contain missing regions, fBNN ﬁts training set perfectly but fails to extrapolate.
This highlights the extrapolation effect of measurement points.
D I MPLEMENTATION D ETAILS D.
I NJECTED
For Gaussian Process prior, p ( f X ) is multivariate Gaussian distribution, which has an explicit formula.
Therefore, we can compute the gradients ∇ f log p
φ ( f X ) analytically.
In practice, we found that the GP kernel matrix suffers from statbility issues.
To stabilize the gradients computation, we propose to inject a small amount of Gaussian noise on the function values, i.e., to instead estimate the gradients of ∇ φ KL[ q φ ∗ p
γ (cid:107) p ∗ p γ ] , where p γ = N (0 , γ 2 ) is the noise distribution.
This is like the instance-noise trick that is commonly used for stabilizing GAN training (Sønderby et al., 2016).
Note that injecting the noise on the GP prior is equivalent to have a kernel matrix
K + γ 2 I , which have more stable properties.
Beyond that, injecting the noise on the parametric variational posterior doesn’t affect the reparameterization trick either.
Therefore all the previous estimation formulas still applies.
D. I MPLICIT P RIORS
Our method is applicable implicit priors.
We experiment with piecewise constant prior and piecewise linear prior.
Concretely, we randomly generate a function f : [0 , 1] → R with the speciﬁc structure.
To sample piecewise functions, we ﬁrst sample n ∼ Poisson(3 . )
, then we have n + 1 pieces within [0 , 1] .
We uniformly sample n locations from [0 , 1] as the changing points.
For piecewise constant functions, we uniformly sample n + 1 values from [0 , 1] as the function values in each piece; For piecewise linear functions, we uniformaly sample n + 1 values for the values at ﬁrst n + 1 locations, we force f (1) = 0 . .
Then we connect together each piece by a straight line.