Bayesian Neural Networks

Bayesian inference allows us to learn a probability distribution over possible neural networks. We can approximately solve inference with a simple modification to standard neural network tools. The resulting algorithm mitigates overfitting, enables learning from small datasets, and tells us how uncertain our predictions are.

What’s wrong with neural networks?

You may have heard deep neural networks described as powerful function approximators. Their power is due to the extreme flexibility of having many model parameters (the weights and biases) whose values can be learned from data via gradient-based optimization. Because they are good at approximating functions (input-output relationships) when lots of data are available, neural networks are well-suited to artificial intelligence tasks like speech recognition and image classification.

But the extreme flexibility of neural networks has a downside: they are particularly vulnerable to overfitting. Overfitting happens when the learning algorithm does such a good job of tuning the model parameters for performance on the training set—by optimizing its objective function—that the performance on new examples suffers. Deep neural networks have a ton of parameters (typically millions in modern models), which essentially guarantees eventual overfitting because the learning algorithm can always do just a little bit better on the training set by tweaking some of the many knobs available to it. The flexibility of neural networks during training time actually makes them brittle at test time. This might sound surprising at first, but let’s look at the training procedure both mathematically and graphically (for a toy problem) to build some intuition around why deep neural networks overfit.

In standard neural network training we want to learn an input-to-output mapping

y \approx f (x, w)

via a network

f

with weights

w

. We use a dataset of labeled examples

D = {x_{i}, y_{i}}

to minimize a loss function

L (D, w)

with respect to the weights

w

L (D, w) : = x_{i}, y_{i} \in D \sum (y_{i} - f (x_{i}, w))^{2} + λ d \sum w_{d}^{2}

Loss of our model on dataset

D

is squared error between the target label

y_{i}

and output of network with input

x_{i}

and weights

w

hyperparameter

λ

sets importance of regularizing large weight values

w_{d}

This turns out to be equivalent to maximum likelihood training, in other words, maximizing the log probability of the data

D

and weights

w

lo g p (D, w) : = x_{i}, y_{i} \in D \sum lo g N (y_{i} ∣ f (x_{i}, w), I) + d \sum lo g N (w_{d} ∣0, \frac{1}{λ})

Log-likelihood of dataset

D

under model

is log-probability of target label

y_{i}

under gaussian distribution

N

with mean given by output of network

f (x_{i}, w)

and variance

1

and prior belief about weights given by log-probability of

w_{d}

under Gaussian

N

with

0

mean and hyperparameter-controlled variance

\frac{1}{λ}

We train on the training dataset

D_{train}

, but evaluate on the test dataset. As we see in the demo below, the training objective function

lo g p (D_{train})

is different from the test objective function

lo g p (D_{test})

. Overfitting happens when we choose network weights

w

that do well on

lo g p (D_{train})

but poorly on

lo g p (D_{test})

. As we add more layers to the network, the problem gets worse.

Luckily, Bayesian neural networks address overfitting by modeling uncertainty in the weights. Plus they can be trained using standard neural net tools using an algorithm called stochastic variational inference, which we cover at the end of this tutorial. But first let’s explore why being Bayesian helps with overfitting.

Training log likelihood

lo g p (D_{train}, w_{1}, w_{2})

The neural network objective

L

is a function of both the weights and the dataset. Here we show

L_{train}

and

L_{test}

as a function of only two of the neuron weights, with the remaining weights fixed to pre-trained values. Redder is better. Use your mouse to change the values of the weights for these neurons. Notice how weight values affect both the training and test objectives (left, right) and the network predictions changes (below).

Training log likelihood

lo g p (D_{test}, w_{1}, w_{2})

Prediction function

\overset{y}{^}_{w} (x)

. Red dots are training data, green are test.

When it comes to describe training, the last figure is somewhat deceptive, in that the background heatmap appear to be fixed. It only appear so because we fixed every weights and biases except the first two. Only two weights are being adjusted here. If we freed all the weights, the colourful background you see would be constantly changing.

Training log likelihood

lo g p (D_{train}, w_{1}, w_{2})

Observe how in actual training, the weights are (gradient) descending a constantly changing landscape.

Training log likelihood

lo g p (D_{test}, w_{1}, w_{2})

Validation Set

Training Set

Training log likelihood

lo g p (D_{train}, w_{1}, w_{2})

The same discrepancy happens in classification tasks. Notice how the contours differ.

Training log likelihood

lo g p (D_{test}, w_{1}, w_{2})

Bayes to the rescue

The problem of overfitting is certainly not unique to neural networks. But because the flexibility of neural networks makes them particularly susceptible, researchers and practitioners have come up with many extensions to the standard learning algorithm (early stopping, weight decay, and dropout, just to name a few) to reduce overfitting [1].

This tutorial focuses on Bayesian inference, a powerful framework that not only helps with overfitting, but also tells us how uncertain our model is about its parameters. Before we used the training data

D_{train}

plus a gradient-based optimizer to tune the weights

w

in order to maximize

L_{train} (w)

. In Bayesian inference, instead of learning the parameter values, we seek to compute

p (w ∣ D_{train})

, the conditional distribution of the weights given the training data.

p (w ∣ D_{train})

