Annealed Softmax Greedy in Many-Armed Bayesian Bandits

arXiv:2605.31034 · cs.LG, cs.AI · Submitted 2026-08-17 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Next we'll be talking about the paper "Annealed Softmax Greedy in Many-Armed Bayesian Bandits".

Jane: The paper was written by William Overman and Mohsen Bayati from Stanford University.

Tom: Stay tuned as we take you through the paper and discuss its implications.

Title: Tom: Welcome back, everyone. Today we're looking at a paper that's been making the rounds on arXiv, and it's called "Annealed Softmax Greedy in Many-Armed Bayesian Bandits" by William Overman and Mohsen Bayati from Stanford.

Jane: And Tom, I have to say, the title alone is a mouthful, but the idea behind it is actually pretty intuitive once you break it down.

Tom: Right, so let's unpack it. A bandit problem is basically a slot machine with multiple arms, and you're trying to figure out which arm gives you the best reward. The twist here is that there are a lot of arms, and the rewards are uncertain.

Jane: Exactly. And the "annealed softmax greedy" part means the algorithm starts by exploring pretty randomly, and then gradually gets more and more greedy, like it's cooling down and committing to what it thinks is best.

Tom: So it's like a temperature schedule. Hot at first, then cooling off. And the paper shows that this simple approach, without any fancy uncertainty tracking, can actually perform really well.

Jane: And that's the surprising part, because there's a classic result showing that this kind of softmax approach can fail badly when you only have a few arms. But here, with many arms, it works.

Tom: So the number of arms changes everything. That's the key insight. And the authors connect this to something called "many-armed Bayesian bandits," which is a fancy way of saying there are tons of options and we have some prior beliefs about how good they are.

Jane: Right, and that prior belief is crucial. The paper assumes something called "β-regularity," which basically means there are plenty of arms that are almost as good as the best one. So even if the algorithm doesn't pick the absolute best arm, it's probably picking a really good one anyway.

Tom: And that's the magic. Because when you have a lot of near-optimal arms, being a little bit wrong doesn't cost you much. The regret, which is the difference between what you got and what you could have gotten, stays small.

Jane: So the title is really about showing that a simple, uncertainty-agnostic strategy can be near-optimal when the action space is rich enough.

Tom: And that has big implications, especially for how we train large language models, which we'll get into later. But for now, let's just say this paper is giving us a new lens on when simple strategies work.

Jane: And when they don't. Because the paper also shows that if you don't have that abundance of near-optimal arms, the whole thing falls apart.

Tom: So it's a conditional result. It works under specific conditions, and understanding those conditions is the real contribution here.

Jane: Exactly. And I'm curious to see how the authors actually prove this, because that's where the real meat is.

Tom: We'll get to that in a moment. But first, let's just appreciate the fact that a simple algorithm can be this powerful when the environment is structured the right way.

Summary: Jane: So Tom, we've talked about the title, but let's dig into what the paper actually does. The authors set up a bandit problem where each arm has a Bernoulli reward, meaning it either gives you a success or a failure, like a coin flip with a weighted coin.

Tom: And they put a prior on the arm means, which is their way of encoding what they expect the arms to be like before they start pulling them. The key assumption is that this prior has a "thick upper tail," meaning there's a decent chance of finding arms that are very close to the best possible reward.

Jane: Right, and that's the β-regularity condition. It's a mathematical way of saying that near-optimal arms are plentiful.

Tom: Then they run their algorithm, which they call ASG, short for Annealed Softmax Greedy. It starts by pulling each arm once to get an initial estimate, and then at each step, it picks an arm with probability proportional to the exponential of the temperature times the empirical mean.

Jane: So it's like a popularity contest where the empirical mean is the popularity score, and the temperature controls how much the crowd follows the leader versus spreading out.

Tom: And the main result is a bound on the Bayes regret, which is the expected difference between the reward you got and the reward you could have gotten if you'd known the best arm from the start.

Jane: And the bound is something like Õ(m + T/m), where m is the number of arms and T is the time horizon. So if you choose the number of arms to be about the square root of the time horizon, you get regret that's about the square root of T.

Tom: And that's near-optimal. It's the same rate that you'd get with a much more sophisticated algorithm that explicitly tracks uncertainty.

Jane: So the message is that you don't need to be smart about exploration if you have enough arms that are good enough. The randomness in the softmax is enough to keep you safe.

Tom: And the proof works by showing that even when the softmax picks a non-optimal arm, that arm is likely to be another near-optimal one, because there are so many of them.

Jane: It's like if you're in a city with a hundred great restaurants, and you randomly pick one that isn't the absolute best, you're still going to have a great meal.

