Masked Thought: Simply Masking Partial Reasoning Steps Can Improve Mathematical Reasoning Learning of Language Models (2024)

Abstract: In reasoning tasks, even a minor error can cascade into inaccurate results, leading to suboptimal performance of large language models in such domains. Earlier fine-tuning approaches sought to mitigate this by leveraging more precise supervisory signals from human labeling, larger models, or self-sampling, although at a high cost. Conversely, we develop a method that avoids external resources, relying instead on introducing perturbations to the input. Our training approach randomly masks certain tokens within the chain of thought, a technique we found to be particularly effective for reasoning tasks. When applied to fine-tuning with GSM8K on Llama-2-7B, this method achieved a 5% improvement in GSM8K accuracy and a 10% improvement in GSM-IC accuracy over standard supervised fine-tuning with a few codes modified. Furthermore, it is complementary to existing methods. When integrated with related explicit data augmentation methods, it leads to improvements across five datasets of various augmentation methods, as well as two different base models. We further investigate the mechanisms behind this improvement through case studies and quantitative analysis, suggesting that our approach may provide superior support for the model in capturing longdistance dependencies, especially those related to questions. This enhancement could deepen understanding of the premises in questions and prior steps. Our code is available at Github.1

heading: 1 Introduction
text: Despite the significant advancements Large Language Models (LLMs) have achieved across various tasks, they continue to encounter difficulties with multi-step reasoning problems (Lightman
* Corresponding authors: Xiting Wang, Rui Yan, and Yongbin Li.
1https://github.com/AlibabaResearch/DAMO-ConvAI, and https://github.com/ChangyuChen347/MaskedThought.
et al., 2023; Huang et al., 2023a; Shi et al., 2023; Chen et al., 2024). A primary challenge in reasoning arises from the fact that even a minor error can disrupt the entire solution process (Lightman et al., 2023). This issue is compounded in state-of-the-art models, which often suffer from hallucinations that lead to such errors (Huang et al., 2023b). Analyses of the error patterns in mathematical reasoning tasks performed by GPT-3 reveal that over 50% of errors stem from either a misunderstanding of the problem or from incoherent reasoning steps (Wei et al., 2022).
To improve the reasoning ability, most training methods focus on acquiring supervised signals obtained through costly means, such as human annotation (Liao et al., 2024), generation of larger models (Yu et al., 2023a; Luo et al., 2023; Liu et al., 2023; Yue et al., 2023), or self-sampling (Yuan et al., 2023; Singh et al., 2023; Wang et al., 2023b; Zelikman et al., 2022). These methods enable the models to learn specific trajectories by precisely outlining the reasoning processes. In contrast, our method illustrates that an entirely different approach is equally effective. Instead of requiring more precise guidance to direct the generative process, our technique achieves comparable results by incorporating random noise into the reasoning steps. For reasoning tasks, which are generally considered to require precise understanding, we find that this noise-based method is surprisingly effective, indicating that the model might have enhanced its performance through a denoising process.
We summarize our contributions as follows: First, we propose a simple and effective Masked thought Fine-Tuning (MFT) method for improving language model reasoning. The implementation of our method is simple: it requires only the substitution of specific tokens with a [mask] in the chain of thought. This is done while maintaining the
5872
same procedures as the standard Supervised FineTuning (SFT). With this minor modification, MFT produces results comparable to those of a more complex data generation pipeline (see Table 2), and it handles context distractors more effectively (see Table 7). Additionally, our method complements other data augmentation techniques using GPT-3.5 or GPT-4 (see Table 1), and the improvements generalize to out-of-domain reasoning datasets when the instruction set is diverse (see Table 3).
Second, we analyze our method from a regularization perspective and introduce two guiding principles for effective regularization. Specifically, we propose a noise injection framework in which our method and many commonly employed regularization methods are special cases, such as dropout (Srivastava et al., 2014) and synonym substitution (Zhang et al., 2015). We conduct a comparative analysis and find that these regularization methods are effective only when two guiding principles are met: (1) A portion of the tokens in the reasoning path should be retained without noise addition; (2) For positions where noise is added, it is crucial to ensure these positions maintain as little semantics as possible (see Sec. 3.3).
Third, we explore the underlying reasons for the effectiveness of our method by revealing its impact on learning dependencies. By conducting both quantitative analyses and case studies, we observed that this method demonstrates an enhanced dependency on the initial mathematical questions and earlier steps. This suggests a possible reason: Instead of depending solely on the local information generated by the model itself from recent steps, which can lead to a higher likelihood of hallucination or misdirection as the sequence extends, our method utilizes more information from the problem statement and earlier steps. These sources are more reliable and less prone to errors. Consequently, this strategy might reduce the risk of misunderstanding the problem and inconsistencies in reasoning (refer to Secs. 2.2 and 3.4).

heading: 2 Methodology

