Multi-Armed Bandits for ML Systems
The exploration-exploitation tradeoff
A multi-armed bandit is the simplest formalization of a fundamental decision problem: you have $K$ options ('arms'), each with an unknown reward distribution, and at each round you pick one arm and observe only its reward. Should you keep pulling the arm that's performed best so far (exploit), or try a different arm that might secretly be better (explore)? The name comes from a row of slot machines ('one-armed bandits') — a gambler must decide which machine to keep playing without ever knowing the true payout rates.
This maps directly onto real ML systems: which ad creative to show, which article headline to display, which recommendation algorithm variant to serve, which price to offer — all are bandit problems where you learn from feedback while that same feedback has real-time cost.
Regret: the formal objective
A bandit algorithm's performance is measured by cumulative regret — the gap between the reward you got and the reward you would have gotten by always pulling the single best arm in hindsight:
$R_T = \sum_{t=1}^{T} \left( \mu^* - \mu_{a_t} \right)$
where $\mu^*$ is the true mean reward of the best arm and $\mu_{a_t}$ is the mean reward of the arm chosen at time $t$. A good algorithm has regret that grows sub-linearly in $T$ — typically $O(\sqrt{T})$ or $O(\log T)$ — meaning its average per-round regret shrinks to zero as it learns which arm is best. A naive algorithm that never explores, or explores forever, incurs linear regret.
Common bandit algorithms
| Algorithm | Mechanism | Regret bound | Notes |
|---|---|---|---|
| Epsilon-greedy | Explore uniformly at random with prob. $\epsilon$, else exploit best known arm | Linear unless $\epsilon$ decays | Simple, but wastes exploration on clearly-bad arms |
| UCB1 | Pick arm maximizing mean + confidence bonus $\sqrt{2\ln t / n_a}$ | $O(\log T)$ | Deterministic, optimism under uncertainty |
| Thompson Sampling | Sample from each arm's posterior, pick the arm with the highest sample | $O(\log T)$ (often better in practice) | Bayesian, naturally handles delayed/batched updates |
| Gradient bandit / softmax | Maintain preference scores, sample via softmax, update via policy gradient | Problem-dependent | Useful when rewards are noisy and non-stationary |
Epsilon-greedy and UCB in detail
Epsilon-greedy is the simplest baseline: with probability $\epsilon$, pick a uniformly random arm; otherwise pick the arm with the highest observed average reward. It's trivial to implement but explores inefficiently — it spends exploration budget equally on arms it already knows are bad.
UCB1 (Upper Confidence Bound) instead picks the arm maximizing $\bar{x}_a + \sqrt{\frac{2\ln t}{n_a}}$, where $\bar{x}_a$ is arm $a$'s observed mean reward and $n_a$ is how many times it's been pulled. The bonus term shrinks as an arm is pulled more (you're more certain about it) and grows with total rounds $t$ (giving under-explored arms a chance). This embodies the principle of optimism under uncertainty: act as if every arm is as good as plausibly possible given the data, and uncertainty itself drives exploration — no random coin flip needed.
import numpy as np
rng = np.random.default_rng(7)
true_means = np.array([0.30, 0.32, 0.45, 0.28]) # ground-truth conversion rates
n_arms = len(true_means)
T = 5000
def epsilon_greedy(epsilon=0.1):
counts = np.zeros(n_arms)
sums = np.zeros(n_arms)
regret = 0.0
for t in range(1, T + 1):
if rng.random() < epsilon or t <= n_arms:
arm = rng.integers(n_arms)
else:
arm = np.argmax(sums / np.maximum(counts, 1))
reward = rng.binomial(1, true_means[arm])
counts[arm] += 1
sums[arm] += reward
regret += true_means.max() - true_means[arm]
return regret
def ucb1():
counts = np.ones(n_arms)
sums = np.array([rng.binomial(1, m) for m in true_means], dtype=float)
regret = 0.0
for t in range(n_arms + 1, T + 1):
bonus = np.sqrt(2 * np.log(t) / counts)
arm = np.argmax(sums / counts + bonus)
reward = rng.binomial(1, true_means[arm])
counts[arm] += 1
sums[arm] += reward
regret += true_means.max() - true_means[arm]
return regret
def thompson_sampling():
alpha = np.ones(n_arms)
beta = np.ones(n_arms)
regret = 0.0
for _ in range(T):
samples = rng.beta(alpha, beta)
arm = np.argmax(samples)
reward = rng.binomial(1, true_means[arm])
alpha[arm] += reward
beta[arm] += 1 - reward
regret += true_means.max() - true_means[arm]
return regret
print("Epsilon-greedy cumulative regret:", epsilon_greedy())
print("UCB1 cumulative regret:", ucb1())
print("Thompson Sampling cumulative regret:", thompson_sampling())
Thompson Sampling: a Bayesian approach
Thompson Sampling maintains a posterior distribution over each arm's true reward rate — for Bernoulli rewards (clicked/didn't-click), a natural choice is a Beta distribution, updated conjugate-style: start with $Beta(1,1)$ (uniform prior) for each arm, and after observing a success or failure, update to $Beta(\alpha + 1, \beta)$ or $Beta(\alpha, \beta+1)$ respectively. At each round, draw one sample from each arm's current posterior and pick the arm with the highest sample. Arms with wide, uncertain posteriors occasionally produce high samples by chance, naturally driving exploration; as evidence accumulates, posteriors narrow and the algorithm converges to exploiting the best arm. It tends to perform very well empirically and extends cleanly to contextual and batched/delayed-feedback settings, which is why it's widely used in production ad and recommendation systems.
Bandits assume stationarity — most production systems aren't stationary
Bandits trade statistical rigor for speed of adaptation
What's Next?
Every arm in a classic bandit is treated as the same decision for every user — there's no notion of context. Real systems almost always have rich per-user, per-session information (device, history, time of day) that should influence which arm is best for *this* specific decision, not just on average. Extending the bandit framework to use that context is the natural next step.
I build these systems professionally.
Whether it's a RAG pipeline, analytics migration, or AI workflow — let's talk.