Improving Math Reasoning through Value-guided Informative Search

Shaohuai Liu; Yuning Wu; Haoran Liu; Enzo Jia; Devin Chen; Kai Wei

cs.AIOct 1, 2026

Improving Math Reasoning through Value-guided Informative Search

Authors: Shaohuai Liu, Yuning Wu, Haoran Liu, Enzo Jia, Devin Chen, Kai Wei

Organizations: Amazon · Department of Electrical and Computer Engineering, Texas A&M University

Abstract

Reinforcement learning with verifiable rewards (RLVR) has substantially improved the mathematical reasoning capabilities of large language models. Recent work introduces search into RLVR rollouts to increase trajectory diversity, but diversity alone does not ensure that the search-induced rollout policy improves upon the current policy. To address this gap, we propose APIVIS, a training-time framework that adapts finite-budget Gumbel search to chunk-level mathematical reasoning. APIVIS combines direct and searched responses within each rollout group, allowing improvements found by search to produce informative relative rewards. It further applies selective supervision to search-improved tokens, preserving a learning signal when uniform group rewards render GRPO ineffective. We show that exact value-guided selection improves the expected verifier reward at each searched state and that this guarantee extends to the complete rollout policy, with a corresponding approximate guarantee under bounded value-estimation error. Experiments on widely recognized mathematical reasoning benchmarks and different model scales demonstrate substantial improvements over competitive search-based methods, validating the effectiveness of APIVIS.

Figures & tables

Figure 1: APIVIS at a glance. (a) A frozen value model scores expanded chunk states while Gumbel sequential halving allocates the search budget and commits the highest-

Q

continuation. (b) Each RLVR group mixes three direct and three searched responses; return contrast trains GRPO, while improved searched tokens provide selective supervision. (c) Mean Avg@16 across MATH-500, AIME-2024, AIME-2025, and OlympiadBench, where APIVIS outperforms all baselines.

Method	MATH-500	AIME-24	AIME-25	Olympiad	Avg.
Qwen3-1.7B-Base	47.60	4.38	2.29	18.61	18.22
+ GRPO	63.50	6.25	3.12	28.16	25.26
+ LATR + GRPO	57.79	5.83	2.50	24.23	22.59
+ MCTS-alpha + GRPO	64.38	6.25	2.29	27.47	25.10
+ APIVIS + InfoSFT	66.17	9.58	4.58	30.80	27.79
Qwen3-4B-Base	45.44	4.58	4.17	21.68	18.97

Table 1: Validation-rank-1 Avg@16 verifier accuracy (%) at two model scales. Each trained row uses one training seed.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Category	Hyperparameter	Value
Training	Base policy	Qwen3-1.7B-Base or Qwen3-4B-Base
Training	Data	Open-RS, 7,000 source prompts; deterministic seed-42 split and order
Training	Seed / optimizer steps	$1$ / $100$
Training	Per-step batch	64 prompts; 6 completions per prompt (384 total)
Training	PPO minibatch / epochs	16 prompts / 1 epoch
Optimization	Actor learning rate	$1\times 10^{-6}$

Appendix

Table 2: Core hyperparameters shared by the Qwen3-1.7B-Base and Qwen3-4B-Base experiments. APIVIS settings refer to the reported fixed 3-direct/3-search C4/S8/K10 configuration.

Scale / method	Steps	Training hours	GPU-hours
4B GRPO	100	6.32	50.57
4B MCTS-alpha	100	18.24	145.92
4B Gumbel	100	19.19	153.50
4B APIVIS	100	20.06	160.51
4B LATR	100	38.64	309.15
4B TS-LLM	100	40.16	321.31

Appendix

Table 3: Logged training time and approximate eight-GPU cost. Summed time excludes downtime between resumed processes.

Figure 2: Qwen3-4B-Base training dynamics for the four methods in Table 1 : validation Avg@4, policy entropy, response length, and PPO KL. Faint traces show logged values, solid lines show centered running means over a 10-step radius, and bands show the local minimum–maximum range within that window.

Method	MATH-500	AIME-24	AIME-25	Olympiad	Avg.
MCTS-alpha + GRPO	78.55	12.92	10.62	43.06	36.29
Gumbel + GRPO	80.70	16.25	13.75	46.60	39.32
Gumbel + GRPO + InfoSFT	81.15	16.46	14.79	47.66	40.02

Appendix

Table 4: Avg@16 verifier accuracy (%) for fixed 3-direct/3-search runs. Each row is one training seed evaluated at its validation-rank-1 checkpoint.

