PaperlayerView exact revision on arXiv

Source-linked AI summary

PETS: A Principled Framework Towards Optimal Trajectory Allocation for Efficient Test-Time Self-Consistency

Zhangyi Liu, Huaizhi Qu, Xiaowei Yin, He Sun, Yanjun Han, Tianlong Chen, Zhun Deng

arXiv:2602.16745v2cs.LGcs.AI

TL;DR

PETS addresses the challenge of making test-time self-consistency sample-efficient under limited budgets. It defines self-consistency against the infinite-budget majority and develops difficulty-aware allocation methods for offline and online settings. Across benchmarks and reasoning models, PETS improves self-consistency and accuracy while reducing required trajectories, with limitations concerning majority errors and online difficulty estimation.

  • Problem

    Sample-efficient test-time self-consistency remains challenging because self-consistency requires many stochastic reasoning traces, while existing allocation methods often rely on heuristic signals or auxiliary reward models.

  • Method

    PETS optimizes self-consistency rate, connects offline trajectory allocation to crowdsourcing, and assigns online budgets using estimated question-difficulty distributions.

  • Results

    PETS consistently increases self-consistency and accuracy while substantially reducing the number of required reasoning trajectories compared with uniform allocation.

  • Takeaways & Limitations

    Difficulty-aware trajectory allocation can improve the efficiency of parallel sampling without sacrificing performance across challenging benchmarks and state-of-the-art reasoning models.

  • Takeaways & Limitations

    Allocation-only approaches provide little benefit when the population majority is systematically wrong, and PETS-Online estimates difficulty through a short warm-up phase rather than directly from the question.

Abstract

from arXiv · show

Test-time scaling can improve model performance by aggregating stochastic reasoning trajectories. However, achieving sample-efficient test-time self-consistency under a limited budget remains an open challenge. We introduce PETS (Principled and Efficient Test-TimeSelf-Consistency), which initiates a principled study of trajectory allocation through an optimization framework. Central to our approach is the self-consistency rate, a new measure defined as agreement with the infinite-budget majority vote. This formulation makes sample-efficient test-time allocation theoretically grounded and amenable to rigorous analysis. We study both offline and online settings. In the offline regime, where all questions are known in advance, we connect trajectory allocation to crowdsourcing, a classic and well-developed area, by modeling reasoning traces as workers. This perspective allows us to leverage rich existing theory, yielding theoretical guarantees and an efficient majority-voting-based allocation algorithm. In the online streaming regime, where questions arrive sequentially and allocations must be made on the fly, we propose a novel method inspired by the offline framework. Our approach adapts budgets to question difficulty while preserving strong theoretical guarantees and computational efficiency. Experiments show that PETS consistently outperforms uniform allocation. On GPQA, PETS achieves perfect self-consistency in both settings while reducing the sampling budget by up to 75% (offline) and 55% (online) relative to uniform allocation. Code is available at https://github.com/ZDCSlab/PETS.

1 Introduction

PETS frames test-time trajectory allocation as an optimization problem grounded in self-consistency, targeting agreement with the infinite-budget majority while distributing samples by question difficulty. It develops offline and online allocation methods and reports improved efficiency and accuracy over uniform allocation.

  • Motivation and framework: Self-consistency rate measures the probability that finite-budget majority voting matches the population majority, providing a principled target for allocation.The formulation supports rigorous analysis of finite-budget allocation policies.
  • Motivation and framework: PETS allocates sampling effort according to question difficulty rather than uniformly, without relying on auxiliary supervision.Different difficulty levels exhibit distinct convergence behaviors under sampling.
  • Offline and online settings: Online PETS estimates the difficulty distribution from training samples, solves a constrained allocation problem, and assigns budgets when questions arrive.The method operates without access to the full future question set.
  • Offline and online settings: Offline PETS connects trajectory allocation to crowdsourcing and uses adaptive allocation machinery to maximize expected self-consistency gains.Reasoning traces are modeled as workers in the offline setting.
  • Empirical findings: PETS consistently reduces trajectories needed for full self-consistency, improves self-consistency under fixed budgets, and achieves higher accuracy than uniform baselines.The reported gains hold across unweighted and weighted cases.

2 Setup

