Intermediate Fine-Tuning Improves Mathematical Reasoning in Smaller Models

Authors: neeraj  gangwar
suma p bhat
nickvash  kani
Venue: NeurIPS_cc_2024_Workshop_MATH-AI
Published: 2024-09-24 17:26:21

Abstract: While large models pre-trained on high-quality data exhibit excellent performance across various reasoning tasks, including mathematical reasoning (e.g. GSM8k, MultiArith), specializing smaller models in mathematical reasoning remains a challenging problem. A common research approach to address this challenge involves distilling knowledge from large pre-trained teacher models into smaller student models. Other techniques include augmenting datasets by rephrasing questions or using multiple views of solutions to improve reasoning performance. In this work, we explore intermediate fine-tuning and show that fine-tuning a model on an arithmetic dataset before fine-tuning it on a reasoning dataset helps improve the model's performance on the reasoning tasks. The arithmetic dataset can be generated programmatically, eliminating the resource-intensive task of dataset creation. We evaluate the impact of intermediate fine-tuning using the original GSM8k training set and an expanded GSM8k training set created through distillation. Our experiments on multiple datasets demonstrate that intermediate fine-tuning leads to average improvements of 6.3% and 14.2% in reasoning tasks using the original and distilled training sets, respectively, with greedy decoding compared to the models fine-tuned directly on these sets.

Main paper: 
 introduction: 
 Scaling the model and data sizes has had a tremendous effect on performance across various natural language processing (NLP) tasks (Chowdhery et al., 2023; Achiam et al., 2023; Touvron et al., 2023b; Jiang et al., 2023) . These pre-trained models can learn from a few demonstrations using in-context learning and do not require task-specific fine-tuning (Brown et al., 2020) . They also benefit from generating a sequence of reasoning steps before arriving at the final answer. These strategies have been particularly effective for mathematical reasoning (Wei et al., 2022c; Nye et al., 2022; Fu et al., 2022; Zhou et al., 2022) .While these large models exhibit excellent performance on mathematical reasoning tasks, adapting smaller models for these tasks remains an open problem (Wei et al., 2022b) . This problem is challenging because the math reasoning datasets, like GSM8k (Brown et al., 2020) , consist of a small number of reasoning problems, generally accompanied by one solution. These datasets do not contain sufficient training examples to capture the complexity of math reasoning. To overcome this issue, a widely explored research direction is to distill knowledge from large pre-trained teacher models into smaller student models. Some methods use questions from existing training datasets and use prompting to generate solutions for fine-tuning smaller models (Ho et al., 2023; Magister et al., 2023; Fu et al., 2023; Hsieh et al., 2023; Yue et al., 2024) . Others use various techniques to rephrase the questions to create more examples (Yu et al., 2024) or multiple views of solutions (Liang et al., 2024) to achieve better reasoning performance. This work focuses on using synthetically generated datasets to improve reasoning performance. We explore if fine-tuning a model on a related dataset before the reasoning datasets helps improve the model's reasoning abilities. Specifically, we programmatically generate a dataset with arithmetic tasks and fine-tune the models on this dataset before fine-tuning them on the reasoning datasets. In transfer learning literature, this is referred to as intermediate fine-tuning (Vu et al., 2020) or supplementary training (Phang et al., 2018) . Empirical observations using several mathematical datasets lead to the following key takeaways:• Models fine-tuned on the arithmetic dataset before a reasoning dataset perform better than the ones directly fine-tuned on the reasoning dataset. The arithmetic dataset can be generated programmatically, eliminating the need for manual resources.• Based on our observations with multiple datasets with varying mathematical reasoning tasks, we find that intermediate fine-tuning results in better out-of-domain generalization.• We evaluate our approach on datasets with relatively small and sufficiently large training sets. Our experiments show that intermediate fine-tuning improves performance in both cases.Our source code and datasets are publicly available.1  
related work: 
 Adapting a pre-trained model for a downstream task has been traditionally done through task-specific fine-tuning. However, this approach does not work for mathematical reasoning tasks because datasets, like GSM8k, do not contain enough examples to capture the complexity of mathematical reasoning. Several works have focused on distilling multi-step reasoning solutions from large teacher models to overcome this limitation. Fu et al. (2023) prompted Codex (Chen et al., 2021) to generate multiple multi-step solutions for the examples in the GSM8k training set and fine-tuned FlanT5 on the ones that led to the correct final answer. Hsieh et al. (2023) used PaLM-540B (Chowdhery et al., 2023) for generating solutions and fine-tuned T5 (Raffel et al., 2020) in a multi-task setting to generate the labels and rationale. Liu et al. (2023) used GPT-3.5-turbo to generate synthetic GSM8k-like examples. Yue et al. (2024) showed that a hybrid of chain-of-thought and program-of-thought solutions performed better than using either format individually and created math generalist models -MAmmoTH. In addition to using LLMs to generate more solutions, Yu et al. (2024) used LLM rephrasing and backward reasoning to augment questions and created a new dataset called MetaMathQA.Transfer learning has played a pivotal role in NLP. Vu et al. (2020) studied the effect of intermediate fine-tuning on the model's performance on a target task. Training on large multi-task mixtures is also a common trend in NLP (Aribandi et al., 2022; Wei et al., 2022a; Chung et al., 2024) . Another research direction explores identifying relevant examples for a given downstream task from a huge collection of datasets, like P3 (Sanh et al., 2021) . These methods create embeddings for all examples of interest using hidden states (Ivison et al., 2023) or gradients (Xia et al., 2024) . Given a task, a small subset of relevant examples are selected based on similarity. These methods have been mainly applied to data-efficient instruction-tuning. 