Tom: That's a great analogy, Jane. And the paper formalizes that intuition with a "leakage" term that measures how much probability mass goes to bad arms, and they show that under β-regularity, that leakage is small.

Jane: So the summary is: simple algorithm, strong assumptions, near-optimal results. And the assumptions are exactly what you'd expect in a world where good options are common.

Tom: And that's what makes this paper exciting, because it suggests that a lot of the complexity in bandit algorithms might be unnecessary in the right settings.

Jane: But we should also mention that the paper is careful to note that this is a stylized model. It's not a direct analysis of, say, GRPO or other language model training methods.

Tom: Right, it's a structural analogy. The math is about bandits, but the intuition carries over.

Jane: And that's what we'll explore next, because the connection to language model training is where things get really interesting.

Improvements: Tom: So Jane, we've covered the basics, but what really gets me excited is the connection to how we train large language models, especially with something called RLVR, which stands for reinforcement learning with verifiable rewards.

Jane: And that's where the paper's "structural analogy" comes in. The idea is that when you train a model to solve math problems or write code, you sample multiple completions, check which ones are correct, and then increase the probability of the correct ones.

Tom: And that's basically what the softmax greedy algorithm does. It's reweighting based on empirical success, without explicitly tracking uncertainty about which completions are good.

Jane: And the paper's insight is that this can work if the base model already has a decent chance of producing correct completions. That's the β-regularity condition in disguise.

Tom: So if the model's output distribution has a thick upper tail, meaning there are plenty of good completions in there, then just reweighting toward the observed good ones is enough.

Jane: And that connects to a debate in the field about whether RLVR actually expands the model's capabilities or just redistributes probability mass within what it already knows.

Tom: The paper doesn't settle that debate, but it provides a mechanism for why redistribution alone can be effective. If the good completions are already there, you don't need to explore new ones.

Jane: And the paper even mentions pass@k, which is a metric that measures the probability that at least one of k samples is correct. That's a direct probe of the upper tail of the model's output distribution.

Tom: So the improvement the paper suggests is not a new algorithm, but a new way of thinking about when simple reweighting works.

Jane: And it also suggests a cautionary note. If the base model doesn't have that thick upper tail, then the same approach can fail, just like the softmax algorithm fails with a small number of arms.

Tom: So it's a double-edged sword. The conditions that make it work are also the conditions that make it unnecessary, because the good stuff is already there.

Jane: But it's not entirely unnecessary, because reweighting can improve the probability of getting a good completion on the first try, which is what pass@one measures.

Tom: Right, so it's about sharpening the distribution, not expanding it. And that's a subtle but important distinction.

Jane: And the paper's simulations back this up. They show that the softmax algorithm matches greedy in the many-armed regime, and that prior-anchored versions can do even better when the prior is informative.

Tom: So the practical takeaway is that if you have a decent base model, simple reweighting can get you a lot of value without complex exploration schemes.

Jane: And that's a message that resonates with practitioners, because simple methods are easier to implement and debug.

Tom: But we should also mention the limitations. The paper is about non-contextual bandits, and language model training is highly contextual and sequential.

Jane: So it's a structural analogy, not a direct proof. The authors are clear about that.

Tom: And that's the honest way to present it. The math is solid for the bandit setting, and the intuition transfers, but the full extension is an open problem.

Jane: Which brings us to the big question: what does this mean for the future of AI training?

Conclusion: Tom: So we've spent this whole episode on "Annealed Softmax Greedy in Many-Armed Bayesian Bandits," and I think we've only scratched the surface.

Jane: We really have. The core message is that a simple, uncertainty-agnostic algorithm can achieve near-optimal regret when there are many near-optimal arms.

Tom: And that's a big deal because it challenges the assumption that you always need sophisticated exploration strategies.

Jane: The paper gives us a clean mathematical framework for understanding when simple reweighting works, and it connects that to the practical world of language model training.

Tom: And while the analogy is structural, not direct, it gives us a useful lens for thinking about why RLVR can be effective even without explicit uncertainty tracking.

Jane: The simulations also point to some interesting directions, like prior-anchored variants that use the prior more aggressively, and the cautionary tale of what happens when the prior is wrong.

Tom: That stress test with the misspecified prior was a good reminder that Bayesian methods are only as good as their priors.

Jane: And that's a lesson that applies beyond bandits. If your model has a strong but wrong prior, it can be worse than having no prior at all.

Tom: So the paper is a nice combination of theory, simulation, and practical intuition. It's not the final word, but it's a valuable contribution.

Jane: And it leaves us with open questions. How does this extend to contextual settings? What happens with sequential decisions? Those are the next steps.