The setup defines finite- and infinite-budget voting over stochastic reasoning traces and uses self-consistency as the probability of recovering the population majority. It formulates allocation as maximizing aggregate self-consistency under a limited total budget in offline and online information settings.

  • Sampling and voting: Each question produces i.i.d. stochastic reasoning traces whose answers are sampled from a finite answer set.The setup treats traces as conditionally independent given the question.
  • Sampling and voting: Majority voting selects the answer with the largest vote count, while weighted voting uses trace-dependent confidence weights.The weighted formulation aggregates confidence-adjusted votes.
  • Self-consistency: The infinite-budget majority vote defines the population consensus toward which finite-budget voting is compared.It formalizes the limiting outcome of aggregating infinitely many trajectories.
  • Self-consistency: Self-consistency rate is the probability that a budget-B vote recovers the population majority label.Its growth with B is governed by the question’s answer distribution and reflects difficulty.
  • Allocation objective: The allocation objective is to choose per-question trace budgets that maximize self-consistency under a limited total reasoning budget.The formulation accommodates policy-dependent posterior information.
  • Information settings: Offline allocation uses the full question set, whereas online allocation assigns each arriving question a budget without seeing future questions.The two settings differ in the information available when budgets are chosen.

3 Offline PETS in the Batch Setting

Offline PETS treats trajectory allocation as a Bayesian decision problem connected to crowdsourcing, using posterior beliefs to allocate limited traces toward higher expected self-consistency.

  • Crowdsourcing connection: Offline allocation models reasoning answers as noisy worker labels and additional traces as labeling effort under a global budget.This connection permits use of Bayesian crowdsourcing methods for learning while allocating.
  • Crowdsourcing connection: Confidence-weighted aggregation extends the crowdsourcing-based framework beyond standard majority voting.The weighted formulation differs from worker-reliability weighting and requires additional care.
  • Bayesian formulation: The Bayesian setup maintains posteriors over each question’s answer distribution and class-conditional confidence means.Answers are modeled categorically, while confidence is modeled with answer-dependent Gaussian means.
  • Bayesian formulation: The offline process is a finite-horizon Bayesian MDP in which each action allocates one additional trace and updates posterior beliefs.The horizon corresponds to the total trace budget, and terminal reward sums posterior self-consistency across questions.
  • Allocation algorithm: PETS-Offline uses Optimistic Knowledge Gradient to approximate intractable dynamic programming by selecting the question with the largest optimistic one-step utility improvement.The exact dynamic program is difficult because the belief space grows exponentially.

4 Online PETS in the Streaming Setting

Online PETS assigns budgets in one shot from estimated question difficulty, using a constrained optimization and greedy allocation procedure for sequential arrivals. Its allocation proportions approach those of offline PETS as budgets grow.

  • Online setting: Online PETS assigns each arriving question a budget from its estimated difficulty without observing future questions or intermediate cross-question feedback.The policy relies on a prior difficulty distribution estimated from additional training data.
  • Execution protocol: The execution protocol estimates problem-parameter distributions, maps incoming questions to difficulty grids, and solves a budget-allocation optimization problem.Training builds difficulty grids and representative self-consistency curves; test time assigns a grid and budget.
  • Execution protocol: A Gaussian-probit surrogate approximates each question’s self-consistency curve before discretizing questions into difficulty grids with prototype curves.The surrogate is fitted from sampled training answers and motivated by a normal approximation to multinomial vote margins.
  • Execution protocol: Four warm-up responses produce a sorted option-count pattern used for deterministic grid assignment before one-shot budget allocation.The construction remains independent of the total number of answer options because at most four options appear in warm-up samples.
  • Optimal allocation: The marginal gain R(θ,n) is the increase in self-consistency from allocating one additional budget unit, and the greedy algorithm repeatedly selects the largest current gain.The procedure starts with the binary-choice case and extends to multi-choice settings.
  • Optimal allocation: Algorithm 1 is optimal in expectation for the discretized online allocation problem when randomized rounding is used.The remaining budget may be insufficient for another full iteration.
  • Comparison with offline allocation: As the average per-question budget increases, offline and online allocation plans converge to nearly identical budget proportions.Figure 2 illustrates this convergence on nine simulated binary-choice questions, where larger θ indicates easier questions.

5 Experiment

