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

Sep 14, 2026

Supervised fine-tuning strengthens language models on mathematical reasoning tasks, yet prior work mainly boosts performance by expanding datasets through augmentation. Such approaches help most in standard single-round question answering, where models focus on forward reasoning. This paper proposes reflective augmentation, embedding explicit problem reflection into each training instance. The method trains models to consider alternative perspectives and use abstractions and analogies, improving both standard and reflective, multi-step scenarios. Extensive experiments confirm its effectiveness and complementarity to existing augmentation.

LanguageEnglish

CategoryResearch & Report

FormatPDF

Pages19

Reading Time48 min

AccessFree to download

Abstract

Introduction

Reflection and prior learning strategies

Reflective augmentation (RefAug)

Inference and training example

Learn Beyond The Answer: Training Language Models with Reflection for Mathematical Reasoning Zhihan Zhang􀀀1†, Tao Ge2 , Zhenwen Liang 1†, Wenhao Yu2 , Dian Yu2 , Mengzhao Jia 1†, Dong Yu2 , Meng Jiang 1 1University of Notre Dame 2Tencent AI Lab, Seattle [zzhang23@nd.edu](zzhang23@nd.edu) Abstract Supervised fine-tuning enhances the problemsolving 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.1 1 Introduction The ability to engage in step-by-step reasoning is pivotal for language models (LMs) to solve mathematical problems (Wei et al., 2022 ; Kojima et al., 2022) . Supervised fine-tuning, particularly on data with detailed reasoning paths, effectively advances the problem-solving performance of LMs (Fu et al., 2023 ; Yue et al., 2023) . To enlarge such benefits, most previous efforts focus on creating additional instances to augment model training (Luo et al., 2023a ; Yu et al., 2024 ; Mitra et al., 2024 ; Li et al., 2024a) . While these data expansion approaches allow LMs to handle a broader range of math problems by increasing the diversity of training data, † This work was done when Zhihan, Zhenwen, and Mengzhao were interns at Tencent AI Lab, Seattle. 1Code and data are available at [https://github.com/](https://github.com/)[ ](https://github.com/)ytyz1307zzh/RefAug. 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 (Liang et al., 2024 ; Wang et al., 2024a) . Similarly, the strategy in human learning is not always to practice an increasing number of problems (Rohrer and Taylor, 2006) . Instead of merely memorizing superficial solutions to more problems, it can be more advantageous to gain a deep understanding of the existing problems (Semerci, 2005) . Reflection, therefore, becomes an essential accompaniment to practice. Stacey et al. (1982) 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 (Kagan et al., 1964 ; Anderson and Fincham, 2014) . 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 mode

FAQs

What problem does the paper target in existing training methods for mathematical reasoning LMs?

Existing data augmentation mainly improves standard single-round QA settings and may not deepen understanding of each problem for reflective, multi-step scenarios.

How does reflective augmentation (RefAug) differ from conventional data augmentation?

RefAug appends a reflective section to the original answer within each training instance, shifting focus from augmenting instance sets to enriching the training sequence.

What capabilities does RefAug aim to improve during training?

It encourages the model to review alternative perspectives and engage with abstractions and analogies to foster deeper reflective reasoning.

Sep 14, 2026

Uploaded by Sophia Brooks