Physics of Language Models: Part 2.1, Grade-School Math and the Hidden Reasoning Process

Introduction

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

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 [cobbe2021training] and its augmentations (e.g., [liu2023tinygsm, zhang2024careful]), we take a more principled approach. We aim to understand the following fundamental questions:

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?

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

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 [kaplan2020scaling]?

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., [liu2023tinygsm, zhang2024careful]). However, this approach has significant limitations:

Data contamination. The pretrain data of existing models mostly come from publicly available internet [gao2020pile], 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 [zhang2024careful]. 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 [AL2023-knowledge, AL2023-knowledgeUB].

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 $\bmod 23$.There is a rich literature studying how well language models can learn arithmetic and length generalization, see [zhou2023algorithms, jelassi2023length] 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 [radford2019language], but replace its positional embedding with rotary embedding (RoPE) [su2021roformer, gpt-neox-20b]. 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 [anil2022exploring, jelassi2023length], 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.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.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 [kaplan2020scaling].

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.

Result 1: Data Generation

Motivation. Recall a standard grade-school math problem in the GSM8K dataset [cobbe2021training] looks like:

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.

Direct dependency ($\heartsuit$): such as $A=5\times(X+Y)$, so $A$ can be computed after $X$ and $Y$.

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.

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.

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.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$) .

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.For example, the total number of backpacks in Riverview High in Figure 1 is calculated as $ip_{1}\times ap_{1}+ip_{2}\times ap_{2}$ where $ip_{1}=\textrm{``Riverview High's number of Dance Studios''}$, $ip_{2}=\textrm{``Riverview High's number of Film Studios''}$, $ap_{1}=\textrm{``each Dance Studio's number of Backpacks''}$, and $ap_{2}=\textrm{``each Film Studio's number of Backpacks''}$, with $ip_{1},ip_{2}$ being instance parameters and $ap_{1},ap_{2}$ abstract parameters. Here, the model must not only retrieve $ip_{1},ip_{2}$ but also compute $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$ more than the difference of [param B ] and [param C ]” for $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.

Problem generation. The problem is articulated by describing the dependency graphs in English, one sentence for each instance parameter.We use simple English sentence templates to describe the problem, and did not worry about grammar mistakes such as singular vs plural forms. There are other randomness besides the dependency graph, such as when parameter $A$ depends on $B,C$ it could be $A+B$ or $A-B$. (Abstract parameters are not described because they are inherited by the structure graph.) We randomly permute the sentence ordering to further increase difficulty. A parameter is selected and asked with a question in the end (or at the beginning). Below is an easy example corresponding to Figure 1; a harder example is in Figure 11. (Problem - Easy) The number of each Riverview High’s Film Studio equals 5 times as much as the sum of each Film Studio’s Backpack and each Dance Studio’s School Daypack. The number of each Film Studio’s School Daypack equals 12 more than the sum of each Film Studio’s Messenger Backpack and each Central High’s Film Studio. The number of each Central High’s Film Studio equals the sum of each Dance Studio’s School Daypack and each Film Studio’s Messenger Backpack. The number of each Riverview High’s Dance Studio equals the sum of each Film Studio’s Backpack, each Film Studio’s Messenger Backpack, each Film Studio’s School Daypack and each Central High’s Backpack. The number of each Dance Studio’s School Daypack equals 17. The number of each Film Studio’s Messenger Backpack equals 13. How many Backpack does Central High have? (2.1)

2 Step 2: Solution Construction (CoT)

Let solution be a sequence of sentences describing the necessary steps towards solving the given problem, where the sentences follow any topological order — also known as Chain-of-Thought, CoT. For each parameter necessary towards answering the final question, we assign to it a random letter among the 52 choices (a..z or A..Z), and use a sentence to describe its computation:There are different ways to format the CoT solution. We noted that starting with “Define [param] as X” instead of [intermediate steps] improves the model’s accuracy, so we have adhered to this CoT format.

Throughout this paper, we consider arithmetics mod $23$ to avoid errors from computation involving large numbers. It is perhaps the easiest to directly see a solution example (corresponding to (2.1)), and a more involved example is in Figure 11: (Solution - Easy) Define Dance Studio’s School Daypack as p; so p = 17. Define Film Studio’s Messenger Backpack as W; so W = 13. Define Central High’s Film Studio as B; so B = p + W = 17 + 13 = 7. Define Film Studio’s School Daypack as g; R = W + B = 13 + 7 = 20; so g = 12 + R = 12 + 20 = 9. Define Film Studio’s Backpack as w; so w = g + W = 9 + 13 = 22. Define Central High’s Backpack as c; so c = B * w = 7 * 22 = 16. Answer: 16. (2.2)

