Beyond Repeated Sampling: Learning Search Policies for LLM Reasoning
Ismail Labiad, Matthieu Kowalski, Marc Schoenauer, Rémi Munos, Julia Kempe
cs.CL, cs.AI
2026-09-23
A GRPO-trained 7B concept generator steers a frozen Qwen2.5-32B answer model, doubling pass@128 from 19.0% to 39.2% on hard DeepMath at the same 128-rollout budget.
The default way to spend extra test-time compute on hard reasoning is repeated sampling: draw many independent answers and hope one is right. That search only happens through local decoding noise, so you mostly get near-duplicate drafts rather than different ideas.
GuidedSampling (Handa et al., 2026) tried to fix this by sampling concepts first (theorems, strategies, hints) and then answering conditioned on them. This paper reruns that pipeline on MATH500 across all 25 Qwen2.5 size pairs from 1.5B to 32B. The original repeated-sampling baseline used a restrictive temperature 0.8 and top-p 0.5. With exploratory decoding (temperature 1.0, top-p 0.95, top-k off), concept guidance is slightly worse than the baseline on almost every pair; a 32B model prompting itself still drops 0.5 pass@50 points. The original loop also emits concepts one at a time, and Qwen averages about one concept per problem, so the promised diversity never shows up.
The real gap is therefore twofold: exploration has to happen at the semantic level, and the concept generator itself can be trained to serve a frozen, possibly closed-source answer model.
The inference protocol is rewritten first. The concept generator is asked, in a single trajectory, to analyze the problem and list at most ten non-duplicate, problem-specific concepts. A fixed answer budget B is then split as evenly as possible across those concepts. Both the baseline and the concept-conditioned answers use the same exploratory decoding, so any gain has to beat a strong sampler.
Only the concept generator is trained. The main setup is Qwen2.5-7B-Instruct as the concept generator and a frozen Qwen2.5-32B-Instruct as the answer model. Training data comes from DeepMath-103k, keeping problems where the answer model scores below 5% with 128 rollouts. Evaluation is 1k held-out problems where that model scored 0/128. Each problem yields G=8 concept trajectories, and each trajectory is scored with B=128 answer rollouts. Correctness is judged by the same 32B model with greedy decoding, because rule-based matching punishes equivalent forms.
Two trajectory-level rewards are compared:
Credit goes to the whole concept trajectory, so the policy is pushed to emit a set of directions at once. Optimization is the GRPO variant used in OLMo3: learning rate 1e-6, KL coefficient 0.001, clip 0.2/0.272, batch 64, 200 steps. Trajectories that parse to zero concepts get a -0.1 penalty.
Even without training, single-trajectory multi-concept sampling already helps on MATH subsets the answer model never solved in 100 rollouts. A 7B concept generator lifts a 32B answer model by +8.8 pass@50; the best pair is 32B concepts for a 7B solver at +9.7. The original iterative protocol, on the same hard subsets, tops out at +6.8 and still regresses on several pairs.
After RL, the DeepMath 1k hard set (answer model = 32B) looks like this:
| Method | pass@1 | pass@64 | pass@128 |
| Naive repeated sampling | 0.23 | 11.35 | 19.00 |
| Ten generic hints | 0.32 | 14.54 | 22.60 |
| Untuned 7B concept generator | 0.77 | 20.62 | 28.90 |
| Untuned 32B concept generator | 0.82 | 23.77 | 33.80 |
| Trained 7B, max-of-max | 1.48 | 26.66 | 35.30 |
| Trained 7B, max-of-mean | 2.18 | 29.64 | 39.20 |
max-of-mean is +20.2 over naive sampling and +10.3 over the untuned 7B, and it beats the untuned 32B generator at 33.8%. On out-of-distribution Omni-MATH 2 (1,382 problems filtered the same way), it reaches 18.60% pass@128 against 11.29% for naive sampling.
The policy also transfers. The 7B generator was trained only against Qwen 32B, yet steering Llama-3.3-70B, which it never saw, yields 34.3% pass@128 on the same DeepMath hard set. That beats Llama's own self-concepts (28.9%) and naive sampling (26.2%).
Ablations pin down the mechanism. Randomly swapping concept trajectories across problems drops pass@128 to 23.9%, barely above generic hints, so about 15 of the 20-point gain is problem-specific. Answer leakage is rare: 0.44% of concepts are flagged as containing the ground truth, and only 3% of problems have any flagged concept. Training the same 7B to answer directly with the same RL loop reaches only 24.7% pass@128, below even the untuned 7B steering the 32B. Holding the 128-rollout budget fixed while raising the concept count from 1 to 10 lifts pass@128 from 22.7% to 39.8%, while pass@1 stays almost flat (2.15-2.44%). Vendi Score on chain-of-thought embeddings, a proxy for the effective number of distinct rollouts, moves from 50.4 (naive) to 73.0 (max-of-mean).
For practitioners this is a small-model search policy in front of a frozen solver. Closed APIs and models too large to train can sit behind a cheap concept generator. One 7B concept trajectory costs about 0.24 of a single 32B answer, under 0.3% of a 128-rollout budget, so inference compute is essentially unchanged.
It also corrects GuidedSampling: those gains were tied to a baseline that was not allowed to explore, plus a protocol that barely produced concepts. What this paper adds is training the concept generator itself as an RL policy.
This is still incremental. The move is to lift exploration from token noise to semantic hints, and to show that a small policy can transfer. It is not a new reasoning paradigm.
Training is expensive. Scoring one problem at one step needs G x B = 1,024 frozen 32B generations plus 1,024 judge calls. A representative step is about 413 seconds, with 300s on answers and 85s on judging (about 94% of the step). The bottleneck is the model that is not being trained. Longer runs do not keep paying off: after about 200 steps both objectives sit in a noisy 0.38-0.40 pass@128 band, and max-of-mean later becomes unstable.
Rewards are trajectory-level, so weak concepts can free-ride on strong ones. Per-position accuracy stays in a narrow 1.3%-2.7% band, with no dead slots, but credit is still coarse. Rollouts are split uniformly across concepts; the appendix shows an oracle would dump the whole budget on the single best concept. No bandit allocator is trained.
Evaluation is almost entirely math, and the judge is the answer model itself, which invites self-bias. The "0%" hard set is a finite-sample filter: a fresh seed of naive sampling still recovers about 19.5%, so the doubled denominator is not "unsolvable." Llama transfer reuses the Qwen-32B-filtered set rather than Llama's own worst problems. Concept quality has no human rating.