Bayesian Inference in the Wolfram Language

dr. ir. Sjoerd Smit, Technical Consultant at Wolfram Research Europe

ABSTRACT

In this presentation I will show you two different methods do Bayesian inference in the Wolfram Language and how these methods are useful for making model predictions that incorporate both variance and parameter uncertainty. 

The first method is based on the Nested Sampling algorithm by John Skilling, which allows you to generate parameter samples from a posterior distribution and also lets you compare how appropriate different mathematical models are for the problem at hand.

The second method uses stochastic neural networks to approximate Bayesian Gaussian process regression/classification. This approach has the advantage that it is very scalable and only uses standard neural network tools in a way that should be quite familiar to anyone who’s already acquainted with neural networks.

Brief introduction to Bayesian inference

The uncertainties involved in making predictions from data

Knowns:

Data

Uncertainties:

Model choice

Model parameters for each a given model

Prediction uncertainty intrinsic to the model

Example of a regression problem: fitting vs Bayesian inference

As an illustrative example, let’s consider polynomial regression because it’s easy to understand. From the standpoint of Bayesian inference it doesn’t really matter if you’re doing a simple linear fit or if you’re working with stochastic differential equations: you always proceed in the same way.

Example data:

ListPlot[testdata,PlotStyle->Directive[Red,PointSize[0.01]]]

Example polynomial models:

y= a + b x + ϵ
ϵ \[Distributed] NormalDistribution[0, σ]
θ = {a, b, σ}

y= a + b x + c (x^2) + ϵ
ϵ \[Distributed] NormalDistribution[0, σ]
θ = {a, b, c, σ}

...

Fitting by point estimation

A frequently use method to get your regression coefficients is maximum likelihood estimation

residuals = testdata[[All,2]]-(a +b testdata[[All,1]])

logLikelihood=Simplify[
LogLikelihood[NormalDistribution[0,σ],residuals]
]

{max,fit}=Maximize[
{logLikelihood, σ>0},
{a,b,σ}
]
y == a + b x +NormalDistribution[0,σ]/.fit

Fitted linear model:

Plot[Evaluate[a+b x+InverseCDF[NormalDistribution[0,σ],{0.01,0.5,0.99}]/.fit],{x,-2.5,2.5},Epilog->{Red,PointSize[0.01],Point[testdata]},ImageSize->500,
PlotLabels->Placed[{"1%", "median", "99%"},{Scaled[1],Above}],
PlotLabel->Style[StringForm["Loglikelihood: `1`", max],Large]
]

The error bands look believable enough, but remember that there is no information about the errors in a, b or σ included into these bands!

With some effort, we can find error margins on a, b or σ, but incorporating these errors into the prediction is not trivial...

Same exercise for (2^nd) order polynomial:

residuals2 = testdata[[All,2]]-(a +b testdata[[All,1]] + c ((testdata[[All,1]])^2));
logLikelihood2=Simplify[
LogLikelihood[NormalDistribution[0,σ],residuals2]
];
{max2,fit2}=Maximize[
{logLikelihood2, σ>0},
{a,b,c,σ}
];
Plot[Evaluate[a+b x+c (x^2)+ InverseCDF[NormalDistribution[0,σ],{0.01,0.5,0.99}]/.fit2],{x,-2.5,2.5},Epilog->{Red,PointSize[0.01],Point[testdata]},ImageSize->500,
PlotLabels->Placed[{"1%", "median", "99%"},{Scaled[1],Above}],
PlotLabel->Style[StringForm["Loglikelihood: `1`", max2],Large]
]

Would you use the (2^nd) order model because it has a higher likelihood?

Would you trust extrapolation based on the (2^nd) order model?

Bayesian inference and Nested Sampling

Bayesian inference is concerned with conditional probabilities. Basically, you write down what you have and then use probability theory to get to the probabilities you need.

The likelihood we used above is the probability of seeing the data when a given model (plus parameters) is assumed:

(Subscript[likelihood, i]) = \[ProbabilityPr][Data\[Conditioned](Subscript[θ, i]),(Subscript[Model, i])]

The likelihood is important to know, but you shouldn’t overestimate it’s value. We know the data, so why are we examining the probability of getting the data when we already have them?

Instead, it is much more interesting to know the probabilities for things we really don’t know, namely:

The probability that (Subscript[Model, i]) is the correct one

