Contextual Bandits: UCB and Thompson Sampling Explained

1. Introduction

The exploration-exploitation dilemma pervades sequential decision-making: how do we balance learning about alternative actions (exploration) with using current knowledge to maximize immediate reward (exploitation)? Multi-armed bandit algorithms give a principled framework for this tradeoff, treating the problem as sequential allocation where each decision also produces information that improves later decisions.

Traditional multi-armed bandits have a well-known limitation: they learn a single global policy applied to every decision. An epsilon-greedy algorithm choosing among five promotional offers learns which offer has the highest average conversion rate across all users, and then offers the same thing to a first-time mobile visitor at 2 PM and to a returning desktop customer at 9 AM. That context-blindness ignores signals that could change which offer is best.

Contextual bandits extend the framework by conditioning action selection on an observed context vector. Instead of one value estimate per action, a contextual algorithm learns a function from context to expected reward. Recommendations can adapt to user preferences, promotions to purchase history, and treatment assignments to individual characteristics.

Scope and Objectives

This paper explains the method and the decisions around using it. It does not present experiments, benchmark numbers, or case studies, and it makes no claim about how much any approach improves any metric. Where an example uses numbers, it says so and the numbers are made up for illustration.

It focuses on two families of contextual bandit algorithm. LinUCB and its variants are the confidence-bound approach: they construct confidence intervals around reward estimates and select the action with the highest upper bound. Thompson Sampling is the Bayesian approach: it maintains a posterior distribution over parameters and selects actions according to their probability of being best.

The objectives are threefold: (1) explain how each algorithm works and why it explores, (2) describe the implementation concerns that apply regardless of algorithm, and (3) give decision criteria for choosing between them based on application requirements.

Why This Matters Now

Several trends make contextual bandits relevant. Personalization is now expected in many consumer applications. Privacy constraints push organizations toward first-party contextual signals. Streaming and cloud infrastructure make online learning practical to build.

The shift from batch supervised learning to online sequential optimization changes the goal. Rather than training a static model on historical data and minimizing prediction error, a contextual bandit learns from the decisions it makes and optimizes cumulative reward. That fits problems where the task is to choose an action under uncertainty, not to predict passively.

2. Background

The Multi-Armed Bandit Problem

The classical problem has K slot machines (arms), each with an unknown reward distribution. At each step the player picks an arm and observes a random reward from that arm's distribution. The goal is to maximize cumulative reward over T rounds, which requires learning which arms pay more (exploration) while mostly pulling those arms (exploitation).

Regret measures performance: the gap between the cumulative reward the algorithm collected and what it would have collected by always playing the best arm. Sublinear regret, meaning regret that grows more slowly than T, indicates the algorithm eventually acts near-optimally. Classical algorithms such as UCB1 and Thompson Sampling have logarithmic regret guarantees under standard assumptions, while epsilon-greedy with a fixed epsilon keeps exploring at a constant rate and so has linear regret.

Limitations of Context-Free Approaches

Standard bandit algorithms assume each arm has a fixed expected reward that does not depend on external factors. In many applications that assumption is wrong. A "recommended for you" module may perform differently for new and returning users, for mobile and desktop sessions, or for morning and evening traffic. A context-free algorithm learns one ranking averaged across all of those populations.

The illustrative example in Section 5 shows how a context-free policy can be forced to treat two offers as equivalent when each is clearly better for a different segment.

Contextual Bandit Formulation

At each round t, the environment presents a context vector x_t in R^d describing the current state. The algorithm selects an action a_t from an action set A and observes a reward r_t drawn from a distribution conditioned on both context and action: r_t ~ P(r | x_t, a_t). The goal is still to maximize cumulative reward, but the best action now depends on context.

This formulation contains the classical problem as the special case where contexts carry no information. It also connects to supervised learning: if rewards are generated by some function r = f(x, a) + noise, then choosing well requires learning f. Contextual bandits can use regression, classification, and function approximation while keeping the online, sequential structure.

Existing Approaches

The literature has developed along two main lines: confidence-bound methods and Bayesian sampling. LinUCB, introduced by Li et al. (2010), assumes a linear reward model r = x^T θ_a + noise for each action a, keeps a ridge regression estimate of θ_a, and adds a confidence bonus derived from the uncertainty of that estimate. It selects the action with the highest upper confidence bound.

