Source-linked AI summary

Mitigating Over-Optimization in PRM-Guided Search in Mathematical Reasoning by Optimizing the Guide

Taejong Joo, Diego Klabjan

arXiv:2608.30051v1cs.AIcs.LG

TL;DR

PRM-guided search can over-optimize noisy step-level rewards, increasingly favoring non-viable prefixes as reasoning depth grows. The paper formulates search as robust optimization over plausible reward perturbations and derives a training-free maximin selector; it reports 17–35% average gains and wins in 14 of 16 settings without verifier adaptation.

  • Problem

    PRM-guided search can over-optimize imperfect step-level rewards, pruning viable trajectories and expanding spurious ones as noisy-score effects accumulate with reasoning depth.

  • Method

    The paper formulates PRM-guided search as robust optimization under plausible reward perturbations and derives a training-free maximin selection procedure.

  • Results

    17–35% average improvement is reported for maximin search, which outperforms outcome- and step-level baselines in 14 out of 16 settings without fine-tuning or online adaptation.

  • Takeaways & Limitations

    Scaling inference-time reasoning requires search algorithms that remain reliable when learned evaluators are uncertain, not only stronger evaluators.

  • Takeaways & Limitations

    The practical calibration of the RKHS error radius B is delicate and relies on calibration data with predicted and reference scores when available.

Abstract

from arXiv · show

Process reward models (PRMs) provide dense step-level guidance for search-based reasoning, enabling inference-time compute to be allocated toward promising partial solutions. However, recent evidence suggests that PRM-guided search can over-optimize imperfect process rewards, pruning viable trajectories while expanding spurious ones. In this work, we theoretically show that directly leveraging PRM score is vulnerable to verifier noise through an extreme-value effect: non-viable prefixes become more likely to receive spuriously high scores as reasoning depth increase. Therefore, we formulate the PRM-guided search as a robust optimization problem over plausible reward perturbations, termed maximin PRM-guided search, leading to a training-free robust process supervision method that preserves promising alternatives when step-level scores are noisy. Maximin PRM-guided search mitigates this failure mode by reducing sensitivity to over-optimized PRM outliers. Without fine-tuning or online adaptation, maximin search consistently improves the PRM-guided search by 17-35\% on average, outperforming outcome- and step-level baselines in 14 out of 16 settings. Our source code is available at https://github.com/tjoo512/maximin-search.

1 Introduction

Inference-time search can improve reasoning by allocating computation across partial trajectories, but PRM-guided search is vulnerable to noisy scores that over-optimize spurious prefixes. The paper responds with robust optimization and a training-free selector that preserves competitive alternatives under reward uncertainty.

  • Inference-time computation can improve reasoning through sampling, search, and verification over explicit intermediate states.Tree of Thoughts is cited as an example of searching partial reasoning states and backtracking from unpromising ones.
  • PRMs provide dense step-level guidance by scoring intermediate reasoning steps, helping search allocate limited computation across partial trajectories.Unlike outcome reward models, PRMs evaluate intermediate reasoning rather than only final solutions.
  • PRM errors can falsely prune promising branches or waste computation on spurious trajectories, producing over-optimization of imperfect process rewards.The cited failure arises from learned reward estimates that may degrade under limited supervision, misspecification, or distribution shift.
  • The robust selector keeps m prefixes balancing high PRM scores with diversity under an RBF kernel before repeatedly expanding the selected prefixes.This selection process is intended to reduce sensitivity to noisy PRM scores.
  • The paper formulates PRM-guided search as robust optimization over plausible reward perturbations rather than treating learned scores as exact.This decision rule preserves trajectories that remain competitive under reward uncertainty.
  • 17–35% average improvement is reported for maximin search, which outperforms outcome- and step-level baselines in 14 out of 16 settings without fine-tuning or online adaptation.The method is presented as a training-free plug-in procedure for PRM-guided decoding.