is called the posterior distribution, or often the posterior for short.

Exact Bayesian inference

Bayes’ rule tells us how to compute the posterior: $$ p(w | D) = \dfrac{p(D|w)p(w)}{p(D)} = \dfrac{p(D|w)p(w)}{\int_{w’} p(D|w’)p(w’) dw’} .

p(w)

, an d t h e * l ik e l ih oo d *

p(D | w)

. Whi l e t h e p os t er i or d i s t r ib u t i o ni se a s i l ye x p resse d, i t i s t y p i c a ll ye x p e n s i v e t oco m p u t e d u e t o t h e p es k y in t e g r a l o v er a llp oss ib l e v a l u eso f

. F or m os t n e u r a l n e tw or k w ec an ’ t e v a l u a t e t hi s in t e g r a l ana l y t i c a ll y . W eco u l d in s t e a d a pp ro x ima t e i t n u m er i c a ll y, b u tt hi s w i ll o n l y b e p r a c t i c a l f ors ma ll n e u r a l n e tw or k s, s in ce

re p rese n t s a llt h e w e i g h t s an d bia ses, so i t b eco m es v ery hi g h d im e n s i o na l f or d ee p n e tw or k s . S o w e w i ll n ee d t o a dd ress t hi sco m p u t a t i o na l i ss u e, w hi c h w e d o in t h e n e x t sec t i o n . I n t h e m e an t im e, l e t ’ sco n s i d er an o t h er u p s i d eo f u s in g t h e p os t er i or d i s t r ib u t i o n o v er w e i g h t s

p(w|D)

in s t e a d o f t h ees t ima t e

w^*

f ro m s t an d a r d n e u r a l n e tw or k t r ainin g . I n p a r t i c u l a r, kn o w in g t h e p os t er i or a ll o w s u s t o d o p ro babi l i s t i c p re d i c t i o nb ye x p ress in g a d i s t r ib u t i o n o v er p re d i c t e d o u tp u t

\hat y

a s a f u n c t i o n o f t h e n e w in p u t

, $$ p(\hat y(x)| D) = \int_{w} p(\hat y(x)| w) p(w | D) dw = \mathbb{E}_{p(w|D)}[p(\hat y(x)|w)], $$ which we call the *predictive distribution*. Notice that we can express this distribution as the expectation of the single network likelihood under the posterior

p(w|D)

. This yields a nice interpretation of the predictive distribution as an *infinite ensemble* of networks \[2\], with each network’s contribution to the overall prediction weighted by the posterior likelihood of its weights given the training data. Also, it implies that we can approximate this infinite ensemble using a finite number of Monte Carlo samples from the posterior. Approximating the predictive distribution with only a single Monte Carlo sample is equivalent to using a single network

\hat y_w(x)

w i t h w e i g h t s

c h ose na t r an d o m f ro m

p(w|D)

; in a slight abuse of terminology, we can think of this process as *sampling networks* from the posterior. ### Modeling uncertainty But before we move on, let’s briefly consider the implications of Bayesian inference, which replaces a *point estimate* of the weights

w^*

an d i t scorres p o n d in g p re d i c t i o n f u n c t i o n

\hat y_{w^*}(x)

f ro m s t an d a r d t r ainin g w i t h * in f erre dd i s t r ib u t i o n s *

p(w|D)

an d

p(\hat y_{w}(x)|D)

. An immediate benefit of specifying distributions over the model parameters and predictions is that we can quantify our *uncertainty* about these things, e.g., by computing their variance. This is especially relevant when learning from small datasets; standard neural net training will overfit for the reasons discussed above, but Bayesian inference will find the best explanation for the model parameters given the available data, which typically have high uncertainty when data is scarce. In the limit of dataset size far exceeding the number of neurons, the inferred distributions sharpen and begin to resemble the solutions from the standard training; for modern networks we don’t typically reach this limit in practice, so having a notion of model uncertainty is helpful. ### Approximate inference We have seen computationally difficult integrals arise during the Bayesian setup in the following two expressions<sup><span id="fn-1" data-hover-ref="dt-fn-hover-box-0">1</span></sup> : - the marginal data likelihood

p(D)

, w hi c h re q u i res ain t e g r a l o v er w e i g h t s p a ce an d i sre q u i re d t oe v a l u a t e t h e p os t er i or

p(w|D)

; - t h e p os t er i or p re d i c t i v e d i s t r ib u t i o n

p(\hat y(x)|D)

, w hi c ha g ainin t e g r a t eso v er w e i g h t s p a ce, t hi s t im e v iaan e x p ec t a t i o n o v er

p(w|D)

. A t f i rs t g l an ce, t h e g l a r in g d i ff i c u lt yo f co m p u t in g t h eseo bj ec t s mi g h t s ugg es tt ha tB a yes iann e u r a l n e tw or k s w i ll b e d i ff i c u lt im pl e m e n t in p r a c t i ce . B u t a s w e w i ll see, i t i s p oss ib l e t o a pp ro x ima t e l yco m p u t eo n eor b o t h o f t h eseo bj ec t s v ia c l e v er a l g or i t hm s, w hi c h c anb e b ro a d l yc a t e g or i ze d a se i t h ers am pl in g - ba se d or v a r ia t i o na l . * S am pl in g m e t h o d s * a pp ro x ima t e l yco m p u t e

