Beyond Repeated Sampling: Learning Search Policies for LLM Reasoning

• Repeated sampling explores through token-level decoding noise and often yields near-duplicate attempts, whereas concept-conditioned answers aim to diversify search at the semantic level.The paper frames effective exploration as a bottleneck for both reinforcement learning and test-time scaling.Source: §1 · p. 1
• The method generates detailed, problem-specific concepts in one trajectory and RL-trains a small concept generator using the downstream success of the frozen answer generator it steers.The reported implementation trains Qwen2.5-7B as a search policy for a frozen Qwen2.5-32B answer generator.Source: §1 · p. 2
• Across all 25 Qwen2.5 concept-generator and answer-generator size pairs on MATH500, GuidedSampling [Handa et al., 2026] gains largely disappear against exploration-enabled repeated sampling.The prior pipeline generated concepts one at a time in an iterative loop and averaged about one concept per problem, limiting diversity and wasting autoregressive generation.Source: §1 · pp. 1–2
• 39.2% pass@128 versus 19.0% for naive repeated sampling at the same answer-generation allocation on 1k held-out hard DeepMath problems.The comparison evaluates a reinforcement-learning-trained concept generator steering a frozen answer generator.Source: §1 · p. 2
• The trained 7B concept generator surpasses an untuned 32B concept generator, transfers without retraining to Llama-3.3-70B, and retains gains on Omni-MATH 2.Source: §1 · p. 2
• The paper analyzes whether the learned generator improves exploration rather than solving problems directly, and studies answer leakage, concept relevance, position effects, and concept count.Source: §1 · p. 2

Inference-Time Search and Concept-Guided Exploration

• Repeated sampling scales with sample count and commonly uses self-consistency or verifier selection, while related methods decompose reasoning into plans, hints, or solution sketches.Source: §2 · p. 2
• Unlike GuidedSampling, this method emits all concepts in one trajectory and trains the generator for downstream success; its reproduced advantage largely disappears when repeated sampling can explore.Source: §2 · p. 2
• Single-trajectory generation diversifies exploration while preserving autoregressive throughput, unlike Tree of Thoughts branching that disrupts continuous batching and KV-cache sharing in vLLM.Source: §2 · p. 2

RL-Driven Abstractions and Small-Model Orchestration

• Unlike RLAD, this method uses no teacher-trace SFT warm start, produces concise problem-specific hints instead of procedural summaries, and trains only the concept generator; like SOAR, it rewards downstream performance.The downstream answer model remains frozen.Source: §2 · pp. 2–3
• Unlike zero-shot Step-Back prompting and Auto-CoT, this method trains its concept generator with trajectory-level rewards from downstream answer success.Auto-CoT gains may stem from problem creation rather than the synthesized examples themselves.Source: §2 · p. 3
• Unlike orchestration systems that coordinate tools and frontier models, this method steers one fixed larger answer generator; a trained 7B concept generator improves a 32B answer generator.Source: §2 · p. 3

Background

• Concept-guided sampling first uses one LLM to generate concepts and a second to produce conditioned answers, aiming for more semantically diverse exploration than answer-only sampling (Handa et al., 2026).The section’s findings motivate the procedure introduced in Section 4.Source: §3 · pp. 3–4

Reproduction and Baseline Sensitivity

• The reproduction tests all 25 combinations of five Qwen2.5-Instruct sizes as concept and answer generators, matching 100 answer rollouts per problem and reporting pass@50.Concept-guided sampling uses the original iterative procedure and prompts; only answer-generator sampling parameters differ from the original protocol.Source: §3 · p. 3

Structural limitations

• The iterative procedure generates concepts one by one, although autoregressive generation can produce newer concepts conditioned on earlier ones within a single trajectory.The method also produces approximately one concept per problem on average, limiting its exploration of multiple solution ideas.Source: §3 · p. 3
• Even with 3.18 concepts per problem, Llama-3.2-3B-Instruct concept guidance reaches 86.7% pass@50, below the repeated-sampling baseline’s 88.1%.The −1.4 point change under exploratory sampling shows that the lost advantage is not explained solely by Qwen models’ low concept count.Source: §3 · p. 3

Reproduction and Baseline Sensitivity

• With exploratory decoding, concept guidance slightly underperforms repeated sampling across nearly all Qwen2.5 concept-generator/answer-generator pairs in the matched-compute comparison.Both conditions use 100 answer rollouts and report pass@50; Table 1 summarizes the baseline and concept-guided differences. The original baseline’s temperature 0.8 and top-p 0.5 restricted exploration.Source: §3 · p. 3, Table 1 · p. 4