our approach: 
 In this work, we fine-tune a model on an intermediate task before specializing it in mathematical reasoning. This is also referred to as intermediate fine-tuning (defined below). In our approach, we use an arithmetic dataset for intermediate fine-tuning. The arithmetic dataset is generated programmatically, thus eliminating the resource-intensive task of dataset creation.Fine-tuning a model on an intermediate task before a downstream task can improve the model's performance on the said downstream task (Phang et al., 2018; Vu et al., 2020) . The downstream task is also referred to as the target task. This is called intermediate fine-tuning and can lead to successful Following prior work on transfer learning (Vu et al., 2020; Neyshabur et al., 2020) , while various mathematical datasets may help improve a model's reasoning performance, we focus on an arithmetic dataset in this work for two reasons. First, arithmetic computations are an integral part of mathematical reasoning. Second, while curating a math reasoning dataset requires considerable resources, a simple arithmetic dataset can be generated programmatically. We leave the exploration of other potential intermediate tasks for future work.We refer to Liu and Low (2023) to programmatically generate an arithmetic dataset. Their work has shown that LLaMA (Touvron et al., 2023a) fine-tuned on a programmatically generated dataset outperforms GPT-4 (Achiam et al., 2023) on arithmetic tasks. While the dataset from Liu and Low (2023) contains the basic arithmetic operations -addition, subtraction, multiplication, and division, we also include fractions and percentages. GSM8k does not require computations over large numbers, hence we limit the number of digits in the operands to seven. Furthermore, we use log-uniform sampling to ensure that the dataset is not skewed towards numbers with greater digits. This dataset contains ∼ 1.3M examples.Training. We use GSM8k for model specialization. As it does not have a validation set, we randomly sample 512 examples from the training set to create a validation set. We use two versions of GSM8k in this work.Original. In the first version, we use the remaining examples from the training set for model specialization. This dataset contains 6961 examples. We refer to this dataset as GSM8k (Orig.).Distilled. We generate a distilled dataset using the questions from GSM8k (Orig.) to evaluate if intermediate fine-tuning benefits tasks with large training datasets. This dataset is generated by prompting Mistral-7B (Jiang et al., 2023) using the prompt from Wei et al. (2022c) . We generate 64 solutions per question and keep the ones that lead to the correct final answer. After removing duplicate solutions, this results in a dataset with ∼175k examples. We refer to this dataset as GSM8k (Dist.).Evaluation We use the GSM8k test set to evaluate the models. We also use three additional datasets -MultiArith (Roy and Roth, 2015) , ASDiv (Miao et al., 2020) , and SVAMP (Patel et al., 2021) -to test out-of-domain generalization. MultiArith contains problems focused on basic arithmetic operations and is relatively simpler than GSM8k. ASDiv focuses on diverse language usage patterns and covers wide problem types taught in elementary school. SVAMP contains problems with varying structures to ensure that a model cannot solve the problems by applying simple heuristics and ignoring question text. 