p(\hat y(x)|D) = \mathbb{E}_{p(w|D)}[p(\hat y(x)|w)

b y g e n er a t in g a f ini t ese t o f n e tw or k w e i g h t s

{w_1, \ldots, w_N}

w h osee m p i r i c a l d i s t r ib u t i o nma t c h es

p(w|D)

in t h e l imi t o f l a r g e

. L oose l ys p e akin g, t h e a l g or i t hmi c d es i g n c ha ll e n g e i s t ore l a t i v e l y q u i c k l y p ro d u ce am o d es t n u mb er (

n o tt oo bi g) o f n e tw or k s am pl es t ha t y i e l d a d ece n t a pp ro x ima t i o n o f

p(\hat y(x)| D)

* Va r ia t i o na l m e t h o d s * in s t e a dd i rec tl y m o d e lt h e p os t er i or

p(w|D)

u s in g a p a r am e t er i ze dd i s t r ib u t i o n

q_\phi(w)

, c a ll e d t h e * a pp ro x ima t e p os t er i or *, t h e ni t er a t i v e l y im p ro v e t h e a pp ro x ima t i o nb yso l v in g a s u i t ab l eo pt imi z a t i o n p ro b l e m s . W e w i ll s h o wt ha t a s u i t ab l eo pt imi z a t i o n p ro b l e m c anb e d er i v e d, so t h ere mainin g d es i g n c ha ll e n g es a re t os p ec i f y

q_\phi

and an optimizer. Thankfully, we can use off-the-shelf optimizers from standard neural network training! While both sampling and variational methods are of practical and historical import, we emphasize that the latter approach admits a learning algorithm, called stochastic variational inference, that closely resembles standard neural network training. We devote a section to this algorithm to underscore that being Bayesian in neural network learning is well worth the modest cost in engineering overhead. TODO: make the point (perhaps in a footnote that variational methods still typically evaluate the posterior predictive via a Monte Carlo estimate, so they are not sampling-free. ## Sampling methods ### Monte Carlo integration Sampling methods rely on *Monte Carlo integration*: the use of a finite set of random samples to approximate an expected value. In particular, they generate a set of neural networks asymptotically distributed according to

p(w|D)

in or d er t o a pp ro x ima t e

p(\hat y(x)|D)

a s t h ee m p i r i c a l e x p ec t a t i o n o f s in g l e - n e tw or k p re d i c t i o n s

p(\hat y(x)|w)

u n d er t h es am pl e d n e tw or k s . M o n t e C a r l o in t e g r a t i o ni s w i d e l y a ppl i c ab l e, b u t i t i s n ’ t a lw a ys p r a c t i c a l . W ec an see t hi s b y d er i v in g anai v e a pp ro x ima t i o n t oo n eo f o u r p re v i o u s l y i d e n t i f i e dd i ff i c u lt in t e g r a l s : t h e ma r g ina l d a t a l ik e l ih oo d

p(D)

. $$ p(D) = \int_{w} p(D|w)p(w) dw = \mathbb{E}_{p(w)} \left[ p(D|w) \right] \approx \frac{1}{N} \sum_{i=1}^{N} p(D|w_i). $$ Here the sample set of network weights

{w_1 \ldots w_N}

a re d r a w nin d e p e n d e n tl y f ro m t h e p r i or

p(w)

. G i v e na g oo d es t ima t e f or

p(D)

t h e n w eco u l d a g ain u se M o n t e C a r l o in t e g r a t i o n u s in g n e tw or k s am pl es

{w_i}

f ro m t h e p r i or t o a pp ro x ima t e t h e p re d i c t i v e p os t er i or :

\begin{aligned} p(\hat y(x)|d) &= \mathbb{e}{p(w|d)} \left[ p(\hat y(x)|w) \right] = \mathbb{e}{p(w)} \left[ p(\hat y(x)|w) \frac{p(d|w)}{p(d)} \right]\ &\approx \frac{1}{n} \sum_{i=1}^{n} p(\hat y(x)|w_i) \frac{p(d|w_i)}{p(d)} \end{aligned}

\begin{aligned} p(\hat y(x)|D) &= \mathbb{E}{p(w|d)} \left[ p(\hat y(x)|w) \right]\ &\approx \frac{1}{N} \sum{i=1}^{n} p(\hat y(x)|w_i) \quad \text{where} \quad w_i \sim p(w|D) \end{aligned}

T h e l a w o f l a r g e n u mb ers a ss u res u s t ha t o u r b o t h M o n t e C a r l oes t ima t es w i ll e v e n t u a ll yco n v er g e t o t h e t r u e q u an t i t i es g i v e n s u ff i c i e n t n u mb ero f s am pl es, b u t in t hi s nai v ese t u p

w i ll n ee d t o b e h ug e d u e t o t h e l a r g e d im e n s i o na l i t yo f

. ### Markov Chain Monte Carlo We saw that Monte Carlo estimates from naively sampling network weights naively from the prior won’t work. We can sacrifice our assumption of sample independence, instead using a Markov Chain to produce a sequence of dependent samples. This is formalized by a *proposal distribution* over choices of next sample in the sequence conditioned on the previous choice

q(w’|w_t)

. T h e M e t ro p o l i s a l g or i t hm u ses t hi s p ro p os a l d i s t r ib u t i o n t o g e n er a t e a se t o f s am pl es t ha t a sy m pt o t i c a ll y a re d i s t r ib u t e d a ccor d in g t o

p(w|D)

. W e u se t h e M a r k o v c hain t o g e n er a t ec an d i d a t es am pl es an d t h e n s t oc ha s t i c a ll y a cce ptt h e m w i t h p ro babi l i t y

, e x p resse d a s t h e a cce pt an cer a t e < s u p >< s p ani d = " f n - 2" d a t a - h o v er - re f = " d t - f n - h o v er - b o x - 1" > 2 < / s p an >< / s u p > i s d e f in e d a s

a = \frac{p(D, w’)}{p(D, w_t)} .

T h e a cce pt r a t e i s a l soe q u i v a l e n tt o t h er a t i oo f p os t er i or p ro babi l i t i es u n d er t h e p ro p ose d an d m os t rece n t s am pl es

\frac{p(w’|D)}{p(w_t| D)}

. T hi s i s b ec a u se

p(D, w) = p(w|D)p(D)

i se q u a lt o t h e p os t er i or u pt o am u lt i pl i c a t i v eco n s t an t

p(D)

t ha t i s t h es am e f or a ll

; h ere w e ’ v e l e v er a g e d t h e f a c tt ha tw ec an co m p u t e

p(D, w)

b u t n o t

p(w|D)

. A simple specification of proposal distribution is a Normal distribution centered at the current sample $$ q(w&#x27;|w_t) = \mathcal{N}(w_t, \sigma^2) $$ . If we were to accept every sample from this proposal distribution, we would be performing a random walk in parameter space with step size equal to

\sigma

. T h e a cce pt / re j ec t s t e p o f t h e M e t ro p o l i s a l g or i t hmmi t i g a t esr an d o m s t e p s in t ore g i o n so f l o wp ro babi l i t y ma ss . W e ’ v e in t ro d u ce d an e w h y p er p a r am e t er

\sigma

t ha tw i ll n ee d t u nin g . G e n er a ll ys p e akin g, w e w an tt o a cce pt n e w s am pl eso f t e nin or d er t oco n v er g e q u i c k l y, an d t h e a cce pt an cer a t e i sse n s i t i v e t o t h e v a l u eo f

\sigma

. MCMC ba se d m e t h o d s a re b e tt er f or hi g h d im e n s i o na l i t y p ro b l e m ss in ce t h ey f oc u so na re a s w i t h re l a t i v e l y hi g h p ro babi l i t y ma ss, w es t i ll d o n ’ t ha v e an y gu a r an t ees ab o u tt h e i rs p ee d o f co n v er g e n ce . TO D O : se gu e w a y t o t h e d e m o G a u ss ian R an d o mWa l k P ro p os a l o n T r ainin g D a t a T h e d e m o ab o v e pl o t s t h ese q u e n ceo ff u n c t i o n s w e g e tw h e n s am pl in g n e tw or k w e i g h t s f ro m t h e p os t er i or u s in g tw os am pl in g a l g or i t hm s : o n e w i t h t h e M e t ro p o l i s p ro p os a l an d o n e w i t h H ami lt o nian d y nami cs . s h o w s t h e p re d i c t i o n f u n c t i o n f or a se q u e n ceo f n e tw or k ss am pl e d in t hi s w a y, a s w e ll a s t h e a pp ro x ima t e p re d i c t i v e d i s t r ib u t i o n co m p u t e d a s t h er u nnin g a v er a g eo f s am pl e d n e tw or k s . T h e w e i g h t e d r u nnin g a v er a g eo f s am pl e df u n c t i o n s i ss ha d e d in re d; t hi scorres p o n d s t o t h e M o n t e C a r l oes t ima t eo f t h e p re d i c t i v e d i s t r ib u t i o n

p(\hat y(x)| D)

. On the other hand, producing *samples* from this unknown distribution is often feasible using algorithms described in the next section, and we can aggregate a finite number of these samples to obtain an approximate posterior. TODO: briefly describe the wealth of other MCMC options that seek to design a better proposal distribution, cite Neal, MacKay, etc.\[3\] ## Variational inference TODO: contrast sampling-based inference with gradient based inference, especially with their ability to scale. Sampling-based inference is good enough for many applications, especially where the sample space is reasonable and achieving an asymptotically exact solution is worth the computational cost. As network and dataset size grow, however, so grows the cost of estimating the posterior via sampling. In these cases we can instead solve a distinct but closely related inference problem that admits a learning algorithm very similar algorithm to what we used for standard neural network training; we’ll pose the problem now and discuss the learning algorithm in the next section. ### Modeling the approximate posterior Generally speaking, when faced with an intractable distribution like the posterior

p(w|D)

, w ec ana pp e a lt o * v a r ia t i o na l m e t h o d s *, w hi c h f i rs t d e f in e a p a r am e t r i ze d an d t r a c t ab l es t an d - in d i s t r ib u t i o n, h erec a ll e d t h e * a pp ro x ima t e p os t er i or *

q_\phi(w)

< s u p >< s p ani d = " f n - 3" d a t a - h o v er - re f = " d t - f n - h o v er - b o x - 2" > 3 < / s p an >< / s u p >, t h e n t u n e t h e p a r am e t ers

\phi

(e . g ., t h es in g l e G a u ss ianin re d) t o a co m pl i c a t e d (so m e t im es in t r a c t ab l e d i s t r ib u t i o n)

(e . g ., t h e mi x t u reo f G a u ss ian s inb l a c k) b yo pt imi z in g t h e d i v er g e n ce b e tw ee n t h ese d i s t r ib u t i o n s . T h es ha d e d co n t o u rre p rese n t s

q \log \frac{q}{p}

an d t h e d iam e t ero f t h ec i rc l e an d i t sco l orre p rese n t s i t s in t e g r a l, t h eK L - d i v er g e n ce

d_{KL}(q||p)

. Ho w s h o u l d w ese l ec tp a r am e t er v a l u eso f

i f w e w an tt o minimi ze

d_{KL}(q||p)

? What about if we want to minimize one of the alternative metrics, reverse KL or Jensen-Shannon distance? We use a quantity from information theory called the *Kullback-Liebler divergence*\[5\] (or KL-divergence) to convey dissimilarity between two distributions, expressed as

\begin{aligned} \quad d_{KL}[q_\phi(w)||p(w|D)] &\coloneqq \int_{w} q_\phi(w) \log \frac{q_\phi(w)}{p(w|D)} \ &= \mathbb{E}{q\theta(w)}[\log q_\phi(w) - \log p(w|D)]. \end{aligned}

S oo u ro pt imi z a t i o n p ro b l e m s h o u l d b eso m e t hin g l ik e “ minimi ze t h eK L - d i v er g e n ce f ro m t h e a pp ro x ima t e p os t er i or t o t h e t r u e p os t er i or ”, w hi c h w e w i ll f or ma l i zes h or tl y . B u t b e f ore p rocee d in g, i t ’ s w or t h w hi l e t oco n s i d erso m e p ro p er t i eso f t h eK L - d i v er g e n ce u s in g t h e ab o v e d e m o, w hi c h f i x es a

t ha tw e ’ d l ik e t o a pp ro x ima t e, an d s h o w s h o wt h e p a r am e t erso f a G a u ss ian

(i t s m e anan d v a r ian ce) a ff ec t

d_{KL}(q||p)

. B y d e f ini t i o n t h eK L - d i v er g e n ce i s n o t sy mm e t r i c a l,

d_{KL}(q||p) \neq d_{KL}(p||q)

, w hi c hi s w h y i t ’ sc a ll e d a d i v er g e n ce in s t e a d o f a d i s t an ce . W ec h oose t h e

q||p

d i rec t i o n f orco m p u t a t i o na l re a so n s : s in ce t h ee x p ec t a t i o ni se x p resse d o v er

w ec an es t ima t e i t u s in g M o n t e C a r l os am pl es . B u tw es h o u l d n o t e t h e im pl i c a t i o n so f t hi sc h o i ceo fd i rec t i o n; w ec an see in t h e d e m o t ha t in t h e

q||p

d i rec t i o n w h ere v er w e d o n ’ t a ss i g n p ro babi l i t y ma ss t o

t h ere i s n o p r i ce p ai df or f ai l in g t o m o d e l ma ss in

. M oreo v er,

m u s t b e 0 w h ere v er

i szero, o t h er w i se

d_{KL}(q||p)

b l o w s u p . I n t h e d e m oyo u c ana l sosee h o wt h e p a r am e t erso f

a ff ec ttw o a lt er na t i v e m e a s u reso fd i ss imi l a r i t y : t h ere v erseK L - d i v er g e n ce

d_{KL}(p||q)

and Jensen-Shannon distance <sup><span id="fn-4" data-hover-ref="dt-fn-hover-box-3">4</span></sup>. ### Approximate inference as optimization Having acquired some intuition about the KL-divergence—-our notion of dissimilarity between the approximate and true posteriors—-we now derive an optimization problem that minimizes this quantity. Regrettably, we can’t directly optimize

d_{KL}(q_\phi(w)||p(w|D)

d u e t o t h e f ami l ia r f r u s t r a t i o n o f n o t b e in g ab l e t oco m p u t e t h e p os t er i or . F or t u na t e l y, w ec an d er i v e a t r a c t ab l e an d c l ose l yre l a t e d q u an t i t y t ha t i se q u a lt o t h eK L - d i v er g e n ce u pt o ana dd i t i v eco n s t an t an d t h u ss u i t ab l e a s a s u rro g a t eo bj ec t i v e .

\begin{aligned} d_{KL} (q_\phi(w)||p(w|D)) &= \mathbb{E}{q\phi(w)} \left[ \log q_\phi(w) \right] - \mathbb{E}{q\phi(w)} \left[ \log p(w|D) \right] \ &= \mathbb{E}{q\phi(w)} \left[ \log q_\phi(w) \right] - \mathbb{E}{q\phi(w)} \left[ \log p(w, D) - \log p(D) \right] \ &= \mathbb{E}{q\phi(w)} \left[ \log q_\phi(w) + \log p(w, D)\right] - \log p(D) \end{aligned}

Here

p(D)

re p rese n t s t h e ma r g ina l d a t a l ik e l ih oo d; w e kn o wt hi s q u an t i t y t o b e in t r a c t ab l e t oco m p u t es in ce i tw a s t h e d e n o mina t or in B a yes ’ r u l e . B u t n o t ab l y t hi s t er m d oes n o t d e p e n d o n t h e p a r am e t erso f

q_\phi

, w hi c hm e an s i t c anb e i g n ore d w hi l eo pt imi z in g w i t h res p ec tt o

\phi

. W e ’ v ee n c a p s u l a t e d a ll o f t h eco m p u t a t i o na l in t r a c t abi l i t y in t h e

\log p(x)

t er m; t akin g t h ee x p ec t a t i o n o f t h ere mainin g t er m so v er t h e d a t a d i s t r ib u t i o n

y i e l d s a t r a c t ab l e an d s u i t ab l es u rro g a t eo bj ec t i v ec a ll e d t h ee v i d e n ce l o w er b o u n d (E L BO) :

\begin{aligned} \tilde L(\phi) &:=\mathbb{E}{D, q\phi(w)} \left[ \log p(w, D) - \log q_\phi(w) \right] \ &= \mathbb{E}{x, y \sim D} \left[ \mathbb{E}{w \sim q_{\phi}(w)} \left[ \log p(\hat y(x) = y|w) + \log p(w) - \log q_{\phi}(w) \right] \right] \end{aligned}

T h u s t h e v a r ia t i o na l in f ere n cecorres p o n d s t oo pt imi z in g t h e E L BOw i t h res p ec tt o t h e v a r ia t i o na lp a r am e t ers :

\underset{\phi}{\text{maximize}} \quad \tilde L(\phi)

T h e p ro babi l i t y d e n s i t y f u n c t i o n o f

q_\phi(w)

, w hi c h d e p e n d so ni t s p a r am e t ers

\phi

ina dd i t i o n t o

, is a design choice that we will need to specify. ## Optimizing with gradients Variational inference re-frames the computation of an integral (the marginal likelihood from exact inference) as the optimization of its lower bound. Notably, this implies that we can now use tools from the optimization literature to approximately solve our inference problem. This includes optimizing with minibatch (a.k.a. “stochastic”) gradients, like we did with standard neural network training. Thus, with a slight tweak, we can adapt stochastic gradient ascent to variational inference in an algorithm called *Stochastic Variational Inference*. As discussed above, an important step of designing a variational inference algorithm is to specify the parameterized variational distribution over the weights,

q_\phi(w)

. I n or d er t oe x pl i c a t e t h e l e a r nin g a l g or i t hm w e d i sc u ss a s im pl ec h o i ceo f v a r ia t i o na l d i s t r ib u t i o n : a d ia g o na l - co v a r ian ce G a u ss ian . Ho w e v er t h e p r a c t i c a l im p or t o f a pp ro p r ia t e l ys p ec i f y in g

s h o u l d b ee m p ha s i ze d; in t h eory w e w an t

t o b e b es u ff i c i e n tl ye x p ress i v e t o m o d e lt h e t r u e p os t er i or

p(w|D)

, w hi c hma y n o t b e t r u eo f t h e d ia g o na l - co v a r ian ce G a u ss ian f or m os t in t eres t in g p ro b l e m s . TO D O : o pp or t u ni t y t oc i t e p a p ers l ik e n or ma l i z in g f l o w s h ere Sp ec i f y in g

a s a d ia g o na l - co v a r ian ce G a u ss ianim pl i es t ha t a ll o f t h e n e tw or k w e i g h t s

a re j o in tl y d i s t r ib u t e d a ccor d in g t o t h e m u lt i v a r ia t e G a u ss ian

q_{\mu, \sigma}(w) = \mathcal{N}(\mu, \sigma^2 I)

, w hi c hha s l o g p ro babi l i t y

\log q_{\mu, \sigma}(w) = -\frac{1}{2}(w-\mu)^T(\sigma^2 I)^{-1}(w - \mu)

N o t e t ha tw e ha v ec h ose nann o t a t i o n t ha t mak ese x pl i c i tt h e v a r ia t i o na lp a r am e t ersco m p r i s in g a llt h e m e an s an d s t an d a r dd e v ia t i o n s

\phi = (\mu, \sigma)

. D u e t o t h es t r u c t u reo f t h eco v a r ian ce ma t r i x, w ec an e q u i v a l e n tl y t hink o f e a c hn e tw or k w e i g h t

w_i

a s a r an d o m v a r iab l e d i s t r ib u t e d a ccor d in g t o i t so w n G a u ss ian d i s t r ib u t i o n w i t hm e an

\mu_i

an d s t an d a r dd e v ia t i o n

\sigma_i

. I n or d er t o d er i v e a g r a d i e n t - ba se d l e a r nin g a l g or i t hm, l e t ’ sco m p a re t h e v a r ia t i o na l in f ere n ceo bj ec t i v e f u n c t i o n

\tilde L(\phi)

(w i t h o u rc h o i ceo f

q_\phi

) t o t h e M A P o bj ec t i v e f u n c t i o n

L(w)

from standard neural network training. MAP neural network training is $$ L(w) = \mathbb{E}_{x, y \sim D} \left[ \log p(\hat y(x) = y|w) + \log p(w) \right]. $$ VI neural network training is

\tilde L(\mu, \sigma) = \mathbb{E}{x, y \sim D} \left[ \mathbb{E}{w \sim q_{\mu, \sigma}} \left[ \log p(\hat y(x) = y|w) + \log p(w) - \log q_{\mu, \sigma}(w) \right] \right]

W e w an tt oo pt imi ze t h e V I o bj ec t i v e w i t h g r a d i e n t s, w hi c hn ecess i t a t esco m p u t in g

\nabla_\mu \tilde L(\mu, \sigma)

an d

\nabla_\sigma \tilde L(\mu, \sigma)

, t h e v ec t oro f p a r t ia l g r a d i e n t so f t h e V I o bj ec t i v e w i t h res p ec tt o t h e m o d e lp a r am e t ers

(\mu, \sigma)

. N o t e t ha tw h ere a s b o t h t h e M A P an d V I o bj ec t i v es in v o l v e an e x p ec t a t i o n o v er t h e d a t a d i s t r ib u t i o n, t h e V I o bj ec t i v e a dd i t i o na ll yre q u i res an e x p ec t a t i o n o v er t h e w e i g h t s d i s t r ib u t i o n . W ec an u se t h es t an d a r d s t oc ha s t i c g r a d i e n t a sce n t a l g or i t hm, w hi c h w e d i sc u ss n e x t, t oco m p u t e g r a d i e n t s t h ro ug h t h ee x p ec t a t i o n o v er d a t a

\mathbb{E}_{x, y \sim D}

u s in g miniba t c h es . T h e t ro u b l e i s w i t h t h eseco n d e x p ec t a t i o n

\mathbb{E}{w \sim q{\mu, \sigma}}

s in ce t h e d i s t r ib u t i o n

q_{\mu, \sigma}

d e p e n d so n t h e p a r am e t ers

\mu

an d

\sigma

that we want to optimize. ### Stochastic gradient ascent Let’s revisit a detail we skipped over in earlier on: how to compute gradients with respect to

o f t h es in g l e n e tw or k l o g l ik e l ih oo d . T h e q u an t i t y w e w an tt oco m p u t e i s

\nabla_w L(w)

. W ec an s h o wt ha tt h e g r a d i e n t o f t h e a v er a g eo v er t r ainin g e x am pl es i se q u i v a l e n tt o t h e a v er a g eo v er t h e

t r ainin g e x am pl eso f in d i v i d u a l g r a d i e n t s a s f o ll o w s :

\begin{aligned} \nabla_w L(w) &= \nabla_w \mathbb{E}{x,y \sim D} \left[ \log p(\hat y(x_i) = y_i|w) + \log p(w) \right] \ &= \nabla_w \frac{1}{N} \sum{i=1}^{N} \left[ \log p(\hat y(x_i) = y_i|w) + \log p(w) \right] \ &= \underbrace{\frac{1}{N} \sum_{i=1}^{N} \nabla_w \left[ \log p(\hat y(x_i) = y_i|w) + \log p(w) \right]}{g{\text{batch}}(w)}. \end{aligned}

W ec a llt hi sco m p u t a t i o n o f t h ee x a c t g r a d i e n t

g_{\text{batch}}(w)

t h e * ba t c h g r a d i e n t * b ec a u se i t in v o l v es a v er a g in g o v er t h ee n t i re ba t c h o f t r ainin g d a t aa t o n ce . W e u se d t h e in d e p e n d e n ceo f t h e d a t a f ro m t h e w e i g h t s t o m o v e t h e g r a d i e n t o p er a t or

\nabla_w

in s i d e t h ee x p ec t a t i o n o v er t r ainin g d a t a

\mathbb{E}_{x,y}

. E a c h

\nabla_w

in s i d e t h es u m c anb eco m p u t e d e ff i c i e n tl y w i t hba c k p ro p a g a t i o n, an d t h eo v er a ll a v er a g e i s p a r a ll e l i z ab l e . T h e a l g or i t hm t ha t i t er a t i v e l y u p d a t es t h e w e i g h t s a ccor d in g t o

w \leftarrow w + \nabla_{w}L(w)

i s t h u sc a ll e d ba t c h g r a d i e n t a sce n t Wh ere a s ba t c h g r a d i e n t a sce n t re q u i res i t er a t in g o v er t h ee n t i re t r ainin g d a t a f ore v ery u p d a t e t o

, w e mi g h t re a so nab l y a ss u m e t ha tt h e g r a d i e n t

\nabla_w L(w)

c anb e * es t ima t e d * b y a v er a g in g o v er a s ma ll ers u b se t o f t r ainin g e x am pl es . B y d i v i d in g t h e t r ainin g d a t ain t o

p a r t i t i o n sc a ll e d “ miniba t c h es ” w ec an co m p u t e t hi s a pp ro x ima t e g r a d i e n t b y a v er a g in g o v er t h e

N_m

s am pl eso f t h e

\- t hminiba t c h :

\nabla_w L(w) \approx \underbrace{\frac{M}{N} \sum_{i=1}^{N_m} \nabla_w \left[ \log p(\hat y(x_i) = y_i|w) + \log p(w) \right]}{g{\text{SGA}}(w)}

E s t ima t in g t h e g r a d i e n t in t hi s w a y i sc a ll e d * s t oc ha s t i c g r a d i e n t a sce n t * (SG A) . I t s p ee d s u pt h e t r ainin g b yro ug h l y a f a c t oro f

so l o n g a s t h e miniba t c h es a re l a r g ee n o ug h so t ha t

g_{\text{SGA}}(w)

i s a re a so nab l y g oo d a pp ro x ima t i o n t o

g_{\text{batch}}(w)

. Interestingly, stochastic gradient updates provide some additional benefits including mitigating overfitting. TODO: cite papers on the theory of SGD here. also maybe talk more about the role of gradient noise as a regularizer, since we will use that concept later on ### The reparameterization gradient estimator Now lets turn our attention to the gradient-based optimization of the variational objective

\tilde L(\mu, \sigma)

. $$ \tilde L(\mu, \sigma) = \mathbb{E}_{x, y \sim D} \left[ \mathbb{E}_{w \sim q_{\mu, \sigma}} \left[ \log p(\hat y(x) = y|w) + \log p(w) - \log q_{\mu, \sigma}(w) \right] \right] $$ We will need to compute the gradients of the objective with respect to the parameters,

\nabla_\mu \tilde L(\mu, \sigma)

an d

\nabla_\sigma \tilde L(\mu, \sigma)

W es a w e a r l i er h o wt h e in d e p e n d e n ceo f t h e d a t a f ro m t h e w e i g h t s l e t s u s m o v e t h e g r a d i e n t o p er a t or

\nabla

inside the outer-most expectation. $$ \nabla_{\mu, \sigma} \tilde L(\mu, \sigma) = \frac{1}{N} \sum_{i=1}^{N} \left[\nabla_{\mu, \sigma} \mathbb{E}_{w \sim q_{\mu, \sigma}} \left[ \log p(\hat y(x_i) = y_i|w) + \log p(w) - \log q_{\mu, \sigma}(w) \right] \right] $$ Unfortunately, that’s as far as we can move

\nabla

, s in ce t h e n e x t e x p ec t a t i o ni so v er a d i s t r ib u t i o n t ha t d e p e n d so n t h e p a r am e t ers

\mu

an d

\sigma

t ha tw e ’ re t ry in g t oo pt imi ze . TO D O : m oree x pl i c a t i o nh ere

w \sim \mathcal{N}(\mu, \sigma I)

w = \mu + \sigma \epsilon

w i t h

\epsilon \sim \mathcal{N}(0, I)

. Putting this all together (including the minibatch trick from before) yields the SVI gradient estimator: $$ \nabla_{\mu, \sigma} \tilde L(\mu, \sigma) \approx \underbrace{\frac{M}{N} \sum_{i=1}^{N_m} \sum_{j=1}^{S} \nabla_{\mu, \sigma} \left[ \log p(\hat y(x_i) = y_i|w(\mu, \sigma, \epsilon_j)) + \log p(w(\mu, \sigma, \epsilon_j)) \right]}_{g_{\text{SVI}}(\mu, \sigma)}. $$ Notice that we’ve introduced a new hyperparameter

w hi c h se t s t h e n u mb ero f s am pl es f ro m

p(\epsilon)

w e d r a wt o M o n t e C a r l oes t ima t e t h ee x p ec t e d l o g l ik e l ih oo d . W ec an t hink o f t hi s a s t h e n u mb ero f s am pl e d n e tw or k s w eco n s i d er d u r in g t r ainin g . P er ha p ss u r p r i s in g l y, w ec an o f t e n g e t a w a y w i t ha s ma ll v a l u eo f

a tt r ainin g t im e (

S=1

is commonly used). At test time, however, we can increase the number of samples used to estimate the posterior predictive distribution. TODO: list the exact assumptions of the reparameterization estimator and mention what we can do if the variational distribution over weights is discrete (cite Luisoz work on binary nnets) ### Stochastic variational inference TODO: put it all together and segue to the algo box 1:procedure SGA(

x, y, w

) 2 : w hi l e t r ainin g n o t co n v er g e dd o 3 :

g = \text{BACKPROP}(x, y, w)

4 :

w = w + \alpha g

5 : e n d w hi l e 6 : e n d p roce d u re 1 : p roce d u re S V I (

x, y, w

) 2 : w hi l e t r ainin g n o t co n v er g e dd o 3 :

{\color{blue} \epsilon \sim \mathcal{N}(0, I)}

4 :

g = \text{BACKPROP}(x, y, {\color{blue}w+\epsilon\sigma})

5 :

w = w + \alpha g

6 :

{\color{blue}\sigma = \sigma + \alpha\epsilon g}

🪴 Anil's Garden

Explorer

Bayesian Neural Networks

What’s wrong with neural networks?

Bayes to the rescue

Exact Bayesian inference

Graph View

Table of Contents

Backlinks