The solution only contain parameters necessary towards calculating the final query parameter.

The solution follows the correct logical order: i.e. all the parameters used in the calculation must have appeared and been computed beforehand.

We break computations to binary ops: $g=12+13+7$ is broken into $g=12+R$ and $R=13+7$ in the above solution. The number of semicolons “;” equals the number of operations. This reduces the arithmetic complexity of the solution, which is not the focus of this paper.Even GPT-4 can make mistakes on calculating “3 * (4+10) + 12 * (5+6)” without using external calculator.

3 Difficulty Control

Although deferring all the pseudocode to Appendix D, we summarize below the main randomness used in the data generation process. This includes the random choice of a hierarchical categorization (i.e., the English part); a structure graph (i.e., the instance parameters); a dependency graph; arithmetic computations on the dependency graph; integer numbers (i.e., the $\mathsf{RNG}$); problem sentence permutation; and the query parameter.

We use two parameters to control data’s difficulty: ip is the number of instance parameters, and op is the number of solution operations; the data’s difficulty is an increasing function over them. We call our dataset iGSM, to reflect the nature that such synthetic dataset can be of infinite size. We use $\textsf{iGSM}^{\textsf{op}\leq op,\textsf{ip}\leq ip}$ to denote the data generated with constraint $\textsf{op}\leq op$ and $\textsf{ip}\leq ip$, and use $\textsf{iGSM}^{\textsf{op}=op,\textsf{ip}\leq ip}$ to denote those restricting to $\textsf{op}=op$. We choose op non-uniformly; for instance, we let $\textsf{op}=\min\{t_{0},t_{1}\}$ for two random draws $t_{0},t_{1}\in[op]$. This ensures that the dataset has more easy data — which makes training faster. (See also similar behavior for arithmetics [jelassi2023length].)

4 Train and Test Datasets

In the iGSM-med data family we use $\textsf{ip}\leq 20$.

In the iGSM-hard data family we use $\textsf{ip}\leq 28$.

Additionally, we use $\textsf{iGSM-med}_{pq}$ to indicate placing the question after the problem and $\textsf{iGSM-med}_{qp}$ the other way (similarly for iGSM-hard). The difficulty of iGSM-med is already quite non-trivial to humans (at least not solvable with few-shot learning using GPT-4/4o, see Figure 2).

Ignoring unused parameters, numerics, sentence orderings, English words, a-z and A-Z letter choices, $\textsf{iGSM-med}^{\textsf{op}=15}$ still has at least $7$ billion solution templates, and $\textsf{iGSM-hard}^{\textsf{op}=21}$ has at least $90$ trillion solution templates.A solution template is created by replacing all numbers with ‘0’, substituting variables (a-z or A-Z) with letters in their appearance order, and changing parameters to their types (instance or abstract). For instance, “Define Owl Forest’s Elephant as y; so y = 11. Define Parrot Paradise’s Raccoon as t; so t = y = 11.” becomes “Define Inst as a; so a = 0. Define Inst as b; so b = a = 0.” We use birthday paradox to estimate the number of solution templates. If $M$ randomly generated problems yield distinct templates, it suggests with good probability that the total number of templates exceeds $\Omega(M^{2})$.

No data contamination. A goal in synthetic math data generation is to prevent data contamination in internet-based math datasets, as noted in [zhang2024careful]. While it may be impossible to certify that models trained on internet data are free from contamination, in our setting, we can certify this:

We perform OOD evaluation such as on $\textsf{op}\geq 28$ while providing only $\textsf{op}\leq 21$ training samples.

We train with data whose hash value of solution template (see Footnote 11) is $<17\pmod{23}$, and test with those $\geq 17$. This ensures no template-level overlap between training and testing.

Result 2-3: Summarize Model’s Behavior Process

