Optimal Bayesian Stopping for Efficient Inference of Consistent LLM Answers

AI summary

Overview

• Research area: Efficient inference for large language models (specifically Adaptive Self-Consistency), Bayesian sequential hypothesis testing, and mode identification under informative priors. Primary category: stat.ML. Keywords listed by the authors: Machine Learning, ICML.
• Technical level: Advanced. The paper is built on Bayesian posterior computation, sequential hypothesis testing, mode-identification theory, and asymptotic (δ → 0) analysis of stopping times; the notation is dense even though the motivating application is practical.
• Scope (one sentence): The paper derives an optimal Bayesian stopping rule for deciding when enough LLM samples have been drawn to confidently output the most consistent answer, and shows that a cheap "L-aggregated" posterior with L = 3 is enough to match the asymptotic sampling efficiency of the exact (intractable) posterior.

Authors and affiliations: Jingkai Huang (Stern School of Business, New York University), Will Ma (Graduate School of Business, Columbia University), Zhengyuan Zhou (Stern School of Business, New York University). Correspondence: zz26@stern.nyu.edu. arXiv:2602.05395v2 [stat.ML], 31 May 2026; published 2026-02-05.

What This Paper Is About

Self-Consistency improves LLM accuracy on math and reasoning tasks by sampling many answers and submitting the most frequently reached one, but it spends a fixed sampling budget on every query, wasting computation on easy questions. Adaptive Self-Consistency tries to stop early once a leading answer emerges, yet existing methods ignore available prior knowledge about whether a query's answer distribution is "peaked" (easy) or "flat" (hard). This paper formalizes stopping as a Bayesian sequential hypothesis test with an informative prior, and solves the resulting computational bottleneck with an approximation that tracks only the top L−1 answer counts.

Key Contributions

1. Bayesian framework for Adaptive Self-Consistency with an informative prior. The stopping decision is formulated as a sequential hypothesis testing problem (testing which observed answer count corresponds to the true most-likely answer), with the posterior computed combinatorially over permutations of answers (Section 2).

2. An efficient L-aggregated posterior approximation. The K candidate answers are aggregated into L ≤ K groups, reducing posterior computation from a naive O(K!) to O(K^L). The authors prove the approximated posterior remains unbiased and achieves asymptotic optimality for any L ≥ 3, meaning the rate of the asymptotic stopping time matches that of the exact posterior as the confidence level 1−δ tends to 1 (Section 3).

3. Theoretical comparison against prior-free baselines and an uncertain-prior extension. For known priors, even the coarsest approximation L = 2 gives a strictly smaller asymptotic stopping time than prior-free ASC baselines, while L = 3 matches the exact posterior. For uncertain priors — where the true prior belongs to a set of M candidate priors — L = 2 can degrade severely, but L ≥ 3 still outperforms prior-free baselines (Section 4).

4. The "three is all you need" takeaway plus empirical validation. L = 3 tracks the most-frequent answer count, the second-most-frequent answer count, and all remaining answers combined, which is the minimal state that captures both the leader's occurrence percentage and its margin over the runner-up. Experiments on real-world datasets are reported to confirm smaller stopping times at high accuracy (Section 5 and appendices).

Main Findings

• Exact posterior is intractable but well defined. With K distinct possible answers and prior π = (p_1, …, p_K) indexed so p_1 > p_2 ≥ … ≥ p_K > 0, the exact posterior P(H_1 | C_n) requires summing over all injective mappings of observed answers to latent labels, giving naive complexity O(K!) once M(n) = K answers have been observed.

• The L-aggregated posterior is unbiased. It satisfies P(H_1 | C_n^L) = E[P(H_1 | C_n) | C_n^L], so it preserves the correct Bayesian belief on average while losing information through compression.

• Asymptotic stopping time for known priors (Theorem 3.2). For L = 2, lim_{δ→0} E[n^{,2}] / log(1/δ) = 1 / D_KL(p_1‖p_2). For L = 3, 4, …, K, lim_{δ→0} E[n^{,L}] / log(1/δ) = 1 / ((p_1 − p_2) · log(p_1/p_2)). Consequently, L = 3 matches the exact posterior (L = K) asymptotically, losing no asymptotic statistical efficiency.

