Thompson Sampling using Conjugate Priors

Multi-Armed Bandits: Part 5b

Mar 9, 2021

Recap

Baby Robot has entered a charging room containing 5 different power sockets. Each of these sockets returns a slightly different amount of charge. We want to get Baby Robot charged up in the minimum amount of time, so we need to locate the best socket and then use it until charging is complete.

This is identical to the Multi-Armed Bandit problem except that, instead of looking for a slot machine that gives the best payout, we’re looking for a power socket that gives the most charge.

Overview

When we previously looked at the theory behind Thompson Sampling we began with the very simple Bernoulli Bandit problem, in which there were only two possible outcomes. When you selected a power socket you’d either get some charge or no charge. In Bandit terms, when you played a bandit, you’d either win or you’d lose. The challenge was to find the socket with the highest probability of winning.

This simple problem was then extended, to take into account the fact that each socket could actually return a different amount of charge, or each Bandit could payout a different amount of cash.

In both of these cases, rather than use the conventional approach to solving the Bandit Problem, where selections are based on the current mean values of the available choices, we took a Bayesian approach and sampled from a probability model.

These probability models were created using Conjugate Prior distributions. By updating the hyper-parameters of these distributions, we were able to create increasingly accurate models of the true distribution.

However, although we examined the problem of each socket returning a variable amount of charge, drawn from a normal distribution, we restricted the problem to use a fixed variance. In real world situations this is unlikely to be the case. Each socket is likely to have a different variance. And so, in this article, we’ll extend our examples to show how Thompson Sampling can be used when neither the mean nor the variance are known.

Using Conjugate Priors to Create Probability Models

When selecting an action, from a set of possible actions, Thompson Sampling takes a Bayesian approach. In our socket problem, we maintain a model of the probability distribution of each socket’s actual output. Then, when we want to decide which socket to try, we randomly choose a value from each of these models and select the socket with the highest returned value.

Obviously, before we have tried any of the sockets, our initial estimate of the true output could be well off the mark. However, as we start collecting data by using the sockets, we can begin to refine our estimate and gradually build up an increasingly accurate model of the true output.

In Bayesian terms our initial estimate is known as the prior probability and, after seeing some data and refining this estimate, we form the posterior probability. When both the prior and posterior have the same type of probability distribution they are said to be conjugate distributions. Additionally, when this happens, the prior is said to be the conjugate prior of the likelihood distribution (which can be thought of as the actual distribution of the data).

When the likelihood probability distribution has one of the more common distributions then the type of its conjugate prior can be found in the table of conjugate priors on Wikipedia. Additionally, for each type of conjugate prior, this table specifies a set of hyper-parameters that can be updated, after seeing some data, to move from the prior to the posterior. In this way the probability model becomes more accurate, letting us make a more informed choice as to the best action to take.

In this article we’ll look at how the conjugate prior can be created for a normally distributed likelihood when:

The variance is known, but the mean is unknown.
The mean is known, but the variance is unknown.
Both the mean and variance are unknown.

Known Variance, Unknown Mean

This is the case we looked at already, when we described the theory behind Thompson Sampling. Each of the sockets in our test system would return a charge from a normal distribution with variance equal to one and an unknown mean.

