Contextual Bandits

Context-dependent rewards, LinUCB, LinTS, offline evaluation, NeuralTS

Standard A/B testing treats all users identically — it asks "which variant is better on average?" But users are different, and the best variant for one user segment may be the worst for another. Contextual bandits extend MAB by observing a feature vector at each round and learning which arm is best as a function of that context, not on average. This is the difference between finding the best treatment on average and finding the best treatment for each patient.

The naive alternative — train a supervised reward model and select greedily — fails because arms underrepresented in the logging policy have poorly calibrated reward estimates and are either permanently avoided or over-trusted. LinUCB's uncertainty bonus √(x^T A^{-1} x) is the key mechanism: it is largest exactly when the current context is far from previously observed data, targeting exploration where knowledge is genuinely lacking.

Key points

• Contextual bandit formulation: at round t, observe context x_t ∈ R^d, choose arm a_t ∈ {1,...,K}, observe reward r_t = f(x_t, a_t) + noise. Goal: minimise Σ_t [f(x_t, a*_t) − f(x_t, a_t)] where a*_t = argmax_a f(x_t, a). Every round is potentially different because the context changes — the same arm may be optimal for one user and suboptimal for another.
• Contextual bandits subsume classical A/B testing. A standard A/B test has context (user features) but ignores it during allocation — both variants are shown randomly regardless of user. A contextual bandit learns which variant is better for which user, discovering heterogeneous treatment effects while running the experiment. The allocation becomes personalised and the learnt policy is more valuable.
• LinUCB (disjoint model): assumes r_t = θ_a^T x_t + ε_t, a separate θ_a per arm. Estimates θ_a via ridge regression over (context, reward) pairs. UCB for arm a at context x = θ̂_a^T x + α√(x^T A_a^{-1} x) where A_a = X_a^T X_a + λI. The exploration bonus x^T A_a^{-1} x measures how far the current context is from previously observed contexts — large for novel contexts, small for familiar ones. This is the right signal: explore when the context is unfamiliar, exploit when it is not. This per-arm model is the *disjoint* variant; a *hybrid* variant shares a parameter vector across arms for common context features when arms have similar response patterns (see LinUCB In Depth).
• LinTS: Gaussian posterior over θ_a with mean θ̂_a and covariance σ²A_a^{-1}. At each round sample θ̃_a ~ N(θ̂_a, σ²A_a^{-1}) and select argmax_a θ̃_a^T x_t. Same regret bounds as LinUCB but empirically better, especially with informative priors. Prior misspecification matters here: a flat prior is the wrong choice when historical data is available to initialise θ̂_a.
• Offline policy evaluation (OPE): logged (context, arm, reward) data from a behaviour policy can evaluate new policies without online deployment. OPE is essential because online deployment is expensive and risky. Key estimators: Direct Method (DM) trains a reward model and evaluates offline — low variance but biased when the reward model is wrong for contexts the behaviour policy did not cover. Importance Sampling (IS) reweights logged rewards by π_e/π_b — unbiased but high variance when policies diverge. Doubly Robust (DR) combines both — consistent if either the reward model or propensities are correct.
• Greedy supervised learning is the wrong baseline. Train a reward model on logs and select the highest-predicted-reward arm each round. This fails because arms underrepresented in the logging policy have reward estimates with high uncertainty — and the greedy policy either perpetually avoids them (if their logged reward is low) or over-trusts them (if they happened to have high reward on the few occasions they were logged). No uncertainty bonus means no correction for this.
• Neural models improve reward estimation but break exact uncertainty quantification. NeuralTS uses last-layer Thompson Sampling: train a neural feature extractor to embed (context, arm) pairs, fix it periodically, run LinTS on the linear head over the embeddings. The neural network provides expressive reward modelling; LinTS provides principled uncertainty quantification on the last layer.
• Regret bounds: LinUCB achieves O(√(dT) ln K) regret where d is the context dimension. Context adds a √d factor over pure MAB O(√(KT)) — learning a linear reward function in d dimensions requires more exploration than learning K scalar arm means. Higher-dimensional or more complex reward functions require proportionally more data before the algorithm can exploit reliably.

Recap

• Formulation: observe context x_t, pick arm a_t, learn which arm is best as a function of context — not on average.
• Subsumes A/B testing: a standard A/B test has context but ignores it; contextual bandits learn heterogeneous treatment effects while running.
• LinUCB: $r_t = θ_a^T x_t + ε_t$; index = θ̂_a^T x + α√(x^T A_a^{-1} x); bonus large for novel contexts.
• LinTS: sample θ̃_a ~ N(θ̂_a, σ²A_a^{-1}); same bounds as LinUCB, empirically better with informative priors.
• Greedy supervised = wrong baseline: no uncertainty bonus → under-logged arms permanently avoided or over-trusted.
• OPE estimators: DM (low variance, biased), IS (unbiased, high variance), DR (consistent if either model is right).
• Regret O(√(dT) ln K): context adds √d over pure MAB — richer reward functions need more exploration.

Check your understanding