Towards Revealing the Mystery behind Chain of Thought: A Theoretical Perspective
Authors: Guhao Feng
Bohang Zhang
Yuntian Gu
Haotian Ye
Di He
Liwei Wang
Venue: NeurIPS 2023 oral (2023)
Keywords: Chain-of-Thought Prompting
Large Language Models
Theory
Circuit Complexity
Dynamic Programming
TL;DR: This paper theoretically and empirically show the utility of CoT in LLMs.

Abstract: 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.

Towards Revealing the Mystery behind Chain of Thought: A Theoretical Perspective
Guhao Fengl.b. Bohang Zhang?.* Yuntian Gu Di He Liwei Wang2...
Haotian Ye
School of EECS, Peking University National Key Laboratory of General Artificial Intelligence,
School Intelligence Science and Technology. Peking University Yuanpei College, Peking University *Center for Machine Learning Research, Peking University PPazhou Lab eeeeluueedd...m edu. dihe@pku. edu.
fenguhaolstu, eku.edu. en haoti anye@pku. edu.
edu.
zhangbohang@pku.
guyuntianlstu, paunnttttuuuu... pku.edu edu. en wanglw@pku. hetpppuu.ed..... edu.
Abstract
Recent studies have discovered that Chain-of-Thougle prompting (CoT) can dra- matically improve the performance of Large Language Models (LL.Ms), particularly when dealing with complex tasks involving mathematics 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, examine the expressivity of LLMs with CoT in solving fundamental mathematical and decision- making problems. By using circuit complexity theory, 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 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 fecision-makinn 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 generate correct solutions step-by step-by-stey step given sufficient CoT demonstrations.
Introduction
Transformer-be 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 [5]] 9 46
54] based on the philosophy that all different tasks can be uniformly treated sequence generation problems. Specifically, given any task, the input along with the task description can be together encoded as sequence of tokens (called the prompt); the answer then generated by predicting subsequent tokens conditioned on the prompt in autoregressive way.
57
Previous studies highlighted that carefully designed prompt greatly matters LLMs' performance (8). In particular, the so-called Chain-of-Though prompting (CoT) [6 has been found crucial for tasks involving arithmetic reasoning, where the correctness generated answers be dramat- ically improved via modified prompt that triggers LLMs to output intermediate derivations. Practi- cally, this can be achieved by either adding special phrases such as "let's think step by step" or by giv- ing few-shot CoT demonstrations [34] 61 56] 44 70 63). However, despite the striking performance, the underlying mechanism behind CoTT remains largely unclear and mysterious. On one hand, are there
Equal contributions.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

indeed inherent limitations of LLMs in directly answering math/reasoning questions? On the other hand, what is the essential reason behind the success of Co Co boosting the performance of LLMs?
This paper takes a step towards theoretically answering the above questions. We begin with studying the capability of L.MS 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 [TO]. 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-pol super-polynomialyy with respect to the input length (Theorems [.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 wide range of reasoning tasks. To gain a systematic understanding of why CoT is beneficial, we next turn to fundamental class of problems known as Dynamic Programming (DP) 6 DP represents a golden framework for solving sequential decision decision-making tasks: it decomposes complex problem into 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) L.MMs 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 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%). 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.
Preliminary
is neural network architecture designed to process a sequence of input tokens and generate tokens for subsequent positions. Given an input sequence 8 of length Transformer operates the sequence as follows. First, each input token s; (i [n]) is converted to a d-dimensional vector = (mbed(s;) E Rd using an embedding layer. To identify the sequence order, there is also a positional embedding p; E TR'd applied to token 8;. The embedded input can be compactly written into a matrix ( = [D_ + ".. Un + pplTT E Raxd Then, L Transformer blocks follow, each of which transforms the input based the formula below:
An (autoregressive) Transformer [58, 50
X (!) = (I-1) + Attn (x (1 (1-1) + FFN() (x((- + Attn (1 (x (x(-))). LE [4], (1)
where (1) (1) and FFN() denote the multi-head self-attention layer and the feed-forward network for the 1-th Transformer block, respectively:
Atinn()) Σ A=1
softmax (WW())) Q
(L,)
(1)T + M) xw W(
(2)
- o (X)
'Throughout 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. 158, 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 I,, width d, and the number of heads
H. The parameters ((l.k) W (((.)) w-(.h) are query, key, value, output matrices of the K
Q
O
h-th head, respectively; and W Wi are two weight matrices in the FFN. The activation o is chosen as GeLU (28), following 21 The matrix M E [-30,0] 0]axa is a causal mask defined as Maj- = -o- oo iff i < j. This ensures that each position can only attend to preceding positions j S i and is the core design for autoregressive generation.
After obtaining ( E Raxd, its last entry X!! E R will be used to predict the next token S,+1 (e.g., via a softmax classifier). By concatenating 4n41 to the end of the input sequence s, the above process can be repeated to generate the subsequent token 4n+ The process continues iteratively until a designated End-of -f-eentence token is generated, signifying the completion of the process.
Chain-of- ofTThuhght prompting. Autoregressive Transformers possess the ability to tackle 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 70 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 will investigate whether autoregressive Transformers are expressive for (i) directly generating the answer, and (ii) generating a CoT solution for the tasks.
reasoning steps before reaching the answer 66 34 36
3 CoT is the Key to Solving Mathematical Problems
Previous studies have observed that Transformer-basedd 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.
3.1 Problem formulation
Arithmetic Expression
Linear Equations
Arithmetic. The first task focuses on evaluat- ing arithmetic expressions. As shown Fig- ure (left), the input of this task is sequence consisting of numbers, addition (+), subtrac- tion (-), multiplication (x), division (+), and brackets, followed by an equal sign. The goal is to calculate the arithmetic expression and generate the correct result. This task has a nat- urrl CoT solution, where each step performs an intermediate computation, gradually reduc- ing one atomic operation at time while copy- ing down other unrelated items. Figure (left) gives an illustration, and the formal definition of the CoT format is deferred to Appendix B
Input:
Input:
3r + 3g + 12: 6;
+ 5g 7
(7 + 5) + (6 44 2X
2r 4g 152
Output:
Output:
12 (6+ 3--1xx-2*4
+ 1y 4: 2:
12+ 12 2x))
62
3;
12 1 2 7)
2/ 7
2:
12 (18 14)
a 2: 1;
42
* 2= 1;
3
32 0:
1;
0:
Figure 1: Illustrations of CoT on two math tasks.
The second task considers solving linear equations. As shown in Figure D (right), the input of this task is a sequence consisting of m linear equations, each of which involves m variables. The input ends with special symbol = 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 certain variable in all but one equations. After m - 1 steps, all equations will have only one variable and the problem is solved. Figure [ (right) gives illustration, and defer the formal definition of the COT format Appendix B
Equation.
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 's input/output format
3

---

and further entails intricate precision considerations when dealing with Seating-poin divisions. To simplity our subsequent analysis, here we turn to a more convenient setting by transitioning to the finite field generated by integers modulo 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 (+, -, x, +) are well-defined and parallel the real number field (see Appendix A.I for details). Therefore, this setting does not lose generalities.
In subsequent sections, we denote by Arithmetic p) the arithmetic evaluation task defined on the fi- nite field modulo 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.
3.2 Theoretical results
We begin by investigating whether Transformers can directly produce answers to the aforemen- tioned problems. This corresponds to generating, for instance, the number "3" or the solution "x = 1; = 1; = 0" in Figure I Jimmediately after the input sequence (without outputting intermedi- ate 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 suffi- cient size (e.g.. the depth width approaches infinity) are already universal function approximators 18 35, 40). Recently such results have been well extended to Transformer models [67] it not hard to show that constant-depth Transformer with sufficient size can solve the above tasks 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- oeg-precision Transformer 42 37 former whose internal neurons can only store floating-point numbers within a finite Olog n) bit precision where is the maximal length of the input sequence (see Appendix A.3] for 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 imput length (e.g.. 2048 in GPT), avoiding the unrealistic (but crucial) assumption of infinite precision made in several prior works [49 20). Further- more, 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 central impossibility result, showing that the required network size must be prohibitively large for both math problems:
it refers to a Trans-
Theorem 3.1. Assume TC" + NC 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 and hidden dimension S (n) can solve the problem Arithmetic( p).
Theorem 3.2. Assume TC" f NC 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
(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 apply- ing circuit complexity theory (2). By framing the finite-precision Transformer as computation model, one can precisely delineate its expressivity limitation through an analysis of its circuit complexity. Here, bounded-depll log- precision Transformers of polynomial size represent a class of shallow cir- cuits with complexity upper bounded by TC 42). On the other hand, we prove that the complexity of both math problems above are lower bounded by NC' by applying reduction from NC² -complete prob- lems. Consequently, they are intrinsically hard to be solved by a well-parallelized Transformer unless the two complexity classes collapse (i.c., TC = NC1), a scenario widely regarded as impossible (65).
How about generating 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.
FFor example, given an input sequence of length n, an n-- head self att-nattenton layer with hidden dimension (n) can extract the information of the entire sequence into the representation of the last position. Then, an MLP operating the last position can universally approximate any (continuous) function over the input sequence.

---

Theorem 3.3. Fix any prime For any integer > O, there exists an autoregressive Transformer defined Section with constant hidden size d (independent of n), depth L = 5, and heads 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 (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 CoI solution defined in Appendix B for all inputs in Equation (m, p). Moreover, all parameter values in the Transformer are bounded by O(poly(m)).
Remark 3.5. The polynomial upper bound for parameters in Theorems 3.3 and 33.4 [344 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.d 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 form of "gaathrlscatter" operator in parallel computing. On the other hand, the FFN in Transformer layer can perform basic computations within each position, such as multiplication (Lemma C.I). 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.I and E. we construct parallel algorithms that can generate CoT sequences for both math problems, thus concluding the proof. O
and C.8 [.8) Here, conditional COPY extracts the content of the unique
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-heas feed- forward 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 Appendix [.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 [.Iaadd 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 significantly deeper circuit with depth proportional to the length of the CoT solution. Note that even the recursive procedure is repeated within a fixed Transformer (or circuit), the expressivity can still be far beyond TC": 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 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 core technique to solve decision-makin problems [55). The basic idea of DP lies breaking down a complex problem into series of small subproblems that can be tackled in 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.
5

---

Formally, a general DP algorithm can be characterized via three key ingredients: state space I, transition function T. and aggregation function A. Given a DP problem with N input sequences s(t),... s(A), denote the problem size be the vector n (,.."),.... (s())). ). Fixing the problem size n., there is an associated state space In C I representing the finite set of decomposed subproblems, where each state E In an index signifying a specific subproblem. The size of the state space In grows with the problem size n. We denote by dp(i) the answer of subproblem i (as well as other information stored the DP process). Furthermore, there is a partial order relation between different states: we say state j precedes state i (denoted as j < i) if subproblem j should be solved before subproblem i, i.e. the value of dp(i) depends on dp(j). This partial order creates directed acyelic graph (DAG) within the state space, thereby establishing reasoning chain where subproblems are resolved in accordance with the topological ordering of the DAG.
The transition function I' characterizes the interconnection among subproblems and defines how a subproblem can be solved based on the results of previous subproblems. It can be generally written as
dp(i) = T (n, i, {(i. dp(j)) : x)));
(4)
where is the concatenation of all input sequences s(1) s(N)) In this paper, we focus a restricted setting where each state only depends on (i) a finite number of tokens in the input
sequence and (ii) finite number of previous states. Under this assumption, into a more concrete form:
rewrite
dp(i) = f (n, (g(m))))) dp(h(n,,))) i)))
(5)
f (n, in (m..): ... (n,,)) dp(h1 (n, ddpp..ddppppiiiii (n, i)))
where functions f. f.g. h fully determine the transition function T and have the following form f: NN x I x x x )² + +y.g N XI -
h N x I (IU {0)) Here, the state space I, input space X. and DP output space y can be arbitrary domains, and J, K are constant integers. If state depends on less than J input tokens or less than K previous states, we use the special symbol 1 to denote placeholder, such that all terms Sy and dp(() are unused in function f.
After solving all subproblems, the aggregation function A is used to combine all results and obtain the final answer. We consider a general class of aggregation functions with the following form: dp(t)) i I, = u (DIeAA,
(6)
where An C In is set of states that need to be aggregated, O is an aggregation function such as min, max, or Σ, and u y + Z is any function where Z denotes the space of possible answers.
A variety of popular DP problems fits the above framework. As examples, the longest increasing subsequence (LIS) and edit distance (ED) are two well-known DP problems presented in the "Introduc- tion Algorithms" book [17] (see Appendix G.TToor problem descriptions and DP solutions). We list the state space, transition function, and aggregation function of the two problems in the table below.
Problem
Longest increasing subsequence
Edit distance
Two strings (1) (2) length n1 (1
Input
A string s length v
and n2 = (2) concatenated together ...... m1) ] [0.... .n2]
State space
(G.k):5 [n], k (0.... -j-1]]
ak bj Ap(j,A) min(dp(j, h: - 1) a, dp(j 1, k) b,
j=
== 0
k=
Transition function
max(dppij k -- 1).
dp(i, k)
dp(k: k- 1)xx
k
I[s s4 +1)+1)
dp(j 1, 1)
otherwise
cll(s (1) (
Aggregation
maxie[n] dp(i, 1)
dp(m1,
function
4.2 Theoretical results
We hegin by investigating whether L.MSs with CoT can solve the general DP problems defined above. We consider a natural CoT generation process, where the generated sequence has the following form:
| sNN)
(in. dp(())
(L/z.)) dp(i zm))) final answer

---

Here, the input sequence consists of N strings separated by special symbols, and (d1. .,,,,)) 1) is a feasible topologiaal ordering of the state space In. We assume that all domains I, X. y. Z belong to the real vector space so that their elements can be effectively represented and handled by a neural network. Each (i, dp(i)) E x y ahove will be represented as a single vector and generated jointly in the CoT output. We further assume that I, X, 3 2 are discrete spaces (e.g., integers) so that the elements can be precisely represented using finite precision. To simplify our analysis, we consider regression setting where each element in the CoT output directly corresponds to the output of the last Transformer layer (without using a softmax layer for tokenization as in Section Bp. Instead, the Transformer output is simply projected to the nearest element in the corresponding discrete space (e.g., I x y or 2). Likewise, each generated output is directly looped back to the Transformer input without using an embedding layer. This regression setting is convenient for manipulating numerical values and has been extensively adopted in prior works [24 I.
Before presenting our main result, we make the following assumptions:
Definition 4.1 (Polynomially-efficeen approximation). Given neural network Pe and target function f x + X2 where X C Rd and o C RR we say f can be approximated by Pa with polynomial efficiency if there exist p > 0, À > 0 such that for any error e > 0 and radius R > 0, there exists parameter 9 satisfying that (i) [x) - Po( + & < c A||ô||ce for all æ E X || S R and all || < pi (ii) all elements of parameter are bounded by O(poly 1/e)).
Assumption 4.2. The size of the state space can be polynomially upper bounded by the problem size n. i.e., |Im| = (ppoly((s)) (I))). Similarly, all input elements, DP values, and answers are polynomially upper bounded by the problem size n.
Assumption 4.3. Each function S. g. h and a in (5) and (6) can be approximated with polynomial efficiency by perceptron of constant size (with GeLUU activation).
Assumption 4.4. The function F : NN x I - I defined as F(n, it) = is+1 for n E N.V k E | 1] can be approximated with polynomial efficiency by a perceptron of constant size (with GeLU activation), where (i1, ... ) is feasible topological ordering of the state space In.
Assumption 4.5. The function F N I + {0, 1} defined as Fnn, i) U[ An] (see 6 can be approximated with polynomial efficiency by a perceptron of constant size (with GeLU activation).
Remark 4.6. All assumptions above are mild. Assumption 4.2 is necessary to ensure that the state vectors, inputs, and DP values can be represented using log- ing-prcisionn and Assumptions 4.3 to 4.5 guarantee that all basic functions that determine the DP process by a composition of finite log-precision Transformer layers of constant size. In Appendix G. we show these assumptions satisfied for LIS and ED problems described above as well as the CFG Membership Testing problem in Theorem 4.8
be well-approximated
We are now ready to present our main result, which shows that LLMs with CoT can solve all DP problems satisfying the above assumptions. We give proof in Appendix G.2
Theorem 4.7. Consider any DP problem satisfying Assumptions 442 to 4.5 For any integer n E N, there exists an autoregressive Transformer with constant depth L, hidden dimension d and attention heads H (independent of n)) such that the answer generated by the Transformer is correct for all input sequences of length no more than Moreover, all parameter values are bounded by Oppoly(1))
To complete the analysis, we next explore whether Transformers can directly predict the answer of DP problems without generating intermediate CoT sequences. We show generally the answer no: many DP problems are intrinsically hard to be solved by a bounded-depth Transformer without COT. One celebrated example is the Context-Free Grammar (CFG) Membership Testing, which take a CFG G and a string v as input and tests whether v belongs to the context-fit free language generated by G. A formal definition of this problem and a standard DP solution are given in Appendix G.1 We have the following impossibility result (see Appendix G.3 for a proof):
Theorem 4.8. Assume TC + NC². There exists a context-free language such that for any depth L and any polynomial Q, there exists a sequence length n E N where log-precision autoregressive transformer with depth L and hidden dimension S Q(n) can generate the correct answer for the CFG Membership Testing problem for all input strings of length n.
In contrast, enabling CoT substantially improves the expressivity of Transformers: as proved in Jones & Laaser 33 the universal CFG Membership Testing is a celebrated P-complete problem and intrinsically hard to be solved by well-parallelized computation model. Combined with these
results, we conclude that CoT plays a critical role in tackling tasks that are inherently difficult.
7

---

