Thompson sampling for A/B tests in 40 lines of JavaScript (and the trap nobody mentions)

#javascript #statistics #abtesting #algorithms

A classic A/B test splits traffic 50/50 until it reaches a fixed sample size. A multi-armed bandit moves traffic toward the version that is winning while the test runs, so fewer visitors see the loser. Thompson sampling is the simplest good way to do it.

Here is a working implementation with zero dependencies, a simulation, and the one trap that makes a lot of bandit "wins" impossible to prove.

The idea in one paragraph

For each version, keep a belief about its conversion rate: a Beta distribution built from its visits and conversions, Beta(1 + conversions, 1 + visits - conversions). For every new visitor, draw one plausible rate from each version's belief and show the version with the highest draw. Versions you are unsure about still get explored; versions that are clearly worse get picked less and less.

The code

JavaScript has no built-in Beta sampler, so we build one from two Gamma draws (Marsaglia-Tsang):

function normal() {
  const u = 1 - Math.random(), v = Math.random();
  return Math.sqrt(-2 * Math.log(u)) * Math.cos(2 * Math.PI * v);
}

// Gamma(k, 1), Marsaglia-Tsang; boost trick for k < 1.
function gamma(k) {
  if (k < 1) return gamma(k + 1) * Math.random() ** (1 / k);
  const d = k - 1 / 3, c = 1 / Math.sqrt(9 * d);
  for (;;) {
    let x, v;
    do { x = normal(); v = 1 + c * x; } while (v <= 0);
    v = v * v * v;
    const u = Math.random();
    if (u < 1 - 0.0331 * x ** 4 || Math.log(u) < 0.5 * x * x + d * (1 - v + Math.log(v))) return d * v;
  }
}

const beta = (a, b) => { const x = gamma(a); return x / (x + gamma(b)); };

// Pick a version for the next visitor.
function pick(arms) {
  let best = 0, bestDraw = -1;
  arms.forEach((arm, i) => {
    const draw = beta(1 + arm.conversions, 1 + arm.visits - arm.conversions);
    if (draw > bestDraw) { bestDraw = draw; best = i; }
  });
  return best;
}

// Chance that B beats A, by simulation.
function chanceBeats(a, b, draws = 20000) {
  let wins = 0;
  for (let i = 0; i < draws; i++)
    if (beta(1 + b.conversions, 1 + b.visits - b.conversions) >
        beta(1 + a.conversions, 1 + a.visits - a.conversions)) wins++;
  return wins / draws;
}

Simulate it

Two versions, true conversion rates 2.0% and 2.5%, 40,000 visitors:

const truth = [0.02, 0.025];
const arms = truth.map(() => ({ visits: 0, conversions: 0 }));
for (let t = 0; t < 40000; t++) {
  const i = pick(arms);
  arms[i].visits++;
  if (Math.random() < truth[i]) arms[i].conversions++;
}
console.log(arms, chanceBeats(arms[0], arms[1]));

Two runs on my machine:

A: 7,745 visits, 160 conversions   B: 32,255 visits, 801 conversions   P(B>A) = 0.984
A: 2,278 visits,  41 conversions   B: 37,722 visits, 975 conversions   P(B>A) = 0.992

The bandit found the better version and sent it most of the traffic. A 50/50 test would have shown the worse page to 20,000 people.

The trap: you starve your control

Look at the second run: the original page got only 2,278 visits. That is fine when you only care about earning more during the test, but it is a problem when someone asks "did the change actually help, and by how much?" With so few visits on the original, the estimate of the original's rate is wide, and seasonality or an ad change can fool you.

The fix is boring and important: keep a floor on the original. Never let it drop below some share of traffic (10% to 20% is common), and only call a winner after a minimum number of visits per version (a few hundred at least) and a high bar (for example 95% chance to beat the original).

function pickWithFloor(arms, floor = 0.1) {
  if (Math.random() < floor) return 0; // arm 0 = your unchanged page
  return pick(arms);
}

When to use which

Fixed 50/50 test: you need a precise size of effect for a big decision (pricing, a redesign) and will reuse the learning for months.
Bandit with a floor on the original: everyday page copy, where earning more during the test matters and you still want proof against the unchanged page.

If you want the sample-size side of this, I wrote that up too: How many visitors does an A/B test actually need? and a longer comparison of bandits vs A/B tests.

Disclosure: I build Luria, which runs Thompson sampling with a floor on the original page for store split tests. The code above is standalone and MIT-style; use it anywhere.

DEV Community