(Subscript[Evidence, i]) ∝ \[ProbabilityPr][(Subscript[Model, i])\[Conditioned]Data]

The probability distributions of the parameters of a given model

(Subscript[Posterior, i]) = \[ProbabilityPr][(Subscript[θ, i])\[Conditioned]Data,(Subscript[Model, i])]

The predictive distributions of future data for a given model

(Subscript[Prediction, i]) = \[ProbabilityPr][new data\[Conditioned]observed data,(Subscript[Model, i])]

All of these probabilities can be calculated by using Baye’s rule, but the problem is usually not analytically tractable. Instead, we can use methods for obtaining samples from the posterior to get approximations:

Markov Chain Monte Carlo (MCMC)

Gibbs sampling

Nested Sampling (which is what I’m using; link to code at the end of the presentation)

Calculates the evidence

Generates samples from the posterior distribution (thus giving us an approximation to this distribution when we wrap EmpiricalDistribution around it)

How does nested sampling work?

Maintain set of “living” samples of the prior (30 in this case)

Sample the prior for a new point; it is required to have a higher likelihood than the worst of the currently living samples

Kill off the worst sample after each iteration

Making predictions from the posterior

If we want to make a prediction for y(x), we can use this posterior to average out over all of the possible values of a, b and σ:

\[ProbabilityPr][y(x)\[Conditioned]Data, (Subscript[Model, i])]==∫\[ProbabilityPr][y(x)\[Conditioned]Data, (Subscript[Model, i]), (Subscript[θ, i])]×\[ProbabilityPr][(Subscript[θ, i])\[Conditioned]Data, (Subscript[Model, i])]dθ

What this practically means is that you do a whole bunch of different fits and assign weights to each one according to how good it is

