KL-UCB for Bernoulli Outcomes

Why Use a Divergence-Based Index?

For a treatment with binary outcomes, the observed success proportion is a natural estimate of its success probability. But an allocation rule also needs to account for uncertainty: a treatment with few observations may warrant more exploration than one with many observations, even when their observed proportions are similar. KL-UCB does this by asking how large the treatment’s success probability could plausibly be under a specified divergence budget.

The resulting quantity is an index, not a posterior mean or a probability that a treatment is best. At each decision, the rule selects a treatment with the largest index. The index is deliberately optimistic: it is the most favorable success probability that remains within the chosen divergence constraint.

Definition: For $p,q\in[0,1]$, the Bernoulli Kullback–Leibler divergence from success probability $p$ to success probability $q$ is $$ d(p,q)=p\log\left(\frac{p}{q}\right)+(1-p)\log\left(\frac{1-p}{1-q}\right). $$ Terms with coefficient zero are defined by continuity. In particular, $d(0,q)=-\log(1-q)$, $d(1,q)=-\log(q)$, and $d(p,p)=0$. If a positive-probability event under $p$ has probability zero under $q$, the divergence is infinite.

Model and Index Definition

Suppose there are $K$ treatments, also called arms. Each outcome is coded as a binary reward: success is 1 and failure is 0. For each arm $a$, assume its successive outcomes are independent Bernoulli random variables with a fixed success probability $\mu_a$. Outcomes from different arms are independent as well. The allocation rule may use past observed outcomes to choose the next arm, but it observes only the outcome for the arm actually assigned.

At decision time $t$, let $N_a(t)$ be the number of observed outcomes for arm $a$ before that decision, and let $S_a(t)$ be its number of successes. When $N_a(t)>0$, the empirical success proportion is $\hat p_a(t)=S_a(t)/N_a(t)$. Assume each arm has been observed at least once before using the index. For definiteness, use the divergence budget $B_t=\log t$, for $t>1$. Other positive budgets are possible, but the budget must be specified when calculating an index.

Definition: The Bernoulli KL-UCB index for arm $a$ at time $t$ is $$ U_a(t)=\sup\left\{q\in[\hat p_a(t),1]: N_a(t)d\bigl(\hat p_a(t),q\bigr)\le B_t\right\}. $$ The KL-UCB allocation rule assigns the next participant to an arm with the largest $U_a(t)$. If several arms tie, a prespecified tie-breaking rule can select among them.

The constraint compares the empirical success probability with a candidate value $q$. The factor $N_a(t)$ makes the constraint more restrictive as an arm accumulates observations, while a larger budget permits a wider range of candidate values. Restricting the search to $q\geq\hat p_a(t)$ is appropriate because the index seeks an upper optimistic value.

Why the Constraint Uses KL Divergence

For $n$ Bernoulli observations with $s$ successes, write $\hat p=s/n$. The likelihood at candidate success probability $q$ is $L(q)=q^s(1-q)^{n-s}$. When $0<s<n$, the likelihood is maximized at $q=\hat p$. Taking the log-likelihood ratio between that maximum and the candidate $q$ gives

\log\left(\frac{L(\hat p)}{L(q)}\right) =s\log\left(\frac{\hat p}{q}\right) +(n-s)\log\left(\frac{1-\hat p}{1-q}\right) =n\,d(\hat p,q).

Thus the constraint $n\,d(\hat p,q)\leq B_t$ keeps candidate probabilities whose log-likelihood loss relative to the empirical fit is no more than the budget. This likelihood identity motivates the construction; it does not by itself make the set an exact confidence interval with a fixed coverage probability.

For $q>\hat p$, differentiation gives

\frac{\partial}{\partial q}d(\hat p,q) =\frac{q-\hat p}{q(1-q)}>0, \qquad 0\leq\hat p<q<1.

Consequently, for $0\leq\hat p<1$, the divergence increases as $q$ moves above $\hat p$. The index is the unique upper solution of $n\,d(\hat p,q)=B_t$ when that solution lies below 1; the divergence tends to infinity as $q$ approaches 1. For $\hat p=1$, every observed outcome was a success and the index is 1. For $\hat p=0$, the constraint simplifies to $-n\log(1-q)\leq B_t$, so

