MARGE: Improving Math Reasoning for LLMs with Guided Exploration

TL;DR

This work tackles the data-efficiency bottleneck and spurious correlations that limit mathematical reasoning in LLMs by introducing MARGE, a guided-exploration framework. MARGE decomposes reasoning into intermediate states and uses hit-guided exploration to complete prefixes from a guide solution, estimating state values via Monte Carlo and updating policies through RL or DPO without external value models. Empirical results across multiple backbones and benchmarks (e.g., MATH, GSM8k, CollegeMath, OlympiadBench) show substantial gains in both single-shot accuracy and exploration diversity (pass@64), with ablations confirming the value of guided exploration and data-generation strategies. The approach scales self-generated training data more effectively, yielding improved reasoning capabilities and broader solution repertoires, and opens avenues for applying guided-exploration principles to other reasoning-intensive tasks.

Abstract

Large Language Models (LLMs) exhibit strong potential in mathematical reasoning, yet their effectiveness is often limited by a shortage of high-quality queries. This limitation necessitates scaling up computational responses through self-generated data, yet current methods struggle due to spurious correlated data caused by ineffective exploration across all reasoning stages. To address such challenge, we introduce \textbf{MARGE}: Improving \textbf{Ma}th \textbf{R}easoning with \textbf{G}uided \textbf{E}xploration, a novel method to address this issue and enhance mathematical reasoning through hit-guided exploration. MARGE systematically explores intermediate reasoning states derived from self-generated solutions, enabling adequate exploration and improved credit assignment throughout the reasoning process. Through extensive experiments across multiple backbone models and benchmarks, we demonstrate that MARGE significantly improves reasoning capabilities without requiring external annotations or training additional value models. Notably, MARGE improves both single-shot accuracy and exploration diversity, mitigating a common trade-off in alignment methods. These results demonstrate MARGE's effectiveness in enhancing mathematical reasoning capabilities and unlocking the potential of scaling self-generated training data. Our code and models are available at \href{https://github.com/georgao35/MARGE}{this link}.

Figures (8)

• Figure 1: The y-axis represents the accuracy on MATH500, and the x-axis represents the number of self-generated responses for training Qwen2-7B-Instruct. The dots and the star show when models start to converge, and the dashed line exhibits their scaling trends as the generation amount for training increases. The top-right zone is preferred, as we can easily scale generation and achieve better performance when queries are limited. By improving the exploration process, MARGE enables the scaling of self-generated responses for training and improves reasoning ability. We discuss more about the scaling trend in Appx. \ref{['Appendix:scaling_trends']}.
• Figure 2: Overview of our method MARGE, which includes four stages. (a): Generate multiple responses from the current policy $\pi^{(i)}$ as candidates for guidance and judge their correctness. Starting from the second iteration, we can directly leverage sampled responses from stage (c). (b): Among all candidate solutions, we select one for each query as guidance according to Sec. \ref{['method:exploration']}. (c): Perform a continuation of all states in the guide solution to complete the exploration (Sec.\ref{['method:exploration']}) and value estimation (Sec.\ref{['method:estimation']}). The collected data is utilized in stage (d) for training and stage (a) in the next iteration as well. (d): Having fully explored the state space in (c), we first form the rollout buffer, then optimize the current policy, and finally acquire the policy $\pi^{(i+1)}$ for the next iteration as described in Sec. \ref{['method:improvement']}.
• Figure 3: Average accuracies when starting from different intermediate states of correct solutions (blue) and incorrect ones (red) with Qwen2-7B-Instruct. A larger state index indicates being closer to the end. On average, completing from a correct (incorrect) state increases the portion of correct (incorrect) answers, which boosts the exploration of more training data.
• Figure 4: Pass@1 (solid) and pass@64 (shaded) of different methods on the MATH500 test set. Pass@64 indicates the ability to explore multiple reasoning paths. The figure displays the improvement pass@64 over pass@1 in the shaded area, symbolizing the models' ability to explore. When MARGE enhances both pass@$1$ and pass@$k$, the performance gap is larger as $k$ grows. It demonstrates that, instead of trading off exploration abilities for pass@$1$ improvement, MARGE also enables a better exploration process than baselines and thus fundamentally improves exploration ability.
• Figure 5: (a): The change of training dataset entropy with number of responses sampled. Our method continues to find new useful pairs when generating more responses. As completing intermediate states does not yield full responses, we convert them based on the number of tokens. (b): Hit-guided explored data improves average state value at every reasoning step compared to vanilla exploration, in particular later ones. The values are estimated over $32$ Monte Carlo simulations.

Theorems & Definitions (7)

• Proposition 3.1
• proof
• Proposition 3.2
• proof
• Proposition 3.3
• proof
• Remark 3.4