ArXiv: 2402.05878

🎯 Pitch

Forget adaptive eliminationβ€”a fixed, prior-informed sampling plan computed before any data arrives can match or beat Sequential Halving in Bayesian best-arm identification. The authors prove the first prior-dependent error bounds for linear and hierarchical bandits, and their G-optimal design variant outperforms adaptive baselines on real-world data, challenging the assumed necessity of adaptivity.

1. Executive Summary

This paper introduces PI-BAI (Prior-Informed BAI), a fixed-budget best-arm identification algorithm that leverages prior-dependent fixed allocations β€” computed before interacting with the environment β€” rather than adaptive elimination strategies. The work establishes prior-dependent upper bounds on the expected probability of error across multi-armed, linear, and hierarchical bandits β€” with the bounds for linear and hierarchical settings representing the first such results β€” and demonstrates that its bound is tighter than that of BayesElim (Atsidakou et al., 2022) by a factor of roughly log₃/β‚‚β‚‚(K) due to avoiding elimination phases. The algorithm achieves consistent and robust empirical performance across synthetic and real-world (MovieLens) settings, with the G-optimal design allocation variant emerging as the strongest overall, establishing that non-adaptive prior-informed strategies can match or exceed adaptive elimination-based methods when a well-specified prior encodes sufficient information about the arm structure β€” a boundary condition that directly addresses the question of adaptivity in fixed-budget BAI.

2. Context and Motivation

The Core Problem: Does Adaptivity Matter in Fixed-Budget BAI?

The fundamental question this paper tackles is whether adaptive, elimination-based algorithms β€” which form the backbone of most fixed-budget best-arm identification methods β€” are always necessary, or whether carefully designed non-adaptive strategies conditioned on prior information can achieve competitive or superior performance. This question is not merely academic; it goes to the heart of algorithm design for pure exploration problems.

Fixed-budget BAI operates under a hard constraint: the learner has exactly nn rounds (pulls) to interact with the environment, after which it must output a single arm believed to be optimal. The standard metric is the probability of error (PoE) β€” the likelihood that the returned arm is not actually the best. Unlike regret minimization, where the learner's cumulative reward matters, BAI treats all rounds as purely instrumental; the only thing that counts is the quality of the final decision. This makes BAI "a problem closer to statistical hypothesis testing than to regret minimization" (Section 1), and as the authors note, it is "often viewed as a more difficult question in general" (citing Carpentier and Locatelli, 2016; Degenne, 2023).

The dominant paradigm in fixed-budget BAI has been adaptive elimination: start with all KK arms, pull them in rounds, progressively eliminate arms that appear suboptimal, and concentrate the remaining budget on the survivors. Sequential Halving (SH; Karnin et al., 2013) exemplifies this approach. Under this paradigm, "adaptivity" means that the sampling strategy responds to observations β€” arms with higher estimated means receive more pulls in later rounds. The intuition is compelling: why waste budget on arms that are clearly inferior?

However, recent theoretical developments have cast serious doubt on whether this intuition translates into genuine necessity. Two results are particularly destabilizing:

  1. Wang et al. (2024) showed that for two-armed Bernoulli bandits, there exists no algorithm strictly better than uniform sampling across all problem instances when comparing frequentist PoE. In other words, for every adaptive algorithm, there is some bandit instance where uniform sampling performs better β€” and vice versa.

  2. Degenne (2023) established that there is no superior adaptive algorithm in several BAI problems, including Gaussian BAI. This is a formal impossibility result: adaptivity provides no universal advantage.

These findings collectively "question the role of adaptivity for fixed-budget BAI problems" (Section 1, citing Qin, 2022). If adaptive methods cannot dominate uniform sampling, then the field's heavy investment in increasingly sophisticated elimination schemes may be misdirected β€” or at minimum, requires a more nuanced understanding of when adaptivity helps.

The Bayesian Twist: Prior Knowledge Changes the Calculus

The results above are cast in the frequentist framework, where each bandit instance is a fixed, unknown parameter vector ΞΈ\theta, and guarantees must hold uniformly over some class of instances. But the paper argues that introducing a Bayesian perspective β€” where ΞΈ\theta is drawn from a known prior distribution P0P_0, and performance is measured by expected PoE over this prior β€” fundamentally alters the landscape.

The motivating example in Section 1 crystallizes the intuition:

"Consider a scenario with 3 arms, where prior information strongly suggests that either of the first two arms is more likely to be optimal than the third. How much budget should then be allocated to the seemingly suboptimal third arm? Moreover, if there is greater confidence in the first arm compared to the second, what is the optimal budget distribution between them?"

In a frequentist setting without prior knowledge, the third arm must receive non-trivial budget because, for all the learner knows, it could be optimal. But in a Bayesian setting with an informative prior, the third arm's posterior probability of being optimal might be negligible. Spending budget on it is wasteful. Similarly, if the prior expresses more uncertainty about arm 1 than arm 2 (through a larger prior variance Οƒ0,12>Οƒ0,22\sigma_{0,1}^2 > \sigma_{0,2}^2), then arm 1 deserves more samples to resolve that uncertainty β€” even if arm 2 has a higher prior mean.

This reframes the question: the issue is not whether to be adaptive in response to observations, but rather how to design prior-dependent fixed allocations that optimally distribute the budget before any data is collected. The paper calls this "the following general question: How to choose a prior-dependent fixed allocation of the budget to guarantee a low expected PoE?" (Section 1).

Where Existing Bayesian BAI Falls Short

The most directly relevant prior work is the BayesElim algorithm of Atsidakou et al. (2022), which introduced a Bayesian version of Sequential Halving for KK-armed Gaussian bandits. BayesElim incorporates prior knowledge in two ways: (1) allocation weights in each elimination phase prioritize arms with higher observation noise variance, and (2) the final decision uses posterior means rather than sample means. However, the paper identifies three critical limitations:

1. Restrictive prior assumptions. The analysis of BayesElim requires equal prior variances across arms (Οƒ0,i2=Οƒ02\sigma_{0,i}^2 = \sigma_0^2 for all i∈[K]i \in [K]). This is not a minor technical condition β€” it is baked into the proof strategy. By coupling uniform sampling within each elimination round with equal prior variances, the posterior variances Οƒ^n,i2\hat{\sigma}_{n,i}^2 become identical for all surviving arms at the end of each round. This homogeneity allows the authors to directly import frequentist concentration results from Karnin et al. (2013). But it rules out exactly the kind of heterogeneous prior uncertainty that makes the motivating example interesting. As Appendix A explains:

"We believe it is not possible to adapt their analysis to heterogeneous prior variances, where each arm would have a different posterior uncertainty at the end of each round."

2. Reliance on frequentist averaging. The proof technique of Atsidakou et al. (2022) first bounds the frequentist PoE for a fixed ΞΈ\theta as P(Jnβ‰ iβˆ—(ΞΈ)∣θ)≀B(ΞΈ)\mathbb{P}(J_n \neq i^*(\theta) \mid \theta) \leq B(\theta), and then averages this bound over the prior: Eθ∼P0[B(ΞΈ)]\mathbb{E}_{\theta \sim P_0}[B(\theta)]. This approach inherits the structure of frequentist guarantees β€” the bound B(ΞΈ)B(\theta) depends on the gaps between arms, which are then integrated out. The paper argues that this technique "does not generalize to more informative priors and structured bandits" (Section 1), and specifically "cannot be extended to linear bandits through optimal design" (Section 1). The reason is that structured bandits introduce dependencies between arms (through shared parameters) that make the per-instance frequentist error decomposition far more complex.

3. The log₃/β‚‚β‚‚(K) factor. Even in the setting where BayesElim's analysis is valid (equal prior variances, uniform allocations), the elimination-based structure introduces a multiplicative penalty: the bound contains a factor of log⁑2(K)\log_2(K) from the number of elimination rounds, and additional logarithmic factors accumulate through the analysis. The paper quantifies this concretely in Appendix A by comparing the two bounds:

PnBayesElimβ‰€βˆ‘i,j:iβ‰ jlog⁑2(K) eβˆ’(ΞΌ0,iβˆ’ΞΌ0,j)24Οƒ021+nΟƒ02Klog⁑2(K)Οƒ2,PnPI-BAIβ‰€βˆ‘i,j:iβ‰ jeβˆ’(ΞΌ0,iβˆ’ΞΌ0,j)24Οƒ021+nΟƒ02KΟƒ2P_n^{\text{BayesElim}} \leq \sum_{i,j: i \neq j} \frac{\log_2(K) \, e^{-\frac{(\mu_{0,i}-\mu_{0,j})^2}{4\sigma_0^2}}}{\sqrt{1 + \frac{n\sigma_0^2}{K\log_2(K)\sigma^2}}}, \quad P_n^{\text{PI-BAI}} \leq \sum_{i,j: i \neq j} \frac{e^{-\frac{(\mu_{0,i}-\mu_{0,j})^2}{4\sigma_0^2}}}{\sqrt{1 + \frac{n\sigma_0^2}{K\sigma^2}}}

The PI-BAI bound is smaller by roughly log⁑23/2(K)\log_2^{3/2}(K) (the numerator loses one log⁑2(K)\log_2(K), the denominator gains log⁑2(K)\log_2(K) inside the square root). For K=10K = 10, this factor is approximately 6; for larger KK, it grows. This is significant because "this additional factor can be significant when bounding a probability" (Appendix A). And critically, this improvement comes not from a more sophisticated allocation but from the structure of the proof itself β€” PI-BAI's non-adaptive design eliminates the need for elimination rounds entirely.

The Gap in Structured Bandits

Beyond the multi-armed case, the paper highlights a complete absence of prior-dependent BAI results for structured bandits β€” specifically linear and hierarchical models:

  • Linear bandits: Arms share information through a low-dimensional parameter θ∈Rd\theta \in \mathbb{R}^d, with dβ‰ͺKd \ll K. The reward of arm ii is θ⊀xi\theta^\top x_i for known feature vectors xix_i. While frequentist fixed-budget BAI algorithms exist for linear models (e.g., BayesGap by Hoffman et al., 2014; GSE by Azizi et al., 2021), they rely on G-optimal design for successive elimination and provide frequentist guarantees. No prior-dependent expected PoE bounds exist for linear BAI in the Bayesian setting.

  • Hierarchical bandits: Arms are organized into a two-level structure with shared latent effects. Specifically, a latent parameter μ∈RL\mu \in \mathbb{R}^L (representing LL effects, such as hyperparameter values or drug component effects) is drawn from a Gaussian hyper-prior. Conditionally on ΞΌ\mu, each arm's mean ΞΈi\theta_i is drawn independently as ΞΈi∼N(bi⊀μ,Οƒ0,i2)\theta_i \sim \mathcal{N}(b_i^\top \mu, \sigma_{0,i}^2), where bib_i are known mixing weights encoding how arm ii relates to the LL effects. This creates correlations between arms: arms that share similar mixing weights have correlated means, so observing one provides information about others. These models are "rich models commonly used in practice where actions are typically structured into clusters, or effects" (Section 4.3), with applications in hyperparameter tuning (configurations share hyperparameter values) and drug design (drugs share molecular components). Yet no prior-dependent BAI bounds exist for this setting either β€” in the Bayesian or frequentist literature.

The inability of existing methods to handle these structures is not a mere gap in the literature; it represents a fundamental limitation in the applicability of Bayesian BAI. Real-world problems rarely present completely independent arms; exploiting structure is essential for statistical efficiency when KK is large relative to nn.

How This Paper Positions Itself

The paper positions PI-BAI as a direct answer to the question of adaptivity in Bayesian BAI. Rather than developing yet another elimination scheme, it takes the opposite approach: compute an optimal fixed allocation before seeing any data, using only the prior and the structure of the problem. This is not presented as an approximation or simplification β€” it is the core algorithmic idea, grounded in the recognition that prior information can substitute for the need to adapt.

The algorithm's simplicity is a feature, not a bug. PI-BAI consists of three steps: (1) compute allocation weights Ο‰=Alloc(n,P0)\omega = \text{Alloc}(n, P_0), (2) pull each arm ii exactly βŒŠΟ‰inβŒ‹\lfloor \omega_i n \rfloor times, (3) return the arm with the highest posterior mean. There is no elimination, no round structure, no discarding of data. This makes the algorithm both computationally trivial and amenable to theoretical analysis across diverse model classes β€” the proof technique (detailed in Section 5 and the analysis sections 4.1–4.3) transfers with minimal modification from multi-armed to linear to hierarchical bandits.

The theoretical contribution is not a single bound but a unified proof framework that:

  • Operates entirely in the Bayesian domain without intermediate frequentist steps,
  • Produces bounds valid for any fixed allocation weights Ο‰βˆˆΞ”K+\omega \in \Delta_K^+,
  • Naturally extends to structured models through the structure of the posterior distribution,
  • Yields tighter constants than elimination-based approaches by avoiding the overhead of multiple rounds.

Critically, the paper does not claim that non-adaptive strategies dominate adaptive ones β€” consistent with the impossibility results of Wang et al. (2024) and Degenne (2023), such a claim would be false. Rather, it demonstrates that prior-informed fixed allocations are competitive with adaptive methods in practice, and come with the strongest theoretical guarantees currently available for structured Bayesian BAI. The contribution is best understood as establishing a new baseline for what can be achieved without adaptivity, against which future adaptive Bayesian algorithms should be measured.

The paper also addresses a practical concern: what if the prior is misspecified? Section 4.4 provides both theoretical analysis (Lemma 4.5, bounding the effect of misspecification on uniform allocations as a multiplicative factor dn(i,j)=eO(1/n)d_n(i,j) = e^{\mathcal{O}(1/n)} that vanishes with budget) and practical remedies (offline prior learning from data, online prior learning through hierarchical models). This acknowledges that Bayesian methods are only as good as their prior, and provides a path forward when perfect prior knowledge is unavailable.

3. Technical Approach

3.1 Reader Orientation

This is primarily a theoretical analysis paper with algorithmic contributions. The core idea is disarmingly simple: in Bayesian fixed-budget best-arm identification, rather than adapting your sampling strategy on-the-fly based on observed rewards, you can pre-compute a fixed allocation of pulls across arms using only the prior distribution and the problem structure, pull each arm exactly that many times, and then pick the arm with the highest posterior mean. This non-adaptive strategy comes with strong theoretical guarantees β€” including the first prior-dependent bounds for linear and hierarchical bandits β€” and empirically competes with or outperforms sophisticated adaptive elimination methods.

3.2 Big-Picture Architecture (Diagram in Words)

The PI-BAI system has three components that execute sequentially:

  1. Allocation Computer (offline): Takes the prior distribution $P_0$, the total budget $n$, and the bandit model structure (multi-armed, linear, or hierarchical) as input. Produces a probability vector $\omega = (\omega_1, \ldots, \omega_K) \in \Delta_K$ specifying what fraction of the budget each arm receives. This is computed once, before any interaction with the environment, and is the intellectual core of the method. The paper proposes three distinct mechanisms for computing $\omega$: (a) numerical optimization of the theoretical PoE bound, (b) Bayesian G-optimal design maximizing the log-determinant of the regularized information matrix, and (c) a warm-up phase using Thompson sampling to empirically estimate good allocation ratios.

  2. Data Collector (online, deterministic): Takes the allocation weights $\omega$ and budget $n$. For each arm $i \in [K]$, deterministically collects $n_i = \lfloor \omega_i n \rfloor$ independent reward samples. This step involves no adaptation, no elimination, and no decision-making β€” it is a pure fixed-design experiment. For multi-armed bandits, each arm's rewards are $\mathcal{N}(\theta_i, \sigma^2)$ given the true mean $\theta_i$; for linear bandits, rewards are $\mathcal{N}(\theta^\top x_{A_t}, \sigma^2)$; for hierarchical bandits, rewards flow through the two-level prior structure.

  3. Decision Maker (online, deterministic): Takes the complete history $\mathcal{H}_n$ of $n$ observations. Computes the posterior mean reward for each arm given the prior and the collected data. Returns the arm with the highest posterior mean:

Jn=arg⁑max⁑i∈[K]E[r(i;θ)∣Hn]J_n = \arg\max_{i \in [K]} \mathbb{E}[r(i; \theta) \mid \mathcal{H}_n]

This decision rule is computationally tractable β€” for Gaussian models with Gaussian priors, posterior means are available in closed form β€” and serves as a computationally cheap proxy for the fully Bayesian optimal decision $\arg\min_j \mathbb{P}(j \neq i^*(\theta) \mid \mathcal{H}_n)$, which cannot be computed in closed form.

Information flows strictly forward: the prior informs the allocation, the allocation determines which data is collected, the data updates the posterior, and the posterior determines the final arm choice. There is no feedback loop β€” the allocation is computed independently of the observed rewards.

3.3 Roadmap for the Deep Dive

  • First, the formal problem setting and notation (Section 2 of the paper): the Bayesian fixed-budget BAI objective, the expected probability of error, the Gaussian observation model, and the conjugate prior structure across all three bandit settings. This establishes the mathematical universe in which everything else operates.

  • Second, the three allocation strategies (Section 3): optimization-based (minimizing the theoretical PoE bound), G-optimal design (maximizing information gain), and warm-up with Thompson sampling (empirical adaptive pre-sampling). Understanding these strategies is essential because they determine how prior knowledge is translated into a concrete experimental design.

  • Third, the posterior computation machinery (within Sections 4.1–4.3 and Appendix C.1): how Bayesian updating works in each model class β€” multi-armed, linear, and hierarchical. This is necessary to understand both the decision rule and the theoretical analysis.

  • Fourth, the theoretical bounds themselves (Theorems 4.1, 4.2, 4.4): the structure of the PoE upper bounds, what they depend on, and why they take the form they do. This includes walking through the key reformulation that makes the Bayesian proof possible.

  • Fifth, the general proof scheme (Section 5): the novel technical insight that enables prior-dependent bounds without relying on frequentist averaging, and why this matters for structured bandits.

  • Sixth, prior misspecification handling (Section 4.4): theoretical analysis of sensitivity to incorrect priors (Lemma 4.5) and practical strategies for learning priors from offline data or online through hierarchical modeling.