Tom: For now, though, we should say goodbye to this paper and thank the authors, William Overman and Mohsen Bayati, for their work.

Jane: And thank you, our listeners, for joining us. We hope this gave you a new perspective on when simple strategies can be surprisingly effective.

Tom: Next up, we'll be looking at a paper on efficient attention mechanisms, so stay tuned.

Jane: Until then, keep exploring, and remember that sometimes the simplest approach is the right one.

Tom: See you next time.

William Overman, Mohsen Bayati

Stanford University

cs.LG, cs.AI

Submitted: 2026-08-17

Updated: 2026-08-18

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 84/100

Key concepts

Many-Armed Bandit Problem
This is a decision problem where an agent must choose among multiple options (arms) to maximize reward. The rewards are uncertain, meaning the agent must balance exploring less-known options with exploiting the best known option.
Annealed Softmax Greedy (ASG)
This algorithm starts by exploring randomly, like a hot temperature, and gradually becomes more committed to good options as it 'cools down.' It selects arms based on a probability related to their observed mean reward.
Beta-Regularity
This is a crucial assumption stating that there are many options whose rewards are very close to the absolute best reward. This abundance of near-optimal choices allows simple, uncertainty-agnostic algorithms to perform effectively.

Terminology

Summary

Summary

This paper investigates whether annealed softmax (Boltzmann) policies, which are agnostic to epistemic uncertainty, can achieve near-optimal Bayes regret in the many-armed Bayesian Bernoulli bandit setting. The motivation comes from reinforcement learning with verifiable rewards (RLVR) and group-based policy optimization methods such as GRPO, which update a stochastic policy by sampling multiple completions per prompt and increasing the policy's probability on those with higher reward, regularized by a KL penalty toward a reference policy. These updates do not include explicit mechanisms that track epistemic uncertainty.

The paper studies a stylized explanation for why such uncertainty-agnostic updates can nevertheless be effective. It analyzes an annealed softmax (Boltzmann) policy that selects actions according to a softmax of empirical mean rewards in a many-armed Bayesian Bernoulli bandit. Under a linear upper-tail condition on the prior (the β = 1 case of β-regularity), which implies an abundance of near-optimal arms, the paper proves that annealed softmax greedy achieves Bayes regret Õ(m + T/m), and in particular Õ(√T) when the number of arms scales as m = Θ(√T). This is the near-optimal Bayes regret rate in this regime, attained also by empirical-mean greedy.

The main contributions are:

  1. Bayes regret guarantee. Under the linear upper-tail condition on the prior (the β = 1 case of β-regularity), annealed softmax achieves Bayes regret Õ(m + T/m), yielding Õ(√T) when m = Θ(√T). This is the near-optimal Bayes regret rate in this regime (also attained by empirical-mean greedy), achieved using only empirical-mean scores, with no optimism term, no posterior sample, and no per-arm confidence interval. The analysis extends to general β > 0, with optimized rate Õ(m + T (log T /m)(1/β)).

  2. Mechanism. Annealed softmax differs from greedy only by placing some probability mass on arms that are not currently empirically best. Under β-regularity, many arms maintain empirical means close to the optimum throughout learning, so when softmax samples an arm other than the empirically best, that arm tends to be another near-optimal one rather than a clearly inferior one. This contrasts with the failure mode studied by Cesa-Bianchi et al. (2017), where the same kind of policy can suffer linear regret with a small number of arms.

  3. Structural analogue in RLVR. The β-regular prior condition has a direct analogue in RLVR: a base policy with non-negligible probability of producing a correct completion plays the role of β-regularity, with pass@k probing the upper tail of the policy's completion distribution. In that setting, repeated sampling surfaces correct completions reliably, and subsequent reweighting can improve pass@1 without an explicit uncertainty-aware exploration rule. The paper notes that this is only a structural analogy and not a formal result about GRPO; the theorem lives in the non-contextual bandit model, and extending the analysis to the contextual or sequential settings where GRPO operates is an open problem.

The paper formalizes the model as follows: Fix a time horizon T and consider m arms indexed by i ∈ [m]. Each arm has an unknown mean µi ∈ [0, 1] drawn i.i.d. from a prior Γ on [0, 1]. Conditional on µi, each arm i has an i.i.d. sequence of Bernoulli rewards. The prior is assumed to be β-regular, meaning there exist constants 0 < c0 ≤ C0 < ∞ and ε0 ∈ (0, 1) such that for all ε ∈ (0, ε0], c0 ε β ≤ Γ([1 − ε, 1]) ≤ C0 ε β.