Exploratory Answer Generation

• Concept-conditioned generation and repeated sampling both use temperature 1.0, top-p = 0.95, and top-k = −1, making the comparison conservative.The baseline is exploration-enabled, so concept guidance must outperform a strong repeated-sampling comparison.Source: §4 · p. 4

Single-Trajectory Concept Generation

• Single-trajectory generation prompts for high-signal, problem-specific, non-duplicative concepts, parses at most ten, and yields 4 to roughly 10 concepts per problem.Stronger generators saturate the ten-concept cap, compared with about one concept per problem in the original procedure.Source: §4 · pp. 4–5

Focusing on Hard Problems

• Each answer generator is tested on its own MATH subset of problems where repeated sampling gets 0% accuracy over 100 rollouts; subset sizes fall from 493 to 178 as model size grows.The subsets are constructed from the MATH test set and become smaller because larger answer generators leave fewer completely unsolved problems.Source: §4 · p. 4

Inference time results

• On model-specific hard MATH subsets, single-trajectory concepts improve nearly every CG/AG pair over repeated sampling by up to +9.7 pass@50 points, versus iterative gains up to +6.8.A 7B concept generator raises the 32B answer generator by +8.8 points, showing that a small generator can guide a much larger one.Source: §4 · pp. 4–5, Table 2 · p. 5

Cross-family check

• On 424 hard MATH problems, Llama-3.2-3B single-trajectory sampling raises pass@50 from 25.2% to 26.1% (+0.9 points), unlike the iterative procedure of Handa et al. (2026).It produces 9.23 concepts per problem versus one to three for the iterative procedure; that procedure degraded the exploratory baseline, while single-trajectory generation makes the effect positive.Source: §4 · p. 5

Motivation

• The answer generator remains frozen while the concept generator is trained, enabling optimization of larger models or API-only systems whose weights are unavailable.The paper describes this as a training handle for otherwise untrainable models.Source: §5 · pp. 5–6

Concept sampling and answer rollouts

• For each problem, the concept generator samples G trajectories, keeps at most 10 parsed concepts per trajectory, and distributes B answer-generator rollouts as evenly as possible across them.When B = 128 and M = 10, some concepts receive 13 rollouts and the rest receive 12.Source: §5 · p. 6

Reward aggregation

• An LLM judge grades each answer rollout for correctness, and the rollout scores are aggregated into a scalar reward for its concept trajectory.The judge is used because rule-based matching can penalize mathematically equivalent answers expressed in different forms or formats.Source: §5 · p. 6
• Max-of-max rewards a trajectory if any rollout is correct, whereas max-of-mean rewards its most accurate concept, distinguishing occasional success from reliably useful concepts.Max-of-max is binary in {0, 1}; max-of-mean is continuous in [0, 1].Source: §5 · p. 6

Trajectory-level credit

• The aggregate reward is assigned to the entire concept trajectory to encourage exploration over concept sets, though this may create position effects or let weak concepts share credit with strong ones.The paper studies these effects in Appendix D.2.Source: §5 · p. 6

Optimization

• The concept-generator policy is trained with a GRPO-style objective (Shao et al., 2024).The full objective and optimization details are provided in Section 6 and Appendix C.Source: §5 · p. 6

6.1 Setup

• Evaluation uses 1k disjoint DeepMath problems where the AG has 0% success in 128 rollouts, plus an equally filtered Omni-MATH 2 set excluded from training.The 7B CG is trained with GRPO-style reinforcement learning using eight concept trajectories per problem and downstream rewards from 128 AG rollouts.Source: §6.1 · p. 7

6.2 Results

• Max-of-mean reaches 39.2% pass@128 versus 19.0% for naive sampling on DeepMath, and 18.6% versus 11.3% on Omni-MATH 2.At pass@64 on DeepMath, the comparison is 29.6% versus 11.4%; Table 3 reports the headline allocations.Source: §6.2 · p. 8, Table 3 · p. 8
• Max-of-mean leads max-of-max at every DeepMath allocation, reaching 39.2% versus 35.3% pass@128, while generic hints reach only 22.6%.On Omni-MATH 2, max-of-max only matches untrained concept generators; the full curves and confidence intervals are reported in Appendix D.7.Source: §6.2 · p. 8
• The trained 7B CG reaches 39.2% pass@128 with max-of-mean, surpassing the untuned 32B CG’s 33.8% and its own max-of-max result of 35.3%.Paired bootstrap tests report a +5.4-point DeepMath and +3.3-point Omni-MATH 2 pass@128 gap over the untuned 32B CG, with both 95% intervals excluding 0.Source: §6.2 · p. 8
• RL training adds +10.3 points over the untuned 7B CG’s 28.9% pass@128, while trained concept guidance produces more diverse reasoning than naive sampling.The diversity comparison uses Vendi Score; qualitative examples show training reshaping concepts into more concrete, problem-specific plans.Source: §6.2 · pp. 8–9

