Improving Math Reasoning through Value-Guided Informative Search
NeurIPS MATH-AI Workshop, 2026 Accepted
This work on search-augmented agentic reinforcement learning was conducted during my Applied Scientist internship at Amazon (June–September 2026) under the mentorship of Yuning Wu.
Reinforcement learning with verifiable rewards (RLVR) improves mathematical reasoning, but greater search-induced trajectory diversity does not by itself guarantee policy improvement. We introduce APIVIS, a training-time framework that adapts finite-budget Gumbel search to chunk-level mathematical reasoning.
APIVIS mixes direct and searched responses within each rollout group, making search improvements visible through informative relative rewards. It also selectively supervises search-improved tokens, preserving a learning signal when uniform group rewards make GRPO ineffective.
We establish policy-improvement guarantees for exact value-guided selection and corresponding approximate guarantees under bounded value-estimation error. Experiments across mathematical reasoning benchmarks and model scales demonstrate improvements over competitive search-based methods.
Accepted to MATH-AI: The 6th Workshop on Mathematical Reasoning and AI at NeurIPS 2026.