Models & Math · ML Systems Lab

Multi-Armed Bandits: Thompson Sampling, UCB, and the Explore-Exploit Trade-off

A/B testing is batch experimentation: you wait, collect, decide. Bandits are online experimentation: you learn and adapt as traffic flows. The distinction matters enormously in high-traffic product environments where weeks of suboptimal routing is unacceptable. Thompson sampling and UCB are the two workhorses — each has a different failure mode.

The multi-armed bandit problem formalises the explore-exploit trade-off: you have K arms (treatments, variants, ads, recommendations), each with an unknown reward distribution. At each timestep you choose an arm, observe a reward, and update your beliefs. The objective is to maximise cumulative reward over T timesteps — which means minimising regret relative to always having chosen the best arm.

Why not just A/B test?

A/B testing assigns equal traffic to all variants for the entire experiment, then picks the winner. This is safe and statistically valid, but every impression on a losing variant during the experiment is wasted. For a website with 1M daily users running a 2-week test, a bad variant receiving 50% of traffic is a large cost. Bandits trade some statistical rigour for lower regret: they shift traffic toward better-performing arms as evidence accumulates.

Regret

Cumulative regret: R(T) = T * μ* − Σ_{t=1}^{T} μ(a_t), where μ* is the mean reward of the best arm. The Lai-Robbins lower bound (1985) shows that for any consistent algorithm, R(T) ≥ Ω(log T). This means some exploration is unavoidable — you cannot learn without occasionally pulling suboptimal arms. Good algorithms achieve O(log T) regret; bad ones achieve O(T) (linear regret = never converging on the best arm).

ε-Greedy (baseline)

With probability ε, explore (pull a random arm); with probability 1-ε, exploit (pull the current best arm by empirical mean). Simple, robust, but constant exploration regardless of uncertainty. ε should decay over time but the decay rate is a hyperparameter with no obvious setting.

Upper Confidence Bound (UCB1)

UCB1 (Auer et al., 2002): at each step, choose the arm that maximises μ̂_a + √(2 ln t / n_a), where μ̂_a is the empirical mean of arm a, t is the total steps, and n_a is the number of times arm a has been pulled. The second term is an uncertainty bonus — arms pulled less often get larger bonuses, driving exploration. UCB1 is deterministic, interpretable, and achieves O(log T) regret with tight constants. Failure mode: the optimism assumption can be violated when reward distributions are heavy-tailed or non-stationary.

Thompson Sampling

Thompson Sampling (Thompson 1933, rediscovered widely post-2010): maintain a Beta(α, β) prior over each arm's success probability (for Bernoulli rewards). After each observation, update the posterior: success → α += 1, failure → β += 1. At each decision point, sample a value from each arm's posterior and pull the arm with the highest sample. This is Bayesian: the randomness is in the posterior sampling, not in an ε parameter. Thompson Sampling achieves O(log T) regret empirically and often outperforms UCB in practice. It generalises to Gaussian, Poisson, and non-conjugate reward models (using Gaussian process or neural network priors for the latter).

Contextual bandits

Pure bandits ignore context. Contextual bandits — the framework behind almost every real recommendation and personalisation system — take a feature vector x_t (user features, item features, context) and choose an arm conditioned on x_t. LinUCB (Li et al., 2010) fits a linear reward model per arm with UCB-style exploration: reward_a(x) = x^T θ_a + α √(x^T A_a^{-1} x). Neural bandits replace the linear model with a neural network; the uncertainty estimate comes from the network's output layer or ensemble spread.

Non-stationarity

When reward distributions shift over time (product changes, seasonality, user drift), fixed posteriors become stale. Solutions: sliding window Thompson Sampling (forget old observations), discounted UCB (weight recent observations more heavily), or change-point detection triggering a posterior reset.

Bandit vs A/B test: when to use which

Use A/B testing when: you need strong statistical guarantees, experiments run for a bounded time, the cost of premature convergence is high (e.g., pharmaceutical trials). Use bandits when: traffic is high, experiment runtime is long, variants are many (multi-armed), and you care more about cumulative reward than causal inference validity. Hybrid: run a short A/B test for 10-15% of traffic to establish priors, then switch to Thompson Sampling for the remaining 85%.

Try on Colab: implement Thompson Sampling and UCB1 on the Bernoulli bandit problem with K=5 arms and true probabilities [0.1, 0.3, 0.5, 0.7, 0.9]. Plot cumulative regret over 10,000 steps for both algorithms plus ε-greedy with ε=0.1. Add a non-stationarity test: at t=5000, swap the rewards of arms 1 and 5. Observe how each algorithm adapts.

Continue interactively
Read this post inside ML Systems Lab — with Simplify toggle, interview Q&As, inline glossary, and the MLE Path forward pointer.
Open in MSL →