Softmax Exploration for Treatment Allocation

From Reward Estimates to Randomized Allocation

A bandit allocation rule must turn what is currently known about treatments into a decision about whom to assign next. A greedy rule gives the next participant to a treatment with the largest estimated reward. That can be attractive when one treatment appears best, but it can also leave apparently inferior treatments with little opportunity to produce new information. Softmax allocation offers a graded alternative: it assigns higher probability to treatments with higher estimated rewards, while retaining some probability for the others.

The amount of randomization is controlled by a parameter called the temperature. A lower temperature makes differences in estimated reward matter more, concentrating allocation on the highest-scoring treatment. A higher temperature makes those differences matter less, spreading allocation more evenly. Temperature does not, by itself, say whether the resulting exploration is sufficient for a clinical trial; it determines how the specified scores are converted into probabilities.

Definition: Suppose $K$ treatments are eligible at a decision time, and let $\hat{r}_{k,t}$ be the estimated reward for treatment $k$ using information available just before that decision. For a finite positive temperature $\tau$, the softmax probability of assigning treatment $k$ is $$ p_{k,t}(\tau)=\frac{\exp(\hat{r}_{k,t}/\tau)}{\sum_{j=1}^{K}\exp(\hat{r}_{j,t}/\tau)}, \qquad k=1,\ldots,K. $$ The next treatment is drawn at random using these probabilities.

This definition assumes that the scores are finite, comparable numerical quantities on a common scale, and that the same positive temperature is used for all eligible treatments at that decision. The $p_{k,t}(\tau)$ are nonnegative and sum to one, so they define a probability distribution. With finite scores and $0<\tau<\infty$, every eligible treatment has a strictly positive probability in exact arithmetic. The rule maps scores to probabilities; it does not specify how the reward estimates were obtained or assert that those estimates are accurate.

In a clinical setting, a reward might be a binary indicator of a favourable outcome, so that its expected value is a success probability. It could also be a prespecified numerical utility. The reward definition and units matter: softmax compares numerical differences between scores, not just their ranking. As in the earlier material on defining rewards for clinical benefit, the reward must represent the clinical objective of interest if the allocation is intended to pursue that objective.

Worked Example: Calculating Three Treatment Probabilities

Suppose three hypothetical treatments have estimated success probabilities $0.70$, $0.55$, and $0.40$, and the selected temperature is $\tau=0.10$. Assume these estimates are available before the decision and are all expressed as success probabilities. Calculate the softmax probabilities.

Directly exponentiating the scores is unnecessary. Subtracting the largest score from every score leaves the probabilities unchanged, because the common exponential factor cancels from numerator and denominator. The adjusted scores are $0$, $-0.15$, and $-0.30$, giving

(p_{1,t},p_{2,t},p_{3,t}) = \frac{(1,\exp(-0.15/0.10),\exp(-0.30/0.10))} {1+\exp(-0.15/0.10)+\exp(-0.30/0.10)}.

Numerically, the unnormalized weights are $1$, $0.2231$, and $0.04979$, whose sum is $1.2729$. Therefore

(p_{1,t},p_{2,t},p_{3,t}) = (1/1.2729,\;0.2231/1.2729,\;0.04979/1.2729) \approx (0.7856,\;0.1753,\;0.03911).

The probabilities sum to approximately $1.000$. The treatment with estimated success probability $0.70$ is most likely to be selected, but the other two still have positive probabilities. These are probabilities for the next allocation under the specified rule, not probabilities that the estimated rewards are correct.

How Temperature Controls Exploration

The temperature is measured in the same units as the reward scores: the exponent $\hat{r}_{k,t}/\tau$ must be dimensionless. This provides a useful way to interpret it. A score difference that is large relative to $\tau$ produces a pronounced probability difference; a score difference small relative to $\tau$ produces probabilities closer to one another.

For any two eligible treatments $i$ and $j$, their probability ratio has a particularly simple form:

\frac{p_{i,t}(\tau)}{p_{j,t}(\tau)} = \exp\left(\frac{\hat{r}_{i,t}-\hat{r}_{j,t}}{\tau}\right).

The shared denominator cancels. Thus, if treatment $i$'s score exceeds treatment $j$'s by $\Delta$, its probability is $\exp(\Delta/\tau)$ times as large. For example, a score gap of $0.20$ gives a probability ratio of $\exp(2)\approx 7.389$ when $\tau=0.10$, but only $\exp(0.4)\approx 1.492$ when $\tau=0.50$. The ratio describes relative probabilities, not the probability of either treatment by itself; the other treatments also affect the normalization.