Thompson Sampling for contextual bandits (Agrawal and Goyal, 2013) keeps a posterior distribution over θ_a, typically Gaussian and updated by Bayesian linear regression. At each round it samples a parameter vector from each posterior and selects the action with the highest predicted reward under the sample.

Neural contextual bandits replace the linear model with a neural network and use techniques such as a neural-linear last layer, dropout-based uncertainty, or ensembles to estimate uncertainty. They can represent non-linear reward functions but need more data and compute, and their uncertainty estimates are harder to get right.

3. Methodology

What This Paper Is, and Is Not

The method of this paper is exposition. It draws on the published literature listed in the references and on standard derivations, and it explains the algorithms, their assumptions, and the engineering issues around them. It does not report an experiment, a deployment, or a comparison run on real traffic. Any statement about which approach "performs better" should be read as a statement about design trade-offs, to be checked on your own data.

LinUCB

For each action a, LinUCB maintains a matrix A_a = λI + Σ x x^T and a vector b_a = Σ r x, summed over the rounds where action a was chosen. The ridge estimate is θ_a = A_a^(-1) b_a. For a new context x, the score of action a is:

score(a) = x^T θ_a + α · sqrt(x^T A_a^(-1) x)

The first term is the predicted reward (exploitation). The second term is a confidence bonus that is large when the context lies in a direction where few observations have been made (exploration). The parameter α sets how much weight exploration gets. The algorithm picks the action with the highest score, so given the same state and context it always makes the same choice. Updates can use the Sherman-Morrison formula to maintain A_a^(-1) incrementally rather than inverting a matrix on each step.

Thompson Sampling

With a Gaussian prior and Gaussian noise, the posterior for θ_a is Gaussian with mean θ_a (the same ridge estimate) and covariance proportional to A_a^(-1). At decision time the algorithm draws one sample of θ_a from each action's posterior and picks the action whose sampled prediction x^T θ_a is highest. Actions with uncertain posteriors sometimes draw optimistic samples and get tried; as data accumulates the posterior tightens and exploration fades without a separately tuned bonus.

For high-dimensional contexts the full covariance matrix is expensive, so a diagonal approximation is a common compromise. It trades some statistical accuracy for lower memory and compute.

Context and Feature Engineering

It helps to sort candidate context features into five types: user attributes (demographics, account characteristics), behavioral history (past actions, engagement), temporal features (time of day, day of week, seasonality), session state (device, referral source, current activity), and action-specific features (item attributes, treatment characteristics). Typical practice is to normalize numeric features, add an intercept term, and add explicit interaction terms (for example user topic preference by item category) when the model is linear, since a linear model cannot form interactions on its own.

Keep the dimension in proportion to the data volume. With few observations per action, a large feature vector produces wide confidence intervals and slow learning, whichever algorithm is used.

Reward Definition and Evaluation

The reward is the single most consequential design choice, because the bandit will optimize exactly what it is given. Clicks are quick to observe but can diverge from the outcome that matters, such as revenue or retention. When the real outcome is delayed, a common pattern is to use a fast proxy for online updates and reconcile against the real outcome later (Section 6).

Useful quantities to track are cumulative reward, regret against a reference policy where one can be estimated, how long the policy takes to stabilize, and the variance of outcomes. Because the policy changes the data it sees, evaluation of a new policy on logged data needs the logged action probabilities and a counterfactual estimator such as inverse propensity scoring or a doubly robust estimator.

4. Key Points

5. Worked Example (Illustrative)

Illustrative example with made-up numbers. Suppose a checkout page can show one of three offers: a 10% discount, free shipping, or a bundled product. Suppose also that customers fall into two segments with opposite preferences, and that the conversion rates and segment sizes below are invented purely to show the mechanism. They are not measurements.

The averages follow from the invented inputs: for the discount, 0.6 × 0.12 + 0.4 × 0.03 = 0.084; for free shipping, 0.6 × 0.04 + 0.4 × 0.15 = 0.084.

In this made-up setting a context-free bandit sees the discount and free shipping as tied at 0.084 and has no basis to prefer either. A contextual bandit that can tell the segments apart would route price-sensitive visitors to the discount and convenience-focused visitors to free shipping, for an overall rate of 0.6 × 0.12 + 0.4 × 0.15 = 0.132 in this example. The point is the mechanism: when the best action differs by context, the best average action can be a poor choice for everyone. The size of the gap in real data depends entirely on how strongly context separates the best actions, which is why it needs to be measured and not assumed.