We use the GPT2 architecture [radford2019language] but replacing its absolute positional embedding with rotary embedding [su2021roformer, gpt-neox-20b], yet still referring to it as GPT2 for short.We also tested with Llama architecture (esp. with gated MLP layers) and did not see any benefit of using it. GPT2-rotary performs no worse than Llama/Mistral for knowledge tasks [AL2024-knowledgeScaling]. We are currently bounded by resources to repeat all experiments in this paper with other architectures that have minor differences from GPT2-rotary. We mostly stick to the 12-layer, 12-head, 768-dim GPT2 (a.k.a. GPT2-small) for experiments, but we explore larger models in Section 6. We use a context length of 768 / 1024 for pretraining on $\textsf{iGSM-med}/\textsf{iGSM-hard}$ and 2048 for evaluation. More details are in Appendix LABEL:app:pretrain.

Result 2: accuracy. After sufficient pre-training, we give the model a problem from the test set (without solution) and let it continue to generate (allegedly a solution followed by an answer). Because we have restricted ourselves to a fixed solution format, language models can learn the format easily, allowing us to write a solution parser to check if the solution is fully correct.We check not only the correctness of the final answer 0..22 but also the calculations and parameter dependencies. Language models can learn very complex syntactics, see [AL2023-cfg] and the references therein. Result 2. Figure 3 shows that GPT2 performs well when pretrained using iGSM-med or iGSM-hard data, even when evaluated out-of-distribution on harder (i.e., larger op) math problems. Thus, the model can truly learn some reasoning skill instead of memorizing solution templates.Llama (of the same model size) gives similar performance, but we refrain from repeating all the experiments with another model. We are not interested in small model differences in this theoretical study; instead, we care more about the general behavior of (autoregressive) language models. This could be reminiscent of language models’ length generalization capability on arithmetic computations [zhou2023algorithms, jelassi2023length]; however, in our case, op captures the “reasoning length” in grade-school math, and our model has never seen any training example of the same length as in test time.Some others such as anil2022exploring start with a transformer pre-trained on internet data; while the transformer may not have seen the same task during training, it’s possible that the model has seen other tasks with the same (or even longer) length and learned to transfer from there.

Such accuracies also indicate that our iGSM data families are indeed good for pretraining purpose, allowing us to investigate how LLMs can solve grade-school math problems.

Result 3: solution redundancy. We examine whether GPT2 achieves high accuracy by

brute-forcedly computing all the parameters during generation (a “level-0” reasoning skill), or

computing only necessary parameters to give shortest solutions (a “level-1” reasoning skill).

Recall our iGSM (pretrain) data only contains necessary solution steps (i.e., CoT) to simulate what we see in textbook solutions for math problems. For instance, if a problem describes X =3+2, E =3+X, Y =X+2 and asks for the value of Y, then a shortest solution would be “X =3+2=5 and Y =X+2 =7” without ever computing E.

This finding is significant as it suggests that, unlike humans who usually rely on “backward reasoning” and a scratch pad to write down necessary parameters by backtracking the dependencies from the question [rips1994psychology], the language model can directly generate shortest solutions without using a scratch pad. But, how does it achieve so? We shall investigate in the next section.

Result 4-5: Discover Model’s Mental Process

To understand how the model learns to solve math problems, we propose studying the following probing tasks, which align closely with human problem-solving strategies:

$\texttt{nece}(A)$: if parameter $A$ is necessary for computing the answer.

$\texttt{dep}(A,B)$: if parameter $A$ (recursively) depends on parameter $B$ given the problem statement.

$\texttt{known}(A)$: if parameter $A$ has already been computed.

$\texttt{can\_next}(A)$: if $A$ can be computed in the next solution sentence (namely, its predecessors have all been calculated). Note that $A$ might not be necessary to answer the question.

$\texttt{nece\_next}(A)$: if parameter $A$ satisfies both $\texttt{can\_next}(A)$ and $\texttt{nece}(A)$.

For a model to generate the shortest solutions, it must identify $\texttt{nece}(A)$ for all $A$’s in its mental process. This is because whether $\texttt{nece}(A)$ is true directly corresponds to whether there is a solution sentence to compute $A$. However, how early does the model recognize this, and how is it stored? Similarly, does it recognize dependencies between parameters (dep)? If so, how early is this mental process completed? Moreover, in the middle of solution generation, does the model keep track of each parameter $A$’s value at all times (value, known)? Does the model mentally know all possible parameters $A$ that are ready to compute in the next sentence (can_next)? Or does it only focus on $A$ that is both ready and necessary (nece_next)?

This section proposes probing technique to answer all of these questions.