The limiting cases make the role of temperature clear. As $\tau$ grows without bound, each exponent approaches one, so the distribution approaches uniform allocation. As $\tau$ approaches zero, treatments with lower scores receive probabilities approaching zero. If there is one highest-scoring treatment, its probability approaches one. If $m$ treatments tie for the highest score, each of those $m$ treatments approaches probability $1/m$, and all lower-scoring treatments approach zero. Temperature zero is not a valid input to the formula.

Justification: Let $M=\max_k\hat{r}_{k,t}$. Rewriting each weight as $\exp((\hat{r}_{k,t}-M)/\tau)$ multiplies every original weight by the same positive constant $\exp(-M/\tau)$, which cancels in the ratio. For $\tau\to\infty$, each adjusted exponent approaches $\exp(0)=1$. For $\tau\to0^+$, the adjusted weight is $1$ for each maximizer and approaches $0$ for every lower score. Normalizing these limiting weights gives the stated uniform limit among maximizers.

Worked Example: Comparing Low and High Temperatures

For a hypothetical decision, suppose the estimated rewards are $0.80$, $0.60$, and $0.50$, measured as expected favourable-outcome proportions. Compare a low temperature of $0.02$ with a higher temperature of $0.50$. The scores are held fixed so the calculation isolates the effect of temperature.

At $\tau=0.02$, subtract $0.80$ from each score. The adjusted weights are $1$, $\exp(-10)\approx0.00004540$, and $\exp(-15)\approx0.000000306$. Their sum is about $1.0000457$, so the probabilities are approximately

(0.999954,\;0.00004540,\;0.000000306).

This is almost greedy allocation: the first treatment receives nearly all the probability. It is not exactly deterministic in exact arithmetic, although very small probabilities can be problematic in finite-precision software.

At $\tau=0.50$, the adjusted weights are $1$, $\exp(-0.20/0.50)\approx0.6703$, and $\exp(-0.30/0.50)\approx0.5488$. Their sum is $2.2191$, giving

(1/2.2191,\;0.6703/2.2191,\;0.5488/2.2191) \approx (0.4506,\;0.3021,\;0.2473).

The highest-scoring treatment remains the most likely, but the allocation is substantially more even. In both calculations, the reward estimates are treated as fixed inputs for the decision; these calculations do not account for their uncertainty.

Properties, Scaling, and Numerical Checks

Adding the same constant $c$ to every score changes no probability: the common factor $\exp(c/\tau)$ cancels. Multiplying every score and the temperature by the same positive constant also leaves the probabilities unchanged. These properties are useful both conceptually and computationally. In particular, changing reward units requires changing the numerical temperature by the same factor if the intended probabilities are to remain the same.

The response of a probability to temperature can also be made precise. Write $\bar{r}_{p,t}=\sum_j p_{j,t}\hat{r}_{j,t}$, the probability-weighted mean of the scores. Differentiating the softmax expression gives

\frac{\partial p_{i,t}}{\partial \tau} = \frac{p_{i,t}(\bar{r}_{p,t}-\hat{r}_{i,t})}{\tau^2}.

A treatment whose score is above the probability-weighted mean therefore loses probability as temperature rises; one whose score is below that mean gains probability. This describes the flattening effect locally. It does not mean that every individual treatment probability must change in the same direction: treatments between the extremes can respond according to their position relative to the weighted mean.

A practical implementation should calculate adjusted scores $(\hat{r}_{k,t}-M)/\tau$, where $M$ is the largest eligible score, before exponentiating. This avoids overflow when scores are large. It does not eliminate underflow for very low scores at a very small temperature. Check that the computed probabilities are finite, nonnegative, and sum to one within numerical tolerance. Also inspect the smallest probabilities: a value that rounds to zero cannot provide the positive allocation probability promised by the mathematical formula.

Worked Example: Tied Scores and a Change of Units

Suppose three hypothetical treatment scores are $0.60$, $0.60$, and $0.20$, with $\tau=0.20$. The first two treatments tie for the highest score. Subtracting $0.60$ gives weights $1$, $1$, and $\exp(-0.40/0.20)=\exp(-2)\approx0.1353$. The sum is $2.1353$, so

