Multi-armed bandit in the DS interview

Why interviewers ask about MAB

The multi-armed bandit is the canonical exploration-vs-exploitation problem, and interviewers love it because it splits candidates in under five minutes. The candidate who only read about A/B tests reaches for fixed 50/50 splits and a two-week timer. The one who shipped ranking or pricing reaches for regret minimization and adaptive allocation.

The question shows up in three flavors. Recommendation teams (Netflix, Spotify, DoorDash) ask contextual bandits because their entire ranker is one. Growth teams (Notion, Linear, Stripe) ask MAB to test whether you understand when not to use it. ML platform teams want the math: UCB1 regret of O(sqrt(K T log T)), the Beta-Bernoulli conjugate update, the discount factor for non-stationary environments.

The setup and the regret you are minimizing

You have K arms. Each arm a has an unknown reward distribution with mean mu_a. At each round t = 1...T you pick arm a_t, observe reward r_t, repeat. Maximizing cumulative reward is equivalent to minimizing regret:

Regret(T) = T * mu_star  -  sum over t of E[r_{a_t}]

where mu_star = max over a of mu_a. The bandit problem is not about finding the best arm — it is about paying as little as possible while you find it.

Pure exploration wastes traffic on bad arms forever and accumulates linear regret. Pure exploitation locks onto whichever arm got lucky in the first 50 pulls. Every algorithm below is a different recipe for splitting the difference.

Epsilon-greedy

The first algorithm everyone implements, and the one every interviewer expects you to dismiss for the wrong reasons. With probability epsilon you pull a uniformly random arm; with probability 1 - epsilon you pull the empirically best arm so far.

import numpy as np

def epsilon_greedy(epsilon, n_arms, n_rounds, pull):
    counts = np.zeros(n_arms)
    means = np.zeros(n_arms)
    for t in range(n_rounds):
        if np.random.rand() < epsilon:
            a = np.random.randint(n_arms)
        else:
            a = int(np.argmax(means))
        r = pull(a)
        counts[a] += 1
        means[a] += (r - means[a]) / counts[a]
    return means, counts

A fixed epsilon = 0.1 keeps exploring 10% forever — fine for a non-stationary world, terrible for a stationary one. The textbook fix is decaying epsilon, typically epsilon_t = min(1, c * K / t), which gives logarithmic regret instead of linear. Most teams ship epsilon = 0.05 and call it a day because operational simplicity beats the theoretical gain.

UCB1

Upper Confidence Bound replaces the "epsilon" knob with an uncertainty bonus. You always pick the arm with the highest mean + bonus, where the bonus shrinks as you pull the arm more often.

UCB_a(t) = mean_a + sqrt( 2 * ln(t) / n_a )

The intuition is optimism in the face of uncertainty. An arm pulled three times has a fat confidence interval; one pulled three thousand has a thin one. UCB1 prefers fat intervals until they shrink, then commits. Auer 2002 shows this gives regret of order sqrt(K T log T) with no tuning parameters, which is why it dominates academic benchmarks.

def ucb1(n_arms, n_rounds, pull):
    counts = np.zeros(n_arms)
    means = np.zeros(n_arms)
    for a in range(n_arms):
        r = pull(a)
        counts[a] = 1
        means[a] = r
    for t in range(n_arms, n_rounds):
        bonus = np.sqrt(2 * np.log(t) / counts)
        a = int(np.argmax(means + bonus))
        r = pull(a)
        counts[a] += 1
        means[a] += (r - means[a]) / counts[a]
    return means, counts

UCB1's hidden weakness: it is deterministic. Two parallel instances with the same history pick the same arm and can stampede traffic. In production you add jitter or use UCB-Tuned, which incorporates per-arm variance.

Thompson sampling

The Bayesian answer, almost always the strongest empirical performer for Bernoulli rewards. Each arm has a posterior over its mean. You sample once from each posterior, pick the arm with the highest sample, observe the reward, update.

def thompson_bernoulli(n_arms, n_rounds, pull):
    alpha = np.ones(n_arms)
    beta = np.ones(n_arms)
    for t in range(n_rounds):
        samples = np.random.beta(alpha, beta)
        a = int(np.argmax(samples))
        r = pull(a)
        alpha[a] += r
        beta[a] += 1 - r
    return alpha, beta

For 0/1 rewards the Beta-Bernoulli conjugate prior makes the update one line. The Beta posterior is wide early (exploration) and tight late (exploitation), and randomness in sampling means parallel instances naturally diverge — no stampede.

Empirical regret of Thompson sampling matches UCB1 in theory and beats it in most simulations, especially when K is moderate (4-50 arms) and rewards are sparse. For dense rewards with large K (1000+ arms), LinUCB usually wins because the linear structure cuts the effective dimensionality.

Contextual bandits

Plain MAB assumes every user is identical. Contextual bandits drop that assumption: at each round you also see a context vector x_t (user features, time of day, device), and your policy is a_t = pi(x_t). Reward distributions now depend on both arm and context.

The canonical algorithm is LinUCB, which assumes the expected reward is linear in the context: E[r | a, x] = x^T theta_a. You maintain per-arm ridge regression coefficients and a confidence interval over x^T theta_a. You pick the arm with the highest upper bound. The regret stays sub-linear in T and scales with the context dimension d rather than the number of arms K, which is why it powers most production ranking systems with thousands of items.