U=1-\exp\left(-\frac{B_t}{n}\right).

Except for the boundary case $\hat p=1$, one can compute the index by bisection on $[\hat p,1]$. The monotonicity just shown guarantees that this search is well behaved.

Worked Example: Inverting the Constraint

Consider a Bernoulli arm with $n=10$ observations and 4 successes, so $\hat p=0.4$. For this calculation, take a divergence budget $B=1.000$. The upper index solves $10d(0.4,q)=1.000$, or $d(0.4,q)=0.1000$.

Substituting into the divergence formula, $$ d(0.4,q)=0.4\log\left(\frac{0.4}{q}\right) +0.6\log\left(\frac{0.6}{1-q}\right). $$ Numerical solution gives $q\approx0.6213$. At this value, $d(0.4,0.6213)\approx0.1000$, so $10d(0.4,0.6213)\approx1.000$. The index is therefore approximately 0.6213, an optimistic success-probability value 0.2213 above the observed proportion.

The constraint can also reject an implausibly large candidate. For example, $d(0.4,0.8287)\approx0.4608$, so $10d(0.4,0.8287)\approx4.608$, well above the budget of 1.000. Thus 0.8287 is not feasible, despite being a possible probability in the abstract.

Worked Example: No Successes Observed

Suppose an arm has $n=4$ observations and no successes. Use a divergence budget $B=0.8000$. Here $\hat p=0$, and the constraint is $4[-\log(1-q)]\leq0.8000$. Dividing by 4 gives $-\log(1-q)\leq0.2000$, hence $$ q\leq1-\exp(-0.2000)=0.1813. $$ The KL-UCB index is 0.1813. Although the observed success proportion is zero, the index is positive because four failures do not establish that the true success probability is exactly zero.

Checking the result, $-\log(1-0.1813)\approx0.2000385$, and multiplying by $n=4$ gives approximately 0.8002, as required.

Worked Example: Many Observed Successes

Suppose another arm has $n=5$ observations, 4 of which are successes, and use budget $B=1.000$. Its empirical proportion is $\hat p=0.8$, and the constraint requires $d(0.8,q)\leq0.2000$. Solving $$ d(0.8,q)=0.2\log\left(\frac{0.2}{1-q}\right) +0.8\log\left(\frac{0.8}{q}\right)=0.2000 $$ gives $q\approx0.9653$. Substitution gives $d(0.8,0.9653)\approx0.2001$, so $5d(0.8,0.9653)\approx1.0003$. This index is high because the observed success proportion is high, and the small sample still permits substantial upward uncertainty.

The value is not a claim that the arm’s true success probability is 0.9653. It is the largest candidate probability permitted by this divergence budget and these data.

Computing and Checking an Index

In practice, an index need not have a closed-form solution except in boundary cases. A stable numerical procedure brackets the answer between the empirical proportion and 1, evaluates $n\,d(\hat p,q)$, and repeatedly narrows the bracket until the constraint is met to the desired tolerance. Handle $q=1$ and $\hat p\in\{0,1\}$ with the boundary definitions rather than evaluating expressions that contain undefined logarithms.

Summarize the observed outcomes.
For the arm, compute $n=N_a(t)$ and $\hat p=S_a(t)/n$.

Set the budget.
Use the specified $B_t$ for the decision time, such as $B_t=\log t$.

Find the largest feasible candidate.
Solve $n\,d(\hat p,q)=B_t$ for $q\geq\hat p$, using the boundary formula when $\hat p=0$ or the index 1 when $\hat p=1$.

Verify the result.
Check that the index is at least $\hat p$, no greater than 1, and that $n\,d(\hat p,U_a(t))$ is approximately $B_t$ unless the boundary case applies.

Several monotonicity checks are useful. For fixed $\hat p$ and $B_t$, increasing $n$ tightens the constraint and cannot increase the upper solution. Increasing the budget cannot decrease the index. These checks are especially useful for catching a misplaced logarithm, an incorrect arm count, or a numerical search that returns the wrong side of the root.

Comparison, Assumptions, and Limitations

