Physics of Language Models: Part 2.1, Grade-School Math and the Hidden Reasoning Process — part 1/6

Cleaned-text code points [0, 15291). Bibliography omitted here; source.md includes it.

# Physics of Language Models: Part 2.1, Grade-School Math and the Hidden Reasoning Process Thanks: Project page + video: https://physics.allen-zhu.com/part-2-grade-school-math/part-2-1 . We would like to thank Lin Xiao, Chunting Zhou for many helpful conversations. We would like to extend special thanks to Lucca Bertoncini, Liao Hu, Caleb Ho, Wil Johnson, Apostolos Kokolis, Parth Malani, Alexander Miller, Junjie Qian and Shubho Sengupta from Meta FAIR; Henry Estela, Rizwan Hashmi, Lucas Noah, and Maxwell Taylor from the Meta Cloud Foundation team; as well as Ian Clark, Gourab De, Anmol Mann, and Max Pfeifer from W&B . Without their invaluable support, the experiments in this paper would not have been possible.

###### Abstract

Recent advances in language models have demonstrated their capability to solve mathematical reasoning problems, achieving near-perfect accuracy on grade-school level math benchmarks like GSM8K. In this paper, we formally study how language models solve these problems. We design a series of controlled experiments to address several fundamental questions: (1) Can language models truly develop reasoning skills, or do they simply memorize templates? (2) What is the model’s hidden (mental) reasoning process? (3) Do models solve math questions using skills similar to or different from humans? (4) Do models trained on GSM8K-like datasets develop reasoning skills beyond those necessary for solving GSM8K problems? (5) What mental process causes models to make reasoning mistakes? (6) How large or deep must a model be to effectively solve GSM8K-level math questions?

Our study uncovers many hidden mechanisms by which language models solve mathematical questions, providing insights that extend beyond current understandings of LLMs.

## 1 Introduction

The field of language models has made significant progress in recent years. Large models like GPT-4 [ 17 ] have shown initial signs of general intelligence [ 8 ] , while smaller models have demonstrated good reasoning abilities by solving challenging coding and math problems [ 15 , 11 , 16 ] .

In this paper, we focus on the ability of small language models to solve grade-school math problems. Unlike previous works that empirically push the accuracy of models on grade-school math benchmarks like GSM8K [ 9 ] and its augmentations (e.g., [ 16 , 22 ] ), we take a more principled approach. We aim to understand the following fundamental questions: 1. How do language models learn to solve grade-school level math problems? Do they just memorize templates, or do they learn reasoning skills similar to humans? Or do they discover new skills to solve the problems?

2. Do models trained solely on grade-school math problems only learn to solve these problems, or do they develop some more general intelligence?

3. How small can a language model be while still solving grade-school math problems? Is depth (number of layers) more important than width (number of neurons per layer), or does only size matter as suggested by practitioners [ 14 ] ?

These questions are fundamental to understanding the intelligence of language models. To study them, it might seem tempting to start with a pre-trained model and fine-tune it on existing datasets like GSM8K or GPT-4 augmented ones (e.g., [ 16 , 22 ] ). However, this approach has significant limitations: • Data contamination. The pretrain data of existing models mostly come from publicly available internet [ 10 ] , which is a pile of mess. We do not know how many math problems are included or their structures. There is significant concern regarding whether the GSM8K benchmark has been leaked to language models’ training datasets [ 22 ] . Even if the exact data is not, the pre-trained model might have seen almost identical questions (e.g., the same problem with different numbers). Thus, this approach cannot answer questions 1-3. We do not know whether a model truly learns the reasoning skills or it simply memorizes problem templates during training. Therefore, we need full control over the model’s pretrain data and must train a language model from scratch. This point has been reiterated recently in [ 3 , 2 ] .

• Solution diversity. The existing fine-tuning data, such as the GSM8K training set, contains only 7.5K grade-school math problems, which is insufficient to train a model from scratch. Although recent works use GPT-4 to augment GSM8K, this is not enough for our purpose. GPT-4 augmented problems might be biased towards a small number of solution templates, since the original GSM8K data has very few (obviously, at most 8K) solution templates. We need a much larger, more diverse set of grade-school math problems .