PETS is evaluated on knowledge and reasoning benchmarks in offline and online allocation settings using self-consistency and accuracy, with uniform and confidence-guided baselines. Across these settings, PETS reduces sampling needed for full self-consistency, improves accuracy, and remains practical for streaming allocation.

  • Evaluation Setup: PETS is evaluated across five benchmarks and five reasoning language models using repeated subsampling of 64 traces from 128 responses per question.The subsampling is repeated 30 times, with mean performance reported.
  • Offline Budget Allocation: PETS-Offline reduces traces required for full self-consistency by up to 75% versus uniform allocation and achieves higher self-consistency at any budget.Offline results compare methods within matched unweighted or confidence-weighted settings.
  • Offline Budget Allocation: Confidence-weighted majority voting shifts the population majority toward the ground truth, resulting in improved accuracy when the population majority aligns with the true answer.The confidence-weighted variant is included at aggregation time and, for PETS-Offline, at sampling time.
  • Online Budget Allocation: PETS-Online reduces traces needed for full self-consistency, improves accuracy, and closely matches PETS-Oracle in both self-consistency and accuracy.PETS-Online makes irreversible allocation decisions using observed traces, while PETS-Oracle accesses latent difficulty θ.
  • Additional Results: Accuracy varies only modestly across online difficulty-bin counts, and coarse grouping is sufficient for the allocation policy in several AIME 2024 settings.The oracle clusters ground-truth parameters, while the practical predictor uses K0=4 warm-up samples.
  • Additional Results: The online predictor adds only a few seconds of preprocessing and budget-assignment overhead, whereas the offline OKG allocator is substantially more expensive.This supports PETS-Online as a practical streaming policy with small decision-time overhead.

6 Related Work

Related work spans test-time scaling and efficient reasoning, as well as crowdsourcing and adaptive allocation. PETS connects trajectory allocation to these lines through sample aggregation, noisy-worker modeling, and adaptive budget assignment.

  • Test-time scaling and efficient reasoning: Test-time scaling extends or samples reasoning trajectories, while self-consistency and Best-of-N aggregate them through voting or selection.Efficient-reasoning methods address the cost of naive scaling by shortening, budgeting, pruning, or filtering trajectories.
  • Crowdsourcing and adaptive allocation: Crowdsourcing methods model noisy workers and item difficulty, while adaptive assignment and knowledge-gradient methods allocate labels using expected value of information.Related work also studies aggregation, adaptive stopping, and compute allocation across queries, tokens, or rollouts.

7 Conclusion

PETS improves parallel-sampling efficiency without sacrificing performance, while its conclusions identify important limits of allocation-only approaches and opportunities for direct difficulty prediction.

  • PETS consistently increases self-consistency and accuracy while substantially reducing required reasoning trajectories across challenging benchmarks and reasoning models.
  • Offline OKG uses sequential Monte Carlo allocation with B=16 per question and nsamples=500, whereas online methods use K0=4 warm-up samples.
  • PETS-Online estimates latent question difficulty through a short warm-up phase rather than predicting it directly from the question before generation.
  • When the population majority is systematically wrong, additional sampling provides little benefit, limiting allocation-only approaches.

B.3 Proof of lemma 3.1

The proof formulates terminal self-consistency as a separable objective across questions and establishes that the streaming policy remains feasible through dynamic replanning and randomized rounding.

  • After allocating H traces, the offline objective is the conditional expected self-consistency rate represented by terminal posterior beliefs.
  • Because each terminal label appears only in its question’s summand, maximizing total self-consistency reduces to independently maximizing every question’s term.
  • The streaming implementation always respects the global budget because remaining budget decreases by exactly the amount spent at each step.
  • The online policy recomputes grid-level allocations using the remaining budget-to-horizon ratio after each arriving question.
  • Randomized rounding executes the next two-trial increment with probability δ/(2pj⋆), preserving feasibility while attaining the optimal convex combination of neighboring allocations.

C.2 Proof of Optimality of Algorithm 1 (Theorem 4.1)

The optimality proof shows that majority-vote allocation has diminishing effective marginal returns, so greedy selection of the largest available increment solves the discretized problem.

  • The allocation objective’s marginal gain and expected cost both contain the grid probability, which cancels when comparing marginal gain per expected-cost unit.
  • Adding two trials from 2m−1 to 2m+1 yields the effective marginal reward R(θ,2m), and these rewards decrease with m.
  • Greedy allocation selects the grid with the largest currently available effective marginal reward, thereby choosing the globally largest feasible rewards.
  • Randomizing an incomplete final two-trial step preserves exact expected-budget feasibility and achieves the optimal convex combination of neighboring integer allocations.