As illustrated in Figure 5, we conduct probing at the end of the problem description for the dep task, and end of the question description nece task.If the problem format is qp (question asked before the problem) then we probe nece and dep both after the problem description. For other tasks, we probe them at the end of every solution sentence (including the start of the first solution sentence).

Recall that standard linear probing involves freezing a pretrained language model and checking if a property is linearly encoded at a hidden layer (usually the last layer) for a given token position. This is done by introducing a trainable linear classifier on the hidden states and performing a lightweight finetuning task for this property (see [hewitt-manning-2019-structural] and references therein).

Our setting is more complex because the properties have one or two conditional variables, $A$ and $B$, described in plain English. To handle this, we truncate the math problems to the probing position and append tokens [START] and [END] around the descriptions of $A$ (or $A,B$). We then probe from the token position of [END] to see if the property is linearly encoded at the last layer.

Unlike standard linear probing, to account for the input change, we introduce a small trainable rank-8 (linear) update on the input embedding layer. We freeze the pretrained language model and finetune both the linear classifier and the rank-8 update for the desired property. We refer to this as V(ariable)-probing and provide details in Appendix B. An illustration of the $\texttt{nece}(A)$ probing task is shown in Figure 6.

We compute the V-probing accuracies on a language model pretrained from iGSM and compare them with the V-probing accuracies on a randomly-initialized transformer model. If the former accuracies are significantly higher, we conclude that the probing signals must have (or be very close to having) come from the pretrained weights, rather than the (lightweight) finetuning stage.

2 Probing Results and Findings

We present our probing results in Figure 7. The probing accuracies are high for all the tasks, compared to majority guess and random-model probing — except for the very hard OOD cases (i.e., for large op where the model’s generation accuracies fall down to 80% anyways in Figure 3),

Result 4: model solves math problems like humans. We make the following observations:

When generating solutions, the model not only remembers which parameters have been computed and which have not ($\texttt{value},\texttt{known}$) but also knows which parameters can be computed next ($\texttt{can\_next},\texttt{nece\_next}$). These abilities ensure that the model can solve the given math problem step by step, similar to human problem-solving skills.

By the end of the problem description, the model already knows the full list of necessary parameters (nece). This indicates that the model has learned to plan ahead, identifying necessary parameters before starting to generate the solution. This aligns with human behavior, except that the model plans mentally while humans typically write this down. This further confirms that the model reaches the “level-1” reasoning skill discussed in Section 3.

The mental process described can be compared to (out-of-context) knowledge manipulation [AL2023-knowledgeUB], which involves retrieving factual knowledge and performing single-step computations (e.g., retrieving two people’s birth dates to determine who was born earlier). AL2023-knowledgeUB found that even single-step computations cannot be performed mentally without a substantial number of pretrain samples. In contrast, this paper studies in-context reasoning and demonstrates that the model can execute very complex mental calculations.

Result 5: model learns beyond human reasoning skills. Remarkably, the model learns $\texttt{dep}(A,B)$ and $\texttt{can\_next}(A)$, even for parameters $A$ not necessary for answering the question, as shown in Figure 7(b). This differs from human problem-solving, where we typically use backward reasoning from the question to identify necessary parameters, often overlooking unnecessary ones [rips1994psychology]. In contrast, language models can pre-compute the all-pair dependency graph $\texttt{dep}(A,B)$ mentally even before a question is asked. We consider this a “level-2” reasoning skill that is very different from human behavior or mental processes.

Thus, although this skill is not needed for solving the math problems and although no pretrain data teaches the model to compute “all-pair dependency” — fitting the data only requires computing necessary parameters — the model still discovers it after training. This enables the model to sort relationships among the things it hears, a skill that can be useful for future tasks (via instruction fine-tuning). To our knowledge, this may be the first evidence of a language model acquiring skills beyond those needed for learning its pretrain data; and this may be a preliminary signal of where the G in AGI can come from (generalizing to skills not taught in the pretrain data).

Corollary: the backward thinking process. A key question for AGI success is whether the “backward thinking process” (e.g., “because I want to compute X, but X depends on Y and Y depends on Z, so let me compute Z first”) needs to be explicitly included in the training data. This differs from CoT, where CoT breaks down complex computations into simpler steps, but planning is still required to decide which step to compute first.

Our findings suggest that, at least for grade-school math problems, with abundant data, this backward thinking process can be autonomously learned through language modeling, without needing to be directly included in the training data.

Result 6: Explain Model’s Mistakes

