Training Language Models with Reflection for Mathematical Reasoning

Reflective Augmentation: A Novel Approach to Enhancing Mathematical Reasoning in Language Models

Update: For a summarization version, click here: Reflection explained in 16 pictures (2-min read)

Introduction

Mathematical reasoning is a critical area for the advancement of language models (LMs). While traditional supervised fine-tuning on detailed reasoning paths can significantly enhance the problem-solving capabilities of LMs, the majority of existing research focuses on data augmentation techniques to broaden the training set. These methods, although effective, predominantly address single-round question-answering (QA) settings and often fall short when dealing with more complex reflective reasoning scenarios.

This whitepaper introduces a novel approach, termed as reflective augmentation (RefAug), designed to embed problem reflection into each training instance. The primary goal is to train models to consider alternative perspectives and engage with abstractions and analogies, ultimately fostering a deeper comprehension through reflective reasoning. This method not only aims to improve performance in standard settings but also equips models to handle more complex mathematical reasoning tasks that require reviewing past steps, addressing follow-up questions, correcting errors, or leveraging external feedback.

Reflective Augmentation: Enhancing Learning Through Reflection

Broadening Understanding Beyond Data Augmentation

Existing data augmentation methods, such as question augmentation (Q-Aug) and answer augmentation (A-Aug), primarily focus on increasing the number of training instances. While these strategies help broaden the range of math problems that a model can handle, they do not necessarily lead to a deeper understanding of each problem. Moreover, the resulting models' scope is often confined to single-round QA settings that mainly require basic forward reasoning skills. Such methods offer limited benefits for more complex reflective reasoning tasks that involve multi-step problem-solving abilities.

The Need for Reflective Reasoning

The principle of reflection is widely recognized in human learning as an essential component of effective education. Instead of merely practicing an increasing number of problems, developing a profound understanding of the problems at hand can be more advantageous. Inspired by this, RefAug harnesses the power of reflection to enhance the training process of LMs.

Reflective reasoning encourages learners to contemplate their previous actions and engage in deeper thinking. Stacy et al. (1982) define reflection as "to review thoughtfully, consider alternatives and follow extensions." This process helps learners identify and simplify the common components of variable expressions by introducing new variables, thereby reducing problem complexity.

Methodology: Integrating Reflection into Training

Alternative and Follow-Up Reasoning

RefAug introduces two types of reflective reasoning into the training data: alternative reasoning and follow-up reasoning. Alternative reasoning involves thinking about the problem from different perspectives and proposing an alternative approach to solve it. Follow-up reasoning associates the initial solution with a broader class of problems, either through abstraction or analogy, encouraging the model to generalize and apply learned methodologies to new contexts.

Data Annotation and Training Implementation

We employed expert models like GPT-4 turbo for data annotation to ensure high-quality reasoning paths with minimal human effort. The annotated reflective sections are incorporated into the training data, immediately following the initial answer to each question. During training, loss is calculated on tokens from both the initial answer and the reflective section, with the sequence trained to predict both the solution and the reflective reasoning components.

During inference, the generation process early stops upon delivering the answer to the input question, without including the reflective section in the final response. This maintains inference efficiency while leveraging the benefits of enhanced learning during training.

Experimental Validation

Standard Math Reasoning

Extensive experiments on various math reasoning tasks demonstrate that RefAug significantly boosts the problem-solving performance of LMs in standard single-round QA settings. The method provides an average accuracy gain of +7.2 over direct fine-tuning, proving its efficacy in enhancing model learning.

RefAug's benefits are complementary to traditional data expansion methods, leading to substantial performance improvements when combined. For instance, augmenting models with both Q-Aug and RefAug results in an overall accuracy improvement of +6.1, underscoring the synergistic advantages of integrating different augmentation techniques.

Moreover, RefAug remains effective even on large datasets, where it continues to elevate model performance by emphasizing deeper comprehension rather than mere memorization of problems.

Reflective Math Reasoning

In reflective math reasoning tasks, which require multi-step problem-solving abilities, RefAug outperforms traditional augmentation methods by a significant margin. The reflective sections help models develop robust reflective reasoning skills, enabling them to handle follow-up questions, correct errors, and make better use of external feedback.

The results highlight that RefAug fundamentally enhances models' capabilities to engage in deeper mathematical reasoning, making it an indispensable tool for developing advanced LMs.

Summary

RefAug addresses the limitations of traditional data augmentation methods by embedding reflective reasoning into training problems. This method fosters a deeper understanding of mathematical concepts and enhances models' problem-solving capabilities not only in standard QA tasks but also in complex reflective reasoning scenarios.

RefAug significantly improves both basic and reflective reasoning skills of LMs, providing substantial accuracy gains and deeper comprehension of mathematical problems. The method introduces additional complexity during the training process, which may extend training time and computational costs. RefAug can be applied to other domains beyond math reasoning, such as code generation, fostering the development of versatile and highly capable LMs.

Overall, RefAug presents a significant advancement in the field of language modeling, providing a robust approach to training models not only to solve problems but to understand them deeply and reflectively.

Whitepaper: Learn Beyond The Answer: Training Language Models with Reflection for Mathematical Reasoning