training details: 
 For our experiments, we use FlanT5 (Chung et al., 2024) which is an instruction-tuned T5 (Raffel et al., 2020) . The base and large versions of FlanT5 are used with 250M and 750M parameters, respectively. We use the AdamW optimizer (Loshchilov and Hutter, 2017) with a learning rate of 10 -4 , a weight decay of 10 -4 , and an effective batch size of 128. For FlanT5-Large, a learning rate warmup of 500 steps is used. All model training processes are initialized using a fixed seed.The intermediate fine-tuning is performed for 10 epochs without validation, and checkpoints are saved every other epoch. To adapt these models for reasoning, we continue the training from these checkpoints on GSM8k. The models are fine-tuned for 20 and 100 epochs on GSM8k (Dist.) and GSM8k (Orig.), respectively. The best checkpoint is selected based on the GSM8k validation performance.We use greedy and consistency decoding at inference. For consistency decoding, nucleus sampling (Holtzman et al., 2019) is used with T = 0.6 and p = 0.9 to sample eight responses, and the most consistent final answer is chosen. As nucleus sampling is a stochastic decoding method, we repeat the evaluation with consistency decoding three times and report the mean accuracy.In-Domain Performance. We first evaluate the models on the GSM8k test set. Table 1 shows the accuracy (%) achieved by different models. We observe that both FlanT5-Base and FlanT5-Large benefit from intermediate fine-tuning, and the performance on GSM8k improves significantly. These results also show that intermediate fine-tuning helps with both GSM8k (Orig.), which has a small training set, and GSM8k (Dist.), which already has a sufficiently large training set.Out-of-Domain Performance. Next, the models fine-tuned on GSM8k are evaluated on MultiArith, ASDiv, and SVAMP, and the results are shown in Table 1 . These results indicate that intermediate fine-tuning does not harm out-of-domain generalization. The models fine-tuned on the arithmetic dataset first generalize better than those directly fine-tuned on GSM8k.Arithmetic vs GSM8k Performance. We use the checkpoints fine-tuned on the arithmetic dataset to initialize the models to be fine-tuned on GSM8k. But does good performance on arithmetic tasks always translate to better GSM8k performance? Our experiments show that this is not always the case. We use the BigBench Arithmetic benchmark (BIG-bench authors, 2023) to evaluate the model's arithmetic abilities and report micro-averaged accuracy across addition, subtraction, multiplication, and division. The results of this experiment are shown in Figure 1 . The models fine-tuned on the arithmetic dataset perform better on the BigBench Arithmetic benchmark as the training progresses (Figure 1a ). However, the models initialized from the checkpoints with better arithmetic performance do not always result in better GSM8k performance (Figure 1b ). These results agree with the findings of Neyshabur et al. (2020) which show, for image datasets, that pre-training performance is not always a faithful indicator of effective transfer learning. This makes it harder to decide when to stop the intermediate fine-tuning. However, it should be noted that all checkpoints from intermediate fine-tuning result in a better performance than directly fine-tuning the model on GSM8k. 
conclusion: 
 In this work, we examined intermediate fine-tuning for mathematical reasoning. Our experiments showed that fine-tuning a model on a programmatically generated arithmetic dataset before a reasoning dataset helped improve the model's performance on the reasoning tasks. We evaluated our approach with small and large datasets, and intermediate fine-tuning resulted in better performance in both cases. Moreover, intermediate fine-tuning did not harm out-of-domain generalization, instead, models fine-tuned on the arithmetic first showed better out-of-domain generalization. Finally, while this work does not offer a method for determining when to stop intermediate fine-tuning, initializing models from any checkpoints during the process yielded better results than fine-tuning them directly on GSM8k. 
limitations: 
 While intermediate fine-tuning improved a model's mathematical reasoning performance, good performance on the arithmetic tasks did not always lead to better reasoning performance. Due to this, deciding the maximum number of epochs for intermediate fine-tuning and selecting a checkpoint from this process to initialize the next model remain open problems. Furthermore, identifying an intermediate task or programmatically generating a dataset for it may not always be feasible, limiting the applicability of this approach. Other mathematical datasets for intermediate fine-tuning may also be explored. Finally, we evaluated our approach on GSM8k. The experiments presented in this work may be extended to include other mathematical reasoning datasets, like MATH (Hendrycks et al., 2021) . We leave these avenues for future research. 
 
 Appendix: https://github.com/neerajgangwar/reasoning-ift 



Reviews: review: title: title: Arithmetic training improves math reasoning performance for SLMs
main_review: main_review: The authors investigate the intermediate fine-tuning(ift) for math reasoning, and perform ift on T5 models to illustrate its effectiveness.

Pros:

- The experiments show that T5 models using ift outperform T5 without ift, with a clear gap.
- IFT does not hurt other OOD test set, and model’s arithmetic performance show strong correlation with math reasoning.

Cons:

- the statement of `OOD` is not very suitable, given the fact that MultiArith, SVAMP, and Asdiv basically contains the same questions with fewer reasoning steps compared with GSM8K.
- it would be better if the author could make some clarification about how to generate ift data, and what is the difference compared with the mentioned “GOAT” method. The current version looks like authors different applies the same GOAT method to generate arithmetic datasets.
- make fair comparison w.r.t FLOPs: the current experiments only consider model’s performance with or without IFT, however, the FLOPs of two methods are quite different. It’s better to train the baseline more epochs to match the FLOPs used in IFT.
title: title: Finetuning on arithmetic data is an interesting idea, but why does better arithmetic not result in better gsm8k scores?
main_review: main_review: The paper proposes to do an intermediate finetuning step on synthetically generated arithmetic dataset and show that it results in better performance on gsm8k than finetuning on that data directly.