We further examine the relationship between our probing results and the model’s generated solutions, focusing on two questions: (1) When does the model answer correctly but include unnecessary parameters? (2) What causes incorrect answers? We aim to determine if such erroneous behavior of the model aligns with errors in the model’s mental process (via probing).

Result 7-8: Depth vs. Reasoning Length

Our controlled dataset enables a systematic exploration of the relationship between a language model’s depth and its reasoning length.

Recent studies have demonstrated that for knowledge storage and extraction, only model size matters (even for 2-layer transformers) [AL2024-knowledgeScaling]. Furthermore, both the seminal scaling-law paper by OpenAI [kaplan2020scaling] and theoretical studies in deep learning [als18dnn] suggest that model depth/width might have a minimal impact universally. Contrary to these findings, we present evidence that Math reasoning data only occupies a tiny fraction of pretraining data for language models, thus one might not observe a difference if we only look at the perplexity as in the original scaling law paper [kaplan2020scaling].

Figure 10 shows our result. It reveals a correlation between the model’s layer hierarchy, reasoning accuracy, and mental reasoning depth. Shallower layers excel at predicting $\texttt{nece}(A)$ for parameters $A$ closer to the query, whereas deeper layers are more accurate and can predict $\texttt{nece}(A)$ for parameters further from the query. This suggests that the model employs layer-by-layer reasoning during the planning phase to recursively identify all parameters the query depends on, and:

We make two disclaimers here. First, if the “backward thinking process” is added as CoT to the data (see the end of Section 4.2), then deep mental thinking is no longer required, reducing the language model’s depth requirement. However, in practice, many such “thinking processes” may not be included in standard math solutions or languages in general.

Second, the above claim does not imply that “a $t$-step mental thinking requires a depth-$t$ transformer”. It is plausible for a single transformer layer (containing many sub-layers) to implement $t>1$ mental thinking steps, though possibly with reduced accuracy as $t$ increases. We refrain from providing an exact correlation in this paper, as it heavily depends on the data distribution.

Conclusion

We use a synthetic setting to demonstrate that language models can learn to solve grade-school math problems through true generalization, rather than relying on data contamination or template memorization. We develop probing techniques to examine the models’ hidden reasoning processes. Our findings reveal that these models can learn math skills aligned with human cognitive processes, as well as “new thinking processes” not present in the training data. Additionally, we propose a method to predict a model’s errors before it begins to solve a problem and to explain why models make mistakes when they occur. Based on this discovery, we write a separate paper to improve language models’ math reasoning accuracy [YXLZ2024-gsm2]. We also provide a principled approach to connect the model’s depth to its capable reasoning length. We believe this research opens doors to study the mathematical reasoning skills of language models from a different angle compared to pushing math benchmarks.

One may argue that iGSM may be very different from the pretrain data that modern LLMs use. While this may be true, we are looking into the future. Recall, even GPT-4/4o of today cannot few-shot learn to solve $\textsf{iGSM-med}^{\textsf{op}=11}$ (see Figure 2). From this perspective, it is reasonable to believe that future versions of LLMs will rely on synthetic math pretrain data to improve their reasoning skills. While one may not directly use iGSM, it is tempting to use existing LLMs (such as Llama-3) to turn iGSM into more natural formats while keeping the logical chains. On the other hand, we have discovered that models trained purely on the iGSM data make similar mistakes compared to GPT-4/4o (see Section 5 and Appendix LABEL:app:gpt-4). This further confirms that our findings do connect to practice, regarding the model’s hidden reasoning process.

Finally, Part 2 of this work series focuses on how language models solve grade-school math problems (including Part 2.2 [YXLZ2024-gsm2]). We also cover how language models learn language structures in Part 1 [AL2023-cfg] (in particular, how they mentally perform dynamical programming), and learn world knowledge in Part 3 [AL2023-knowledge, AL2023-knowledgeUB, AL2024-knowledgeScaling].

Appendix A Result 1 — An Example in iGSM-hard with op=21fragmentsop21\textsf{op}=21

Appendix B Results 4-5 — Details on V-probing

Recall that we wish to conduct probing at the end of the problem description for the nece and dep tasks (before the solution for nece; before the solution or even the question for dep). For other tasks, we probe at the end of every solution sentence (including the start of the first solution sentence). The goal is to freeze a pretrained language model, then introduce a very small number of additional trainable parameters on top of it, and finetune them for each probing task.