C.3 2-parameter Approximation in Multi-Choice Case

For multi-choice questions, PETS approximates self-consistency with a two-parameter Gaussian-probit curve whose error has a 1/√n rate and preserves the diminishing-returns structure needed for greedy allocation.

  • The multi-choice difficulty parameter is high-dimensional, making direct optimization costly; PETS therefore uses a Gaussian surrogate family Φ(a√n+b).
  • The approximation is implemented by replacing SC(θ;n) with Φ(a√n+b), with parameters fit through regression in practice.
  • Proposition C.2 establishes a Gaussian-probit approximation to multinomial majority-vote success with a 1/√n error rate.
  • The surrogate is concave on the relevant greedy budget range, so its marginal gains are nonincreasing and Algorithm 1 applies unchanged.
  • Figure 5 compares exact or Monte Carlo curves with fitted probit curves for M∈{2,4,10}, which closely track across regimes.
  • Online estimation uses four warm-up responses to assign questions to one of five count-pattern grids before allocating additional samples.

D.1 Convergence of Budget Proportions for the Online Case

The online allocation algorithm samples every problem type indefinitely and converges to stable budget proportions. Its limiting allocation is inversely related to the difficulty-dependent quantity p(θi).

  • Asymptotic coverage: Every problem type is sampled infinitely often as the total budget B approaches infinity.The proof uses diminishing marginal gains: a type sampled only finitely often would retain a positive gain and eventually be selected.
  • Asymptotic proportions: The sampling ratio mi(B)/B converges to a constant for each problem type as B approaches infinity.This establishes asymptotically stable allocation proportions across types.
  • Limiting allocation: The limiting allocation satisfies B ∝ 1 p(θi).Thus, asymptotic budget proportions are determined by the inverse of the difficulty-dependent quantity p(θi).
  • Proof mechanism: The algorithm repeatedly selects the type with the minimum potential and increases its potential while leaving other potentials unchanged.The potential-range argument underlies convergence of the allocation counts.

D.2 Convergence of Budget Proportions for Offline Case

The offline analysis derives asymptotic budget proportions for equal-weight Offline PETS using Beta-tail identities and binomial asymptotics. It proves convergence and identifies an inverse difficulty-dependent allocation rule.

  • Beta-tail analysis: The derivation establishes a closed form for Beta-tail probabilities and verifies it through matching derivatives and boundary values.The proof uses polynomial identities, term-by-term cancellation, and the fundamental theorem of calculus.
  • Convergence result: Proposition D.6 states that offline sampling ratios mi(B)/B converge for every question as B approaches infinity.The result applies to equal-weight Offline PETS with Beta(1,1) when all questions and their difficulty labels are available upfront.
  • Limiting proportions: The limiting offline allocation obeys B ∝ 1 p(θi).The same inverse p(θi) relationship characterizes the asymptotic budget allocation stated for the offline case.
  • Asymptotic approximation: The asymptotic analysis uses binomial sampling paths and Stirling’s formula to connect empirical outcome proportions with the latent success probability.The binomial proportions converge almost surely to θ/(1−θ), supporting the limiting allocation analysis.
  • Offline–online distinction: Offline and online asymptotic limits are not exactly identical because the two settings use different statistical viewpoints.The offline setting infers difficulty parameters during allocation and therefore requires Bayesian posterior updating for optimal decisions.

E.3.2 Full PETS-Online Results

The full online-results section covers GPQA, AIME 25, AIME 24, HMMT, and BRUMO through dataset-specific figures and an online-results table. The comparison includes realistic and oracle online settings.

  • Dataset coverage: The online evaluation includes GPQA, AIME 25, AIME 24, HMMT, and BRUMO.These datasets are represented by Figures 11–15 for the online setting.
  • Evaluation settings: Table 8 reports PETS-Online results and distinguishes oracle access to latent θ from realistic online estimation.In the realistic setting, θ is unavailable and must be learned from a training dataset.
  • Comparison summary: Figure 17 summarizes the detailed online comparison across datasets and models using Pass@1 and MV@128 as reference points.Pass@1 is a lower-bound reference, while MV@128 is a finite proxy for infinite-budget performance.
Loading 2602.16745v2…