heading: 2.1 Masked thought Fine-Tuning
text: Our approach, named Masked thought Fine-Tuning,
falls under the umbrella of SFT. It maintains the
simplicity of standard SFT implementation. To fa-
cilitate comparison with other regularization meth-
ods, we present a general framework of token noise
injection. This framework considers both the posi-
tion and the type of noise introduced:
LMFT = − T∑
t=1
log p(wgtt |ws1:t−1)
wsi = { f(wgti , ni), if Mi = 1 wgti , otherwise
ni ∼ token noisy distribution N Mi ∼Bernoulli(p) for each target/source token wsi
(1) Similar to SFT, our method employs the unchanged wgt (the ground truth solution to question q) as the label. The distinction lies in our introduction of noise ni to the input. The implementation of f varies depending on the choice of token noise distribution N (see Tab. 5). We compare various noise settings and ultimately choose [mask] as our method design. The comparisons will be discussed in detail in the experiment section (Sec. 3.3).
(a) Both tokens are in the chain of thought. Supervised Fine-Tuning (SFT) exhibits a pronounced reliance on close tokens, whereas Masked thought Fine-Tuning (MFT) demonstrates a heightened dependence on tokens further along the sequence.
(b) The first token is in the mathematical question, and the second is in the chain of thought. MFT places greater emphasis on the mathematical question compared to SFT.

heading: 2.2 Enhancing Dependency Learning with MFT Regularization
text: Observation. To understand the underlying reason for the effectiveness of MFT, we conduct a detailed analysis and find that models tend to shift towards longer-distance dependencies, particularly enhancing dependency on questions as shown in Fig. 1. Specifically, we iterate over pairs of numeric tokens in a sequence and disturb the first token to see if the second token changes, indicating dependency (Schwab and Karlen, 2019). As illustrated, Fig. 1(a) compares the dependency between two tokens within a chain of thought, showing that SFT has a greater dependency on short distances, whereas MFT shifts towards longer-distance dependencies. In Fig. 1(b), we examine the effect of disturbances in questions on the chain of thought, observing that MFT enhances all dependencies of questions compared to SFT. Then, We check this pattern of other models and datasets in Appendix D. We also observe decreasing trends of short-distance dependencies within the chain. The degree of this decrease differs among various datasets. A clear pattern is that there’s a consistently increased dependence on questions. Distance Bias. SFT may cause the model to learn the distance bias from the pattern in ground truth data, where previous steps are correctly annotated, allowing the model to extract useful information from nearby thoughts and calculations, sometimes perhaps even overlooking the original question description. However, during inference, it generates steps on its own, moving away from the accurately verified ground truth. With each new step, the likelihood of errors increases. Even without obvious mistakes, inappropriate rephrasing or omissions of question details could mislead subsequent reasoning. This risk is heightened when the model excessively focuses on nearby tokens while neglecting the question. This phenomenon can also be referred to as exposure bias (Bengio et al., 2015), where no safe reasoning is guaranteed during inference. Dependency Regularization. MFT is helpful in addressing these dependency issues. During stepwise reasoning, an intermediate step might utilize multiple premises from previous steps and questions, forming a directed acyclic graph of dependencies. For human behavior, in the process of solving mathematical problems, we need to refer both to the initial conditions given in the problem and to many results obtained in previous steps frequently. Although Transformers allow for full connectivity between every token pair, enabling datadriven learning of such dependencies, they risk learning shortcuts through pattern-matching (Liu et al., 2022; Khona et al., 2024; Dziri et al., 2023). MFT randomly prunes the fully connected graph during training to mask potential shortcuts, encouraging the model to build up more connections to more robust features, such as information from questions and earlier, less error-prone steps. We maintain the mathematical question unmasked, assuming it is error-free, to prompt the model to comprehend the question while leveraging the provided premises. Empirical evidence also shows that masking parts of the questions do not improve performance.

heading: 3 Experiments

heading: 3.1 Datasets and Models
text: Training Dataset. In our main experiment, we utilized five training datasets covering math problems of varying difficulty levels and data quality. GSM8K: This dataset comprises 7,500 manually annotated grade school math problems (Cobbe et al., 2021). MATH: This dataset contains 7,500 math competition problems from high school, which are more challenging than those in GSM8K (Hendrycks et al., 2021). GSM8K-RFT: Yuan et al. (2023) generates additional reasoning paths for each query in GSM8K by utilizing an SFT model that was trained on the original dataset. The process further includes a step of filtering out those reasoning paths that lead to incorrect answers. GSM8K-RFT-X differentiates the sources of these reasoning paths. GSM8K-RFT100 sources from the SFT model itself, sampling 100 responses for each query, while GSM8K-RFTU33 integrates reasoning paths from multiple models of various sizes to create a larger dataset. MetaMathQA: According to Yu et al. (2023a), GPT-3.5-Turbo is employed to rephrase queries and responses in both MATH and GSM8K datasets and to introduce reverse engineering problems. MathInstruct: MAmmoTH (Yue et al., 2023) introduces a dataset named MathInstruct, which incorporates a variety of math problem datasets, including but not limited to GSM8K, MATH, AQUA (Ling et al., 2017), and Camel-Math (Li et al., 2023a). It also leverages GPT-4 to generate reasoning paths using both the Program-of-
Thought (POT) (Chen et al., 2022b) and Chain-ofThought (COT) methods.
We also conduct experiments on a non-reasoning dataset in Appendix I.
Evaluation Dataset. These datasets are mainly augmented from the training sets of GSM8K and MATH, and we test on the corresponding test sets. The GSM8K dataset sets aside 1,319 entries for testing, while the MATH dataset reserves 5,000 entries as a test set. Grade-School Math with Irrelevant Context (GSM-IC ) (Shi et al., 2023) is a variant of GSM8K testsets containing 58K noisy questions. We extract the final answers by specific templates (e.g., “the answer is”) in the predicted reasoning paths following (Yue et al., 2023; Yu et al., 2023a). Subsequently, we calculate accuracy as the ratio of predicted answers matching the ground truth. We test on the Out-of-Domain (OOD) science reasoning dataset following (Yue et al., 2024) to check the generalization ability.
Models. We use two types of base models: Llama2 and Mistral-7B. Llama-2 has a series of widely used base models, and we conduct experiments on
its Llama-2-7B 2 and Llama-2-13B 3. Mistral-7B 4 is a relatively newer base model with better performance than Llama-2-7B. The default experimental subjects are fine-tuned Llama-2-7B on the GSM8K training set unless otherwise noted in the captions. We pick the checkpoint obtained after a default number of training epochs (3 for SFT following previous works and 10/5 for MFT, as we observe it requires more training steps). For MFT, we default to 10 epochs for training sets smaller than 10K and 5 epochs for other larger datasets. Other hyperparameters are identical for SFT and MFT as listed in implementation details (Appendix C).

heading: 3.2 Main Results
text: Improvements Compared with SFT. The results presented in Tab. 1 illustrate that our approach improves performance no matter whether augmentation data is used, demonstrating effectiveness across datasets of varying quality and under two different base models. Notably, this enhancement is most significant in simpler or smaller datasets
2https://huggingface.co/meta-llama/Llama-2-7b-hf 3https://huggingface.co/meta-llama/Llama-2-13b-hf 4https://huggingface.co/mistralai/Mistral-7B-v0.1
where the available useful signals in the data may be limited. Our approach demonstrates enhancements in the performance of the MathInstruct dataset, which incorporates Program-of-Thought Data. This suggests that our method remains effective when including programming tasks.
Generalization Ability on OOD Reasoning Datasets. To verify whether the improvements brought by our method can generalize to other tasks, we first conduct validations on the singletask training set in Tab.1, i.e., training using GSM8K/MATH and testing on MATH/GSM8K. The experiments show that neither RFT nor MFT contributes to generalization when trained with GSM8K. Additionally, MATH appears to facilitate generalization more effectively compared to GSM8K. We also conduct comparative experiments with multi-source datasets on science reasoning tasks. We select four models trained with multisource data for comparison. The results in Table 3 indicate that models trained on MathInstruct data show relatively greater improvements compared
to those trained on MetaMathQA after MFT training. In contrast, when trained on MetaMathQA with Mistral-7B, there was even a slight decline in performance after MFT training. The difference between these two datasets is that MathInstruct includes 13 sources of mathematical reasoning data, whereas MetaMathQA only includes two datasets. Since MFT is still data-driven, it may require using multi-task data to generalize the knowledge learned from MFT to other reasoning tasks. Sample Efficiency. Within an equivalent number of training samples consumed, MFT yields a better performance than SFT. However, this does not imply that our approach requires fewer steps to converge. On the contrary, our method demands additional steps to reach a higher convergence value, as illustrated in Fig 2. This may be attributed to the exploration through masking, enabling us to gather new information from a limited dataset continually. Conversely, extending the training epochs for SFT introduces variability without yielding superior results. In our implementation details, we allow MFT to train for an extended number of epochs to ensure convergence of results as detailed in Tab. 9. Mask Ratio and Scheduling. We employ a conser-
vative linear warmup strategy over several epochs for stability. In Fig. 3, we investigate the influence of two major hyperparameters in MFT: mask ratio and warm-up duration. Our analysis reveals that incorporating a warm-up period contributes to enhanced final performance for a large mask ratio, while a small mask ratio without a warmup also works. Within the GSM8K dataset, varying mask ratios exhibit minimal sensitivity, as the problem-solving procedures within it lack dense formulas. Consequently, this dataset permits a relatively higher mask ratio. In our primary experiment, we set this ratio to 0.4. However, for datasets with more intricate formulas, such as MATH, we opt for smaller mask ratios, typically 0.1 or 0.2. As for the warm-up period, we adopt a conservative duration without excessive tuning of this hyperparameter. Generally, we allocate around two-thirds of the total training steps for warm-up.
Comparison with Self-Sampling. RFT (Yuan et al., 2023) collects the reasoning paths sampled by the SFT model itself and filters out those paths that contain incorrect answers5. Our experiments in Tab. 2 demonstrate that our method performs comparably to the explicit generation of reasoning paths, with even better outcomes when both approaches are applied concurrently. One primary factor of this comparable performance may be that RFT does not sample iteratively, yet the critical steps in the reasoning paths remain largely unchanged from the ground truth. The primary distinctions of sampled new paths lie in relatively minor semantic variances, without significantly
5LRFT = − ∑T
t=1 log p(wt|w1:t−1; q),where w1:T ∼ psft(.|q), and the outcome of w1:T is correct . w1:T is sampled from a fixed SFT model psft(.|q) and added to the original data for the second SFT stage.
increasing the new observed connections among local variables (Prystawski et al., 2023). It allows us to attain comparable effects by employing masked reasoning paths, where we sample new skip-connection by dropping the intermediate steps. On the other hand, if we compare from the perspective of addressing exposure bias, RL’s self-sampling provides careful supervision for the error-prone inference distribution, whereas MFT aims to directly prevent the model from learning the bias present in the ground truth. We provide an additional introduction about RL and RFT in Appendix B.
Despite the performance, for training efficiency, our model is trained on 7K data for 10 epochs, while RFT-100 requires training on 46K data for 3 epochs to achieve similar results. Notably, our approach achieves comparable performances without requiring extra machine computing resources.

heading: 3.3 Comparison with Other Regularization Techniques
text: There are inquiries into whether alternative approaches can yield comparable outcomes. Thus, we delve into different techniques for introducing noise and empirically evaluate their effectiveness. We explore the effects of adding noises to different positions, as well as the noise type introduced
Table 5: We compared MaskedLM (our main implementation of MFT) with other regularization techniques. The probability (p) for the sequence noise distribution is fixed at 0.4. The table outlines the framework of Masked thought Fine-Tuning, characterized by: 1. Masking partial tokens during reasoning steps as opposed to introducing noise across all tokens; 2. Prioritizing the minimization of information conveyed by noisy tokens, as indicated
by superior outcomes with empty and random strategies. The experiments are conducted on Llama-2-7b with GSM8K.
Sequence Noise Method Token Noise N Hyperparameter GSM8K - SFT - - 41.6
Mi = 1 for each token i
Dropout embi = embi ⊙Di, Di ∼ Bernoulli(q)
q=0.1 40.3 q=0.5 39.4 q=0.9 32.9
NEFTune embi += U(−αc , αc ) α=1 40.6 α=5 41.3 α=100 1.8
Mi = 1 for each target token i
Dropout embi = embi ⊙Di, Di ∼ Bernoulli(q)
q=0.1 39.9 q=0.5 43.1 q=0.9 40.4
Mi ∼ Bernoulli(p) for each source token i,
Dropout embi = embi ⊙Di, Di ∼ Bernoulli(q)
q=1,p=0.2 36.3 q=1,p=0.6 30.6
Mi ∼ Bernoulli(p) for each target token i,
NEFTune embi += U(−αc , αc ) α=100 42.0 α=1,000 47.2
Dropout embi = embi ⊙Di, Di ∼ Bernoulli(q)
q=0.1 39.0 q=0.95 42.9 q=1 46.5
Scheduled Sampling seqi ∼ softmax( Oi−1 τ )
τ=0.1 41.7 τ=1 41.1 τ=3 45.4 τ=10 47.0 τ=100 47.1
MaskedLM seqi = { [mask] prob m, w ∼ U(V ) prob r,
m=1,r=0 46.7 m=0,r=1 47.1
by NEFTune, Dropout, Scheduled Sampling, and MaskedLM (our primary implementation). The variables of compared methods are introduced in Tab. 4. Detailed descriptions of these methods are provided in the Appendix K. Results. Our experiments yield the following insights: 1. In terms of the placement of noise, we find that masking partial reasoning steps in the target yields the best performance when compared to introducing noise across all tokens or specifically within the question. 2. We observe that introducing more huge noise to remove the original semantic is advantageous. Specifically, setting the temperature τ = 100 for Scheduled Sampling and the noise scale a = 1000 for NEFTune enhances the model’s performance. In contrast, the strategy of replacing synonyms with τ = 1, 2 in Scheduled Sampling does not surpass the effectiveness of random re-
placement. When the model is trained with MaskedLM using r = 1, i.e., replacing with a random token, the model encounters tokens unrelated to the current context, leading it to disregard the semantics of these tokens, akin to the role of [mask]. In the experiments in Tab. 5, r = 1 yields slightly higher results than m = 1. While r = 1 demonstrates higher levels of randomness in comparison to m = 1, it has the potential to better mitigate overfitting. However, on average, the distinction between r = 1 and m = 1 is marginal for larger datasets. With the exception of using r = 1 for GSM8K, we default to utilizing m = 1 for other datasets in Tab. 1. We also explore the impact of adding attention masking at the masked positions within each transformer layer beyond just the input masking. We find that its effectiveness is on par
with that of using input masking alone. Another hypothesis is that the model could benefit from predicting multiple tokens ahead when doing MFT, so we conduct an additional comparison. We observe that even when solely computing losses on non-noisy tokens in MFT where no multiple-token prediction occurs, the performance still surpasses that of standard SFT (45.0, +3.4). Nonetheless, it exhibits lower performance compared to the computation of losses across all tokens (47.1, +5.5). Predicting the next tokens for masked positions poses a greater complexity than for unmasked ones, as they lack adjacent unmasked tokens from which to gather local information. We believe this increased complexity further improves the performance.

heading: 3.4 Investigating the Impact of MFT on Language Modeling
text: use correct premises even when misleading errors occur in the chain of thoughts. we observe that the dependency changes are consistent across different samples and present one as an example . In Fig. 4, we vary a digit within the context between 0 and 9 to see whether it influences the predictions of subsequent steps. If the model consistently predicts 16 − 3, this indicates that the model’s prediction is not influenced by the digit change but rather relies on the information provided in the question and prior steps, with the word “three” highlighted in orange. Conversely, if the model’s prediction shifts to 16 − n, it implies a dependency on immediate, adjacent tokens, suggesting a reliance on distance shortcuts. The frequency of 16− 3 predictions, as summarized in Tab. 6, reveals that MFT establishes longer dependencies to locate correct premises rather than relying only on recent intermediate results in the chain of thoughts.
Conversely, we also modify the premise within the question as shown in Fig. 5, where a prediction of 16−n indicates that, when faced with conflicting premises, the model prioritizes the question. The results of altering the question demonstrate that MFT excels in interpreting the questions, while SFT tends to ignore changes within the question.
Dealing with Distractors. Some previous studies have also indicated that adding some disturbance to questions, such as inserting an irrelevant context (Shi et al., 2023) or changing the order of
premises within the question (Chen et al., 2024), can decrease the reasoning ability of large language models. We test on the GSM-IC (Shi et al., 2023) dataset, which inserts some irrelevant distractor sentences into the original questions in GSM8K. The results in Tab. 7 show that although MFT and RFT are roughly comparable on the original dataset, MFT shows a greater improvement on the GSM-IC dataset. This suggests that there are differences in the mechanisms by which the two methods improve performance and indicates that MFT can effectively locate useful premises within the disturbed questions, even though MFT does not mask the questions during training.
False Positive Analysis. Despite the premises in these experiments being artificially disturbed, we would like to know if MFT might generate false positives due to the reduced use of neighboring features. False positives in this context refer to instances where intermediate steps are incorrect, yet the final result is accurately predicted. This analysis is conducted using GPT-46, which achieves an accuracy of 92% on the specified dataset. The findings, as detailed in Tab. 8, indicate that MFT does not result in a significant rise in the occurrence of false positives. MFT not only improves the accuracy of the final answer but also preserves the correctness of the intermediate steps.

heading: 4 Related work
text: Mathematical Reasoning of LLM. To improve the mathematical reasoning ability through supervised fine-tuning, the approach generally involves leveraging larger models or manual efforts to collect high-quality data (An et al., 2023; Yu et al., 2023a; Yue et al., 2023; Luo et al., 2023; Tao et al., 2024; Gou et al., 2023). MetaMath (Yu et al., 2023a) augments its dataset by rewriting mathematical queries and providing answers with the help of an LLM. MAmmoTH (Yue et al., 2023) compiles
6To help GPT-4 judge the correctness of solutions, we include the ground-truth reasoning steps in the prompts. We provide the prompt of GPT-4 and the cases of detected false positives in Appendix F.
diverse mathematical datasets for training, enriching them with Chains of Thought (CoT) and Programs of Thought (PoT) rationales for reasoning, resulting in the MathInstruct Dataset. RFT (Yuan et al., 2023) is a simplified offline method of standard online RL. It enhances its training dataset with accurate reasoning paths by generating and collecting them via an SFT model.
Data Augmentation for NLP. Noise-based strategies are widely used in the field of data augmentation for NLP. Common strategies include the integration of random continuous noise (Cheng et al., 2018; Jain et al., 2023; Kedzie and McKeown, 2019), token dropout (Xie et al., 2016; Gal and Ghahramani, 2015), token swapping (Wei and Zou, 2019), and more advanced mixup strategies (Guo et al., 2020a; Xie et al., 2022). We adopt the discrete token noise through token dropout (Xie et al., 2016; Gal and Ghahramani, 2015) to reasoning tasks, finding that it works well, which may complement the dependency-rich structure characteristic of reasoning data.
We provide a more thorough discussion of related works in Appendix A.

heading: 5 Conclusion
text: We propose Masked thought Fine-Tuning as a regularization technique during fine-tuning. Our findings indicate that MFT not only boosts the reasoning performance but also improves the model’s ability to capture long-distance dependencies. This study could pave the way for future research on regularization strategies for reasoning tasks and further exploration of the interplay between language dependency and reasoning capabilities.

heading: Acknowledgments
text: This work was supported by the National Natural Science Foundation of China (NSFC Grant No. 62122089), Beijing Outstanding Young Scientist Program NO. BJJWZYJH012019100020098, Intelligent Social Governance Platform, Major Innovation & Planning Interdisciplinary Platform for the "Double-First Class" Initiative, Renmin University of China, the Fundamental Research Funds for the Central Universities, the Research Funds of Renmin University of China, and Alibaba Group through Alibaba Innovative Research Program.
Limitations
The discussion regarding the root cause of the observed performance enhancement primarily stems from empirical experiments. We offer an explanation based on a notable phenomenon—the transfer of dependencies. However, we have not proven that dependency transfer is the sole factor in improving model performance; there may still be other influencing factors. This suggests that further empirical and theoretical research into the effective reasons behind the method’s success is warranted.
This method is effective for reasoning data while demonstrating no significant effect on a general instruction tuning dataset which lacks multi-step reasoning. We have only provided a qualitative explanation so far without proposing a method that includes quantitative metrics to determine which types of data are best suited for this approach.
The experiments indicate that the mask ratios that different datasets can accommodate may not be entirely the same. As for mixed data, this paper has not yet explored strategies for adaptively fitting mask ratios to different datasets. Additionally, we have not yet investigated whether this approach, when applied to mixed reasoning data and general instruction tuning data, affects the general capabilities or diminishes the observed enhancement of reasoning.

heading: A Related Work
text: A.1 Mathematical Reasoning of LLM
Pre-trained Large Language Models have demonstrated impressive reasoning capabilities and have shown corresponding improvements on various downstream tasks (Kojima et al., 2022; Huang et al., 2022; Yang et al., 2024; Wang et al., 2023c; Liu et al., 2024). To further enhance the capabilities of pre-trained Large Language Models (Touvron et al., 2023; Jiang et al., 2023; Achiam et al., 2023; Azerbayev et al., 2023) on reasoning tasks, the primary strategies involve fine-tuning and various training-free post-processing techniques. The post-processing techniques include approaches like prompting (Wei et al., 2022; Fu et al., 2022; Kojima et al., 2022; Wang et al., 2023a; Gao et al., 2024), ensembling (Wang et al., 2022b; Li et al., 2022), and reranking (Uesato et al., 2022; Cobbe et al., 2021). A unique technique specific to LLMs is to leverage the model’s reasoning capabilities through prompting. This entails crafting zero-shot prompts (Kojima et al., 2022; Wang et al., 2023a) or designing a small set of demonstrations to facilitate in-context learning (Wei et al., 2022; Fu et al., 2022). Additionally, there are methods focused on designing intricate pipelines utilizing various rules or algorithms for enhanced reasoning. These methods include search (Yao et al., 2023; Besta et al., 2023), iterative feedback collection (Yao et al., 2022), and employing the Program-of-Thought approach for generating code solutions (Chen et al., 2022b).
In terms of fine-tuning, whether through Supervised Fine Tuning (SFT) or Reinforcement Learning (RL), the primary aim is to gather high-quality supervisory signals for model improvement. For SFT, the approach generally involves leveraging larger models or manual efforts to collect high-quality data (An et al., 2023; Yu et al., 2023a; Yue et al., 2023; Luo et al., 2023; Tao et al., 2024; Gou et al., 2023). MetaMath (Yu et al., 2023a) augments its dataset by rewriting mathematical queries and providing answers with the help of an LLM. It also leverages backward reasoning to enhance the model’s capabilities. MAmmoTH (Yue et al., 2023) compiles diverse mathematical datasets for training, enriching them with Chains of Thought (CoT) and Programs of Thought (PoT) rationales for reasoning. RL-based methods strive to gather superior training signals through self-sampling and reward supervision. Lightman et al. (2023) proposes to use process supervision to give more precise signals than outcome supervision. Some strategies adopt the process supervision into the online RL training (Wang et al., 2023b; Xi et al., 2024; Luo et al., 2023; Yu et al., 2023b). RFT (Yuan et al., 2023) is a simplified offline method of standard online RL. It enhances its training dataset with accurate reasoning paths by generating and collecting them via an SFT model. This approach is further extended by ReST (Singh et al., 2023) through an iterative process.
Our method is parallel to the above fine-tuning methods. after obtaining supervisory signals through different approaches, one can consider using this regularization method in the final optimization step to further aid in training.
A.2 Data Augmentation for NLP
In the area of Natural Language Processing (NLP), an extensive range of data augmentation techniques has been introduced for diverse applications, such as translation, summarization, and dialogue (Chen et al., 2023b; Fadaee et al., 2017; Xie et al., 2016; Lample et al., 2017; Guo et al., 2020b; Xie et al., 2022; Wang et al., 2018; Lin et al., 2022; Chen et al., 2022a). The noise-based strategies commonly involve the integration of random continuous noise (Cheng et al., 2018; Jain et al., 2023; Kedzie and McKeown, 2019), token dropout (Xie et al., 2016; Gal and Ghahramani, 2015), token swapping (Wei and Zou, 2019), and more advanced mixup strategies (Guo et al., 2020a; Xie et al., 2022). Recently, in the context of regularization for LLMs, Jain et al. (2023) has demonstrated that the application of continuous noise across all tokens during the instruction tuning can significantly improve the length and detail of the generated responses. However, this technique does not improve the reasoning ability. We adopt the discrete token noise through token dropout (Xie et al., 2016; Gal and Ghahramani, 2015) to reasoning tasks, finding that it works well, which may complement the dependency-rich structure characteristic of reasoning data.

heading: B Discussion of the Relation to RL
text: Since MFT and RFT exhibit similar performance, and RL is also considered a highly effective method for reasoning tasks, we discuss the relationship between MFT and RL and try to explain the reasons for the similar performance of MFT and RFT.
B.1 Supervised Fine-tuning (SFT)
The pre-trained language model is often enhanced for specific downstream tasks through Supervised Fine-Tuning (SFT) as follows:
LSFT = − T∑
t=1
log p(wgtt |wgt1:t−1) (2)
B.2 Reinforcement Learning (RL) Fine-tuning
After SFT, to further improve the model, RL is applied to align the model to the specific goal (Wang et al., 2023b, 2022a, 2021). The model automatically samples responses, and a reward model scores these generated responses to distinguish their quality. For mathematical reasoning, it has been verified that rewards based on intermediate processes yield better results compared to rewards based solely on final results (Wang et al., 2023b; Lightman et al., 2023).
Eq. (3) shows a simplified policy gradient formulation of process supervision (omitting some techniques like the KL penalty used for regularization (Schulman et al., 2017)). For a standard online RL process, the model samples a reasoning path for a question, and then the reward model is used to compute the reward for each step.
LRL = − T∑
t=1
R(wt) log p(wt|w1:t−1; q)
w1:T ∼ p(.|q)
R(wt) =    1, if w1:t is correct & wt is the end of one step,
0, otherwise.
(3)
Where q is the question of a math problem. w1:T denotes the sampled answer with a length of T . p is the policy model.
B.3 Rejection sampling Fine-Tuning (RFT)
RFT is a simplified offline version of RL with:
LRFT = − T∑
t=1
log p(wt|w1:t−1; q)
w1:T ∼ psft(.|q) & the outcome of w1:T is correct (4)
Where w1:T is sampled from a fixed SFT model psft in the first stage, and the new data is added to the original data for the second SFT stage. While p in Eq. (3) is updated online. Singh et al. (2023) also proposes to do this two-stage method for multiple iterations.
B.4 Discussion
Relation to RL Fine-Tuning. Our method can be regarded as a special case of Eq. (3) with the following assumption: 1) We relax the reward computation in Eq. (3) and assume that every token in the ground-truth data can provide a positive process reward. 2) Sampling masked position has a similar effect with sampling new paths.
LMFT = − T∑
t=1
log p(wgtt |ws1:t−1)
= − T∑
t=1
R(wt) log p(wt|w1:t−1; q),
R(wt) =
{ 1, if wt = w gt t ,
0, otherwise.
(5)
For assumption 1, since the ground truth is correct, there is no need to use the reward model for labeling, and most tokens are supposed to be associated with positive reward signals. For assumption 2, we think that it may hold when sampling of online RL is insufficient:
Differences in Sampling Compared to Online RL. Standard auto-regressive sampling can bring more reasoning paths to one query. However, we find that the sampled reasoning paths have minor differences in Fig. 6, such as word substitution (orange part), rewriting of formulas (red part), rearrangement of syntax (brown part), but the correct intermediate steps are consistent (green part).
Therefore, for queries in the dataset, each of them already includes a ground truth response, which can provide rich process supervision signals of correct intermediate steps. MFT does not sample new correct reasoning steps but rather different preceding reasoning steps, allowing the model to predict the correct steps given different contexts. In short, the online RL sample new actions and states in a Markov Decision Process, while we sampled the state to make it partially observed. Adding noise to the state is verified to be helpful in generalizing for offline RL (Sinha et al., 2022), such as continuous noise on state embedding. We compare with other noise in Tab. 5, and we find that the discrete noise of MFT is more efficient.
Sample Efficiency. For online RL, it is necessary to sample diverse and better reasoning paths to achieve better results than using a static dataset. However, after SFT for initialization, the sampled paths of
the trained query will be very similar to the ground truth. If the sampling is sufficient and the reward model is strong enough to find high-quality reasoning paths that may have better generality and conciseness, RL can achieve better results. This can be the reason why RFT, which uses much more samples, performs similarly to ours, as RFT only samples for multiple paths at once and only uses outcome supervision. On the contrary, online RL can gain more benefits (Wang et al., 2023b).
For MFT, the sample efficiency is related to the mask ratio. A small mask ratio gives the model a chance to exploit current high-confidence context tokens, while a higher mask ratio requires the model to explore the use of other distant tokens.
C Implementation Details

heading: D Dependency Statistics of Other Datasets and Models

heading: E Case Study of Dependency
text: Method Is 16-n Is 16-n=16-n
SFT 1 1 MFT(m=0,r=1) 1 1 MFT(m=1,r=0) 1 1
We also give visualizations of attention score for the case in Fig. 5. There is a bit brighter attention at the token “three” when predicting the next token in the position of “-” for MFT.

heading: F False Positives

heading: G Cases

heading: H Error Analysis
text: We also perform a quantitative error analysis to understand the impact of MFT on the final results by analyzing the types of errors. We adopt the following error types: Calculator Error (CE), One Step Missing Error (OSME), Semantic Understanding Error (SUE), and Incoherent Steps Error (ISE). CE and OSME are minor errors that can be fixed with one step, while SUE and ISE are errors of misunderstanding the question or previous steps that need substantial fixing as defined by Wei et al. (2022): Calculator Error (CE). The chains of thought were completely correct except for a calculator error—in other words, applying an external calculator to equations would make the chain of thought correct. One Step Missing Error (OSME). Chains of thought were correct except that they were missing a single step. These chains of thoughts could be rewritten to be correct by adding in an additional reasoning step that was missed. Semantic Understanding Error (SUE). There are errors in the semantic understanding of the problem. Incoherent steps error (ISE). Incoherent chain of thoughts. Some steps in the generated chain of thought did not follow from prior ones. We use GPT-4 as a proxy for the judgment. Tab. 15 demonstrates that MFT has fewer SUE and ISE errors. This may be due to the benefits obtained by enhancing the model’s question dependence. Additionally, the model experiences an increase in OSME and CE compared to SFT. The increased ratios might be minor errors that arise after solving major errors in question understanding.

heading: I Experiments on Alpaca

heading: J The Impact on Length
text: We investigate whether MFT has a preference for problems of different solution lengths. The horizontal axis represents the ground truth length. The experiment shows that MFT and RFT are essentially consistent, demonstrating improvements for solutions of different lengths.

heading: K Detail of Compared Regularization Methods
text: Dropout: Dropout is a widely used Method to reduce overfitting (Srivastava et al., 2014). We compared MFT with Embedding Dropout which drops out each dimension of the token embedding. For a fair comparison, we only add it to the input embedding. NEFTune: Adding random noise to the token embedding vectors. Its experiments show that a small scale of noise can enhance the capability of LLM conversation and improve Alpaca-eval scores (Jain et al., 2023). In its original setup, noise is added to all input embeddings, including both source and target. Scheduled Sampling: We utilize the Scheduled Sampling technique for Transformers (Mihaylova and Martins, 2019), which entails two forward passes. The first forward pass predicts the next token for each position, and the input for the second forward pass incorporates some of the predicted tokens from the first pass. In determining which tokens to replace, we employ a sampling strategy characterized by the parameter τ , representing temperature; a higher τ introduces more randomness. At lower temperatures, there is a higher likelihood of replacing tokens with their synonyms, as opposed to random replacement. MaskedLM. In our main experiments, we implement MFT by this method, adhering to the methodologies of prior Masked Language Models as (Bao et al., 2020; Devlin et al., 2019). Unlike encoder-based models (Devlin et al., 2019), which typically predict the current token at masked positions, we predict the next token at masked positions (Xie et al., 2016; Chen et al., 2023a). Upon determining the positions for masking, we are presented with two options: employing a [mask] token or substituting with a random token. The probability of selecting the [mask] token is denoted by m, and r represents the probability of opting for a random token. Su et al. (2022) set r to 0, while Devlin et al. (2019) set r to 0.1.

heading: L Detail of OOD Reasoning Datasets
text: We follow (Yue et al., 2024) to use several science reasoning datasets to evaluate the generalization ability. MMLU-S (Hendrycks et al., 2020). MMLU-S is a STEM subset of the full MMLU with 3.13K problems covering genetics, computer science, astronomy, etc. GPQA (Rein et al., 2023). Questions in biology, physics, and chemistry, crafted by experts. We use the diamond subset, which consists of 198 challenging problems. ARC (Clark et al., 2018). Questions sourced from diverse grade-level science exams. We utilize the challenge subset, comprising 1,172 test questions.
BBH (Suzgun et al., 2022). BIG-Bench Hard (BBH) contains a total of 6,511 challenging problems for LLM. During inference, we use 5/5/8/3-shot COT prompting for MMLU-S/GPQA/ARC/BBH separately, and the prompt we use is in alpaca format.

heading: Future Work
text: This study could pave the way for future research on regularization strategies for reasoning tasks and further exploration of the interplay between language dependency and reasoning capabilities