Specifically, we take a pretrained language model, e.g., pretrained from the iGSM-hard training data. We freeze its parameters completely except for adding a trainable rank-$r$ update on the embedding layer to account for the task change (from next-token prediction to probing). Throughout this paper we use a small value $r=8$. We feed this network with training data that are the same as iGSM-hard, but truncated at exactly the position we wish to probe. Importantly, we append such inputs with a special starting token [START] along with a parameter name (or two names, if it is the $\texttt{dep}(A,B)$ task). We then extract the hidden states of the last token position at the last transformer layer, and add a trainable linear layer (a.k.a. linear head) to perform classification for one of the six probing tasks.

This probing method is illustrated in Figure 13. We call it V(ariable)-Probing, because it can take an arbitrary number of variables (i.e., parameters in this paper) to allow us to perform functional probing inside the transformer.

Note, if it were only a trainable linear head such probing would be called linear probing [hewitt-manning-2019-structural]. Unlike traditional linear probing, we are adding a small low-rank update on the model’s embedding layer. This is arguably the minimum change needed (to account for the task change, for special tokens like [START] [MID] [END], etc.) in order to perform any non-trivial probing. This is related but different from the nearly-linear probing methods introduced in AL2023-cfg, AL2023-knowledge, because they do not support taking variables as probing inputs.In AL2023-cfg, AL2023-knowledge, the authors are interested in probing the model’s behavior via fixed classification tasks (such as a 100-class classification task) given data that are identical or nearly-identical to the pretrain data. In this paper, we are interested in the model’s behavior with respect to given variables (such as parameter names, which can have $\sim 100k$ possibilities); and we append such variable names to the input to make the training inputs appear very different from the original pretrain data.

Appendix C Result 8 — Additional Figure

Appendix D Result 1 Details — Math Data Generation

Our math data generation process consists of first generating the structure graph (see Figure 1 and 11 left), which defines the set of parameters we shall use; then generating the dependency graph (see Figure 1 and 11 right), which defines the arithmetic relationship between the parameters; and finally generating the English problem and solution descriptions.

Notations. In this section, to make the description concise, when we say “randomly sampling” in the pseudocode, we mean uniform random unless otherwise noted. Whenever we consider a (directed) graph $G$, slightly abusing notation, we write $a\in G$ to indicate that $a$ is a vertex in $G$ and $(a\to b)\in\G$ to indicate that there is an edge from $a$ to $b$ in $G$.

Recall the structure graph (see Figure 1 and 11 left) describes the set of possible items (nodes) and instance parameter (edges) that we shall rely on to construct our math problem.

Specifically, suppose $l_{i}\in\{w_{0},w_{0}+1,\dots,w_{1}\}$ represents the number of items for each layer $i$. In this configuration, one must have at least $e^{-}=l_{2}+\cdots+l_{d}$ edges to ensure the graph is “connected”, and at most $e^{+}=l_{1}l_{2}+\cdots+l_{d-1}l_{d}$ edges. Using this formula, we first randomly choose a configuration $(l_{1},\dots,l_{d})$ so that $e^{-}\leq e\leq e^{+}$ for the given parameter $e$. Then, after the configuration is chosen, we randomly generate edges accordingly. Details are given in Algorithm 1.

As described in Section 2.1, we have prepared 4 predefined hierarchical categorizations, each of them with 4 total layers of categories: {python}[basicstyle=] [ ["District", "Supermarket", "Product", "Ingredient"], ["Zoo", "Enclosure", "Animal", "Bone"], ["School", "Classroom", "Backpack", "Stationery"], ["Ecosystems", "Creatures", "Organs", "Cells"] ] In each of the above 16 categories, we have prepared around 100 items (further decomposed into 5 sub-categories). Below is a showcase of them: {python}[basicstyle=] "District": "Residential Districts": [...], "Commercial Districts": [ "Shopping District", "Business District", "Financial District", "Industrial District", "Warehouse District", "Market District", "Restaurant District", "Entertainment District", "Arts District", "Fashion District", "Silicon Valley", "Wall Street", "Tech Park", "Automotive District", "Jewelry District", "Medical District", "Legal District", "Media District", "Research Park", "Manufacturing District" ], "Historical Districts": [...], "Educational Districts": [...], "Government Districts": [...] , "Supermarket": ..., "Product": "Canned Foods": [...], "Snack Foods": [ "Potato Chips", "Pretzels", "Popcorn", "Candy Bars", "Gummy Candy", "Cookies", "Crackers", "Granola Bars", "Fruit Snacks", "Cheese Puffs", "Nuts", "Trail Mix", "Beef Jerky", "Rice Cakes", "Yogurt Covered Raisins", "Chocolate Covered Pretzels", "Tortilla Chips", "Salsa", "Hummus", "Dried Fruit" ], "Beverages": [...], "Baked Goods": [...], "Dairy Products": [...] , "Ingredient": ..., "Zoo": ..., "Enclosure": ..., "Animal": ..., "Bone": ..., "School": ..., "Classroom": ..., "Backpack": ..., "Stationery": ..., "Ecosystems": ..., "Creatures": ..., "Organs": ..., "Cells": ...

