Thompson sampling

Overview

Thompson sampling chooses the next action by sampling from a posterior and acting on the sample. The posterior describes how good each option is, and a sample is one plausible version of that. Options that look good are chosen more often, while uncertain options still get tried 2310.11324 §3.2 2205.14090 §3.

Every use runs the same loop: hold a posterior, draw one sample, act on it, observe the result, and update the posterior. Uses differ in four things:

• The option. A finite arm, such as a prompt format 2310.11324 §3.2, or a point in a continuous domain, chosen by maximising a sampled function 2205.14090 §3.1. In HdSafeBO the option is a point in a latent space, limited to a trust region and a safe set 2412.20350 alg. 1, lines 6–8.
• The posterior. A Beta posterior per arm 2310.11324 app. A.4, a weighted mixture of Gaussian processes 2205.14090 §3, or one GP for the objective with a second GP for safety 2412.20350 §3.
• How far one observation reaches. FormatSpread updates only the arm that was pulled. EGP-TS reweights and updates every mixture component after each query 2205.14090 §3.1. HdSafeBO also resizes its trust region after each trial 2412.20350 §4.2.
• The candidate set. A finite list of arms (FormatSpread), the whole continuous domain (EGP-TS), or a box around the best safe point, filtered by an optimistic safety bound (HdSafeBO) 2412.20350 alg. 1, lines 6–7.

The FormatSpread paper introduces its search as Bayesian optimisation 2310.11324 §1, but its method is a multi-armed bandit 2310.11324 §3.2. For the general family of model-based search, see Bayesian optimization.

Variants

Beta-Bernoulli Thompson sampling, two-phase (FormatSpread)

For a finite set of arms. Each arm has a Beta posterior over its accuracy. A pull scores a mini-batch of B examples not yet seen for that arm, and the outcome is the number of correct answers. The prior is Beta(a, 5) with a = max(5x / (1 - x), 1.1), where x is the accuracy of the original format 2310.11324 app. A.4. The search runs twice with the total budget E split in half. The first run looks for the highest-accuracy arm and the second for the lowest, and the second run reuses the counts from the first 2310.11324 §3.2 2310.11324 app. A.4. Card: Quantifying Language Models' Sensitivity to Spurious Features in Prompt Design.

Baselines in the same comparison: UCB and naive sampling

The paper compares Thompson sampling with two other allocation rules. UCB scores each arm by its empirical accuracy plus c * sqrt(log(t) / N_i), with c = 2 2310.11324 §3.2. Naive sampling skips the bandit and gives every arm E / n pulls 2310.11324 §3.2. Both are baselines here, not Thompson variants, and are listed so that a reader can see what the Thompson result is compared against. Card: Quantifying Language Models' Sensitivity to Spurious Features in Prompt Design.

Mixture-of-GP posterior sampling for Bayesian optimization (EGP-TS)

For a continuous domain with an expensive black-box function. The surrogate is a weighted mixture of Gaussian processes, one per kernel in a fixed dictionary. The experiments use a uniform prior weight of 1/M 2205.14090 §3, §5. Each round the method draws a component from the mixture weights, draws a function from that component's posterior, and queries the maximiser of that function 2205.14090 §3, eq. 10, eq. 12. The draw uses random features, with D = 50 per component in the experiments 2205.14090 §3.1, §5. After the observation, every component's weight is multiplied by the likelihood of the new value and renormalised, and every component's posterior is updated 2205.14090 §3.1, eq. 13–15b. The paper refits kernel hyperparameters every few iterations, and every 50 in its experiments 2205.14090 §3.1 2205.14090 §5.

The paper also gives synchronous and asynchronous versions for several workers, with regret bounds for each 2205.14090 §3, alg. 1, alg. 2 2205.14090 §4. Card: Surrogate modeling for Bayesian optimization beyond a single Gaussian process.

Trust-region search with an optimistic safety filter (HdSafeBO)

For a controller with hundreds to thousands of parameters, where each trial must stay safe. The method maps the data into a latent space (PCA or an isometric VAE), fits one GP to the objective and one to safety, and searches a box of side l_t around the best safe latent point 2412.20350 §4 2412.20350 alg. 1, lines 4–6. A candidate stays only if its optimistic safety bound, μ + β σ, reaches the threshold 2412.20350 alg. 1, line 7. Thompson sampling then draws one sample of the objective over the surviving candidates and takes its maximum 2412.20350 alg. 1, line 8. The paper's experiments use β = 2 and a batch of 10 2412.20350 app. B.3.

The side length grows after enough successes and shrinks after enough failures. A success is a safe trial that improves the incumbent 2412.20350 §4.2 2412.20350 alg. 2. The paper gives a one-step probability bound when Φ(β) ≤ 1 − α, where Φ is the standard normal CDF 2412.20350 §4.1. It does not state α for its experiments 2412.20350 app. B.3. Card: Safe Bayesian Optimization for the Control of High-Dimensional Embodied Systems.

When to use which

The variants answer different questions. FormatSpread fits a finite list of about 320 options that you can only score through model outputs. EGP-TS fits a continuous domain, and the paper's benchmarks use 2 to 5 input dimensions 2205.14090 §5 2205.14090 tab. 3. HdSafeBO fits a continuous controller with a safety limit on each trial, in dimensions the paper puts in the hundreds to thousands 2412.20350 §3 2412.20350 §5.1. The choices below follow from the constraints each paper states.

Implementation sketch

The core is the same for every use: draw from the posterior, act on the draw, observe, update. The FormatSpread variant has one posterior per arm, so one pull updates one arm. The first block is that core. The two-phase wrapper is the FormatSpread setup, and the per-arm prior is the only input that depends on the original arm's accuracy.