This is an interesting idea and the results in Table 1 demonstrate significant gains both in and out-of-domain. Yet Figure 1 is really puzzling to me. What does the model learn during the intermediate finetuning if doing arithmetic better doesn't correspond to better gsm8k accuracy? I feel that without studying this more, this work is not fully conclusive. E.g., is it important to use arithmetic dataset or can some other datasets be used to achieve the same effects (e.g. some general instruction following data)? Also, does intermediate finetuning results in a model that makes fewer arithmetic mistakes on gsm8k validation/test set? You could try to include the exact calculations from gsm8k validation questions into the training dataset and see if at least those are being solved 100% correctly, since they become part of the training data during intermediate finetuning stage. If that doesn't happen (as Figure 1 potentially indicates), then I feel this needs more analysis and perhaps the take-away message that "arithmetic intermediate finetuning is helpful" is somewhat misleading as "arithmetic" might not even be necessary here.

Other than that, the paper is well written and presents an interesting idea that improves accuracy and efficiency of finetuning small, specialized models on mathematical reasoning tasks.
title: title: Review for paper 54: Intermediate Fine-Tuning Improves Mathematical Reasoning in Smaller Models
main_review: main_review: **Summary of contributions:** 

In this study the authors explore the challenge to tailor smaller models specifically for mathematical reasoning. To address this challenge, they delve into intermediate fine-tuning, a process where a model is fine-tuned on an arithmetic dataset before further fine-tuning on a reasoning dataset. The authors report intermediate fine-tuning step significantly improves the model’s performance on reasoning tasks.

**Relevance and Clarity:** 

Relevant to the guiding theme of the proposed workshop and the experimental results are explicitly mentioned and well explained for easy understanding.

**Originality:** 

The authors come up with a novel methodology for intermediate fine-tuning which boosts the performance of the fine-tuned model on the final target task and at the same time preserves prior knowledge of the model which is a rare phenomenon in the world of fine-tuned models.

**Strengths** 

1. *Improvement in reasoning tasks without loss of prior knowledge:*

The experiments involving intermediate fine-tuning reveal considerable improvements in performance on the final reasoning tasks for multiple diverse datasets thus proving the model’s robustness to varying reasoning tasks without loss of its prior knowledge.

2. *Generation of synthetic dataset:*

The authors’ endeavor to programmatically generate a dataset with arithmetic tasks overcomes the challenge of resource intensive and time-consuming dataset creation. Also, including fractions and percentages as well in addition to the basic arithmetic operations in the dataset is noted.

3. *Discussion of related work:*

All relevant works are discussed and properly cited within the main body of the paper, with clear relationship to submission.

**Weakness** 

1. As already reported by the authors, this research does not identify a methodology to determine the optimal stopping point for intermediate fine-tuning to ensure maximum accuracy on the final target task.

2. The research could have addressed one significant concern of fine-tuning which is fine-tuning can inherit biases present in the training data. Perhaps a better clarity on how the generated data for intermediate fine-tuning was tested for potential biases will make this research more engaging.

**Significance**

 This paper demonstrates the potential of intermediate fine-tuning on programmatically generated arithmetic dataset as a game changer to increase the mathematical reasoning ability of models and can serve as an alternative solution to knowledge distillation from large language models to smaller versions to achieve the same goal.
title: title: Official Review of Reviewer dCvS
main_review: main_review: The paper introduces an effective method for enhancing mathematical reasoning in smaller models via intermediate fine-tuning. This resource-efficient and scalable approach yields significant performance improvements. However, the focus solely on addressing arithmetic deficiencies lacks novelty. Further exploration of alternative tasks, checkpoint selection, and broader applications could enhance the paper's impact.

**Strengths:**

The proposed method is straightforward, and the authors conduct extensive experiments on GSM8K math reasoning tasks to validate its effectiveness.

**Weaknesses:**

1.	*Lack of Novelty:* Previous works have already tackled the issue of LLMs' poor performance on arithmetic operations, employing tools like programmatic thinking, symbolic languages, and fine-tuning to address these challenges.
2.	*Checkpoint Selection Issues:* The paper notes that strong performance on arithmetic tasks doesn't necessarily lead to improved reasoning performance. The absence of a clear method for selecting the optimal checkpoint during intermediate fine-tuning limits the approach's consistency.
3.	*Dataset Dependence:* The method heavily relies on the GSM8K dataset and a programmatically generated arithmetic dataset. Extending this approach to other reasoning tasks outside mathematics may require significant modifications or different intermediate datasets, which the paper does not address.
4.	*Unexplored Model Architectures:* The study does not investigate whether the advantages of intermediate fine-tuning are confined to specific model architectures like FlanT5. It would be valuable to determine if this technique generalizes across various architectures.