5 Experiments
In previous sections, we proved by construction that LLMs exhibit sufficient expressive power to solve mathematical and decision-making tasks. On the other hand, it is still essential to check whether a Transformer model can learn such ability directly from training data. Below, we will complement our theoretical results with experimental evidence, showing that the model can easily learn underlying task solutions when equipped with CoT training demonstrations.
5.1 Experimental Design
Tasks and datasets. We choose four tasks for evaluation: Arithmetic, Equation, LIS, and ED, The first two tasks (Arithmetic and Equation) as well as their input/CoT formats have been illustrated in Figure T For the LIS task, the goal is find the length of the longest increasing subsequence of given integer sequence. For the ED task. the goal is to calculate the minimum cost required (called edit distance) convert one sequence to another using three basic edit operations: insert, delete and replace. All input sequences. CoT demonstrations, and answers in LIS and ED are bounded-rand -ragge integers and can therefore be tokenized (similar to the first two tasks). We consider two settings: (i) CoT datasets, which consist of <problem, CoT steps, answer> samples; (ii) Direct datasets, which are used to train models that directly predict the answer without CoT steps. These datasets are constructed by removing all intermediate derivations from the CoT datasets.
For each task, we construct three datasets with increasing difficulty. For Arithmetic, we build datasets with different numbers of operators ranging from {4, 5, 63. For Equation, we build datasets with different numbers of variables ranging from {3, 4, 5). For LIS, we build datasets with different input sequence lengths ranging from (50, 80, 100}. For ED, we build datasets with different string lengths, where the average length the two strings ranges from {12, 16, 20). Creating these datasets allows us to investigate how model performance varies with the increase of problem size. We generate 1M sam- ples for each training dataset and M.IMM for testing while ensuring that duplicate samples between train- ing and testing are removed. More details about the dataset construction can be found in Appendix H
Model training and inference. For all experiments, we use standard Transformer models with hidden dimension d = 256, heads H = 4, and different model depths L. We adopt the Adam optimizer 39 with 33 0.9, 32 0.999, Ir 10-4, and weight decay 0.01 in all experiments. We use a fixed dropout ratio of 0.1 for all experiments to improve generalization. For CoT datasets, we optimize the negative log- liklililiooo loss on all tokens in the CoT steps and answers. This is similar to the so-called process supervision proposed in a concurrent work (36). For direct datasets, we optimize the negative log- likellhoo loss on answer tokens. All models are trained on 4 V100 GPUs for 100 epochs. During inference, models trained on the direct datasets are required to output the answer directly, and models trained on CoT datasets will generate the whole CoT process token-by by-token (using greedy search) until generating the End-of- Sooff-Sentence token, where the output in the final step is regarded as the answer. We report the accuracy as the evaluation metric. Please refer to Appendix H for more training configuration details.
5.2 Experimental Results
Main results. All results are shown in Figure 2 where each subfigure corresponds to a task with x-axis representing the difficulty level and y-axis representing the test accuracy (%). We repeat each experiment five times and report the error bars. In each subfigure, the purple bar and blue bars indicate the performance of the model trained on the CoT and direct datasets, respectively. The model depths are specified in the legend. From these results, one can easily see that 3-layer Transformers with CoT already achieve near- perfect performance on all tasks for all difficulty levels, which is consistent to Theorems [.3] 5.4 and 4.7 and may further imply that they can learn better solutions with fewer model depth than our constructions. Moreover, it is worth noting that, for some difficult datasets such as Equation of 5 variables, a perfect accuracy would imply that the model can precisely generate a correct CoT with a very long length of roughly 500 for all inputs. In contrast, models trained on direct datasets perform much worse even when using larger depths (particularly on the Equation task). While increasing the depth usually helps the performance of direct prediction (which is consistent with our theory), the performance drops significantly when the length of the input sequence grows. All these empirical findings verify our theoretical results and clearly demonstrate the benefit of CoT in autoregressive generation.
8

---

Arithmetic Expression
Linear Equation
100
100
Direct (Lm3)
90
Direct (L=4)
80
Direct
Rccuracy (%)
Accuracy (%)
CoT (L=3)
60
(L=3)
Direct (L=4)
Direct (L=5)
30
20
Number of Operators
Number of Variables
Longest Increasing Subsequence
Edit Distance
100
100
90
90
Accuracy (%)
80
Accuracy (%6)
80
70
60
60
Direct (L
Direct (L=3)
(L=4)
Direct
40
Direct (L=5)
CoT
(L=3)
CoT (L=3)
50
80
100
12
16
20
Input Sequence Length
Length of the First String
Figure 2: Model performance on different tasks. For all tasks and various difficulty levels, autore- gressive Transformers with CoT consistently outperform Transformers trained direct datasets. In particular, 3-layer Transformers already succeed in these tasks with almost perfect accuracy, while deeper Transformers
5) trained on the direct datasets typically fail.
Robustness to data quality. Unlike the synthetic datasets constructed above, real-world training datasets are not perfect and often involve corruption or miss intermediate steps. This calls into question whether the model can still perform well when training low-quality datasets. To investigate this question, we construct corrupted datasets for the arithmetic task Appendix []] Surprisingly, our results show that the 3-layer Transformer model can still achieve more than 95% accuracy even with 30% of the data missing an intermediate CoT step and involving a single-token corruption. This clearly demonstrates the robustness of CoT training low- low-quality datasets.
Length extrapolation. We finally study whether the learned autoregressive models can further extrapolate to data with longer lengths. We construct a CoT training dataset for the arithmetic task with the number of opera- tors ranging from to 15, and test the model on expres- sions with the number of operators in {16, 17, 18). As shown in Figure 3 our three-layer Transformer model still performs well longer sequences, suggesting that the model indeed learns the solution to some extent (in- stead of memorizing data distributions). Potentially, believe models trained on more data with varying lengths eventually reveal the complete arithmetic rules.
Arithmetic Expression Extrapolation
100
Accuracy (%)
20
Col (L=3)
16
18
Number Operators
Figure 3: Performance of length extrapola- tion experiment, tested on sequences that are longer than those in training.
Related Work
Owing to the tremendous success of Transformers and Large Language Models across diverse domains, there has been a substantial body of works dedicated theoretically comprehending their capabilities and limitations. Initially, researchers primarily focused on exploring the expressive power of Transformers the context of function approximation. Yun al. [67 proved that Transformers with sufficient size can universally approximate arbitrary continuous sequence-fo-sequenee functions on a compact domain. Recently, universality results have been extended model variants such as Sparse Transformers (68) and Transformers with relative positional encodings (RPE) [4].

---

More relevant to this paper, another line of works investigated the power of Transformers from computation perspective. Early results have shown that both standard encoder-decoder Transformers [58] and looped Transformer encoders are Turrnngcompllte [49
47
20 g
However, these results
depend on the unreasonable assumption of infinite precision, yielding a quite unrealistic construction that does not match practical scenarios. Recently, Giannou [25]
demonstrated that a constant-depth looped Transformer encoder can simulate practical computer programs. Wei et al. [60] showed that finite-precision encoder-decoder Transformers can approximately simulate Turing machines with bounded computation time. Liu et al. 37 considered a restricted setting of learning automata,
for which shallow non-recursive Transformer provably suffices. Yao et al. [66] demonstrated that Transformers can recognize or generate bounded- -eeth Dyck language [30], a specific type
of context-freee language. Besides affirmative results, other works characterized the expressivity limitation of Transformers via the perspective of modeling formal languages [26] 62 66 [4] or
However, none of these works (except (66]) explored the setting of autoregressive Transformers typically adopted in L.Ms, which we study in this paper. Moreover, we consider more practical setting that targets the emergent ability of LLMs in solving basic reasoning
simulating circuits (27 43 42
problems via a readable CoT output, which aligns well with real-world scenarios.
Recently, the power of Transformers has regained attention due to the exceptional in-context learn- ability exhibited by LLMs [9. Garg et al. [24] demonstrated that autoregressive Transformers can in-context learn basic function classes (c.g., linear functions, MLPs, and decision trees) via input sample sequences. Subsequent works further revealed that Transformers can implement learning algorithms such linear regression [.] gradient descent I [59 19.. and Bayesian inference [64], and a broad class of machine learning algorithms 3 The works of [23 45) studied in-context learning via the concept of "induction heads". All the above works investigated the power of (autoregressive) Transformer models from an expressivity perspective, which shares similarities to this paper. Here, we focus on the reasoning capability of Transformers and underscore the key role of CoT in improving the power of LLMs.
7 Limitations and Future Directions
In this work, from a model-capacity perspective, we theoretically analyze why Chain-offThouuht prompting is essential in solving mathematical and decision- decision-making problems. Focusing on two basic mathematical problems well as Dynamic Programming. we show that a bounded-depth Transformer without CoT struggles with these tasks unless its size grows prohibitively large. In contrast our negative results, we prove by construction that when equipped with CoT, constant-size Transformers are sufficiently capable of addressing these tasks by generating intermediate derivations sequentially. Extensive experiments show that models trained on CoT datasets can indeed learn solutions almost perfectly, while direct prediction always fails. We further demonstrate that CoT has the potential generalize to unseen data with longer lengths.
Several foundational questions remain to be answered. Firstly, while this paper investigates why CoT enhances the expressivity of LLMs, we do not yet answer how the CoT generation process is triggered by specific prompts. Revealing the relation between prompts and outputs is valuable for better harnessing L.MMs. Secondly, it has been empirically observed that scaling the model size significantly improves the CoT ability [61]. Theoretically understanding how model size plays a role in CoT would be an interesting research problem. Thirdly, this paper mainly studies the expressivity of LLMs in generating CoT solutions, without theoretically thinking about their generalization ability. Given our experimental results, we believe is an important future direction for theoretically studying how LLMs can generalize from CoT demonstrations (even in the cattof-distributiin setting, e.g.. length extrapolation (Figure (3) 60 [ Finally, from a practical perspective, it is interesting to investigate how models can learn CoT solutions when there are only limited CoT demonstrations in training (or even purely from direct datasets). We would like to leave these questions as future work, which we believe are beneficial to better reveal the power and limitations of L.Ms..
Acknowledgement. This work is supported by National Key R&D Program of China (2022ZD01 14900) and National Science Foundation of China (NSFC62276005), and partially funded by Microsoft Research Asia Collaborative Research Project. The authors are grateful to David Chiang, who pointed out mistake regarding the P-completeness of CFG Membership Testing the carly version of this paper. The authors also thank all reviewers for their valuable suggestions.
10

---

Author Contributions
Guhao Feng proposed to analyze the expressivity of Transformers using circuit complexity and proved all impossibility results in this paper, including Theorems 3.1 the initial idea of Dynamic Programming. He and Bohang Zhang formalized the DP framework and
3.2 and 4.8 He came up with
proved all positive results in this paper, including Theorems 3.3
3.4 and 4.7
He contributed to the
paper writing of Appendices[C Ct
Bohang Zhang supervised all undergraduate students. He raised the the problem of linear equation. He and Guhao Feng formalized the DP framework and proved all positive results this paper, including Theorems 3.3 3.4 and 4.7 In the experimental part, he helped generate datasets with Yuntian Gu, conducted
tuning, and finalized all experiments in Section 5 He was responsible for writing the majority of this paper and checking/correcting all proofs.
Yuntian Gu responsible for the experimental part. He wrote the entire code, including the model details, training pipeline, and evaluation. He created the datasets for all four tasks in Section [ and conducted extensive experimental exploration during this project (with the help of Bohang Zhang).
Haotian Ye participated in regular discussions, raised ideas, and helped check the code in experiments.
Di He initiated the problem of studying the capability of Transformers in basic tasks like evaluating arithmetic expressions. Di He and Liwei Wang led and supervised the research, suggested ideas and experiments, and assisted in writing the paper.
References
[1] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learn- ing algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations, 2023.
[2] Sanjeev Arora and Boaz Barak. Computational complexity: a modern approach. Cambridge University Press, 2009.
[3] Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers statis- ticians: Provable in-context learning with in-context algorithm selection. arXiv preprint arXiv:2306. 2366077706.46637, 2023.
[4] David Barrington. Bounded-width polynomial-size branching programs recognize exactly those languages in nc. In Proceedings of the eighteenth annual ACM symposium on Theory of computing, pages 1-5, 1986.
[5] David A Mix Barrington and Denis Therien. Finite monoids and the fine structure of ne. Journal of the ACM (JACM), -544)9 1988.
[6] Richard Bellman. The theory of dynamic programming. Bulletin of the American Mathematical Society, 5006):503-515, 1954.
[7] Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal. On the ability and limitations of trans- formers to recognize formal languages. In Proceedings of the 2020 Conference on Empirical Methods Natural Language Processing (EMNLP), pages 7096-7116, 2020.
[8] Satwik Bhattamishra, Arkil Patel, and Navin Goyal. On the computational power of trans- formers and its implications in sequence modeling. In Proceedings of the 24th Conference on Computational Natural Language Learning, pages 455-475, 2020.
[9] Tom Brown, Benjamin Mann. Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et Language models are few-shot learners. In Advances in neural information processing systems, volume 33, pages 1877-1901, 2020.
[10] Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, et al. Sparks of artificial general intelligence: Early experiments with GPT-4. arXiv preprint arXiv:2303. 12712, 2023.
11

---

[11] Collin Burns, Haotian Ye, Dan Klein, and Jacob Steinhardt. Discovering latent knowledge in language models without supervision. arXiv preprint arXiv:2212.03827) 2022.
[12] Samuel R Buss. The boolean formula value problem is in alogtime. In Proceedings of the nineteenth annual ACM symposium on Theory of computing, pages 123-131, 1987.
[13] Stephanie C.Y. Chan, Adam Santoro, Andrew Kyle Lampinen, Jane X Wang, Aaditya K Singh, Pierre Harvey Richemond, James McClelland, and Felix Hill. Data distributional properties drive emergent in-context learning in transformers. In Advances in Neural Information Processing Systems, 2022.
[14] David Chiang, Peter Cholak, and Anand Pillay. Tighter bounds on the expressivity of transformer encoders. Proceedings of the 40th International Conference on Machine Learning, pages 2023.
[15] Andrew Chiu, George Davida, and Bruce Litow. Division logspace-uniform ncl. RAIRO- Theoretical Informatics and Applications, 35(3):259-275, 2001.
[16] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sehastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.023/1 2022.
[17] Thomas H Cormen, Charles E Leiserson, Ronald L Rivest, and Clifford Stein. Introduction to algorithms. MIT press, 2022.
[18] George Cybenko. Approximation by superpositions of a sigmoidal function. Mathematics of control, signals and systems, 2(4):303-314, 1989.
[19] Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Zhifang Sui, and Furu Wei. Why can gpt learn in-context? language models secretly perform gradient descent as meta optimizers. arXiv preprint arXiv:22/2. 10559, 2022.
[20] Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Uni- versal transformers. In International Conference on Learning Representations, 2019.
[21] Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidircctional transformers for language understanding. In Proceedings of the 2019 Confer- ence of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Volume I (Long and Short Papers), pages 4171-4188 Association for Computational Linguistics, 2019.
[22] Benjamin L. Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and vari- able creation in self-attentio mechanisms. In International Conference on Machine Learning. PMLR, 2022.
pages
[23] Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dode Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. https:/transfommr-cccciiss http:///sssfommmrrrrriiriii////p http:a/ssuuuaoo//e/////oo///aea /wwp///wrk/idneee..inaa htmml.
[24] Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. In Advances in Neural Information Processing Systems, 2022.
[25] Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D Lee, and Dim- itris Papailiopoulos. Looped transformers as programmable computers. arXiv preprint arXiv:2301.13196 2023.
[26] Michael Hahn. Theoretical limitations of self-attention in neural sequence models. Transactions of the Association for Computational Linguistics, 8:156-171, 2020.
12

---

[27] Yiding Hao, Dana Angluin, and Robert Frank. Formal language recognition by hard atten- tion transformers: Perspectives from eircuit complexity. Transactions of the Association for Computational Linguistics, 10:800- 810, 2022.
[28] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint 1606. 08415, 2016.
[29] William Hesse. Division is in uniform tc0, In International Colloquium on Automata, Languages. and Programming, pages 104-114. Springer, 2001.
[30] John Hewitt, Michael Hahn, Surya Ganguli, Percy Liang, and Christopher D Manning. Rnns can generate bounded hierarchical languages with optimal memory. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 1978-2018 2020.
[31] IEEE Computer Society. leee standard for floating-pont point arithmetic. IEEE Std 754-2019, 2019.
[32] Zhengbao Jiang, Frank F Xu, Jun Araki, and Graham Neubig. How can we know what language models know? Transactions of the Association for Computational Linguistics, 8:423-438, 2020.
[33] Neil D Jones and William T Laaser. Complete problems for deterministic polynomial time. In Proceedings of the sixth annual ACM symposium on Theory of computing, pages 40-46, 1974.
[34] Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa. Large language models are zero-shot reasoners. In Advances in Neural Information Processing Systems, 2022.
[35] Moshe Leshno, Vladimir Ya Lin, Allan Pinkus, and Shimon Schocken. Multilayer feedforward networks with a nonpolynomial activation function can approximate any function. Neural networks, 6(6):861-867, 1993.
[36] Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's verify step by step. arXiv preprint 2.Xiv:2305.20055 2023.
[37] Bingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Trans- formers learn shortcuts to automata. In The Eleventh International Conference on Learning Representations, 2023.
[38] Pengfei Liu, Weizhe Yuan, Jinlan Fu, Zhengbao Jiang, Hiroaki Hayashi, and Graham Neubig. Pre-train, prompt, and predict: A systematic survey of prompting methods in natural language processing. ACM Computing Surveys, 55(9):1-35,, 2023.
[39] Ilya Loshchilov and Frank Hutter. Fixing weight decay regularization in adam. 2017.
[40] Zhou Lu, Hongming Pu, Feicheng Wang, Zhiqiang Hu, and Liwei Wang. The expressive power of neural networks: A view from the width. Advances in neural information processing systems, 30, 2017.
[41] Shengjie Luo, Shanda Li, Shuxin Zheng, Tie- Yan Liu, Liwei Wang, and Di He. Your transformer may not be as powerful as you expect. In Advances in Neural Information Processing Systems, 2022.
[42] William Merrill and Ashish Sabharwal. The parallelism tradeoff: Limitations of log-precision transformers. Transactions of the Association for Computational Linguistics, 2023.
[43] William Merrill, Ashish Sabharwal, and Noah A Smith. Saturated transformers are constant- depth threshold circuits. Transactions of the Association for Computational Linguistics, 10:843- 856. 2022.
[44] Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowyez, Maarten Bosma, David Luan, Charles Sutton, and Augustus Odena. Show your work: Scratchpads for intermediate computation with language models. In Deep Learning for Code Workshop, 2022.
13

---

[45] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernande7, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread, 2022. heads/index.htm.
[46] OpenAl. Gpt-4 technical report. arXiv preprint arXiv:303.08877, 2023.
[47] Jorge Pérez, Pablo Barceló, and Javier Marinkovic. Attention is turing complete. The Journal of Machine Learning Research, 22(1):
2021.
[48] Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference on Learning Representations, 2022.
[49] Jorge Pérez, Javier Marinković, and Pablo Barceló. On the turing completeness of modern neural network architectures. In International Conference on Learning Representations, 2019.
[50] Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. 2018.
[51] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog, (8):9, 2019.
[52] Jack Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:2112.11444 2021.
[53] Itiroo Sakai. Syntax in universal translation. In Proceedings of the International Conference on Machine Translation and Applied Language Analysis, 1961.
[54] Teven Le Scao, Angela Fan, Christopher Akiki, Ellie Pavlick, Suzana Ilié, Daniel Hesslow, Roman Castagné, Alexandra Sasha Luccioni, François Yvon, Matthias Gallé, et al. Bloom: A 1 rrbb-parameter open-access multilingual language model. arXiv preprint 1.05500 2022.
[55] Richard Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
[56] Mirac Suzgun, Nathan Scales, Nathanael Schärli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung, Aakanksha Chowdhery, Quoe V Le, Ed H Chi, Denny Zhou, et al. Challenging big- bench tasks and whether chain-of-thought can solve them. arXiv preprint 10.0926,, 2022.
[57] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anee Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 13971, 2023.
[58] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, volume 30, 2017.
[59] Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mord- vintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. arXiv preprint arXiv:2212.0761 2022.
[60] Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: a case study on approximating turing machines with transformers. In Advances Neural Information Processing Systems, volume 35, pages 12071-12088 2022.
14

---

[61] Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, brian ichter, Fei Xia, Ed H. Chi, Quoe V Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. In Advances in Neural Information Processing Systems, 2022.
[62] Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like transformers. In International Conference on Machine Learning, pages 11080-11090. PMLR, 2021.
[63] Noam Wies, Yoav Levine, and Amnon Shashua. Sub-task decomposition enables learning in sequence to sequence tasks. In The Eleventh International Conference on Learning Representa- lions, 2023.
[64] Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. In International Conference on Learning Representations, 2022.
[65] Andrew C Yao. Circuits and local computation. In Proceedings of the twenty-first annual ACM symposium on Theory of computing, pages 186-196, 1989.
[66] Shunyu Yao, Binghui Peng, Christos Papadimitriou, and Karthik Narasimhan. Self-attention networks can process bounded hierarchical languages. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume I: Long Papers), pages 3770-3785, 2021.
[67] Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequence functions? In International Conference on Learning Representations, 2020.
[68] Chulhee Yun, Yin-Wen Chang, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. O (n) connections are expressive enough: Universal approximability of sparse transformers. Advances in Neural Information Processing Systems, volume 33, pages 13783-13794, 2020.
[69] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained 2022.
transformer language models. arXiv preprint arXiv:2205.0111
[70] Zhuosheng Zhang, Aston Zhang, Mu Li, and Alex Smola. Automatic chain of thought prompting in large language models. In The Eleventh International Conference on Learning Representa- lions, 2023.
15

---

Appendix
The Appendix is organized follows. ( additional mathematical background and useful notations. Appendix B presents formal definitions and CoT solutions of the arithmetic expression task and the linear equation task. Appendix gives several technical lemmas, which will be frequently used in our subsequent proofs. The formal proofs for arithmetic expression, linear equation, and dynamic programming tasks given in Appendices[ and G respectively. Discussions on other architectures such as encoder-decoder Transformers and RNNs are presented in Appendix F Finally, we provide experimental details in Appendix
A
Additional Background and Notation
A.1
Finite field
Intuitively, a field is a set F on which addition, subtraction, multiplication, and division are defined and behave as the corresponding operations on rational and real numbers do. Formally, the two most basic binary operations in field is the addition (+) and multiplication (x), which satisfy the following properties:
Associativity: for any a, b, c F, (a - == a (b + c) and (a x b) x == a x (b x c);
Commutativity: for any o, b E F, a - b= b+ and x b b= b x
Identity: there exist two different elements 0, E F such that 0 and 1 for all a E F;;
Additive inverses: for any F, there exists an element in F, denoted as -a, such that a + (-a) = 0;
Multiplicative inverses: for any a E F and a * 0, there exists an element in F, denoted as a such that a x x - 1:
Distributivity of multiplication over addition: for any a, b, c F,ax = (axb)+
Then, subtraction (-) is defined by a-b b= = a + (-b) for all a, b E F; division (+) is defined by + b for all a, E F, o 0.
Two most widely-used fileds are the rational number field Q and the real number field R, both of which satisfy the above properties. However, both fields contain infinite number of elements. In this paper, we consider a class of fields called finite fields, which contain finite number of elements. Given prime number p. the finite field ž, is the field consisting of p elements, which can be denoted 1. In 2p. both addition and multiplication are defined by simply adding/multiplying
as 0, .,............... two input integers and then taking the remainder modulo p. It can be easily checked that the two operations satisfy the six properties described above. Thus, subtraction and division can be defined accordingly. Remarkably, a key result in abstract algebra shows that all finite fields with p elements are isomorphic, which means that the above definitions of addition, subtraction, multiplication, and division are unique (up to isomorphism).
As an example, consider the finite field Z5. We have that 2 + equals 0, since (2 + +3) mod = 0. Similarly, 2 3 equals 1; 2 equals 4; and 3 equals 4.
In Section 3 we utilize the field Zp to address the issue infinite tokens. Both tasks of evaluating
arithmetic expressions and solving linear equations (Section 3.1 are well-defined in this field.
A.2 Circuit complexity
In circuit complexity theory, there are several fundamental complexity classes that capture different levels of computation power. Below, we provide a brief overview of these classes; however, for comprehensive introduction, we recommend readers refer to Arora & Barak (2).
The basic complexity classes we will discuss this subsection are NC", AC TC NC1, and These classes represent increasing levels of computation complexity. The relationships between these
16

---

classes can be summarized as follows:
NC _ AC" E TC C NC C
Moreover, in the field of computational theory, it is widely conjectured that all subset relations in the hierarchy are proper subset relations. This means that each class is believed to capture a strictly larger set of computational problems than its predecessor in the hierarchy. However, proving some of these subset relations be proper remains a critical open question in computational complexity theory. For example, NC¹ P will imply P NP, which is widely regarded as impossible but is still a celebrated open question in computer science.
To formally define these classes, we first introduce the concept of Boolean circuits. A Boolean circuit with n input bits is a directed acyclic graph (DAG), in which every node is either an input bit or an internal node representing one bit (also called a gate). The value of each internal node depends on its direct predecessors. Furthermore, several internal nodes are designated as output nodes, representing the output of the Boolean circuit. The in-degree of node is called its fan-in number, and the input nodes have zero fan-in.
A Boolean circuit can only simulate a computation problem of fixed number of input bits. When the input length varies, a series of distinct Boolean circuits will be required, each designed to process a specific length. In this case, circuit complexity studies how the circuit size (e.g., depth, fan-in number, width) increases with respect to the input length for a given computation problem. We now describe each complexity class as follows:
NC is the class of constant-depth, constant-fan-in, polynomial-sized circuits consisting of AND, OR, and NOT gates. NCº circuits is the weakest class in the above hierarchy with limited expressive power because they cannot express functions that depend on a growing number of inputs as the input size increases. For example, the basic logical-ANN with an arbitrary number of input bits is not in NC In [22], the authors considered restricted version of the Transformer model with constant depth and a constant degree sparse selection construction, which can be characterized by this complexity class.
function
AC the class of constant-depth, unbounded -in--nn,, polynomial sized circuits consisting of AND, OR, and NOT gates, with NOT gates allowed only at the inputs. It is strictly more powerful than NC mainly because the fan-in number can (polynomially) depend on the input length. However, there are still several fundamental Boolean functions that are not in this complexity class, such as the parity function or the majority function (see below).
TC is an extension of AC that introduces an additional gate called MAJ (i.e., the majority). The MAJ gate takes an arbitrary number of input bits and evaluates to false when half or more of the input bits are false, and true otherwise. Previous work 42 43 showed that the log-precision Transformer is in this class.
NC² is complexity class that consists of constant-fan-in, polynomial-sized circuits with a logarithmic depth of (log where n is the input length. Similar to NC the basic logical gates are AND, OR, and NOT. Allowing the number of layers to depend on the input length significantly increases the expressiveness of the circuit. On the other hand, the logarithmic dependency still enables a descent parallelizability. Indeed, NC¹ is widely recognized as an important complexity class that captures efficiently parallelizable algorithms.
P is the complexity class that contains problems that can be solved by a Turing machine in polynomial time. It contains a set of problems that do not have an efficient parallel algorithm. For example, the Context-Free FreemGrrmmmr Membership Testing is in this class and is proved to be P-complete B3
A.3 Log-precision
In this work, we focus on Transformers whose neuron values are restricted to be floating point numbers of finite precision, and all computations operated on flottin-pooin numbers will be finally truncated, similar to how a computer processes real numbers. In practice, the two most common formats to store real numbers are the fixed-point format and floating-point format (e.g., the IEEE-754 standard (3I). Likewise, there are several popular truncation approaches (also called rounding), such as round-to-tinn the-------e--e-esstt. nearest, round-to-7ero, round-up, and round-down. Our results in this paper hold for both formats and all these truncation approaches.
17

---

Specifically, the log- preciisecion assumption means that we can use (loo((n) bits to represent a real number, where the length of the imput sequence is bounded by n. For any floating-point format described above with (n)gg(n) bits, an important property is that it can represent all real numbers of magnitude O(poly(n) (n)) within ((poly(1/n)) truncation error. We next analyze how the truncation error will propagate and magnify in log-precision Transformer from the input to the output layer. Note that since the functions represented by Transformers are continuous, the approximation error in a hidden neuron will smoothly influence the approximation error of subsequent neurons in deeper layers. This impact can be bounded by the Lipschitz constant of the Transformer, which depends on its basic layers. In particular, the softmax function (in attention) is I-Lipschitz, the GeLU activation is 2-Lipschitz, and the Lipschitz constant of a linear layer depends on the scale of its weight parameters. Combining these together leads to the following result: given bounded-deptte log-precision Transformer of polynomial size, when all parameter values of the Transformer are further bounded by (poly(nn)) all neuron values only yield an ((poly(1/n) approximation error compared to the rnfinite-prccision counterpart. Therefore, if a problem can be solved by a bounded- peenite-rrecisoo Transformer with selynomially-boundee parameters, it can also be solved by a log-prece fog-precision Transformer of the same size. This finding is helpful for understanding Theorems 3.3 3.4 and 4.7
depth polynomial-size
Finally, we point out that a key property of log- pre-piiecionn Transformer is that each neuron can only hold ((log(n))-an information and thus cannot store the full information of the entire input sequence. Therefore, the log- pog-rrccision assumption captures the idea that the computation must be somehow distributed each token, which well-resembles practical situations and the way Transformers work.
B Formal Definitions of CoT in Section 3
In this section, we will formally define the CoT derivation formats for the two math problems described in Section 3
Arithmetic expression. In an arithmetic expression that contains operators, there exists at least one pair of neighboring numbers connected by an operator that can be calculated, which we refer as a handle. More precisely, one can represent an arithmetic expression into a (binary) syntax tree where each number leaf node and each operator is an internal node that has two children. In this case, pair of neighboring numbers is a handle if they share the same parent in the syntax tree, For instance, consider the arithmetic formula
(6 + 5 4 x 5). Then, 6+ and x 5 are two handles.
An important observation is that we can determine whether a pair of numbers a and can form a handle with the operator op by examining the token before a and the token after b, where these tokens are either operators, brackets, or empty (i.e., approaching the beginning/ending of the sequence, including the equal sign '='). Specifically, given subsequence 81 op b 82, have that a op b forms handle iff one of the following conditions holds:
. op - and 8 { empty). $2 & (x,+);
. E +) and 81 & {x +J..
In the proposed chain of thought (CoT), an autoregressive Transformer calculates one handle at each step. If there are multiple handles, the leftmost handle is selected. The subsequence a op b is then replaced by the calculated result. For the case of s ( (and s2 ), there will be a pair of redundant brackets and thus the two tokens are removed. It is easy to see that the resulting sequence is still valid arithmetic expression. By following this process, each COT step reduces one operator and the formula is gradually simplified until there is only one number left, yielding the final answer.
System of linear equations. Assume that we have a system of m linear equations with variables + ...4
Tm. The i-th equation in the input sequence is grammatically written as a,14 + 0,242 - 0m = by where a;j E {o.-.. .P- - 1} and b, E {0. (0..... p 1}. For simplicity, we do not
omit the token 0:j or 0,y in the input sequence when a.y E (1.
We can construct the following CoT to solve the equations by using the Gaussian elimination algorithm. At each step i, we select an equation l satisfying the following two conditions:
* The coefficient of z; is nonzero.
* The coefficients of "1,-" " 24-1 are all zero
18

---

Such an equation must exist, otherwise the solution is not unique or does not exist. If there are multiple equations satisfying the above conditions, we choose the k-th equation with the smallest index k. We then swap it with equation i, so that the i-th equation now satisfy the above conditions.
We then eliminate the variable x; in all other equations by leveraging equation i. Formally, denote the i-th equation at the i-th step as
(4) d + 1 + aim Fm bi.
(4)
and denote the coefficient of x4 in the j-th equation (j * i) as Oj (a())-- >- (8) and add the resulting equation to the j-th equation. This will eliminate the term Fy in the j-th equation. We further normalize equation so that the coefficient a becomes 1.
(8) We can multiply
7 by
(1)
Depending on whether j S i or j > i, the resulting equation in the CoT output will have the following grammatical form:
j S i, the j-th equation will be written as 2, + 01.++17441 ........ ammmm +..ammmmm = =bj . If > i, the j-th equation will be written as 03.1417441 + +... + @jm:m = =by.
Note that remove all zero terms dz for lo S i, # in the CoT output and also remove the coefficient 1 in TR for l S i, similar to how human write solutions (see Figure | for an illustration). However, to simplify our proof, we reserve the coefficient 0 or I (i.e., outputting Ozg, or 1xx)) when k > i since cannot be determined easily before computing the coefficient. The above process is repeated for steps, and after the final step we obtain the solution.
C Technical Lemmas
C.1 Technical lemmas for MLP
In this subsection, we will demonstrate the representation efficiency of two-layer MLPs in performing several basic operations, such as multiplication, linear transformation, conditional selection, and look-up table. These operations will serve as building blocks in performing complex tasks.
We first show that a two-layer MLP with GeLU activation can efficiently approximate the scalar multiplication, with all weights bounded by (ppolyLL/e) where e is the approximation error.
Lemma C.1. Let f R22 R be two-layer MLP with GeLU activation, and the hidden dimension 4. Then, for any E > and M > O, there exist MLP parameters with fox norm upper bounded by Oppoy((M, 1/c)) such that (f b) abl holds for all a, b E [- M,MM. M).
Proof. Denote the input vector to the MLP as (a, b) E RR. After the first linear layer, it is easy to construct weight matrix such that the hidden vector is (a bb -a b, - b, -a + b) R4, where is an arbitrary scaling factor. Let a be the GeLU activation. We can similarly construct a weight vector such that the final output of the MLP is
f(a, 6) 22TA?
(<(2>))= +o a
((a+b)
(4-b)- ( (=a))(((=t))))
(-atb))-
a
8
We will prove that the above MI.P satisfies the theorem by picking appropriate A. By definition of GeLU activation, (Ix) <<x(x:) where (xx is the standard Gaussian cumulative distribution
function. We thus have o' d'(0) = 0.5 and d"(0)= = 22/T.. Applying Taylor's 's formula and assuming A 2M, we have
a
a
b)
Sab
6
0
X
X
X
X
v 22
2
+
b
2
a
a
2
a
+
b)
)
2
8ab
+
2
A
A
v 2
4
(2M)3
+
max
(3)
3!
X3
16M3
x2
80M3
max
3x3
4x)
exp(
313
xE[-1,1]
2
2
3 2x A3
19

---

Therefore.] b)-abl <
3e and then we can obtain [s b)-abl < Moreover, each weight element in the MLP is upper bounded by O()²), which is clearly O(poly 1/e)). L
33 10 Set X 2 10.²
Next, we will demonstrate that a two-layer MLP with GeLU activation can efficiently approximate two-layer MLP with ReLU activation, with all weights upper bounded by O(poly(1/e)) This result is useful in proving subsequent lemmas.
Lemma C.2. Let : Rdd + Rd be a two-layer ML. with ReLU activation, and all parameter values are upper bounded by M. Then, for any ε O, there exists two-layer MLP f of the same size with GeLU activation and parameters upper bounded by O(poly(M, 1/e)) in the &oo norm, such that for all a E RR,,, we have [f (z) g(=)]l C.
Proof. Let g(x) = W2 (cLU(W)a)) æ:). We construct f(x) - W2 GeLU(AW where À >> a sufficiently large constant. To prove that [[f(x) (())) so I e for all æ E R it suffices to prove that [W((RUUU()) Gel.UU(2)))l S e for z E RRd where d the hidden size. Since
W2 ReILU())
1 (oLLU(A))
< W ReIUU()) GoLU(A2))
.dd ReLU(z) GoLU(Az))
suffices to consider the scalar setting and prove that GeLU(())) - ReLU(u)) & e/ for all y E R.. By definition of ReLU and GeLU, have
GeLU(Xv) ReLU(v) =
1
ReLUUAg)
- -f-a Ay
Ay
exp
t²
dt
(8)
V2T
2
When y 0, becomes
00
Ay
22
Ay
Ay
t
+00
xy
t²
exp
di
exp
de
exp
dl.
2
2
2m
2
2
Combined with the case of
< 0.
can be consistently written as
Aa2l
t2
GoI.U(y ReI.U(v)
exp
dt
V2TT
2
Abol
+00
t²
1
(- X2y²
Lexp
dt
exp
(-
V2T
Alvl
V2TA
2
V2TA
Picking À Md. yields the desired result and completes the proof. V2ne
Equipped with the above result, perform linear transformation and conditional selection.
prove that two-layer MLP with GeLU activation can
Proposition C.3. Let f : RRd > Rd2 a two-layer MLP with GeLU activation, and the hidden dimension is 2dz. Let W Rd2xd, be any matrix and denote M = maxij [Wiz) Then, for any
e 0, there exist MLP parameters with foo norm bounded by O 1/e)), such that any Rdd we have [ Wolloo S.
Proof. We following construction:
two-layer MLP with ReLU activation implement g(#) Wa by the
Wx ReLU( + ReLU(- Wx)
Combined with Lemma C.2 we can also implement (ar) by a two-layer MLP with GeLU activation.
D
20

---

Lemma C.4. Define the selection function RR x R x R + R as follows:
g( y. t) = { x if 2 0,
(9)
y ift <
Let f F. x Rd x R + - Rd be two-layer MLP with GeLU activation, and the hidden dimension 2d +2 Then, for any c > O, a > O, and M > O, there exist MLP parameters with foo norm bounded by (ppoly(M,1/a, 1/c),, such that for all a E [-M, MJ y E I-MM, and t t I=000,
U la, too,, we have [., y.t) g(z, y, t)ll S.
Proof. We can simply use two-layer MLP with ReLU activation to implement by the following construction:
h(æ () = (h1, h2, hy ha) == (a + a M1 a Mi Mt, - M1) E 2d+2+
fax y, = Re)LU(h) ReLL(hs))l + (h2) - ReLU(h4)ld:
where 1, is the all-one vector of dimension. It is easy to check that, for all x E [-M, MJ y E
MM, Mids, and E [-oo, -ai U [a, +ool. we have f(x, y. = g(x, y. y. Morcover, all parame- ters are bounded by (MM/a). /a). Therefore, by using Lemma C.2 can also implement g(x) by
two-layer MIP with GeL.U activation and all parameters are bounded by ((poly 1/a, 1/c)). M
We final show that two-layer MLP efficiently represent look-up table. Consider a k- dimensional table of size d', where each element in the table is an integer ranging from I to d. Denote the set D fe, [d]]. where e; is d-dimensional one-hot vector with the i-th element being 1. The above look-up table can thus be represented as discrete function g DD - 'D. The following lemma shows that g can be implemented by a two-layer MLP with GeLU activation.
Lemma C.5. Let g D + DD be any function defined above, and let f. R&xd - Rd be two-layer MLP with GeLU activation and bias, and the hidden dimension is d' Then, for any > there exist MI.P parameters with to norm bounded by O(poly (k. 1/c)), such that for all x E DD C R
and all perturbation ô E 1-1/2k, 1/2kj/xd, we have [ 5) g(x) S E - 2k| ||oo0, where || is the vector Los -nom applied to the flattended matrix 5.
Proof. We can simply use a two-layer MLP with ReLU activation to implement g by the following construction. Denote the index of the MLP hidden layer (f1, (f,,"",, ; E [d]+ We can construct the weights of the first MLP layer such that
h (o) (z) = 2(x,, + Jatis ... + ((k 2k + 1.
We can then construct the weights of the second layer such that the final output of the MLP is
fz(x)
Σ
ReLU(h
9, (e )=1
One can check that f(x)= = g(x) holds for all a E DR c Raxk Furthermore, we have
[I + 6) - g(x) = f + 6) - f(x)
Σ
A0 + 8)) ReLU(h (x))
max
JE(4]
93 (ei Roi))
max (f1 E(d]
I (a) (x + 5) - hig di)) (x)) < |
for all perturbations ō E 11/22k 1/2k|axd. Thus by using Lemma C.2
also implement g(x))
by a two-layer MLP with GeLU activation and all parameters are bounded by Oppol((k, 1/c)).
C.2 Technical lemmas for the attention layer
In this subsection, we will introduce two special operations that can be performed by the attention layer (with causal mask). Below, let n E N be an integer and let 1 .2.... Ly, be a sequence of
vectors where 2; = (2,, ri, 1) E [-M, M] d-2 T; E Rd r; E IR, and M is a large constant. Let
K, V. E Rd' x (d+2) be any matrices with V Ixx < 1, and let 0 < 5 < M be any real numbers. Denote g = Qa,, k2 Kæ,,, n2 v,, and define the matching set S, = (j S i: (g: ky S p)- Equipped with these notations, we define two basic operations as follows:
21

---

COPY: The output is a sequence of vectors " with = ppoa(i), where pos(i)
argmaxjes; "j-
MEAN: The output is a sequence of vectors u1, U, with u, m ty-
The output u, is undefined when S; = 0. We next make the following regularity assumption:
Assumption C.6. The matrices Q. K, V and scalars p, ō satisfy that for all considered sequences
21, 22, En the following hold:
* For any i E [n]] either |q: ky S qi kg --8.
For any i.j E [n]] either - or (r Tyl 8.
Assumption C.6 that there are sufficient gaps between the attended position (e.g., pos(i)) and other positions. The two lemmas below show that the attention layer with casual mask can implement both COPY operation and MEAN operation efficiently.
Lemma C.7. Assume Assumption C.6 holds with p S 831 For any e > O, there exists an at- tention layer with embedding size Oga) and one causal attention head that can approximate the COPY operation defined above. Formally, for any considered sequence of vectors $1. 22,..... En denote the corresponding attention output as 01, O2,...,, OO. Then, we have l|o; uillo S for all i E [n] with S; * 9. Moreover, the too norm of attention parameters is bounded by (ppol((M, 1/8, log(n), log(1/e))).
Proof. The purpose of the attention head is to focus only on the vector that needs to be copied. To achieve this, we construct the key, query, and value vectors as follows (by assigning suitable key. query, and value weight matrices in the attention head):
Query: (Agss (1) RR+11
Key: (ki, r:) E Rd+1
Value: 8; E Rdd
where A and ju are constants that will be defined later. Denote dij as the attention score, then
exp(A(q: k2) + pry)
exp(A(q: kz))
Ot = ES exp(A(q ky ur..) En exp(X(q: k2) + plry r3))"
Since S SM and M > 8, we have 6 - p 1ō. By setting À = SM SM (n(2)))) 2nd
3 2n
and
(which are bounded by (ppoy((M, 1/8, log()) log(1/)))) we have
exp(-Ap)
ds.pos(1)
exp( (xp(-Ap)= + (n 1) exp(max Ao 2M Ap 118)
(10)
1+ (n 1) exp (omnax(-((6- p) 2M 2Ap nO0
> emx -X(6 p) + 2M 2hp -10))
(11)
MM
2n ) - In 2nn
- nexp
max
(-6 0
In
/2nM
2 cxx
In
(12)
2M
where in 10 we use Assumption C.6 which implies that whenever j' * pos(t), either qi ky S -
or (q. by S and ry- rú < 0); in I1b
we use the inequality 1 > 1 - for all x > 0: in 12
we use the fact that M 2 5. We thus have
(1-dppane))
I|o. uilloo =
Σ aiju, Upas(s)
<MIVV
+
Σ aid
sepos(()
MI -20),po((i) S.C.
which concludes the proof.
22

---

In( z.m For any O < < M, there exists an atlention layer with embedding size (dd) and one causal attention head that can approximate the MEAN operation defined above. Formally, for any considered sequence of vectors æ1; :2,-...... En. denote the attention output as 01, O2...... On Then, we have - S for all i E [n] with S o Moreover, the lo norm of attention parameters bounded by O(poly(M 1/6, log(n)) log(/e)))
Lemma C.8. Assume Assumption holds with p
6e
Proof. The purpose of the attention head is to average across all tokens that satisfy the condition qi kj R O. To achieve this, construct the key, query, and value vectors follows:
Query: Aq R
.
Key: k Rd
Value: 82 E IRd
where À is constant which will be defined later. Denote ai as the attention score, then
exp(>q; kj)
Ej exp(>q: ky))
1/6, By setting À In (AM)) (which is bounded by (10p0(((((,10..oo((()))0/(///))) we have:
(n (S exp(-10)
Σ jæS.
atj S
(13)
(n - S.() exp(=Xp + |S]
1 S. STT exp( (>(pp - -8)
< mexp(Ap) exp(-Add) -A6))
(14)
E
4Mn
S nexp
exp
-Inn In
16M
(15)
3M
where in (13) Assumption C.6 which implies that qi --yk S - for all j & S₁ and qi ky 2 -p
for all j E S;; in 14) we the inequality 14z 1 for all x > 0: in 15) use that the assumption that M and the fact that cx161/1 4/3.
Similarly, for any j E Si, we have
1
1
exp(-pX)
exp(pA)
0ij
S
max
(16)
S
[
[S exp((A)) + (n [S,]) exp --5) |S (exp(pp))
[S
1
max exp(2pA) + n cxp( A(6 p)) exp(2p2) -1)
|S
1
<
max (exp(2pA) + exp( 1(6 - p)), exp(2pd) - 1)
(17)
S
1
(cxp(2pA) - 1 + n exp( 2(6 p)))
|S,|
1
<
exp
-1+
(18)
[S,]
8MM
3M
(19)
3M
where in 16) we Assumption C.6,siillarly as before; in 17
use the inequality 1 1 < 1
for all x > 0; in (18) we use the incquality previously derived in (15) in (19 we use the inequality
exp(x) < 1 + 2x for all 0 < < 1. We thus obtain
||o; uolloo
Σ Faij
1
Σ JES,
2y
<
Σ JES.
oij +
Σ
Oij
1
Sc.
S
JES.
which concludes the proof.
23

---

D Arithmetic Formula
In this section, we prove that the autoregressive Transformer can evaluate arithmetic expressions when equipped with CoT, whereas the it cannot solve this task without COT.
D.1 Proof of Theorem 3.3
Before proving this theorem, there is one point that needs to be clarified: all residual connections in the attention/MI_P layers can be replaced by concatenation, in the sense that both architectures have the same expressive power. Formally, consider an MLP (or an attention layer) denoted as f F:R - RR, and let y = fx).. It is easy to see that we can construct another MLP (or attention layer) denoted as g IR2d + R such that g(x, 0) + (x, 0) (0, y) (xx, 0) (az, y), namely, the residual connection can implement concatenation. Conversely, concatenation can implement residual connection by using a linear projection. Based on the equivalence, we can use the concatenation operation instead of residual connection in all subsequent proof' presented in Appendices D Badd Similarly, the output of multi-head attention can be replaced by the concatenation of the output of
each head (instead of performing aggregation via matrices willk) defined in (2). For clarity, we further omit the unnecessary parts in the concatenated outputs and only retain the outputs that are used in subsequent layers.
We now present the proof of Theorem 3.3 For ease of reading, we restate Theorem 3.3 below:
Theorem D.I. For any prime and integer > O, there exists autoregressive Transformer defined in Section 2 with hidden size d = (poly(p) (independent of n., depth L 5, and 5 heads in each layer that 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(poly(n)).
Proof sketch. The intuition behind our construction is that when the CoT output proceeds a certain position, the Transformer can read the context related to this position and determine whether it should copy token or perform calculation. Remarkably, the context only contains a fixed number of tokens (as discussed in Appendix B). Based on the key observation, we can construct our five-layer transformer as follows. The first layer collects important positional information. The second and third layers determine whether to perform a calculation by examining the context related to the current token, which contains five tokens. The fourth layer and the fifth layers are used to generate the output via three cases: before/at/after the position that performs calculation. For the first and the last cases, the output simply copies a previous token with position computed by the two layers. For the middle case, the outcome is computed via a look-up table that stores the arithemtic rules (t, +).
Proof. We construct each layer as follows.
Token Embeddings. Assume that we have a sequence of tokens ......, 8, and we want to generate
the next token 8/+1. For any j E [m]] let id(s,) be the index of token sy the embedding dictionary. with values ranging from 1 to the number of tokens. We can embed the token 8j by x) (0) (eia(s,) J. 1) E ignum_Jukems=22 where e, is one-hot vector with the j-th element being 1, j E N4 is the positional embedding, and the constant embedding 1 is used as a bias term.
Layer 1. The first layer of the autoregressive Transformer uses two attention heads to perform the following tasks:
1. Count the number of equal signs ('=') in previous tokens, denoted as n² i.e., n = =Ki- S i:
8)
'.
2. Copy the position of the last equal sign, denoted as iie, P = max(j - = If the set (j j < i, 8j '='} is empty, define p. 0.
3. Compute ??
Based on Appendix (Lemma C.8
can use the first attention head to perform the MEAN
operation that counts the percentage of equal signs in the preceding sentences (i.e.,,, /i). This can
be achieved by setting Q = 0, K = 0, V = (ea())) 0, 0) (defined in Appendix C.2 so that
q: 0, k2 0, 2y eld(=) eid(s,) L[ss = '='], S, = [..
24

---

Similarly. we can use the second attention head to perform a COPY operation that copies the position index of the last equal sign (by Lemma C.7 This can be achieved by setting Q (0.0, 0, 1)
K (Eid(=)) 0, 1) V = (0, j (defined in Appendix C.2 so that
9: 1, = ==']-1 1, by = S Si::s:) =
It is easy to check that the above construction outputs 24, max j : 1, '=' when S o D. Note that i, may not equal to p;- when S; = 0.
Using the residual connection to perform concatenation, the output of the attention layer has the form (Eia(s,): i, 1, n, /t. max j S s. =' We can then use an MLP multiply n; /i and to obtain nz and use another MI.P to compute 12 according to Lemma C.1 Simultaneously, we can
compute the value Pj using the following way:
max j S s :===== ='}
if nf /1 /1 1/n, if n5 /1 / 0,
Pt
which is a conditional selection operation and can be implemented by an MLP (Lemma C.4 Also note that the gap in Lemma CammaaC.ijpa is a 1/2ni, which can be implemented within log-precision. The final output of the first layer has the form x; (1) = (eia(s.); i, i, m .PT,1) 1).
Layer 2. The second layer of the Transformer does some tricky preparation work for the next layer.
1. Compute the distance to the nearest and the last equal sign, denoted as di and di. respectively. Formally, dj
max j S i, 8 '='), di i max : i, 8j '='}. If the nearest/last equal sign does not exist, define di or d; = i. The relation between di d7. and pF can be expressed as di = i PT d. = PP-11
2. Count the number of equal signs in strictly previous tokens, denoted as n, i.e., n. = 1g <
s; = "
3. Compute (n.)², (n. (d;)?, and (d;
The first and the second tasks can be done using the COPY operation by setting
q. = Qm (1) = i-,,,,, kj = K (1) = (-1,2j, Tj = 0, 0; = (n --p +1).
Under the above construction, we have q kj = -i 1)2, and thus S {i 1}. namely, the output (ñ,, d We then use an MLP calculate (d5)2 (2 (n ) and (n 2 by using Lemma
The output of the second layer is
a;
(2) (ela(e)) i, n., AT d5 dF. (n,
(n
( (
Layer 3. The third Transformer layer judges whether the calculation should be performed at the current position and computes the result when needed. Based on the CoT format given in Appendix we need to extract five previous tokens related to this position. Formally, we need five attention heads to perform the following tasks:
1. Copy the embedding esd(sy) located at position j such that n7 n,j - and dj di for E {1, 3, 4, 5). as shown in Figure
2. Check if the copied expression can be evaluated at the current position according to the rule given in Appendix B If it can be evaluated, compute the result and determine how much sentence length will be reduced after this calculation (see Appendix B for details on how the reduced sentence length depends brackets); otherwise, keep the token Gid(sy) with
カT n and dj d 1
We can use the multi-head attention to perform the COPY operation five times in parallel. For each t, we construct the matrices Q. K, V the COPY operation such that
qi Qx (2) (n,)2 2n + 1,
1,
n
(d 2td, + t²
1.
d - t]
k, Kx ; =|
-1
(n
2n
1,
-(d7)2
2d L
vj Cid(e;);
25

---

dj
di
1 1+ 2 xi
(3
4)
++2
+
Figure 4: IIlustration of the proof of Theorem 3.3
and
(2) x
Qx)
(mn n2 - -12- (dī - dj + +1)².
Therefore, Ka Ka,
Q2
only when 27 n 1 and di dī t, and Kx, Qm -1
otherwise. It is easy to see that:
Whenever n. 0, for any t, the number of indices j satisfying qi kj 0 is at most one (i.e., unique).
Whenever n, > 0, for any f, the index j satisfying q: ky 0 exists, unless there is a t' < such that the copied token at an equal sign ('=').
In other words, based on Lemma C.7 the above property guarantees that we can copy the desired tokens until reaching equal sign, after which the copied tokens are invalid as illustrated in Figure4rris right). The output the attention layer can be written as
(Cid(a,)) Gj Gjz ejs: Eja: CJ i, (m,)) di. ( (hT) ,d,d, (d.)).
where we slightly abuse notation and use ej. to denote the embedding we copied by the t-th attention heads.
We can then use an MLP to perform the second task. Note that whether the current position can be calculated or not depends on the following six tokens (Cia(ai), Ej Gjz Gjs: Cy., ej,). Concretely, there are several cases:
esd(sy) corresponds to the embedding of a number or a right bracket. In this case, the current position should simply output ej (which is operator right bracket).
Gid(s,) corresponds to the embedding of left bracket, an operator, or the equal sign '='. In this case, ej, corresponds to the embedding of number or a left bracket. There are two subcases:
Ej corresponds to the embedding of a number. In this case, whether the current posi- tion can be evaluated depends on (Eia(s.) Gj, Cju Cjs e,,) according to Appendix B Ej, corresponds to the embedding of left bracket. In this case, whether the cur- rent position can be evaluated simply depends on whether ej corresponds to the embedding of a right bracket.
When all embeddings Cj, are one-hot vectors, whether the expression at the current position can be calculated or not forms look-up table. Therefore, it can be implemented by two-layer MLP with hidden dimension (z))) according to Lemma C.5 Similarly, the computed result and how much sentence length will be reduced after this calculation can also be implemented as look-up tables. However, some the embeddings ej. may not be one-hot vectors when reaching an equal similarly implement extra look-up up tables that take
sign (as discussed above). In this case, fewer number of inputs, and the result which lookup table will be used depends on the position of the equal sign. This corresponds to multivariate conditional selection operation with multiple
Boolcan conditions I[e,, +Eid(=") 1 (for E 21, 3, 4, 5)), which can be similarly implemented by an MLP by extending Lemma
Moreover, note that the composition of look-up tables and the multivariate conditional selection operation can be merged in just one MLP by following the
construction in Lemmas C.4 and C.5 (we omit the details for clarity).
The final output of the third layer is represented by
(3) (ej, n4 (mz 2 ., (n, ) (d, 2 outcome _1,
reduce
Here, fa is a Boolean value recording whether the next output at position i is a computed value, e eoutzomme is the one-hot embedding of the outcome when is true, and n reduce records the reduced
length after calculation when fi is true. When fi is false, e, eoucoome and n. redues undefined.
26

---

Layer 4. Note that in an arithmetic expression there can be multiple expressions that can be calculated (or handles defined in Appendix BB) all of which are processed in the last layer. Therefore, the fourth layer of the Transformer should keep only the leftmost calculation and discard other calculations (according to Appendix B). Meanwhile, for subsequent positions i after the position that has been calculated, this layer finds the related token that should be copied for position i based on the reduced setence length. Formally, we need two attention heads to perform the following tasks:
1. Check whether there is an index j S i such that nj n. and fj is true. Denote the answer as f. Ej= fi. where f. 2 1 means the answer is yes, and f: = 0 otherwise.
2. If the answer is yes (f: 2 1), copy the value necduce the leftmost position satisfying fy is true and nj n.. Denote the result as iij n
ni If f. 0, figeluce is undefined.
3. Filter the outcome: if f: is true, then maintain eq otherwise set e, eqaccome to Cj
Similar to the construction of the third layer, can construct the matrices Q. K and V follows. For the first task, we leverage the MEAN operation with
q. = (3) (1, .m Kn Ka
(3) (-
f/ fz.
We have q; k, (n," n p. Therefore, qi k; 0 iff n.j = and qi kj -1 otherwise. This attention head thus outputs प. Ej=r-d fz. For the second task, we leverage the COPY operation with
q: = Qz; - (1, (nz 2n 1,, 1), kj Kæ (3) = (-
f1: -1: Uj ny
reduce
We have q: k; = (n, n n² fj Therefore, q: kj = 0 iff nj = n and fj = and q k; < -1 otherwise. Moreover, we set r- -j in the COPY operation, by which the attention head copies n nudduce where j [j nj = n, fj = 1}. as desired. The output of the attention layer has the form
f di +
outcome
Bii (m ,, (17) di .F.,, Si,
(dF)
reduce
We next use an MLP to perform the third task, which is a conditional selection operation and can be done according to Lemma C.4 We can simultaneously obtain ft by multiplying 4=+I with (df + 1).
We also compute (dī + f reulcce
) which will be used in the next layer. The final output of the fourth
layer is represented by
(Egutcome
f.. (n (nT) nT, (n,) 2, . fr h. reduce (F, + fi
(reduce)? ??,,
x
where aqutcome either e.utcomm or Ej--
Layer 5. The final layer of the Transformer uses one attention head to copy the corresponding token for generating the output when f. 2 and f, is false. Similar to previous layers, can copy the embedding Eid(e;) located at position j such that n n 1 and di di + ,iedue,, The output of the attention layer is (Zoutcome Ci. ft). We then use an MLP to obtain the output: if f 2 1, then output ei otherwise output equicome, This corresponds to a conditional selection operation and can be implemented by an MLP according to Lemma C.4 Finally, pass the output through a softmax layer to generate the next token $,+1-
Now it remains to conduct an eгror analysis and determine the scale of parameters. Note that we can tolerate (11) error of the final layer output in the sense that the generated token 8;41 is still correct. Based on Lemmas C.I C.4 C.5
C.7 and C.8 we can guarantee that when all parameters of the Transformer are bounded by Oppolynn, 1/e)), all intermediate neurons will only induce an error below (Also note that Assumption C.6lin Lemmas C.7 and C.8 is satisfied when is small
enough.) Therefore, by picking fixed small ((1), all parameter values in the Transformer are bounded by (poly(n)))
D
27

---

D.2 Proof of Theorem 3.1
We now prove that evaluating arithmetic expressions without CoT is extremely difficult for bounded- depth autoregressive Transformers. We will make the widely-believed assumption that TC & NC² (see Appendix [A.2] for definitions of these complexity classes). We further need the notion of uniformity: informally, this condition says that there exists an efficient algorithm to construct the circuits. For rigorous definition, we refer readers to Arora & Barak [2].
We first present some lemmas on the expressive power of the TC" circuits, based on which we will give reduction proof of Theorem 3.1
Lemma D.2. For any n. n-bils binary integers 01, 01,02,".. ". ans let f {0, 13, {0. 1) be boolean function such that f(11, "g,- ,an) an) = -ie(j) ai Then, f can be implemented by the uniform TC circuits.
This lemma demonstrates that computing the sum of n n-bit integers (known as iterated addition or simply summation) is in uniform TC The detailed proof of this lemma can be found in the previous works 29 [5].
Lemma D.3. Consider any string s = $182 5.. Sn of length n containing brackets '(', '), and other characters, and all brackets in are paired. Let f be boolean function taking as input and output pairs of integers defined as follows:
(-1,5)
if s; a left bracket and 8t, 83 are paired.
f((s) =
(1)
if s; is a right bracket and s,, 8 are paired.
kj
if s; is not a bracket, and 8j, 8k is the nearest pair of matching brackets containing
Then f can be implemented by the TC circuits.
Note that the input characters and the output integers are all encoded in binary. This lemma demon- strates that the task of bracket matching is in TC
Proof. Given the input string 8, we first get the index i of each character 8,, which can be hard-coded into the circuits. Then for each character 8;, calculate the following term:
Σ Mss [(] M(s, ()) I[s '().
Then, the output of f((s) can be written as the following form:
(-1, min > i: sy = Y,,r, = + fi(s) (max(j < 8;:85- = 'C,ry=:ph - -
if = 'C if 'Y,
(max(j < <i:s;= =C,rr = r min >i = - s; not bracket.
is simple to verify that the construction can be implemented by TC cireuits.
Theorem D.4. Assume TC + uniform NC!. For any prime number p. any integer L, and any poly- nomial Q there exists problem size such that no log-precision autoregressive Transformer defined in Section 2 with depth L and hidden dimension d S Q(n) can solve the problem Arithmetic p).
Proof. Our proof is based on leveraging the NC¹ completeness of a classic problem: Boolean Formula Evaluation. According to the Buss reduction 02). calculating whether Boolean formula is true false is complete for uniform NC'. Based on this theorem, it suffices to prove that the Boolean Formula Evaluation problem can be reduced to evaluating the arithmetic expression. This will yield the conclusion by using the result that bounded-depth log-precision Transformers with polynomial
size are in TC 42 as well as the assumption that TC * uniform NC!
Formally, let Σ 10, v, (.) be the alphabet. A Boolean formula is string defined on alphabet Σ by the following recursive way:
"While the authors only proved that predicting a binary label using Transformer encoder can be implemented by TC circuits, it is straightforward to extend the result to autoregressive Transformers predicting label in a finite vocabulary.
28

---

0 and 1 are Boolean formulae;
If yp is Boolean formula, then (-p) is Boolean formula;
If (Pi, P22 are two Boolean formulae, then both (5p1 A A42) and (501 V (22) are Boolean formulae.
The Boolean Formula Evaluation problem aims to compute whether a Boolean formula is true (1) or false (0). We now show that we can translate this problem into the problem of evaluating arithmetic expressions, Given a Boolean formula s, the translation function f generates the corresponding arithmetic expression (s) that has the same result under evaluation. The translation is recursively defined as follows:
f(0) 0 and f(1) =
For any Boolean formula yp, /(-4)= and [((+)) = (f(i)):
For any Boolean formulae (P1, 92, f(p1 A 92) = f(91) x f(p2);
For any Boolean formulae PLL, 92, f(pi (22 1 f(pi) x xI-
is easy to see that for any Boolean formula s, the length of f(s) upper-bounded by O(Is)). More- over, the translation function can be simply implemented using the circuits within TC complexity. To do so, we first replace the symbols -, A, and V with 1-, x, and x. respectively. Furthermore, for each operator V, we must insert - (1- after the nearest left bracket containing the operator V, insert a right bracket before x, insert '(1-¹ after x, and insert right bracket before the nearest right bracket containing the operator V. According Lemmas D.2 and and D.3 all of these operations can be implemented by the TC circuits. Therefore, this translation function can be simply implemented by TC" circuits. Also, note that the above construction does not depend on the modulus p. Therefore, by reduction, we obtain that the problem of evaluating arithmetic expressions is -hadh
E System of Linear Equations
In this section, will prove that the autoregressive Transformer equipped with CoT can solve system of linear equations, whereas the autoregressive Transformer without CoT cannot solve it.
E.1 Proof of Theorem
For ease of reading, we restate Theorem 3.3 below:
Theorem E.1. For any prime p and integer m > there exists an autoregressive Transformer defined in Section 2 with hidden size d = (ppoly()) (independent of m.) depth L = 4, and 5 heads in each layer That can generate the CoT solution defined in Appendix Blor all inputs Equation(m, p)- Moreover, all parameter values in the Transformer are bounded by (pply(m)).
We recommend readers to read the proof Theorem 3.3 first as we will omit redundant details in the subsequent proof. Below, without abuse of notation, we use x to denote the output at position i after the 1-th Transformer layer, and use 2, denote the i-th variable linear equations. We also note that m the upper bound on the number of variables, and we will construct Transformer parameters such that the Transformer can solve all linear equations with the number of variables no more than m.
Proof. The proof technique is similar to that of Theorem 3.3
St and we want to Token Embeddings. Assume that we have a sequence of tokens 81, 82, (0) (eid(s.) 4, i, 1): generate the next token s1+1 We can embed the token s; using the format x;
1. The vector e; represents the one-hot vector with the i-th element being 1, and id(s;) is the index of token s; in the vocabulary. Since we hope the embedding dimension is a constant and does not depend on the number of variable tokens m, we consider representing them using unified (single) encoding and distinguishing them via the term and sy are two different variables, we have id(s;) = id(s;) and I, * L,-
4, This means that
1 E R is vector used to distinguish between different variables. Its first element, denoted var(s;), represents the index of the variable s,. If the token 8; is not variable, then
29

---

4 (0, 0) and var(s;) O. If it is the variable Iy for some E [m], then var(s;) and
4 (j. m² sin( ) m² cos( 24 ))
3. is the positional embedding, representing the position of the token in the sequence.
4. The constant embedding is used as a bias term.
Layer 1. The first layer of the Transformer uses three attention heads to record some basic informa-
tion:
1. Count the number of ';" (i.e., equations) in previous tokens, denoted as n eq = |[j S : 8,
=
2. Count the number of "=' in previous tokens, denoted as n; = 16 Si: Namely, the current position belongs the n -th CoT step.
HH
3. Determine the number of variables in the system of linear equations. This can be done by copying var(s,) for index such that sj is a variable and var(s;) is the largest. Denote the result as n. n Note that according to the input format, n2; is correct whenever n; eq > 1.
Similar to the proof of arithmetic expression, the first and the second tasks can be implemented by two attention heads, which perform the MEAN operation to obtain the fraction of ";' and '=>" tokens in all previous tokens. The last attention head perform the COPY operation with S; = { :: sy is a variable), ry = var(s,), ). and 1U = var(s;). Note that while rij = T' may hold for different positions J1. their values are the same (i.e., V/, 0j2), so the COPY operation still works and obtains n:" (when n2; 1).
Then, we use MLPs in parallel to calculate n; Lemma C.I Besides, we use MLP compute the auxiliary term i2 that will be used in the next layer. Therefore, the output of the first layer is
59 /1) i and n. = (n /2) i based on
(1)
og 1,,,,,,," var m m. (eia(s,): 4,,,,,,,,,
Layer 2. As described in Appendix B each CoT step eliminates one variable, and thus at the current position we are eliminating variable To By the uniqueness of the solution, there must exist an equation with nonzero coefficient for variable Tnan. In the second Transformer layer, we can determine which equation satisfies this condition. More precisely, we record whether the current equation will be used to eliminate the variable T,,0+1 in the next CoT step n + 1. We also use additional attention heads to perform some auxiliary calculations that will be used subsequent layers. Concretely, the second layer uses four attention heads to perform the following tasks:
1. Copy the value n with position j corresponding to the nearest "=" token s, (j S i). Clearly, the value well-defined when n; 2 1, and we define the value to be 0 if n = 0.
e4 + 1. which corresponds to the index of the current equation in the 2. Compute ded = m "i current CoT step.
3. Copy the embedding Edd(s) with the smallest j satisfying ny 21 = 2l and 8y is a number. Note that eid(s;) is well-defined when sy "=".
4. Compute Boolean flag (denoted as fi), which is true only when esd(s,) + eid(u). d > n and 8 "=". The definition of Si means that in the n2 CoT step, we only focus on the j-th equation when j > n2t and check whether the first number in the equation is non-zero.
If it is non-zero, set the flag to true at the specific position corresponding to token '=',
5. Copy the embeddings (Eid(s)))): 4-11) and (eia(s 42-2) of the (i 1)-th and (i 2)-th token.
The first task can be implemented by attention head via the COPY operation obtain n when n
gol
1. For the third task, we construct the matrices Q. K, V of the COPY operation such that
q = Qz = 1,1,1), k2 = Ka
eq 1,
Σ
1|s =a], )
aE[p)
0. = ) and ry = -j. By construction, q: k2 = (nz) - n, + > 4|>s; a] - 1, and thus q: k5 = 0 only when n. n,," and sy is number, and q, kg < -1 otherwise. Furthermore, the
30

---

choice of ry guarantees that the leftmost position satisfies q: k; = 0 is copied. This exactly solves the third task. For the fifth task, we use two attention heads perform the COPY operation. We only give the construction of the first head that copies (e:d(s.-):=4i)) 4-1) The matrices Q, K, V of the COPY operation constructed such that
q: = x ((i - 1) 1,, k; =j=Kx] =
(--1,2j, - fj = (eid(s;)) (44,,
and qi K; O iff i
We next use an MLP to correct the value of ng when ni and compute the second task, which a linear operation. We also compute an auxiliary flag L[ = dfth via an MLP. Regarding the fourth task, it multivariate conditional selection operation and can be similarly implemented by MLP by extending Lemma C.4 Note that we can compute the second task and the fourth task in parallel using two-layer MLP because both tasks correspond (multivariate) conditional selection and can be merged. We finally use multiplication to compute the auxiliary terms ), (var(s;)) and (var(s;-1))² The output of the MLP is
(Eid(s)) 4, 2, 1, ni (n ce fi:
Eid(s; 1) 4-1, Cid(s;- 4-2, (var(s:-1)) (var(s)) 1mm d.l)).
Layer The third layer of the Transformer uses two attention heads perform the following tasks:
1. Copy the embedding d't with the smallest j satisfying f, = and noa the answer as d.
=n? 1. Denote
2. Determine whether the next token 8:41 is a number. Denote the result as s.um
3. Determine the output of the next token 8:411 if 8:41 is not number. We denote its embedding as eq Also, we need to determine the variable index var(s;41) of the next token the next token is variable.
4. Determine the token 8,+2 if the next token 1 1 is number. There are two cases: 8:42 is a variable, and 8;42 is the token ';'. Denote the result as exex2 and var(s:+2) and compute (var($:+2))².
5. If the current token 8; is a variable, copy the embedding Gid(ay-1) (which is a number) for index j satisfying n = ni n = d, and var(s,) = var(s;). Denote the answer as c. eooimum When is not variable or d ni c. e is undefined.
We can use an attention head to perform the COPY operation that completes the first task. The construction is similar to the fourth layer in arithmetic expression and we omit it for clarity. The second attention head performs the fifth task, which can also done via the COPY operation. Regarding the second task, whether the next token is number can be purely determined by d, n, and the current token s;. Specifically, $141 is a number if s; = '+', or s; = '=', or (s; = and de9 > nz Whether the output of the next token is variable can also be purely determined by the previous tokens 8;- and also d and n Specifically, 8;+ is variable if s; '+' and is number, or s:-1 ";' and is number, or (s, or The variable index can be determined by either 2ar(8;-2) or d.a When the next token is neither a similarly determine the token
= ') and d < n
variable nor number (i.e., the symbols '+', '=', ":,, or "=", by checking 8i-1, 8i, d.,, and var When the next token is a number, 8;+2 can be determined by checking the variable s:-1 via three cases: (i) if 8:-1
variable and var(s;-1) yar then 8,-2 is a variable and var($:+2) = var(si-1) + 1; (ii) if 8;-1 is a variable and var(s;- = n: then 8142 = ::: (iii) otherwise. 8:-1 is number, then 8:42 variable and var(s142) n + 1.
All these tasks can implemented by MLPs that performs the conditional selection or look-up table based on Lemmas C.4 and C.5 Moreover, the composition of conditional selection and look- up table can be merged into a single two-layer MLP (as shown in the construction of the third layer in arithmetic expression). We next multiplication to compute the auxiliary terms (d.9),
(d.4+ + [[++ ':'1)2, and (2. However, to compute (var(*++2))² we cannot use multiplication directly as the composition of multiplication and conditional selection will require deeper MLP. Instead, note that ( (var(*) (a)(()2))) )) linearly depends on (var(s;-1))² and var(s;-1), or linearly depends on (145 2 and no.0 all of which is already computed. Therefore, can compute (var(s142))
31

---

without multiplication. The output of this layer has the form
(Eid(s,): 4, 2, n, efcu, eqox2,, var(s;), var($:+1), var($;+2), (var(s;))** (var))) (var(s;42))
var ni (n; d. (di)), (o + 1[s422 = =;1) di,, (d.t)) fium
Cid(s;-1) e num)
Layer& 4. The fourth layer of the Transformer performs the the core calculation of equation coefficients when the next token is number. There are two equations related to the calculation: the d," equation in the last CoT step, and the -th equation in the last CoT step. There are also two variables related to the calculation: the variable Tvar(s=+2) and Zngs. Specifically, we need to copy four coefficients de mo 0 var(s)))) dn mor, var(2)))))) defined as follows:
The d th equation: + 2 mg Eng + ................................ + a22 varls+a)) + ................................ 6,y0 + al nn Enze +... + a& var(#.42) (var((s4))
The df--th equation:
bus
For the case of s1+2 = "i, we need to copy coefficients 6azh and b.d. To unify the two cases, this Transformer layer uses four attention heads to perform the following tasks (note that we define var(s;) = when 8j is not a variable):
1. Copy the embedding Eid(s)-1) for position j satisfying nith = ni - 1, dht = d, 3j is a variable, and var(s;) (sy) = n.
2. Copy the embedding 1sd(s5-1) for position j satisfying n = n2th - 1, dg = dia + l||1422
:.'. eid(s;) =e?e and var(s,) var(s:+2).
3. Copy the embeddings eid(s)-1) and e eum for position satisfying ng neaa - 1, 690 d.
8j is variable, and var(s;) ni
4. Copy the embedding eid(sy-1) and e, e for position j satisfying mj = n2 1, de
de + M[s+2 ':',, Eid(,;) eex1, and var(s;) var(8:+2).
Note that for each task, there is exactly one index j satisfying the condition, and thus the copied embeddings contain the four coefficients defined above. Then, we can use an MLP to compute the "d''.ng
desired output ad45 var((::2)) - adz war().+2) /42z .n adp,ng (or baz sor--bag / / 2n m (dg.,ng)) mgr)), which can be implemented as a look-up table (according to Lemma C.5). However, there are several special cases we have consider:
dfd = n..h. In this case, the coefficient is simply computed by normalizing the - equation, which can also be implemented via look-up table.
d. and d.n + ng, In this case, the - equation and the -t equation are swapped according to Appendix B and the coefficient should be instead computed by ang" war(s.+2)
(f?.va((===)// 133,, ,. an mnp.. Fortunately, the embeddings _ in the third and the fourth tasks contain exactly var(44))())))) and apgr .n!!".
Overall, the coefficient can be computed by composition of look-up tables and (multivariate) conditional selection operations, which can be merged in single two-layer MLP.
Now two more things remain to be done. The first is to obtain the 3-dimensional embedding 4++1 when 8:41 is a variable, while currently we have only obtained var(si+1). However, we cannot m
compute the remaining two dimensions m² sin( 2var($1+=) ) and m² cos( 2var(*++1)))) since we do not assume that the MLP can approximate sin and cos functions. Nevertheless, this can be done by directly copying the embedding L, for any j such that sj is the variable Fvar(s.41) by using an attention head. Finally, the output is conditioned on the flag fnum; when f.um is true, this layer outputs the computed coefficient embedding; otherwise, it outputs eqext and 4;+1- We denote the output of this layer as x;
m
(4) = (eeet, I.").
Linear projection and softmax layer. Finally, pass it through softmax layer to predict the next token 8,+1- Unlike the proof of arithmetic expression, here the embedding 1S"H is not one-hot (which contains var(s1+1)), so we need to additionally prove the following result: let the output logit corresponding to token& & (before softmax) be (e(e,)) w: TeT, TTT + b,, where w; and b, are
32

---

parameters of the linear projection for logit t. Then, there exist parameters {w,, bo such that for any two tokens and with i +
2var(t)T
2var(()
))
Gap: == (dd()), var(I), m sin
.m²
cos
m
m
2var
2var (t)r )) 2 (1))
Ca(a)) var(t), m² sin
m cos
m
m
To prove the above result, simply set wt We have
Eid(t) var(l), m² sin
2vartt))rr
m² cos
(var(t)
Gap =1 + (var(1))² + m [id(()) id(0) var(I)var())
(2var(() m
(2var(()T]
2var(() 2var(()
.m²
sin
sin
cos
m
m (2(var(t) (f))
-(1 - =(1-I]id(() id)) + var(1) var(()) + m
cos
m
When var(t) = var(t), we have id(t) + id()) and thus Gap = 1. Otherwise,
Gap 2 - m² + m (1 cos (2x/mm) = m² + m sin 2 (T mm) 2 1,
where use the fact that sin(x) 2 x/T whenever
n/2.
Now it remains to conduct an error analysis and determine the scale parameters. Similar to the proof of arithemetic expression, we can prove that all parameter values in the Transformer are bounded by (ppoly(n)).
E.2 Proof of Theorem
We will now prove that solving system of linear equations without CoT is extremely difficult for hounded-dephl autoregressive Transformers.
Theorem E.2. Assume TC + NC². For any prime number p, any integer L, and any polynomial Q, there exists problem size such that no log-precision autoregressive Transformer defined in Section) with depth and hidden dimension
Q(m) can solve the problem Equation(m.pp)
Proof. Our proof is based on leveraging the NC -Co-compleseness of a classic problem: Unsolvable Automaton Membership Testing. According to Barrington's 's theorem 4 31. given a fixed unsolvable automaton, judging whether the automaton accepts input is complete in NC-. Below, we will prove that solving the system of linear equations is -hadhard by demonstrating that the Unsolvable Automaton Membership Testing problem is NC reducible to the problem of solving a system of linear equations. This will yield the conclusion since bounded-deptle log pog-prccision Transformers with polynomial size in T 42
Let = (2, &. 90) be any automaton, where is a set of states, Σ is a set of symbols (alphabet), 8 Q x the initial state. For any input string W142 ..-.. ... whether accepts the string can be reduced into solving a system of linear equations defined follows. The system of linear equations has (n + 1) Q + variables, which we denote as x= and Fi.g (i ......., .nn... E Q). The equations defined as follows:
Q is the transition function, I C Q is set of accept states, and go
x²
EgET In.q
20.40
Fo.4
for q E @i(go} for O < <i 2i
Fi.g
Eo(raa)) Ti-1r
It is easy see that Fi.g 1 iff the automaton arrives at state & when taking the substring wiw2 as input. Therefore, x: iff the automaton accepts the input string. Note that the above solution does not depend on the modulus p. and the solution of these equations always exists and is unique.
33

---

Furthermore, the coefficient of each equation only depends on at most one imput symbol. This implies that these equations can be efficiently constructed using a highly parallelizable algorithm within complexity of NC Therefore, by reduction, we obtain that the problem judging whether there exists solution such that x¹ 1 is NC'-
Now consider solving linear equations using a Transformer without COT. While the output of the Transformer contains multiple tokens, has to output the value of x" first. The parallel complexity of outputting the first token is bounded Therefore, cannot judge whether there exists a solution satisfying
arrange the order of variables such that the Transformer
by TC according to (42) r L.
F
Discussion on Other Architectures
F.1 Encoder-Decoder Transformer
In this paper, we choose the autoregressive Transformer is simple and the de facto standard for LLMs (c.g.. GPT). However, all of the theoretical results of this paper can be easily transferred to an encoder-decoder Transformer (such T5) using the following arguments.
Given a bounded-depth polynomial-size log- reg-precision encoder-decoder Transformer, its computation complexity is still bounded by TC so our negative results hold.
Any finite-depth autoregressive Transformer can be mimicked by finite-depth encoder- decoder Transformer of roughly the same size, since causal masking for the input sequence can be mimicked through the use of positional encoding and joint attention can be mimicked by the integration of cross aossnaattntionn and self-attention. Thus, all positive results in this paper are not exclusive to autoregressive Transformers.
F22 Recurrent Neural Network(RNN)
Although the theorems in this paper can be naturally extended other popular Transformer ar- chitectures, such as the encoder-decoder Transformer (e.g., T5), RNNs do not have this property. Theorem F. shows that RNNs cannot generate the CoT sequence using the same format proposed in our paper for the arithmetic formula task and the linear equation task unless the hidden dimension of the RNN is at least 52( lo ). where n the length of the input sequence.
Theorem F.1. For any integer p. any log-precision RNN with constant layers can neither generate
the CoT solution for the problem Arithmetic(n, p) with hidden dimension of of log (n))) nor generate
the CoT solution for the problem Eqution(n,,) p) with the hidden dimension of of log
Proof. Given an RNN model, if the hidden embeddings in each layer of the last symbols of two different input sequences the same, the two output sequences will also be exactly the same. When the RNN finishes processing the input sequence, the input sequence has been compressed into a hidden state of O(D log n) bits, where D is the hidden dimension and each scalar is represented by Olog n)) bits (by definition of fog-precisiio). Therefore, an RNN with a hidden dimension of D can only generate at most n2 different output sequences. However, for the problem Arithmetic(n, p). the first step of the CoT needs to output sequence of length (n), which contains ((n) numbers. Therefore, there are at least ((2((n) different solution sequences for the problem Arithmetic(n, p). Similarly, there are at least ((20())) different solution sequences for the problem Eqution(n.,P) p). Thus, by the Pigeon Hole Principle, to be able to generate all 20(n) different output sequences, we must have a hidden dimension of D 20 for the problem Arithmetic(n, p) and hidden for the problem Eqution(n.,p)
dimension of = 22 n² log
Due to the great difference in the model architecture between the RNN and the transformer, most of the theorems in this paper cannot be simply generalized to the RNN. Moreover, To the best of our knowledge, there is very little work to investigate the chain of thought in RNN model, either from the theoretical or the empirical aspect.
34

---

G Dynamic Programming
G.1 Examples
Longest Increasing Subsequence (LIS). The LIS problem aims to compute the length of the longest increasing subsequence given an input sequence E N Formally, 8 is subsequence
if there exists indices 1 S I 12 < S i|j I such that & Sú, holds for all & E [| A sequence called increasing if 81 < 82 < .... < ā- The LIS problem aims to find an increasing subsequence
of with maximal length. A standard DP solution is to compute the length of the longest increasing
subsequence that ends at each position i, which we denote as dp(i). is easy to write the transition function follows:
dp = ++
dp(j).
(20)
jKbb,y <3,
The final answer will be max dp(i).
However, the above DP transition function does not match the form of ( since dp(i) may depend on (an unbounded number of) all previous dp(j) (j < i). Nevertheless, this issue can be casily addressed by using different DP formulation. Let dp(j, k) be the longest increasing subsequence that ends at
position j and the second last position no more than ls (k < j). In this case, is easy to write the transition function follows:
1
k
dp(j. k)
(21)
max(dp(j.k 1). dp(k, 1). 1[s sk] 11)
i
The final answer will be max dp(2, i- 1). This DP formulation fits our framework
Edit Distance (ED). The ED problem aims to find the minimum operation cost required to convert a sequence u E 5h to another sequence v E _ There are three types of operations: inserting letter into any position, deleting letter from any position, and replacing a letter any position by a new one. The costs of insert, delete, and replace are a, b, and respectively. These operations sequentially executed and the total operation cost is the summation of all costs of individual operations.
A standard DP solution is to compute the minimum operation cost to convert the substring U1 U2 Uj to the substring U1 th Uks which we denote as dp(j, k). It is easy to write the transition function as follows:
ak bj min min(dp(j,k-2) k-1)j a,
if j
dp(j,A) -{
if A: = otherwise
+ cI(s)) # s
(21
(22)
The final answer will be dp(n1, n2). This DP formulation fits our framework 5
CFG Membership Testing. A context-free grammar (CFG) is 4-tuple, denoted as G
(V, Σ, S), where
V is a finite set of non-terminal symbols.
Σ is finite set of terminal symbols, disjoint from V.
a finite set of production rules, where each rule has the form + 3, with V and
3 E (V U Σ)* (the asterisk represents the Kleene star operation).
. S E V is the start symbol.
The non-terminal symbols in represent (abstract) syntactic categories, while the terminal symbols in Σ represent the actual words/tokens of the language. The production rules in It specify how non- terminal symbol can be replaced by sequences of terminal and non-terminal symbols concatenated together. Below, focus on canonical version of CFG [33]. in which the term 3 in each rule is either an empty string or contains exactly two symbols:
* A - where A V and is the empty string:
+ BC, where B, C V.
By introducing additional non-terminal symbols to split complex rules, any CFG can be easily translated into the canonical version expressing the same language.
35

---

The (universal) CFG Membership Testing problem is defined as follows: given a canonical CFG G and string v. judge whether v can be generated from G. This problem is known be P-complete The CYK algorithm 153, which is based on DP, is classic algorithm to solve the CFG Membership Testing problem. Here, we consider a variant of the CYK algorithm that can also handle rules of the form A + c. Denote R = (R1 ,mn) and let o/. The state space of the DP process is defined as
R5
[vm.m, = (14,k, :0 :0 S M. 1 S S A S m)
where dp i, k. r) stores whether the substring 11+1 ... By can be generated by nonterminal A by first using rule R; A + BC with F S and there exists index k S k such that the substring V++1 ... 2 can be generated by and the substring Ui+1
can be generated by C. The index t represents the number of iterations (which will be detailed later). For the boundary setting when i ==j= = kt dp(0, i, i,1, A, r) simply stores whether nonterminal A can generate an empty string by first using rule Rş with F S r. The transition function can be formally written
i,,
if = k, 0, 0, if = j = k, 0, 2 1,
T [ddp(1, J. k, - or R cc
0
if < j. = 0.
dp(t - L.i.-,,,jj.j,mm
if > 0, = 0, k = i
ift > 0, 0, i,
A,r, -V A BC for some BC s.t.
(B E = (B dp(t, i, k k. B) 1). (CEEAj=k+iA) = C dp(t, k.j.j, = 1))l if > 0,r
The final answer will be dp(()). n, n, m). This DP formulation fits our framework f
Regarding Remark 4.6 It can be easily verified that the state spaces of the three problems mentioned above are of polynomial size, satisfying Assumption 4.2, Additionally, the MLP with the ReLU activation function can implement (the composition of) the following functions:
max(a, b) and min(a, b). where a, I;
I[a + b). I[a < b]. I[ > b). where a, Z:
a b, where E b {0, 1}:
linear transformation;
conditional selection (Lemma C.4 for example,
prclect (x) { (x) i (x()) x
where f² and fc are functions that can be implemented by MLPs with ReLU activation, and x E Z.
This implies that the MLP with ReLU activation approximate the functions f. g. h in the transition function for the above three DP problems. According to Lemma C.2 these functions can be efficiently approximated by perceptron of constant size with GeLU activation. Similarly, the topological ordering can also be efficiently implemented by MLPs with GeLU activation:
* LIS: (1,A) { (j,++1
if < j- if = j-
1,+)
ED: (1.b) {
(1,++1)
if < n2
(+1,0)
if m2
"There is subtle mismatch between DP formulation and our framework 5 that terms and C depend on the input (rather than purely determined by the state). Nevertheless, is easy to extend Theorem 4.7 to this setting. Specifically, autoregressive Transformer can first extract B and C' using one layer and then extract dp(t, k. k. B, m) and dp(t, j.j. C, m)) using additional layer.
36

---

CFG Membership Testing:
(1, hi A, +1)
if < m if m, A + last(V)
(t,i,j, kt next(A).
1,
if = m, A last last(V). <
1,, firstt) (+t1,i1.1j+,itt (V) (0.0.--i+1,0ffis((())
if m, A. last (2),()), = j
if r = m, A last k= < [
if m, A
k=j=,,, M
where we order the set V and denote first(V) as the first element of V. last(V) the last element of V, and denote next(A) the successor element of A.
Therefore, Assumptions 4.3 to 4.5
are satisfied and all three problems can be solved by autoregressive
Transformers with CoT.
G.2 Proof of Theorem 4.7
In this subsection, we will give proof of the Theorem 4.7
Theorem G.1. Consider any DP problem satisfying Assumptions 20 For any integer n E N, there exists an autoregressive Transformer with constant depth T., hidden dimension d and attention heads H (independent of n), such that the answer generated by the Transformer is correct for all input sequences of length no more than n. Moreover, all parameter values are bounded by O(poy((n)
Proof. Input Format. Assume that we have a sequence of tokens "1, 8 and we want to generate the next token $141- We embed the token $y by
ax (e
inpar state
Ek e% ex
answer
sep ki 1),
where each part of the embedding is defined follows:
1. input is the embedding of the input token 8% E X. If the current position does not represent
e%
an input token, then ex
2. ej
state is the embedding of the DP state in I at position k. If the current position corresponds an input token the final answer, then ex 0. We also assume that for all E I, the embedding of state i is non-zero.
3. 6g dp is the embedding of the DP value in y at position k. If the current position corresponds to
input token the final answer, then e dp 0.
4. e the embedding the answer token in 2, and e if the current position corresponds to an input token or an intermediate DP position.
5. ej is the embedding of the separator separating different input sequences. We set ex e; if the current token 9y is the j-th separator, where ej is the one-hot vector with the j-th element begin
6. The position embedding & indicates the index of the token in the sequence.
Block 1. The first block of the autoregressive Transformer contains several layers. It first uses N attention heads to perform the following task:
Copy the positional embedding of the N separators sep.1
Pk Sep.N
E.
Similar to the previous proofs, this can achieved via the COPY operation with S, = {j S k cep e,} fort [N] and by = j. Then, several MLPs follow, which perform the following tasks:
Calculate the problem size na (wi sep.1 1, 1,
sep.2
sep,1 1,
sep.NN
Nep..N 1).
Pk
PP
Pk
Obtain the next state e stxt_state If the current state already the last state, set e
stat
The first task is a linear transformation, which can clearly be processed by an MLP Proposition According Assumption |4.4] we can use an MLP compute the embedding of the next state e mext_state based on the embedding of the current state eslate and the problem size n. When the required
use L - 1 Transformer layers MLP in Assumption 4.4 has multiple layers (i.c., L layers),
37

---

to implement a L-layer MLP. This can be achieved by just zero the weight matrices in the attention layers while maintaining the input using residual connections. The output this block is
(e),
Ck stxtt
c
spp Ck
Tok, k 1).
The second layer of the Transformer does not use attention heads. It only uses the MLP to perform the following tasks:
Block 2.
Calculate h h(ng; C e and g( 6x sestsstae) We assume that the embedding of 0 is O.
Set the flag f representing whether current state ex is the last state.
Set the flag fim representing whether current state ex is in the set A, i... used in the aggregation function.
Similar to the first block, we stack several two-layer perceptrons to implement a multilayer perceptron.
According to Assumptions 4.3 and 4.5]
MLP to complete the first and the last tasks. next state
The second task can be done by checking whether exe + 0 and &g
O. We also compute the auxiliary quantities (f( ek. pext_stalle) ))² (g(nx: nexi state (ef 2 and h², which are elementwise
square operations and can be implemented by MLP (Lemmaa C.I. The output of this block is
(2) k (ex cu EE next state dp
.ex ex ngg, hnn;, Ck (h( ex etxttsttte))" (g(n), ex mexttsta)))? (estule)) 2 f 1 samswer k k2 1).
next_state stats nest sttuutt (Tkk: E
Block 3. The third block of the Transformer uses K + J heads to perform the following tasks (where K and J are defined in (5)):
*
Copy the input token embeddings corresponding to $g (i), (1 (g)(1) where corresponds to ex cheut state When 9e (i) 0, we set (t) to be special token.
Copy the DP value embeddings corresponding to dp(hi (i)) dp(hx (i)) for i corresponds Stale When he (i 0. we set dp(h+ be a special value.
ex
Calculate the output dp(i) for corresponds to ex
next_statee denoted as ex next_dp
The first two tasks can be done via the COPY operation. To copy DP values, the attention head attends positions with state matching hx(i) for E [K]. To copy input tokens, the attention head attends positions j = ge (1) for E [J]. To handle the special token/value, it is simply a conditional selection operation and can be handled by an MLP (Lemma C.4). According to Assumption 4.3 we can calculate the function f (defined (5)) using an MLP. The output of this layer is
(3) x
(e nextsstate Cg mexi_dp niks Jh.
polate
ff canswer 1).
Block The fourth block of the autoregressive transformer contains one Transformer layer. De- pending the aggregation function. it uses one attention head for the operation max or min, or two attention heads for the operation Σ. This block performs the following tasks:
Aggregate the DP values according to the aggregation function Equation 6
Generate the output based on the flag famwer
For the first task, if the aggregation function is max min, we use one attention head simply copy the embedding e for index j such that fanswer = 1 and e, is the largest/smallest, according to Lemma C.7 the aggregation function is E, we use two attention heads, where one attention head computes the mean of e; the fraction of elements in the sequence such that fanswer = 1. Finally, the second task is a conditional selection operation and thus be implemented by MLP (Lemma C.4).
dp for index j such that fanswer = 1, and the other attention head calculates
G.3 Proof of the Theorem 4.8
Theorem G.2. Assume TCD + NC'. There exists a context-free language such that for any depth and any polynomial Q there exists a sequence length n E N where loe-precision autoregressive transformer with depth and hidden dimension S Q(n) can generate the correct answer for the CFG Membership Testing problem for all input strings of length n.
38

---

Proof. According to Barrington's 's theorem [, given fixed unsolvable automaton, judging whether the automaton accepts an input is complete in NC which is special case for the CFG Membership Testing problem. Therefore, the CFG Membership Testing problem NC With the assumption that TC + NC', the CFG Membership Testing problem is out of the capacity of the log- oeg-precisoon autoregressive transformer.
H Experimental Details
In this section, we present the experimental details.
H.1 Datasets
We set the number field 11 in the math experiments. In the LIS experiment, we set the number of different input tokens to 150; in the ED experiment, we set the number of different input tokens to 26. The vocabulary is constructed by including all symbols. For all tasks and settings (direct v.s. CoT), the size of the training and testing dataset is 1M and O.IM respectively. The constructions of different datasets are introduced below.
Arithmetic Expression. All arithmetic expression problems generated according to Algorithm [ In Algorithm [ we first create number that serves as the answer to the problem. We then decompose the number using sampled operators sequentially, serving the problem, until the maximum number of operators met. The CoT procedure is precisely defined by reversing this problem generation process. For example, sample in the direct dataset looks like
1 (1 2) =
while the corresponding sample in the CoT data looks like
1 + 5 x (1 2) 2) = 1
10 10
Algorithm 1: Arithmetic Expression Problem Generation
Input :Number of Operators
Input +Vocabulary of numbers V (0, 10/ number field = 11
Output : Arithmetic expression
Sample the first number uniformly from V:
= []:
3 Append / to 8; for i +- 1 to do Sample p uniformly from {0, 1, 1,..,, len(s) 1}. satisfying spp] is a number;
Sample uniformly from {+, -+--, X, +};
Sample numbers t1, 12, satisfying the result of 12) equals s(pl;
if' spp 1|= or {+, -} and slp E {-, x) (o {+, -} and
s(p 1] E [xx +)) then
pop s spp:;
10
insert [(. [t1]. [o]. [tzl. D] sequentially into s(p]:
11
else
12
pop sppl;
13
insert [f]], [o]. [tz] sequentially into s(p]:
end 14
end
Linear Equation. All linear equation problems are generated according to Algorithm In Algorithm 2
we consider the linear systems that only have a unique solution. Given sampled linear system that satisfies this condition, we "translate" it to sequence by concatenating all the equations (separated by commas), which serves as the problem. The answer to the problem is also sequence consisting of variables and the corresponding values. The CoT solution of cach problem is
39

---

the calculation process of the Gaussian elimination algorithm applied to each variable sequentially. For example, a sample the direct dataset looks like
2x1 322 + 3x3 la + 7x2 + Oang = 0, Oz + 2x² + 1xz 1, [SEP] 21 4, 322 1, 10,
while the corresponding sample in the CoT dataset looks like
2 + 3x2 + 323 8, 1æ + 7 + Or3 0, 0r + 2x² + 123 = 1,
[SEP] 7 + 7 4, 0n 4mg 7,2r2
123
[SEP] 21 + 9xn = 6, + 62z = 6, 423 = 7,
[SEP] I 4,42 1, 10,
Algorithm 2: Linear Equation Data Generation
Input :Number of Variable n Input :Vocabulary of numbers V = 10,/ number field p 11 Output:Lineet Equation s
Sample b uniformly from Vnx1;
2 do
Sample A uniformly from Vaxm;
4 while A is not invertible;
"A11r1 ... + Ann bi...., Am1z1 ... + Annon bn.
Longest Increasing Subsequence. All input sequences (i.e., problems) generated according to Algorithm 3 To make the task challenging enough, we first concatenate several increasing subsequences of given length, and then randomly insert numbers into the whole sequence. The inputs has 150 different tokens, ranging from 101 to 250 to avoid token overlap with DP array. The CoT solution to the problem is the DP array plus the final answer, which is defined in (20). Here, we consider the DP formulation (20) because the CoT output length is much shorter than the one corresponding formulation
(2)). This allows us consider more challenging input sequences with longer length. While this DP formulation does not precisely obey the theoretical assumption given in Assumption 4 we found that the Transformer can still learn easily.
For example, a sample the direct dataset looks like
103 107 109 112 101 103 105 107 115 109 111 113 102 [SEP]
while the corresponding sample in the CoT dataset looks like
103 107 109 112 101 103 105 107 115 109 111 113 102
[SEP] 2 4 234556
[SEP] 7
Edit Distance. All input sequences (i.e., problems) are generated according to Algorithm In Algorithm[4] [ we generate the first string randomly. For the generation of the second string, we use two methods. In the first method, we generate the second string randomly, corresponding to large edit distance. In the second method, to small edit distance. The two strings are concatenated by 'I",, and the concatenation used as the model input. For the calculation of edit distance, assign different costs to different operators. The costs for the ADD and DELETE operators are set to 2, while the REPLACE operator is assigned a cost of 3, since REPLACE should be more costly than ADD/DELETE while less costly than their summation. The CoT procedure is also the DP array, defined in Section 4.1 For example, a sample in the direct dataset looks like
copy the first string with random corruption, corresponding
p pass [SEP] 4
while the corresponding sample in the CoT dataset looks like
a pass
[SEP] 3 2 4 6
[SEP] 4 2 4
[SEP] 4
40

---

Algorithm 3: LIS Data Generation
Input :Sequence Length n
Input Vocabulary of numbers V {101, 150...200
Output :Sequence 8
:utput:Sequence
Sample uniformly from (3.4...m
2 Sample uniformly from {1,2, []:
4 push 0 to a:
if then
Sample uniformly from (1/2...[//2 + 1}:
push j to
else if then
Sample uniformly from (1.2...11/3) + +1]: Sample & uniformly from {1.2... j)/2 +1;;
11
push j to a
push j
a
push / to
14 s f- Sample numbers from V;
for i +- to f do
16
Sort s(ali 1] a[t]]:// This process makes sure the LIS of the generated sequence is least [t/t].
17 end
18 r f Sample n numbers from V;
Randomly insert into s;
Algorithm ED Data Generation
Input :Length of the First String n
Input :Alphabet V {aa b...2]
Output:SSquencce 81,
Sample uniformly from (3, 4...10
T + Sample letters from V:
81 < Sample letters uniformly from T
Sample p uniformly from [0, 1];
5 p < 0.4 then
Sample uniformly from {n - 3, 2,..., n + 2};
82 < Sample letters uniformly from T
8 else
do
10
82 -- 81
for
+ 1 to n do
12
Sample uniformly from (0, I...len(s2) - -1:: Sample uniformly from T;
13
14
Randomly conduct one of the followings: pop 82 [p]. substitute s2[p] with 1, insert into s2 s2[p::
15
end
16
while (en(s2) not in [n 3, n + 2]:
17 end
41

---

H.2 Model training
We use the minGPT implementation for all the experiments, where the detailed Transformer layer are listed below.
x = LayerNorm ([01 + P1, ........ On + pml'))
(23)
Attn (X) = Σ
(sottmax (XW( (XW ( ( T M)) xww..w
(24)
h=1
FFN (X) = <ox W
(25)
= Attn (4) (LayerNorm( X (1-1))
(26)
x -0= + FFN (LayerNorm (1))))
(27)
We use sinusoidal positional embedding and use Xavier initialization for all the parameters. The activation function is chosen to be GeLU. The dimension of the embedding is set to 256, and the number heads is set to The hidden size in the FFN layer is set 1024.
We use the same hyperparameter configuration for all experiments, i.e., the performance comparison between the models trained on the direct and CoT datasets of Arithmetic, Equation, LIS, and ED tasks, and the additional length extrapolation experiments (which we use relative positional encodings 48 instead of absolute positional encodings). In detail, use AdamW optimizer with B1. 82= 0.9, 0.999. The learning rate is set to le-4, and the weight decay is set to 0.01. We set the batch size to 512 during training with linear learning rate decay scheduler. We use learning rate warm-up, and the warm-up stage is set to 5 epochs. The dropout rate is set to 0.1. The total number of training epochs is set 100.
I Robustness in Training Data Quality
In the training of language models, the quality of data, especially intermediate steps, is crucial for optimal performance. However, real-world training datasets may contain corrupted or missing intermediate steps. This highlights the importance of evaluating model resilience under imperfect conditions. To investigate this, we conducted series of experiments focused on arithmetic tasks, introducing varying rates of corruption and omission, encapsulated by the parameter y. For instance, when 7 = 0.1, it denotes that 10% the intermediate steps are omitted and 10% of the remaining steps corrupted through a single-token random replacement. Our motivation for introducing this metric is to simulate potential imperfections in real-world datasets, thus assessing the robustness of our model more thoroughly. The number of operators the arithmetic task is set 10.
Table 1: Experimental results of robustness.
0
0.1
0.2
0.3
Accuracy
100.0%
98.5%
97.6%
95.8%
Table 2: Model performance on Arithmetic task with 3-layer Transformers.
Our experimental results are presented in Table I the rate of corruption and omission increases, there is slight decrease in accuracy. These results elucidate the inherent robustness of training CoT demonstrations. It is noteworthy that the accuracy of the model remains commendably high despite significant distortions present in the training data, highlighting its potential utility in real-world scenarios where data quality may be compromised.
42


Reviews:
[{"review_id": "JgRIVMxGoT", "reviewer": "Reviewer_C8fV", "summary": "This paper studies the theoretical power of the Chain-of-Thought (CoT) prompting. In particular, this paper mathematically confirms that two well-chosen tasks (e.g., arithmetic and equation) and the problem of Dynamic Programming are beyond bounded-depth Transformer models without CoT (unless their size grows prohibitively large); with CoT and generated intermediate derivations, those problems become solvable. The mathematical derivations are under mild assumptions, and empirical results confirms the mathematical study (on four representative tasks).", "questions": "Line 245, the success of solving Dynamic Programming problems critically depends on the input sequences being laid out in a topological order. In the two DP experiments (Longest Increasing Subsequence and Edit Distance), the topological order appears to be learnable from the CoT dataset. Could there be any theoretical study about how hard it is to learn this order?", "limitations": "The authors have adequately addressed the limitations.", "rating": 9, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "Given the prevalence of CoT, theoretical study of the limit of CoT becomes extremely valuable. This paper overcomes several limits in previous studies (e.g., assuming infinite precision) and focuses on the setting of autoregressive Transforms, which is close to the scenario of real-world LLMs. Moreover, the proposed empirical tasks are illustrative and easily reproducible.", "weaknesses": "I don't find any noticeable weakness in this paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank Reviewer C8fV for the positive feedback, appreciation for our work, and insightful questions.\n\n> **Question**: Line 245, the success of solving Dynamic Programming problems critically depends on the input sequences being laid out in topological order. In the two DP experiments (Longest Increasing Subsequence and Edit Distance), the topological order appears to be learnable from the CoT dataset. Could there be any theoretical study about how hard it is to learn this order?\n\nWe thank the reviewer for the insightful question. In our paper, we mainly study CoT from an expressivity perspective, i.e., whether there exists a model that can generate the topological ordering of the DP states and find the correct solution, and we prove that the answer is yes. Your question relates to the generalization ability of CoT training, i,e., why the model behaves well on unseen data. For this question,  we can offer some intuitive insights into why generating the topological ordering is easy for an autoregressive Transformer. In our formulation, the topological ordering can be determined by a function $F$ that takes the current state $i_k$ and the problem scale $n$ as inputs and outputs the next state $i_{k+1}$, i.e., $F(i_k,n)=i_{k+1}$. Assumption 4.4 guarantees that $F$ can be efficiently approximated by a constant-sized perceptron with GeLU activation, by which **the topological order can be easily generated in an autoregressive way by this simple function $F$**. In the Appendix, we give an explicit expression of the function $F$ for each DP problem considered in this paper, all of which are simple:\n$$\n\\text{LIS}: F((j,k),n)=\\begin{cases}\n(j,k+1)\\ \\ \\text{if}\\ \\ k<j-1\\\\\\\\\n(j+1,0)\\ \\ \\text{if}\\ \\ k=j-1\n\\end{cases}\n$$\n$$\n\\text{ED}: F((j,k),(n_1,n_2))=\\begin{cases}\n(j,k+1)\\ \\ \\text{if}\\ \\ k<n_2\\\\\\\\\n(j+1,0)\\ \\ \\text{if}\\ \\ k=n_2\n\\end{cases}\n$$\nIt is easy to see that $F$ can be represented by a small MLP with ReLU (or GeLU) activation. Therefore, it may not be difficult for the model to learn the topological order given sufficient training demonstrations.\n\nOn the other hand, we believe theoretically understanding the generalization ability of CoT learning process is extremely important and we leave it as future work. Although there are some works trying to explain the generalization of modern deep learning models [1,2], they may not provide useful insights for the CoT setting. Studying the generalization ability of large language models from massive data around their reasoning ability should be a very fundamental and practical problem that deserves more attention in the future.\n\n[1] Li, Yingcong, et al. \"Transformers as algorithms: Generalization and implicit model selection in in-context learning.\" \n\n[2] Xu, Keyulu, et al. \"How neural networks extrapolate: From feedforward to graph neural networks.\" \n"}, {"review_id": "3T9UQmjhWy", "reviewer": "Reviewer_QoTB", "summary": "This paper presents various separation results, showing that a Transformer with CoT can solve certain formally-defined reasoning tasks, but a Transformer *without* CoT cannot (assuming bounded depth). This sheds light on the power of CoT. The formal results are supplemented with empirical results that support the claims.", "questions": "- See CoT naming comment above. Do you agree?\n- See notation comment above. Is there a way to simplify the notation in Section 4.1? Is there a reason it has to be this complicated? If it seems necessary, maybe there's a simpler version that can be presented in the body of the paper, with the full version moved to the appendix?\n- My main question is also discussed in the Weaknesses section. Is there any conceptual takeaway from this theoretical analysis beyond \"generating more tokens is more powerful\"? Even without that, I think this is a strong paper. However, I think this could be an especially powerful paper if the intuitions from the theoretical analysis could be further synthesized in this direction (and made accessible to readers).", "limitations": "The authors do a nice job of discussing some of the limitations of this work. There is no discussion of societal impact, but I do not think this is a problem for this particular paper.", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Understanding CoT is an important and timely question. This paper is therefore tackling a significant question.\n- To the best of my knowledge, this is the first paper to provide a theoretical explanation of the power of CoT (i.e., originality).\n- In general, the paper is quite clear and high-quality, though I think the notation could be improved (see below).", "weaknesses": "- I think calling it CoT but focusing on \"CoT generation\" is perhaps misleading. Many people think of CoT as a prompting technique. In my opinion, the generation aspect that this paper focuses on is actually bigger/more important than CoT suggests, and I actually think that - although buzzwordy - CoT diminishes the general power of generation that this paper is getting at (i.e., this paper transcends CoT). I would suggest looking into alternative, more general titles.\n- Section 4.1 has a lot of notation, including double subscripts. I think there's a significant lack of accessibility created by the heavy notation. I would really encourage the authors to think about whether the notation could be simplified. I think doing so could substantially increase the long-term impact of this paper.\n- Relatedly, equations (4) and (5) would be clearer if their new objects were introduced and explained conceptually before jumping into equations (4) and (5).\n- Overall, I think the paper could do a better job explaining the significance of the theoretical results. Although this is an important problem to study theoretically, and this paper does a good job of initiating that study, it's not entirely clear what can be gained conceptually from this analysis. I think many people already intuitively grasp that generating more tokens gives an LLM more power, and various methods that allow more tokens to be generated before arriving at a final answer are more powerful. It would be nice to understand whether there's a conceptual message here that goes any deeper than the aforementioned intuition.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank Reviewer QoTB for the careful reading, positive feedback, valuable suggestions regarding presentations, and detailed comments. Below, we would like to give detailed responses to each of your comments and questions.\n\n**Regarding the scope of the paper**. Thanks for the suggestion. We connect our work to \"CoT\" because our theory suggests that to solve math/reasoning problems, \"generating intermediate deviations in an autoregressive way\" is easier and much more parameter-efficient. As this generation process is regarded as \"Chain of Thought\" by the community, we leverage the concept during writing. On the other hand, we fully agree that intrinsically, this work is more about the way of using LLMs rather than how to develop specific CoT prompts. We will make this clear in the introduction and are happy to illustrate more and try to come up with a more accurate and appropriate title.\n\n**Regarding notations in Section 4**. We appreciate the constructive suggestions and will revise our manuscript accordingly. In particular, we intend to simplify the notations in Section 4.1 by using a vectorized notation. For example, in equation (5), we will use $s_{\\mathbf{g}(i)}$ to represent the vector $(s_{g_1(i)},\\cdots,s_{g_J(i)})$. In this way, equation (5) can be rewritten in a more concise form, avoiding the problem of double subscript. The original equation and the updated equation are shown below:\n$$\n\\mathsf{dp}(i)=f(i,s_{g_1(i)},\\cdots,s_{g_n(i)},\\mathsf{dp}(h_1(i)),\\cdots,\\mathsf{dp}(h_K(i)))\n$$\n\n$$\n\\mathsf{dp}(i)=f(i,s_{\\mathbf{g}(i)},\\mathsf{dp}(\\mathbf{h}(i)))\n$$\n\n> **Question**: Is there any conceptual takeaway from this theoretical analysis beyond \"generating more tokens is more powerful\"? ...... However, I think this could be an especially powerful paper if the intuitions from the theoretical analysis could be further synthesized in this direction (and made accessible to readers).\n\nThanks for raising this good question. It makes us realize that there is still much room for us to improve the writing of the paper. Indeed, our theoretical analysis has a bunch of conceptual takeaways beyond the surface conclusion that \"generating more tokens is more powerful\". We detail these takeaways below:\n\n* **The key role of self-attention in CoT generation**. Our analysis points out two key components that enable CoT generation, which we call **COPY** and **MEAN**. Specifically, the COPY operation extracts the hidden information of a specific previous position that satisfies certain conditions, e.g., extracting the hidden embedding of the last equal sign (=), or the last non-number token. The MEAN operation averages the hidden embedding for a set of previous positions that satisfy certain conditions, e.g., the average hidden embedding of all *number* tokens between the last equal sign and the current token. We prove that **both COPY and MEAN can be realized by a self-attention layer**. We then use exactly the two operations to build entire **parallel algorithms** that can generate CoT sequences for all math and DP problems. Our result highlights the crucial role played by the self-attention and Transformer architecture, and may inspire future research on architectural design in Large Language Models (e.g., more efficient LLMs).\n\n* **Regarding the length of CoT generation**. Our theoretical analysis also gives insights into how many intermediate steps are needed in the CoT generation. In particular, when the Transformer model generates a CoT sequence, we have to ensure that the complexity of generating each step is within $\\mathsf{TC}^0$. Using this argument, one can easily check whether the length of a specific CoT format is sufficient. For more complex problems, we need to decompose it into more subproblems (or more steps), so that each step is within the complexity of $\\mathsf{TC}^0$. Moreover, we give a standard criterion to check whether the $\\mathsf{TC}^0$ expressivity is satisfied: if each CoT step can be represented by a finite composition of **COPY** and **MEAN** operations, then each CoT step will be within the complexity of $\\mathsf{TC}^0$.\n\n* **Regarding dynamic programming**. We theoretically show that CoT allows Transformers to solve general DP problems. Thus, as a direct consequence, Transformers are even capable of solving extremely hard problems that are $\\mathsf{P}$-complete. This provides a deeper understanding of why popular large language models can be so powerful in reasoning. We believe our proposed dynamic programming framework may also be useful for studying and measuring the expressive power of other architectures in the future.\n\nWe will incorporate these discussions into the next version of our paper.\n\nWe hope our response can clarify your concerns and will improve the paper writing according to your suggestions and questions. We are happy to go into more detail regarding any of the above questions and we look forward to your reply.\n"}, {"review_id": "Cbc0rwbafK", "reviewer": "Reviewer_P6n5", "summary": "The paper contributes to theoretical and empirical understanding of Chain-of-thought, i.e. intermediate process generation to assist desired output generation. In the theory part, authors show\n- log-presicion Transformer with bound depth cannot solve simple math tasks (calculate, solve linear equations) unless the model size grows super-polynomially w.r.t. input length. the proof is based on circuit complexity and a bottleneck of parallel complexity (assuming $TC^0 \\neq NC^1$). However, constant-size Transformer can solve both by generating common math intermediate steps.\n- for dynamic programming (DP), Transformers with CoT can generate the correct answer intuitively. In contrast, it is proven that Context-Free Grammar Membership Testing cannot be solved with a bounded-depth Transformer with polynomial depth.\n\nExperiments then validate these results using two math and two DP tasks, showing CoT > no CoT (with more layers even). ", "questions": "- Seems \"Self-attention networks can process bounded hierarchical languages\" proved that (2-layer) Transformers can recognize Dyck, yet this papers proves CFG recognition is in general hard? Some discussion would be nice.\n\n- Would results hold for RNN or other architectures? Some discussion would make the paper stronger.", "limitations": "Authors write about limitations fairly.", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- CoT has been important in language processing, and its theoretical and empirical understanding is important and timely.\n\n- As far as I can tell, both the theory and experiment parts of the paper is solid, and well-written.\n\n- I like the connection to circuit complexity theory. Though the conclusion and \"intuition\" is not hard to grasp, the theory part is technical and non-trivial to establish.", "weaknesses": "- As noted by authors, it's still limited to expressivity (not learning with large corpora, large model).\n\n- Some missing references around Dyck language recognition:\n\nSelf-attention networks can process bounded hierarchical languages. ACL 2021\nRNNs can generate bounded hierarchical languages with optimal memory. EMNLP 2020\n ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank Reviewer P6n5 for the positive feedback, valuable suggestions, and two insightful questions regarding the related work and other architectures. Below, we would like to give detailed responses to each of your comments and questions.\n\n**Regarding related work**. Thanks for the valuable suggestions on the related work. We will follow your advice and include these two papers in our related work section. These two papers focus on the Dyck language, a special case of the CFG. The first one shows that the hard-attention transformers can recognize Dyck language with bounded depth. Moreover, a two-layer soft-attention transformer can generate Dyck language. The second paper tries to use RNN to generate bounded-depth Dyck language. This paper demonstrates that RNN can generate the Dyck language with hidden units of reasonable size. Furthermore, the paper proves that the size of the hidden units is optimal. The conclusions of these two papers complement our theorems on the general CFG.\n\nWe also thank you for posing two insightful questions that deserve careful discussion.\n\n> **Question**: Seems \"Self-attention networks can process bounded hierarchical languages\" proved that (2-layer) Transformers can recognize Dyck, yet this papers proves CFG recognition is in general hard?\n\nYes, the problem of general CFG recognition is much harder than the special case of the Dyck language. The paper “Self-attention networks can process bounded hierarchical languages” employs a Transformer encoder that is similar to the auto-regressive Transformer without CoT, except for the causal mask. The complexity of both Transformer encoder and autoregressive Transformer without CoT is upper bounded by the circuit complexity $\\mathsf{TC}^0$. For the special case of Dyck language, it has been proved that the problem of recognizing Dyck language is actually in the complexity class $\\mathsf{TC}^0$ [1]. Therefore, it is possible for a transformer model to recognize the Dyck language. However, general CFG recognition is $\\mathsf{P}$-complete, which is intrinsically hard to be solved by a Transformer without CoT.\n\n[1] On the relative complexity of some languages in NC1. Barrington et al.\n\n> **Question**: Would results hold for RNN or other architectures?\n\nFirst, most of the theorems in this paper can be naturally extended to some popular settings with Transformer, such as the encoder-decoder architectures (e.g., T5). However, as for the RNN, it is actually not the case. We can show that RNNs cannot generate the CoT sequence using the same format proposed in our paper for the arithmetic formula task and the linear equation task unless the hidden dimension of the RNN is at least $O(\\frac{n}{\\log n})$, where $n$ is the input length. The reason is as follows. When the RNN generates the first equal sign, it has to compress the input sequence into a hidden state of $O(D\\log n)$ bits, where $D$ is the hidden dimension and each element is represented by $O(\\log n)$ bits (by definition of log-precision). On the other hand, the first step of the CoT needs to output a sequence of length $O(n)$, which contains $O(n)$ numbers. Therefore, there are at least $2^{O(n)}$ different output sequences, and thus by the Pigeon Hole Principle, to be able to generate all $2^{O(n)}$ different output sequences, we must have $D=\\Omega(\\frac{n}{\\log n})$.\n"}, {"review_id": "M2nngsdVt0", "reviewer": "Reviewer_7MZe", "summary": "The paper mainly focuses on theoretically proving the effectiveness of the CoT in autoregressive Transformer models for solving fundamental mathematical and decision problems through generating intermediate steps. It demonstrates that any finite-depth Transformer model cannot directly output correct answers to these tasks unless the model size grows super-polynomially with the input length. The paper also includes experimental validation using a constructed dataset.", "questions": "1. Is there a specific reason for choosing autoregressive Transformer? Were other models, such as the encoder-decoder architecture like T5 model be considered? Would CoT be equally effective in those models?\n\n2. In the experiments, a Transformer with three layers achieves near-perfect accuracy on each task, but the proofs utilize five layers. How much does the number of layers affect the performance of CoT, and are there necessary layer numbers for different tasks? Additionally, the difficulty of the datasets for now seem insufficient to explore this question.\n\n3. In the Length Extrapolation section, only Arithmetic Expression Extrapolation is explored. Does CoT also possess the ability to learn the underlying mechanisms for other tasks? The conclusion seems to lack sufficient experimental verification and analysis.\n\n4. Have you considered the impact of the quality of intermediate steps on the performance, such as inserting some invalid intermediate steps or omitting a certain number of intermediate steps?", "limitations": "Section 7 of the paper provides a thorough discussion of the limitations. Additionally, while CoT demonstrates excellent performance across various tasks, this paper mainly focuses on mathematical tasks and could further explore other types of tasks, such as logic reasoning tasks represented in natural language.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The paper provides a solid theoretical proof of the effectiveness of CoT, making a valuable contribution to further exploring the underlying mechanisms of CoT operation.\n\n2. The theoretical proofs presented in the paper are comprehensive and well-organized.\n\n3. The research problem addressed in the paper is clearly significant. CoT has shown strong empirical performance but lacks theoretical analysis. Thus, this study is timely and necessary.\n", "weaknesses": "1. The experimental section seems to be insufficient, for example, the length extrapolation experiment only provides results for one task.\n\n2. The exploration of model architectures is lacking. (see my questions)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank Reviewer 7MZe for the careful reading, thoughtful inquiries, and positive feedback. Below, we are happy to provide further elaboration on each of the points you raised:\n\n> Q1: Is there a specific reason for choosing autoregressive Transformer? Were other models, such as the encoder-decoder architecture like T5 model be considered? Would CoT be equally effective in those models?\n\nThanks for the question. We choose autoregressive Transformer as it is the de facto standard for LLMs (e.g., GPT). It is also simpler than encoder-decoder architectures, which makes our analysis and proof cleaner. Yet, our theoretical results can be easily transferred to an encoder-decoder model (like T5) using the following argument. (a) On one hand, given a bounded-depth polynomial-size log-precision encoder-decoder Transformer, its parallel complexity is still bounded by $\\mathsf{TC}^0$, so our negative results hold. (b) On the other hand, any finite-depth autoregressive Transformer can be mimicked by a finite-depth encoder-decoder Transformer of roughly the same size, since causal masking for the input sequence can be mimicked through the use of positional encoding and joint attention can be mimicked by an integration of cross-attention and self-attention. Thus, all results in this paper are not exclusive to autoregressive Transformers. We can add those discussions into the paper if you think they are helpful.\n\n> Q2: How much does the number of layers affect the performance of CoT, and are there necessary layer numbers for different tasks? Additionally, the difficulty of the datasets for now seem insufficient to explore this question.\n\n\nThanks for the question. We study simple tasks in the paper and found that from both theoretical and practical perspectives, a shallow LLM with CoT already has enough capacity to solve them. We agree and believe deeper LLMs are required to solve more complicated tasks with CoT. For example, it can be easily imagined that a deeper LLM is needed to complete a task if it composes of solving linear equations and arithmetic with dynamic programming. We are working on more general and advanced reasoning tasks and investigating the dependency between task and LLM parameter complexity.\n\n\n> Q3: Does CoT also possess the ability to learn the underlying mechanisms for other tasks beyond Arithmetic Expression Extrapolation? \n\nThanks for the question. During the tight discussion period, we have completed experiments for longest increasing subsequence task. The results of the LIS task are shown in the table below. Here, we train the autoregressive Transformer model with various input length ranging from 1 to 80, and test the model using longer input sequence lengths unseen during training. It can be seen that the model can still extrapolate well to longer sequences. \n\n   | **Length**   | 82 | 84 | 86 | 88 | 90 | \n   | ------------ | -- | -- | -- | -- | -- |\n   | **Accuracy** | 94.3% | 92.8% | 90.7% | 89.6% | 87.4% | \n\nWe are currently running experiments on the linear equation task. But it cannot be finished before the rebuttal deadline given very limited GPU memory resources (the required CoT length is very long even for solving linear equations with 7 variables). We will report the performance when it is finished using more advanced GPUs.\n\n> Q4: Have you considered the impact of the quality of intermediate steps on the performance, such as inserting some invalid intermediate steps or omitting a certain number of intermediate steps?\n\nThis is a good catch as real-world language model training data often involves corrupted or omitted intermediate steps. To assess this, we conducted experiments on arithmetic task with varying rates of corruption and omission, denoted as γ. Specifically, γ = 0.1 indicates that we skip 10% intermediate steps and corrupt 10% steps with a single-token random replacement. The table below presents the experimental results:\n\n   | γ | Accuracy |\n   | ---- | -------- |\n   | 0.1 | 98.5% |\n   | 0.2 | 97.6% |\n   | 0.3 | 95.8% |\n\nThe results clearly demonstrate the robustness of training CoT demonstrations, showcasing its ability to maintain high accuracy even in the presence of imperfect intermediate steps in the training datasets. We will add those robustness evaluation into the final version of the paper.\n\nWe hope these clarifications can address your questions satisfactorily and we are happy to delve further into any \nof these aspects."}]

Author rebuttal:
We would like to express our sincere thanks to the reviewers and the area chair for taking the time to review our paper. We have responded to each reviewer's comments separately and will incorporate their suggestions into the next version of our paper. We hope that our response can adequately addresse the reviewers' concerns, and we're happy to provide more details about any questions they may have.

Meta-review: This paper provides new theoretical understandings of the Chain-of-Thought (CoT) technique, by proving rigorous separation results between the CoT (autoregressive) mode and the direct mode for transformers to solve arithmetic and linear equation solving problems. The positive results for CoT are further extended to dynamical programming problems. The paper also provides empirical validations.

All reviewers are strongly positive about the research direction and the contributions of the results, which I agree with. I believe the paper contains multiple results (lower bounds via circuit theory, concrete transformer constructions, the dynamical programming problem framework, and the experiments) that are worthy to be highlighted in the community and may inspire future research. Therefore, I recommend acceptance with oral presentation, and congratulate the authors for the nice work.
Decision: Accept (oral)