Now, given a constructed structure graph $G_{\mathsf{s}}$, we first randomly pick one of the four categorizations, then randomly pick $d\in\{2,3,4\}$ consecutive layers of categories, next randomly pick one of the five subcategories, and finally pick $l_{i}$ random item names in this subcategory for each layer $i$.

At this point, we have constructed $G_{\mathsf{s}}$ as well as added English names to each of its node, just like Figure 1 and 11 (left).

D.2 Generate Dependency Graph

A structure graph $G_{\mathsf{s}}$ defines the set of possible parameters we consider, while a dependency graph defines how these parameters depend on each other. We use an edge $a\to b$ to indicate that parameter $b$ depends on $a$; there is a special vertex $\mathsf{RNG}$ and it can happen that $\mathsf{RNG}\to b$. What an abstract parameter depends on is inherited from the structure graph $G_{\mathsf{s}}$. For each instance parameter, we shall randomly add edges to indicate what parameters it depends on.

High-level plan. We shall use $G_{\mathsf{d}}$ to denote the dependency graph, we start from an empty graph and then add vertices/edges incrementally and randomly. Our process is as follows:

Generate a necessary dependency graph $G_{\mathsf{d}}^{\mathsf{nece}}$ which covers all the vertices and nodes that are necessary for the computation of the query parameter.

Generate necessary abstract parameters (and add parameters they depend on); call this graph $G_{\mathsf{d}}^{\mathsf{nece1}}$.

Generate necessary instance parameters and add them to $G_{\mathsf{d}}^{\mathsf{nece1}}$; call this graph $G_{\mathsf{d}}^{\mathsf{nece2}}$.

Generate a topological order for parameters $G_{\mathsf{d}}^{\mathsf{nece2}}$ and ensure all of them are necessary towards computing the query parameter (which is the last one in this tropologic order). During this process, we shall add additional edges from $G_{\mathsf{d}}^{\mathsf{nece2}}$ to create $G_{\mathsf{d}}^{\mathsf{nece3}}$.

Generate additional necessary edges and add them to $G_{\mathsf{d}}^{\mathsf{nece3}}$; call this graph $G_{\mathsf{d}}^{\mathsf{nece}}$.

Add to $G_{\mathsf{d}}^{\mathsf{nece}}$ all the remaining (unnecessary) parameters and edges to form $G_{\mathsf{d}}$.

At a high level, our problem description shall solely depend on $G_{\mathsf{d}}$— by describing each instance parameter in it using a sentence, and our solution description shall solely depend on $G_{\mathsf{d}}^{\mathsf{nece}}$— by describing the computation of each parameter in it using a sentence.

Before we proceed with the construction let us formally introduce:

Given any dependency graph $G_{\mathsf{d}}$,

In our final design of $G_{\mathsf{d}}$, we shall ensure that each parameter (except the special vertex $\mathsf{RNG}$) has in-degree at least $1$; however, during the construction process since we add edges incrementally, some (instance) parameter may temporarily have in-degree . For notation simplicity, we still say $\mathsf{op}_{G_{\mathsf{d}}}(a)=\max\{1,-1\}=1$ in such a case.

Hyperparameters. We use hyperparameters $1\leq n\leq m\leq s$ to control the difficulty of $G_{\mathsf{d}}$.

we shall ensure $\mathsf{op}(G_{\mathsf{d}}^{\mathsf{nece1}})\leq n$ and is as close as possible to $n$;

we shall ensure $\mathsf{op}(G_{\mathsf{d}}^{\mathsf{nece3}})=\mathsf{op}(G_{\mathsf{d}}^{\mathsf{nece2}})\leq m$ and is as close as possible to $m$;