# Beta-Bernoulli Thompson sampling over K arms  [§3.2] [alg. 1]
# Input:  arms, budget E, batch size B, objective 'max' or 'min',
#         prior (a, b) per arm, counts S (successes) and N (trials)
# Output: updated S and N

def thompson(arms, E, B, objective, prior, S, N):
    # Each step scores one mini-batch of B unseen examples  [§3.2]
    for step in range(E // B):
        # Draw one plausible accuracy per arm from its posterior  [§3.2]
        draws = {}
        for arm in arms:
            a, b = prior[arm]
            draws[arm] = sample_beta(a + S[arm], b + N[arm] - S[arm])
        # Play the arm with the highest draw for 'max', lowest for 'min'  [alg. 1]
        if objective == 'max':
            arm = argmax(draws)
        else:
            arm = argmin(draws)
        # Score the next unseen batch on the chosen arm  [§3.2]
        correct = score_batch(arm, B)
        S[arm] += correct
        N[arm] += B
    return S, N

# Two-phase search for the best and worst arm  [§3.2] [app. A.4]
# Input:  arms, original accuracy x, budget E, batch size B

def two_phase(arms, x, E, B):
    # Prior from the original arm's accuracy x  [app. A.4]
    a = max(5 * x / (1 - x), 1.1)
    prior = {arm: (a, 5) for arm in arms}
    S = {arm: 0 for arm in arms}
    N = {arm: 0 for arm in arms}
    # Phase 1: spend E/2 to find the highest-accuracy arm  [§3.2]
    S, N = thompson(arms, E // 2, B, 'max', prior, S, N)
    best = argmax_empirical(S, N)
    # Phase 2: reuse phase-1 counts, spend E/2 on the lowest  [app. A.4]
    S, N = thompson(arms, E // 2, B, 'min', prior, S, N)
    worst = argmin_empirical(S, N)
    return best, worst

The paper does not say how the final arm is picked. The official code takes the arm with the best or worst empirical accuracy among those scored, not the posterior mean Card (see "Paper vs code" on that card). Seed the random generator that draws the Beta samples. The official code seeds Python's random but not the NumPy generator that does the draws, so repeated runs can pick different arms Card (see "Pitfalls").

The EGP-TS variant differs in three places: the draw is a component plus a function sample, the argmax runs over a continuous domain, and the update reweights and conditions every component. One round, in our own words:

# One EGP-TS round, one query per round  [§3.1] [eq. 10–15b]
# comps: M Gaussian-process components, each with random features,
#        posterior mean theta and covariance Sigma
# w: mixture weights over the components, sum to 1  [§3]

def egp_ts_round(comps, w, domain, objective):
    # Sample a component from the mixture weights  [eq. 10]
    m = sample_categorical(w)
    # Draw feature weights from that component's posterior  [eq. 10]
    theta = sample_gaussian(comps[m].theta, comps[m].Sigma)
    # Query the maximiser of the sampled function over the domain  [eq. 12]
    x = argmax_over(domain, lambda z: comps[m].features(z) @ theta)
    y = objective(x)
    for k in range(len(comps)):
        # Reweight by the predictive likelihood of y  [eq. 13]
        mean, var = comps[k].predict(x)
        w[k] *= gaussian_pdf(y, mean, var)
        # Rank-one update of this component's posterior  [eq. 14, eq. 15a–b]
        comps[k].condition(x, y)
    w = w / sum(w)
    return x, y

The paper's text says refits rebuild every component from all acquired data 2205.14090 §3.1. The released code differs in how it refits, so the card lists the details Card (see "Paper vs code").

HdSafeBO adds a safe-candidate filter and a trust-region update around the same draw-and-observe loop. Its full loop is in the card's Method section, which is our own pseudocode of 2412.20350 alg. 1 and 2412.20350 alg. 2 Card.

Contradictions

• Batch size. The caption of the best-and-worst table gives B = 10, while the caption of the efficiency figure gives B = 20 2310.11324 §4.5. Pick one value and state it.
• UCB constant and form. The paper's UCB uses c * sqrt(log(t) / N_i) with c = 2 2310.11324 §3.2. The official code uses a different expression, and the UCB path is disabled in the main run. The two are not the same rule Card (see "Paper vs code").
• Number of tasks. The model comparison uses 53 tasks and the feature and probing analyses use 31 2310.11324 §4.3 2310.11324 §4.4. Check which count applies before comparing numbers.
• EGP-TS text and code. These are not disagreements between papers, but the EGP-TS paper's text and its released code differ in refits, initial weights and the robot task. Reproducing its numbers depends on which one you follow Card (see "Paper vs code").
• HdSafeBO success rule, paper and code. The paper counts a success as a safe trial that improves the incumbent 2412.20350 §4.2 2412.20350 alg. 2. The released code requires every sample in a batch to be safe, and its improvement test, as the card reads it, never passes Card (see "Paper vs code"). The card's reading was not run.

Papers

• Surrogate modeling for Bayesian optimization beyond a single Gaussian process: posterior sampling from a weighted mixture of Gaussian processes, with random features, for picking each query in Bayesian optimization.
• Quantifying Language Models' Sensitivity to Spurious Features in Prompt Design: a Beta-Bernoulli Thompson-sampling bandit with a two-phase search finds the best and worst of 320 prompt formats within a call budget, compared with UCB and naive sampling.
• Safe Bayesian Optimization for the Control of High-Dimensional Embodied Systems: HdSafeBO runs a trust-region Thompson-sampling search in a PCA or isometric-VAE latent space, keeping candidates whose optimistic safety bound passes a threshold.