When the likelihood value (i.e. the actual value from the socket) is drawn from a normal distribution, and the variance is known but the mean value is not, then the table of conjugate priors on Wikipedia tells us that the Conjugate Prior is also a normal distribution with parameters mean ‘ μ₀‘ and total precision ‘ τ₀‘ (remember that the precision ‘ τ‘ (tau), is just one over the variance (precision τ = 1/variance). We use the precision version of the equation in preference to the variance version just because it simplifies things a bit.

And the update functions for the hyper-parameters are:

where:

‘ τ‘ is the true precision of the actual socket output, which in this case is fixed.
‘ n‘ is the number of times the socket has been tested.
‘ xᵢ‘ is the output received at each test ‘ i’ of this socket.
‘ μ₀‘ is the estimated mean (the mean of the distribution used to model the output).
‘ τ₀‘ is the total precision (the sum of all individual precisions).

A Python implementation of these update equations is shown below. This updates the hyper-parameters for the conjugate prior of a normal distribution with known variance and unknown mean:

Note, to avoid clutter, I’ve only shown the "update" function in isolation. In the actual code this a member function of a class which defines and initialises the variables etc. The full class can be found in the github repository.

One other thing to note is that we previously just used a variance of 1 (since this is the default in the numpy randn function). Now we allow any variance to be used although, since we’re considering the situation of a known variance, we just set the variance of the posterior distribution to be the value defined for the underlying distribution of the socket.

Using a single socket, with its underlying reward distribution initialised with a set mean and variance, we repeatedly sample from the true distribution and supply this value to the update function. The hyper-parameters of the conjugate prior are then updated and the posterior distribution moves to become a more accurate representation of the true distribution.

This is shown below for a socket with a true mean of 8 and a variance of 5. Here it can be seen how the estimated mean moves to become a very accurate estimate of the true mean. (The width of the distribution is already correct, since we know the variance.)

In the plots shown above, the estimated distribution is created using the mean value that has been calculated in the update function, when the posterior’s hyper-parameters are updated. However, the variance that is used is the socket’s true variance and not the variance of the conjugate prior (the update function actually computes the precision for the conjugate prior, but it’s simply a case of taking the reciprocal of this to get to the variance).

It’s worth looking at how the variance of the conjugate prior changes as the hyper-parameters are updated and this is shown below:

As we’ve already seen, as the number of trials increases, the estimated mean moves closer to the true mean (which in this case has a value of 8). In terms of the socket problem, each time a socket is tried, its estimate of the true mean output for that socket becomes more accurate.

We could just exploit this knowledge, and choose the socket with the current highest estimate, as done in the greedy algorithm. However, this would then suffer from the same limitation as the greedy algorithm: the lack of exploration. Thompson sampling overcomes this problem, and introduces exploration, by sampling from the prior probability distribution.

The variance of the distributions shown above, represented by the width of the curve, is a measure of the confidence in each prediction of the mean. As the distribution becomes narrower, so we become more confident in the accuracy of the prediction. Similarly, in the bandit system, as the confidence increases so the level of exploration will decrease.

Initially we want a lot of exploration and so we set the starting variance to produce an almost flat, uniform, distribution. This makes all values equally likely and so all sockets will be tried during the first time steps.

As actual data starts to be gathered from the socket the hyper-parameters of the conjugate prior are updated and the prior probability distribution will start to narrow.

The chances of selecting a badly performing socket will reduce since their distribution will shrink towards the left of the graph, while the distribution of sockets that perform well will start to move to the right. Therefore, the probability decreases of a sample value from a bad socket being higher than a sample value from a good socket.

Unknown Variance, Known Mean

So far we haven’t seen much that we didn’t already cover when we first looked at Thompson Sampling. However, the point of this extra article is to see how we can use Thompson Sampling when we know neither the mean nor the variance of the sockets, so let’s take our first step in that direction. In this case we’ll use a fixed value for the mean and use an unknown variance.

Consulting our trusty table of conjugate priors on Wikipedia, we see that a normal likelihood distribution, with a known mean, has a conjugate prior of either a gamma or inverse-gamma, depending on whether you want the model parameter to be the precision or the variance. As before, working with the precision is slightly easier, so we’ll go with the gamma distribution for our conjugate prior. The update equations for this are shown below.

As with the beta distribution, which we used for the Bernoulli bandit, the gamma distribution takes two parameters, ‘ α‘ (alpha) and ‘ β‘ (beta). In the beta distribution both of these parameters are termed ‘ shape‘ parameters since, unsurprisingly, they control the shape of the distribution. Whereas, in the gamma distribution, only ‘ α‘ is a shape parameter, with ‘ β‘ being defined as a ‘ rate‘ parameter which is related to the distribution’s variance. From the update equations it’s easy to see that ‘ α‘ is a count of the number of values sampled and ‘ β‘ is a measure of the sample variance.

Additionally, the beta distribution is only defined over the range zero to one (which makes it good for modelling events that have 2 possible outcomes), whereas the gamma distribution is defined over the range of all positive numbers.

As we did for the normal distribution with fixed variance and unknown mean, we can easily take these equations and form a Python update function, as shown below:

Note that here, the local variable ‘ n’ represents the number of sample values seen since the last time step, so is just set to a value of 1, whereas ‘self.n’ represents the total number of sample values seen since the start of the run, which is equivalent to ‘ ν‘ in the update equations.