we shall ensure $\mathsf{op}(G_{\mathsf{d}}^{\mathsf{nece}})=s$ is exact.

In other words, hyperparameter $s$ controls exactly how many operations are needed to compute the query parameter, which is the primary factor controlling the problem’s difficulty.

Given a structure graph $G_{\mathsf{s}}$, recall its edges represent all the instance parameters we shall use. Its abstract parameters are those ones that describe quantities across 1 or multiple layers: for instance in Figure 1, Central High’s number of Classrooms is across 1 layer, and Central High’s number of Backpacks is across 2 layers. We define this number as the difficulty level of abstract parameters.

With this notion, our construction of $G_{\mathsf{d}}^{\mathsf{nece1}}$ and $G_{\mathsf{d}}^{\mathsf{nece2}}$ are described together in Algorithm 2.

At a high level, we try to incrementally and randomly add abstract parameters to $G_{\mathsf{d}}^{\mathsf{nece1}}$ while maintaining $\mathsf{op}(G_{\mathsf{d}}^{\mathsf{nece1}})\leq n$. We cannot make this exact equality because when adding a single abstract parameter requires also (recursively) adding all the other parameters it may depend on. We tried to prioritize adding abstract parameters with higher difficulty levels. Once we finish constructing $G_{\mathsf{d}}^{\mathsf{nece1}}$, we randomly add additional instance parameters from $G_{\mathsf{s}}$ to make it $G_{\mathsf{d}}^{\mathsf{nece2}}$.

D.2.2 Construction of G𝖽𝗇𝖾𝖼𝖾𝟥fragmentsG𝖽𝗇𝖾𝖼𝖾𝟥G_{\mathsf{d}}^{\mathsf{nece3}}

Our goal next is to select a random $\mathtt{query}$ parameter in $G_{\mathsf{d}}^{\mathsf{nece2}}$ and construct a random topological ordering $\mathtt{Topo}$ for all the parameters in $G_{\mathsf{d}}^{\mathsf{nece2}}$, so as to ensure that all the parameters are necessary towards the computation of $\mathtt{query}$.

We start with $\mathtt{Topo}=[\mathtt{query}]$ and append parameters to its left one by one. During this process, we may also introduce new edges randomly; we start with $G_{\mathsf{d}}^{\mathsf{nece3}}=G_{\mathsf{d}}^{\mathsf{nece2}}$ and add edges incrementally. This process may not always succeed — sometimes the created topological ordering cannot make all the parameters necessary towards the computation of the $\mathtt{query}$. If this happens we declare a failure.The outside pseudocode, which comes later, shall go back to regenerate the structure graph and start again.

We introduce two notions (we use $G_{\mathsf{d}}^{\mathsf{nece3}}\setminus\mathtt{Topo}$ to denote the set of vertices in $G_{\mathsf{d}}^{\mathsf{nece3}}$ that are not in $\mathtt{Topo}$):

Intuitively, if $a\not\in\mathsf{Next1}(\mathtt{Topo})$ then we cannot immediately append $a$ to the front of $\mathtt{Topo}$, because it is not yet necessary towards the computation of $\mathtt{query}$.

Intuitively, if $a\not\in\mathsf{Next2}_{G_{\mathsf{d}}^{\mathsf{nece3}}}(\mathtt{Topo})$ then we cannot immediately append $a$ to the front of $\mathtt{Topo}$, because some other parameter depends on it and is not yet added to $\mathtt{Topo}$. (Obviously we always have $\mathsf{Next2}_{G_{\mathsf{d}}^{\mathsf{nece3}}}(\mathtt{Topo})\neq\varnothing$ unless $G_{\mathsf{d}}^{\mathsf{nece3}}\setminus\mathtt{Topo}=\varnothing$ so we are done.)

Our generation algorithm is now easy to describe: we keep adding parameters that are in $\mathsf{Next1}_{G_{\mathsf{d}}^{\mathsf{nece3}}}(\mathtt{Topo})\cap\mathsf{Next2}_{G_{\mathsf{d}}^{\mathsf{nece3}}}(\mathtt{Topo})$ to the front of $\mathtt{Topo}$; and if we get stuck, we introduce new edges to $G_{\mathsf{d}}^{\mathsf{nece3}}$ (or declare failure). The pseudocode is in Algorithm 3.