A square-root UCB rule, derived from concentration inequalities, adds an uncertainty term to the empirical proportion. KL-UCB instead uses the Bernoulli likelihood geometry: the allowed increase depends on the divergence between the observed proportion and each candidate probability. The distinction can matter near 0 or 1, where a symmetric square-root adjustment may not reflect the asymmetric range of Bernoulli probabilities. Both are optimistic-index approaches, but their indices and assumptions should not be treated as interchangeable.

Approach	How optimism is represented	Important distinction
KL-UCB for Bernoulli outcomes	Largest $q$ satisfying $n\,d(\hat p,q)\leq B_t$	Uses the Bernoulli divergence and respects the probability boundaries.
Square-root UCB	Empirical proportion plus an uncertainty allowance	Uses a concentration-based radius rather than the Bernoulli KL constraint.
Empirical-proportion ranking	Ranks arms by $\hat p$ alone	Does not incorporate a separate uncertainty allowance into the ranking.

The standard model requires a stable success probability for each arm and conditionally independent Bernoulli outcomes within each arm. Adaptive assignment is compatible with the model when each arm’s outcome sequence remains Bernoulli with that fixed probability; the next assignment can depend on observed history. In a clinical setting, the binary reward and its timing must be defined consistently, and the rule must receive the outcome information it needs before the next decision.

If success probabilities change over calendar time, participants have dependence not captured by the model, outcomes are delayed or missing, or the endpoint is not binary, the stated index may no longer be appropriate without modification. The rule also does not by itself impose clinical safety constraints, account for covariates, or guarantee that every treatment continues to be sampled. Those are separate design requirements.

Common Mistakes and Diagnostics

Using the divergence in the wrong direction. The constraint here is $d(\hat p,q)$, not $d(q,\hat p)$. KL divergence is generally asymmetric, so reversing the arguments changes the index.
Forgetting the sample size. The constraint is $n\,d(\hat p,q)\leq B_t$. Omitting $n$ can produce a much wider index than intended.
Calling the index a confidence probability. The index is a constrained upper value. It is not the probability that the arm is best, nor does the construction alone imply that the true parameter lies below it with a specified probability.
Accepting a numerical root without checking it. Substitute the proposed index into $n\,d(\hat p,q)$. If the result substantially exceeds the budget, the candidate is infeasible.
Applying Bernoulli KL-UCB to an undefined reward. The success event must be specified, and the assumed binary outcomes and stable arm probabilities must be defensible for the decision being made.

Key takeaway: For Bernoulli outcomes, KL-UCB assigns each arm the largest success probability $q$ whose divergence from the empirical proportion fits within the sample-size-scaled budget. The index is computed by a one-dimensional monotone search and interpreted as an optimistic score, not as a posterior probability.

Check Your Understanding

Answer each question, then compare your reasoning with the solutions below.

An arm has 6 successes in 10 observations. What empirical proportion enters its KL-UCB constraint?
For $\hat p=0$, $n=5$, and budget $B=0.5$, calculate the index.
For fixed $\hat p$ and budget, what happens to the index when the sample size increases, and why?
Does an index of 0.7 mean there is a 70% probability that the treatment is best? Explain.

Solutions. (1) $\hat p=6/10=0.6$. (2) The boundary formula gives $U=1-\exp(-B/n)=1-\exp(-0.1)\approx0.09516$. (3) The index cannot increase: at any candidate $q>\hat p$, increasing $n$ increases $n\,d(\hat p,q)$, so fewer upper candidates satisfy the fixed budget. (4) No. The index is the largest candidate success probability allowed by the divergence constraint; it is not a probability that the treatment is best.

Approach	How optimism is represented	Important distinction
KL-UCB for Bernoulli outcomes	Largest \(q\) satisfying \(n\,d(\hat p,q)\leq B_t\)	Uses the Bernoulli divergence and respects the probability boundaries.
Square-root UCB	Empirical proportion plus an uncertainty allowance	Uses a concentration-based radius rather than the Bernoulli KL constraint.
Empirical-proportion ranking	Ranks arms by \(\hat p\) alone	Does not incorporate a separate uncertainty allowance into the ranking.

What You'll Learn