$K$	MATH-500	AIME-24	AIME-25	Olympiad	Avg.
5	81.04	15.42	14.38	45.98	39.20
10	81.15	16.46	14.79	47.66	40.02
20	80.39	18.75	12.50	47.72	39.84

Appendix

Table 5: Avg@16 verifier accuracy (%) as the search-chunk budget

K

varies. Each setting uses one training seed and a best-validation checkpoint;

K=20

uses a validation-co-best checkpoint because the strict earliest tied checkpoint lacks a completed full evaluation.

Method	Avg.	Chunks	Step (s)	GPU-h
GRPO	36.82	–	227.57	50.57
Adaptive C4/S8	39.98	0.52	441.75	98.17
Adaptive C8/S16	40.06	1.07	883.91	196.42
Fixed C4/S8	40.02	3.91	722.30	160.51

Appendix

Table 6: Validation-rank-1 Avg@16 and full 100-step training telemetry for Qwen3-4B-Base. Chunks are searched chunks per designated search trajectory; GPU-hours use eight GPUs.

Figure 3: Per-step searched-minus-direct terminal reward for the reported APIVIS runs. Thin lines show raw training-step values, thick lines show centered five-step means, and dashed lines show run means. Positive values indicate that searched completions outperform direct completions within the same training rollout groups.

Explore similar work

Right Makes Might: Aligning Verified Hidden States Empowers RL Reasoning
Jun 2, 2026Ziyue Wang, Aomufei Yuan, Yongfu Zhu +10Reinforcement Learning With Verifiable Reward Mathematical Reasoning Benchmarks
1Peking University · 2JINGDONG · 3Shanghai Innovation Institute +2
LEEPS: Latent-Guided Explore-Exploit Prompt Sampling for Efficient RLVR in Large Language Models
Jul 30, 2026Shuang Liang, Haoyang Zhou, Yifan Gong +2Reinforcement Learning With Verifiable Reward LLM Reasoning Strategies
Renmin University of China · AntGroup
Enhanced LLM Reasoning by Optimizing Reward Functions with Search-Driven Reinforcement Learning
May 3, 2026Arash Ahmadi, Sarah Sharif, Yaser +1Large Language Model Reinforcement Learning LLM Reasoning Strategies
Mike · School of Electrical and Computer Engineering, University of Oklahoma, Norman, OK 73019 USA · Intelligent Neuromorphic and Quantum Understanding for Innovative Research and Engineering (INQUIRE) Laboratory, University of Oklahoma, Norman, OK, USA

Jun 2, 2026, cs.LGJ/K move · Enter open · S save

Right Makes Might: Aligning Verified Hidden States Empowers RL Reasoning

Ziyue Wang, Aomufei Yuan, Yongfu Zhu, Shuai Dong+9

1Peking University · 2JINGDONG · 3Shanghai Innovation Institute+2

Reinforcement Learning With Verifiable Reward Mathematical Reasoning Benchmarks Verifiable Rewards Value Alignment Stance

Open paper AskPDF arXiv

Reinforcement Learning from Verifiable Rewards (RLVR) has become the dominant approach for improving mathematical reasoning in large language models, yet current methods reduce each correct rollout to a single reward bit, ignoring the geometric structure shared among their hidden states. Investigating this structure, we find that at the anchor token (the position immediately before the answer marker), correct rollouts converge naturally because they must produce the same answer (cosine similarity ~0.84), yet each retains residual variance from its unique reasoning path. Encouraging full alignment at this point pushes the model to extract a unified "correct decision" representation, reducing sensitivity to which reasoning path was taken. Based on this observation, we propose Hidden-Align, an auxiliary loss function that aligns the last-layer hidden states of correct rollouts at the anchor token during RL training, with zero overhead in both training and inference. On eight mathematical reasoning benchmarks, Hidden-Align improves average pass@1 over the DAPO baseline by 3.8, 6.2, and 5.4 percentage points on Qwen3-1.7B, 4B, and 14B respectively, with consistent pass@k gains across all three scales, supported by ablations on loss type, anchor position, layer depth, and loss weight.

Abstract

Explore similar work

Right Makes Might: Aligning Verified Hidden States Empowers RL Reasoning

LEEPS: Latent-Guided Explore-Exploit Prompt Sampling for Efficient RLVR in Large Language Models

Enhanced LLM Reasoning by Optimizing Reward Functions with Search-Driven Reinforcement Learning