posterior=nestedSampling[
(* pass the log-likelihood as a pure function *)
Function[{a,b, σ},Evaluate@logLikelihood], 
(* specify the priors on a, b and σ *)
{"LocationParameter","LocationParameter","ScaleParameter"}, 
(* specify the parameters of the model and the range over which they're allow to vary *)
{{a,-100,100},{b,-100,100},{σ,0.01,100}},
"MonteCarloSteps"->300,
"SamplePoolSize"->100,
"MaxIterations"->5000
];
posterior["LogEvidence"]
posterior["EmpiricalPosteriorDistribution"]

covarianceMatrixPlot[posterior]
posteriorBubbleChart[posterior,{a,b},0.99 (* Show the most relevent 99% of the posterior *)]
posteriorBubbleChart[posterior,{a,b,σ},0.99]

Predictions of the model based on the calculated posterior distributions. Note how the evidence of the (1^st) order model is higher than that of the (2^nd) order model: the evidence naturally takes model complexity into account

(1^st) order (log-evidence = -29.75):

({{, }})

(2^nd) order (log-evidence = -33.10):

({{, }})

Approximate Bayesian inference with neural networks

Basic idea:

The following method was proposed by Yarin Gal (see references at the end of the presentation) and has been presented on the (HyperlinkURL["Wolfram blog", http://blog.wolfram.com/2018/05/31/how-optimistic-do-you-want-to-be-bayesian-neural-network-regression-with-prediction-errors/]):

Build neural network like normal

Add DropoutLayers to randomise the the network

Train with L2 regularisation (basically the prior on the network weights)

Keep the DropoutLayers active (using NetEvaluationMode -> “Train”) while you make predictions and sample the network several times to get a feel for the prediction errors

Advantages of this method:

Free-form method like Gaussian process regression

Easy to implement in the Wolfram neural network framework

Good performance on large amounts of data

Building the network

We need a network that predicts a y value and an error bar for each x value.

λ2 = 0.01; (* L2 regularisation coefficient *)
pdrop = 0.5; (* dropout probability *)
nUnits = 50; (* dimensionality of the hidden layers *)
activation =Ramp; (* Ramp non-linearity gives us a piecewise-linear regression method *)
regressionNet= NetGraph[
<|
"regression"-> NetChain[{
LinearLayer[nUnits],(* map x to a high-dimensional vector space *)
ElementwiseLayer[activation],
DropoutLayer[pdrop],
LinearLayer[nUnits],
ElementwiseLayer[activation],
DropoutLayer[pdrop]
}],
"number1"-> LinearLayer[], (* Project output vector back to single numbers *)
"number2"->LinearLayer[]
|>,
{
NetPort["Input"]->"regression",
"regression"->"number1"-> NetPort["Mean"],
"regression"->"number2"-> NetPort["LogPrecision"]
},
"Input"->"Real",
"Mean"->"Real",
"LogPrecision"->"Real"  (* logτ = Log[1/(σ^2)]*)
]

Training the network

To train the network, we minimize the negative log-likelihood of the observed data (multiply by 2 to make the formula a bit nicer):

FullSimplify[-2LogLikelihood[NormalDistribution[yPredicted,σ],{yObserved}]/.σ-> (1/(Sqrt[Exp[logτ]])),Assumptions->logτ∈Reals]

loss=Function[{yObserved,yPredicted,logPrecision},
(((yPredicted-yObserved))^2)*Exp[logPrecision]-logPrecision
];
lossNet =NetInitialize@ NetGraph[<|
"reg"-> regressionNet,
"negLoglikelihood"->ThreadingLayer[loss]
|>,
{
NetPort["x"]->"reg",
{NetPort["y"],NetPort[{"reg","Mean"}],NetPort[{"reg","LogPrecision"}]}-> "negLoglikelihood"->NetPort["Loss"]
},
"y"->"Real",
"Loss"->"Real"
]

trainedNet = NetTrain[
lossNet,
<|"x"->testdata[[All,1]],"y"->testdata[[All,2]]|>,
LossFunction->"Loss",
Method->{"ADAM","L2Regularization"->λ2},
TimeGoal->20
];

Making predictions

First we need to get our regression network back out of the network we trained:

trainedRegressionNet=NetExtract[trainedNet,"reg"]

Here’s what happens when you apply the network to a list of x-values:

xValues=Range[-2,2,0.025];
singleSample = trainedRegressionNet[xValues,NetEvaluationMode->"Train"];
Short/@singleSample

errorPlot[mean_List,stdev_List,xValues_List,col:_?ColorQ:Blue]:=ListPlot[
TimeSeries[
Transpose[{
mean-stdev,
mean,
mean+stdev
}],
{xValues},
ValueDimensions->3
],
errorStyle[col]
];
Show[
errorPlot[singleSample["Mean"],Exp[-0.5*singleSample["LogPrecision"]],xValues],
ListPlot[testdata,PlotStyle->Red],
ImageSize->Large
]

This function samples the net several times and aggregates the results into a predictive mean and standard deviation:

sampleNet[net :(_NetChain|_NetGraph),xvalues_List,sampleNumber_Integer?Positive]:=
With[{regressionNet=NetExtract[net,"reg"]},
With[{
(* Evaluate the stochastic network a given number of times on the xvalues *)
samples = Select[Table[regressionNet[xvalues,NetEvaluationMode->"Train"],{i,sampleNumber}],AssociationQ]
},
<|
"Mean"-> Mean[samples[[All,"Mean"]]],
"StandardDeviation" -> Sqrt[Variance[samples[[All,"Mean"]]]+Mean[Exp[-samples[[All,"LogPrecision"]]]]]
|>
]
];

sampledPredictions = sampleNet[trainedNet,xValues,1000];
Short/@sampledPredictions

Show[
errorPlot[sampledPredictions["Mean"],sampledPredictions["StandardDeviation"],xValues],
ListPlot[testdata,PlotStyle->Red],
ImageSize->600,
PlotRange->All,
PlotLabel->"Prediction + 1σ error bands"
]

Compare this with the “default” (non-stochastic) prediction of the network:

defaultPrediction = trainedRegressionNet[xValues];
Short/@defaultPrediction

Legended[
Show[
ListPlot[testdata,PlotStyle->Red],
errorPlot[defaultPrediction["Mean"],Exp[-0.5*defaultPrediction["LogPrecision"]],xValues,Orange],
errorPlot[sampledPredictions["Mean"],sampledPredictions["StandardDeviation"],xValues,Blue],
ImageSize->Large,
PlotRange->All
],
SwatchLegend[{Blue,Orange},{"Sampled stochastic network","Standard prediction"}]
]

Notice how the extrapolation error of the standard prediction seems very optimistic because it doesn’t incorporate any uncertainty information about the network weights

References

My Wolfram Blog post (HyperlinkURL["http://blog.wolfram.com/2018/05/31/how-optimistic-do-you-want-to-be-bayesian-neural-network-regression-with-prediction-errors/", http://blog.wolfram.com/2018/05/31/how-optimistic-do-you-want-to-be-bayesian-neural-network-regression-with-prediction-errors/])

E. T. Jaynes, Probability Theory - The Logic of Science, Cambridge University Press (2003)

John Skilling, Nested sampling for general Bayesian computation, Bayesian Anal. Volume 1, Number 4 (2006), 833-859, DOI:10.1214/06-BA127

David J.C. MacKay, Information Theory, Inference, and Learning Algorithms, Cambridge University Press (2003). Freely available at (HyperlinkURL["http://www.inference.org.uk/mackay/itila/", http://www.inference.org.uk/mackay/itila/])

Yarin Gal, Uncertainty in Deep Learning, 2016, PhD thesis. Available at (HyperlinkURL["http://www.cs.ox.ac.uk/people/yarin.gal/website/thesis/thesis.pdf", http://www.cs.ox.ac.uk/people/yarin.gal/website/thesis/thesis.pdf])

Yingzhen Li and Yarin Gal, Dropout Inference in Bayesian Neural Networks with Alpha-divergences, 2017. Available from arXiv: (HyperlinkURL["https://arxiv.org/pdf/1703.02914.pdf", https://arxiv.org/pdf/1703.02914.pdf])

Github repository with code that implements Bayesian inference using Nested Sampling: (HyperlinkURL["https://github.com/ssmit1986/BayesianInference", https://github.com/ssmit1986/BayesianInference])

Initialisation

$HistoryLength=10;
<<BayesianInference`;
stockdata =Once[FinancialData["GE","Jan. 1, 2010"],"Local"];

BlockRandom[
SeedRandom[1];
testdata={#,(#+2 RandomReal[])}&/@RandomReal[{-1,1},20];
];

sdeTestData=;
residuals = testdata[[All,2]]-(a +b testdata[[All,1]]);
logLikelihood=Simplify[
LogLikelihood[NormalDistribution[0,σ],residuals]
];
errorStyle[col_?ColorQ] := Sequence[
Joined->True,
Filling->{1->{2},3->{2}},
PlotStyle->{Lighter[col],col,Lighter[col]}
];
errorPlot[mean_List,stdev_List,xValues_List,col:_?ColorQ:Blue]:=ListPlot[
TimeSeries[
Transpose[{
mean-stdev,
mean,
mean+stdev
}],
{xValues},
ValueDimensions->3
],
errorStyle[col]
];

nestedSamplingAnimation[result_,logLikelihood_]:=Module[{
samples=Values/@Values@Part[
SortBy[#LogLikelihood&]@result["Samples"],
All,{"Point","LogLikelihood"}],
subset,
nLiveSamples=result["SamplePoolSize"],
dist = result["GeneratingDistribution"]
},
Table[
subset=TakeDrop[Take[samples,UpTo[n]],-nLiveSamples];
ContourPlot[
logLikelihood[a,b],
{a,-10,10},{b,-10,10},
FrameLabel->(Style[#,"Text",16]&)/@{"a","b"},ImageSize->Large,
Epilog->
{
{Black,PointSize[0.01],Point[Most@subset[[2,All,1]]],PointSize[0.025],Point[Last@subset[[2,All,1]]]},
{Red,PointSize[0.01], Point[Most@subset[[1,All,1]]],PointSize[0.025],Point[Last@subset[[1,All,1]]]}

},
Contours->subset[[2,All,2]],
PlotLegends->Grid[{{Graphics[{Red,Disk[]},ImageSize->10],"Alive"},{Graphics[{Black,Disk[]},ImageSize->10],"Dead"}},BaseStyle->"Text",Alignment->Left]
],
{n,nLiveSamples+1,100}
]
];

loglikeFunction=Function[{a,b},Evaluate[logLikelihood/.σ->1]];
animationResult=nestedSampling[
loglikeFunction,
{"LocationParameter","LocationParameter"},
{{a,-10,10},{b,-10,10}},
"MonteCarloSteps"->300,
"SamplePoolSize"->30,
"MaxIterations"->5000
];

animation=nestedSamplingAnimation[animationResult,Function[{a,b},Evaluate[logLikelihood/.σ->1]]];

ListAnimate[animation,AnimationRate->1]