• Strict dominance over prior-free baselines (Remark 3.4). The prior-free rate from Shah et al. (2020) and Jain et al. (2022) is 1 / (p_1 · log(2p_1/(p_1+p_2)) + p_2 · log(2p_2/(p_1+p_2))). For 0 < p_2 < p_1 < 1 with p_1 + p_2 < 1, the ordering is E[n^{,f}] > E[n^{,2}] > E[n^{,3}] = … = E[n^{,K}] = E[n^*], so even L = 2 strictly beats the prior-free rule.

• Computational complexity is exponential in L, not K. The total complexity of computing P(H_1 | C_n^L) is O(K! · n̄_{L(n)}^2 / (K−L)!), dominated by O(K^L · n̄_{L(n)}^2), where n̄_{L(n)} is the total count mass assigned to aggregated "other" answers. Iterating over injective mappings costs O(K!/(K−L+1)!), and evaluating the aggregation constant for each mapping costs O((K−L+1) · n̄_{L(n)}^2) via dynamic programming.

• Uncertain priors change the picture (Theorem 4.1). When the true prior is one of M candidate priors Π^M with hyper-prior weights λ_m, the L = 2 rate becomes 1 / J_{2,m^†}(p_1), with J_{2,m}(p_1) = D_KL(p_1‖p_{2,m}) + (1−p_1)·1{ρ < ρ^†, ρ < 1}·D_KL(ρ‖ρ^†), where ρ = p_1/(1−p_1) and ρ^† = p_{1,m}/(1−p_{2,m}). For L = 3, …, K the rate is 1 / J^L_{2,m^†}(p_L).

• A concrete failure case for L = 2 with uncertain priors (Remark 4.2). With M = 2, π^1 = (0.51, 0.10, …), π^2 = (0.49, 0.48, …), and true prior m^* = 1, the paper reports 1 / J_{2,m^†}(p_1) = 555.22, which is far larger than the prior-free asymptotic rate of 6.64. Thus L = 2 can perform worse than the prior-free baseline because it only tracks the leading answer and cannot distinguish whether the leader is the true winner under π^1 or π^2.

• L ≥ 3 still beats prior-free under uncertain priors, but is no longer exactly optimal (Remark 4.3). The ordering becomes E[n^{,f}] > E[n^{,3}{Π^M}] ≥ … ≥ E[n^{*,K}{Π^M}] ≥ E[n^]. When the mixture's selected prior equals the true prior (m^† = m^), the rate collapses to the known-prior quantity (p_{1,m^} − p_{2,m^}) · log(p_{1,m^}/p_{2,m^}).

• Empirical claim about L = 3. The paper states that L = 3 replicates the optimal sample efficiency of L = K starting at confidence level 1−δ = 0.8, without replicating the posterior computation time of L = K (Table 1). Specific dataset names, dataset sizes, accuracy figures, and wall-clock timings are not reported in the available content, which is truncated before the experiments section.

Methodology in Plain English

The researchers model an LLM answering a question as repeatedly drawing an answer from an unknown probability distribution over K possible answers. Because the model does not know the identities of the answers in advance, the only information it can use is how many times each observed answer has appeared — summarized as a "count-of-counts" object that records which frequencies occurred and how many distinct answers had each frequency. For example, observing the sequence B, A, B, D gives n = 4 samples, M(n) = 3 distinct answers, counts (2, 1, 1), and the compressed representation {(2,1), (1,2)}, meaning one answer appeared twice and two answers appeared once.

The stopping rule is simple in principle: keep sampling until the posterior probability that the currently leading answer is the true most-likely answer reaches a target confidence level 1−δ, then stop. Computing that posterior exactly requires summing over every way the observed distinct answers could map onto the K latent answer labels, which costs O(K!) and is unusable in real time when K is large.

The fix is aggregation. Instead of tracking every answer's count, the method keeps explicit track only of the L−1 most frequent answers and lumps everything else into a single "other" bucket. The authors introduce a weighting correction to avoid double-counting when several answers share the same frequency at the head/tail boundary, and they compute the "other" bucket's likelihood contribution via dynamic programming. The result is a posterior whose cost scales as O(K^L) rather than O(K!). Choosing L = 3 means tracking only the most-frequent count, the second-most-frequent count, and everything else combined — enough to capture both how dominant the leader is and how far ahead it is. The authors then analyze the expected stopping time as δ shrinks to zero, and extend the whole construction to the case where the true prior is not known exactly but is known to be one of M candidate priors drawn from a hyper-prior.