7.1 Do Problem-Relevant Concepts Matter?

• 39.2% pass@128 with matched concepts exceeds 23.9% with mismatched concepts and 19.0% with naive sampling; mismatched concepts are near the 22.6% generic prompt-modification baseline.Problem relevance accounts for roughly 15 of the 20 percentage-point improvement over naive sampling.Source: §7.1 · p. 9

7.2 Transfer Across Datasets and Models

• 34.3% pass@128 on DeepMath with the transferred CG beats naive sampling at 26.2% and Llama self-generated concepts at 28.9%.The 7B CG was trained only against Qwen2.5-32B-Instruct; evaluation used common hard sets filtered using Qwen scores and Qwen correctness judgments.Source: Table 4 · p. 9, §7.2 · p. 9

7.3 Training and Inference Compute

• Under 0.3% of the B = 128 answer-rollout allocation is added by concept generation, making inference compute essentially the same as naive repeated sampling.A CG call costs about 0.24 of one AG rollout. In a representative 413-second training step, AG rollouts and judging take about 94% of compute; training is run only once.Source: §7.3 · pp. 9–10

7.4 Additional Ablations

• 39.8% pass@128 with ten concepts exceeds 22.7% with one concept under the fixed B = 128 rollout budget, while pass@1 remains nearly unchanged.Other ablations find direct training of the 7B model to answer performs worse than using it to guide the frozen 32B model; accuracy varies narrowly by trajectory position, whose positions 2–10 have broadly overlapping bootstrap intervals. Only 0.44% of concepts are flagged as containing the answer, and only 3% of problems have any flagged concept.Source: §7.4 · p. 10

8 Conclusion

• Roughly doubles pass@k at the same answer-rollout allocation, beats concepts from much larger untuned models, and transfers to an answer generator from a different model family.The learned policy uses downstream success from a frozen answer generator to train the small concept generator.Source: §8 · p. 10
• Prior concept-guided gains largely disappear against an exploratory baseline, motivating a procedure that generates multiple detailed concepts in one trajectory.The procedure trains a small concept generator using the downstream success of a frozen answer generator.Source: §8 · p. 10
• Training is compute heavy, rewards are assigned at the trajectory level, and evaluation focuses on mathematical reasoning.Future work should reduce reward-estimation compute, improve credit assignment, develop judge-robust evaluation and bandit-based rollout allocation, and extend the approach to broader domains.Source: §8 · p. 10

A Formal Setting and Analysis of the Reward Objectives

• Conditioning improves pass@k only when average per-concept yield exceeds the baseline; a below-average concept dilutes the allocation, while higher mean solve rate is sufficient but not the exact criterion.Each concept receives k/M rollouts under the balanced allocation; the exact improvement criterion is based on average yield.Source: A Formal Setting and Analysis of the… · p. 13
• The evaluation targets problems with weak baselines because pass@k has little headroom when the unconditional solve rate is high; uniform allocation is retained for throughput, with adaptive allocation left to future work.The analysis also notes that conditioning narrows the answer generator’s solution-space coverage, limiting potential gains when its unconditional coverage is already strong.Source: A Formal Setting and Analysis of the… · p. 14
• The max-of-max reward is unbiased for deployed passc@k but provides only a one-bit signal, whereas max-of-mean is more informative yet favors oracle allocation and can miss gains from combined moderate concepts.Max-of-mean retains O(log(k/M)) bits, but has upward bias from maximizing noisy estimates; the paper adopts it because its added informativeness outweighs that bias in the evaluated regime.Source: A Formal Setting and Analysis of the… · p. 14

B.1 Concept Generator Prompt

• The concept-generation prompt asks for problem analysis followed by 1–10 specific, high-signal, nonduplicated concepts or techniques, while forbidding a solution or final answer.Each concept is formatted on its own line beginning with ####.Source: §B.1 · pp. 15–16

B.2 Concept-Conditioned Answer Generator Prompt

• The conditioned answer prompt requests a step-by-step solution using one concept as a hint; the repeated-sampling baseline uses the same answer prompt without that hint.The answer prompt asks for the final answer in boxed form.Source: §B.2 · p. 16, §B.3 · p. 16

B.4 Generic Prompt Baseline

