Accelerating LLM Reasoning via Early Rejection with Partial Reward Modeling (2025)

Abstract: Large Language Models (LLMs) are increasingly relied upon for solving complex reasoning tasks in domains such as mathematics, logic, and multi-step question answering. A growing line of work seeks to improve reasoning quality by scaling inference time compute particularly through Process Reward Models (PRMs), used to reward the reasoning at intermediate steps. While effective, these methods introduce substantial computational overhead, especially when generating large numbers of solutions in parallel. In this paper, we investigate whether PRMs can be used mid-generation to provide early signals that enable the rejection of suboptimal candidates before full generation of step is complete. We introduce the hypothesis that PRMs are also Partial Reward Models, meaning that the scores they assign to partially completed reasoning step are predictive of final output quality. This allows for principled early rejection based on intermediate token-level signals. We support this hypothesis both theoretically, by proving that the risk of discarding optimal beams decreases exponentially with generation length and empirically, by demonstrating a strong correlation between partial and final rewards across multiple reward models. On math reasoning benchmarks, our method achieves up to 1.4 \times – 9 \times reduction in inference FLOPs without degrading final performance. These results suggest that early rejection is a powerful mechanism for improving the compute-efficiency of reasoning in LLMs.

## introduction large language models (llms) are at the forefront of ai capabilities due to their emerging ability to perform complex reasoning tasks @xcite . prior works. as the scaling of model parameters and pretraining data has started to become a bottleneck, recent efforts have shifted toward increasing compute at inference time to improve the reasoning capabilities of llms @xcite . improving the reasoning capabilities of llms by scaling compute at inference time has been pursued through multiple strategies. a prominent approach leverages outcome reward models, which train a separate evaluator to score the final output of the llm based on correctness or quality @xcite @xcite . another approach uses, process reward models (prms) to evaluate intermediate steps or reasoning trajectories generated during inference @xcite @xcite . under this paradigm, the model generates multiple candidate reasoning paths, which are then evaluated by the prm that assigns rewards at the end of each step. this step-wise evaluation enables the selection of promising trajectories for further expansion while allowing the rejection of less promising ones, thereby guiding the reasoning process more efficiently. techniques such as beam search, monte carlo tree search (mcts) guided by value models, and prm-guided methods, exemplify this strategy @xcite . in this 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 partial rewards 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 final rewards r 2 = 0.631 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 partial rewards 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 final rewards r 2 = 0.716 figure 2: linear relationship between partial rewards (reward calculated at half step completion) and full rewards (rewards calculated at full step completion) with (left) llemma-metamath-7b and (right) mathshepherd-mistral-7b as reward models. challenges. while scaling test-time compute using methods like prms can significantly enhance performance, it also introduces substantial computational overhead, especially for long sequences where many generated beams contribute little value @xcite @xcite . to be competitive with state-of-theart post-training approaches, the number of beams often needs to be scaled up to 1000-60000, resulting in a large number of output tokens @xcite . these output tokens are not only computationally expensive to generate, but are also produced sequentially, leading to considerable latency @xcite . a natural solution is to reject unpromising candidates early in the decoding process, after only a few tokens, before committing to full step generation, a strategy we refer to as early rejection.however, a major challenge in early rejection is making sure that decisions based on only part of the output don't accidentally discard high-quality completions. this is difficult because the overall quality of a reasoning trace often depends on its full structure, which might not be obvious from the first few tokens. as a result, developing a reliable method that can make early yet accurate decisions about which traces to keep, based on partial generations, remains an important and open research problem. ## method beam search for reasoning. beam search is a widely used decoding strategy in llms for structured generation tasks such as mathematical problem solving and multi-step reasoning @xcite @xcite . at each decoding step, the model expands a fixedwidth set of n candidate beams by sampling multiple possible continuations and retaining only the top-scoring ones based on a predefined heuristic (e.g., log-probability or reward score). this iterative expansion and rejecting process allows the model to explore a larger space of possible outputs than greedy decoding, while remaining tractable compared to exhaustive search. in prm-guided reasoning, each beam is scored at the end of every reasoning step by a prm, which evaluates the coherence or correctness of the generated step. the highest scoring beams are then selected for further expansion, enabling the model to gradually construct a valid multi-step reasoning trace. ## theoretical guarantees background and notation. at each decoding step, which we define as a block of τ tokens, a width of n beams is maintained. for beam i, let p i denote its partial reward after the first τ tokens and f i its final reward after completing the step. our preliminary results in figure 2 indicate that the final reward is related to the partial reward via a monotonic mapping with added noise: guarantee under noisy, nonlinear conditions. although the mapping between p i and f i need not be linear, we assume (i) the noise terms η i are independent and σ-sub-gaussian, and (ii) the expected partial scores preserve the ordering of the expected final scores. let the shared first τ tokens drive the entire covariance: as τ → l the correlation approaches 1 meaning the partial score is an almost perfect proxy, whereas as τ → 0 it vanishes. figure 4 shows that this τ /l trend, tightening toward 1 as τ increases, is also true empirically. for example, attaining ρ * = 0.8 demands τ ≥ 0.64 l. ## experiments we evaluate our method on three challenging math-reasoning benchmarks, math-500 (lightman et al., 2023b), sat-math from agieval @xcite , and aime 2024. for generation we use the instruct variants of two open-source llms, llama-3.2-3b (meta, 2024) and qwen-2.5-3b @xcite , selected for their strong reasoning ability at modest scale. process evaluation is performed with two prms of different capacities, mathshepherd-mistral-7b @xcite and skywork-prm-1.5b, allowing us to study the impact of early rejection for prms of different sizes. early rejection is triggered after a prefix of τ ∈ 32, 64, 128 tokens. these thresholds are motivated by preliminary analysis (figure 4 ), which shows that partial-reward scores at these lengths are already highly correlated with final rewards. at each decoding step we sample @xmath0, 8, 16, 32, 64 candidate beams and retain the top @xmath1, mirroring prior prm-guided search settings @xcite . we compare our early-rejection decoder with the conventional pipeline that scores only fully completed beams, reporting average answer accuracy and total inference flops for each run. all experiments are conducted on an hpc cluster, with each run executed using four nvidia a100 gpus (40 gb memory each). ## conclusion we demonstrate that prms can be effectively repurposed as partial reward models, enabling a single mid-generation evaluation to provide a reliable accept or reject signal. this allows weak beams to be pruned early, well before full reasoning steps are completed, thereby reducing unnecessary computation without sacrificing final accuracy. under mild noise assumptions, we provide theoretical guarantees showing that the probability of mistakenly discarding the optimal beam decays exponentially with prefix length, offering formal safety for early rejection. extensive experiments across sat-math, math-500, and aime confirm the practical benefits: early rejection reduces inference-time flops by 1.4×-9× when using a mid-sized prm (7b parameters), with no degradation in task performance. even with a smaller prm (1.5b), we observe 1.5×-4× compute savings, highlighting that lightweight evaluators are sufficient for effective and efficient reasoning. together, these findings establish early rejection as a simple, model-agnostic plug-in that narrows the gap between compute-heavy tree search and fast single-pass decoding, offering state-of-the-art compute efficiency without compromising solution quality.