With these points in mind, we introduce our framework to generate a large set of diverse grade-school math (GSM) problems and use the dataset to train (from scratch) and test a GPT2-like language model. In the framework, we focus on the “logical reasoning” aspect of grade-school math problems, which involves the dependency of parameters in the problem statement, such as “Alice’s apple is three times the sum of Bob’s orange and Charles’s banana.” We use synthetic sentences to reduce the difficulty arising from Common Sense , like “a candle burned for 12 hours at 1 inch per hour” (implying the candle is reducing in length). We also remove the difficulty from pure arithmetic: we only consider integers and arithmetic mod 23 \bmod 23 . 1 1 1 There is a rich literature studying how well language models can learn arithmetic and length generalization, see [ 23 , 13 ] and the references therein. Modern language models are also equipped with retrieval-augmented generation (RAG), allowing arithmetic computations to be delegated to a calculator.

Moreover, our framework ensures that the generated math problems are highly diverse and do not come from a small subset of templates. Even ignoring all the arithmetic, English, variable names, and unused parameters, our problems still have more than 90 trillion solution templates (see Proposition 2.2 ), much larger than the size of GPT2-small (100M). Thus, language models cannot solve the math problems in our case by simply memorizing the solution templates.

In this paper, we use the GPT2 model [ 18 ] , but replace its positional embedding with rotary embedding (RoPE) [ 20 , 7 ] . We still call it GPT2 for brevity. We summarize our main contributions: – Result 2. We demonstrate that the GPT2 model, pretrained on our synthetic dataset, not only achieves 99% accuracy in solving math problems from the same distribution but also out-of-distribution generalizes, such as to those of longer reasoning lengths than any seen during training. This is similar to length generalization in arithmetics [ 6 , 13 ] , however, in our case, the model has never seen any training example of the same length as in test time. This signifies that the model can truly learn some reasoning skill instead of memorizing solution templates.

– Result 3. Crucially, the model can learn to generate shortest solutions, almost always avoiding unnecessary computations. This suggests that the model formulates a plan before it generates, in order to avoid computing any quantities that are not needed towards solving the underlying math problem.

– Result 4. We examine the model’s internal states through probing, introducing six probing tasks to elucidate how the model solves math problems. For instance, we discover the model (mentally!) preprocesses the full set of necessary parameters before it starts any generation. Likewise, humans also do this preprocess although we write this down on scratch pads.

– Result 5. Surprisingly, the model also learns unnecessary, yet important skills after pretraining, such as all-pair dependency. Before any question is asked, it already (mentally!) computes with good accuracy which parameters depend on which, even though some are not needed for solving the math problem . Note that computing all-pair dependency is a skill not needed to fit all the solutions in the training data. To the best of our knowledge, this is the first evidence that a language model can learn useful skills beyond those necessary to fit its pretraining data. 2 2 2 In our case, one can solve all the math problems without computing all-pair dependency. Our pretraining data never includes such information — all the solutions only compute necessary variables. This may be a preliminary signal of where the G in AGI can come from. 3 3 3 Indeed, the skill to sort relationships among in-context objects is a general skill, which may lead to — via instruction fine-tuning — skills for solving other tasks, such as discovering causal relationships, determining the influence of parameter changes, etc.

– Result 6. We explain why mistakes occur. For instance, the model makes systematic errors that can be explained by probing its internal states. Sometimes, these mistakes can be predicted before the model generates answers, making them independent of the random generation process. We connect this to practice, noting that GPT-4/4o also makes similar errors (though we cannot probe their internal states).

– Result 7+8. The depth of the language model is crucial for its reasoning ability. For example, a 16-layer, 576-dim transformer solves harder problems (in reasoning length) than a 4-layer, 1920-dim one, despite the latter being twice as large. This holds even when Chain-of-Thought (CoT) is used. We explain this necessity in depth by the complexity of the mental processes involved. We advocate for the use of controlled, synthetic data as a more principled approach to derive such claims, contrasting with predictions like “only size matters” based on training loss using internet pretrain data [ 14 ] .

While we refrain from overstating that our findings directly apply to foundation models like GPT-4 or more challenging mathematical reasoning tasks, we believe our work significantly advances the understanding of how language models develop their mathematical reasoning skills, and this has to be done in a way different from pushing benchmarks .

## 2 Result 1: Data Generation

Motivation. Recall a standard grade-school math problem in the GSM8K dataset [ 9 ] looks like: Betty is saving money for a new wallet which costs 100. Betty has only half of the money she needs. Her parents decided to give her 15 for that purpose, and her grandparents twice as much as her parents. How much more money does Betty need to buy the wallet? This problem involves multiple parameters whose values are connected through various equalities, such as “Betty’s current money = 0.5 × \times cost of the wallet” and “money given by grandparents = 2 × \times money given by parents.” Motivated by this, we build a GSM8K-like math dataset through a synthetic generation pipeline that captures the dependencies of parameters. We wish to capture at least the following three types of dependencies. 1. Direct dependency ( ♡ \heartsuit ): such as A = 5 × ( X + Y ) A=5\times(X+Y) , so A A can be computed after X X and Y Y .

