VIB is Half Bayes
Abstract
In discriminative settings such as regression and classiﬁcation there are two random variables at play, the inputs X and the targets Y .
Here, we demonstrate that the Variational
Information Bottleneck can be viewed as a compromise between fully empirical and fully
Bayesian objectives, attempting to minimize the risks due to ﬁnite sampling of Y only.
We argue that this approach provides some of the beneﬁts of Bayes while requiring only some of the work.
1. Introduction
Big models, big data, and maximum likelihood training are a proven recipe for learning powerful and generalizable neural network models.
But training such large models on small data results in overﬁtting and poor performance.
How can we achieve good performance from limited data?
Bayesian inference presents one such mechanism.
Bayesian inference can be seen as minimizing a PAC-style upper bound on generalization performance from ﬁnite data (Morningstar et al., 2020; Masegosa, 2020;
Germain et al., 2016).
However, exactly performing
Bayesian inference is costly, requiring careful tuning of MCMC methods or expressive variational distributions to match the posterior (Betancourt, 2018; Zhang et al., 2018).
Here, we show that training on multiple outputs Y for each input X can be beneﬁcial, and derive a training objective which provides these beneﬁts without actually having to collect multiple outputs for each input.
The resulting objective matches the Variational
Information Bottleneck (VIB) (Alemi et al., 2016), and provides a tractable alternative to
Bayesian inference that loses some of the guarantees but retains much of the qualitative and quantitative performance.
2. Preliminaries
Consider training a neural network with parameters θ to output a stochastic representation z ∼ q
( Z | x, θ ) for each input x ∈ X .
From the representation, we can predict the target y with a ﬁxed (parameter-free) classiﬁer or regressor p ( y | z ).
At test time, we will form the predictive distribution q ( y | x, θ ) =
(cid:82) d z q ( z | x, θ ) p ( y | z ).
Assuming the true data comes from some joint distribution
ν
( X, Y ), we aim to learn a predictive distribution that is as close as possible to the true conditional distribution ν
( Y | X ), as measured by the expected conditional KL divergence:
[Equation]
where H [ ν ( Y | X )] is the true conditional entropy and we call P ( θ ) the true predictive risk :
[Equation]
Invoking Jensen’s inequality, we can upper bound P ( θ ) with the true classiﬁcation risk :
[Equation]
C ( θ ) measures how well we can predict the targets given a sample from our representation z of each input x ∈ X .
Unlike P ( θ ) which contains a log of an expectation, we can compute unbiased estimates of C using Monte-Carlo.
While we don’t know the true distribution, we have access to n paired samples from this distribution, a dataset D n = { ( x i , y i ) } ni .
We need ways to approximate the true classiﬁcation risk (eq. (3)) while only using a ﬁnite number of samples.
Maximum Likelihood (ML) tries to minimize the empirical classiﬁcation risk :
[Equation]
approximating the expectation with respect to the true distribution with an average over the observed samples.
From the perspective of variational optimization, ML can concentrate on the deterministic representation q
( Z | x, θ )
=
δ ( z − f θ ( x )) which best predicts the observed target y for each input x .
Unfortunately, ML with ﬁnite samples provides no guaranteed relationship to the true classiﬁcation risk.
In other words, eq.
(4) is neither an upper nor a lower bound on eq.
As an illustration, in ﬁg. 1( a ) we show what happens if we try to minimize eq.
(4) using a neural network with a two dimensional representation z
= ( µ, σ 2 ).
This is used to parameterize the mean and standard deviation of a conditional normal distribution y ∼ Normal ( µ ( x ) , σ 2 ( x )).
The true model in this case consists of x values uniformly distributed from -5 to 5, and y s that are cubic in the x s with ﬁxed standard deviation: y ∼ Normal ( x 3 / 100 , 0 .
This true distribution is shown in orange.
The 10 sampled ( x, y ) pairs the model was trained on are shown as the blue dots.
The neural network quickly learns to set its predictive standard deviation to a small value and overﬁts to the samples.
By approximating the true distribution in eq.
(3) with the average over the empirical samples, instead of concentrating on the true distribution ν
( X, Y ), we’ve had our network attempt to model the empirical distribution ˆ ν n ( x, y )
=
1 n (cid:80) ni δ ( X − x i ) δ
( Y − y i ).
The neural network did well at the task it was asked to do, but not at the task we wanted.
The traditional way to improve the situation is to try to ﬁt a Bayesian neural network.
Instead of assuming that the parameters of the neural network take on particular values,
[Figure]
we treat the parameters θ as a random variable themselves and compute the posterior over the parameters.
This noticeably improves the predictive distribution, as can be seen in ﬁg. 1( b ), but this improvement comes at great computational cost.
In this example, doing
Hamiltonian Monte Carlo to generate samples from the posterior for the 2178 parameters of the neural network took 8 times longer than training the maximum likelihood model.
3. Multiple Target Samples
In the typical setup of a discriminative task, we have a ﬁnite sample of pairs ( x, y ) from the true joint distribution ν
( X, Y ).
This amounts to a single sample for the target for each input, a single draw from each of the conditional distributions
ν ( Y | x ).
Hypothetically, what would happen if we kept the same 10 X samples we used above in ﬁgs.
1( a ) and 1( b ) but collected many Y samples for each?
In ﬁg. 1( c ) we show the result of training precisely the same neural network as in ﬁg. 1( a ) in this new setup.
With access to many Y samples, the network learns to match the true predictive distribution nearly exactly at those sampled points.
At the same time, the neural network does a reasonable job of interpolating between the sampled points while maintaining a good degree of predictive uncertainty.
When asked to extrapolate outside of the data, the quality of the predictive distribution diminishes noticeably.
Overall, asking the neural network to match a half-empirical distribution ν
( Y | X )ˆ ν n ( X ) =
ν ( Y | X ) 1 n (cid:80)
ni
δ ( X − x i ) has produced reasonably good results.
This multiple target setup is similar to problem settings that use soft targets , such as teacher-student learning setups (Hinton et al., 2014), which have proven eﬀective.
Similarly, for image classiﬁcation tasks, using multiple label samples can lead to improved ﬁts (Peterson et al., 2019).
Is there some way to provide the sorts of guarantees Bayesian inference provides, but only with respect to the ﬁnite number of Y samples for each X ?
Is there some kind of compromise position we could adopt that achieves performance similar to that in ﬁg. 1( c ) without requiring actually collecting additional target samples for each input?
4. VIB as PAC-Bayes
One way to view the source of the Bayesian guarantees is that Bayesian inference optimizes a PAC style upper bound on the true risk (Masegosa, 2020; Morningstar et al., 2020).
By penalizing the posterior from being too distinct from the prior, Bayesian inference probably won’t overﬁt
(PAC stands for Probably Approximately Correct, and the bound that Bayesian inference optimizes holds with high probability even with a ﬁnite training sample).
We can invoke the same PAC-Bayes bound as in the Bayesian case, but only on the inner expectation over targets demonstrating that, for any positive eﬀective temperature τ , with probability at least 1 − ξ X (see appendix A for proofs):
[Equation]
While this is a lot to unpack, notice that all of the terms on the second line are constants with respect to the neural network’s representation q ( Z | X, θ ), and so can be dropped in its objective.
Equation (5) (nearly always) provides an upper bound on the true classiﬁcation risk ,
However, it is still intractable as it includes an expectation over ν ( X ).
With this observation, we could instead adopt a mixed approach.
Why not take the
Bayesian strategy of minimizing an upper bound with respect to the conditional expectation of targets ν
( Y | X ) while using the bold Maximum Likelihood strategy of a drop-in Monte
Carlo estimate for the expectation over inputs X ?
Doing so gives us:
[Equation]
This objective is equivalent to the Variational Information Bottleneck (VIB) objective of Alemi et al.
The VIB objective was originally motivated as being a variatonal upper bound on the Information Bottleneck objective (Tishby et al., 1999):
[Equation]
I ( X ; Y ) def =
E p ( X,Y ) (cid:104) log p ( X,Y ) p ( X ) p ( Y ) (cid:105) is the mutual information between X and Y .
The Information Bottleneck aims to ﬁnd a representation Z that is as maximally informative about the target Y as possible ( I ( Z ; Y )), subject to a constraint on how expensive that representation is, measured by how many bits about the input it retains ( I ( Z ; X )).
The eﬀective temperature τ controls the tradeoﬀ between the two.
Here we have stumbled upon an alternative motivation of the same objective, showing that the VIB objective can be seen as half Bayesian .
VIB attempts to protect against overﬁtting on a ﬁnite number of sampled targets for each input without addressing potentially overﬁtting to the ﬁnite number of sampled inputs themselves.
It tries to concentrate on the half-empirical distribution of ﬁg. 1( c ).
The VIB objective does not itself provide any bound on the true classiﬁcation risk, just as Maximum Likelihood does not.
Yet, VIB style objectives have been shown to improve model’s generalization and robustness (Fischer and
Alemi, 2020).
Where building a traditional Bayesian neural network requires a distribution over all of the parameters of the network, solving eq.
(6) only requires a distribution over the output activations of the network.
This is a much lower dimensional space and much easier to deal with computationally.
In the VIB setup, the output of the neural network is made an explicit distribution on the representation space, often chosen to be a Gaussian distribution for simplicity.
Notice that in this interpretation, we are not allowed to learn either the classiﬁer distribution p ( Y | Z ) or the prior or marginal r ( Z ) using eq.
(5), as both of those distributions appear in the second line but are dropped in the objective (eq.
In this way this halfBayesian interpretation diﬀers from the existing VIB literature, where both p ( Y | Z ) and r ( Z ) are thought to be variational approximations that are free to be ﬁt simultaneously with the representation q ( Z | X, θ ).
If the data were split, or there were additional holdout data, these could be used to reﬁne either p ( Y | Z ) or r ( Z ) similar to the setup in Dziugaite et al.
If we wanted to generate a fully valid bound on the true classiﬁcation risk, we could continue the road we are on and simple apply another PAC-Bound to eq.
(5), this time with respect to the parameters of the encoding distribution q
( Z | X, θ ).
See appendix A for the full details, but dropping the constant terms with regards to the objective we obtain a fully
Bayesian variational information bottleneck , now with an added parameter prior r (Θ) and new eﬀective temperature γ :
[Equation]
Realizing eq.
(8) could be as simple as adding weight decay to the parameters of the representation in eq.
Objectives of this sort also appeared in Alemi and Fischer (2018), where again they were motivated from an alternative, information theoretical perspective.
5. Demonstration
To illustrate that this can work, in ﬁg. 2 we show the result of ﬁtting the VIB objective (eq.
(6)) to the same 10 data points (with only a single target sample for each point) as in ﬁgs.
1 and 1( c ), using the same random network initialization.
We emphasize that the VIB objective only requires a single sample of each target The results are sensitive to the choice of τ , so we show several values near the best performing models.
Full experimental details can be found in appendix B. Figure 2( b ) in particular has a similar predictive distribution to ﬁg. 1( c ), while only having access to a single target sample for each of the 10 input samples shown. [Figure]
This qualitative sense that the VIB methods are doing well can be veriﬁed quantitatively.
In table 1 we show the computed KL divergences between the true conditional distribution
ν
( Y | X ) and the predictive distributions q ( Y | X, θ )
= (cid:82) d
z q ( z | X, θ ) p ( Y | z ) for each method.
This conditional KL can then be computed in expectation both with respect to the empirical
X distribution ˆ ν ( X ) (simply the average on the 10 samples), or in expectation with respect to the true ν ( X ), marginalizing from x = − 5 to x = 5 uniformly.
This assesses how well the methods did at learning the predictive distribution both on the values they were given ( KL ) as well as on all values ( KL ).
The VIB approaches are competitive with the fully Bayesian model, while being signiﬁcantly cheaper to optimize. The VIB models did not take noticeably longer to train than the ML model. We give additional experimental results on MNIST classiﬁcation in appendix C.
[Table]
KLs are measured in bits.
The large value for KL for Determ is due to the fact that we know the true σ 2 for ν – even interpolating the sampled points doesn’t protect against a large empirical risk.
6. Conclusion
We’ve demonstrated that on a simple problem we can provide most of the beneﬁts of
Bayesian inference for signﬁcantly less work.
The Variational Information Bottleneck method of Alemi et al.
(2016) can be thought of as a half-Bayesian approach that oﬀers some assurance that it won’t too severely overﬁt, but only with regards to the ﬁnite sampling of the targets in a discriminative modeling task. Acknowledgments
We would like to thank Mallory Alemi, Sergey Ioﬀe, Rif A. Saurous for feedback on the draft.
Appendix A. Theory
In this appendix we prove the claims in the paper.
Suppose ( X, Y ) n iid ∼ ν
( Y | X ) ν
( X ); write their empirical distributions as,
[Equation]
[Equation] For notational simplicity, we regard ˆ ν n ( Y | X ) as 0 (for all y ) for X (cid:54)∈ { x i } ni .
Depending on its context, the symbol δ denotes either the Dirac or Kronecker delta function.
(These two caveats are our only notational abuses in the paper.)
Theorem 1 For all q (Θ) absolutely continuous with respect to r (Θ) , q ( Z | X, Θ) absolutely continuous with respect to r ( Z ) for all { Θ ∈ T : q (Θ) > 0 } and { X ∈ X : ν ( X ) > 0 } , ( X, Y ) n iid ∼ ν
( Y | X ) ν
( X ) , β X , β ∈ R + , n ∈ N , and ξ X , ξ ∈ (0 , 1] , then with probability at least 1 − max( ξ X , ξ ) : [Equation] [Equation]
[Equation]
where,
[Equation]
and where ∆ X , ∆ are deﬁned by Equations 21 and 22 and ν (ˆ ν n ( Y | X )) is the true probability of the empirical conditional measure.
Note that neither ψ X nor ψ are a function of q (Θ) and that quantities 13, 14, andare not a function of the unknowable true data generating distribution,
ν
( Y | X ) ν ( X ) .
Proof
First, with probability at least 1 − ξ X we have:
[Equation]
Inequality 19 follows from Jensen’s inequality and inequality 20 holds with probability at least 1 − ξ X and follows from applying Lemmas 2 and 3 to:
[Equation]
Again applying Lemmas 2 and 3 to:
[Equation]
we conclude that with probability at least 1 − max( ξ X , ξ ) : [Equation 20]
[Equation]
The proof is completed by expanding occurrences of E ˆ ν n as a summation.
The proof of Theorem 1 is similar to twice applying the technique of Morningstar et al.
(2020) (with m = 1).
The Morningstar et al.
(2020) proof followed arguments similar to
Masegosa (2020) which itself followed arguments similar to Germain et al. (2016).
Note that in the text body we used τ = 1 β X and γ = 1 β .
We also note that one can use Lemma 5 to rewrite E ˆ ν n expectations as conditional averages (e.g., this could be done to eq.
A.1. Lemmas
In this section we present several Lemmas used to simplify this paper’s proofs.
The Lemmas are well-known and are given here for the reader’s convenience.
Lemma 2 (Compression)
If p (Θ) is absolutely semicontinuous wrt r (Θ) and E r (Θ) [ e f (Θ) ] < ∞ , then E p (Θ) [ f (Θ)] ≤ KL [ p (Θ) , r (Θ)]
+ log E r (Θ) [ e f (Θ) ] .
Proof Write q (Θ) def = r (Θ) e f (Θ) E r (Θ)
[ e f (Θ) ] and note that Lemma 4 implies, 0 ≤ KL [ p (Θ) , q (Θ)] =
KL [ p (Θ) , r (Θ)] − E p (Θ)
[ f (Θ)]
+ log E r (Θ) [ e f (Θ) ] .
Proof due to Banerjee (2006); Zhang (2006).
Lemma 3 (Log Markov Inequality)
For any ξ ∈ (0 , 1] and random variable Z ∼ p with p ( Z ≤ 0) =
then p (log Z ≤ log E p [ Z ] − log ξ ) ≥ 1 − ξ.
Proof Markov’s inequality states that p ( Z > t ) ≤ E p [ Z ] t for non-negative random variable
Z ∼ p and t > 0 .
Substituting t =
E p
[ Z ] ξ implies p ( Z > E p [ Z ] ξ ) ≤ ξ .
Combining this with the fact that log is a non-decreasing bijection implies p (log Z > log E p
[ Z ] − log ξ ) ≤ ξ.
Examining the complement interval completes the proof.
Lemma 4 (Gibb’s Inequality)
If p (Θ) is absolutely semicontinuous wrt r (Θ) , then KL [ p, q ] ≥ 0 .
Proof KL [ p, q ] = − E p
( x ) (cid:104) log q ( x ) p ( x ) (cid:105) ≥ − log E p ( x ) (cid:104) q ( x ) p ( x ) (cid:105)
= − log 1 = 0 where the inequality is Jensen’s.
Lemma 5 (Conditional Empirical Expectation)
Assuming X ∈ { x i } ni , then:
[Equation]
Proof E
[Equation]
The δ function is either the Dirac or Kronecker delta function, depending on whether measure µ is continuous or discrete.
Appendix B. Experimental Details
For ﬁgs.
1( a ) to 1( c ) and 2( a ) to 2( c ) all experiments were done with JAX (Bradbury et al., 2018).
The true data distribution was taken to be
[Equation]
Ten samples were taken for the data distribution.
The predictive network consisted of two fully connected layers with 32 hidden units followed by an elu activation.
The ﬁnal layer was a linear layer with 2 outputs, the ﬁrst of which was taken as the mean, and the second generated the standard deviation of the predictive model with a softplus activation and a minimum value of 0.01: ( σ 2 = 0 .
01 + softplus ( x )).
A standard Lecun style truncated normal initialization scheme was used for the kernels, and the biases were initialized to be zero.
The initial parameter variance was increased by a factor of 5, which was found to be important to get the one dimensional networks to converge well on the range and domain of the toy problem.
All problems used the same initial parameters and the same adabelief optimizer with a cosine decay schedule on the learning rate starting at 10 − 3 and ending at 0 after 100k steps, the length of each optimization run.
To sample from the Bayesian neural network, tensorﬂow probability ’s JAX backend Hamiltonian Monte Carlo sampler was used (Dillon et al., 2017; Lao et al., 2020).
In particular 1000 results were generated from the chain with 10k burn-in steps, dual averaging step size adaptation with a step size of 10 − 3 and 100 leapfrog steps, 9000 adaptation steps and a target acceptance probability of 0.7.
The initialization distribution used for the neural network experiments was taken to be the prior distribution, both for the kernel and bias parameters.
For the VIB experiments, the same neural network as above was used to form the representation q ( Z | X ).
The classiﬁer network p ( Y | Z ) was taken to be a Normal distribution with a small ﬁxed variance: Normal ( z, 0 . 01 2 ).
The marginal r ( Z ) was set to be a ﬁxed unit
Normal: Normal (0 , 1).
All KL divergences were computed using 100k samples of data points and 10k samples from the any intermediate distributions as required.
Appendix C. MNIST Experiments
In the main paper, we demonstrate that VIB learns a model which retains uncertainty about the targets Y , even if it provides no guarantees about generalization in X .
For an additional illustrative comparison, we train models on the MNIST dataset.
Speciﬁcally, we compare the predictive models of a deterministic deep neural network and a VIB model having the same architecture.
We use a parameter free decoder, with a categorical likelihood, multivariate normal prior, and multivariate normal posterior.
For the prior, we assume zero mean and identity covariance, while we predict the full covariance matrix for the posterior.
The deterministic model replicates this setup, but simply predicts the categorical likelihood rather than the posterior, and it has no prior.
We train models for 50 epochs with a learning rate of 0.001, which is decayed by half every 5000 steps.
We further use a batch size of 128.
For the VIB model, we evaluate the objective using 4 samples from the posterior and use τ = 0 .
005 as a weighting for the KL penalty.
We evaluate models using the classiﬁcation accuracy as well as the log-likelihood of the test set.
To compute the log-likelihood of VIB models, we marginalize over 1000 samples from the posterior.
Both models produce similar ﬁnal test set accuracies (99.2% for the deterministic model, and 99.4% for the VIB model).
We further ﬁnd that the test set log-likelihood for the VIB model is higher (-508 versus -786 for the deterministic model).
However, this later ﬁnding is heavily dependent on the τ used in training:
We ﬁnd that
the
VIB model can measure a lower log-likelihood if the τ multiplier is larger.
Note that the accuracy is more robust to this hyperparameter, and we ﬁnd that the VIB model consistently observes higher accuracy than the deterministic model over the range of τ we explored.
In addition to showing that VIB leads to models which have a higher test log-likelihood and accuracy, we also examined if the resulting predictive models learned by VIB incorporate more uncertainty into the labels than do those which are trained with a deterministic network.
For this, in ﬁg. 3 we show 10 images from the test set, each of which is classiﬁed incorrectly by either the VIB model or by the deterministic model.
Below each image, we show the class probabilities predicted by the deterministic and VIB models, and indicate the true class with a vertical line.
Our main observation is that many of the deterministic models overpredict the probability of a label, with 7 out of 10 images being assigned a class probability greater than 95%.
We also ﬁnd that in these same situations, VIB often folds additional probability into other classes one of which is typically the correct class.
This, when combined with the higher test set log-likelihood agrees with our ﬁndings that VIB facilitate generalization over the label distribution.