Learn Beyond The Answer: Training Language Models with Reflection for Mathematical Reasoning (2024)

Abstract: Supervised fine-tuning enhances the problem-solving abilities of language models across various mathematical reasoning tasks. To maximize such benefits, existing research focuses on *broadening* the training set with various data augmentation techniques, which is effective for standard single-round question-answering settings. Our work introduces a novel technique aimed at cultivating a *deeper* understanding of the training problems at hand, enhancing performance not only in standard settings but also in more complex scenarios that require reflective thinking. Specifically, we propose **reflective augmentation**, a method that embeds problem reflection into each training instance. It trains the model to consider alternative perspectives and engage with abstractions and analogies, thereby fostering a thorough comprehension through reflective reasoning. Extensive experiments validate the achievement of our aim, underscoring the unique advantages of our method and its complementary nature relative to existing augmentation techniques.

## introduction the ability to engage in step-by-step reasoning is pivotal for language models (lms) to solve mathematical problems @xcite . supervised fine-tuning, particularly on data with detailed reasoning paths, effectively advances the problem-solving performance of lms @xcite . to enlarge such benefits, most previous efforts focus on creating additional instances to augment model training @xcite @xcite . while these data expansion approaches allow lms to handle a broader range of math problems by increasing the diversity of training data, stacking more training instances does not necessarily lead to a deeper understanding of each problem. moreover, the scope of resulting models is confined to single-round question-answering (qa) settings that primarily require basic forward reasoning skills. consequently, these methods provide limited benefits for more complex reflective reasoning scenarios that involve reviewing past steps for further reasoning, such as addressing follow-up questions, correcting errors, or leveraging external feedback @xcite . similarly, the strategy in human learning is not always to practice an increasing number of problems @xcite . instead of merely memorizing superficial solutions to more problems, it can be more advantageous to gain a deep understanding of the existing problems @xcite . reflection, therefore, becomes an essential accompaniment to practice. @xcite define reflection as "to review thoughtfully, consider alternatives and follow extensions", which encourages learners to contemplate their previous actions to engage in deeper reasoning, thereby fostering reflective thinking capabilities @xcite . inspired by such human cognition, we propose a novel training strategy for lms that integrates reflection into each math problem. unlike traditional data expansion methods which operate on the instance dimension by adding more training examples (see figures 1b & 1c ), our approach targets a complementary direction, i.e., the sequence dimension of the training data. we introduce reflective augmentation (refaug), which appends a reflective section to the original answer of each training instance, advancing model learning beyond mere answer generation (see figure 1d ). such a design not only strengthens the model's understanding of the associated knowledge and methodologies in training problems, but also maintains the inference efficiency as the model ceases generation before corresponding training sequences are shown in an (input, output) format, where augmented parts are in red. training question: solve the equation 𝑥 -99 𝑥 -101 = 8. let 𝑡 = 𝑥 -100 . then the equation becomes 𝑡 -1 (𝑡 + 1) = 8 ⇒ 𝑡 ! -1 = 8. therefore, 𝑡 = ±3, and accordingly we get 𝑥 = 97 or 103. let 𝑡 = 𝑥 -100. … [omitted, same as standard training] … we get 𝑥 = 97 or 103. ## approach refaug extends each training sequence with a reflective section that encourages the lm to reflect on its initial reasoning process to engage in further math reasoning. figure 1 contrasts refaug with traditional augmentation methods, and its detailed implementation is elaborated below. reflection types following the definition by @xcite to "review thoughtfully, consider alternatives and follow extensions", we consider two types of reflection in composing the reflective section: alternative reasoning and followup reasoning. alternative reasoning involves thinking about the problem from different perspectives @xcite . therefore, besides the initial solution, we annotate an alternative approach that also effectively solves the problem. this helps the model master related methodologies and develop critical thinking skills. follow-up reasoning associates the initial solution to a broader class of problems @xcite . to fit various contexts, we consider two options: abstraction and analogy. abstraction refers to creating a generalized form of the original problem, thereby encouraging the model to reduce dependency on specific numerical values. analogy challenges the model in applying methodologies of solving the original problem to a more complex situation. learning to design follow-up scenarios enables the model to understand the associated math concepts and principles better and apply them flexibly in new contexts. the relationship between the initial instance and components of the reflective section is illustrated in figure 3 . ## experiments we test refaug in a variety of mathematical tasks that cover both standard single-round qa and reflective reasoning scenarios. we mainly evaluate two aspects: the influence of refaug on lms' math reasoning abilities and its interaction with existing augmentation techniques. besides, we extend our approach to code generation tasks and perform comprehensive analyses. ## conclusion this paper proposed reflective augmentation (re-faug) for math reasoning, a method that incorporates reflection into training problems and is complementary to existing data augmentation approaches. we proved the efficacy of refaug in not only enhancing lms' basic problem-solving skills on single-round math problems but also in cultivating their capabilities to solve more complex reflective reasoning tasks. we further verified the effectiveness of refaug in code generation tasks and its scalability, along with ablation studies and analyses of the methodological choices, such as the impact of data sequencing and the stability of the annotation process.