Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Optimal Stopping vs Best-of-$N$ for Inference Time Optimization

Yusuf Kalayci, Vinod Raman, Shaddin Dughmi

TL;DR

This work reframes LLM inference-time optimization as a Pandora's Box stopping problem, where each generated sample incurs a cost and reveals a random reward. It delivers a principled adaptive strategy by developing a UCB-style Pandora's Box algorithm that learns the unknown reward distribution, plus a Bradley–Terry-based normalization to handle cross-prompt scaling, enabling practical, on-the-fly stopping thresholds. The authors provide theoretical guarantees in the known-distribution limit via Weitzman's optimal policy, and finite-time guarantees under distributional families (notably Exponential) when the distribution is unknown. Empirically, the approach achieves the same reward as non-adaptive Best-of-$N$ while using 15–35% fewer generations on average across diverse prompts, models, and reward functions, thus offering meaningful efficiency gains for deployed LLM systems.

Abstract

Large language model (LLM) generation often requires balancing output quality against inference cost, especially when using multiple generations. We introduce a new framework for inference-time optimization based on the classical Pandora's Box problem. Viewing each generation as opening a costly "box" with random reward, we develop algorithms that decide when to stop generating without knowing the underlying reward distribution. Our first contribution is a UCB-style Pandora's Box algorithm, which achieves performance that is provably close to Weitzman's algorithm, the optimal strategy when the distribution is known. We further adapt this method to practical LLM settings by addressing reward scaling across prompts via a Bradley-Terry inspired transformation. This leads to an adaptive inference-time optimization method that normalizes rewards and learns stopping thresholds on the fly. Experiments on the AlpacaFarm and HH-RLHF datasets, using multiple LLM-reward model pairs, show that our adaptive strategy can obtain the same performance as non-adaptive Best-of-N sampling while requiring 15-35 percent fewer generations on average. Our results establish a principled bridge between optimal stopping theory and inference-time scaling, providing both theoretical performance bounds and practical efficiency gains for LLM deployment.

Optimal Stopping vs Best-of-$N$ for Inference Time Optimization

TL;DR

This work reframes LLM inference-time optimization as a Pandora's Box stopping problem, where each generated sample incurs a cost and reveals a random reward. It delivers a principled adaptive strategy by developing a UCB-style Pandora's Box algorithm that learns the unknown reward distribution, plus a Bradley–Terry-based normalization to handle cross-prompt scaling, enabling practical, on-the-fly stopping thresholds. The authors provide theoretical guarantees in the known-distribution limit via Weitzman's optimal policy, and finite-time guarantees under distributional families (notably Exponential) when the distribution is unknown. Empirically, the approach achieves the same reward as non-adaptive Best-of- while using 15–35% fewer generations on average across diverse prompts, models, and reward functions, thus offering meaningful efficiency gains for deployed LLM systems.

Abstract

Large language model (LLM) generation often requires balancing output quality against inference cost, especially when using multiple generations. We introduce a new framework for inference-time optimization based on the classical Pandora's Box problem. Viewing each generation as opening a costly "box" with random reward, we develop algorithms that decide when to stop generating without knowing the underlying reward distribution. Our first contribution is a UCB-style Pandora's Box algorithm, which achieves performance that is provably close to Weitzman's algorithm, the optimal strategy when the distribution is known. We further adapt this method to practical LLM settings by addressing reward scaling across prompts via a Bradley-Terry inspired transformation. This leads to an adaptive inference-time optimization method that normalizes rewards and learns stopping thresholds on the fly. Experiments on the AlpacaFarm and HH-RLHF datasets, using multiple LLM-reward model pairs, show that our adaptive strategy can obtain the same performance as non-adaptive Best-of-N sampling while requiring 15-35 percent fewer generations on average. Our results establish a principled bridge between optimal stopping theory and inference-time scaling, providing both theoretical performance bounds and practical efficiency gains for LLM deployment.

Paper Structure

This paper contains 32 sections, 8 theorems, 69 equations, 7 figures, 2 algorithms.

Table of Contents

  1. Introduction
  2. Our Contributions.
  3. Related Works
  4. Preliminaries
  5. Notation
  6. Test-time Steering and Best-of-$N$ sampling
  7. Optimal Stopping, Pandora's Box, and Weitzman's Algorithm
  8. Pandora's Box with Unknown Reward Distributions
  9. Main Theoretical Result
  10. Example: Exponential distribution with unknown parameter $\lambda$
  11. Applying UCB Pandora's Box to Inference-time Optimization
  12. The Reward Scaling Challenge
  13. The Adaptive Algorithm
  14. Target Acceptance Rate Variant
  15. Experimental Results
  16. ...and 17 more sections

Key Result

Theorem 5

Let $\mathcal{F}$ be a family of distributions and $c > 0$ be the cost value for each sample. Let $\tau^{+}_{\delta}(n, v_{1:n})$ be an anytime upper confidence bound with deterministic width $\sigma_{\delta, \tau}(n)$ for $\mathcal{F}$ according to Definition def:csfaircap. i.e., for every distribu Then for every $D \in \mathcal{F}$ and $\delta \in (0, 1)$, we have that where $F_D$ is the CDF of

Figures (7)

  • Figure 1: Our algorithm (red) matches optimal non-adaptive performance across varying cost ratios.
  • Figure 2: Win rate of the adaptive algorithm compared to non-adaptive Best-of-$N$. Adaptive stopping leverages computation more effectively, particularly at lower costs.
  • Figure 3: Adaptive algorithm achieves specified target acceptance rates while saving 15--35% of samples compared to non-adaptive Best-of-$N$.
  • Figure 4: Median rewards across $960$ generations for $100$ prompts from the AlpacaEval dataset.
  • Figure 5: Profit performance ratios for four LLM generators across varying cost/utility ratios. Each panel represents a different dataset-reward model combination, with four cost/utility scenarios per panel. Higher ratios indicate better profit efficiency relative to best non-adaptive algorithm.
  • ...and 2 more figures

Theorems & Definitions (24)

  • Definition 1: Fair-cap value
  • Definition 2: Weitzman's Algorithm for infinitely-many identical boxes
  • Definition 3: Additive Sub-optimality Gap
  • Definition 4: Anytime Valid Upper Confidence Bound on the Fair-cap Value
  • Theorem 5: Upper bound on Additive Sub-optimality
  • Theorem 6: Anytime-valid Upper Confidence Bound for Exponential Fair-cap value
  • Corollary 7: Sub-optimality gap upper bound for Exponential distribution
  • Definition 8: Benchmark Reward
  • Definition 9: Acceptance Rate
  • Definition 10: Utility Function
  • ...and 14 more