Why This Matters

Impact on research. The paper reframes adaptive self-consistency as a new problem in sequential hypothesis testing: optimal Bayesian stopping for mode identification with a prior that does not reveal the mode. Prior mode-identification work (Shah et al., 2020; Jain et al., 2022) studied the prior-free case, and Jain et al. (2022) supplied the theoretical backbone for the ASC literature. This paper shows that LLM answer distributions provide a legitimate new form of Bayesian prior for mode identification, and that a very coarse posterior approximation (L = 3) suffices for asymptotic optimality — a result the authors contrast with MCMC and normal approximations, which they describe as too slow for real-time use and ill-suited to discrete data structures.

Real-world applications (derived from the paper's framing).

• Serving LLMs for math and multi-step reasoning queries, where sampling many Chain-of-Thought paths is standard practice but per-query cost is the bottleneck.
• Reducing latency and operational cost in production inference pipelines by stopping easy (peaked-distribution) queries early rather than spending a fixed budget on every query.
• Test-time scaling workflows where a fixed sampling budget is currently pre-determined for every question.
• Reusing historical evaluation data from the same LLM on related questions as a learned prior, so that per-query decisions inherit information from prior tasks.

Industry relevance. The paper targets the operational cost of self-consistency: fixed-budget sampling is described as leading to significant latency and operational costs due to redundancy. A stopping rule that is cheap to compute (exponential in a small constant L rather than in K) and that reduces LLM calls is directly relevant to anyone paying per-token or per-request for test-time scaling. The uncertain-prior analysis is also practically framed around the common situation where a new question's answer distribution is unknown but related tasks have been evaluated before.

Future Directions

• Learn the prior empirically and validate it. The paper's organization states that Section 5 justifies how the prior distribution can be experimentally learned from a small subset of historical data; the truncated content does not report those results, leaving the accuracy of learned priors as an open empirical question.
• Handle the uncertain-prior gap for L = 3. Theorem 4.1 shows L = 3 is not asymptotically optimal under a mixture of candidate priors, even though it still beats the prior-free baseline. Characterizing exactly when L = 3 closes the gap to L = K (for instance, when m^† = m^*) and whether a larger but still constant L can restore optimality is an explicit open question.
• Choose L adaptively. The paper frames L as a tradeoff knob between statistical complexity (sample count) and computational complexity (posterior evaluation time); how to pick L per query at runtime is not resolved.
• Test the theory outside the LLM setting. Because the framework is stated generally as optimal Bayesian stopping for mode identification with an informative prior, applying it to other sequential decision problems with unknown answer-support structure is a natural extension.

Target Audience

Researchers and practitioners working on LLM inference efficiency, test-time scaling, and adaptive decoding strategies; applied probabilists and statisticians interested in sequential hypothesis testing, Bayesian stopping rules, and mode identification; and engineers building cost-sensitive LLM serving systems who need stopping rules that are cheap enough to run per query. Readers need comfort with Bayesian posteriors, KL divergence, and asymptotic analysis; the paper is not beginner-friendly, though the L = 3 rule itself is easy to implement as "track the top two counts plus a combined remainder."

Authors’ abstract

A simple strategy for improving LLM accuracy, especially in math and reasoning problems, is to sample multiple responses and submit the answer most consistently reached. In this paper we leverage Bayesian prior information to save on sampling costs, stopping once sufficient consistency is reached. Although the exact posterior is computationally intractable, we further introduce an efficient "L-aggregated" stopping policy that tracks only the L-1 most frequent answer counts. Theoretically, we prove that L=3 is all you need: this coarse approximation is sufficient to achieve asymptotic optimality, and strictly dominates prior-free baselines, while having a fast posterior computation. Empirically, this identifies the most consistent (i.e., mode) LLM answer using fewer samples, and can achieve similar answer accuracy while cutting the number of LLM calls (i.e., saving on LLM inference costs) by up to 50%.