• The generic prompt baseline fills the concept slot with one of ten fixed, problem-agnostic hints instead of a problem-specific concept.The hints include trying another angle, breaking the problem into subproblems, and looking for symmetry or invariants.Source: §B.4 · p. 17

B.5 Judge Prompt

• An LLM judge returns a binary correctness verdict, while a separate leakage judge checks whether a concept states or reveals the ground-truth answer.The leakage judge treats methods, strategies, intuitions, reformulations, and partial setups that do not state the answer as nonleaking.Source: §B.5 · p. 17, §B.6 · p. 18

C Training Details

• The training loop is implemented on top of verl (Sheng et al., 2024).The appendix provides the full training and generation configuration and points to Section 6 for the setup summary.Source: C Training Details · p. 18

C.1 Training Objective

• The concept generator uses a token-mean, group-relative clipped policy-gradient objective with a KL penalty to its fixed pre-RL reference policy.The trajectory-level advantage comes from binary answer-generator correctness rewards aggregated by max-of-max or max-of-mean; the KL coefficient is β = 0.001, and advantages are not standard-deviation normalized.Source: §C.1 · pp. 18–19

C.2 Dataset Filtering

• Training uses 10,776 DeepMath-103k problems where Qwen2.5-32B-Instruct succeeds on at most 6 of 128 rollouts, while evaluation holds out 1,000 problems with zero successes.The training filter is below 5% success, preserving a learnable signal; evaluation problems are disjoint and filtered more strictly.Source: §C.2 · p. 19
• For out-of-distribution evaluation, the same filtering procedure selects 1,382 Omni-MATH-2 problems with zero successes in 128 rollouts.Omni-MATH-2 is itself a filtered version of Omni-MATH; these problems are used only for evaluation (Ballon et al., 2026; Gao et al., 2025).Source: §C.2 · p. 19

C.3 Concept Parsing

• Concept trajectories use #### delimiters and lenient regex parsing, retaining at most ten concepts and truncating excess response tokens on-policy.Text before the first delimiter is ignored, whitespace-only spans are dropped, and malformed text yields no matches rather than an error.Source: §C.3 · p. 19
• Empty parsed lists receive a −0.1 training reward and no answer rollouts, while validation substitutes one empty concept; duplicate concepts remain separate entries.Empty trajectories stay at roughly 0% during training, and concepts that exceed the remaining prompt allowance are truncated to fit.Source: §C.3 · p. 20

D Additional Results

• Additional empirical results include sanity checks, a direct-training comparison, full pass@k curves with confidence intervals, and analyses of concept position and count.These results were omitted from the main paper for space.Source: D Additional Results · p. 20

D.1 Training the Smaller Model Directly

• The trained 7B concept generator steering the 32B reaches 28.9% pass@128, already exceeding the trained 7B direct-answer model.The authors attribute the approach’s value to injecting exploration into a stronger frozen answer generator, rather than having the small model solve problems alone.Source: §D.1 · p. 20, Figure 2 · p. 20

D.2 Position Bias Under Trajectory-Level Rewards

• Per-position concept accuracy ranged from 1.3% at the first position to 2.7% at the fifth, with no position dead or dominant.Positions 2–10 had broadly comparable estimated accuracies and overlapping bootstrap confidence intervals, supporting trajectory-level credit.Source: §D.2 · pp. 20–21, Figure 3 · p. 21

D.3 Answer Leakage

• Only 0.44% of generated concepts were flagged as containing the ground-truth answer, and just 3% of held-out problems had any flagged concept.An LLM judge checked whether concepts stated or gave a mathematically equivalent value to the answer, avoiding exact-string matching.Source: §D.3 · p. 21

D.4 Effect of the Number of Concepts

• 39.80% pass@128 with ten concepts versus 22.70% with one shows that more concepts improve performance under a fixed 128-rollout allocation.The concept sets are nested, and gains remain positive through ten concepts; pass@1 stays within 2.15–2.44%.Source: §D.4 · pp. 21–22, Figure 4 · p. 21

D.5 Iterative Guided-Sampling Protocol on the Hard Subsets

• The iterative protocol’s largest pass@50 gain was +6.8 points, below the single-trajectory method’s +9.7, and some CG/AG pairs still regressed.The iterative protocol produced at most 1.5 concepts per problem on average for the 32B generator, compared with up to about 10 for the single-trajectory procedure.Source: §D.5 · p. 22, Table 6 · p. 22

D.6 Why Naive Resampling Is Nonzero on the “0%” Held-Out Set

• 19.46% unbiased and 19.47% Monte Carlo pass@128 estimates show that fresh naive sampling solves some problems in the filtered “0%” set.The filter used 128 answers, so zero observed successes estimates finite-sample performance rather than guaranteeing a true success probability of zero.Source: §D.6 · pp. 22–23, Table 7 · p. 23

