SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Guhao Feng $^{1,5,*}$  Bohang Zhang $^{2,*}$  Yuntian Gu $^{3,*}$  Haotian Ye $^{3,*}$  Di He $^{2,\boxtimes}$  Liwei Wang $^{2,4,\boxtimes}$
$^{1}$ School of EECS, Peking University  $^{2}$ National Key Laboratory of General Artificial Intelligence, School of Intelligence Science and Technology, Peking University  $^{3}$ Yuanpei College, Peking University  $^{4}$ Center for Machine Learning Research, Peking University  $^{5}$ Pazhou Lab fenguhao@stu.pku.edu.cn zhangbohang@pu.edu.cn guyuntian@stu.pku.edu.cn haotianye@pu.edu.cn dihe@pu.edu.cn wanglw@pu.edu.cn
Recent studies have discovered that Chain-of-Thought prompting (CoT) can dramatically improve the performance of Large Language Models (LLMs), particularly when dealing with complex tasks involving mathematics or reasoning. Despite the enormous empirical success, the underlying mechanisms behind CoT and how it unlocks the potential of LLMs remain elusive. In this paper, we take a first step towards theoretically answering these questions. Specifically, we examine the expressivity of LLMs with CoT in solving fundamental mathematical and decision-making problems. By using circuit complexity theory, we first give impossibility results showing that bounded-depth Transformers are unable to directly produce correct answers for basic arithmetic/equation tasks unless the model size grows super-polynomially with respect to the input length. In contrast, we then prove by construction that autoregressive Transformers of constant size suffice to solve both tasks by generating CoT derivations using a commonly used math language format. Moreover, we show LLMs with CoT can handle a general class of decision-making problems known as Dynamic Programming, thus justifying their power in tackling complex real-world tasks. Finally, an extensive set of experiments show that, while Transformers always fail to directly predict the answers, they can consistently learn to generate correct solutions step-by-step given sufficient CoT demonstrations.
Transformer-based Large Language Models (LLMs) have emerged as a foundation model in natural language processing. Among them, the autoregressive paradigm has gained arguably the most popularity [51, 9, 46, 69, 57, 16, 52, 54], based on the philosophy that all different tasks can be uniformly treated as sequence generation problems. Specifically, given any task, the input along with the task description can be together encoded as a sequence of tokens (called the prompt); the answer is then generated by predicting subsequent tokens conditioned on the prompt in an autoregressive way.
Previous studies highlighted that a carefully designed prompt greatly matters LLMs' performance [32, 38]. In particular, the so-called Chain-of-Thought prompting (CoT) [61] has been found crucial for tasks involving arithmetic or reasoning, where the correctness of generated answers can be dramatically improved via a modified prompt that triggers LLMs to output intermediate derivations. Practically, this can be achieved by either adding special phrases such as "let's think step by step" or by giving few-shot CoT demonstrations [34, 61, 56, 44, 70, 63]. However, despite the striking performance, the underlying mechanism behind CoT remains largely unclear and mysterious. On one hand, are there
indeed inherent limitations of LLMs in directly answering math/reasoning questions? On the other hand, what is the essential reason behind the success of  $\mathrm{CoT}^2$  in boosting the performance of LLMs?
This paper takes a step towards theoretically answering the above questions. We begin with studying the capability of LLMs on two basic mathematical tasks: evaluating arithmetic expressions and solving linear equations. Both tasks are extensively employed and serve as elementary building blocks in solving complex real-world math problems [10]. We first provide fundamental impossibility results, showing that none of these tasks can be solved using bounded-depth Transformer models without CoT unless the model size grows super-polynomially with respect to the input length (Theorems 3.1 and 3.2). Remarkably, our proofs provide insights into why this happens: the reason is not due to the (serialized) computational cost of these problems but rather to their parallel complexity [2]. We next show that the community may largely undervalue the strength of autoregressive generation: we prove by construction that autoregressive Transformers of constant size can already perfectly solve both tasks by generating intermediate derivations in a step-by-step manner using a commonly-used math language format (Theorems 3.3 and 3.4). Intuitively, this result hinges on the recursive nature of CoT, which increases the "effective depth" of the Transformer to be proportional to the generation steps.
Besides mathematics, CoT also exhibits remarkable performance across a wide range of reasoning tasks. To gain a systematic understanding of why CoT is beneficial, we next turn to a fundamental class of problems known as Dynamic Programming (DP) [6]. DP represents a golden framework for solving sequential decision-making tasks: it decomposes a complex problem into a sequence (or chain) of subproblems, and by following the reasoning chain step by step, each subproblem can be solved based on the results of previous subproblems. Our main finding demonstrates that, for general DP problems of the form (5), LLMs with CoT can generate the complete chain and output the correct answer (Theorem 4.7). However, it is impossible to directly generate the answer in general: as a counterexample, we prove that bounded-depth Transformers of polynomial size cannot solve a classic DP problem known as Context-Free Grammar Membership Testing (Theorem 4.8).
Our theoretical findings are complemented by an extensive set of experiments. We consider the two aforementioned math tasks plus two celebrated DP problems listed in the "Introduction to Algorithms" book [17], known as longest increasing subsequence (LIS) and edit distance (ED). For all these tasks, our experimental results show that directly predicting the answers without CoT always fails (accuracy mostly below  $60\%$ ). In contrast, autoregressive Transformers equipped with CoT can learn entire solutions given sufficient training demonstrations. Moreover, they even generalize well to longer input sequences, suggesting that the models have learned the underlying reasoning process rather than statistically memorizing input-output distributions. These results verify our theory and reveal the strength of autoregressive LLMs and the importance of CoT in practical scenarios.
An (autoregressive) Transformer [58, 50] is a neural network architecture designed to process a sequence of input tokens and generate tokens for subsequent positions. Given an input sequence  $s$  of length  $n$ , a Transformer operates the sequence as follows. First, each input token  $s_i$  ( $i \in [n]$ ) is converted to a  $d$ -dimensional vector  $\mathbf{v}_i = \operatorname{Embed}(s_i) \in \mathbb{R}^d$  using an embedding layer. To identify the sequence order, there is also a positional embedding  $\mathbf{p}_i \in \mathbb{R}^d$  applied to token  $s_i$ . The embedded input can be compactly written into a matrix  $X^{(0)} = [\mathbf{v}_1 + \mathbf{p}_1, \dots, \mathbf{v}_n + \mathbf{p}_n]^\top \in \mathbb{R}^{n \times d}$ . Then,  $L$  Transformer blocks follow, each of which transforms the input based on the formula below:
$$
\boldsymbol {X} ^ {(l)} = \boldsymbol {X} ^ {(l - 1)} + \operatorname {A t t n} ^ {(l)} \left(\boldsymbol {X} ^ {(l - 1)}\right) + \operatorname {F F N} ^ {(l)} \left(\boldsymbol {X} ^ {(l - 1)} + \operatorname {A t t n} ^ {(l)} \left(\boldsymbol {X} ^ {(l - 1)}\right)\right), \quad l \in [ L ], \tag {1}
$$
where  $\mathrm{Attn}^{(l)}$  and  $\mathrm{FFN}^{(l)}$  denote the multi-head self-attention layer and the feed-forward network for the  $l$ -th Transformer block, respectively:
$$
\operatorname {A t t n} ^ {(l)} (\boldsymbol {X}) = \sum_ {h = 1} ^ {H} \operatorname {s o f t m a x} \left(\boldsymbol {X} \boldsymbol {W} _ {Q} ^ {(l, h)} \left(\boldsymbol {X} \boldsymbol {W} _ {K} ^ {(l, h)}\right) ^ {\top} + \boldsymbol {M}\right) \boldsymbol {X} \boldsymbol {W} _ {V} ^ {(l, h)} \boldsymbol {W} _ {O} ^ {(l, h)}, \tag {2}
$$
$$
\operatorname {F F N} ^ {(l)} (\boldsymbol {X}) = \sigma \left(\boldsymbol {X} \boldsymbol {W} _ {1} ^ {(l)}\right) \boldsymbol {W} _ {2} ^ {(l)}. \tag {3}
$$
2Throughout this paper, we use the term CoT to refer to the general framework of the step-by-step generation process rather than a specific prompting technique. In other words, this paper studies why an LLM equipped with CoT can succeed in math/reasoning tasks rather than which prompt can trigger this process.
Here, we focus on the standard setting adopted in Vaswani et al. [58], namely, an  $H$ -head softmax attention followed by a two-layer pointwise FFN, both with residual connections. The size of the Transformer is determined by three key quantities: its depth  $L$ , width  $d$ , and the number of heads  $H$ . The parameters  $\boldsymbol{W}_Q^{(l,h)}, \boldsymbol{W}_K^{(l,h)}, \boldsymbol{W}_V^{(l,h)}, \boldsymbol{W}_O^{(l,h)}$  are query, key, value, output matrices of the  $h$ -th head, respectively; and  $\boldsymbol{W}_1^{(l)}, \boldsymbol{W}_2^{(l)}$  are two weight matrices in the FFN. The activation  $\sigma$  is chosen as GeLU [28], following [51, 21]. The matrix  $M \in \{-\infty, 0\}^{n \times n}$  is a causal mask defined as  $M_{ij} = -\infty$  iff  $i < j$ . This ensures that each position  $i$  can only attend to preceding positions  $j \leq i$  and is the core design for autoregressive generation.
After obtaining  $\mathbf{X}^{(L)} \in \mathbb{R}^{n \times d}$ , its last entry  $\mathbf{X}_{n;}^{(L)} \in \mathbb{R}^d$  will be used to predict the next token  $s_{n+1}$  (e.g., via a softmax classifier). By concatenating  $s_{n+1}$  to the end of the input sequence  $s$ , the above process can be repeated to generate the subsequent token  $s_{n+2}$ . The process continues iteratively until a designated End-of-Sentence token is generated, signifying the completion of the process.
Chain-of-Thought prompting. Autoregressive Transformers possess the ability to tackle a wide range of tasks by encoding the task description into a partial sentence, with the answer being derived by complementing the subsequent sentence [9]. However, for some challenging tasks involving math or general reasoning, a direct generation often struggles to yield a correct answer. To address this shortcoming, researchers proposed the CoT prompting that induces LLMs to generate intermediate reasoning steps before reaching the answer [61, 34, 56, 44, 70, 11]. In this paper, our primary focus lies in understanding the mechanism behind CoT, while disregarding the aspect of how prompting facilitates its triggering. Specifically, we examine CoT from an expressivity perspective: for both mathematical problems and general decision-making tasks studied in Sections 3 and 4, we will investigate whether autoregressive Transformers are expressive for (i) directly generating the answer, and (ii) generating a CoT solution for the tasks.
3 CoT is the Key to Solving Mathematical Problems
Previous studies have observed that Transformer-based LLMs exhibit surprising math abilities in various aspects [46, 10]. In this section, we begin to explore this intriguing phenomenon via two well-chosen tasks: arithmetic and equation. We will give concrete evidence that LLMs are capable of solving both tasks when equipped with CoT, while LLMs without CoT are provably incapable.
Arithmetic. The first task focuses on evaluating arithmetic expressions. As shown in Figure 1 (left), the input of this task is a sequence consisting of numbers, addition (+), subtraction (-), multiplication  $(\times)$ , division  $(\div)$ , and brackets, followed by an equal sign. The goal is to calculate the arithmetic expression and generate the correct result. This task has a natural CoT solution, where each step performs an intermediate computation, gradually reducing one atomic operation at a time while copying down other unrelated items. Figure 1 (left) gives an illustration, and the formal definition of the CoT format is deferred to Appendix B.
Figure 1: Illustrations of CoT on two math tasks.
Equation. The second task considers solving linear equations. As shown in Figure 1 (right), the input of this task is a sequence consisting of  $m$  linear equations, each of which involves  $m$  variables. The input ends with a special symbol  $\Longrightarrow$ . The goal is to output the value of these variables that satisfies the set of equations (assuming the answer exists and is unique). A natural CoT solution is the Gaussian elimination algorithm: at each step, it eliminates a certain variable in all but one equations. After  $m - 1$  steps, all equations will have only one variable and the problem is solved. Figure 1 (right) gives an illustration, and we defer the formal definition of the CoT format to Appendix B.
Number field. Ideally, for both tasks, the input sequences involve not only symbol tokens but also floating-point numbers. This complicates the definitions of the model's input/output format
and further entails intricate precision considerations when dealing with floating-point divisions. To simplify our subsequent analysis, here we turn to a more convenient setting by transitioning to the finite field generated by integers modulo  $p$  for a prime number  $p$ . Importantly, the finite field contains only  $p$  numbers (ranging from 0 to  $p - 1$ ) and thus can be uniformly treated as tokens in a pre-defined dictionary (like other operators or brackets), making the problem setting much cleaner. Moreover, arithmetic operations  $(+, -, \times, \div)$  are well-defined and parallel the real number field (see Appendix A.1 for details). Therefore, this setting does not lose generalities.
In subsequent sections, we denote by Arithmetic  $(n,p)$  the arithmetic evaluation task defined on the finite field modulo  $p$ , where the input length does not exceed  $n$ . Similarly, we denote by Equation  $(m,p)$  the linear equation task defined on the finite field modulo  $p$  with no more than  $m$  variables.
We begin by investigating whether Transformers can directly produce answers to the aforementioned problems. This corresponds to generating, for instance, the number "3" or the solution " $x = 1; y = 1; z = 0$ " in Figure 1 immediately after the input sequence (without outputting intermediate steps). This question can be examined via different theoretical perspectives. One natural approach is to employ the classic representation theory, which states that multi-layer perceptrons with sufficient size (e.g., the depth or width approaches infinity) are already universal function approximators [18, 35, 40]. Recently, such results have been well extended to Transformer models [67]: it is not hard to show that a constant-depth Transformer with sufficient size can solve the above tasks<sup>3</sup>. However, the above results become elusive when taking the representation efficiency into account, since it says nothing about the required model size for any specific task. Below, we would like to give a more fine-grained analysis of how large the network needs to be by leveraging the tool of complexity theory.
We focus on a realistic setting called the log-precision Transformer [42, 37]: it refers to a Transformer whose internal neurons can only store floating-point numbers within a finite  $O(\log n)$  bit precision where  $n$  is the maximal length of the input sequence (see Appendix A.3 for a formal definition). Such an assumption well-resembles practical situations, in which the machine precision (e.g., 16 or 32 bits) is typically much smaller than the input length (e.g., 2048 in GPT), avoiding the unrealistic (but crucial) assumption of infinite precision made in several prior works [49, 20]. Furthermore, log-precision implies that the number of values each neuron can take is polynomial in the input length, which is a necessary condition for representing important quantities like positional embedding. Equipped with the concept of log-precision, we are ready to present a central impossibility result, showing that the required network size must be prohibitively large for both math problems:
Theorem 3.1. Assume  $\mathsf{TC}^0\neq \mathsf{NC}^1$ . For any prime number  $p$ , any integer  $L$ , and any polynomial  $Q$ , there exists a problem size  $n$  such that no log-precision autoregressive Transformer defined in Section 2 with depth  $L$  and hidden dimension  $d\leq Q(n)$  can solve the problem Arithmetic  $(n,p)$ .
Theorem 3.2. Assume  $\mathsf{TC}^0\neq \mathsf{NC}^1$ . For any prime number  $p$ , any integer  $L$ , and any polynomial  $Q$ , there exists a problem size  $m$  such that no log-precision autoregressive Transformer defined in Section 2 with depth  $L$  and hidden dimension  $d\leq Q(m)$  can solve the problem Equation  $(m,p)$ .
Why does this happen? As presented in Appendices D.2 and E.2, the crux of our proof lies in applying circuit complexity theory [2]. By framing the finite-precision Transformer as a computation model, one can precisely delineate its expressivity limitation through an analysis of its circuit complexity. Here, bounded-depth log-precision Transformers of polynomial size represent a class of shallow circuits with complexity upper bounded by  $\mathsf{TC}^0$  [42]. On the other hand, we prove that the complexity of both math problems above are lower bounded by  $\mathsf{NC}^1$  by applying reduction from  $\mathsf{NC}^1$ -complete problems. Consequently, they are intrinsically hard to be solved by a well-parallelized Transformer unless the two complexity classes collapse (i.e.,  $\mathsf{TC}^0 = \mathsf{NC}^1$ ), a scenario widely regarded as impossible [65].
How about generating a CoT solution? We next turn to the setting of generating CoT solutions for these problems. From an expressivity perspective, one might intuitively perceive this problem as more challenging as the model is required to express the entire problem solving process, potentially necessitating a larger model size. However, we show this is not the case: a constant-size autoregressive Transformer already suffices to generate solutions for both math problems.
For example, given an input sequence of length  $n$ , an  $n$ -head self-attention layer with hidden dimension  $O(n)$  can extract the information of the entire sequence into the representation of the last position. Then, an MLP operating on the last position can universally approximate any (continuous) function over the input sequence.
Theorem 3.3. Fix any prime  $p$ . For any integer  $n > 0$ , there exists an autoregressive Transformer defined in Section 2 with constant hidden size  $d$  (independent of  $n$ ), depth  $L = 5$ , and 5 heads in each layer that can generate the CoT solution defined in Appendix B for all inputs in Arithmetic( $n, p$ ). Moreover, all parameter values in the Transformer are bounded by  $O(\mathrm{poly}(n))$ .
Theorem 3.4. Fix any prime  $p$ . For any integer  $m > 0$ , there exists an autoregressive Transformer defined in Section 2 with constant hidden size  $d$  (independent of  $m$ ), depth  $L = 4$ , and 5 heads in each layer that can generate the CoT solution defined in Appendix B for all inputs in Equation  $(m, p)$ . Moreover, all parameter values in the Transformer are bounded by  $O(\mathrm{poly}(m))$ .
Remark 3.5. The polynomial upper bound for parameters in Theorems 3.3 and 3.4 readily implies that these Transformers can be implemented using log-precision without loss of accuracy. See Appendix A.3 for a detailed discussion on how this can be achieved.
Proof sketch. The proofs of Theorems 3.3 and 3.4 are based on construction. We begin by building a set of fundamental operations in Appendix C that can be implemented by Transformer layers. Specifically, the softmax attention head can perform two types of operations called the (conditional) COPY and MEAN (Lemmas C.7 and C.8). Here, conditional COPY extracts the content of the unique previous position that satisfies certain conditions, while Conditional MEAN averages the values of a set of previous positions that satisfy certain conditions. These two operations can be seen as a form of "gather/scatter" operator in parallel computing. On the other hand, the FFN in a Transformer layer can perform basic computations within each position, such as multiplication (Lemma C.1), conditional selection (Lemma C.4), and lookup tables (Lemma C.5). With these basic operations as "instructions" and by treating autoregressive generation as a loop, it is possible to write "programs" that can solve fairly complex tasks. As detailed in Appendices D.1 and E.1, we construct parallel algorithms that can generate CoT sequences for both math problems, thus concluding the proof.  $\square$
Several discussions are made as follows. Firstly, the constructions in our proof reveal the significance of several key components in the Transformer design, such as softmax attention, multi-head, feedforward networks, and residual connection. Our proofs offer deep insights into the inner workings of Transformer models when dealing with complex tasks, significantly advancing prior understandings such as the "induction head" mechanism [45]. Moreover, our results identify an inherent advantage of Transformers compared to other sequence models like RNNs: indeed, as shown in Appendix F.2, constant-size RNNs cannot solve any of the above math tasks using the same CoT format. Secondly, we highlight that in our setting, the CoT derivations of both math problems are purely written in a readable math language format, largely resembling how humans write solutions. In a broad sense, our findings justify that LLMs have the potential to convey meaningful human thoughts through grammatically precise sentences. Finally, one may ask how LLMs equipped with CoT can bypass the impossibility results outlined in Theorems 3.1 and 3.2. Actually, this can be understood via the effective depth of the Transformer circuit. By employing CoT, the effective depth is no longer  $L$  since the generated outputs are repeatedly looped back to the input. The dependency between output tokens leads to a significantly deeper circuit with depth proportional to the length of the CoT solution. Note that even if the recursive procedure is repeated within a fixed Transformer (or circuit), the expressivity can still be far beyond  $\mathsf{TC}^0$ : as will be shown in Section 4, with a sufficient number of CoT steps, autoregressive Transformers can even solve P-complete problems.
4 CoT is the Key to Solving General Decision-Making Problems
The previous section has delineated the critical role of CoT in solving math problems. In this section, we will switch our attention to a more general setting beyond mathematics. Remarkably, we find that LLMs with CoT are theoretically capable of emulating a powerful decision-making framework called Dynamic Programming [6], thus strongly justifying the ability of CoT in solving complex tasks.
4.1 Dynamic Programming
Dynamic programming (DP) is widely regarded as a core technique to solve decision-making problems [55]. The basic idea of DP lies in breaking down a complex problem into a series of small subproblems that can be tackled in a sequential manner. Here, the decomposition ensures that there is a significant interconnection (overlap) among various subproblems, so that each subproblem can be efficiently solved by utilizing the answers (or other relevant information) obtained from previous ones.