Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

BEACON: Bayesian Optimal Stopping for Efficient LLM Sampling

Guangya Wan, Zixin Stephen Xu, Sasa Zorc, Manel Baucells, Mengxuan Hu, Hao Wang, Sheng Li

TL;DR

This work introduces BEACON (Bayesian Efficient Adaptive Criterion for Optimal N-stopping), a principled adaptive sampling framework grounded in Sequential Search with Bayesian Learning, and establishes both theoretical optimality guarantees and practical tractability.

Abstract

Sampling multiple responses is a common way to improve LLM output quality, but it comes at the cost of additional computation. The key challenge is deciding when to stop generating new samples to balance accuracy gains against efficiency. To address this, we introduce BEACON (Bayesian Efficient Adaptive Criterion for Optimal N-stopping), a principled adaptive sampling framework grounded in Sequential Search with Bayesian Learning. BEACON sequentially generates responses from the policy LLM, updates posterior belief over reward distributions in real time without further training, and determines when to stop by weighing expected gains against computational cost. Sampling terminates once the marginal utility of further exploration no longer justifies the expense. We establish both theoretical optimality guarantees and practical tractability, and show empirically that BEACON reduces average sampling by up to 80% while maintaining response quality. We further demonstrate BEACON's utility for cost-efficient preference data generation and outline practical extensions, offering actionable insights for future researchers.

BEACON: Bayesian Optimal Stopping for Efficient LLM Sampling

TL;DR

This work introduces BEACON (Bayesian Efficient Adaptive Criterion for Optimal N-stopping), a principled adaptive sampling framework grounded in Sequential Search with Bayesian Learning, and establishes both theoretical optimality guarantees and practical tractability.

Abstract

Sampling multiple responses is a common way to improve LLM output quality, but it comes at the cost of additional computation. The key challenge is deciding when to stop generating new samples to balance accuracy gains against efficiency. To address this, we introduce BEACON (Bayesian Efficient Adaptive Criterion for Optimal N-stopping), a principled adaptive sampling framework grounded in Sequential Search with Bayesian Learning. BEACON sequentially generates responses from the policy LLM, updates posterior belief over reward distributions in real time without further training, and determines when to stop by weighing expected gains against computational cost. Sampling terminates once the marginal utility of further exploration no longer justifies the expense. We establish both theoretical optimality guarantees and practical tractability, and show empirically that BEACON reduces average sampling by up to 80% while maintaining response quality. We further demonstrate BEACON's utility for cost-efficient preference data generation and outline practical extensions, offering actionable insights for future researchers.
Paper Structure (56 sections, 8 theorems, 22 equations, 9 figures, 7 tables)

This paper contains 56 sections, 8 theorems, 22 equations, 9 figures, 7 tables.

Table of Contents

  1. Introduction
  2. Modeling LLM Sampling as Sequential Search Problem
  3. Methodology Overview
  4. Model Setup
  5. Optimal Sampling Policy
  6. BEACON Framework
  7. Experiments
  8. Main Results
  9. Extensions and Applications
  10. Related Work
  11. Conclusion
  12. Experimental Setup and Details
  13. Main Experiment Setup
  14. Policy Models and Reward Models
  15. Datasets and Evaluation Benchmarks
  16. ...and 41 more sections

Key Result

Theorem 1

After generating initial samples $\{y_1, y_2, \ldots, y_{k_0}\}$ to establish valid posterior parameters, the optimal Bayesian policy at each step $k \geq k_0$ is to continue sampling if and only if: The stopping time $K = \min\{k \geq k_0 : h_{n,k}(\hat{z}_k) \leq c/\sigma_k\} \wedge n$ maximizes the expected net gain $\mathbb{E}[z_K - K \cdot c]$.

Figures (9)

  • Figure 1: Comparison of BEACON adaptive sampling versus Best-of-N fixed sampling. BEACON adaptively determines sample size by learning the reward distribution and determine if additional sampling is worth the cost. Intuitively, BEACON stops earlier for consistent samples and continues sampling to find better solutions for variable reward samples, while Best-of-N always uses fixed samples (in this case, 8) regardless.
  • Figure 2: BEACON framework: The top layer shows sequential sampling of LLM responses with reward model evaluation. The bottom layer illustrates the optimal stopping mechanism, which updates Bayesian posterior beliefs about reward distribution parameters after each sample and determines when to stop based on optimal sampling policy, comparing the index-based threshold to the sampling cost.
  • Figure 3: Impact of the sampling cost ($c$) on the value optimization and the optimal sample size. Higher $c$ results in earlier stopping to reach Bayesian optimality.
  • Figure 4: Comparative performance trajectories on Best-of-N vs. BEACON. When forcing the same number of samples, BEACON achieves better performance; When restricting on a threshold performance, BEACON would require less number of samples.
  • Figure 5: Iterative DPO performance on MATH500 with Qwen2.5-7B-instruct. BEACON (accuracy: green squares) achieves competitive accuracy relative to BoN (fixed $n{=}16$, accuracy: blue circles) across seven DPO iterations while reducing the average number of samples required per prompt ($\overline{K}$: red dashed line, right y-axis).
  • ...and 4 more figures

Theorems & Definitions (10)

  • Definition 1
  • Theorem 1: Optimal Sampling Policy
  • Lemma 1
  • Definition 2
  • Theorem 2
  • Theorem 3
  • Lemma 2: Basic properties of $R_t$
  • Lemma 3: Strictly decreasing “continue” margin
  • Theorem 4: Reservation--cost policy
  • Proposition 1