(p_{1,t},p_{2,t},p_{3,t}) \approx (1/2.1353,\;1/2.1353,\;0.1353/2.1353) \approx(0.4683,\;0.4683,\;0.06336).

The tied treatments receive equal probabilities, but not one-half each, because the lower-scoring treatment retains some probability. Now express the same scores in percentage points: $60$, $60$, and $20$. To preserve the same probabilities, use $\tau=20$, also in percentage points. The score-to-temperature ratios, and hence every softmax probability, are unchanged. Keeping $\tau=0.20$ after multiplying the scores by $100$ would instead make the distribution extremely concentrated. Temperature is not a unit-free tuning constant.

When Softmax Is Useful—and When It Is Not Enough

Softmax is useful when a design calls for a smooth randomized preference: estimated rewards influence allocation, yet nonmaximizing treatments retain positive probability at every finite temperature in exact arithmetic. Compared with greedy allocation, it avoids assigning probability zero solely because a treatment currently has a lower point estimate. Compared with uniform randomization, it uses the ranking and size of score differences. Compared with confidence-based rules such as UCB, however, the basic softmax formula uses reward estimates alone; it does not automatically increase a treatment's score because that estimate is uncertain.

The distinction matters in a clinical trial. A treatment with little data might have a modest estimated reward and high uncertainty, but softmax will not recognize that uncertainty unless it is explicitly incorporated into the score by a separately specified method. Conversely, a high estimated reward driven by a small, noisy sample can receive a large probability at low temperature. The formula is an allocation transformation, not a guarantee of good reward estimation or good statistical performance.

Temperature also does not enforce safety constraints or minimum allocation shares. If a treatment must not be assigned, it must be excluded from the eligible set before probabilities are computed; if a treatment must receive at least a specified share, ordinary softmax does not guarantee that condition. In particular, mathematical positivity is not an operational minimum: a probability such as $10^{-12}$ is positive but may be effectively zero over a finite trial.

Useful diagnostics include checking probabilities at plausible score vectors, assessing how much they change under small score perturbations, and examining sensitivity to the reward scale and temperature. If one treatment receives nearly all probability, determine whether that concentration is intended and whether the estimates supporting it are reliable. If probabilities are nearly uniform despite large clinically meaningful score differences, the temperature may be too high relative to the score scale. Neither pattern alone proves the design is wrong; each signals a choice that should be justified against the trial's objectives and constraints.

Common mistake: “A higher temperature means the best treatment gets a higher probability because the algorithm explores more.” This reverses the effect. Increasing temperature flattens the distribution, generally reducing the probability of an above-average-scoring treatment and increasing the probability of below-average-scoring treatments. Exploration is not the same as increasing the best treatment's probability.

Define the scores.
Specify the clinical reward and calculate comparable treatment reward estimates from information available at the decision.

Set the temperature on the reward scale.
Choose a finite $\tau>0$; interpret score gaps relative to $\tau$, rather than treating temperature as a probability.

Calculate and inspect probabilities.
Subtract the largest eligible score before exponentiation, normalize the weights, and check the resulting allocation probabilities and constraints.

Softmax is therefore best understood as a transparent score-to-probability map with one tuning parameter. It can express a continuum from near-greedy allocation to near-uniform allocation, but the appropriate point on that continuum depends on the score scale, uncertainty, safety requirements, and trial objective. Those design considerations cannot be settled by the softmax formula alone.

Check Your Understanding

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

Two treatments have estimated rewards $0.70$ and $0.50$. What is their softmax probability ratio when $\tau=0.10$?
For scores $0.40$ and $0.40$, what probabilities does softmax assign at any common positive temperature?
If every score is multiplied by $10$, what change to temperature preserves the probabilities?
Why does a positive softmax probability not ensure that a treatment receives a practically adequate number of participants?

Solutions. (1) The ratio is $\exp((0.70-0.50)/0.10)=\exp(2)\approx7.389$; the first treatment is about 7.389 times as likely as the second. (2) Equal scores give equal weights, so each treatment receives probability $0.5$. (3) Multiply the temperature by $10$ as well, because each score-to-temperature ratio remains unchanged. (4) A probability can be positive but arbitrarily small. In a finite trial, the expected number assigned at probability $p$ over $n$ decisions is only $np$ if that probability remains constant, and a small $np$ does not ensure adequate sampling. Softmax imposes no minimum allocation share.

Softmax Exploration for Treatment Allocation

What You'll Learn