C loze M ath: Improving Mathematical Reasoning in Language Models by Learning to Fill Equations (2025)

Abstract: The capabilities of large language models (LLMs) have been enhanced by training on data that reflects human thought processes, such as the Chain-of-Thought format. However, evidence suggests that the conventional scheme of next-word prediction may not fully capture how humans learn to think. Inspired by how humans generalize mathematical reasoning, we propose a new approach named ClozeMath to fine-tune LLMs for mathematical reasoning. Our ClozeMath involves a text-infilling task that predicts masked equations from a given solution, analogous to cloze exercises used in human learning. Experiments on GSM8K, MATH, and GSM-Symbolic show that ClozeMath surpasses the strong baseline Masked Thought in performance and robustness, with two test-time scaling decoding algorithms, Beam Search and Chain-of-Thought decoding. Additionally, we conduct an ablation study to analyze the effects of various architectural and implementation choices on our approach.

## introduction to mimic human reasoning in mathematical settings, current large language models (llms) are prompted or trained to generate intermediate thinking steps before reaching a conclusion @xcite . some efforts focus on generating data in this format by collecting responses from humans or llms @xcite or by implicitly learning task-specific representations @xcite . however, while these outputs resemble the surface form of human thought, the current language modeling objective may not align with how humans in text infilling, masked equation spans are replaced with special mask tokens (<x>, <@xmath0. these mask tokens ensure alignment of token positions during training. a @xmath1, using bidirectional and causal attention masks, respectively. learn to think @xcite . when learning from a mathematical solution, it is easier for a human to generalize the main approach before working on the details, rather than memorizing which step follows another @xcite . considering the input problem in figure 1 , the textual part of the target solution describes its general approach (i.e., rationales), while the equations are problem-specific. in light of this observation, we propose enhancing the training of language models for mathematical reasoning by introducing an additional text infilling objective @xcite . in this approach, models learn to predict masked equations alongside the standard language modeling objective. this method reinforces the model's ability to infer mathematical relationships by predicting equations based on the surrounding rationale, encouraging a more structured understanding of problem-solving steps. while traditional text infilling objectives involve masking and predicting random text spans, we find that this approach disrupts the logical coherence of the unmasked parts available to the model, ultimately harming reasoning performance. regarding model architecture, prefixlm @xcite , which applies a bidirectional attention mask to the model's prompt, has been shown to work well with text infilling in t5 @xcite . however, in the context of mathematical reasoning, we demonstrate that only our equation masking strategy benefits from this architectural choice. from a different perspective, masked thought @xcite , a recent method with strong empirical results, regularizes a model to look further back into a math problem's definition by randomly corrupting the problem's solution during language model training. however, mathematical solutions may include auxiliary derived steps that are not directly connected to the problem definition. we demonstrate a common example in which masked thought may teach the model spurious correlations in section 3.2. empirically, our method also substantially outperforms masked thought. in summary, our contributions are as follows: (1) we introduce a simple yet highly effective finetuning strategy, namely clozemath, that enhances the mathematical reasoning abilities of llms. (2) we conduct extensive experiments on two standard benchmarks, gsm8k @xcite and math @xcite , including a study on clozemath scalability on test-time compute, and perform robustness analyses on the newly released benchmark @xcite . we also show that clozemath outperforms the recent strong baseline masked thought. (3) we present a comprehensive ablation study that explores the impact of various architectural choices on the effectiveness of our proposed method. finally, our clozemath approach trains to minimize the two objectives simultaneously: since the number of terms in l tf depends on the number of equations to be masked within each dataset, in practice, we balance the two training objectives by duplicating samples optimized for l lm to maintain a roughly 50:50 sample ratio. we implement clozemath with prefixlm (liu et al., 2018) on pre-trained decoder-only models. in prefixlm, the prompt part's attention is bidirectional, while the target sequence uses a causal mask. to let the model learn to use the two attention patterns correctly, there is a @xmath2. during inference, we employ the conventional next-token prediction. ## conclusion to better align with human-like mathematical reasoning, we introduce clozemath, an approach that trains language models to predict masked equations using a text-infilling objective alongside the standard language modeling objective. we highlight the pitfall of learning spurious correlations with maskedthought, a strong current baseline, and demonstrate the consistently superior performance of clozemath across various tests of performance and robustness.