2 Related Work

Prior work uses learned evaluators to guide step-level reasoning search and studies over-optimization mainly in outcome-level reranking. This paper focuses instead on a robust inference-time treatment of PRM misspecification in sequential search and pruning.

  • Step-level search methods use generators and learned evaluators to allocate computation selectively across candidate reasoning chains.Examples include self-evaluation-guided beam search, AlphaZero-style tree search, step-level beam search, and actor–critic decoding.
  • Related self-improvement methods address reward misspecification through iterative policy, value-model, or reward updates.The cited approaches include ReST-MCTS*, V-STaR, and PRIME.
  • The paper studies robust inference with a fixed generator and fixed PRM instead of retraining or test-time adaptation.It presents this as a comparatively overlooked perspective on PRM-guided mathematical reasoning.
  • Outcome-level over-optimization research shows that increasing optimization pressure against a misspecified reward can eventually induce reward hacking and degrade the true objective.This literature primarily concerns Best-of-N decoding and related reranking methods.
  • PRM-guided tree search can underperform Best-of-N because inaccurate step-value estimates compound with depth and distribution shift.In step-level guidance, reward errors affect both final selection and allocation of search compute across intermediate trajectories.
  • The work claims the first mathematical treatment of over-optimization in PRM-guided step-level search.Its analysis formalizes how PRM misspecification interacts with sequential search and pruning decisions.

3 Background

The paper models reasoning as trajectories generated step by step by a policy, with search expanding candidate prefixes and a PRM ranking them for retention or pruning. Beam search provides the central concrete setting, while the decomposition extends to broader search procedures.

  • A reasoning model generates a structured sequence of intermediate thoughts followed by a final answer for input x.The analysis considers inference on one input and represents trajectories as sequences of reasoning steps.
  • The thought-level policy π defines the conditional distribution of the next thought given a partial trajectory and induces a distribution over complete trajectories.The final output is obtained through a deterministic answer-extraction map.
  • PRM values estimate the future utility of continuing from an intermediate trajectory, while an oracle value function represents its utility.The oracle may encode ground-truth-answer probability or correctness of intermediate thoughts.
  • The intended role of PRM-guided search is to filter unpromising candidates so retained trajectories have higher expected oracle value than the full candidate set.The paper uses step-level beam search as the focus of its theoretical analysis and examines the gap under an imperfect verifier.
  • At each step, active trajectories are expanded by sampling or enumerating candidate next thoughts from the base model.The resulting candidates form the set considered for subsequent selection.
  • The PRM scores candidate prefixes, after which the next beam retains the top trajectories under that score.The framework reduces to Best-of-N when the branching factor is one and to step-level greedy decoding when beam width is one.
  • The same generator-and-evaluator decomposition extends beyond beam search to greedy best-first search and Monte Carlo tree search.In this decomposition, the generator proposes steps while the evaluator controls which partial trajectories receive further computation.

4 Closer look at an over-optimization in PRM-guided search

