Latency-Aware Adaptive Best-of-N

Best-of-$N$ sampling is a widely deployed inference-time compute method for language models, but it can be computationally inefficient when applied uniformly across prompts without adjusting for problem difficulty. Recent adaptive variants allocate compute based on prompt-level signals such as observed rewards, curtailing generation when additional generations or tokens are unlikely to improve the final response. These methods, however, assume that the cost of additional compute is a fixed constant – an assumption violated in real serving systems, where the wall-clock time to generate a batch varies, growing with current server load, so a method using a single fixed parameter spends too little compute when generation is fast and too much when it is slow. We propose a two-stage adaptive Best-of-$N$ algorithm that adapts jointly to prompt difficulty and to generation latency: it draws an initial batch of completions in parallel, then uses the observed rewards and the time taken to generate them to decide whether to draw a second batch and at what size by weighing expected improvement in reward against predicted cost. Concretely, the algorithm is less likely to draw a second batch when the first batch was slow, and more likely when it was fast. We use a CDF-based reward transformation – mapping each observed reward to its estimated quantile under the prompt's reward distribution – to obtain a closed-form expression for the expected improvement that requires no Monte Carlo estimation. Because reward distributions have arbitrarily different scales across prompts, language models, and reward models, the CDF transformation also allows the same tradeoff parameters to be reused across tasks without retuning. We evaluate the method on a preference-based alignment benchmark and show that, under fixed latency, it improves over parallel Best-of-$N$ on the reward–compute Pareto frontier, and that, under variable latency, conditioning on the measured first-batch latency gives a higher *user reward* – the reward minus the cost of latency – than fixed-threshold (latency-unaware) baselines at matched numbers of generations, with the advantage growing as latency variability increases.