SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Let  $D = (\mathcal{Q}, \Sigma, \delta, \mathcal{F}, q_0)$  be any automaton, where  $\mathcal{Q}$  is a set of states,  $\Sigma$  is a set of symbols (alphabet),  $\delta: \mathcal{Q} \times \Sigma \to \mathcal{Q}$  is the transition function,  $\mathcal{F} \subset \mathcal{Q}$  is a set of accept states, and  $q_0$  is the initial state. For any input string  $\omega_1\omega_2\dots \omega_n$ , whether  $D$  accepts the string can be reduced into solving a system of linear equations defined as follows. The system of linear equations has  $(n + 1)|\mathcal{Q}| + 1$  variables, which we denote as  $x^{*}$  and  $x_{i,q}$  ( $i \in \{0,\dots ,n\} , q \in \mathcal{Q}$ ). The equations are defined as follows:
$$
\left\{ \begin{array}{l l} x ^ {*} = \sum_ {q \in \mathcal {F}} x _ {n, q} \\ x _ {0, q _ {0}} = 1 \\ x _ {0, q} = 0 & \text {f o r} q \in \mathcal {Q} \backslash \{q _ {0} \} \\ x _ {i, q} = \sum_ {\delta (r, \omega_ {i}) = q} x _ {i - 1, r} & \text {f o r} 0 <   i \leq n, q \in \mathcal {Q} \end{array} \right.
$$
It is easy to see that  $x_{i,q} = 1$  iff the automaton arrives at state  $q$  when taking the substring  $\omega_1\omega_2\dots \omega_i$  as input. Therefore,  $x^{*} = 1$  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.
Furthermore, the coefficient of each equation only depends on at most one input symbol. This implies that these equations can be efficiently constructed using a highly parallelizable algorithm within a complexity of  $\mathsf{NC}^0$ . Therefore, by reduction, we obtain that the problem of judging whether there exists a solution such that  $x^{*} = 1$  is  $\mathsf{NC}^1$ -hard.
Now consider solving linear equations using a Transformer without CoT. While the output of the Transformer contains multiple tokens, we can arrange the order of variables such that the Transformer has to output the value of  $x^{*}$  first. The parallel complexity of outputting the first token is bounded by  $\mathsf{TC}^0$  according to [42]. Therefore, it cannot judge whether there exists a solution satisfying  $x^{*} = 1$ .
F Discussion on Other Architectures
F.1 Encoder-Decoder Transformer
In this paper, we choose the autoregressive Transformer as it is simple and the de facto standard for LLMs (e.g., GPT). However, all of the theoretical results of this paper can be easily transferred to an encoder-decoder Transformer (such as T5) using the following arguments.
- Given a bounded-depth polynomial-size log-precision encoder-decoder Transformer, its computation complexity is still bounded by  $\mathsf{TC}^0$ , so our negative results hold.
- 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 the integration of cross-attention and self-attention. Thus, all positive results in this paper are not exclusive to autoregressive Transformers.
F.2 Recurrent Neural Network(RNN)
Although the theorems in this paper can be naturally extended to other popular Transformer architectures, such as the encoder-decoder Transformer (e.g., T5), RNNs do not have this property. Theorem F.1 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  $\Omega\left(\frac{n}{\log n}\right)$ , where  $n$  is 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  $o\left(\frac{n}{\log n}\right)$  nor generate the CoT solution for the problem Equation $(n, p)$  with the hidden dimension of  $o\left(\frac{n^2}{\log n}\right)$ .
Proof. Given an RNN model, if the hidden embeddings in each layer of the last symbols of two different input sequences are 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  $O(\log n)$  bits (by definition of log-precision). Therefore, an RNN with a hidden dimension of  $D$  can only generate at most  $n2^{D}$  different output sequences. However, for the problem Arithmetic $(n,p)$ , 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  $O(2^{O(n)})$  different solution sequences for the problem Arithmetic $(n,p)$ . Similarly, there are at least  $O(2^{O(n^2)})$  different solution sequences for the problem Equation $(n,p)$ . Thus, by the Pigeon Hole Principle, to be able to generate all  $2^{O(n)}$  different output sequences, we must have a hidden dimension of  $D = \Omega\left(\frac{n}{\log n}\right)$  for the problem Arithmetic $(n,p)$  and a hidden dimension of  $D = \Omega\left(\frac{n^2}{\log n}\right)$  for the problem Equation $(n,p)$ .
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.
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  $s \in \mathbb{N}^n$ . Formally,  $\tilde{s}$  is a subsequence of  $s$  if there exists indices  $1 \leq i_1 \leq i_2 \leq \dots \leq i_{|\tilde{s}|} \leq n$  such that  $\tilde{s}_k = s_{i_k}$  holds for all  $k \in [| \tilde{s} |]$ . A sequence  $\tilde{s}$  is called increasing if  $\tilde{s}_1 < \tilde{s}_2 < \dots < s_{|\tilde{s}|}$ . The LIS problem aims to find an increasing subsequence of  $s$  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  $\mathrm{dp}(i)$ . It is easy to write the transition function as follows:
$$
\mathrm {d p} (i) = 1 + \max  _ {j <   i, s _ {j} <   s _ {i}} \mathrm {d p} (j). \tag {20}
$$
The final answer will be  $\max_{i\in [n]}\mathsf{d}\mathsf{p}(i)$
However, the above DP transition function does not match the form of (5), since  $\mathrm{dp}(i)$  may depend on (an unbounded number of) all previous  $\mathrm{dp}(j)$  ( $j < i$ ). Nevertheless, this issue can be easily addressed by using a different DP formulation. Let  $\mathrm{dp}(j,k)$  be the longest increasing subsequence that ends at position  $j$  and the second last position is no more than  $k$  ( $k < j$ ). In this case, it is easy to write the transition function as follows:
$$
\mathsf {d p} (j, k) = \left\{ \begin{array}{l l} 1 & \text {i f} k = 0 \\ \max  (\mathsf {d p} (j, k - 1), \mathsf {d p} (k, k - 1) \cdot \mathbb {I} [ s _ {j} > s _ {k} ] + 1) & \text {i f} k > 0 \end{array} \right. \tag {21}
$$
The final answer will be  $\max_{i\in [n]}\mathsf{d}\mathsf{p}(i,i - 1)$ . This DP formulation fits our framework (5).
Edit Distance (ED). The ED problem aims to find the minimum operation cost required to convert a sequence  $\mathbf{u} \in \Sigma^{n_1}$  to another sequence  $\mathbf{v} \in \Sigma^{n_2}$ . There are three types of operations: inserting a letter into any position, deleting a letter from any position, and replacing a letter at any position by a new one. The costs of insert, delete, and replace are  $a, b$ , and  $c$ , respectively. These operations are 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  $u_1u_2 \cdots u_j$  to the substring  $v_1v_2 \cdots v_k$ , which we denote as  $\mathsf{dp}(j,k)$ . It is easy to write the transition function as follows:
$$
\mathrm {d p} (j, k) = \left\{ \begin{array}{l l} a k & \text {i f} j = 0 \\ b j & \text {i f} k = 0 \\ \min  (\mathrm {d p} (j, k - 1) + a, \mathrm {d p} (j - 1, k) + b, \mathrm {d p} (j - 1, k - 1) + c \mathbb {I} [ s _ {j} ^ {(1)} \neq s _ {k} ^ {(2)} ]) & \text {o t h e r w i s e} \end{array} \right. \tag {22}
$$
The final answer will be  $\mathrm{dp}(n_1, n_2)$ . This DP formulation fits our framework (5).
CFG Membership Testing. A context-free grammar (CFG) is a 4-tuple, denoted as  $G = (\mathcal{V}, \Sigma, R, S)$ , where
-  $\mathcal{V}$  is a finite set of non-terminal symbols.
-  $\Sigma$  is a finite set of terminal symbols, disjoint from  $\nu$ .
-  $R$  is a finite set of production rules, where each rule has the form  $A \rightarrow \beta$ , with  $A \in \mathcal{V}$  and  $\beta \in (V \cup \Sigma)^*$  (the asterisk represents the Kleene star operation).
-  $S \in \mathcal{V}$  is the start symbol.
The non-terminal symbols in  $\mathcal{V}$  represent (abstract) syntactic categories, while the terminal symbols in  $\Sigma$  represent the actual words/tokens of the language. The production rules in  $R$  specify how a non-terminal symbol can be replaced by sequences of terminal and non-terminal symbols concatenated together. Below, we focus on a canonical version of CFG [33], in which the term  $\beta$  in each rule is either an empty string or contains exactly two symbols:
-  $A \to \epsilon$ , where  $A \in \mathcal{V}$  and  $\epsilon$  is the empty string;
-  $A \to BC$ , where  $B, C \in \mathcal{V} \cup \Sigma$ .
By introducing additional non-terminal symbols to split complex rules, any CFG can be easily translated into the canonical version expressing the same language.
The (universal) CFG Membership Testing problem is defined as follows: given a canonical CFG  $G$  and a string  $v$ , judge whether  $v$  can be generated from  $G$ . This problem is known to be P-complete [33]. The CYK algorithm [53], which is based on DP, is a 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 \rightarrow \epsilon$ . Denote  $R = \{R_1, \dots, R_m\}$  and let  $n = |v|$ . The state space of the DP process is defined as
$$
\mathcal {I} _ {| \mathcal {V} |, m, n} = \left\{\left(t, i, j, k, A, r\right): 0 \leq t \leq | \mathcal {V} |, 0 \leq i \leq k \leq j \leq n, A \in \mathcal {V}, 0 \leq r \leq m \right\},
$$
where  $\mathsf{dp}(t,i,j,k,A,r)$  stores whether the substring  $v_{i + 1}\dots v_j$  can be generated by nonterminal  $A$  by first using rule  $R_{\tilde{r}}:A\to BC$  with  $\tilde{r}\leq r$ , and there exists index  $\tilde{k}\leq k$  such that the substring  $v_{i + 1}\dots v_{\tilde{k}}$  can be generated by  $B$  and the substring  $v_{\tilde{k} +1}\dots v_j$  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 = k$ ,  $\mathsf{dp}(0,i,i,i,A,r)$  simply stores whether nonterminal  $A$  can generate an empty string by first using rule  $R_{\tilde{r}}$  with  $\tilde{r}\leq r$ . The transition function can be formally written as
$$
\begin{array}{l} \mathrm {d p} (t, i, j, k, A, r) = \\ \left\{ \begin{array}{l l} 0 & \text {i f i = j = k , t = 0 , r = 0}, \\ \mathbb {I} \left[ \mathsf {d p} (t, i, j, k, A, r - 1) = 1 \text {o r} R _ {r}: A \to \epsilon \right] & \text {i f i = j = k , t = 0 , r \geq 1}, \\ 0 & \text {i f i <   j , t = 0}, \\ \mathsf {d p} (t - 1, i, j, j, A, m) & \text {i f t > 0 , r = 0 , k = i} \\ \mathsf {d p} (t, i, j, k - 1, A, m) & \text {i f t > 0 , r = 0 , k > i}, \\ \mathbb {I} [ \mathsf {d p} (t, i, j, k, A, r - 1) = 1 \vee (R _ {r}: A \to B C \text {f o r s o m e} B C \text {s . t .}} \\ (B \in \Sigma \wedge k = i + 1 \wedge v _ {k} = B) \vee (B \in \mathcal {V} \wedge \mathsf {d p} (t, i, k, k, B, m) = 1), \\ (C \in \Sigma \wedge j = k + 1 \wedge v _ {j} = C) \vee (C \in \mathcal {V} \wedge \mathsf {d p} (t, k, j, j, C, m) = 1)) ] & \text {i f t > 0 , r > 0}. \end{array} \right. \\ \end{array}
$$
The final answer will be  $\mathrm{dp}(|\mathcal{V}|, 0, n, n, S, m)$ . This DP formulation fits our framework  $^{5}$ .
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, b \in \mathbb{R}$ ;
-  $\mathbb{I}[a \neq b]$ ,  $\mathbb{I}[a < b]$ ,  $\mathbb{I}[a > b]$ , where  $a, b \in \mathbb{Z}$ ;
-  $a \times b$ , where  $a \in \mathbb{R}$ ,  $b \in \{0,1\}$ ;
- linear transformation;
- conditional selection (Lemma C.4), for example,
$$
f ^ {\mathrm {s e l e c t}} (x) = \left\{ \begin{array}{l l} f ^ {>} (x) & \text {i f} x \geq 0, \\ f ^ {<  } (x) & \text {i f} x <   0, \end{array} \right.
$$
where  $f^>$  and  $f^<$  are functions that can be implemented by MLPs with ReLU activation, and  $x \in \mathbb{Z}$ .
This implies that the MLP with ReLU activation can 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 a perceptron of constant size with GeLU activation. Similarly, the topological ordering can also be efficiently implemented by MLPs with GeLU activation:
LIS:  $(j,k)\to \left\{ \begin{array}{ll}(j,k + 1) & \mathrm{if}k <   j - 1\\ (j + 1,0) & \mathrm{if}k = j - 1 \end{array} \right.$
- ED:  $(j,k) \to \left\{ \begin{array}{ll}(j,k + 1) & \text{if } k < n_2 \\ (j + 1,0) & \text{if } k = n_2 \end{array} \right.$
There is a subtle mismatch between the DP formulation and our framework (5), in that the terms  $B$  and  $C$  depend on the input (rather than purely determined by the state). Nevertheless, it is easy to extend Theorem 4.7 to this setting. Specifically, an autoregressive Transformer can first extract  $B$  and  $C$  using one layer and then extract  $\mathsf{dp}(t,i,k,k,B,m)$  and  $\mathsf{dp}(t,k,j,j,C,m)$ ) using an additional layer.
CFG Membership Testing:
$$
(t, i, j, k, A, r) \to \left\{ \begin{array}{l l} (t, i, j, k, A, r + 1) & \text {i f} r <   m \\ (t, i, j, k, \mathsf {n e x t} (A), 0) & \text {i f} r = m, A \neq \mathsf {l a s t} (\mathcal {V}) \\ (t, i, j, k + 1, \mathsf {f i r s t} (\mathcal {V}), 0) & \text {i f} r = m, A = \mathsf {l a s t} (\mathcal {V}), k <   j \\ (t, i + 1, j + 1, i + 1, \mathsf {f i r s t} (\mathcal {V}), 0) & \text {i f} r = m, A = \mathsf {l a s t} (\mathcal {V}), k = j <   n \\ (t + 1, i + 1, j + 1, i + 1, \mathsf {f i r s t} (\mathcal {V}), 0) & \text {i f} r = m, A = \mathsf {l a s t} (\mathcal {V}), k = j = n, t <   | \mathcal {V} | \\ (0, 0, j - i + 1, 0, \mathsf {f i r s t} (\mathcal {V}), 0) & \text {i f} r = m, A = \mathsf {l a s t} (\mathcal {V}), k = j = n, t = | \mathcal {V} | \end{array} \right.
$$
where we order the set  $\mathcal{V}$  and denote first  $(\mathcal{V})$  as the first element of  $\mathcal{V}$ , last  $(\mathcal{V})$  as the last element of  $\mathcal{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 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  $s$  of length no more than  $n$ . Moreover, all parameter values are bounded by  $O(\mathrm{poly}(n))$ .
Proof. Input Format. Assume that we have a sequence of tokens  $s_1, \dots, s_t$  and we want to generate the next token  $s_{t+1}$ . We embed the token  $s_k$  by
$$
\boldsymbol {x} _ {k} ^ {(0)} = \left(\boldsymbol {e} _ {k} ^ {\text {i n p u t}}, \boldsymbol {e} _ {k} ^ {\text {s t a t e}}, \boldsymbol {e} _ {k} ^ {\text {d p}}, \boldsymbol {e} _ {k} ^ {\text {a n s w e r}}, \boldsymbol {e} _ {k} ^ {\text {s e p}}, k, 1\right),
$$
where each part of the embedding is defined as follows:
1.  $e_k^{\mathrm{input}}$  is the embedding of the input token  $s_k \in \mathcal{X}$ . If the current position does not represent an input token, then  $e_k^{\mathrm{input}} = 0$ .
2.  $e_k^{\mathrm{state}}$  is the embedding of the DP state in  $\mathcal{I}$  at position  $k$ . If the current position corresponds to an input token or the final answer, then  $e_k^{\mathrm{state}} = 0$ . We also assume that for all  $i \in \mathcal{I}$ , the embedding of state  $i$  is non-zero.
3.  $e_k^{\mathrm{dp}}$  is the embedding of the DP value in  $\mathcal{V}$  at position  $k$ . If the current position corresponds to an input token or the final answer, then  $e_k^{\mathrm{dp}} = 0$ .
4.  $e_k^{\text{answer}}$  is the embedding of the answer token in  $\mathcal{Z}$ , and  $e_k^{\text{answer}} = 0$  if the current position corresponds to an input token or an intermediate DP position.
5.  $e_k^{\mathrm{sep}}$  is the embedding of the separator | separating different input sequences. We set  $e_k^{\mathrm{sep}} = e_j$  if the current token  $s_k$  is the  $j$ -th separator, where  $e_j$  is the one-hot vector with the  $j$ -th element begin 1.
6. The position embedding  $k$  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  $p_k^{\mathrm{sep},1}, \dots, p_k^{\mathrm{sep},N} \in \mathbb{N}$ .
Similar to the previous proofs, this can be achieved via the COPY operation with  $S_{k} = \{j\leq k:$ $e_j^{\mathrm{sep}} = e_t\}$  for  $t\in [N]$  and  $v_{j} = j$ . Then, several MLPs follow, which perform the following tasks:
Calculate the problem size  $\pmb{n}_k = (p_k^{\mathrm{sep},1} - 1, p_k^{\mathrm{sep},2} - p_k^{\mathrm{sep},1} - 1, \dots, p_k^{\mathrm{sep},N} - p_k^{\mathrm{sep},N-1} - 1)$ .
- Obtain the next state  $e_k^{\text{next_state}}$ . If the current state is already the last state, set  $e_k^{\text{next_state}} = 0$ .
The first task is a linear transformation, which can clearly be processed by an MLP Proposition C.3. According to Assumption 4.4, we can use an MLP to compute the embedding of the next state  $e_k^{\mathrm{next\_state}}$  based on the embedding of the current state  $e_k^{\mathrm{state}}$  and the problem size  $n$ . When the required MLP in Assumption 4.4 has multiple layers (i.e.,  $\tilde{L}$  layers), we can use  $\tilde{L} - 1$  Transformer layers
to implement a  $\tilde{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 of this block is
$$
\boldsymbol {x} _ {k} ^ {(1)} = \left(\boldsymbol {e} _ {k} ^ {\text {i n p u t}}, \boldsymbol {e} _ {k} ^ {\text {s t a t e}}, \boldsymbol {e} _ {k} ^ {\text {n e x t - s t a t e}}, \boldsymbol {e} _ {k} ^ {\text {d p}}, \boldsymbol {e} _ {k} ^ {\text {s e p}}, \boldsymbol {n} _ {k}, k, 1\right).
$$
Block 2. The second layer of the Transformer does not use attention heads. It only uses the MLP to perform the following tasks:
Calculate  $h(\pmb {n}_k,\pmb {e}_k^{\mathrm{next\_state}})$  and  $\pmb {g}(\pmb {n}_k,\pmb {e}_k^{\mathrm{next\_state}})$  . We assume that the embedding of  $\varnothing$  is 0.
- Set the flag  $f_{k}^{\mathrm{state}}$  representing whether current state  $e_{k}^{\mathrm{state}}$  is the last state.
- Set the flag  $f_{k}^{\text{answer}}$  representing whether current state  $e_{k}^{\text{state}}$  is in the set  $\mathcal{A}$ , i.e., used in the aggregation function.
Similar to the first block, we stack several two-layer perceptrs to implement a multilayer perceptron. According to Assumptions 4.3 and 4.5, we can use an MLP to complete the first and the last tasks. The second task can be done by checking whether  $e_k^{\text{state}} \neq 0$  and  $e_k^{\text{next_state}} = 0$ . We also compute the auxiliary quantities  $(h(n_k, e_k^{\text{next_state}}))^2$ ,  $(g(n_k, e_k^{\text{next_state}}))^2$ ,  $(e_k^{\text{state}})^2$ , and  $k^2$ , which are elementwise square operations and can be implemented by an MLP (Lemma C.1). The output of this block is
$$
\begin{array}{l} \boldsymbol {x} _ {k} ^ {(2)} = \left(\boldsymbol {e} _ {k} ^ {\text {i n p u t}}, \boldsymbol {e} _ {k} ^ {\text {s t a t e}}, \boldsymbol {e} _ {k} ^ {\text {n e x t - s t a t e}}, \boldsymbol {e} _ {k} ^ {\text {d p}}, \boldsymbol {e} _ {k} ^ {\text {s e p}}, \boldsymbol {n} _ {k}, \boldsymbol {h} \left(\boldsymbol {n} _ {k}, \boldsymbol {e} _ {k} ^ {\text {n e x t - s t a t e}}\right), \boldsymbol {g} \left(\boldsymbol {n} _ {k}, \boldsymbol {e} _ {k} ^ {\text {n e x t - s t a t e}}\right), \right. \\ \left. \left(\boldsymbol {h} \left(\boldsymbol {n} _ {k}, \boldsymbol {e} _ {k} ^ {\text {n e x t - s t a t e}}\right)\right) ^ {2}, \left(\boldsymbol {g} \left(\boldsymbol {n} _ {k}, \boldsymbol {e} _ {k} ^ {\text {n e x t - s t a t e}}\right)\right) ^ {2}, \left(\boldsymbol {e} _ {k} ^ {\text {s t a t e}}\right) ^ {2}, f _ {k} ^ {\text {s t a t e}}, f _ {k} ^ {\text {a n s w e r}}, k, k ^ {2}, 1\right). \\ \end{array}
$$
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  $s_{g_1(i)}, \dots, s_{g_J(i)}$  where  $i$  corresponds to  $e_k^{\text{next_state}}$ . When  $g_t(i) = \emptyset$ , we set  $s_{g_t(i)}$  to be a special token.
- Copy the DP value embeddings corresponding to  $\mathrm{dp}(h_1(i),\dots ,\mathrm{dp}(h_K(i))$  for  $i$  corresponds to  $\pmb{e}_k^{\mathrm{next\_state}}$ . When  $h_t(i) = \emptyset$ , we set  $\mathrm{dp}(h_t(i))$  to be a special value.
- Calculate the output  $\mathrm{dp}(i)$  for  $i$  corresponds to  $e_k^{\mathrm{next\_state}}$ , denoted as  $e_k^{\mathrm{next\_dp}}$ .
The first two tasks can be done via the COPY operation. To copy DP values, the attention head attends to positions  $j$  with  $e_j^{\mathrm{state}}$  matching  $h_t(i)$  for  $t \in [K]$ . To copy input tokens, the attention head attends to positions  $j = g_t(i)$  for  $t \in [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 in (5)) using an MLP. The output of this layer is
$$
\boldsymbol {x} _ {k} ^ {(3)} = \left(\boldsymbol {e} _ {k} ^ {\text {n e x t - s t a t e}}, \boldsymbol {e} _ {k} ^ {\mathrm {d p}}, \boldsymbol {e} _ {k} ^ {\text {n e x t - d p}}, \boldsymbol {n} _ {k}, f _ {k} ^ {\text {s t a t e}}, f _ {k} ^ {\text {a n s w e r}}, k, 1\right).
$$
Block 4. The fourth block of the autoregressive transformer contains one Transformer layer. Depending the aggregation function, it uses one attention head for the operation max or min, or two attention heads for the operation  $\sum$ . This block performs the following tasks:
- Aggregate the DP values according to the aggregation function Equation (6).
- Generate the output based on the flag  $f_{k}^{\text{answer}}$ .