2. Instance dependency ( ♠ \spadesuit ): such as “every classroom has X chairs, and there are Y classrooms.” Here, the model must infer the total number of chairs by multiplying X by Y.

3. Implicit dependency ( ♣ \clubsuit ): such as “Bob has 3 times more fruits than Alice. Alice has 3 apples, 4 eggs and 2 bananas.” Here, the model must learn that apples and bananas are fruits and egg is not, and “Alice’s fruits” is an abstract parameter derived from the problem statement.

### 2.1 Step 1: Graph Construction and Problem Generation

Hierarchical categorization. We use a layered structure of categories , each contains possible items . For instance, categories = ( School , Classroom , Backpack ) has three layers; category School = { Central High , Riverview High , …}; category Classroom = { Dance Studio , Film Studio , …}; category Backpack = { School Daypack , Messenger Backpack , …}. We prepare 4 predefined hierarchical categorizations, each with 4 layers and 100 items in each layer; this represents the world knowledge.

Structure graph. In each math problem, only specific items exist, leading to a structure graph that outlines what sub-items can appear under what item, see Figure 1 (left). For instance, • Connecting Dance Studio and School Daypack with an edge signifies an instance parameter , ‘‘the number of school daypacks in each dance studio,’’ which is a quantifiable variable that can be assigned. 4 4 4 Even though Central High and Rivierside High can both have (possibly multiple) Dance Studios, for simplicity, we assume that each Dance Studio has the same number of School Daypacks. This captures the instance dependency ( ♠ \spadesuit ) as mentioned above.

• Abstract parameters , like “the total number of classrooms in Central High,” cannot be assigned and are excluded from the structure graph. They reflect implicity dependency ( ♣ \clubsuit ) .

###### Remark 2.1 .

Rather than using simple objects like Alice’s apple or fake items like Items A/B/C/D , this structure allows us to describe abstract parameters and adds 2 levels of complexity to the data: • The model must implicitly learn English concepts, such as a classroom category includes 100 different classroom types. These concepts cannot be derived from individual math problems, as only a limited selection of classrooms will be mentioned in each problem.

• The model is required to hierarchically access multiple items to calculate abstract parameters, as opposed to a straightforward retrieval of ‘‘Alice’s apple’’ in the context. 5 5 5 For example, the total number of backpacks in Riverview High in Figure 1 is calculated as i ​ p 1 × a ​ p 1 + i ​ p 2 × a ​ p 2 ip_{1}\times ap_{1}+ip_{2}\times ap_{2} where i ​ p 1 = “Riverview High’s number of Dance Studios” ip_{1}=\textrm{``Riverview High's number of Dance Studios''} , i ​ p 2 = “Riverview High’s number of Film Studios” ip_{2}=\textrm{``Riverview High's number of Film Studios''} , a ​ p 1 = “each Dance Studio’s number of Backpacks” ap_{1}=\textrm{``each Dance Studio's number of Backpacks''} , and a ​ p 2 = “each Film Studio’s number of Backpacks” ap_{2}=\textrm{``each Film Studio's number of Backpacks''} , with i ​ p 1 , i ​ p 2 ip_{1},ip_{2} being instance parameters and a ​ p 1 , a ​ p 2 ap_{1},ap_{2} abstract parameters. Here, the model must not only retrieve i ​ p 1 , i ​ p 2 ip_{1},ip_{2} but also compute a ​ p 1 , a ​ p 2 ap_{1},ap_{2} hierarchically.

Dependency graph. The dependency graph is a directed acyclic graph that outlines the dependency among parameters. For each instance parameter , we choose a random set of (up to 4) parameters it can depend on — including possibly a special vertex 𝖱𝖭𝖦 \mathsf{RNG} representing a random number generator. For instance, if “[param A ] is X X more than the difference of [param B ] and [param C ]” for X X being randomly generated, then we draw edges from B, C and 𝖱𝖭𝖦 \mathsf{RNG} to parameter A. The dependency of abstract parameters is implied by the dependency of instance parameters. This captures direct dependency ( ♡ \heartsuit ) as mentioned above. We give an examples on the right side of Figure 1 , and details for how we randomly generate such dependency graph are in Appendix D.2 .