This section characterizes how imperfect PRM scores can cause greedy and beam search to select non-viable prefixes. The analysis shows that verifier noise, candidate-set imbalance, and repeated selection across depth compound the risk of irreversible failure.

  • Setup: PRM-guided search maximizes average estimated trajectory scores, but imperfect verifier estimates can select non-viable prefixes despite viable continuations.The setup models viable prefixes as those that can still reach the correct answer and analyzes selection under a noisy learned verifier.
  • Greedy over-optimization: Theorem 1 identifies a cutoff event where every viable prefix scores below a threshold while at least one non-viable prefix exceeds it, forcing erroneous selection.This mechanism formalizes greedy reward hacking under verifier noise.
  • Greedy over-optimization: Verifier over-optimization is driven by the upper tail of non-viable errors after maximization, especially when weak generators produce few viable and many non-viable candidates.A larger non-viable population creates more opportunities for an extreme positive error, while fewer viable prefixes reduce the beam’s defense.
  • Verifier uncertainty: The failure bound worsens as the effective signal-to-noise ratio decreases, because larger verifier uncertainty makes deterministic separation less protective.The effective gap Γt(Ct) measures a conservative separation scale, while σt controls stochastic uncertainty.
  • Generator coverage: Increasing candidate width improves generator coverage but also creates more non-viable opportunities for verifier hacking, so larger test-time compute may not resolve over-optimization.Theorem 2 separates failure into missing viable continuations and verifier selection of non-viable candidates.
  • Error accumulation across reasoning depth: Depth compounds local errors multiplicatively because one irreversible coverage or selection failure prevents trajectory-level success; beam width mitigates but does not remove this risk.For width wt, at least wt non-viable outliers must exceed the viable frontier, but the path-level accumulation argument remains.
  • Implications for more sophisticated guided search: More sophisticated search can improve separation or reduce verifier uncertainty, yet residual overvaluation remains amplified by maximization and repeated selection.The qualitative failure structure persists whenever non-viable prefixes retain a nonzero chance of being overvalued.

5 Mitigating Over-Optimization in Step-Level Guided Decoding

The method replaces pure PRM score maximization with robust, diversity-aware optimization that limits sensitivity to structured verifier errors and preserves alternatives. It combines a kernel penalty with approximate submodular optimization and calibration-based radius selection.

  • Robust optimization retains high-scoring trajectories that remain competitive under plausible reward perturbations, rather than treating PRM scores as exact.
  • The method penalizes pairwise trajectory similarity, discouraging beam collapse onto regions where correlated verifier errors can compound.The RBF kernel is larger for trajectories close in representation space, so the penalty favors geometric diversity.
  • Under RKHS-bounded verifier errors, the kernel regularizer equals the worst-case aggregate estimation-error radius for a selected subset.Theorem 4 identifies the regularizer as the worst-case error over smooth admissible error functions.
  • The robust objective trades off estimated reward against coherent local overestimation, preventing a spuriously high cluster from dominating the selected beam.Figure 2 illustrates how diversity can recover a higher-value point from a broad true peak when the estimate creates a spurious sharp-region peak.
  • The RKHS radius is delicate to estimate, so the practical method uses a calibration-based scaling rule when predicted and reference scores are available.The radius is treated as a hyperparameter and approximated through a minimum-norm interpolating function.
  • The algorithm uses a variational surrogate with randomized greedy selection and bounded pairwise-swap refinement to obtain an approximate solution with controlled error.The first stage has a provable approximation bound, while the second stage reaches a 1-swap local optimum; refinement is capped at 30 steps.

6 Experiments

Experiments evaluate maximin search across mathematical benchmarks, matched inference budgets, reward noise, and algorithmic design choices. Maximin generally improves PRM-guided search and is more robust to noisy scores, while adaptive uncertainty and local improvement contribute to performance.

  • Benchmark evaluation: Experiments compare maximin search with unguided, outcome-level, and step-level baselines across GSM8K, MATH-500, AIME’24, and AIME’25.Evaluation uses matched budgets of N ∈ {16, 64} trajectories, with three random seeds for the 30-sample AIME datasets.
  • Benchmark results: 17%–35% average improvement over SBS and best or tied-best performance in 14 of 16 generator–dataset–budget blocks demonstrate the benefit of maximin aggregation.At N = 16, runtime overhead increases by only 0.27%.
  • Benchmark results: 70.2% versus 66.4% on MATH-500 at N = 64 shows maximin’s gain over SBS with Phi-3.5-mini-instruct.On AIME’24 at N = 64, maximin increases accuracy from 6.7% to 13.3%.
  • Reward-noise robustness: At ε = 1.0, maximin degrades 1.4% from its noise-free baseline, compared with 12.7% for SBS.At ε = 3.0, degradation is 5.4% for maximin versus 16.9% for SBS.
  • Algorithmic design choices: Removing local improvement reduces accuracy from 79% to 77%, supporting its contribution to optimizing the robust objective.The ablation contrasts further robust-objective optimization with naive PRM-score optimization.
  • Algorithmic design choices: Fixed RKHS radii yield accuracies of 71%, 78%, 77%, and 74% for B̃ ∈ [0.03, 0.1, 0.3, 1], while the best radius depends on the PRM, generator, and dataset.The estimated radius also varies across reasoning depth and problem instance.
  • Algorithmic design choices: Maximin improves over beam search by +8.0, +4.6, and +4.0 percentage points at temperatures 0.3, 0.7, and 1.0, respectively.The largest gain occurs at the lowest temperature.