At round t:
  Observe context x_t.
  For each arm a:
    Compute mean estimate  m_a = x_t^T theta_a
    Compute upper bound    u_a = m_a + alpha * sqrt(x_t^T A_a^-1 x_t)
  Pull arm with max u_a, observe reward r_t.
  Update A_{a_t} and b_{a_t} for ridge regression.

Neural bandits replace the linear payoff with a deep network and use techniques like NeuralUCB or Neural Thompson sampling. The trade-off is the same as everywhere in ML — more flexibility, more data needed, harder to debug. In an interview, name LinUCB first and only mention neural variants if asked.

MAB vs A/B testing

This comparison is the single most-asked follow-up. Memorize the framing.

The cleanest one-liner: A/B tests answer "which is better?", bandits answer "right now, which should I serve?". If your PM needs a green-or-red ship decision in two weeks, you want an A/B test with a power calculation, MDE, and a single Z-test at the end. If you are picking which of 12 push notification templates to send right now and you will keep doing it forever, you want a bandit.

A subtler answer that wins senior rounds: bandits bias your estimate of treatment effect, because high-performing arms get more traffic, which inflates their observed lift versus their true lift. If a stakeholder asks "how much better was arm B?" after a bandit run, the empirical mean is not an unbiased estimator. You either need inverse propensity weighting or you have to run a confirmatory A/B test on the winner.

Common pitfalls

The most damaging mistake is running a bandit when you needed an A/B test. A PM asks "did the new checkout convert better?" and a junior DS spins up Thompson sampling. The bandit will route traffic to the winner, but it will not produce an unbiased treatment effect or a p-value. The fix is boring — use 50/50 with a power calculation, then ship a bandit on top of the winner if you want continuous optimization.

Reward delay breaks the basic loop. Bandits assume you observe r_t before choosing a_{t+1}. If your reward is a 7-day subscription conversion, your bandit is deciding on traffic from a week ago. The fix is either a fast proxy reward (click within an hour, add-to-cart within a session) and accept the bias, or offline batch learning with importance weighting. Wrong delay strategy is the most common production failure of bandit systems.

Ignoring non-stationarity is the next failure mode. Reward distributions drift — what worked at 9am does not work at midnight, what worked before a competitor launched does not work after. Vanilla UCB1 commits to a stale winner and never revisits. The fix is a sliding window over the last W observations or exponential discounting with gamma = 0.99, which trades steady-state regret for responsiveness.

Skipping context is the pitfall that separates a recommendation engineer from a data scientist. "Users on average prefer template A" hides that iOS users prefer A and Android users prefer B by a 3x margin. A plain MAB pools across segments and ships a 60/40 compromise that loses to either segment-aware policy. If you have features, use contextual bandits — even a 5-feature LinUCB beats a 12-arm plain MAB on most personalization workloads.

The last trap is mis-tuned exploration. A fixed epsilon = 0.3 wastes 30% of traffic forever; epsilon = 0.001 locks onto a local optimum in the first hour. UCB1 has no tuning knob; Thompson sampling is self-tuning through the posterior. If you must use epsilon-greedy in production, decay it as c * K / t with c between 1 and 10.

Related reading

• A/B testing complete guide
• Bayesian methods in DS interviews
• Collaborative filtering interview prep
• Bayesian A/B testing
• Sequential testing explained simply

If you want to drill MAB and A/B questions every day with timed feedback, NAILDD ships 500+ DS interview problems including bandit math, regret bounds, and production trade-offs.

FAQ

When should I pick MAB over an A/B test?

Pick MAB when the cost of serving a known-worse variant is high and a slightly biased treatment effect estimate is acceptable. Classic examples are content ranking, push notification templates, ad creative rotation, and dynamic pricing within a band. Pick A/B when you need a clean ship decision, an unbiased effect size, or compliance review — pricing changes that touch revenue forecasts, onboarding redesigns, anything that goes to legal.

Why is Thompson sampling so popular if UCB1 has the same regret bound?

Two reasons. Thompson sampling is randomized, so parallel workers do not stampede onto the same arm and traffic distributions stay smooth. And in practice Thompson sampling has lower regret on most realistic workloads despite the same theoretical bound, because the bound is worst-case and the posterior adapts faster to the actual reward distribution. The trade-off is needing a conjugate prior or sampling from a more expensive posterior.

How do I handle delayed rewards?

Three options, in order of preference. Use a faster proxy reward correlated with the true reward — click instead of purchase, session start instead of D7 retention. Batch updates on a daily or hourly cadence and accept that you are optimizing yesterday's traffic. Or switch to an offline contextual bandit with logged propensities, training on logs and deploying a fresh policy daily.

What is regret and why does it matter more than accuracy?

Regret is the cumulative gap between what you earned and what you would have earned with the best arm. Accuracy is binary — did you pick the right arm? — but in MAB you might pick wrong a thousand times and still have low regret if all arms are close, or pick right 99% of the time and have huge regret if your 1% mistakes are on a much-worse arm. The algorithm optimizes regret, so "accuracy" is a category error here.

Can I use MAB with thousands of arms?

Plain MAB scales linearly in K, so 1000 arms means 1000 pulls just to try each once. The fix is to cluster arms or use contextual bandits where arms are described by features and the policy generalizes across them. LinUCB and neural bandits handle 10k+ arms cleanly because they parameterize the reward function instead of memorizing per-arm statistics.

Is this official information?

No. The recipes summarize standard literature — Auer et al. 2002 on UCB1, Thompson 1933 on Bayesian sampling, Li et al. 2010 on LinUCB, Russo and Van Roy 2018 on Thompson analysis. Production combines these with engineering constraints the papers do not cover.