The main theorem (Theorem 5.2) states: Assume Bernoulli rewards, and assume the prior Γ is 1-regular with regularity constants (c0, C0, ε0). Fix any δ ∈ (0, δ0] with δ0 ≤ min 1/8, ε0/8. Run ASG with any nonnegative nondecreasing schedule ηt. Then there exist constants c, C > 0 such that BRT,m(ASG) ≤ C [m + Tδ + m(1 + log(1/δ)) + (1/δ) Σ t=1 T exp(−ηt δ) + T exp(−c m δ)]. In particular, choosing δ = min δ0, A log(T ∨ 2)/m and ηt = (cη/δ) log(t ∨ 2) with cη > 1 and A > 1/c, one obtains BRT,m(ASG) = Õ(m + T/m), and for m = Θ(√T) one obtains BRT,m(ASG) = Õ(√T).

The proof decomposes regret into several interpretable terms: the terms m and Tδ are the same coarse baseline contributions that already appear in greedy-style many-armed analyses; the logarithmic term m(1 + log(1/δ)) reflects the cost of integrating over the upper tail; the summation term (1/δ) Σ exp(−ηt δ) is the softmax-specific leakage term, controlled by the cooling schedule; and the exponentially small term T exp(−cmδ) corresponds to the bad event that too few of the m arms are near-optimal. In this sense, ASG is essentially greedy up to a leakage penalty that becomes small under suitable annealing.

The paper also includes simulation experiments on many-armed Bernoulli bandits. The simulations test three prior regimes: uninformative (uniform Beta(1,1) per arm), informative (prior mean is a monotone function of the true mean), and misspecified (prior built from a noisy version of the true mean). The results show that Scheduled ASG matches Pure Greedy across the range where the theorem operates, which is the direct empirical validation of Theorem 5.2. The paper also introduces prior-anchored variants of Greedy and Scheduled ASG that replace the empirical mean with the Beta posterior mean, which attain near-zero regret under an informative prior but fail badly under an uninformative prior. A stress test with a strongly misspecified prior shows that Thompson Sampling with a confidently wrong prior can be worse than with no prior at all.

The paper concludes that the analysis identifies a different regime from the classical negative results for Boltzmann exploration: the many-armed Bayesian one with sufficient upper-tail mass, in which the same softmax policy achieves near-greedy Bayes regret. The difference is the geometry of the action space induced by the prior, not the policy class. Two limitations remain: the model is narrow (i.i.d. Bayesian Bernoulli bandits, with no context, no shared structure across actions, and no sequential state dynamics), and the guarantees are Bayesian rather than minimax, relying on upper-tail regularity of the prior.

Improvements for AI systems

Based on the paper, here are specific improvements to AI systems and what the improved system can do:

Improvement: Modify GRPO/RLVR training to use an annealed softmax over empirical reward scores instead of hard argmax or fixed-temperature softmax, with a logarithmic cooling schedule (η t = (c η/δ) log(t∨2)).

What the improved system can do:

  • Achieve near-optimal policy improvement without explicit epistemic uncertainty tracking, provided the base policy has non-negligible probability of correct completions (the β-regular condition)

  • Automatically balance exploration vs. exploitation during training: early in training, the softmax spreads probability across many completions; later, it concentrates on high-reward ones

  • Avoid the failure mode of premature convergence to suboptimal completions when many near-optimal solutions exist in the base model's support

Abstract

Reinforcement learning with verifiable rewards (RLVR) and group-based policy optimization methods such as GRPO update a stochastic policy by sampling multiple completions per prompt and increasing the policy's probability on those with higher reward, regularized by a KL penalty toward a reference policy. These updates do not include explicit mechanisms that track epistemic uncertainty. This paper studies a stylized explanation for why such uncertainty-agnostic updates can nevertheless be effective. We analyze an annealed softmax (Boltzmann) policy that selects actions according to a softmax of empirical mean rewards in a many-armed Bayesian Bernoulli bandit. Under a linear upper-tail condition on the prior (the beta=1 case of beta-regularity), which implies an abundance of near-optimal arms, we prove that annealed softmax greedy achieves Bayes regret (m + T/m), and in particular (sqrt T) when the number of arms scales as m = (sqrt T). This is the near-optimal Bayes regret rate in this regime, attained also by empirical-mean greedy. Under beta-regularity, many arms maintain empirical means close to the optimum throughout learning, so when softmax samples an arm other than the empirically best, that arm tends to be another near-optimal one rather than a clearly inferior one. By contrast, with a small number of arms, the same kind of softmax policy can suffer linear regret. The result also provides a structural analogy to RLVR, where a base policy with a non-negligible probability of producing a correct completion plays the role of beta-regularity.

Sources

Related papers