Self-Explore: Enhancing Mathematical Reasoning in Language Models with Fine-grained Rewards (2024)

Abstract: Training on large amounts of rationales (i.e., CoT Fine-tuning) has been found effective for improving mathematical reasoning of large language models (LLMs). However, acquiring human-authored solutions or augmenting rationales from proprietary models is costly and not scalable. In this paper, we study the problem of whether LLMs could self-improve mathematical reasoning capabilities. To this end, we propose Self-Explore, where the LLM is tasked to explore the first wrong step (i.e., the first pit) within the rationale and use such signals as fine-grained rewards for further improvement. On the GSM8K and MATH test set, Self-Explore achieves 11.57% and 2.89% improvement on average across three LLMs compared to supervised fine-tuning (SFT). Our code is available here]9.

## introduction recent works have shown that large language models (llms) can solve complex reasoning tasks with chain-of-thought (cot) prompting, which involves generating a rationale before its final prediction @xcite @xcite . such ability is especially evident for mathematical reasoning, where many times precise reasoning over multiple steps is required to reach the correct answer @xcite . meanwhile, relatively smaller models have shown limited performance, and thus prior works have focused on augmenting rationales from proprietary llms and distilling them to smaller models @xcite @xcite @xcite @xcite . however, acquiring high-quality solutions remains challenging. for humans, hand-crafting detailed step-by-step rationale annotations is timeconsuming and costly @xcite . on the other hand, using close-sourced models through apis incurs high expenses and distillation-based methods are inherently limited by the performance of their teacher model, which acts as the upper bound @xcite . hence, such strategies are limited in advancing frontier models @xcite . one potential solution to address this issue is to improve general capability of llms through self-training @xcite @xcite . inspired by prior works that focus on aligning llms to user preferences through self-training, we propose self-explore, a training method designed to self-improve the mathematical reasoning capabilities of llms by extracting granular learning signals from its own generated rationales. specifically, a target model conducts step-level exploration to identify the first wrong step (i.e., first pit) within each rationale by sampling multiple continuations. then, we construct a pair-wise dataset by sorting the rationales into positive and negative samples at a step level. finally, by applying an arbitrary preference learning objective (e.g., direct preference optimization (dpo) @xcite ) on a step-level, we relatively increase the probability of generating positive rationales and lower the probability of generating negative ones in a fine-grained manner. through experiments, we find that self-explore constantly improves the performance across different base models (mistral-7b @xcite , llemma-7b @xcite , and deepseek-math 7b @xcite ) without any distillation from proprietary models. for each model, we observe a 13.19%, 10.23%, and 11.30% improvement on gsm8k @xcite and a 1.98%, 3.16%, and 3.54% improvement on math @xcite compared to supervised fine-tuning (sft). also, we find that constructing a pair-wise dataset in a granular man- d pair target model c 1 ŷ correct ŷ incorrect problem c 2 c 3 c 4 ... therefore, it will take her 120/2.66 = <<120/2.66=44>>44 hours to read 120 pages. the answer is 44 ner based on a step-by-step basis (i.e., identifying the first pit) outperforms a naive approach of constructing based on the correctness of the final prediction, leading to a 3.64% and 2.76% margin on the gsm8k and math dataset, respectively. ## preliminaries given a language model π θ and a dataset d, self-training algorithms comprise of two stages: (1) dataset growth, where the dataset d is augmented with a π θ 's generations, and (2) policy improvement, where the pre-trained model improves human-alignment through preference learning followed by supervised fine-tuning @xcite @xcite . here, we describe two relevant methods that are employed in our framework for self-training. ## conclusion in this paper, we propose self-explore where llms can self-improve from given initial human curated dataset using fine-grained supervision. by utilizing automatic self-exploratory annotation, self-explore effectively integrates the roles of the annotator, target, and reward models into a single system. on mathematical reasoning datasets gsm8k and math, our method outperforms traditional supervised fine-tuning (sft) method by 11.57% and 2.89% in average across three different models, respectively. furthermore, we demonstrate that our method introduces minimal computation overhead (see appendix h). we hope our work could motivate future works to explore self-training methods that could more robustly generalize to a broader reasoning space across various domains, with ends of advancing the frontier of llm reasoning.