7 Conclusion

The conclusion identifies over-optimization of imperfect PRM estimates as an extreme-value failure mode and addresses it with training-free robust selection. The resulting approach improves guided search without verifier training or test-time adaptation.

  • Conclusion: Noisy PRM scores can increasingly favor non-viable prefixes with reasoning depth, pruning promising trajectories and wasting inference-time compute.The conclusion frames this as an extreme-value failure mode of imperfect step-level reward estimates.
  • Conclusion: Maximin search formulates process supervision as robust optimization under plausible reward perturbations and improves guided search without verifier training or test-time adaptation.The method is presented as a training-free robust selection procedure for PRM-guided decoding.
  • Conclusion: Scaling inference-time reasoning requires search algorithms that remain reliable when learned evaluators are uncertain, alongside stronger learned evaluators.

A Proof of claims

The proofs characterize how noisy verifier scores can cause greedy search to select non-viable prefixes and how such failures accumulate across generated reasoning depth. They also establish a worst-case aggregate error bound under RKHS smoothness.

  • Theorem 1: Theorem 1 establishes a deterministic-cutoff bound for greedy reward hacking under the stated assumptions.
  • Proof mechanism: The proof lower-bounds failure by intersecting events in which no viable continuation appears and a non-viable prefix receives a sufficiently high verifier score.Independence of the relevant noise variables allows the event probabilities to be combined.
  • Theorem 2: Theorem 2 combines candidate-generation assumptions with one-step verifier failure analysis to bound greedy selection failure when viable and non-viable continuations coexist.The proof conditions on the number of viable candidates and uses independent candidate generation.
  • Theorem 2: Theorem 2’s bound applies across mixed candidate sets, with zero verifier-selection failure when all candidates are viable.

A.3 Proof of Corollary 3

The corollary extends one-step greedy failure across the reasoning horizon, while the RKHS result bounds aggregate estimation error for bounded-norm perturbations.

  • Proof of Corollary 3: Corollary 3 extends the one-step failure analysis across reasoning depth under Markovian generation.The proof recursively applies the conditional survival bound and takes complements to obtain the cumulative failure lower bound.
  • RKHS error bound: Theorem 4 provides a worst-case aggregate error bound when estimation error is induced by an RKHS function with norm at most B.The proof uses the reproducing property followed by Cauchy–Schwarz.

A.5 Proof of Theorem 5

Theorem 5 establishes an end-to-end approximation guarantee for the algorithm’s output under its stated assumptions. The proof combines the grid approximation with the randomized greedy and post-processing guarantees.

  • A.5 Proof of Theorem 5: Theorem 5 gives an end-to-end approximation guarantee for the output q† relative to the optimal solution q⋆.The theorem compares the algorithm’s final output with q⋆ under the stated assumptions.
  • A.5 Proof of Theorem 5: The proof uses a multiplicative grid to select a grid point ηt⋆ close to the optimal η⋆.The grid ensures η⋆≤ηt⋆≤(1+ξ)η⋆.
  • A.5 Proof of Theorem 5: The inner solver’s randomized greedy guarantee and objective-improving post-processing are applied to the candidate returned at ηt⋆.These steps provide the intermediate inequalities used in the final guarantee.
  • A.5 Proof of Theorem 5: Combining the intermediate inequalities and taking expectations yields the claimed end-to-end bound.

