Predicting and improving test-time scaling laws via reward tail-guided search

TL;DR

It is theoretically prove that SLG achieves vanishing regret compared to perfect-information oracles, and achieves expected rewards that would otherwise require a polynomially larger compute budget required when using BoN.

Abstract

Test-time scaling has emerged as a critical avenue for enhancing the reasoning capabilities of Large Language Models (LLMs). Though the straight-forward ''best-of-$N$'' (BoN) strategy has already demonstrated significant improvements in performance, it lacks principled guidance on the choice of $N$, budget allocation, and multi-stage decision-making, thereby leaving substantial room for optimization. While many works have explored such optimization, rigorous theoretical guarantees remain limited. In this work, we propose new methodologies to predict and improve scaling properties via tail-guided search. By estimating the tail distribution of rewards, our method predicts the scaling law of LLMs without the need for exhaustive evaluations. Leveraging this prediction tool, we introduce Scaling-Law Guided (SLG) Search, a new test-time algorithm that dynamically allocates compute to identify and exploit intermediate states with the highest predicted potential. We theoretically prove that SLG achieves vanishing regret compared to perfect-information oracles, and achieves expected rewards that would otherwise require a polynomially larger compute budget required when using BoN. Empirically, we validate our framework across different LLMs and reward models, confirming that tail-guided allocation consistently achieves higher reward yields than Best-of-$N$ under identical compute budgets. Our code is available at https://github.com/PotatoJnny/Scaling-Law-Guided-search.

Figures (5)

• Figure 1: Empirical analysis of reward distributions. We visualize reward samples ($N=5000$) for two distinct partial responses generated by Llama-3.2-1B-Instruct on an AIME 2024 problem, scored by Skywork-Reward-V2. (Left) The standard density histograms show that the distribution body may contain irregularities or deviations from a perfect bell curve. (Right) The log-density plots reveal that despite body noise, the upper tail (highlighted in color) strictly follows a parabolic decay, confirming that the tail behavior is well-approximated by Gaussian.
• Figure 2: Main Results on Test-Time Scaling. We compare SLG (Solid) vs. BoN (Dashed) across three datasets (Columns) and two model scales (Rows). Here, 1B and 7B denote Llama-3.2-1B-Instruct and Qwen2.5-7B-Instruct. SLG demonstrates superior efficiency across all configurations.
• Figure 3: Ablation Studies. (a) Impact of search width $K$ and estimation budget $m$, showing a concave trend with a clear optimum. (b) Comparison against a mean-based selection baseline, confirming the necessity of tail-targeting. (c) Verification using the Qwen-based reward model, demonstrating robustness to the feedback signal.
• Figure 4: Detailed Q-Q Plot Analysis.Top Row: Q-Q plots for the entire reward distribution against a standard Normal. The high density of points and slight deviations at the ends indicate body irregularities. Bottom Row: Q-Q plots for the tail region ($> R_{0.2}$) against a Truncated Normal. The strictly linear alignment and near-perfect $R^2$ scores confirm the validity of the Gaussian Tail Assumption for extrapolation.
• Figure 5: Broad Visual Validation of Gaussian Tails. We visualize the reward densities for 9 randomly sampled query instances distinct from the main text examples. Grey Bars: The distribution body ($< 80$th percentile). Blue Bars: The tail region ($> 80$th percentile). Dashed Line: A Gaussian distribution fitted specifically to the tail statistics. Despite the irregular and diverse shapes of the distribution bodies, the tail regions consistently follow a predictable Gaussian profile.

Theorems & Definitions (29)

• Theorem 1
• Definition 1: Full-information oracle $\mathcal{A}^*$
• Definition 2: Best-of-$N$ baseline $\mathcal{A}_{\text{BoN}}$
• Theorem 2
• Proposition 1
• Corollary 1
• Lemma 1
• proof : Proof of Lemma \ref{['lemma:threshold-error']}
• Lemma 2
• proof : Proof of Lemma \ref{['lemma:bias-error']}