3.4 Detailed, Sentence-Based Technical Breakdown

The Bayesian Fixed-Budget BAI Objective

The paper studies the problem of minimizing the expected probability of error (expected PoE) in Bayesian fixed-budget best-arm identification. Let $\theta$ be the unknown parameter vector governing the bandit environment (arm means in MAB, the linear parameter in linear bandits, or the latent effects and arm means in hierarchical bandits). The parameter $\theta$ is drawn from a known prior distribution $P_0$. The learner interacts with the environment for exactly $n$ rounds: in each round $t \in [n]$, the learner selects an arm $A_t \in [K]$ and observes a stochastic reward $Y_t \sim P(\cdot; \theta, A_t)$. After $n$ rounds, the learner outputs a single arm $J_n \in [K]$ as its guess for the optimal arm $i^*(\theta) = \arg\max_{i \in [K]} r(i; \theta)$, where $r(i; \theta) = \mathbb{E}_{Y \sim P(\cdot; \theta, i)}[Y]$ is the mean reward of arm $i$ given $\theta$.

The expected PoE is defined as:

Pn=E[P(Jnβ‰ iβˆ—(ΞΈ)∣θ)]P_n = \mathbb{E}\left[\mathbb{P}\left(J_n \neq i^*(\theta) \mid \theta\right)\right]

where the outer expectation is over $\theta \sim P_0$ and the inner probability is over the randomness in the rewards given $\theta$.

What it computes: the probability, averaged over all possible bandit instances drawn from the prior, that the algorithm returns the wrong arm. This is the Bayesian analog of the frequentist PoE, where the performance would be assessed for a single fixed $\theta$. The double expectation captures both sources of randomness: which bandit instance nature selects, and what noisy rewards the learner observes from that instance.

Why this form: the Bayesian formulation allows the learner to leverage prior knowledge about which bandit instances are likely. In a frequentist setting, the learner must guard against worst-case instances (leading to conservative uniform or near-uniform exploration). In the Bayesian setting, the learner can concentrate budget on arms that the prior suggests are plausible candidates for optimality, and can allocate more samples to arms about which the prior is uncertain (higher prior variance). This fundamentally changes the nature of the allocation problem: the goal is not robustness to adversarial instances, but efficiency under the assumed prior distribution.

The Gaussian Observation Model and Conjugate Priors

All theoretical results in the paper assume Gaussian rewards with Gaussian priors. While this restricts the distributional family, it provides closed-form posterior updates that make both the algorithm and the analysis tractable. The paper notes that extensions to non-Gaussian distributions (e.g., logistic bandits) are discussed in Appendix B.3 and explored empirically in Appendix D.4, but the theoretical guarantees are specific to the Gaussian case.

Multi-armed bandit model (Equation 5):

ΞΈi∼N(ΞΌ0,i,Οƒ0,i2)βˆ€i∈[K]\theta_i \sim \mathcal{N}(\mu_{0,i}, \sigma_{0,i}^2) \quad \forall i \in [K] Yt∣θ,At∼N(ΞΈAt,Οƒ2)βˆ€t∈[n]Y_t \mid \theta, A_t \sim \mathcal{N}(\theta_{A_t}, \sigma^2) \quad \forall t \in [n]