D.7 Extended Pass@k Curves and Confidence Intervals

• Figure 5 shows full pass@k curves with 95% bootstrap confidence bands across rollout allocations for DeepMath and Omni-MATH 2; paired-bootstrap intervals quantify method gains over naive sampling.Intervals are estimated by resampling held-out problems; gain intervals use the same resampled problems for both methods.Source: §D.7 · pp. 23–24, Figure 5 · p. 24, Table 8 · p. 24
• At pass@128, trained max-of-mean gains +20.2 points on DeepMath and +7.3 points on Omni-MATH 2 OOD over naive sampling; generic prompt modification is indistinguishable from naive OOD.On DeepMath, every method significantly improves over naive at all reported allocations, with gains increasing from under two points at pass@1 to double digits at pass@128.Source: §D.7 · p. 23, Table 8 · p. 24
• The trained 7B max-of-mean concept generator beats the untuned 32B generator at every allocation on both datasets, by +5.4 points on DeepMath and +3.3 on Omni-MATH 2 at pass@128.The respective 95% paired-bootstrap intervals are [+2.4, +8.6] and [+1.2, +5.4].Source: §D.7 · p. 24, Table 9 · p. 25
• RL training significantly improves the 7B max-of-mean policy at every allocation on DeepMath and Omni-MATH 2, while max-of-max is indistinguishable from untuned 7B OOD.At pass@128, max-of-mean gains +10.3 points on DeepMath and +3.6 on Omni-MATH 2; max-of-max gains +6.4 points on DeepMath.Source: §D.7 · pp. 24–25, Table 10 · p. 25

D.8 Diversity of the Generated Reasoning

• Semantic diversity is measured on DeepMath reasoning-only chains of thought using Qwen3-Embedding-8B embeddings and Vendi Score, interpreted as the effective number of distinct rollouts.The problem statement and final answer are excluded from the embedded text; scores are computed per problem and averaged over the held-out set.Source: §D.8 · p. 25
• Figure 6 sweeps the RBF kernel bandwidth; methods are most separated at σ = 0.166, where the average-score spread reaches 22.7 points.This bandwidth is used for the Table 11 comparison.Source: §D.8 · p. 25, Figure 6 · p. 26
• At σ = 0.166, Vendi Score rises from naive (50.4) and generic perturbation (53.5) to untuned 7B/32B concepts (63.3/62.9), trained max-of-max (66.3), and max-of-mean (73.0).Max-of-mean is 45% higher than naive in effective distinct rollouts; this ordering mirrors the pass@k gains.Source: §D.8 · pp. 25–26, Table 11 · p. 26

D.9 Effect of Longer Training

• At 200 steps, max-of-mean reaches 0.386 pass@128 versus 0.353 for max-of-max; max-of-max catches up by about 300 steps.The paper attributes max-of-max’s slower early progress to its binary trajectory reward, which makes credit assignment harder than max-of-mean’s graded per-trajectory success rate.Source: §D.9 · p. 26
• 0.38–0.40 pass@128 is the noisy plateau reached after the first few hundred training steps, with no sustained gains from extending training to 750 or 900 steps.Figure 7 compares max-of-max through 750 steps and max-of-mean through 900 steps.Source: §D.9 · p. 26, Figure 7 · p. 27
• ≈0.28 pass@128 at step 750 marks max-of-mean’s sharp collapse, as truncation rises to ≈60% near step 720 amid entropy and KL growth.Over-long trajectories hit the generation-length limit before a stop token and score poorly; max-of-max stays within the plateau band through 750 steps, supporting the 200-step schedule as a compute-saving way to capture essentially all attainable gain.Source: §D.9 · pp. 26–27

D.10 Qualitative Examples: Concepts Before and After Training

• In three held-out examples—bounded ball distribution, repeated integrals, and elevator motion—trained max-of-mean concepts succeed where untuned 7B concepts fail.Ground-truth answers are shown for reference only and are not included in either generator’s prompts.Source: §D.10 · pp. 27–29
• Across the examples, trained concepts give concrete, problem-specific substitutions, constructions, or lines of attack, while untuned concepts tend toward generic strategies or theorem names without usable plans.Each problem presents three untuned concepts followed by three trained concepts.Source: §D.10 · p. 27
• For the constrained ball-distribution problem, a trained concept constructs the generating function as (1+x+x^2+x^3+x^4+x^5+x^6+x^7)^5.The concept specifies the one-juggler polynomial and raises it to the fifth power for five jugglers.Source: §D.10 · p. 28