B Additional theoretical results

The additional results characterize when PRM-guided beam search can lose viable prefixes and establish guarantees for shifted surrogate objectives. They also show that nonnegativity shifts affect multiplicative guarantees even though implementation need not compute the shift.

  • Beam-search failure: The failure probability is analyzed using the distribution of non-viable prefixes exceeding a viable-prefix cutoff.Independence of non-viable noises enables a binomial stochastic-dominance argument.
  • Beam-search failure: Beam search can fail when at least wt non-viable prefixes receive verifier scores no lower than every viable prefix.This condition can fill the entire top-wt beam with non-viable prefixes.
  • Maximin correction: For a failed PRM beam satisfying the stated edge-gap conditions, the maximin objective can prefer a replacement set over the failed beam.The result shows that the failed beam is not a maximizer of Jt under those conditions.
  • Shifted nonnegative surrogates: Corollary 8 extends the end-to-end guarantee to shifted nonnegative surrogates when the original surrogate objectives may be negative.The shifted objective is analyzed over candidate sets containing at most m elements.
  • Shifted nonnegative surrogates: The multiplicative guarantee gains an additional additive term (1−ρ)LH because shifting changes the approximation scale.The term disappears only when the unshifted surrogates are already nonnegative, and the implementation does not need to compute Lη.

C Comparison with MBR and diversity-based decoding

Maximin search uses candidate geometry to robustify intermediate pruning against verifier misspecification, unlike MBR’s final-output consensus and generic diversity methods’ redundancy control.

  • Comparison with MBR: Maximin search regularizes intermediate pruning, whereas MBR selects a representative completed output.Maximin preserves partial prefixes for future expansion rather than choosing one central final candidate.
  • Comparison with MBR: MBR can favor a dominant incorrect reasoning mode when many sampled chains share the same misconception.Centrality is useful only when it correlates with correctness.
  • Comparison with MBR: Maximin preserves optionality by avoiding concentration of the search frontier in regions where the verifier may be coherently wrong.Its target is robustness of the pruning step under verifier uncertainty.
  • Comparison with diversity-based decoding: Generic diversity methods discourage redundant candidates, but their diversity terms are not derived from a verifier-error model.Diverse beam search and DPP methods typically use lexical, semantic, volume, or near-orthogonality criteria.
  • Robustness interpretation: The maximin penalty represents worst-case coherent verifier error and uses λ as an uncertainty radius rather than merely a diversity weight.The similarity kernel K measures relationships between prefixes, and the penalty bounds aggregate misspecification error.

D Experimental details

The experiments use fixed generation, search, kernel, prompting, and optimization settings across mathematical reasoning benchmarks. AIME results are reported with averages and corresponding standard deviations.

  • Kernel construction: PRM embeddings use the generator’s last hidden state, with a kernel based on normalized embedding distance.The kernel bandwidth is determined from the median pairwise distance.
  • Kernel construction: The median distance and σ2 are clamped at 10^-6 for numerical stability.
  • Generation and search settings: Generation uses top-p sampling with p=0.9, temperature 0.7, maximum step length 128 tokens, and maximum reasoning depth 30.
  • Generation and search settings: PRM-guided search uses branching factor 4, maintains m parent trajectories, and halts trajectories beyond 2,048 generated tokens.The maximum number of trajectories is 4m.
  • Optimization settings: The inner solver uses randomized, deterministic, and random initializations followed by up to 30 one-swap local-search passes.The best fixed-η surrogate value is retained among the resulting local optima.
  • Evaluation reporting: AIME benchmark results are averaged across three random seeds, with their standard deviations reported separately in Table 2.Table 2 reports standard deviations corresponding to the AIME results in Table 1.
Loading 2608.30051v1…