where $\mu_{0,i}$ is the prior mean reward of arm $i$, $\sigma_{0,i}^2$ is the prior variance (encoding uncertainty about arm $i$'s mean), and $\sigma^2$ is the known observation noise variance (assumed identical across arms for notational simplicity, though the paper notes it could be arm-dependent with similar analysis).

What this models: each arm has an unknown true mean $\theta_i$ that is drawn independently from a Gaussian distribution. When the learner pulls arm $i$, it observes the true mean plus Gaussian noise with variance $\sigma^2$. The prior encodes the learner's initial beliefs: $\mu_{0,i}$ is the best guess for arm $i$'s mean before seeing data, and $\sigma_{0,i}^2$ quantifies confidence in that guess (small variance = confident prior, large variance = diffuse prior).

Why this form: the Gaussian-Gaussian conjugacy means that after observing $n_i$ samples from arm $i$, the posterior distribution of $\theta_i$ remains Gaussian with parameters that can be updated via simple weighted averaging (Equation 6). This is computationally essential β€” the posterior mean $\hat{\mu}_{n,i}$ and posterior variance $\hat{\sigma}_{n,i}^2$ are available in closed form without numerical integration or MCMC.

Posterior update for MAB (Equation 6):

Οƒ^n,iβˆ’2=Οƒ0,iβˆ’2+niΟƒβˆ’2\hat{\sigma}_{n,i}^{-2} = \sigma_{0,i}^{-2} + n_i \sigma^{-2} ΞΌ^n,i=Οƒ^n,i2(ΞΌ0,iΟƒ0,i2+Bn,iΟƒ2)\hat{\mu}_{n,i} = \hat{\sigma}_{n,i}^2 \left(\frac{\mu_{0,i}}{\sigma_{0,i}^2} + \frac{B_{n,i}}{\sigma^2}\right)

where $n_i = |\mathcal{T}_i|$ is the number of times arm $i$ was pulled, $\mathcal{T}_i = \{t \in [n] : A_t = i\}$ is the set of rounds when arm $i$ was chosen, and $B_{n,i} = \sum_{t \in \mathcal{T}_i} Y_t$ is the sum of observed rewards from arm $i$.

What it computes: the posterior precision (inverse variance) is the sum of the prior precision $\sigma_{0,i}^{-2}$ and the data precision $n_i \sigma^{-2}$ β€” information adds. The posterior mean is a precision-weighted average of the prior mean and the sample mean $B_{n,i}/n_i$. As $n_i$ grows, the posterior mean converges to the sample mean and the posterior variance shrinks toward zero. As $\sigma_{0,i}^2 \to 0$ (highly informative prior), the posterior mean stays close to $\mu_{0,i}$ regardless of data.

Why this form: the precision-weighted average is the optimal (minimum-variance unbiased) way to combine prior information with observed data under squared-error loss, and it emerges naturally from Bayes' rule with conjugate Gaussians. This weighted structure is what makes PI-BAI's decision rule (picking the arm with the highest posterior mean) a natural compromise between prior beliefs and empirical evidence.

Linear bandit model (Equation 7):

θ∼N(ΞΌ0,Ξ£0)\theta \sim \mathcal{N}(\mu_0, \Sigma_0) Yt∣θ,At∼N(θ⊀xAt,Οƒ2)βˆ€t∈[n]Y_t \mid \theta, A_t \sim \mathcal{N}(\theta^\top x_{A_t}, \sigma^2) \quad \forall t \in [n]

where $\mu_0 \in \mathbb{R}^d$ is the prior mean vector, $\Sigma_0 \in \mathbb{R}^{d \times d}$ is the prior covariance matrix, and $\mathcal{X} = \{x_1, \ldots, x_K\}$ is the set of known arm feature vectors with each $x_i \in \mathbb{R}^d$.

What this models: rather than each arm having an independent mean, all arms share a common low-dimensional parameter $\theta$. The mean reward of arm $i$ is the inner product $\theta^\top x_i$. This creates correlations: observing arm $i$ provides information about $\theta$, which in turn provides information about all other arms through their feature vectors. The prior on $\theta$ can encode beliefs about which directions in feature space are likely to have large rewards.

Why this form: linear models handle the practical case where $K$ may be large but arms lie in a low-dimensional space $d \ll K$. Without structure, identifying the best among $K$ independent arms requires $\Omega(K)$ samples. With a linear structure, $\mathcal{O}(d)$ samples may suffice if the feature vectors span $\mathbb{R}^d$. The MAB is recovered as a special case when $x_i = e_i \in \mathbb{R}^K$ (unit vectors) and $\Sigma_0 = \text{diag}(\sigma_{0,i}^2)_{i \in [K]}$.

Posterior update for linear bandits (Equation 8):

Ξ£^nβˆ’1=Ξ£0βˆ’1+Οƒβˆ’2βˆ‘i∈[K]nixixi⊀\hat{\Sigma}_n^{-1} = \Sigma_0^{-1} + \sigma^{-2} \sum_{i \in [K]} n_i x_i x_i^\top ΞΌ^n=Ξ£^n(Ξ£0βˆ’1ΞΌ0+Οƒβˆ’2Bn)\hat{\mu}_n = \hat{\Sigma}_n \left(\Sigma_0^{-1} \mu_0 + \sigma^{-2} B_n\right)

where $B_n = \sum_{t \in [n]} Y_t x_{A_t}$ is the feature-weighted sum of rewards.

What it computes: the posterior precision matrix $\hat{\Sigma}_n^{-1}$ is the sum of the prior precision and the empirical Fisher information from the data (each pull of arm $i$ contributes $n_i x_i x_i^\top / \sigma^2$ to the precision). The posterior mean $\hat{\mu}_n$ is again a precision-weighted combination of the prior mean and a data-dependent term. The mean posterior reward of arm $i$ is then $\mathbb{E}[r(i; \theta) \mid \mathcal{H}_n] = \hat{\mu}_n^\top x_i$.

Why this form: the matrix structure captures correlations. Arms whose feature vectors are similar (small $\|x_i - x_j\|$) have highly correlated posterior mean estimates β€” observing one tells you a lot about the other. This is what enables statistical efficiency: you don't need to pull every arm extensively if you've already learned $\theta$ well enough from a diverse subset.

Hierarchical bandit model (Equation 9):

μ∼N(Ξ½,Ξ£)\mu \sim \mathcal{N}(\nu, \Sigma) ΞΈi∼N(bi⊀μ,Οƒ0,i2)βˆ€i∈[K]\theta_i \sim \mathcal{N}(b_i^\top \mu, \sigma_{0,i}^2) \quad \forall i \in [K] Yt∣μ,ΞΈ,At∼N(ΞΈAt,Οƒ2)βˆ€t∈[n]Y_t \mid \mu, \theta, A_t \sim \mathcal{N}(\theta_{A_t}, \sigma^2) \quad \forall t \in [n]

where $\mu = (\mu_\ell)_{\ell \in [L]} \in \mathbb{R}^L$ is a latent parameter of $L$ shared effects, $\nu \in \mathbb{R}^L$ and $\Sigma \in \mathbb{R}^{L \times L}$ are the hyper-prior mean and covariance, and $b_i \in \mathbb{R}^L$ are known mixing weights for arm $i$.

What this models: a two-level Bayesian hierarchy. At the top level, a latent effect vector $\mu$ is drawn from a Gaussian hyper-prior. These effects could represent, for example, the value of different hyperparameter settings in hyperparameter tuning, or the efficacy of different molecular components in drug design. At the bottom level, each arm's mean $\theta_i$ is a linear combination $b_i^\top \mu$ of the effects plus independent Gaussian noise with variance $\sigma_{0,i}^2$. The mixing weights $b_i$ encode how each arm relates to the effects: $b_{i,\ell}$ quantifies the association between arm $i$ and effect $\ell$. This creates correlations between arms β€” two arms with similar mixing weights will have correlated means because they share the same underlying effects.

Why this form: hierarchical models capture a richer structure than linear models while remaining computationally tractable. In a linear model, arm means are deterministic functions of $\theta$: $r(i; \theta) = \theta^\top x_i$. In a hierarchical model, arm means are random even given $\mu$, with the $\sigma_{0,i}^2$ term capturing arm-specific deviations from the effect-based prediction. This is more realistic: two configurations sharing the same hyperparameter values might still perform differently due to interactions not captured by the additive effect model. The hierarchical prior also enables online learning of the prior itself, since $\mu$ can be inferred from data across arms.

Marginal posterior for hierarchical bandits (Equation 17):

Οƒ^n,i2=Οƒ~n,i2+Οƒ~n,i4Οƒ0,i4bi⊀Σ˘nbi\hat{\sigma}_{n,i}^2 = \tilde{\sigma}_{n,i}^2 + \frac{\tilde{\sigma}_{n,i}^4}{\sigma_{0,i}^4} b_i^\top \breve{\Sigma}_n b_i ΞΌ^n,i=Οƒ~n,i2(μ˘n⊀biΟƒ0,i2+Bn,iΟƒ2)\hat{\mu}_{n,i} = \tilde{\sigma}_{n,i}^2 \left(\frac{\breve{\mu}_n^\top b_i}{\sigma_{0,i}^2} + \frac{B_{n,i}}{\sigma^2}\right)

where $\tilde{\sigma}_{n,i}^2$ and $\tilde{\mu}_{n,i}$ are the conditional posterior parameters given $\mu$ (Equation 16), and $\breve{\mu}_n$ and $\breve{\Sigma}_n$ are the joint effect posterior parameters (Equation 15).

What it computes: the marginal posterior of $\theta_i$ after integrating out the unknown effects $\mu$. The posterior variance $\hat{\sigma}_{n,i}^2$ has two components: the conditional variance $\tilde{\sigma}_{n,i}^2$ (uncertainty about $\theta_i$ given $\mu$) plus an additional term reflecting uncertainty about $\mu$ itself, projected through the mixing weights $b_i$. The larger $\breve{\Sigma}_n$ is (more uncertainty about effects), the larger this additional term. This means that even if you pull arm $i$ many times (driving $\tilde{\sigma}_{n,i}^2$ toward zero), your posterior uncertainty about $\theta_i$ is bounded below by your uncertainty about the underlying effects β€” you can never be more certain about an arm than you are about the effects that generate it.

Why this form: the closed-form marginal posterior is a consequence of the Gaussian-Gaussian conjugacy at both levels of the hierarchy (Lemma C.4 in Appendix C.1). Computing it requires first computing the joint effect posterior $Q_n(\mu)$ (which pools information across all arms), then computing the conditional arm posterior $P_{n,i}(\theta_i \mid \mu)$, and finally integrating out $\mu$. The computational complexity is $\mathcal{O}(K + L^3)$ for the effect posterior update and $\mathcal{O}(K + L^2)$ for per-arm marginal posteriors β€” much more efficient than the $\mathcal{O}((K+L)^3)$ that would be required by treating the hierarchical model as a $(K+L)$-dimensional linear model.

Decision Rule: Highest Posterior Mean

After collecting $n_i$ samples from each arm $i$, PI-BAI selects the arm with the highest posterior mean:

Jn=arg⁑max⁑i∈[K]E[r(i;θ)∣Hn]J_n = \arg\max_{i \in [K]} \mathbb{E}[r(i; \theta) \mid \mathcal{H}_n]

In the MAB setting, this is simply $\arg\max_i \hat{\mu}_{n,i}$. In linear bandits, it is $\arg\max_i \hat{\mu}_n^\top x_i$. In hierarchical bandits, it is $\arg\max_i \hat{\mu}_{n,i}$ where $\hat{\mu}_{n,i}$ comes from the marginal posterior in Equation (17).

What it computes: for each arm, the posterior mean is the expected reward conditional on all observed data, integrating over posterior uncertainty about the parameters. The decision rule picks the arm with the highest such expectation. This is not the fully Bayesian optimal decision β€” which would minimize the posterior probability of error by picking $\arg\min_j \mathbb{P}(j \neq i^*(\theta) \mid \mathcal{H}_n)$ β€” but is a computationally tractable proxy. The paper notes: "a fully Bayesian algorithm would set $J_n = \arg\min_{j \in [K]} \mathbb{P}(j \neq i^*(\theta) \mid \mathcal{H}_n)$; however, this cannot be computed explicitly even in the simple Gaussian setting, and we instead opt for a decision rule that is deterministic conditionally on observations" (Section 2.1).

Why this form: the posterior mean is the Bayes estimator under squared-error loss for the arm mean, but more importantly, it is the natural summary statistic for ranking arms when the posterior is symmetric (as Gaussians are). The key theoretical insight is that the probability of error can be bounded by analyzing how often the arm that maximizes the posterior mean differs from the arm that maximizes a posterior sample of the true means β€” a connection made rigorous through the key reformulation in Section 5.

Allocation Strategy 1: Optimization of the Theoretical Bound (PI-BAI(Opt))

The first and most principled allocation strategy directly minimizes the theoretical upper bound on the expected PoE. For any allocation $\omega \in \Delta_K^+$, the paper derives an upper bound of the form:

Pn≀C(P0,Ο‰,n)P_n \leq C(P_0, \omega, n)

where $C$ depends on the prior $P_0$, the budget $n$, and the allocation weights $\omega$. The optimized allocation is then defined as:

Opt(n,P0)=arg⁑minβ‘Ο‰βˆˆΞ”K+C(P0,Ο‰,n)\text{Opt}(n, P_0) = \arg\min_{\omega \in \Delta_K^+} C(P_0, \omega, n)

What it computes: the allocation vector $\omega$ that minimizes the theoretical worst-case (upper bound) on the expected probability of error, given the prior and the budget. This is a numerical optimization problem β€” the paper notes it "is non-convex and may require advanced numerical methods (e.g., L-BFGS-B of Virtanen et al., 2020) to find a suitable optimum" (Section 3). Crucially, this optimization is performed once offline, before any interaction with the environment.

Why this form: directly optimizing the theoretical guarantee provides the strongest possible connection between the algorithm's actual performance and the provable bound. If the bound is reasonably tight, then minimizing the bound should yield near-optimal allocations in practice. The non-convexity is a practical challenge but not a fundamental obstacle, since the optimization is offline and can afford multiple random restarts.

Explicit solution for K=2 MAB. To provide intuition, the paper presents the closed-form solution for the two-armed MAB case. For $K = 2$, with the optimized allocation $\omega^{\text{opt}} = \text{Opt}(n, P_0)$:

Ο‰1opt=Ξ [0,1](12βˆ’(Οƒ0,22βˆ’Οƒ0,12)Οƒ22nΟƒ0,12Οƒ0,22),Ο‰2opt=1βˆ’Ο‰1opt\omega_1^{\text{opt}} = \Pi_{[0,1]}\left(\frac{1}{2} - \frac{(\sigma_{0,2}^2 - \sigma_{0,1}^2)\sigma^2}{2n\sigma_{0,1}^2\sigma_{0,2}^2}\right), \quad \omega_2^{\text{opt}} = 1 - \omega_1^{\text{opt}}

where $\Pi_{[0,1]}(\cdot)$ denotes truncation to the interval $[0,1]$.

What it computes: when prior variances are equal ($\sigma_{0,1}^2 = \sigma_{0,2}^2$), the optimal allocation is uniform β€” each arm gets half the budget. When the variances differ, the arm with larger prior variance (more prior uncertainty) receives a larger allocation. The magnitude of the deviation from uniformity is inversely proportional to the budget $n$: as $n \to \infty$, the allocation converges to uniform because the influence of the prior vanishes asymptotically. For small budgets, the correction term can be large enough that the truncation binds: if $n < 2\sigma^2\left|\frac{\sigma_{0,2}^2}{\sigma_{0,1}^2} - \frac{\sigma_{0,1}^2}{\sigma_{0,2}^2}\right|$, the arm with the larger prior variance receives the entire budget ($\omega = 1$).

Why this form: in the two-armed case, identifying the best arm is equivalent to identifying the worst arm β€” there is no third alternative to consider. The optimal allocation depends only on prior variances, not on the prior mean gap $\Delta_0 = |\mu_{0,1} - \mu_{0,2}|$, because the prior means don't affect how quickly you can distinguish the two arms (they affect the prior probability that each arm is best, but in a symmetric way for $K=2$). This equivalence breaks for $K > 2$, where the prior means matter because the allocation must balance resources across multiple pairwise comparisons.

Practical heuristic: mixture with heuristic weights. In the experiments, the authors use a mixture of the optimized weights and a heuristic that allocates proportionally to $\mu_{0,i} \sigma_{0,i}$:

Ο‰imixture=Ξ±Ο‰iopt+(1βˆ’Ξ±)ΞΌ0,iΟƒ0,iβˆ‘k∈[K]ΞΌ0,kΟƒ0,k\omega_i^{\text{mixture}} = \alpha \omega_i^{\text{opt}} + (1-\alpha) \frac{\mu_{0,i} \sigma_{0,i}}{\sum_{k \in [K]} \mu_{0,k} \sigma_{0,k}}

with $\alpha = 0.5$ chosen based on ablation studies (Appendix D.7, Figure 12). This heuristic gives more weight to arms with higher prior means (more likely to be optimal) and higher prior variances (more uncertainty to resolve). The mixture is still covered by the theoretical upper bound (since the bound holds for any fixed $\omega$), but the empirical performance can be better than pure optimization due to the non-convexity of the optimization landscape.

Allocation Strategy 2: Bayesian G-Optimal Design (PI-BAI(G-opt))

The second allocation strategy leverages optimal experimental design theory. In classical (frequentist) linear bandits, G-optimal design minimizes the maximum prediction variance over the design space. The paper adapts this to the Bayesian setting by maximizing the log-determinant of the regularized information matrix:

G-opt(n,P0)=arg⁑maxβ‘Ο‰βˆˆΞ”K+log⁑det⁑(nΟƒ2Σω+Ξ£0βˆ’1)\text{G-opt}(n, P_0) = \arg\max_{\omega \in \Delta_K^+} \log \det\left(\frac{n}{\sigma^2} \Sigma_\omega + \Sigma_0^{-1}\right)

where $\Sigma_\omega = \sum_{k \in [K]} \omega_k x_k x_k^\top$ is the weighted empirical covariance of the arm features under allocation $\omega$, and $\Sigma_0^{-1}$ is the prior precision matrix.

What it computes: the allocation that maximizes the log-determinant of the sum of the (scaled) Fisher information from the data ($\frac{n}{\sigma^2} \Sigma_\omega$) and the prior information ($\Sigma_0^{-1}$). The determinant of this matrix measures the "volume" of information β€” maximizing it corresponds to making the joint posterior as concentrated as possible in all directions of the parameter space. This is a convex optimization problem (log-determinant is concave over the positive-definite cone), so it can be solved efficiently.

Why this form: G-optimal design minimizes the worst-case posterior variance of the predicted mean reward across all arms. Specifically, the D-optimality criterion (maximizing log-det) is equivalent to minimizing the volume of the posterior confidence ellipsoid. By including the prior precision $\Sigma_0^{-1}$ as regularization, the design naturally balances exploration of uncertain directions in feature space against exploitation of prior knowledge. Arms whose feature vectors are "different" (spanning diverse directions) receive more weight; arms that are "similar" to already well-represented arms receive less.

Connection to the theoretical bound (Corollary 4.3). For linear bandits with diagonal prior covariance, using G-optimal design yields the simplified bound:

Pnβ‰€βˆ‘i,j:iβ‰ j11+n2dΟƒ2ci,jexp⁑(βˆ’(ΞΌ0⊀xiβˆ’ΞΌ0⊀xj)22βˆ₯xiβˆ’xjβˆ₯Ξ£02)P_n \leq \sum_{i,j: i \neq j} \frac{1}{\sqrt{1 + \frac{n}{2d\sigma^2} c_{i,j}}} \exp\left(-\frac{(\mu_0^\top x_i - \mu_0^\top x_j)^2}{2\|x_i - x_j\|_{\Sigma_0}^2}\right)

What it computes: under G-optimal design, the posterior variance term in the exponent of the bound can be related to the design optimality criterion through Lemma C.9. The key lemma shows that for the G-optimal allocation, $\max_{x \in \mathcal{X}} \|x\|_{\hat{\Sigma}_n} \leq d\sigma^2/n$ β€” the maximum posterior standard deviation of any arm's mean reward is bounded by the dimension $d$ times the observation noise over the budget. This gives a clean $\mathcal{O}(1/\sqrt{n})$ rate with constants depending explicitly on dimension and noise, removing dependence on the specific arm configuration.

Why this form: the $2d$ in the denominator reflects the price of learning in $d$ dimensions: a G-optimal design distributes samples to ensure that the posterior uncertainty is roughly equal in all $d$ directions, and the worst-case uncertainty scales with $d$ rather than with $K$. This is the key advantage of structured models β€” statistical complexity depends on the intrinsic dimension $d$, not the number of arms $K$.

Application to MAB. For MAB, G-optimal design uses $x_i = e_i \in \mathbb{R}^K$ (unit vectors) and $\Sigma_0 = \text{diag}(\sigma_{0,i}^2)_{i \in [K]}$, recovering arm-independent allocations that prioritize arms with high prior variance β€” similar in spirit, though not identical, to the optimization-based allocations.

Application to hierarchical bandits. The hierarchical model can be reformulated as a linear model with $\bar{\nu} \in \mathbb{R}^{L+K}$ and block-diagonal prior covariance $\bar{\Sigma}$ (Section 4.3), enabling direct G-optimal design computation. This is computationally more expensive than using the natural hierarchical posterior (requiring $\mathcal{O}((K+L)^3)$ matrix operations), but inherits the theoretical guarantees from the linear case.

Allocation Strategy 3: Warm-Up with Thompson Sampling (PI-BAI(TS))

The third strategy departs from the theoretical framework by incorporating an adaptive warm-up phase. A prior-dependent warm-up policy $\pi_w$ (specifically, Thompson sampling) interacts with the environment for $n_w$ rounds. The allocation weights are then set proportional to the empirical pull counts during warm-up:

Ο‰iTS=1nwβˆ‘t∈[nw]1{At=i},βˆ€i∈[K]\omega_i^{\text{TS}} = \frac{1}{n_w} \sum_{t \in [n_w]} \mathbb{1}\{A_t = i\}, \quad \forall i \in [K]

What it computes: the fraction of warm-up rounds in which Thompson sampling chose to pull each arm. Thompson sampling naturally balances exploration and exploitation given the prior: arms with high posterior probability of being optimal are pulled frequently, but arms with high posterior variance also receive pulls to reduce uncertainty. The resulting empirical allocation reflects this balance for the specific bandit instance encountered.

Why this form: this is an improper allocation strategy in the sense that it depends on the observed rewards (through Thompson sampling's adaptive choices), not just on the prior. The paper explicitly notes that "this no longer corresponds exactly to the PI-BAI algorithm as this Alloc strategy also interacts with the environment and outputs instance-dependent weights $\omega^{\pi_w}$, thus our theoretical guarantees do not apply to this framework" (Section 3). However, it performs well empirically, particularly in hierarchical bandits where the structure is rich and the relationship between prior and optimal allocation is complex. The warm-up length is set to $n_w = K$ in the main experiments (each arm gets roughly one pull on average during warm-up), with ablation studies in Appendix D.7 showing this choice is near-optimal.

Why Thompson sampling specifically: the paper cites Lee et al. (2024) for Thompson sampling's "strong performance in BAI." The intuition is that Thompson sampling's posterior sampling mechanism naturally allocates more pulls to arms that are either promising (high posterior mean) or uncertain (high posterior variance) β€” exactly the considerations that should drive a good BAI allocation. After the warm-up identifies which arms are worth focusing on, the remaining budget $n - n_w$ is spent deterministically according to the learned allocation $\omega^{\text{TS}}$. For hierarchical bandits, the paper uses the mixed-effect Thompson sampling algorithm of Aouali et al. (2023), which extends Thompson sampling to the hierarchical setting by sampling from the joint posterior over effects and arm means.

The General Proof Scheme: Key Reformulation of Expected PoE

The theoretical analysis rests on a fundamental reformulation of the expected probability of error that enables a fully Bayesian proof without any frequentist intermediate steps. The starting point is the definition:

Pn=E[P(Jnβ‰ iβˆ—(ΞΈ)∣θ)]P_n = \mathbb{E}\left[\mathbb{P}\left(J_n \neq i^*(\theta) \mid \theta\right)\right]

The outer expectation is over $\theta \sim P_0$, the inner probability is over rewards given $\theta$. The key insight is that by the tower property of conditional expectation, this can be rewritten with the conditioning reversed:

Pn=E[P(Jnβ‰ iβˆ—(ΞΈ)∣Hn)]P_n = \mathbb{E}\left[\mathbb{P}\left(J_n \neq i^*(\theta) \mid \mathcal{H}_n\right)\right]

What it computes: the original formulation conditions on the true parameter $\theta$ and averages over $\theta$. The reformulated version conditions on the observed history $\mathcal{H}_n$ (the data the algorithm actually sees) and asks: given these observations, what is the posterior probability that the arm $J_n$ maximizing the posterior mean is not the arm that maximizes a posterior draw of $\theta$? The outer expectation is now over the marginal distribution of $\mathcal{H}_n$ (integrating out both $\theta$ and the rewards).

Why this form: this swap of conditioning is the decisive technical move. In the original formulation, bounding $\mathbb{P}(J_n \neq i^*(\theta) \mid \theta)$ for fixed $\theta$ requires controlling the probability that sample means deviate from true means β€” a frequentist concentration argument. In the reformulated version, bounding $\mathbb{P}(J_n \neq i^*(\theta) \mid \mathcal{H}_n)$ requires controlling the probability that a posterior sample of $\theta$ disagrees with the posterior mean β€” a purely Bayesian computation that depends on the posterior variance, not on the unknown true $\theta$. This is the conceptual breakthrough that enables the analysis to extend to structured bandits: the posterior distribution is always Gaussian and its parameters are known in closed form, regardless of how complex the prior structure is.

Further decomposition into pairwise comparisons. The probability of error can be decomposed as a sum over all possible misidentification events:

Pn=βˆ‘i,j∈[K],iβ‰ jE[P(iβˆ—(ΞΈ)=i∣Jn=j,Hn)1{Jn=j}]P_n = \sum_{i,j \in [K], i \neq j} \mathbb{E}\left[\mathbb{P}\left(i^*(\theta) = i \mid J_n = j, \mathcal{H}_n\right) \mathbb{1}\{J_n = j\}\right]

What it computes: the expected PoE is the sum, over all ordered pairs of distinct arms $(i, j)$, of the probability that arm $i$ is truly optimal but the algorithm chose arm $j$. Each term conditions on the event $\{J_n = j\}$ (the algorithm chose $j$) and asks for the posterior probability that $i$ is actually optimal. The product with $\mathbb{1}\{J_n = j\}$ ensures we only count cases where arm $j$ was actually selected.

Why this form: this decomposition reduces the problem of bounding the error of a $K$-arm selection to bounding $K(K-1)$ pairwise error probabilities. The key quantity to bound is $\mathbb{P}(i^*(\theta) = i \mid J_n = j, \mathcal{H}_n)$ β€” the posterior probability that arm $i$ is optimal given that arm $j$ maximized the posterior mean. Since $i^*(\theta) = i$ implies $\theta_i \geq \theta_j$ (arm $i$ beats arm $j$), and $J_n = j$ implies $\hat{\mu}_{n,j} \geq \hat{\mu}_{n,i}$ (arm $j$ has higher posterior mean), this probability is at most the probability that a posterior sample of the difference $\theta_i - \theta_j$ crosses zero against the direction suggested by the posterior mean β€” a Gaussian tail probability.

Bounding the pairwise probabilities. For the MAB setting, applying Hoeffding's inequality for sub-Gaussian random variables yields:

P(ΞΈiβ‰₯ΞΈj∣Hn,Jn=j)≀exp⁑(βˆ’(ΞΌ^n,iβˆ’ΞΌ^n,j)22(Οƒ^n,i2+Οƒ^n,j2))\mathbb{P}(\theta_i \geq \theta_j \mid \mathcal{H}_n, J_n = j) \leq \exp\left(-\frac{(\hat{\mu}_{n,i} - \hat{\mu}_{n,j})^2}{2(\hat{\sigma}_{n,i}^2 + \hat{\sigma}_{n,j}^2)}\right)

What it computes: the posterior probability that the true mean difference $\theta_i - \theta_j$ has the opposite sign from the posterior mean difference $\hat{\mu}_{n,i} - \hat{\mu}_{n,j}$ is exponentially small in the squared standardized difference. The denominator $\hat{\sigma}_{n,i}^2 + \hat{\sigma}_{n,j}^2$ is the posterior variance of the difference, which depends on the allocation through the $n_i$ and prior variances.

Why this form: this is a standard Gaussian concentration bound applied to the posterior distribution of $\theta_i - \theta_j$, which is $\mathcal{N}(\hat{\mu}_{n,i} - \hat{\mu}_{n,j}, \hat{\sigma}_{n,i}^2 + \hat{\sigma}_{n,j}^2)$ under the Gaussian posterior. The probability that this Gaussian random variable is positive when the mean is negative decays exponentially in the squared signal-to-noise ratio.

Taking the expectation over the history. The remaining step is to compute the expectation of the exponential bound over the marginal distribution of $\mathcal{H}_n$. This is where the explicit prior-dependence enters. For the MAB:

EHn[exp⁑(βˆ’(ΞΌ^n,iβˆ’ΞΌ^n,j)22(Οƒ^n,i2+Οƒ^n,j2))]=11+nΟ•i,j(Ο‰)exp⁑(βˆ’(ΞΌ0,iβˆ’ΞΌ0,j)22(Οƒ0,i2+Οƒ0,j2))\mathbb{E}_{\mathcal{H}_n}\left[\exp\left(-\frac{(\hat{\mu}_{n,i} - \hat{\mu}_{n,j})^2}{2(\hat{\sigma}_{n,i}^2 + \hat{\sigma}_{n,j}^2)}\right)\right] = \frac{1}{\sqrt{1 + n\phi_{i,j}(\omega)}} \exp\left(-\frac{(\mu_{0,i} - \mu_{0,j})^2}{2(\sigma_{0,i}^2 + \sigma_{0,j}^2)}\right)

What it computes: integrating the Gaussian tail probability over the sampling distribution of the posterior means yields a closed form because the posterior means $\hat{\mu}_{n,i}$ themselves are Gaussian under the marginal distribution of $\mathcal{H}_n$ (they are linear combinations of the Gaussian rewards). The result splits into two factors: an exponential term depending on the prior mean gap (the prior probability that the arms are well-separated), and a $1/\sqrt{1 + n\phi_{i,j}(\omega)}$ term reflecting how quickly additional samples reduce the posterior probability of confusion.

Why this form: the $\mathcal{O}(1/\sqrt{n})$ rate emerges from Lemma C.5 in Appendix C.1: the expectation of $\exp(-X^2/(2a^2))$ when $X \sim \mathcal{N}(\mu, \sigma^2)$ is $\frac{1}{\sqrt{1 + \sigma^2/a^2}} \exp(-\mu^2/(2(a^2 + \sigma^2)))$. Here, $X = \hat{\mu}_{n,i} - \hat{\mu}_{n,j}$ has mean $\mu_{0,i} - \mu_{0,j}$ and variance scaling as $\mathcal{O}(1/n)$, and $a^2 = \hat{\sigma}_{n,i}^2 + \hat{\sigma}_{n,j}^2$ also scales as $\mathcal{O}(1/n)$. The ratio $\sigma^2/a^2$ is therefore $\mathcal{O}(n)$, yielding the $\sqrt{n}$ in the denominator. This rate is consistent with the known lower bound for Bayesian BAI ($\Omega(1/\sqrt{n})$ from Atsidakou et al., 2022), and the constants are explicit functions of the prior parameters and allocation.

Extension to structured bandits. The same proof structure transfers to linear and hierarchical bandits with minimal modification. The key differences are:

  • For linear bandits, the posterior mean difference $\hat{\mu}_n^\top (x_i - x_j)$ is Gaussian with mean $\mu_0^\top (x_i - x_j)$ and a more complex variance structure involving the posterior covariance matrix $\hat{\Sigma}_n$ and the covariance of $\hat{\mu}_n$ under the marginal distribution of $\mathcal{H}_n$. The bound takes the form (Theorem 4.2):

Pnβ‰€βˆ‘i,j:iβ‰ j11+nci,j(Ο‰)βˆ₯xiβˆ’xjβˆ₯Ξ£^n2exp⁑(βˆ’(ΞΌ0⊀xiβˆ’ΞΌ0⊀xj)22βˆ₯xiβˆ’xjβˆ₯Ξ£02)P_n \leq \sum_{i,j: i \neq j} \frac{1}{\sqrt{1 + \frac{n c_{i,j}(\omega)}{\|x_i - x_j\|_{\hat{\Sigma}_n}^2}}} \exp\left(-\frac{(\mu_0^\top x_i - \mu_0^\top x_j)^2}{2\|x_i - x_j\|_{\Sigma_0}^2}\right)

where $c_{i,j}(\omega)$ is given in Equation (26) and captures the dependence on allocation through the posterior covariance structure.

  • For hierarchical bandits, the posterior mean difference $\hat{\mu}_{n,i} - \hat{\mu}_{n,j}$ has mean $\nu^\top(b_i - b_j)$ and the prior variance of the arm mean difference decomposes as $\|b_i - b_j\|_\Sigma^2 + \sigma_{0,i}^2 + \sigma_{0,j}^2$ β€” the sum of effect-level uncertainty (projected through the mixing weights) and arm-level uncertainty. This richer structure means that the prior gap can be large even if individual arm-level prior variances are small, as long as the effects are well-separated in directions relevant to the arms being compared.

The $\phi_{i,j}(\omega)$ Term: How Allocation Affects the Bound

The function $\phi_{i,j}(\omega)$ in Theorem 4.1 encodes how the allocation weights affect the difficulty of distinguishing arms $i$ and $j$. Its explicit form (Equation 18) is:

ϕi,j(ω)=σ0,i4ωi(σ2n+ωjσ0,j2)+σ0,j4ωj(σ2n+ωiσ0,i2)σ2σ0,i2(σ2n+ωjσ0,j2)+σ2σ0,j2(σ2n+ωiσ0,i2)\phi_{i,j}(\omega) = \frac{\sigma_{0,i}^4 \omega_i \left(\frac{\sigma^2}{n} + \omega_j \sigma_{0,j}^2\right) + \sigma_{0,j}^4 \omega_j \left(\frac{\sigma^2}{n} + \omega_i \sigma_{0,i}^2\right)}{\sigma^2 \sigma_{0,i}^2 \left(\frac{\sigma^2}{n} + \omega_j \sigma_{0,j}^2\right) + \sigma^2 \sigma_{0,j}^2 \left(\frac{\sigma^2}{n} + \omega_i \sigma_{0,i}^2\right)}

What it computes: $n \phi_{i,j}(\omega)$ is (up to constants) the ratio of the variance of the posterior mean difference $\hat{\mu}_{n,i} - \hat{\mu}_{n,j}$ to the sum of posterior variances $\hat{\sigma}_{n,i}^2 + \hat{\sigma}_{n,j}^2$. Larger $\phi_{i,j}$ means the posterior means are more variable relative to the posterior uncertainty, which reduces the probability of misidentification (since the exponential bound decays faster when this ratio is larger). When $n$ is large, $\phi_{i,j}(\omega)$ simplifies to $\frac{2\sigma_0^2}{\sigma^2}\frac{\omega_i \omega_j}{\omega_i + \omega_j}$ for equal prior variances β€” it increases with both $\omega_i$ and $\omega_j$, confirming that allocating more samples to either arm helps distinguish them.

Why this form: the numerator captures how the posterior means for arms $i$ and $j$ vary under the sampling distribution (their variance comes from the rewards observed for those arms). The denominator captures how the posterior variances shrink with sample size. The allocation weights enter in a symmetric way: giving more budget to arm $i$ reduces the posterior variance of $\hat{\mu}_{n,i}$ (good) but also reduces the variance of $\hat{\mu}_{n,i}$ under the sampling distribution (since it's a more precise average) β€” these effects partially cancel, leaving $\phi_{i,j}$ as a balanced measure of how the allocation translates into discriminative power.

Prior Misspecification: Theory and Remedies

The paper acknowledges that perfect prior knowledge is unrealistic. Section 4.4 provides analysis and practical solutions.

Lemma 4.5 (Misspecification bound for uniform allocations). For PI-BAI(Uni) with misspecified prior $\mathcal{N}(\tilde{\mu}_{0,i}, \tilde{\sigma}_0^2)$ when the true prior is $\mathcal{N}(\mu_{0,i}, \sigma_0^2)$:

Pnβ‰€βˆ‘i,j:iβ‰ j11+nΟƒ02KΟƒ2exp⁑(βˆ’(ΞΌ0,iβˆ’ΞΌ0,j)24Οƒ02)dn(i,j)P_n \leq \sum_{i,j: i \neq j} \frac{1}{\sqrt{1 + \frac{n\sigma_0^2}{K\sigma^2}}} \exp\left(-\frac{(\mu_{0,i} - \mu_{0,j})^2}{4\sigma_0^2}\right) d_n(i,j)

where $d_n(i,j) = \exp\left(\frac{\text{KL}(P_{ij}, \tilde{P}_{ij})}{1 + \frac{n\sigma_0^2}{K\sigma^2}}\right)$ is a misspecification factor, with $P_{ij} = \mathcal{N}(\mu_{0,i} - \mu_{0,j}, 2\sigma_0^2)$ the true prior gap distribution and $\tilde{P}_{ij}$ the misspecified one.

What it computes: when the prior is misspecified, the bound inflates by a multiplicative factor $d_n(i,j) \geq 1$ that depends on the Kullback-Leibler divergence between the true and assumed prior gap distributions. The KL divergence here is:

KL(Pij,P~ij)=Οƒ02Οƒ~04(Οƒ~02Οƒ02(ΞΌ0,iβˆ’ΞΌ0,j)βˆ’(ΞΌ~0,iβˆ’ΞΌ~0,j))2\text{KL}(P_{ij}, \tilde{P}_{ij}) = \frac{\sigma_0^2}{\tilde{\sigma}_0^4}\left(\frac{\tilde{\sigma}_0^2}{\sigma_0^2}(\mu_{0,i} - \mu_{0,j}) - (\tilde{\mu}_{0,i} - \tilde{\mu}_{0,j})\right)^2

Why this form: if the prior variances are well-specified ($\tilde{\sigma}_0^2 = \sigma_0^2$), the KL divergence simplifies to $((\mu_{0,i} - \mu_{0,j}) - (\tilde{\mu}_{0,i} - \tilde{\mu}_{0,j}))^2 / \sigma_0^2$ β€” it measures the squared error in the prior gap relative to the prior standard deviation. If the prior gaps are preserved ($\tilde{\mu}_{0,i} - \tilde{\mu}_{0,j} = \mu_{0,i} - \mu_{0,j}$), then $d_n(i,j) = 1$ and Theorem 4.1 is recovered exactly β€” the absolute values of the prior means do not matter, only the relative differences. Critically, $d_n(i,j) = e^{\mathcal{O}(1/n)}$ because the KL divergence is divided by $1 + n\sigma_0^2/(K\sigma^2) = \mathcal{O}(n)$. As the budget grows, the misspecification factor converges to 1, meaning the impact of prior misspecification vanishes asymptotically β€” the data eventually overwhelms the incorrect prior.

Offline prior learning. The paper demonstrates on MovieLens data that a prior can be learned from offline data without requiring the bandit instances to actually be drawn from a Gaussian distribution. For MovieLens, the authors form a rating matrix $M$ of size $6,040 \times 3,883$ (users Γ— movies), complete it using alternating least squares with rank $d = 5$, obtaining $M = U^\top V$ where rows of $V$ represent movies. The prior mean is set to $\mu_0 = \frac{1}{3883} \sum_j V_j$ (average movie embedding) and the prior covariance to $\Sigma_0 = \text{diag}(v)$ where $v$ is the empirical variance of the embeddings along each dimension. This prior is fixed across all experiments and does not depend on the specific bandit instance. Despite the mismatch (bandit instances are not Gaussian), PI-BAI with this learned prior performs very well (Figure 2a, fourth column).

Online prior learning through hierarchical models. When no offline data is available, the paper proposes using the hierarchical model itself to learn the prior online. Consider the MAB setting where the true prior means $\mu_{0,i}$ are unknown. By adopting a hierarchical model with $\mu_0 \sim \mathcal{N}(\nu, \Sigma)$ (treating $\mu_0$ as a latent parameter with its own hyper-prior), the learner can infer $\mu_0$ from the observed rewards across arms while simultaneously identifying the best arm. Figure 2b shows that this online learning approach converges to the performance of PI-BAI with the correct prior as the budget grows, and substantially outperforms using a misspecified prior.

4. Key Insights and Innovations

Innovation 1: A Fully Bayesian Proof Framework That Eliminates the Need for Frequentist Intermediaries

The paper's most fundamental intellectual contribution is not the PI-BAI algorithm itself β€” which is, after all, exceptionally simple β€” but rather the novel proof technique that makes its theoretical guarantees possible. Prior work on Bayesian fixed-budget BAI, specifically BayesElim (Atsidakou et al., 2022), relied on a two-step proof structure: first bound the frequentist error probability for a fixed instance $\theta$ using concentration inequalities, then average that frequentist bound over the prior. This approach inherits all the limitations of frequentist analysis β€” it requires the posterior variances to be homogeneous across arms (forcing $\sigma_{0,i}^2 = \sigma_0^2$ for all $i$), it cannot naturally handle arm correlations from structured models, and it accumulates logarithmic factors from elimination rounds.

The key reformulation in Section 5 β€” swapping the order of conditioning to write $P_n = \mathbb{E}[\mathbb{P}(J_n \neq i^*(\theta) \mid \mathcal{H}_n)]$ rather than $\mathbb{E}[\mathbb{P}(J_n \neq i^*(\theta) \mid \theta)]$ β€” is deceptively simple but represents a genuine conceptual shift. Rather than asking "given the true parameters, what's the probability the algorithm makes a mistake?" (a frequentist question), it asks "given the data the algorithm actually observed, what's the posterior probability that the arm with the highest posterior mean isn't truly optimal?" (a Bayesian question). This transforms the problem from one about concentration of sample means around unknown true values β€” which depends on the true gaps and is hard to integrate β€” into one about the behavior of the Gaussian posterior distribution, whose parameters are known in closed form regardless of model complexity.

This is a fundamental shift, not an incremental refinement. It is what enables the paper to derive bounds that are simultaneously: (a) valid for any fixed allocation weights $\omega \in \Delta_K^+$, (b) free of assumptions about equal prior variances, (c) naturally extensible to linear and hierarchical bandits where frequentist decomposition is intractable, and (d) tighter than elimination-based bounds by roughly $\log_2^{3/2}(K)$ due to the removal of round-structure overhead. The evidence for this last point is explicit in Appendix A, where the paper compares its bound term-by-term against BayesElim's for the special case of homogeneous prior variances and uniform allocation. The $\log_2^{3/2}(K)$ improvement (a factor of ~6 for $K=10$) comes entirely from the proof structure, not from any algorithmic advantage β€” PI-BAI with uniform allocation is essentially a simpler algorithm than BayesElim, yet it achieves a stronger guarantee.

The broader significance of this proof technique extends beyond the specific bounds in this paper. It provides a template for analyzing Bayesian BAI algorithms that operate by posterior mean maximization, suggesting that similar analyses could be developed for other decision rules, other prior families (beyond Gaussian-Gaussian conjugacy), and potentially even adaptive algorithms β€” though Section 5 explicitly notes that "its application to adaptive algorithms could be challenging, particularly due to the complexity of taking expectations under the history $\mathcal{H}_n$ in that case." This honesty about limitations strengthens rather than weakens the contribution: it precisely delineates where the technique works and where new ideas are needed.

Innovation 2: Prior-Dependent Fixed Allocations as a Competitive Alternative to Adaptive Elimination

The paper makes a substantive algorithmic argument that challenges the dominant paradigm in fixed-budget BAI. For over a decade, the standard approach has been adaptive elimination: start with all arms, pull them in rounds, discard the worst, concentrate budget on survivors. This approach β€” exemplified by Sequential Halving (Karnin et al., 2013) and extended to structured settings through optimal design (Azizi et al., 2021; Katz-Samuels et al., 2020) β€” treats adaptivity as obviously beneficial. Why would you not want to stop pulling clearly suboptimal arms?

PI-BAI provides a concrete counter-argument: when you have a meaningful prior distribution, you can pre-compute a fixed allocation that optimally balances exploration across arms, and this non-adaptive strategy can match or exceed the performance of adaptive methods. The evidence is both theoretical (the bounds in Theorems 4.1, 4.2, 4.4 hold for any fixed $\omega$ and can be optimized) and empirical (Figure 2a shows PI-BAI variants consistently performing at or above the level of adaptive baselines like Sequential Halving, Successive Rejects, and BayesElim).

What makes this more than a "simple methods sometimes work" observation is the explicit connection to impossibility results. The paper grounds its approach in the recent findings of Wang et al. (2024) β€” no algorithm strictly dominates uniform sampling for two-armed Bernoulli bandits β€” and Degenne (2023) β€” no superior adaptive algorithm exists in Gaussian BAI. These results establish that adaptivity cannot provide universal advantages. PI-BAI can be understood as the Bayesian complement to these frequentist impossibility results: while adaptivity cannot dominate uniformly in the frequentist setting, prior-dependent fixed allocations can provide systematic advantages in the Bayesian setting when the prior is informative. The "motivating example" in Section 1 (3 arms, one clearly suboptimal a priori) crystallizes this: an adaptive method must still pull the suboptimal arm to confirm it's bad, while a prior-informed fixed allocation can simply ignore it.

This reframes the role of adaptivity from a presumed necessity to a design choice with tradeoffs. The paper does not claim fixed allocations dominate adaptive ones β€” that would contradict the impossibility results β€” but rather establishes that they are a competitive and theoretically well-understood alternative. This is a conceptual contribution that should influence how researchers think about algorithm design for BAI: the question is no longer "how can we be more adaptive?" but "given our prior knowledge, what is the right balance between fixed design and adaptivity?"

Innovation 3: The First Prior-Dependent Bounds for Structured Bayesian BAI, Revealing a Dimension-Dependent Complexity

Before this work, Bayesian best-arm identification with theoretical guarantees was limited to multi-armed bandits with restrictive assumptions (equal prior variances). The paper establishes the first prior-dependent upper bounds on expected PoE for linear bandits (Theorem 4.2) and hierarchical bandits (Theorem 4.4). This is not merely an incremental extension β€” it reveals a fundamental structural insight about how the statistical complexity of BAI depends on problem dimensionality.

For linear bandits, Corollary 4.3 shows that under G-optimal design, the bound depends on the ambient dimension $d$ through a factor of $2d\sigma^2$ in the denominator of the rate term: $P_n = \mathcal{O}(1/\sqrt{1 + n/(2d\sigma^2)})$. This means that the effective sample size is $n/d$, not $n/K$ β€” the statistical difficulty scales with the intrinsic dimension of the parameter space, not the number of arms. This is the Bayesian analog of the well-known frequentist result that linear bandits require $\mathcal{O}(d)$ samples rather than $\mathcal{O}(K)$, but the paper makes this explicit in the PoE bound through the G-optimal design analysis (Lemma C.9, Appendix C.3). The significance is that it provides a theoretically principled way to allocate budget in high-dimensional structured problems, with explicit constants that can guide practice.

For hierarchical bandits, the bound reveals a richer structure. The prior variance of the arm mean difference decomposes as $\|b_i - b_j\|_\Sigma^2 + \sigma_{0,i}^2 + \sigma_{0,j}^2$ in the exponential term (Theorem 4.4) β€” uncertainty about the shared effects (first term) plus arm-specific uncertainty (remaining terms). This decomposition shows that learning about effects helps all arms simultaneously: even if you never pull arm $j$, your posterior uncertainty about its mean decreases as you pull other arms that share the same effects, because $\breve{\Sigma}_n$ shrinks and pulls down $\hat{\sigma}_{n,j}^2$ through the $b_j^\top \breve{\Sigma}_n b_j$ term in Equation (17). This correlation structure is the statistical engine that makes hierarchical models efficient, and the bound makes it quantitatively explicit for the first time.

The computational complexity insight is also notable as a practical contribution: the natural hierarchical posterior update has complexity $\mathcal{O}(K + L^3)$, compared to $\mathcal{O}((K+L)^3)$ if one naively treats the hierarchical model as a linear model with $K+L$ features. This is not a theoretical advance per se, but it is a valuable practical observation for anyone implementing these methods, and it justifies treating hierarchical bandits as a distinct model class rather than merely a special case of linear bandits.

Innovation 4: Verifier Over-Optimization as a Non-Issue β€” The Bayesian Prior as a Natural Regularizer Against Overfitting to Noise

While not explicitly framed this way in the paper, a cross-cutting insight emerges from juxtaposing PI-BAI against the broader literature on test-time compute and search: the Bayesian prior acts as a natural regularizer that prevents the kind of over-optimization that plagues verifier-based search methods. In the BAI literature, elimination-based methods can suffer from "peeking" β€” eliminating an arm too early because of an unlucky sample, then never recovering. Bayesian methods with fixed allocations avoid this because there is no data-dependent elimination: every arm gets its pre-allocated budget regardless of observed rewards, and the final decision uses the posterior mean which shrinks noisy estimates toward the prior.

This is philosophically similar to the role of priors in preventing overfitting in Bayesian statistics generally, but the paper operationalizes it in the specific context of BAI. The theoretical bound reflects this regularization explicitly: the exponential term $\exp(-(\mu_{0,i} - \mu_{0,j})^2 / (2(\sigma_{0,i}^2 + \sigma_{0,j}^2)))$ discounts the error contribution from pairs with large prior gaps, regardless of what the data says. Even if the observed sample means suggest arm $i$ is worse than arm $j$, if the prior strongly favors arm $i$, the posterior mean may still select it β€” and the bound accounts for this as a low-probability error because the prior gap suggests separation. This is the Bayesian analog of the "exploration bonus" in frequentist methods, but it emerges automatically from the posterior computation rather than requiring explicit optimism.

The significance of this insight extends beyond BAI. It suggests that in any pure exploration problem where a meaningful prior exists, fixed-design Bayesian methods may be more robust to noise than adaptive frequentist methods β€” not because they are more clever, but because they refuse to let noisy early samples derail the allocation. This is a conceptual contribution to the broader debate about when Bayesian vs. frequentist approaches are preferable, with BAI serving as a concrete case study.

Innovation 5: A Unified Analysis Across Model Classes Through Posterior Structure Rather Than Algorithm Structure

A subtle but distinguishing feature of this work is that the theoretical analysis is unified by the structure of the posterior distribution, not by the structure of the algorithm. In prior work, algorithms for different bandit models (multi-armed, linear, contextual) typically require fundamentally different analyses β€” the proof for Sequential Halving in MAB (Karnin et al., 2013) shares little with the proof for GSE in linear bandits (Azizi et al., 2021), because the algorithms themselves are structurally different.

PI-BAI's analysis, in contrast, follows the exact same proof template across all three model classes (Section 5, then specialized in Appendices C.2, C.3, C.4):

  1. Decompose the expected PoE into pairwise posterior error probabilities.
  2. Bound each pairwise term using Gaussian tail inequalities on the posterior.
  3. Integrate over the marginal distribution of the history using Lemma C.5.
  4. Compute the required moments (mean and variance of posterior mean differences) for the specific model.

The only thing that changes between model classes is Step 4 β€” the computation of $\mathbb{E}[\hat{\mu}_{n,i}]$ and $\mathbb{V}(\hat{\mu}_{n,i} - \hat{\mu}_{n,j})$ β€” which depends on the prior and likelihood structure. For MAB, this is a simple variance calculation with independent arms. For linear bandits, it involves the posterior covariance matrix $\hat{\Sigma}_n$ and the covariance of $B_n$. For hierarchical bandits, it requires propagating uncertainty through two levels of the prior and computing cross-arm covariance terms (the formidable Equation 30). But the proof architecture remains identical.

This is significant because it demonstrates that algorithm simplicity enables analytical unity. Because PI-BAI does the same thing regardless of model β€” pull arms according to fixed weights, compute posterior means, pick the highest β€” the analysis can be modular: the probabilistic argument (Steps 1–3) is model-agnostic, and the model-specific computation (Step 4) is a self-contained Gaussian calculation. This is in stark contrast to adaptive elimination methods, where the interplay between data-dependent decisions and error probabilities entangles the analysis with the specific elimination schedule and model structure.

The practical implication is that extending PI-BAI-style analysis to new model classes (e.g., generalized linear bandits, as discussed in Appendix B.3) reduces to the tractable problem of computing posterior moments β€” potentially via Laplace approximation or variational inference β€” rather than the much harder problem of designing and analyzing a new adaptive algorithm from scratch. The paper does not fully deliver on this promise (the logistic bandit extension in Appendix D.4 is empirical only, without theoretical bounds), but the framework provides a clear roadmap for future theoretical work.

5. Experimental Analysis

Evaluation Methodology

  • Dataset. The paper uses three categories of experiments: (a) Synthetic data generated from the Gaussian priors described in each bandit model (MAB, linear, hierarchical), with parameters chosen to create informative but not trivial settings (e.g., prior means sampled from $\mathcal{U}([0,1])$, prior standard deviations evenly spaced between 0.1 and 0.5). For the "Fixed" setting (Appendix D.2), prior means are set deterministically (e.g., evenly spaced between 0 and 1) rather than sampled, creating a single consistent problem instance across runs. (b) MovieLens 1M (Lam and Herlocker, 2016): a real-world dataset with 1 million ratings from 6,040 users for 3,883 movies, forming a sparse rating matrix $M$ of size $6,040 \times 3,883$. The authors complete $M$ using alternating least squares with rank $d = 5$, extract $M = U^\top V$, and construct a linear Gaussian prior with $\mu_0 = \frac{1}{3883} \sum_j V_j$ (mean movie embedding) and $\Sigma_0 = \text{diag}(v)$ where $v$ is the empirical variance of the embeddings. A subset of $K = 100$ randomly subsampled movies is used in each bandit instance. (c) Logistic bandit experiments (Appendix D.4) use the same synthetic prior generation but with Bernoulli rewards and logistic link, testing generalization beyond the Gaussian setting.

  • Base model(s). All experiments use Gaussian bandit instances with observation noise fixed at $\sigma = 1$. The bandit instances are sampled from the priors described in the text: for MAB, $K = 10$ arms with $\mu_{0,i} \sim \mathcal{U}([0,1])$ and $\sigma_{0,i}$ evenly spaced in $[0.1, 0.5]$ ("Random setting"); for linear bandits, $K = 30$ arms, $d = 4$ dimensions, with $\mu_{0,i} \sim \mathcal{U}([0,1])$ and $\Sigma_0$ diagonal with entries evenly spaced in $[0.1, 0.5]$; for hierarchical bandits, $K = 60$ arms, $L = 10$ effects, mixing weights $b_i$ sampled uniformly in $[0,1]$ and normalized to sum to 1, $\nu_i \sim \mathcal{U}([-1,1])$, and both $\Sigma$ and $\Sigma_0$ diagonal with entries evenly spaced in $[0.1^2, 0.5^2]$. These choices produce non-trivial identification problems where the prior provides meaningful but imperfect information. The model sizes are consistent with those used in prior BAI work and allow computational tractability for $10^4$ repetitions.

  • Metrics. The primary metric throughout is the expected probability of error (PoE) β€” the fraction of the $10^4$ runs in which the arm returned by the algorithm after $n$ rounds is not the true optimal arm $i^*(\theta)$ for the sampled instance $\theta$. This is an empirical estimate of $P_n = \mathbb{E}[\mathbb{P}(J_n \neq i^*(\theta) \mid \theta)]$. Standard errors are computed across the $10^4$ runs and are reported as "narrow" (Section 6), indicating statistical reliability. For Figure 2b, the metric is again average PoE, but the comparison is between PI-BAI variants with different prior knowledge (correct, misspecified, learned online). Appendix D.6 additionally reports Bayesian simple regret as $\mathbb{E}[\theta_{i^*(\theta)} - \theta_{J_n}]$ (not $r(i; \theta)$ notation β€” the simple regret is the difference in true means between the optimal arm and the chosen arm, integrated over the prior and the algorithm's randomness). The authors discuss in Appendix A why simple regret and PoE may have different rates in the Bayesian setting (Komiyama et al., 2023 showed $\mathcal{O}(1/n)$ for Bayesian simple regret vs. $\mathcal{O}(1/\sqrt{n})$ for PoE from Atsidakou et al., 2022).

  • Baselines. The paper compares against a comprehensive set of both Bayesian and frequentist methods. For multi-armed bandits: BayesElim (Atsidakou et al., 2022) β€” the only prior Bayesian BAI algorithm with theoretical guarantees, using elimination rounds with posterior-based arm discarding; Top-Two Thompson Sampling (TTTS) (Jourdan et al., 2022; Russo, 2016) β€” an adaptive Bayesian method that resamples arms from the posterior and compares the top two candidates, denoted as TS-TCI with $\beta = 0.5$ in the paper (Appendix D.5); Sequential Halving (SH) (Karnin et al., 2013) β€” the canonical frequentist elimination algorithm that halves the active arm set each round; Successive Rejects (SR) (Audibert et al., 2010) β€” a frequentist elimination method with a more aggressive rejection schedule. For linear bandits: BayesGap (Hoffman et al., 2014) β€” a Bayesian method using G-optimal design for successive elimination; GSE (Azizi et al., 2021) β€” a frequentist method using G-optimal design for successive elimination, representing "the current (tractable) state-of-the-art approaches" (Section 6). The paper notes that Katz-Samuels et al. (2020) has tighter bounds but is computationally intractable, and other elimination-based methods (Alieva et al., 2021; Yang and Tan, 2022) have lower performance, so they are excluded. For hierarchical bandits: the MAB frequentist baselines (SH, SR) are given a marginalized Gaussian MAB prior (Equation 12) since they cannot directly use the hierarchical structure. The Bayesian baselines (BayesElim, TTTS) also receive this marginalized prior because they "are not suitable for a hierarchical prior" (Section 6). Linear bandit baselines are excluded from hierarchical experiments because treating the hierarchical model as a $(K+L)$-dimensional linear model leads to poor performance "due to their inefficient representation of the structure" (Section 6). For MovieLens: BayesGap and GSE are used as linear bandit baselines, since the learned prior induces a linear Gaussian model. For logistic bandits (Appendix D.4): only GSE is included as a baseline (BayesGap is not tested), and the comparison focuses on PI-BAI variants.

  • Generation budget / compute accounting. The budget is measured in total number of arm pulls $n$ (rounds). All methods β€” both adaptive and non-adaptive β€” are compared at the same total budget. For PI-BAI, the allocation vector $\omega$ is computed offline (no budget cost for optimization or design calculations), and $n_i = \lfloor \omega_i n \rfloor$ pulls are performed. For PI-BAI(TS), the warm-up phase consumes $n_w$ pulls (set to $n_w = K$ throughout the main experiments), which are included in the total budget $n$ β€” the remaining $n - n_w$ pulls are allocated according to $\omega^{\text{TS}}$. For adaptive baselines, the full budget $n$ is spent according to each algorithm's internal logic. Budgets are swept across a range appropriate to each setting: for MAB, $n \in [50, 300]$ roughly; for linear bandits, $n \in [100, 1000]$; for hierarchical bandits, $n \in [150, 550]$; for MovieLens, $n \in [200, 1800]$ (all visible from Figure 2a x-axes). The observation noise is $\sigma = 1$ throughout, making budgets directly comparable across settings in terms of signal-to-noise ratio per pull.

  • Cross-validation / statistical protocol. No cross-validation is needed because PI-BAI's allocation weights are computed from the prior, not learned from data that needs to be held out. For the optimized allocation PI-BAI(Opt), the numerical optimization of the theoretical bound (Equation 3) is performed once offline on the known prior β€” there are no hyperparameters selected using test data. For PI-BAI(TS), the warm-up length $n_w = K$ was selected based on ablation studies in Appendix D.7 (Figure 10) that sweep $n_w$ across multiple values for different $K$ and budgets; the choice $n_w = K$ is near-optimal across settings. The mixture parameter $\alpha = 0.5$ for PI-BAI(Opt) was selected based on Figure 12 in Appendix D.7, which shows that $\alpha \approx 0.5$ minimizes PoE across both Fixed and Random settings. For MovieLens, the prior is learned from the entire MovieLens dataset using low-rank factorization β€” this prior is fixed across all MovieLens bandit experiments and does not use bandit rewards, so there is no data leakage. Statistical reliability is ensured by $10^4$ independent repetitions: for each repetition, a new bandit instance $\theta$ is sampled from the prior (for synthetic experiments), the algorithm is run on that instance, and the outcome is recorded. Standard errors are computed across these repetitions. For the "Fixed" setting (Appendix D.2), the prior parameters are deterministic, so the $10^4$ repetitions sample only the observation noise β€” testing algorithm performance on a single hard bandit configuration.

Main Quantitative Results

The experiments are organized into four blocks: comparison to baselines across settings (Figure 2a), prior misspecification analysis (Figure 2b, Figure 5 in Appendix D.1), ablation and hyperparameter studies (Appendix D), and extension to logistic bandits (Appendix D.4).

Multi-Armed Bandit Results (Figure 2a, first column)

Headline: All four PI-BAI variants (Uni, Opt, G-opt, TS) outperform the best baseline across most budget levels. At the highest budget shown ($n \approx 250$), PI-BAI(G-opt) achieves the lowest PoE (~0.22-0.23), followed closely by PI-BAI(Opt) and PI-BAI(TS), while the best baseline (BayesElim) reaches approximately 0.28-0.30. The gap widens at intermediate budgets: at $n \approx 100$, PI-BAI variants achieve ~0.30-0.35 PoE while BayesElim is at ~0.40.

Side-by-side at matched budgets: At $n = 50$, PI-BAI(Uni) achieves ~0.55 PoE, PI-BAI(TS) ~0.52, PI-BAI(Opt) ~0.48, PI-BAI(G-opt) ~0.46, while BayesElim is at ~0.60, TTTS at ~0.58, SH at ~0.62, and SR at ~0.63. At $n = 150$, the ordering is: PI-BAI(G-opt) ~0.28, PI-BAI(Opt) ~0.30, PI-BAI(TS) ~0.31, PI-BAI(Uni) ~0.33, BayesElim ~0.37, TTTS ~0.38, SH ~0.42, SR ~0.44. At $n = 250$, PI-BAI(G-opt) and PI-BAI(Opt) converge to ~0.22-0.23, while BayesElim reaches ~0.28-0.30 and the frequentist methods remain above 0.35.

Key pattern: Uniform allocation (PI-BAI(Uni)) β€” which is the simplest possible non-adaptive strategy β€” already outperforms all adaptive frequentist baselines (SH, SR) and is competitive with TTTS at higher budgets. The Bayesian baselines (BayesElim, TTTS) use the same informative prior but perform worse, suggesting that (a) BayesElim's elimination structure discards useful information (the paper notes it forces discarding of all previous observations at the end of each round), and (b) TTTS's adaptive exploration does not translate to better final decisions than a well-chosen fixed allocation. This is the empirical counterpart to the theoretical $\log_2^{3/2}(K)$ improvement in the bound.

Linear Bandit Results (Figure 2a, second column)

Headline: PI-BAI variants substantially outperform the linear-specific baselines BayesGap and GSE across the entire budget range ($n = 100$ to $n = 1000$). PI-BAI(G-opt) is the clear winner, with PI-BAI(TS) close behind. The gap is largest at intermediate budgets: at $n \approx 400$, PI-BAI(G-opt) achieves ~0.30 PoE while BayesGap is at ~0.45 and GSE at ~0.52.

Side-by-side at matched budgets: At $n = 100$, PI-BAI(G-opt) achieves ~0.57 PoE, PI-BAI(TS) ~0.58, PI-BAI(Opt) ~0.60, PI-BAI(Uni) ~0.62, while BayesGap is at ~0.68 and GSE at ~0.72. At $n = 400$, PI-BAI(G-opt) reaches ~0.30, PI-BAI(TS) ~0.32, PI-BAI(Opt) ~0.34, PI-BAI(Uni) ~0.40, BayesGap ~0.45, GSE ~0.52. At $n = 1000$, all PI-BAI variants converge to ~0.18-0.22, with PI-BAI(G-opt) at ~0.18, PI-BAI(TS) at ~0.19, while BayesGap is at ~0.30 and GSE at ~0.35.

Key pattern: The advantage of PI-BAI is more pronounced here than in the MAB setting β€” the gap between PI-BAI(G-opt) and the best baseline (BayesGap) at $n = 1000$ is roughly 12 percentage points, compared to roughly 6 points in MAB. This validates the paper's claim that "our algorithms and proofs are naturally applicable to structured problems" (Section 1): the baselines, which use G-optimal design for successive elimination, are structurally well-suited to linear bandits, yet PI-BAI's Bayesian approach extracts more value from the prior. The fact that PI-BAI(G-opt) β€” which uses the same G-optimal design criterion as BayesGap and GSE but applies it to a non-adaptive design rather than elimination rounds β€” substantially outperforms both suggests that the elimination structure itself is the performance bottleneck, not the design criterion.

Hierarchical Bandit Results (Figure 2a, third column)

Headline: PI-BAI(TS) is the standout performer, substantially outperforming all other methods. At the highest budget ($n \approx 550$), PI-BAI(TS) achieves ~0.55 PoE, while PI-BAI(G-opt) reaches ~0.62, PI-BAI(Opt) ~0.65, PI-BAI(Uni) ~0.70, and the best baseline (BayesElim) is at ~0.72. The MAB baselines (TTTS, SH, SR) all perform worse, staying above 0.75 at all budgets.

Side-by-side at matched budgets: At $n = 150$, PI-BAI(TS) achieves ~0.75 PoE, PI-BAI(G-opt) ~0.80, PI-BAI(Opt) ~0.82, PI-BAI(Uni) ~0.85, while BayesElim is at ~0.87 and TTTS at ~0.90. At $n = 350$, PI-BAI(TS) reaches ~0.64, PI-BAI(G-opt) ~0.72, PI-BAI(Opt) ~0.76, PI-BAI(Uni) ~0.78, BayesElim ~0.82, TTTS ~0.84. At $n = 550$, the ordering is: PI-BAI(TS) ~0.55, PI-BAI(G-opt) ~0.62, PI-BAI(Opt) ~0.65, PI-BAI(Uni) ~0.70, BayesElim ~0.72, TTTS ~0.75, SR ~0.78, SH ~0.80.

Key pattern: The hierarchical setting reveals a striking reversal: PI-BAI(Opt), which performed near the top in MAB and linear settings, drops to third place among PI-BAI variants. The authors note that "this method reaches its limits due to the complexity of the bound to be optimized" (Section 4.3) β€” the hierarchical bound in Theorem 4.4 involves the complex covariance expression in Equation (30), and numerical optimization of this high-dimensional non-convex function appears to find suboptimal local minima. PI-BAI(TS) excels because Thompson sampling naturally navigates the hierarchical structure: the mixed-effect Thompson sampling of Aouali et al. (2023) samples from the joint posterior over effects and arm means, allocating pulls to arms that are promising given the inferred effects. This warm-up phase effectively learns a good allocation for the specific bandit instance without needing to optimize the unwieldy theoretical bound.

The baselines' poor performance confirms that "the MAB frequentist baselines do not require priors, while the Bayesian ones need a Gaussian MAB prior since they are not suitable for a hierarchical prior" (Section 6). Marginalizing the effects (Equation 12) discards the correlation structure that PI-BAI exploits β€” arms that share effects provide information about each other, and this information is captured in the hierarchical posterior but lost in the marginalized MAB prior.

MovieLens Results (Figure 2a, fourth column)

Headline: Despite the learned prior being a Gaussian approximation to non-Gaussian data, PI-BAI variants perform well. PI-BAI(G-opt) achieves the lowest PoE across budgets, reaching ~0.18 at $n = 1800$, compared to ~0.25 for BayesGap and ~0.30 for GSE. PI-BAI(TS) is competitive at lower budgets but flattens out above $n \approx 1000$.

Side-by-side at matched budgets: At $n = 200$, PI-BAI(Uni) and PI-BAI(TS) achieve ~0.38, PI-BAI(G-opt) ~0.35, BayesGap ~0.42, GSE ~0.48. At $n = 1000$, PI-BAI(G-opt) reaches ~0.22, PI-BAI(TS) ~0.24, PI-BAI(Opt) ~0.25, PI-BAI(Uni) ~0.27, BayesGap ~0.31, GSE ~0.38. At $n = 1800$, the gap widens: PI-BAI(G-opt) ~0.18, PI-BAI(TS) ~0.22, BayesGap ~0.25, GSE ~0.30.

Key pattern: The paper emphasizes that "the bandit instances are not sampled from a Gaussian prior, yet our algorithm employs a learned Gaussian prior, highlighting the robustness of our method to prior misspecification" (Section 6). This is significant because the theoretical guarantees assume the prior matches the data-generating process. The empirical success suggests that a reasonably well-learned Gaussian prior β€” even when the true data distribution is not Gaussian β€” provides sufficient information to guide allocation effectively. The flattening of PI-BAI(TS) at high budgets may reflect the warm-up phase's finite-sample limitation: with $n_w = K = 100$ warm-up pulls and a total budget of $n = 1800$, the warm-up phase represents only ~5.5% of the budget, and the remaining ~94.5% is spent according to the empirical allocation. If the warm-up phase misidentifies which arms are promising, the subsequent fixed allocation locks in that error.

Prior Misspecification: Online Learning vs. Misspecified Prior (Figure 2b)

Headline: PI-BAI with the correct prior outperforms both the misspecified-prior variant and the online-learning variant at small budgets, but the online-learning variant converges to the performance of the correct-prior variant as $n$ grows (across all four allocation strategies: Uni, Opt, G-opt, TS). The misspecified-prior variant (with $\tilde{\mu}_0 = \mu_0 + 0.5\epsilon$, $\epsilon \sim \mathcal{N}(0,1)$) shows a permanent gap that narrows but does not close.

Detailed findings (Figure 2b):

  • Uniform allocation (leftmost panel): At $n = 30$, correct prior achieves ~0.45 PoE, online learning ~0.52, misspecified ~0.56. By $n = 300$, correct ~0.28, online learning ~0.30, misspecified ~0.34. The online-learning curve converges to within ~2 percentage points of the correct prior by $n \approx 150$.

  • Optimized allocation (second panel): At $n = 30$, correct ~0.42, online learning ~0.50, misspecified ~0.54. By $n = 300$, correct ~0.25, online learning ~0.28, misspecified ~0.33. The convergence behavior is similar.

  • G-optimal allocation (third panel): At $n = 30$, the spread is wider: correct ~0.44, online learning ~0.54, misspecified ~0.62. By $n = 300$, correct ~0.26, online learning ~0.30, misspecified ~0.38. The misspecified prior performs notably worse under G-optimal design than under optimized or uniform allocation at small budgets.

  • TS warm-up allocation (rightmost panel): At $n = 30$, correct ~0.40, online learning ~0.52, misspecified ~0.58. By $n = 300$, correct ~0.24, online learning ~0.30, misspecified ~0.42. The misspecified prior shows the largest and most persistent gap under TS warm-up β€” the authors attribute this to the fact that "the misspecified prior is used twice: by the warm-up policy TS and by PI-BAI" (Section 6). The double use compounds the error: Thompson sampling explores the wrong arms during warm-up, and PI-BAI then allocates based on those misleading pull counts.

Key pattern: The hierarchical online learning approach β€” treating the prior means as unknown and learning them through a hierarchical model with $\nu \sim \mathcal{U}([0,1])$ and $\Sigma = I_K$ β€” works remarkably well. By $n \approx 150-200$ (roughly 15-20 pulls per arm on average), it has essentially recovered the correct prior from data. This validates the practical strategy proposed in Section 4.4: "In the absence of offline data, one approach to address prior misspecification online is to use hierarchical models." The convergence rate is consistent with Lemma 4.5's prediction that $d_n(i,j) = e^{\mathcal{O}(1/n)}$ β€” the impact of misspecification decays as $\mathcal{O}(1/n)$ in the exponent.

Additional Setting: Fixed Prior (Appendix D.2, Figure 6)

The "Fixed setting" removes randomness in the prior, testing algorithm performance on a single hard bandit configuration. The qualitative patterns are preserved: PI-BAI variants outperform all baselines across MAB, linear, and hierarchical settings. In the MAB Fixed setting (Figure 6, first column), the gap between PI-BAI and BayesElim is smaller than in the Random setting, with PI-BAI(G-opt) achieving ~0.22 at $n = 300$ vs. BayesElim at ~0.26. In the linear Fixed setting (second column), PI-BAI(G-opt) reaches ~0.32 at $n = 1000$ vs. BayesGap at ~0.52 β€” a larger gap than in the Random setting. In the hierarchical Fixed setting (third column), PI-BAI(TS) dominates, reaching ~0.55 at $n = 550$ while BayesElim stays above 0.75. The consistency across both Fixed and Random settings indicates that PI-BAI's advantage is not an artifact of averaging over easy instances.

Logistic Bandit Extension (Appendix D.4, Figure 7)

Headline: PI-BAI with G-optimal design allocations (computed as if the bandit were linear-Gaussian) generalizes to logistic bandits with Laplace-approximated posteriors. In the Random setting with $K = 30$, $d = 4$, PI-BAI(G-opt) achieves ~0.42 PoE at $n = 1000$ vs. GSE at ~0.60. The gap is narrower in lower-dimensional settings ($d = 3$: PI-BAI(G-opt) ~0.60 vs. GSE ~0.72 at $n = 1000$).

This is a purely empirical result with no theoretical guarantees β€” the paper explicitly states in Appendix B.3 that "proving an upper bound on the expected PoE of this algorithm is challenging." The Laplace approximation introduces error in the posterior mean estimates, and the G-optimal design is computed under the (incorrect) assumption of Gaussian rewards. The fact that performance remains strong despite these approximations is encouraging for practical applications but does not come with the same level of theoretical support as the Gaussian results.

Ablation Studies and Robustness Checks

  • Warm-up length $n_w$ for PI-BAI(TS) (Appendix D.7, Figure 10): The performance is relatively flat across warm-up lengths in the range $n_w \in [K, 2K]$. For $K = 10$, $n = 100$, PoE varies from ~0.36 at $n_w = K$ to ~0.34 at $n_w = 2K$ (Fixed setting) and from ~0.42 to ~0.40 (Random setting). For $K = 20$, $n = 200$, the optimal is $n_w = K = 20$ in the Fixed setting (~0.38 vs. ~0.40-0.42 at other values) and $n_w = 2K = 40$ in the Random setting (~0.36 vs. ~0.38 at $n_w = K$). The choice $n_w = K$ is near-optimal in most configurations and is used throughout the main experiments. The flatness suggests that the warm-up phase is not critical β€” as long as it's long enough to sample each arm at least once, the subsequent fixed allocation dominates performance.

  • Choice of warm-up policy: TS vs. T3C vs. TSTCI (Appendix D.7, Figure 11): For the MAB Fixed setting, TS and T3C perform nearly identically (~0.25-0.28 PoE at $n = 250$), with TSTCI slightly worse (~0.30-0.33). For the MAB Random setting, TS is the clear winner (~0.30 at $n = 250$ vs. ~0.34-0.36 for T3C/TSTCI). For the hierarchical Fixed setting, the mixed-effect TS of Aouali et al. (2023) dominates (~0.55 at $n = 550$ vs. ~0.60-0.63 for T3C/TSTCI). For the hierarchical Random setting, meTS reaches ~0.55 vs. ~0.58-0.60. The paper's choice of TS for MAB and meTS for hierarchical settings is empirically justified.

  • Mixture parameter $\alpha$ for PI-BAI(Opt) (Appendix D.7, Figure 12): Sweeping $\alpha \in [0, 1]$ for $K = 5$ and $K = 10$ with various budgets $n$. In the Fixed setting, performance improves monotonically as $\alpha$ decreases from 1 (pure optimized weights) to 0 (pure heuristic $\mu_{0,i}\sigma_{0,i}$ weights), with $\alpha \approx 0$ achieving the lowest PoE (e.g., for $K = 10, n = 100$, PoE drops from ~0.38 at $\alpha = 1$ to ~0.33 at $\alpha = 0$). In the Random setting, the pattern reverses: pure optimized weights ($\alpha = 1$) outperform pure heuristic ($\alpha = 0$), with the minimum typically at $\alpha \approx 0.5$ (e.g., $K = 10, n = 100$: PoE ~0.48 at $\alpha = 0$, ~0.44 at $\alpha = 0.5$, ~0.46 at $\alpha = 1$). The paper adopts $\alpha = 0.5$ as a compromise. This is a non-obvious finding: the heuristic weights help when problem instances are uniformly difficult (Fixed setting) but hurt when there is instance-to-instance variation that the optimized weights are designed to handle (Random setting).

  • Elimination vs. no elimination on top of PI-BAI (Appendix D.5, Figure 8): Adding elimination rounds (as in BayesElim) to PI-BAI does not improve performance and often degrades it. For PI-BAI(Opt) in the MAB Fixed setting, PI-BAI(Opt) without elimination achieves ~0.24 PoE at $n = 250$, while PI-BAI(Opt) with elimination achieves ~0.26 β€” a small but consistent degradation. For PI-BAI(G-opt), the gap is larger: ~0.22 without elimination vs. ~0.28 with elimination. For PI-BAI(TS), elimination hurts substantially: ~0.25 without vs. ~0.35 with. The paper notes that elimination methods "discard all observations at the end of each round" (Appendix A), which wastes information. The non-elimination version retains all data for the posterior computation, which is particularly important for PI-BAI(TS) where the warm-up phase already provides a good allocation and subsequent elimination would discard useful observations.

  • Effect of prior misspecification severity (Appendix D.1, Figure 5): Sweeping the misspecification magnitude $\alpha$ (prior mean shift or prior variance inflation) from 0 to 0.5 shows that performance degrades monotonically with $\alpha$ for all allocation strategies, but the effect diminishes as $n$ grows β€” confirming Lemma 4.5's prediction. For uniform allocation with mean misspecification: at $n = 30$, PoE ranges from ~0.45 ($\alpha = 0$) to ~0.58 ($\alpha = 0.5$); at $n = 300$, the range narrows to ~0.28 ($\alpha = 0$) to ~0.34 ($\alpha = 0.5$). For TS warm-up with mean misspecification, the gap is larger and more persistent: at $n = 30$, ~0.40 to ~0.66; at $n = 300$, ~0.24 to ~0.42. Variance misspecification (second row of Figure 5) shows similar patterns but with smaller absolute degradation. The key takeaway: theoretically-grounded allocations (Uni, Opt, G-opt) are more robust to prior misspecification than TS warm-up, because the misspecification only affects PI-BAI once (through the allocation computation), whereas for TS it affects both the warm-up exploration and the final allocation.

  • Simple regret metric (Appendix D.6, Figure 9): Under simple regret, the qualitative ordering of methods is preserved: PI-BAI variants outperform baselines. In the MAB Fixed setting at $n = 300$, PI-BAI(G-opt) achieves simple regret ~0.02, PI-BAI(Opt) ~0.03, while BayesElim is at ~0.07 and SH at ~0.10. Note that simple regret values are much smaller than PoE values because even when the wrong arm is selected, it is typically close in mean to the optimal arm. The convergence rate appears faster than $\mathcal{O}(1/\sqrt{n})$ visually β€” consistent with the $\mathcal{O}(1/n)$ asymptotic lower bound of Komiyama et al. (2023) for Bayesian simple regret. The paper does not provide theoretical bounds on simple regret, noting in Appendix A that "upper bounding $\Delta_{\max}$ by a pure constant would lead to a suboptimal rate of $\mathcal{O}(1/\sqrt{n})$ for the simple regret."

  • Posterior covariance comparison: hierarchical vs. marginalized MAB (Appendix B.2, Figure 3): Demonstrates the statistical efficiency gain from using the hierarchical model directly rather than marginalizing to an MAB prior. For setting 1 ($\Sigma = 2I_L$, large effect uncertainty), the average posterior variance across arms under the hierarchical model drops from ~0.28 to ~0.02 as $n$ goes from 0 to 200, while under the marginalized MAB model it drops from ~0.28 to only ~0.14 β€” the hierarchical model achieves roughly half the posterior variance for the same budget. For setting 2 ($\Sigma = 10^{-3}I_L$, small effect uncertainty), the gap is narrower (~0.24 vs. ~0.20 at $n = 200$), consistent with intuition: when effects are nearly deterministic, the hierarchical model provides less benefit over a well-chosen MAB prior. This explains why PI-BAI(TS) with the mixed-effect Thompson sampling excels in the hierarchical setting (Figure 2a, third column) β€” the posterior variance reduction translates directly to lower PoE.

Critical Assessment

Claim 1: "PI-BAI provides tighter prior-dependent upper bounds than existing approaches"

The comparison in Figure 1 directly supports this claim for the MAB setting with $K = 3$. The PI-BAI bound (computed analytically from Theorem 4.1) is substantially below the BayesElim bound across all budgets $n \in [50, 200]$ and for both homogeneous and heterogeneous prior variance settings. The gap is quantitatively large: at $n = 50$ in the homogeneous setting, the PI-BAI(Opt) bound is ~0.48 while BayesElim's is ~0.72 β€” roughly a 50% reduction. At $n = 200$, PI-BAI(Opt) bound is ~0.42 vs. BayesElim's ~0.55.

However, three important caveats apply:

  • The bound comparison is only shown for $K = 3$ in MAB. The $\log_2^{3/2}(K)$ improvement factor is derived theoretically in Appendix A (comparing the formulas term-by-term) but is not validated empirically for larger $K$. For $K = 100$, the theoretical factor would be $\log_2^{3/2}(100) \approx (6.64)^{1.5} \approx 17.1$, but Figure 1 only shows $K = 3$. The empirical gap in the bound might not scale exactly as theory predicts because the bounds are sums over $K(K-1)$ pairs and the dominant terms may differ between the two bounds.
  • The bound comparison is for upper bounds, not actual performance. A tighter upper bound does not guarantee better empirical performance β€” it could simply mean that PI-BAI's bound is less conservative. The empirical performance comparisons in Figure 2a show that PI-BAI indeed performs better, but the bound tightening and the performance improvement are separate claims.
  • The bound for PI-BAI(Opt) is computed by numerically optimizing the upper bound itself. This introduces a potential circularity: the bound being evaluated is the same one being minimized. A fairer comparison would evaluate PI-BAI(Opt)'s empirical PoE against BayesElim's bound, or compare both algorithms' actual performance (which Figure 2a already does).

Claim 2: "PI-BAI achieves consistent and robust empirical performance across synthetic and real-world settings"

This claim is well-supported by Figure 2a, which covers four distinct settings, and by the additional Fixed setting results in Appendix D.2 (Figure 6). The consistency is genuine: in every single setting, across all budget levels, at least one PI-BAI variant is at or near the top, and no baseline consistently outperforms PI-BAI. The robustness is demonstrated by:

  • The MovieLens results showing strong performance with a learned (misspecified) prior.
  • Figure 2b showing that online learning of the prior recovers most of the performance of the correct prior.
  • The logistic bandit extension (Figure 7) showing that the approach generalizes beyond Gaussian rewards (empirically, without theory).

Genuine weaknesses:

  • The best PI-BAI variant varies across settings (G-opt wins in MAB, linear, and MovieLens; TS wins in hierarchical). This means that a practitioner must choose which variant to use without clear guidance β€” the paper does not provide a principled selection rule. The fact that PI-BAI(Opt) underperforms in the hierarchical setting despite being theoretically grounded (it directly optimizes the bound) is a practical concern: if the bound optimization fails due to non-convexity, users need to know when to fall back to G-opt or TS.
  • The baselines for hierarchical bandits are weakened by being given a marginalized MAB prior rather than being adapted to the hierarchical structure. The paper acknowledges this (Section 6: "the Bayesian ones need a Gaussian MAB prior since they are not suitable for a hierarchical prior"), but this means the comparison is not fully fair β€” it shows PI-BAI exploiting hierarchical structure against baselines that cannot. A fairer comparison would adapt BayesElim or TTTS to use the hierarchical posterior, though the paper argues this is not straightforward.
  • The MovieLens prior is learned from the entire dataset, which contains all user ratings for all movies β€” vastly more data than a typical BAI scenario. The prior is fixed across all bandit instances, but it is learned from what amounts to a complete observation of the reward matrix. In a genuine BAI application, such rich offline data might not be available. The learned prior may be unrealistically informative, inflating PI-BAI's apparent advantage.

Claim 3: "Non-adaptive prior-informed strategies can match or exceed adaptive elimination-based methods"

This is the paper's central algorithmic claim, and Figure 2a provides strong evidence. However, the claim needs qualification:

  • "Can match or exceed" β€” yes, under the specific prior distributions tested. The priors are informative (prior standard deviations of 0.1-0.5 with observation noise $\sigma = 1$ means the prior provides substantial information: a prior SD of 0.5 means the prior is roughly equivalent to 4 observations per arm in terms of precision). For less informative priors (larger $\sigma_{0,i}$), the advantage of PI-BAI would shrink because the allocation would converge toward uniform (as shown by the $K=2$ closed form where $\omega^{\text{opt}} \to 1/2$ as $n \to \infty$ or as $\sigma_{0,i}^2$ becomes large). The paper does not explore the boundary where the prior becomes uninformative β€” this would be a valuable ablation to understand when PI-BAI's advantage disappears.
  • The warm-up variant PI-BAI(TS) IS adaptive β€” it uses Thompson sampling to adapt to the specific bandit instance during the warm-up phase. So the claim that "non-adaptive" strategies work well applies most cleanly to PI-BAI(Uni), PI-BAI(Opt), and PI-BAI(G-opt), which indeed use zero data-dependent adaptation. The fact that the warm-up variant often performs best (especially in hierarchical settings) suggests that a small amount of adaptivity (just enough to learn a good allocation) is beneficial, even if full elimination-based adaptivity is not.
  • Missing baseline: adaptive Bayesian BAI with the same posterior decision rule. All adaptive baselines (BayesElim, TTTS, SH, SR) have different internal mechanisms than PI-BAI β€” not just different adaptivity, but different ways of using data and making decisions. A cleaner ablation would be: PI-BAI with fixed allocation vs. an algorithm that adaptively re-allocates budget online but uses the same posterior mean decision rule. This would isolate the effect of adaptivity from other confounding factors (elimination discarding data, TTTS's sampling-based exploration, etc.).

Claim 4: "First prior-dependent bounds for linear and hierarchical Bayesian BAI"

This is a factual claim about the literature, and it appears to be correct based on the paper's review. The paper provides explicit bounds in Theorems 4.2 and 4.4. However, the empirical validation of these bounds is incomplete:

  • The bounds themselves are not plotted for linear or hierarchical settings β€” Figure 1 only shows MAB bounds. It would be valuable to see how the linear and hierarchical bounds behave numerically (e.g., how they vary with $d$, with the structure of $\Sigma_0$, with the mixing weights $b_i$).
  • The bound for hierarchical bandits (Theorem 4.4) is extraordinarily complex (Equation 30 spans nearly a page). The paper does not discuss whether this bound is computable in practice for reasonable $K$ and $L$ β€” the covariance calculation involves $\mathcal{O}(K^2 L^2)$ terms. For the $K = 60$, $L = 10$ setting in the experiments, this would require computing roughly $60^2 \times 10^2 / 2 \approx 1.8 \times 10^5$ covariance terms across all arm pairs β€” potentially expensive but feasible. For larger problems, the bound may be computationally intractable, limiting its practical use for guiding allocation (though the algorithmic variants G-opt and TS remain applicable).
  • The hierarchical bound was not used for optimization (PI-BAI(Opt) underperformed), so the theoretical guarantee, while existent, did not translate to an effective allocation strategy in this setting. This is a gap between theory and practice that the paper acknowledges but does not resolve.

Missing Experiments That Would Strengthen the Paper

  • Scaling with $K$ and $d$. All experiments use fixed, relatively small values ($K = 10, 30, 60$; $d = 4$; $L = 10$). How does the gap between PI-BAI and baselines scale with $K$? The theoretical $\log_2^{3/2}(K)$ improvement suggests the gap should widen. For linear bandits, how does performance scale with $d$? Corollary 4.3 predicts $P_n = \mathcal{O}(1/\sqrt{1 + n/(2d\sigma^2)})$ β€” testing this over a range of $d$ values would validate the dimension-dependence.
  • Less informative priors. How does PI-BAI perform when $\sigma_{0,i} \gg \sigma$ (prior is very diffuse)? The theory predicts convergence to uniform allocation and equivalence with frequentist methods. Testing this would establish the boundary of PI-BAI's advantage.
  • Comparison to the asymptotic optimal allocation. For $K = 2$, the exact optimal allocation is known (the closed form in Section 4.1). Does PI-BAI(Opt) recover it numerically? For $K > 2$, the optimal allocation is unknown, but one could compare against brute-force search over allocations for very small $K$ (e.g., $K = 3, 4$) to assess how close PI-BAI(Opt) comes to the true optimum.
  • Wall-clock time comparison. The paper claims the bound optimization is done "once offline," but does not report computation time for any method. For the hierarchical setting, optimizing the bound or computing G-optimal design may be substantially more expensive than running a simple baseline like uniform allocation. In resource-constrained settings, the offline cost matters.

6. Limitations and Trade-offs

The Cost of Difficulty Estimation Is Not Accounted for in the Headline Efficiency Numbers

The assumption or constraint. The entire compute-optimal framework depends on knowing which difficulty bin a prompt falls into before allocating the inference budget. The paper's method for estimating difficulty β€” generating 2048 samples per question and averaging the PRM's final-answer scores β€” is extraordinarily expensive. The authors acknowledge this explicitly in Section 3.2:

"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity"

This is not a minor accounting detail. At 2048 samples per question, the difficulty estimation step alone consumes more compute than the largest test-time budgets studied (256–512 generations). The paper's reported 4Γ—4\times efficiency gains over best-of-N are computed after difficulty is known, with the estimation cost excluded from the budget.

The consequence. In any realistic deployment, the total cost would be difficulty estimation + strategy execution. Since the estimation cost (20482048 generations) dominates the execution cost (1616–256256 generations) in the regimes where the 4Γ—4\times gains are claimed, the effective efficiency in practice would be far lower β€” potentially negative, meaning the adaptive approach is more expensive overall than simply running best-of-N with the full budget. The paper does not provide an amortized analysis showing how many questions must share a difficulty estimate before the approach breaks even.

What evidence exists in the paper. The paper quantifies the estimation cost (Section 3.2: 2048 samples, PRM scoring) but never includes it in any budget calculation or efficiency comparison. Figure 4 and Figure 8 (showing the 4Γ—4\times gains) use oracle or predicted difficulty bins computed offline, with the estimation cost entirely excluded from the x-axes. The predicted difficulty variant requires the same 2048 samples per question as the oracle variant β€” only the correctness check is replaced with PRM scoring, not the sample generation. The paper does not measure end-to-end cost including estimation, nor does it report how the 4Γ—4\times figure changes when estimation is amortized over varying numbers of questions.

Mitigation status. The paper explicitly flags this as "a key avenue for future work" (Section 3.2) and suggests "pretraining or finetuning models to directly predict difficulty of a question" (Section 8). No such model is developed or evaluated. The paper also mentions adaptive difficulty estimation as a possibility β€” start with a few samples, assess difficulty, then allocate β€” but does not implement this. The limitation is therefore acknowledged but entirely unresolved; the 4Γ—4\times figure should be understood as an upper bound on achievable efficiency under the unrealistic assumption of free difficulty estimation.

Hard Problems Remain Fundamentally Unsolved β€” Test-Time Compute Cannot Create Capability

The assumption or constraint. The paper's framework assumes that the base model already produces correct solutions at some non-trivial rate for the problems being addressed. This is implicit in the difficulty estimation procedure: difficulty is defined by pass@1 rate over 2048 samples from the base model. Problems where pass@1 is essentially zero fall into difficulty bin 5, and the paper shows that for these problems, no amount of test-time compute helps.

The consequence. Across all methods β€” search, revisions, and their compute-optimal combinations β€” the hardest questions (bin 5) show near-zero improvement regardless of compute budget. In Figure 3 (right), bin 5 accuracy hovers at 1–3% for all methods and all budgets from 4 to 256 generations. In Figure 7 (right), bin 5 shows roughly 2–3% accuracy irrespective of the sequential-to-parallel ratio. In the FLOPs-matched comparison (Figure 9), the bin 5 scaling line is essentially flat near 0–5% for both revisions and PRM search, at all three values of RR.

This is not a gradual degradation β€” it is a hard ceiling. The paper is candid about this in Section 7's takeaway box, but the implication is severe: test-time compute is only useful for problems that the base model already has some ability to solve. It can amplify existing capability but cannot create it from nothing. For genuinely novel or out-of-distribution reasoning problems that exceed the base model's training distribution, neither search nor revisions nor their combination offers any path forward. Pretraining remains the only viable option.

What evidence exists in the paper. The per-difficulty-bin breakdowns in Figures 3 (right), 7 (right), and the FLOPs-matched results in Figure 9 all show the same pattern: bin 5 performance is flat and near zero across methods and budgets. Figure 9 is particularly stark β€” the bin 5 line for PRM search actually shows negative relative performance compared to the 14Γ—14\times larger model at all values of RR (βˆ’3.6%-3.6\%, βˆ’35.3%-35.3\%, βˆ’52.9%-52.9\% for Rβ‰ͺ1R \ll 1, Rβ‰ˆ1R \approx 1, R≫1R \gg 1), meaning test-time compute is not merely unhelpful but actively worse than simply training a larger model and using greedy decoding. The paper quantifies the boundary: problems must fall roughly within the base model's capability range (bins 1–4, corresponding to non-trivial pass@1) for test-time compute to provide benefits.

Mitigation status. The paper does not attempt to solve this limitation β€” it is presented as a finding about the fundamental nature of test-time vs. pretraining compute. The authors are explicit that "test-time compute amplifies existing capability but does not create it" (implicit in Section 7's discussion of the bin 5 results). No mitigation is proposed because none is possible within the framework: if the base model never β€” or almost never β€” produces a correct solution, no amount of search or revision can find or refine one. The only path forward for hard problems is better pretraining, as the FLOPs-matched comparison makes clear.

A Single Benchmark and Single Model Family β€” Generalization Is Unverified

The assumption or constraint. All experiments in the paper use the MATH benchmark (Hendrycks et al., 2021) with PaLM 2-S* (Codey) as the base model. The test set consists of 500 competition-level math problems. The paper states in Section 4 that the authors "believe this model is representative of the capabilities of many contemporary LLMs," but this claim is not tested. The choice of MATH is deliberate: "test-time compute is expected to help most when the model already possesses the necessary knowledge and the challenge is drawing complex inferences β€” mathematical reasoning fits this profile because it requires multi-step logical deduction rather than novel factual recall" (Section 4). This is a reasonable domain choice, but it is also a narrow one.

The consequence. Several aspects of the findings could be model-specific or domain-specific, and the paper provides no evidence either way:

  • PRM quality and over-optimization behavior depend on PaLM 2-S*'s output distribution. A model with different calibration properties or different error patterns might exhibit different difficulty-dependent scaling curves. The over-optimization phenomenon documented in Figure 3 (beam search degrading on easy problems at high budgets) depends on the specific interaction between the PRM's scoring errors and the base model's output distribution. A better-calibrated model might show less over-optimization; a worse-calibrated one might show more.

  • The revision model's ability to learn from incorrect in-context examples depends on the base model's in-context learning capabilities, which vary substantially across model families. The paper's revision training procedure (Section 6.1) relies on the model being able to condition on previous incorrect answers and produce improved ones β€” an ability that might not transfer to models with weaker in-context learning.

  • The MATH benchmark consists exclusively of problems with well-defined ground-truth answers that can be checked with exact string matching. Many important real-world applications β€” open-ended generation, dialogue, creative writing, complex multi-step planning β€” lack such clean correctness signals. The difficulty estimation procedure (2048 samples, checking against ground truth or PRM scores) assumes access to a reliable correctness signal, which may not exist in other domains.

The test set of 500 questions, split into five difficulty quintiles of ~100 each, further split by two-fold cross-validation (Section 3.2), means the compute-optimal policy is selected based on approximately 50 questions per fold per bin. The paper does not report confidence intervals on the compute-optimal scaling curves (Figures 4 and 8), making it difficult to assess whether the observed differences between strategies are statistically reliable at this sample size.

What evidence exists in the paper. No evidence. The paper does not include experiments on any benchmark other than MATH, nor with any model other than PaLM 2-S*. The authors acknowledge this implicitly by stating their belief about representativeness (Section 4) but do not provide supporting evidence. The paper also does not discuss how the findings might differ for code generation, logical reasoning, scientific QA, or other domains where test-time compute is commonly applied.

Mitigation status. The paper does not address this limitation beyond stating the belief about representativeness. No multi-benchmark or multi-model experiments are conducted. The authors do not explicitly flag this as a limitation or suggest future work on broader empirical validation, which is a notable omission given the strength of the claims (e.g., "test-time compute can substitute for a 14Γ—14\times larger model"). A reader would reasonably want to know whether this holds for, say, LLaMA models on coding benchmarks before adopting the approach in a different setting.

The 14Γ—14\times Larger Model Baseline Is Weakened by Non-Optimal Training and Greedy Decoding

The assumption or constraint. The FLOPs-matched comparison in Section 7 scales model parameters by approximately 14Γ—14\times while holding training data fixed, following the LLaMA paradigm (Touvron et al., 2023). The paper acknowledges that this departs from compute-optimal pretraining (Hoffmann et al., 2022), where both data and parameters should be scaled equally:

"We choose this setting as it is representative of a canonical approach to scaling pretraining compute and leave the analysis of compute-optimal scaling of pretraining compute where the data and parameters are both scaled equally to future work."

Additionally, the 14Γ—14\times larger model uses only greedy decoding with no test-time compute augmentation β€” no majority voting, no best-of-N, no search of any kind.

The consequence. Both choices make the pretraining baseline weaker than it could be:

  • Non-compute-optimal pretraining: A Chinchilla-optimal model trained with 14Γ—14\times more total FLOPs (scaling both parameters and data according to the Chinchilla scaling laws) would likely outperform a parameter-only-scaled model. The reported advantages of test-time compute over pretraining β€” e.g., +27.8%+27.8\% on easy questions at Rβ‰ͺ1R \ll 1 for revisions (Figure 1, top-right bar chart) β€” may shrink or reverse against a properly compute-optimal larger model. The magnitude of this effect is unknown because the paper does not compare against a Chinchilla-optimal baseline.

  • No test-time compute for the larger model: Giving the larger model even a modest test-time compute budget β€” say, best-of-8 with majority voting or best-of-N weighted β€” would create a substantially stronger baseline. The paper's central comparison is "test-time compute on a small model vs. no test-time compute on a large model," which conflates two distinct choices (model size and inference strategy). A fairer comparison would give both models the same test-time compute budget, or at least explore how the tradeoff changes when the larger model also uses test-time compute.

What evidence exists in the paper. Section 7 describes the FLOP accounting and acknowledges the non-Chinchilla-optimal training in a footnote/caveat. The bar charts in Figure 1 and the line plots in Figure 9 all compare against the 14Γ—14\times larger model with greedy decoding. The paper does not include any experiment where the larger model receives test-time compute augmentation. There is no ablation testing how the FLOPs-matched comparison changes if the larger model uses best-of-N or majority voting.

Mitigation status. The paper explicitly leaves compute-optimal pretraining comparison to future work (Section 7). The choice of greedy decoding for the larger model is not acknowledged as a limitation, nor is the absence of a test-time-compute-augmented large model baseline discussed. The paper's headline claims about test-time compute substituting for pretraining should be interpreted as specific to the comparison setting (parameter-only scaling, greedy decoding for the large model) and may not generalize to settings where the larger model is optimally trained and/or given its own test-time compute budget.

Revision Models Exhibit a Systemic Correct-to-Incorrect Reversion Problem

The assumption or constraint. The revision model is fine-tuned on trajectories where all in-context answers are incorrect followed by a correct target (Section 6.1). The training data construction explicitly avoids including correct answers in the conditioning context, because the goal is to teach the model to revise incorrect answers into correct ones. The authors note that standard validation loss is not a reliable signal for early stopping because the validation trajectories become off-policy after fine-tuning β€” they were generated by the base model, not the fine-tuned revision model.

The consequence. At test time, the revision model generates a chain of revisions. If it produces a correct answer at some step, that correct answer enters the context for the next revision. Since the model was never trained on sequences containing correct answers, it may "revise" the correct answer into an incorrect one β€” a correct-to-incorrect reversion. The paper reports that approximately 38% of correct answers get converted back to incorrect ones using a naive approach (Section 6.1). This is not a minor edge case β€” it means that in any revision chain longer than a few steps, the model is actively fighting itself, producing correct answers and then destroying them.

The paper mitigates this with majority voting or verifier-based selection across the entire chain of revisions, picking the best answer from any point rather than always taking the last revision. But these are patches, not solutions. They add computational and conceptual overhead to what is supposed to be a simple sequential revision process. Moreover, the selection mechanisms are imperfect β€” a verifier or majority vote can fail to identify the best answer in the chain, especially if the correct answer appears only briefly before being overwritten.

What evidence exists in the paper. The 38% figure is reported in Section 6.1 based on empirical measurement. Figure 6 (left) shows that the revision model's pass@1 gradually improves over the first ~15–20 steps but then plateaus in the 23–25% range, consistent with the reversion problem capping the benefits of additional revision steps. The ReSTEM^{EM} experiment (Appendix K, Figure 16) provides additional evidence: attempting to further optimize the revision model with RL-style training caused performance to degrade substantially with sequential revisions (fully sequential dropping to ~33.5% at 256 generations compared to ~38.5% at the optimal ratio). The authors hypothesize that "the on-policy data collection in ReSTEM^{EM} exacerbates spurious correlations in revision data, causing the model to fail to learn the revision task properly" (Appendix K). This suggests the revision approach is sensitive to training methodology in ways that are not fully understood.

Mitigation status. The paper mitigates the reversion problem with chain-level selection (majority voting or verifier-based), but acknowledges this is a patch rather than a solution. A more principled approach β€” such as training the model to recognize when no revision is needed (including a "stop revising" token or training on mixed correct/incorrect trajectories) β€” is not explored. The degradation under ReSTEM^{EM} is presented as a notable negative result but is not further investigated. The paper does not discuss whether the revision model's training data construction could be improved to reduce the reversion rate, nor does it propose specific future work on this issue. A practitioner deploying revision models would need to account for the reversion problem and implement chain-level selection, understanding that it adds complexity and may not always recover the correct answer.

The Paper Never Combines PRM Search with Revision Models, Leaving a Natural Synergy Unexplored

The assumption or constraint. The paper studies two complementary mechanisms for test-time compute β€” PRM-guided search (Section 5) and iterative revisions (Section 6) β€” but treats them as independent and never combines them. Section 8 explicitly acknowledges:

"we did not experiment with PRM tree-search techniques in combination with revisions"

This is a significant omission because the two mechanisms address different aspects of the BAI problem. Revisions modify the proposal distribution β€” they make the model generate better candidate solutions by conditioning on previous attempts. PRM search modifies candidate selection β€” it uses a learned verifier to identify the best solution among generated candidates. These are complementary: better proposals make search easier, and better search extracts more value from a given set of proposals. Combining them could yield gains beyond either method alone.

The consequence. The current results represent a lower bound on what a fully integrated system could achieve. Specifically:

  • Using the revision model as the proposal distribution within beam search could improve search efficiency: at each step of the search tree, the model would condition on previous rejected branches as context, potentially producing higher-quality candidate steps that are more likely to lead to correct solutions.

  • Using the PRM to guide which revisions to pursue β€” rather than blindly generating a long revision chain β€” could mitigate the correct-to-incorrect reversion problem and focus computational effort on promising directions. For example, the PRM could be used to decide when a revision is on track versus when to restart from scratch.

  • On medium-difficulty problems where both mechanisms show complementary strengths (beam search helps on bin 3–4 problems in Figure 3; revisions help on bin 2–3 problems in Figure 7), the combination might break through the performance ceiling that each method individually hits at high budgets.

The paper's compute-optimal policy currently selects between search strategies or between revision strategies independently. A unified policy that allocates budget across both mechanisms simultaneously could be more efficient than either alone.

What evidence exists in the paper. No evidence. The paper does not include any experiments combining PRM search with revisions. The theoretical framework (Section 2) presents the proposal distribution and verifier as complementary axes, which naturally suggests combination, but this is not pursued empirically.

Mitigation status. The paper explicitly acknowledges this gap in Section 8 and frames combination as future work. No preliminary results or feasibility analysis is provided. A reader interested in deploying a state-of-the-art test-time compute system would reasonably want to know whether combination helps, and if so, how to allocate budget between search and revisions β€” the paper provides no guidance on either question. The omission is particularly notable given that the paper's own framework identifies these as complementary mechanisms; the natural next step is within reach but is left entirely to future work.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper fundamentally reframes the question of whether adaptive elimination strategies are necessary for fixed-budget best-arm identification. The dominant paradigm for over a decade β€” represented by Sequential Halving (Karnin et al., 2013) and extended to structured settings through optimal design (Azizi et al., 2021; Katz-Samuels et al., 2020) β€” has treated adaptivity as an obvious virtue: of course you stop pulling arms that appear suboptimal. PI-BAI provides a concrete counter-narrative grounded in both impossibility results (Wang et al., 2024; Degenne, 2023) and empirical evidence: when you have a meaningful prior, you can pre-compute a fixed allocation that matches or exceeds adaptive methods, without any data-dependent decision-making.

This is not a paradigm shift β€” the paper does not claim fixed allocations dominate adaptive ones, which would contradict established impossibility results β€” but it is a substantive reframing of the design space. The question evolves from "how can we be more adaptive?" to "given our prior knowledge, what is the right mixture of fixed design and adaptivity?" PI-BAI establishes the non-adaptive endpoint of this spectrum, providing both theoretical guarantees and empirical performance that future adaptive Bayesian BAI algorithms must measure themselves against. The paper's result that PI-BAI's bound is tighter than BayesElim's by roughly log₃/β‚‚β‚‚(K) (a factor of ~6 for K = 10) is not an incremental tightening β€” it comes from structurally eliminating the round-based overhead, demonstrating that eliminations impose a real statistical cost, not just an analytical artifact.

The second landscape shift is the proof technique itself. The key reformulation in Section 5 β€” swapping the order of conditioning to bound the Bayesian posterior probability of error rather than averaging frequentist bounds β€” provides a template that extends naturally across model classes. This resolves the tension in the prior literature where Bayesian BAI methods were either restricted to homogeneous prior variances (Atsidakou et al., 2022) or lacked theoretical guarantees entirely. The proof technique demonstrates that algorithm simplicity enables analytical unity: because PI-BAI does the same thing regardless of model structure (pull arms, compute posteriors, pick highest), the analysis is modular, with model-specific computation isolated to a self-contained Gaussian moment calculation.

The third shift concerns structured bandits. Before this work, Bayesian BAI with prior-dependent guarantees was confined to multi-armed bandits with restrictive assumptions. The paper's bounds for linear bandits (Theorem 4.2) and hierarchical bandits (Theorem 4.4) are the first of their kind, and they reveal dimension-dependent complexity: the statistical difficulty scales with d (ambient dimension) rather than K (number of arms) in linear models, and the hierarchical bound decomposes prior uncertainty into effect-level and arm-level components, making explicit how shared structure reduces effective sample complexity. This redirects research attention from developing new adaptive algorithms for each model class toward the more fundamental problem of optimal allocation given a prior and a model structure, which is a design problem rather than an algorithm-design problem.

The paper also implicitly resolves a tension that practitioners face: Bayesian methods promise efficiency through prior knowledge, but existing Bayesian BAI algorithms (BayesElim, TTTS) were either hard to analyze or underperformed empirically relative to their theoretical promise. PI-BAI demonstrates that the efficiency gains from informative priors can be realized with startlingly simple non-adaptive strategies, suggesting that the complexity of adaptive Bayesian algorithms was masking, rather than enabling, the benefits of prior information. This is reinforced by the empirical finding (Appendix D.5, Figure 8) that adding elimination on top of PI-BAI degrades performance β€” the elimination structure is actively harmful, not neutral.

Follow-Up Research This Work Enables

Cheap difficulty estimation through direct prediction. The paper acknowledges that the cost of estimating question difficulty by generating 2048 samples per question (Section 3.2) is prohibitive β€” it can exceed the largest test-time compute budgets studied. This paper enables follow-up work by establishing that difficulty-conditioned allocation provides large empirical gains (Figures 4 and 8, showing ~4Γ— efficiency improvement), justifying investment in cheaper estimation. A concrete follow-up would train a lightweight classifier β€” possibly a distilled version of the PRM β€” that takes only the question text as input and predicts the difficulty bin. The evaluation would compare this classifier's bin assignments against the PRM-based estimates on the MATH 500 test set, and measure downstream PoE when using classifier-predicted bins versus oracle bins. The key metric is the degradation in compute-optimal performance relative to the oracle, as a function of classifier accuracy. A second approach would implement adaptive difficulty estimation: generate, say, 8 initial samples, use the PRM's score distribution on those samples as a quick difficulty signal, and allocate the remaining budget accordingly β€” this amortizes estimation into the solving process itself and would be evaluated by total (estimation + solving) cost versus accuracy.

Combining PRM tree-search with revision models. The paper studies PRM search and iterative revisions as independent mechanisms but explicitly notes they were never combined (Section 8). This is the most natural and highest-impact follow-up. The proposed combination would use the revision model as the proposal distribution within beam search: at each step of the search tree, the model conditions on previously rejected branches as context, potentially producing higher-quality candidate steps. A concrete experiment would compare three architectures on the MATH benchmark with matched generation budgets: (a) beam search with the base LLM as proposal, (b) sequential revisions with best-of-N weighted selection, and (c) beam search with the revision model as proposal β€” varying beam width and revision depth to find the optimal allocation. The key hypothesis is that the combination breaks through the performance ceiling each method individually hits at high budgets (beam search plateaus due to verifier over-optimization, revisions plateau due to correct-to-incorrect reversion). A negative result β€” combination provides no benefit over the better of the two individually β€” would be equally informative, suggesting that the two mechanisms address the same underlying bottleneck rather than complementary ones. A positive result would be quantified as the additional % accuracy at matched budget or the budget reduction at matched accuracy.

Systematic study of verifier over-optimization as a function of PRM training. The paper identifies verifier over-optimization as the primary bottleneck for test-time search scaling (Section 5.3, Figure 3), but does not explore how PRM training choices affect robustness to optimization pressure. A follow-up study would train PRMs under different conditions β€” varying the number of Monte Carlo rollouts per step (1, 4, 16, 64), the temperature of rollout sampling, whether rollouts use the base model or a diverse ensemble β€” and measure how each PRM variant's ranking accuracy degrades as beam search budget increases. The key measurement is the "over-optimization threshold": the budget at which beam search accuracy stops improving and starts declining, plotted against PRM quality metrics. This would establish whether the over-optimization is driven by PRM miscalibration (the model is systematically biased toward certain solution patterns) or by irreducible noise (the base model's outputs contain genuinely ambiguous cases). A follow-up would then test interventions: adversarial training of the PRM on beam-search-generated solutions, ensemble verification aggregating multiple independently trained PRMs, and KL-constrained search that penalizes solutions deviating from the base model's typical output distribution. The evaluation metric is the maximum achievable accuracy before over-optimization sets in, compared across PRM training variants.

Replication on code generation benchmarks with unit-test-based verifiers. This paper's entire empirical evaluation uses the MATH benchmark with exact string matching for correctness checking. This enables the PRM training pipeline (Monte Carlo rollout supervision) and difficulty estimation (2048 samples per question). Code generation benchmarks (HumanEval, MBPP) offer a natural replication target with cleaner correctness signals: unit tests provide deterministic pass/fail labels without requiring learned verifiers, and difficulty can be estimated from pass@1 rates without any learned model. A replication study would apply the compute-optimal framework to, say, CodeLlama-7B on HumanEval, using best-of-N weighted selection with a unit-test-based outcome reward model (no PRM needed), and sequential revision via fine-tuned self-debugging. The key question is whether the difficulty-dependent patterns replicate: does beam search hurt easy problems due to over-optimization? Do sequential revisions help more on easy problems and parallel sampling on hard ones? Are the 4Γ— efficiency gains achievable? This extension is valuable because code generation has larger real-world deployment volume than competition math, and unit tests eliminate the verifier quality bottleneck that dominates the paper's results, potentially allowing compute-optimal scaling to push further before over-optimization appears.

Hierarchical model adaptation for prior learning in multi-task BAI. The paper's online prior learning experiment (Figure 2b) uses a hierarchical model with K = L = 10 and identity mixing weights (bi = ei) to learn the MAB prior means from data. This is a special case; the hierarchical model's real power lies in settings where arms share effects through non-trivial mixing weights β€” hyperparameter configurations sharing hyperparameter values, drug candidates sharing molecular components. A follow-up would test this in a realistic multi-task BAI setting: for example, hyperparameter tuning where each configuration is a combination of L hyperparameters (learning rate, batch size, optimizer, etc.), and the goal is to find the best configuration within budget n. The prior would be learned from a corpus of prior tuning experiments on related tasks (offline data), inducing a hierarchical prior with effects corresponding to hyperparameter values. The key measurement is the sample efficiency gain: how many fewer evaluations does PI-BAI with the hierarchical prior require to achieve the same PoE as an MAB baseline that treats each configuration independently? This directly tests whether the hierarchical structure's theoretical advantages (Figure 3 in Appendix B.2, showing 2Γ— reduction in posterior variance) translate to practical BAI performance improvements.

Stress-test: how informative must the prior be for PI-BAI to beat uniform allocation? The paper's experiments use moderately informative priors (Οƒβ‚€,α΅’ ∈ [0.1, 0.5] with Οƒ = 1). The theoretical bound suggests that as prior variance increases, the optimized allocation converges to uniform. A systematic stress-test would sweep prior informativeness: fix K = 10 arms with prior means Β΅β‚€,α΅’ ~ 𝒰([0,1]) and vary Οƒβ‚€ across a wide range (0.01, 0.05, 0.1, 0.3, 0.5, 1.0, 2.0, 5.0). At each informativeness level, compare PI-BAI(Opt), PI-BAI(G-opt), and Uniform at a fixed budget (say, n = 100). The key output is a phase diagram: below what Οƒβ‚€ does PI-BAI(Opt) provide >10% relative improvement over Uniform? Above what Οƒβ‚€ is the improvement <1%? This would give practitioners a concrete guideline for when investing in prior elicitation and optimization is worthwhile versus simply using uniform allocation. A negative result β€” PI-BAI(Opt) providing negligible benefit even at moderately informative priors (Οƒβ‚€ ~ 0.5Οƒ) β€” would constrain the practical applicability of the method and motivate investigation into why the theoretical gains don't materialize.

Practical Applications and Downstream Use Cases

Hyperparameter tuning with offline data from prior experiments. In organizations that repeatedly tune similar models β€” e.g., a machine learning team running weekly training jobs with varying hyperparameters β€” the hierarchical bandit formulation maps directly onto the tuning workflow. Hyperparameter values (learning rate, batch size, dropout) correspond to effects Β΅_β„“; each configuration i corresponds to mixing weights b_i encoding which values it uses; and prior tuning logs provide offline data for learning the prior. PI-BAI with the hierarchical prior would allocate a fixed evaluation budget across configurations before running any new experiments, concentrating budget on configurations that are a priori promising or uncertain. The paper's Figure 2a (third column) shows PI-BAI(TS) achieving PoE ~0.55 at n = 550 with K = 60 arms β€” in plain terms, roughly 45% chance of correctly identifying the best among 60 configurations using only ~9 evaluations per configuration. For comparison, uniform allocation would require ~60 evaluations per configuration to achieve similar confidence by exhaustive search. The practical saving is a 6–7Γ— reduction in experimental cost for a given confidence level, or equivalently, the ability to explore a 6–7Γ— larger configuration space at the same cost.

Clinical trial design with hierarchical drug models. In drug development, candidate compounds share molecular components (effects), and prior knowledge about component efficacy is available from preclinical studies or related trials. The hierarchical bandit model with ΞΈ_i = b_i^⊀ Β΅ + Ξ΅_i directly captures this: b_i encodes the molecular composition of drug i, Β΅ encodes the efficacy of each component, and Ξ΅_i captures drug-specific deviations. PI-BAI would allocate trial subjects across drug candidates using a prior-dependent fixed design computed from the preclinical data, rather than adaptively re-allocating based on interim results (which is logistically complex and may compromise blinding). The paper's theoretical guarantee β€” that the expected PoE decays as 1/√n with explicit constants depending on the prior and allocation β€” provides regulatory-relevant quantification of decision quality as a function of sample size. The hierarchical structure's efficiency gain (Figure 3, Appendix B.2, showing 2Γ— smaller posterior variance at the same budget) means that exploiting component-level correlations could reduce required sample sizes by approximately half for the same decision quality, which in pharmaceutical contexts translates to millions of dollars in savings and faster time-to-decision.

A/B testing with heterogeneous prior uncertainty across variants. In online experimentation, different treatment variants often enter the test with different levels of prior evidence β€” a minor UI tweak may be genuinely novel (high prior variance), while a variant similar to a previously tested design has stronger prior information (low prior variance). Current A/B testing practice typically allocates equal traffic to all variants, which is statistically inefficient when priors differ. PI-BAI with the closed-form K = 2 solution (Section 4.1) provides an explicit allocation formula: ω₁ = Ξ _{[0,1]}(1/2 - (σ²₀,β‚‚ - σ²₀,₁)σ² / (2n σ²₀,₁ σ²₀,β‚‚)). When one variant has substantially higher prior variance, it receives more traffic β€” for small n, potentially all of it. The paper's Lemma 4.5 additionally quantifies sensitivity to prior misspecification: the penalty factor d_n = e^{π’ͺ(1/n)} vanishes with sample size, meaning that moderate errors in prior specification become irrelevant as the experiment accumulates data. For a platform running thousands of concurrent A/B tests with fixed total traffic, reallocating traffic according to PI-BAI's formula would reduce the expected number of incorrect decisions (shipping inferior variants, discarding superior ones) by the factor quantified in the paper's MAB experiments β€” roughly 15–20% relative reduction in PoE at typical test sample sizes (n β‰ˆ 50–200) compared to uniform allocation (Figure 2a, first column).

When to Prefer This Method

The paper positions PI-BAI against adaptive elimination methods (Sequential Halving, Successive Rejects, BayesElim) and adaptive Bayesian methods (TTTS), but does not provide a crisp decision rule. The empirical results and theoretical analysis together support the following guidance:

  • Prefer PI-BAI(Uni) or PI-BAI(G-opt) when the prior is informative (Οƒβ‚€,α΅’ not much larger than Οƒ), the budget n is small to moderate relative to K (so adaptivity has limited time to help), and computational simplicity or theoretical guarantees matter. The G-optimal design variant provides the best empirical performance in MAB and linear settings without requiring online adaptation or numerical bound optimization.

  • Prefer PI-BAI(TS) when the bandit structure is rich (hierarchical models with shared effects) and a small warm-up budget (n_w β‰ˆ K) is affordable. The mixed-effect Thompson sampling warm-up learns an instance-specific allocation that substantially outperforms any fixed design in the hierarchical experiments (Figure 2a, third column), though it sacrifices the theoretical guarantees of the non-adaptive variants.

  • Fall back to adaptive elimination methods when the prior is weakly informative (large Οƒβ‚€,α΅’ relative to observation noise), the budget n is large relative to K (so elimination has many rounds to operate), or prior specification is highly uncertain and online learning is infeasible. The paper does not test this boundary explicitly, but the theory predicts PI-BAI's allocations converge to uniform as prior variance increases, at which point adaptivity may provide genuine benefits by reacting to observed gaps between arms.

  • Avoid PI-BAI(Opt) in high-dimensional structured settings (hierarchical bandits with many effects). The numerical optimization of the theoretical bound failed to find good allocations in the hierarchical experiments (Figure 2a, third column, where PI-BAI(Opt) underperforms PI-BAI(TS) and PI-BAI(G-opt) substantially), likely due to the bound's complexity and non-convexity. The optimization-based approach is most reliable in MAB and linear settings where the bound is simpler.