6. Production Considerations

This section is guidance on what a production system around a contextual bandit needs. It describes what to build, not what any particular system achieved.

When the Algorithm Choice Matters

Thompson Sampling is a natural fit when the application can tolerate stochastic decisions, exploration should taper off on its own, and feedback arrives quickly enough that the posterior is updated while it is still relevant. LinUCB is a natural fit when decisions must be explained or reproduced, when stakeholders want an explicit exploration knob, or when offline replay needs to be simple. Traffic volume, context dimension, and reward delay all affect how fast either method learns, so they are worth estimating before committing.

A small expected difference between algorithms does not justify added complexity by itself. Weigh any expected gain against the extra work: sampling infrastructure, posterior calibration, and harder offline testing for Thompson Sampling, or tuning of the exploration parameter for LinUCB.

Infrastructure and Operations

Logging and replay: Log context, chosen action, the probability with which it was chosen, and a timestamp for every decision. Rewards arrive later and must be joined back to the decision. Without logged propensities, counterfactual evaluation of a new policy is not possible. Plan storage for the log for as long as replay evaluation will be needed.

Model update pipelines: Choose between per-decision, micro-batch, and periodic updates according to how quickly the environment changes and how costly an update is. Micro-batches are a common middle path. Slow updates mean the policy lags behind changing preferences.

Delayed rewards: When the true outcome takes hours or days, a staged design helps: use an immediate proxy (a click, an engagement signal) for online updates, and apply the delayed outcome (a conversion, revenue) in periodic recalibration. Check regularly that the proxy and the real outcome still move together, since they can diverge.

Distribution shift monitoring: Context distributions drift as cohorts, seasons, and catalogs change. Monitor feature distributions against a baseline period, and reward distributions per context cluster, and define in advance what triggers a retrain or a feature pipeline review.

Technical Debt and Maintenance

Contextual bandits carry ongoing costs. Feature pipelines need review as data sources change. Parameters grow with the number of actions, since each action has its own model. Thompson Sampling with full covariance stores a d by d matrix per action, which grows quadratically with the number of features unless a diagonal approximation is used.

Offline evaluation systems, such as inverse propensity scoring and doubly robust estimators, become dependencies in their own right, and a bug in logging or replay can quietly corrupt the numbers used to judge a policy. Budget engineering time for maintenance, monitoring, and debugging reward feedback loops, not only for the first implementation.

Common Pitfalls

• Optimizing the wrong reward. A proxy such as clicks can improve while the outcome that matters does not.
• Missing propensities. Without the logged probability of each action, logged data cannot be reused to evaluate other policies.
• Overconfident initialization. A warm start with tight uncertainty stops exploration too early.
• Feedback loops. The policy shapes the data it later learns from, so a bias in early decisions can persist.
• Ignoring drift. A model fitted to last season's behavior keeps acting on it until something tells it otherwise.
• Too many features for the data. High dimension with few observations per action leads to wide intervals and slow learning.
• Interference between users. If one user's treatment affects another's outcome, standard bandit assumptions no longer hold.

7. Recommendations

These are suggestions drawn from the design considerations above, not conclusions from a study.

8. Conclusion

Contextual bandits move from learning one global policy to learning a context-dependent decision function. When the best action differs by context, that is what allows personalization. This paper has explained how that works, not measured how much it helps in any particular setting.

LinUCB and Thompson Sampling both explore in proportion to uncertainty, one through an explicit confidence bonus and the other through posterior sampling. The practical differences are mostly operational: determinism, explanation, offline replay, and the cost of calibration. The choice between them is best made on those criteria and then checked on your own data.

Feature engineering and reward design deserve attention before algorithm selection, because a bandit can only act on the signal it is given and optimizes the reward it is given. Cold start, delayed rewards, logging with propensities, drift monitoring, and offline evaluation are the supporting systems that determine whether the method is usable outside a notebook.

Contextual bandits treat uncertainty as information to act on, rather than as a problem to remove. That shift, from static prediction to sequential decision-making under uncertainty, is the core of the approach.