Title: Every Rollout Counts: Optimal Resource Allocation for Efficient Test-Time Scaling
Venue: NeurIPS 2025
Accepted: yes
Keywords: Test-Time Scaling

Resource Allocation

Large Language Model
Abstract: Test-Time Scaling (TTS) improves the performance of Large Language Models (LLMs) by using additional inference-time computation to explore multiple reasoning paths through search. Yet how to allocate a fixed rollout budget most effectively during search remains underexplored, often resulting in inefficient use of compute at test time. To bridge this gap, we formulate test-time search as a resource allocation problem and derive the optimal allocation strategy that maximizes the probability of obtaining a correct solution under a fixed rollout budget. Within this formulation, we reveal a core limitation of existing search methods: solution-level allocation tends to favor reasoning directions with more candidates, leading to theoretically suboptimal and inefficient use of compute. To address this, we propose Direction-Oriented Resource Allocation (DORA), a provably optimal method that mitigates this bias by decoupling direction quality from candidate count and allocating resources at the direction level. To demonstrate DORA’s effectiveness, we conduct extensive experiments on challenging mathematical reasoning benchmarks including MATH500, AIME2024, and AIME2025. The empirical results show that DORA consistently outperforms strong baselines with comparable computational cost, achieving state-of-the-art accuracy.  We hope our findings contribute to a broader understanding of optimal TTS for LLMs.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors address a fundamental  aspect of TTS, specifically the optimal allocation of computational resources (rollouts) to maximise the success probability for LLM during inference. Existing approaches, including beam search, DVTS, and REBASE, tend to allocate resources at the solution level, unintentionally favouring reasoning directions with more candidate solutions. To rectify this inefficiency, the authors propose DORA, a method that allocates resources based on distinct reasoning directions rather than individual candidate solutions. DORA leverages semantic embeddings and a soft clustering approach to identify and balance resources across reasoning directions, theoretically proven optimal under certain assumptions. Extensive experiments conducted on challenging mathematical reasoning tasks (MATH500, AIME2024, AIME2025) demonstrate DORA’s superiority over baseline methods to some extent.\n\nSTRENGTHS AND WEAKNESSES:\n### **Pros:**\n\n1. The derivation of optimal strategies through Bayesian probabilistic modeling provides solid theoretical backing.\n2. Well-written and well-structured paper.\n\n### **Cons:**\n\n1. The claim \"DORA is the most effective parallel search method\" appears overly strong. The paper evaluates DORA exclusively using a single reward model (Qwen2.5-Math-PRM-7B), three relatively small-scale LLMs (<7B), and on a narrow set of benchmarks limited to mathematical reasoning datasets (MATH500, AIME2024, and AIME2025). Such a restricted evaluation setup raises concerns about the generality and robustness of the proposed method. To strengthen this claim, the authors should include additional experiments using alternative reward models and more diverse datasets. \n\n2. While DORA introduces embedding-based clustering to achieve optimal allocation, the computational overhead of embedding calculations, especially in scenarios involving a large number of candidates, is unclear.\n\n3. The reliance on soft clustering and temperature parameters (e.g., allocation temperature) suggests potential sensitivity to hyperparameter settings, yet detailed sensitivity analysis is missing.\n\nQUESTIONS:\n1. How sensitive is DORA’s performance to the choice of hyperparameters (e.g., temperature parameters)? Is there any guidance or robustness analysis for selecting these parameters?\n2. Could you provide insights into the scalability of DORA, particularly how the embedding overhead scales with very large numbers of candidate solutions?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates the problem of test-time rollout allocation for large language models (LLMs) under a fixed computational budget. The authors formalize the rollout allocation process as a resource optimization problem, and propose a novel method called DORA. Experiments on challenging mathematical benchmarks show consistent performance gains over existing methods.\n\nSTRENGTHS AND WEAKNESSES:\nThe paper tackles a practically important yet underexplored problem, the proposed direction-aware allocation strategy is novel and effectively addresses a real shortcoming in rollout allocation, showing state-of-the-art performance on multiple mathematical reasoning tasks and demonstrates superior compute efficiency.\n\nQUESTIONS:\n(1)\tAssumption of soft clustering quality: DORA relies on semantic embeddings and soft clustering via cosine similarity. However, it is unclear how robust this semantic grouping is across tasks and models. How sensitive is DORA to the choice of embedding model and the temperature parameters $T_s$?\n\n(2)\tScalability concerns: While the reported FLOPs and latency gains are encouraging, DORA introduces an extra embedding + affinity matrix computation step. How does this overhead scale with very large numbers of candidate trajectories (e.g., >10,000) which may lead to a failure case?\n\n(3)\tAssumption of shared PRM score per direction: The theoretical justification of DORA assumes that all solutions in a direction $j$ share the same PRM score $R_j$. In practice, PRM scores are noisy and not strictly aligned. How is this mismatch handled, and what is the empirical impact when intra-directional variance in PRM scores is high?\n\n(4)\tNo ablation on clustering effect: The paper does not include ablation studies comparing with other clustering methods at all. This would help isolate whether the gains come from semantic direction-aware allocation, or merely from reduced redundancy.\n\n(5)\tLimited discussion of generalization: The method is evaluated only on mathematical reasoning tasks. While these are valid and difficult benchmarks, it remains unclear whether DORA generalizes to other domains like code generation, etc.\n\nLIMITATIONS:\nWhile this paper tackles an important problem with an elegant solution and promising results, several key assumptions require further clarification and justification.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper tackles the problem of efficiently allocating a fixed computational budget during test-time search with Large Language Models. The authors reframe Test-Time Scaling as a constrained resource-allocation problem, providing a novel theoretical lens through which to analyze existing search strategies. Using a Bayesian framework with a Beta prior to model uncertainty in reward estimates, they derive the form of the optimal allocation. This analysis reveals that popular methods like Beam Search and REBASE are suboptimal because they allocate resources at the individual solution level, favoring \"reasoning directions\" that are overrepresented with many similar candidates.\n\nThe paper's core contribution is Direction-Oriented Resource Allocation (DORA), which computes semantic embeddings of candidate solutions to identify redundant paths, and then re-weights the rollout allocation to separate a direction's quality from the number of candidates representing it. Extensive experiments on challenging mathematical reasoning benchmarks show that DORA consistently achieves higher accuracy than strong baselines and requires 3.5x fewer FLOPs and 4x lower latency than the next-best method for comparable performance.\n\nSTRENGTHS AND WEAKNESSES:\nI found this to be a very strong paper, with its primary strengths being its principled theory, the novelty of the method, and a compelling set of experiments.\n\nThe paper’s theoretical grounding is a major strength. By moving from simple heuristics to a formal resource allocation problem, the authors offer a unified view of the test-time search landscape. This framework elegantly explains existing methods as special cases and provides a clear motivation for why they might fail. The proposed method, DORA, is a direct and clever solution to the problem identified by the theory. Using semantic similarity to compute a \"uniqueness\" score is a fresh approach that moves beyond simple diversity-promoting heuristics.\n\nMost importantly, the method seems to work very well. The empirical results are a highlight. The gains in accuracy are consistent, but the efficiency gains are what truly make the contribution stand out. A 4x speedup, as shown in Table 1, is a significant practical achievement. The analysis is thorough, covering multiple models, benchmarks, and budget sizes, which gives me confidence in the results. The paper is also exceptionally well-written and easy to follow.\n\nThe weaknesses are minor in comparison but worth noting. The \"provably optimal\" claim rests on the simplifying assumption that all solutions within a \"direction\" have identical PRM scores, a condition that won't hold perfectly in practice. The paper doesn't explore how the guarantees degrade as this assumption is relaxed. Additionally, the method's performance is now tied to new components, the embedding model and a key temperature hyperparameter ($T_s$), whose robustness and sensitivity are not explored in an ablation study. While the chosen settings work well, understanding their stability is important for real-world deployment.\n\nQUESTIONS:\nThis is an excellent submission, and my questions are mainly aimed at exploring the boundaries of the claims and the method's robustness. A strong response here could easily convince me to raise my score.\n\n1. Exploring the Limits of Optimality: Your optimality proof hinges on some key assumptions. I'd be interested to hear more about what happens when they are pushed.\n(i) How does DORA's performance hold up as the variance of PRM scores within a direction increases, violating the core assumption of Theorem 1?\n(ii) What happens in the high-confidence regime $(k → \\infty)$? When the PRM is highly reliable, the best strategy is to exploit the top candidate. Does DORA gracefully recover this behavior, or could its diversity-promoting mechanism become a liability by steering compute away from a single, truly optimal path?\n\n2. Robustness to Hyperparameters: The performance seems to depend on a very specific temperature, $T_s = 0.01$. This value makes the softmax behave almost like a hard assignment. How sensitive is the method to this choice? An ablation plot showing performance as $T_s$ is varied (e.g., across {0.001, 0.01, 0.1, 1.0}) would be very helpful for understanding if the method is stable or requires careful tuning.\n\n3. The Clustering Mechanism: DORA's effectiveness relies on its soft clustering. To better understand why it works, I have two questions:\n(i) Have you considered a simpler \"hard clustering\" baseline? For example, one could run K-means on the embeddings and then allocate the budget at the cluster level. Comparing against this would clarify whether the benefit comes simply from accounting for redundancy, or if the soft-reweighting scheme is truly superior.\n(ii) Could you provide a more direct, empirical measure of how well the clustering works? For instance, measuring the success-rate correlation between pairs of solutions before and after they are grouped by DORA would be a great way to quantify how effectively it identifies and mitigates dependencies.\n\nLIMITATIONS:\nYes, the authors do a good job of discussing limitations. I'd encourage them to also briefly touch on the following points to make the section even stronger:\n- The gap between the theoretical assumption of identical PRM scores and the practical reality.\n- The introduction of new components (embedding model, T_s) and the potential for their selection to induce new biases or require careful tuning.\n- A brief note that the empirical validation, while strong, is confined to mathematical reasoning, and generalization to other domains (e.g., code generation) remains an open and interesting question.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a critical analysis of rollout allocation efficiency in Test-Time Scaling for Large Language Models (LLMs) focused on mathematical reasoning, leveraging Process Reward Models (PRMs). It identifies limitations in prior methods by demonstrating that the assumption of independence between rollout allocations fails in practical scenarios, leading to suboptimal performance. To address this, the authors propose a theoretical framework that incorporates embedding similarity scores to mitigate the breakdown of independence, thereby achieving optimality under the proposed conditions. The framework is rigorously evaluated on three standard mathematical reasoning benchmarks using a 7B-parameter PRM and 1 to 3B-parameter generative models, demonstrating its effectiveness in improving scalability and performance for mathematical tasks. The work contributes a novel approach to optimizing test-time scaling in LLMs through both theoretical justification and empirical validation.\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths\n1. Theoretical Justification: The method derives the sub-optimal conditions for REBASE and proposes more effective method based on an additional score from embedding similarity. The theoretical result establishes the optimality of the direction-level allocation schema proposed in the paper.\n\n2. Performance: The proposed framework achieves superior performance at lower computer cost in comparison to the baseline, for example, 3.5 times  fewer flop in MATH500, resulting in a lower latency as well.\n\n### Weaknesses\n1. The generalization of the presented results are limited. First, the  improvement is mainly visible for MATH500, but for AIME task, there is a good amount of overlap with baselines for statistical significance. Next, the range of policy models are limited as it covers only 1-3B models. Evaluating larger model classes such as 8B and 32B could further establish the significance of the proposed methods. Similarly, only single PRM is evaluated. The results for larger PRMs (72B) and other families (e.g., Skywork) could clarify the scalability of the  proposed method.\n\n2. Motivation of  some design choices are  unclear. For example, semantic similarity temperature is chosen as 0.01. It is unclear how is it tuned and the sensitivity of the results due  to this parameter.\n\n3. The proposed method introduces an additional model dependency for Test-Time Scaling which could hinder its future extension.  For example, it the retriever model is not  trained with code data, the scaling method will  face bottleneck for coding benchmarks.\n\n4. Table 1 does not include the error bars (deviation) and statistics for AIME tasks. Including these metrics could further clarify the scalability of the proposed framework.\n\nQUESTIONS:\n1. How much sensitive could the proposed method be towards the choice of retriever (E5/GTE instead of  BGE)?\n2. Which chat template is used for Llama models as  there seems to be an issue related to it[1]?\n3. The quality score parameter T_b is not well ablated in the REBASE paper. How much are the results in the proposed paper sensitive to this parameter?\n\n\nReferences:\n\n1. https://huggingface.co/spaces/HuggingFaceH4/blogpost-scaling-test-time-compute\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}