SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Formally, a general DP algorithm can be characterized via three key ingredients: state space  $\mathcal{I}$ , transition function  $T$ , and aggregation function  $A$ . Given a DP problem with  $N$  input sequences  $\boldsymbol{s}^{(1)}, \dots, \boldsymbol{s}^{(N)}$ , denote the problem size to be the vector  $\boldsymbol{n} = (|\boldsymbol{s}^{(1)}|, \dots, |\boldsymbol{s}^{(N)}|)$ . Fixing the problem size  $\boldsymbol{n}$ , there is an associated state space  $\mathcal{I}_n \subset \mathcal{I}$  representing the finite set of decomposed subproblems, where each state  $i \in \mathcal{I}_n$  is an index signifying a specific subproblem. The size of the state space  $\mathcal{I}_n$  grows with the problem size  $\boldsymbol{n}$ . We denote by  $\mathrm{dp}(i)$  the answer of subproblem  $i$  (as well as other information stored in the DP process). Furthermore, there is a partial order relation between different states: we say state  $j$  precedes state  $i$  (denoted as  $j \prec i$ ) if subproblem  $j$  should be solved before subproblem  $i$ , i.e., the value of  $\mathrm{dp}(i)$  depends on  $\mathrm{dp}(j)$ . This partial order creates a directed acyclic graph (DAG) within the state space, thereby establishing a reasoning chain where subproblems are resolved in accordance with the topological ordering of the DAG.
The transition function  $T$  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
$$
\mathrm {d p} (i) = T \left(\boldsymbol {n}, \boldsymbol {s}, i, \{(j, \mathrm {d p} (j)): j \prec i \}\right), \tag {4}
$$
where  $s$  is the concatenation of all input sequences  $s^{(1)},\dots ,s^{(N)}$ . In this paper, we focus on a restricted setting where each state  $i$  only depends on (i) a finite number of tokens in the input sequence  $s$  and (ii) a finite number of previous states. Under this assumption, we can rewrite (4) into a more concrete form:
$$
\begin{array}{l} \mathrm {d} \mathfrak {p} (i) = f \left(\boldsymbol {n}, i, s _ {\boldsymbol {g} (\boldsymbol {n}, i)}, \mathrm {d} \mathfrak {p} (\boldsymbol {h} (\boldsymbol {n}, i))\right) \\ \begin{array}{l} = f \left(\boldsymbol {n}, i, s _ {\boldsymbol {g} (\boldsymbol {n}, i)}, \mathrm {d p} (\boldsymbol {h} (\boldsymbol {n}, i))\right) \\ = f \left(\boldsymbol {n}, i, s _ {g _ {1} (\boldsymbol {n}, i)}, \dots , s _ {g _ {J} (\boldsymbol {n}, i)}, \mathrm {d p} (h _ {1} (\boldsymbol {n}, i)), \dots , \mathrm {d p} (h _ {K} (\boldsymbol {n}, i))\right), \end{array} \tag {5} \\ \end{array}
$$
where functions  $f, g, h$  fully determine the transition function  $T$  and have the following form  $f: \mathbb{N}^N \times \mathcal{I} \times \mathcal{X}^J \times \mathcal{Y}^K \to \mathcal{Y}, g: \mathbb{N}^N \times \mathcal{I} \to (\mathbb{N} \cup \{\emptyset\})^J, h: \mathbb{N}^N \times \mathcal{I} \to (\mathcal{I} \cup \{\emptyset\})^K$ . Here, the state space  $\mathcal{I}$ , input space  $\mathcal{X}$ , and DP output space  $\mathcal{Y}$  can be arbitrary domains, and  $J, K$  are constant integers. If state  $i$  depends on less than  $J$  input tokens or less than  $K$  previous states, we use the special symbol  $\emptyset$  to denote a placeholder, such that all terms  $s_{\emptyset}$  and  $\mathrm{dp}(\emptyset)$  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:
$$
A \left(\{(i, \mathrm {d} \mathfrak {p} (i)): i \in \mathcal {I} _ {\boldsymbol {n}} \}\right) = u \left(\square_ {i \in \mathcal {A} _ {\boldsymbol {n}}} \mathrm {d} \mathfrak {p} (i)\right), \tag {6}
$$
where  $\mathcal{A}_n\subset \mathcal{I}_n$  is a set of states that need to be aggregated,  $\square$  is an aggregation function such as min, max, or  $\sum$ , and  $u:\mathcal{V}\to \mathcal{Z}$  is any function where  $\mathcal{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 "Introduction to Algorithms" book [17] (see Appendix G.1 for 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
Input	A string s of length n	Two strings s(1), s(2) of length n1 = |s(1)| and n2 = |s(2)|, concatenated together
State space	{(j,k):j∈[n],k∈{0,···,j-1}}	{0,···,n1} × {0,···,n2}
Transition function	dp(j,k)={1 if k=0 max(dp(j,k-1), dp(k,k-1)× if k>0 [I[sj>s_k]+1)} dp(j,k)={ak if j=0 bj if k=0 min(dp(j,k-1)+a, dp(j-1,k)+b, dp(j-1,k-1) otherwise +cI[sj(1)≠sk(2)]) dp(n1,n2)
Aggregation function
4.2 Theoretical results
We begin by investigating whether LLMs 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:
$$
\mathbf {s} ^ {(1)} \quad | \quad \dots \quad | \quad \mathbf {s} ^ {(N)} \quad | \quad (i _ {1}, \mathrm {d p} (i _ {1})) \quad \dots \quad (i _ {| \mathcal {I} _ {n} |}, \mathrm {d p} (i _ {| \mathcal {I} _ {n} |})) \quad \text {f i n a l a n s w e r}
$$
Here, the input sequence consists of  $N$  strings separated by special symbols, and  $(i_1,\dots ,i_{|\mathcal{I}_n|})$  is a feasible topological ordering of the state space  $\mathcal{I}_n$ . We assume that all domains  $\mathcal{I}$ ,  $\mathcal{X}$ ,  $\mathcal{Y}$ ,  $\mathcal{Z}$  belong to the real vector space so that their elements can be effectively represented and handled by a neural network. Each  $(i,\mathrm{dp}(i))\in \mathcal{I}\times \mathcal{V}$  above will be represented as a single vector and generated jointly in the CoT output. We further assume that  $\mathcal{I}$ ,  $\mathcal{X}$ ,  $\mathcal{Y}$ ,  $\mathcal{Z}$  are discrete spaces (e.g., integers) so that the elements can be precisely represented using finite precision. To simplify our analysis, we consider a 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 3). Instead, the Transformer output is simply projected to the nearest element in the corresponding discrete space (e.g.,  $\mathcal{I}\times \mathcal{V}$  or  $\mathcal{Z}$ ). 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, 1].
Before presenting our main result, we make the following assumptions:
Definition 4.1 (Polynomially-efficient approximation). Given neural network  $P_{\theta}$  and target function  $f: \mathcal{X}^{\mathrm{in}} \to \mathcal{X}^{\mathrm{out}}$  where  $\mathcal{X}^{\mathrm{in}} \subset \mathbb{R}^{d_{\mathrm{in}}}$  and  $\mathcal{X}^{\mathrm{out}} \subset \mathbb{R}^{d_{\mathrm{out}}}$ , we say  $f$  can be approximated by  $P_{\theta}$  with polynomial efficiency if there exist  $\rho > 0$ ,  $\lambda > 0$  such that for any error  $\epsilon > 0$  and radius  $R > 0$ , there exists parameter  $\theta$  satisfying that (i)  $\| f(\boldsymbol{x}) - P_{\theta}(\boldsymbol{x} + \boldsymbol{\delta}) \|_{\infty} < \epsilon + \lambda \| \boldsymbol{\delta} \|_{\infty}$  for all  $\boldsymbol{x} \in \mathcal{X}^{\mathrm{in}}$ ,  $\| \boldsymbol{x} \|_{\infty} \leq R$  and all  $\| \boldsymbol{\delta} \|_{\infty} < \rho$ ; (ii) all elements of parameter  $\theta$  are bounded by  $O(\mathrm{poly}(R, 1 / \epsilon))$ .
Assumption 4.2. The size of the state space can be polynomially upper bounded by the problem size  $n$ , i.e.,  $|\mathcal{I}_n| = O(\mathrm{poly}(|s|))$ . Similarly, all input elements, DP values, and answers are polynomially upper bounded by the problem size  $n$ .
Assumption 4.3. Each function  $f$ ,  $g$ ,  $h$  and  $u$  in (5) and (6) can be approximated with polynomial efficiency by a perceptron of constant size (with GeLU activation).
Assumption 4.4. The function  $F: \mathbb{N}^N \times \mathcal{I} \to \mathcal{I}$  defined as  $F(\pmb{n}, i_k) = i_{k+1}$  for  $\pmb{n} \in \mathbb{N}^N$ ,  $k \in [|\mathcal{I}_n| - 1]$  can be approximated with polynomial efficiency by a perceptron of constant size (with GeLU activation), where  $(i_1, \dots, i_{|\mathcal{I}_n|})$  is a feasible topological ordering of the state space  $\mathcal{I}_n$ .
Assumption 4.5. The function  $F: \mathbb{N}^N \times \mathcal{I} \to \{0,1\}$  defined as  $F(\pmb{n},i) = \mathbb{I}[i \in \mathcal{A}_n]$  (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-precision, and Assumptions 4.3 to 4.5 guarantee that all basic functions that determine the DP process can be well-approximated by a composition of finite log-precision Transformer layers of constant size. In Appendix G.1, we show these assumptions are satisfied for LIS and ED problems described above as well as the CFG Membership Testing problem in Theorem 4.8.
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 a proof in Appendix G.2.
Theorem 4.7. Consider any DP problem satisfying Assumptions 4.2 to 4.5. For any integer  $n \in \mathbb{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  $n$ . Moreover, all parameter values are bounded by  $O(\mathrm{poly}(n))$ .
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 is 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-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  $\mathsf{TC}^0\neq \mathsf{NC}^1$ . There exists a context-free language such that for any depth  $L$  and any polynomial  $Q$ , there exists a sequence length  $n\in \mathbb{N}$  where no log-precision autoregressive transformer with depth  $L$  and hidden dimension  $d\leq 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 is intrinsically hard to be solved by a well-parallelized computation model. Combined with these results, we conclude that CoT plays a critical role in tackling tasks that are inherently difficult.
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 1. For the LIS task, the goal is to find the length of the longest increasing subsequence of a given integer sequence. For the ED task, the goal is to calculate the minimum cost required (called edit distance) to 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-range 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,6\}$ . 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 of 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 samples for each training dataset and 0.1M for testing while ensuring that duplicate samples between training 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 AdamW optimizer [39] with  $\beta_{1} = 0.9$ ,  $\beta_{2} = 0.999$ ,  $\mathrm{lr} = 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-likelihood 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-likelihood 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-token (using greedy search) until generating the End-of-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.
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.3, 3.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.
Figure 2: Model performance on different tasks. For all tasks and various difficulty levels, autoregressive Transformers with CoT consistently outperform Transformers trained on direct datasets. In particular, 3-layer Transformers already succeed in these tasks with almost perfect accuracy, while deeper Transformers ( $L = 3, 4, 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 on low-quality datasets. To investigate this question, we construct corrupted datasets for the arithmetic task in Appendix I. 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 on 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 operators ranging from 1 to 15, and test the model on expressions with the number of operators in  $\{16, 17, 18\}$ . As shown in Figure 3, our three-layer Transformer model still performs well on longer sequences, suggesting that the model indeed learns the solution to some extent (instead of memorizing data distributions). Potentially, we believe models trained on more data with varying lengths can eventually reveal the complete arithmetic rules.
Figure 3: Performance of length extrapolation experiment, tested on sequences that are longer than those in training.
Owing to the tremendous success of Transformers and Large Language Models across diverse domains, there has been a substantial body of works dedicated to theoretically comprehending their capabilities and limitations. Initially, researchers primarily focused on exploring the expressive power of Transformers in the context of function approximation. Yun et al. [67] proved that Transformers with sufficient size can universally approximate arbitrary continuous sequence-to-sequence functions on a compact domain. Recently, universality results have been extended to model variants such as Sparse Transformers [68] and Transformers with relative positional encodings (RPE) [41].
More relevant to this paper, another line of works investigated the power of Transformers from a computation perspective. Early results have shown that both standard encoder-decoder Transformers [58] and looped Transformer encoders are Turing-complete [49, 47, 20, 8]. 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 a shallow non-recursive Transformer provably suffices. Yao et al. [66] demonstrated that Transformers can recognize or generate bounded-depth Dyck language [30], a specific type of context-free language. Besides affirmative results, other works characterized the expressivity limitation of Transformers via the perspective of modeling formal languages [26, 7, 62, 66, 14] or simulating circuits [27, 43, 42]. However, none of these works (except [66]) explored the setting of autoregressive Transformers typically adopted in LLMs, which we study in this paper. Moreover, we consider a more practical setting that targets the emergent ability of LLMs in solving basic reasoning 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 learnability exhibited by LLMs [9]. Garg et al. [24] demonstrated that autoregressive Transformers can in-context learn basic function classes (e.g., linear functions, MLPs, and decision trees) via input sample sequences. Subsequent works further revealed that Transformers can implement learning algorithms such as linear regression [1], gradient descent [1, 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-of-Thought prompting is essential in solving mathematical and decision-making problems. Focusing on two basic mathematical problems as 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 to 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 to 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 LLMs. 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 it is an important future direction for theoretically studying how LLMs can generalize from CoT demonstrations (even in the out-of-distribution setting, e.g., length extrapolation (Figure 3)) [60, 13]. 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 LLMs.
Acknowledgement. This work is supported by National Key R&D Program of China (2022ZD0114900) and National Science Foundation of China (NSFC62276005), and is partially funded by Microsoft Research Asia Collaborative Research Project. The authors are grateful to David Chiang, who pointed out a mistake regarding the P-completeness of CFG Membership Testing in the early version of this paper. The authors also thank all reviewers for their valuable suggestions.