SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Proof. The proof technique is similar to that of Theorem 3.3. 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_{i}^{(l)}$  to denote the output at position  $i$  after the  $l$ -th Transformer layer, and use  $x_{i}$  to denote the  $i$ -th variable in linear equations. We also note that  $m$  is 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$ .
Token Embeddings. Assume that we have a sequence of tokens  $s_1, s_2, \ldots, s_t$  and we want to generate the next token  $s_{t+1}$ . We can embed the token  $s_i$  using the format  $\boldsymbol{x}_i^{(0)} = (\boldsymbol{e}_{\mathrm{id}(s_i)}, \boldsymbol{l}_i, i, 1)$ :
1. The vector  $e_i$  represents the one-hot vector with the  $i$ -th element being 1, and  $\operatorname{id}(s_i)$  is the index of token  $s_i$  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 a unified (single) encoding and distinguishing them via the term  $l_i$ . This means that if  $s_i$  and  $s_j$  are two different variables, we have  $\operatorname{id}(s_i) = \operatorname{id}(s_j)$  and  $l_i \neq l_j$ .
2.  $l_{i} \in \mathbb{R}^{3}$  is a vector used to distinguish between different variables. Its first element, denoted as  $\operatorname{var}(s_i)$ , represents the index of the variable  $s_i$ . If the token  $s_i$  is not a variable, then
$l_{i} = (0,0,0)$  and  $\operatorname {var}(s_i) = 0$  . If it is the variable  $x_{j}$  for some  $j\in [m]$  , then  $\operatorname {var}(s_i) = j$  and  $l_{i} = \big(j,m^{2}\sin (\frac{2j\pi}{m}),m^{2}\cos (\frac{2j\pi}{m})\big).$
3.  $i$  is the positional embedding, representing the position of the token in the sequence.
4. The constant embedding 1 is used as a bias term.
Layer 1. The first layer of the Transformer uses three attention heads to record some basic information:
1. Count the number of  $i$ ; (i.e., equations) in previous tokens, denoted as  $n_i^{\mathrm{eq}} = |\{j \leq i : s_j = '';\}|$ .
2. Count the number of  $\Rightarrow$  in previous tokens, denoted as  $n_i^{\mathrm{cot}} = |\{j \leq i : s_j = ' \Rightarrow '\}|$ . Namely, the current position belongs to the  $n_i^{\mathrm{cot}}$ -th CoT step.
3. Determine the number of variables in the system of linear equations. This can be done by copying  $\operatorname{var}(s_j)$  for index  $j$  such that  $s_j$  is a variable and  $\operatorname{var}(s_j)$  is the largest. Denote the result as  $n_i^{\mathrm{var}}$ . Note that according to the input format,  $n_i^{\mathrm{var}}$  is correct whenever  $n_i^{\mathrm{eq}} \geq 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  $\cdot$  and  $\Rightarrow$  tokens in all previous tokens. The last attention head performs the COPY operation with  $S_{i} = \{j:j\leq i:s_{j}$  is a variable},  $r_j = \mathrm{var}(s_j)$ , and  $v_{j} = \mathrm{var}(s_{j})$ . Note that while  $r_{j_1} = r_{j_2}$  may hold for different positions  $j_{1},j_{2}$ , their values are the same (i.e.,  $v_{j_1} = v_{j_2}$ ), so the COPY operation still works and obtains  $n_i^{\mathrm{var}}$  (when  $n_i^{\mathrm{eq}}\geq 1$ ).
Then, we use MLPs in parallel to calculate  $n_i^{\mathrm{eq}} = (n_i^{\mathrm{eq}} / i) \cdot i$  and  $n_i^{\mathrm{cot}} = (n_i^{\mathrm{cot}} / i) \cdot i$  based on Lemma C.1. Besides, we use an MLP to compute the auxiliary term  $i^2$  that will be used in the next layer. Therefore, the output of the first layer is
$$
\pmb {x} _ {i} ^ {(1)} = (\pmb {e} _ {\mathrm {i d} (s _ {i})}, \pmb {l} _ {i}, i, i ^ {2}, 1, n _ {i} ^ {\mathrm {v a r}}, n _ {i} ^ {\mathrm {e q}}, n _ {i} ^ {\mathrm {c o t}}).
$$
Layer 2. As described in Appendix B, each CoT step eliminates one variable, and thus at the current position we are eliminating variable  $x_{n_i^{\mathrm{cot}}}$ . By the uniqueness of the solution, there must exist an equation with a nonzero coefficient for variable  $x_{n_i^{\mathrm{cot}}}$ . 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  $x_{n_i^{\mathrm{cot}} + 1}$  in the next CoT step  $n_i^{\mathrm{cot}} + 1$ . We also use additional attention heads to perform some auxiliary calculations that will be used in subsequent layers. Concretely, the second layer uses four attention heads to perform the following tasks:
1. Copy the value  $n_j^{\mathrm{eq}}$  with position  $j$  corresponding to the nearest  $\Rightarrow$  token  $s_j$  ( $j \leq i$ ). Clearly, the value is well-defined when  $n_i^{\mathrm{cot}} \geq 1$ , and we define the value to be 0 if  $n_i^{\mathrm{cot}} = 0$ .
2. Compute  $d_{i}^{\mathrm{eq}} = n_{i}^{\mathrm{eq}} - n_{j}^{\mathrm{eq}} + 1$ , which corresponds to the index of the current equation in the current CoT step.
3. Copy the embedding  $e_{\mathrm{id}(s_j)}$  with the smallest  $j$  satisfying  $n_j^{\mathrm{eq}} = n_i^{\mathrm{eq}}$  and  $s_j$  is a number. Note that  $e_{\mathrm{id}(s_j)}$  is well-defined when  $s_i = ' = '$ .
4. Compute a Boolean flag (denoted as  $f_{i}$ ), which is true only when  $e_{\mathrm{id}(s_j)} \neq e_{\mathrm{id}(0)}, d_i^{\mathrm{eq}} > n_i^{\mathrm{cot}}$  and  $s_i = \text{'=}$ . The definition of  $f_{i}$  means that in the  $n_i^{\mathrm{cot}}$ -th CoT step, we only focus on the  $j$ -th equation when  $j > n_i^{\mathrm{cot}}$  and check whether the first number in the equation is non-zero. If it is non-zero, we set the flag to true at the specific position corresponding to token '='.
5. Copy the embeddings  $(e_{\mathrm{id}(s_{i - 1})},l_{i - 1})$  and  $(e_{\mathrm{id}(s_{i - 2})},l_{i - 2})$  of the  $(i - 1)$ -th and  $(i - 2)$ -th token.
The first task can be implemented by an attention head via the COPY operation to obtain  $n_j^{\mathrm{eq}}$  when  $n_i^{\mathrm{cot}} \geq 1$ . For the third task, we construct the matrices  $Q, K, V$  of the COPY operation such that
$$
\boldsymbol {q} _ {i} = \boldsymbol {Q} \boldsymbol {x} _ {i} ^ {(1)} = (- n _ {i} ^ {\mathrm {e q}}, 1, 1, 1), \quad \boldsymbol {k} _ {j} = \boldsymbol {K} \boldsymbol {x} _ {j} ^ {(1)} = \left(1, n _ {j} ^ {\mathrm {e q}}, \sum_ {a \in [ p ]} \mathbb {I} [ s _ {j} = a ], - 1\right),
$$
$\pmb{v}_j = \pmb{e}_{\mathrm{id}(s_j)}$ , and  $r_j = -j$ . By construction,  $\pmb{q}_i \cdot \pmb{k}_j = (n_j^{\mathrm{eq}} - n_i^{\mathrm{eq}}) + \sum_{a \in [p]} \mathbb{I}[s_j = a] - 1$ , and thus  $\pmb{q}_i \cdot \pmb{k}_j = 0$  only when  $n_j^{\mathrm{eq}} = n_i^{\mathrm{eq}}$  and  $s_j$  is a number, and  $\pmb{q}_i \cdot \pmb{k}_j \leq -1$  otherwise. Furthermore, the
choice of  $r_j$  guarantees that the leftmost position satisfies  $\pmb{q}_i \cdot \pmb{k}_j = 0$  is copied. This exactly solves the third task. For the fifth task, we use two attention heads to perform the COPY operation. We only give the construction of the first head that copies  $(e_{\mathrm{id}(s_{i-1})}, l_{i-1})$ . The matrices  $Q, K, V$  of the COPY operation is constructed such that
$$
\boldsymbol {q} _ {i} = \boldsymbol {Q} \boldsymbol {x} _ {i} ^ {(1)} = ((i - 1) ^ {2}, i - 1, - 1), \quad \boldsymbol {k} _ {j} = \boldsymbol {K} \boldsymbol {x} _ {j} ^ {(1)} = (- 1, 2 j, j ^ {2}), \quad \boldsymbol {v} _ {j} = (\boldsymbol {e} _ {\mathrm {i d} (s _ {j})}, \boldsymbol {l} _ {j}),
$$
and  $\pmb{q}_i\cdot \pmb{k}_j = 0$  iff  $j = i - 1$
We next use an MLP to correct the value of  $n_j^{\mathrm{eq}}$  when  $n_i^{\mathrm{cot}} = 0$  and compute the second task, which is a linear operation. We also compute an auxiliary flag  $\mathbb{I}[n_i^{\mathrm{cot}} = d_i^{\mathrm{eq}}]$  via an MLP. Regarding the fourth task, it is a multivariate conditional selection operation and can be similarly implemented by an MLP by extending Lemma C.4. Note that we can compute the second task and the fourth task in parallel using a two-layer MLP because both tasks correspond to (multivariate) conditional selection and can be merged. We finally use multiplication to compute the auxiliary terms  $(n_i^{\mathrm{cot}})^2$ ,  $(\mathrm{var}(s_i))^2$  and  $(\mathrm{var}(s_{i - 1}))^2$ . The output of the MLP is
$$
\begin{array}{l} \boldsymbol {x} _ {i} ^ {(2)} = \left(\boldsymbol {e} _ {\mathrm {i d} (s _ {i})}, \boldsymbol {l} _ {i}, i, i ^ {2}, 1, n _ {i} ^ {\text {v a r}}, n _ {i} ^ {\text {c o t}}, \left(n _ {i} ^ {\text {c o t}}\right) ^ {2}, d _ {i} ^ {\text {e q}}, f _ {i}, \right. \\ \left. \boldsymbol {e} _ {\operatorname {i d} \left(s _ {i - 1}\right)}, \boldsymbol {l} _ {i - 1}, \boldsymbol {e} _ {\operatorname {i d} \left(s _ {i - 2}\right)}, \boldsymbol {l} _ {i - 2}, \left(\operatorname {v a r} \left(s _ {i - 1}\right)\right) ^ {2}, \left(\operatorname {v a r} \left(s _ {i}\right)\right) ^ {2}, \mathbb {I} \left[ n _ {i} ^ {\text {c o t}} = d _ {i} ^ {\text {e q}} \right]\right). \\ \end{array}
$$
Layer 3. The third layer of the Transformer uses two attention heads to perform the following tasks:
1. Copy the embedding  $d_j^{\mathrm{eq}}$  with the smallest  $j$  satisfying  $f_{j} = 1$  and  $n_j^{\mathrm{cot}} = n_i^{\mathrm{cot}} - 1$ . Denote the answer as  $\hat{d}_i^{\mathrm{eq}}$ .
2. Determine whether the next token  $s_{i+1}$  is a number. Denote the result as  $f_i^{\mathrm{num}}$ .
3. Determine the output of the next token  $s_{i+1}$  if  $s_{i+1}$  is not a number. We denote its embedding as  $e_i^{\mathrm{next}}$ . Also, we need to determine the variable index  $\operatorname{var}(s_{i+1})$  of the next token if the next token is a variable.
4. Determine the token  $s_{i+2}$  if the next token  $s_{i+1}$  is a number. There are two cases:  $s_{i+2}$  is a variable, and  $s_{i+2}$  is the token ‘;’. Denote the result as  $e_i^{\mathrm{next2}}$  and  $\operatorname{var}(s_{i+2})$  and compute  $(\operatorname{var}(s_{i+2}))^2$ .
5. If the current token  $s_i$  is a variable, copy the embedding  $e_{\mathrm{id}(s_j - 1)}$  (which is a number) for index  $j$  satisfying  $n_j^{\mathrm{cot}} = n_i^{\mathrm{cot}}$ ,  $n_j^{\mathrm{cot}} = d_j^{\mathrm{eq}}$ , and  $\operatorname{var}(s_j) = \operatorname{var}(s_i)$ . Denote the answer as  $e_i^{\mathrm{cot\_num}}$ . When  $s_i$  is not a variable or  $d_i^{\mathrm{eq}} \leq n_i^{\mathrm{cot}}$ ,  $e_i^{\mathrm{cot\_num}}$  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 be done via the COPY operation. Regarding the second task, whether the next token is a number can be purely determined by  $d_{i}^{\mathrm{eq}}$ ,  $n_{i}^{\mathrm{cot}}$ , and the current token  $s_i$ . Specifically,  $s_{i+1}$  is a number if  $s_i = ' +'$ , or  $s_i = ' ='$ , or  $(s_i = ' '$  and  $d_{i}^{\mathrm{eq}} > n_{i}^{\mathrm{cot}}$ ). Whether the output of the next token is a variable can also be purely determined by the previous tokens  $s_{i-1}, s_i$  and also  $d_{i}^{\mathrm{eq}}$  and  $n_{i}^{\mathrm{cot}}$ . Specifically,  $s_{i+1}$  is a variable if  $s_{i-1} = ' +'$  and  $s_i$  is a number, or  $s_{i-1} = ' '$  and  $s_i$  is a number, or  $(s_i = ' ;'$  or  $s_i = ' \Rightarrow '$ ) and  $d_{i}^{\mathrm{eq}} \leq n_{i}^{\mathrm{cot}}$ . The variable index can be determined by either  $\operatorname{var}(s_{i-2})$  or  $d_{i}^{\mathrm{eq}}$ . When the next token is neither a variable nor a number (i.e., the symbols '+, ' =', ' ', or ' => ' ), we can similarly determine the token by checking  $s_{i-1}, s_i, d_{i}^{\mathrm{eq}}$ , and  $n_{i}^{\mathrm{var}}$ . When the next token is a number,  $s_{i+2}$  can be determined by checking the variable  $s_{i-1}$  via three cases: (i) if  $s_{i-1}$  is a variable and  $\operatorname{var}(s_{i-1}) < n_{i}^{\mathrm{var}}$ , then  $s_{i+2}$  is a variable and  $\operatorname{var}(s_{i+2}) = \operatorname{var}(s_{i-1}) + 1$ ; (ii) if  $s_{i-1}$  is a variable and  $\operatorname{var}(s_{i-1}) = n_{i}^{\mathrm{var}}$ , then  $s_{i+2} = ' ;'$ ; (iii) otherwise,  $s_{i-1}$  is a number, then  $s_{i+2}$  is a variable and  $\operatorname{var}(s_{i+2}) = n_{i}^{\mathrm{cot}} + 1$ .
All these tasks can be implemented by MLPs that perform the conditional selection or look-up table based on Lemmas C.4 and C.5. Moreover, the composition of conditional selection and lookup table can be merged into a single two-layer MLP (as shown in the construction of the third layer in arithmetic expression). We next use multiplication to compute the auxiliary terms  $(d_i^{\mathrm{eq}})^2$ ,  $(d_i^{\mathrm{eq}} + \mathbb{I}[s_{i + 2} = {}^{\prime};])^{2}$ , and  $(\hat{d}_i^{\mathrm{eq}})^2$ . However, to compute  $(\mathrm{var}(s_{i + 2}))^2$ , we cannot use multiplication directly as the composition of multiplication and conditional selection will require a deeper MLP. Instead, note that  $(\mathrm{var}(s_{i + 2}))^2$  linearly depends on  $(\mathrm{var}(s_{i - 1}))^2$  and  $\mathrm{var}(s_{i - 1})$ , or linearly depends on  $(n_i^{\mathrm{cot}})^2$  and  $n_i^{\mathrm{cot}}$ , all of which is already computed. Therefore, we can compute  $(\mathrm{var}(s_{i + 2}))^2$
without multiplication. The output of this layer has the form
$$
\begin{array}{l} \pmb {x} _ {i} ^ {(3)} = (\pmb {e} _ {\mathrm {i d} (s _ {i})}, \pmb {l} _ {i}, i, 1, n _ {i} ^ {\mathrm {v a r}}, n _ {i} ^ {\mathrm {c o t}}, (n _ {i} ^ {\mathrm {c o t}}) ^ {2}, d _ {i} ^ {\mathrm {e q}}, (d _ {i} ^ {\mathrm {e q}}) ^ {2}, (d _ {i} ^ {\mathrm {e q}} + \mathbb {I} [ s _ {i + 2} = ‘; ’ ]) ^ {2}, \hat {d} _ {i} ^ {\mathrm {e q}}, (\hat {d} _ {i} ^ {\mathrm {e q}}) ^ {2}, f _ {i} ^ {\mathrm {n u m}}, \\ \pmb {e} _ {i} ^ {\mathrm {n e x t}}, \pmb {e} _ {i} ^ {\mathrm {n e x t 2}}, \mathrm {v a r} (s _ {i}), \mathrm {v a r} (s _ {i + 1}), \mathrm {v a r} (s _ {i + 2}), (\mathrm {v a r} (s _ {i})) ^ {2}, (\mathrm {v a r} (s _ {i + 2})) ^ {2}, \pmb {e} _ {\mathrm {i d} (s _ {i - 1})}, \pmb {e} _ {i} ^ {\mathrm {c o t \_ n u m}}). \\ \end{array}
$$
Layer 4. The fourth layer of the Transformer performs the core calculation of equation coefficients when the next token is a number. There are two equations related to the calculation: the  $d_{i}^{\mathrm{eq}}$ -th equation in the last CoT step, and the  $\hat{d}_{i}^{\mathrm{eq}}$ -th equation in the last CoT step. There are also two variables related to the calculation: the variable  $x_{\mathrm{var}(s_{i+2})}$  and  $x_{n_i^{\mathrm{cot}}}$ . Specifically, we need to copy four coefficients  $a_{\hat{d}_i^{\mathrm{eq}}, n_i^{\mathrm{cot}}}, a_{\hat{d}_i^{\mathrm{eq}}, \mathrm{var}(s_{i+2})}, a_{a_i^{\mathrm{eq}}, n_i^{\mathrm{cot}}}, a_{d_i^{\mathrm{eq}}, \mathrm{var}(s_{i+2})}$  defined as follows:
$$
\text {T h e} \hat {d} _ {i} ^ {\mathrm {e q}} - \text {t h e q u a t i o n :} \quad \dots + a _ {\hat {d} _ {i} ^ {\mathrm {e q}}, n _ {i} ^ {\mathrm {c o t}}} x _ {n _ {i} ^ {\mathrm {c o t}}} + \dots + a _ {\hat {d} _ {i} ^ {\mathrm {e q}}, \operatorname {v a r} (s _ {i + 2})} x _ {\operatorname {v a r} (s _ {i + 2})} + \dots = b _ {\hat {d} _ {i} ^ {\mathrm {e q}}}
$$
$$
\text {T h e} d _ {i} ^ {\mathrm {e q}} \text {- t h e q u a t i o n :} \dots + a _ {d _ {i} ^ {\mathrm {e q}}, n _ {i} ^ {\mathrm {c o t}}} x _ {n _ {i} ^ {\mathrm {c o t}}} + \dots + a _ {d _ {i} ^ {\mathrm {e q}}, \operatorname {v a r} (s _ {i + 2})} x _ {\operatorname {v a r} (s _ {i + 2})} + \dots = b _ {d _ {i} ^ {\mathrm {e q}}}
$$
For the case of  $s_{i+2} = {}^{\prime \prime}$ , we need to copy coefficients  $b_{\hat{d}_i^{\mathrm{eq}}}$  and  $b_{d_i^{\mathrm{eq}}}$ . To unify the two cases, this Transformer layer uses four attention heads to perform the following tasks (note that we define  $\operatorname{var}(s_j) = 0$  when  $s_j$  is not a variable):
1. Copy the embedding  $\mathbf{e}_{\mathrm{id}(s_{j-1})}$  for position  $j$  satisfying  $n_j^{\mathrm{cot}} = n_i^{\mathrm{cot}} - 1$ ,  $d_j^{\mathrm{eq}} = d_i^{\mathrm{eq}}$ ,  $s_j$  is a variable, and  $\operatorname{var}(s_j) = n_i^{\mathrm{cot}}$ .
2. Copy the embedding  $e_{\mathrm{id}(s_{j-1})}$  for position  $j$  satisfying  $n_j^{\mathrm{cot}} = n_i^{\mathrm{cot}} - 1$ ,  $d_j^{\mathrm{eq}} = \hat{d}_i^{\mathrm{eq}} + \mathbb{I}[s_{i+2} = ',']$ ,  $e_{\mathrm{id}(s_j)} = e_i^{\mathrm{next2}}$ , and  $\operatorname{var}(s_j) = \operatorname{var}(s_{i+2})$ .
3. Copy the embeddings  $e_{\mathrm{id}(s_{j-1})}$  and  $e_j^{\mathrm{cot\_num}}$  for position  $j$  satisfying  $n_j^{\mathrm{cot}} = n_i^{\mathrm{cot}} - 1$ ,  $d_j^{\mathrm{eq}} = d_i^{\mathrm{eq}}$ ,  $s_j$  is a variable, and  $\operatorname{var}(s_j) = n_i^{\mathrm{cot}}$ .
4. Copy the embedding  $\pmb{e}_{\mathrm{id}(s_j - 1)}$  and  $\pmb{e}_j^{\mathrm{cot\_num}}$  for position  $j$  satisfying  $n_j^{\mathrm{cot}} = n_i^{\mathrm{cot}} - 1$ ,  $d_j^{\mathrm{eq}} = d_i^{\mathrm{eq}} + \mathbb{I}[s_{i + 2} = {}^{,};]$ ,  $\pmb{e}_{\mathrm{id}(s_j)} = \pmb{e}_i^{\mathrm{next2}}$ , and  $\operatorname{var}(s_j) = \operatorname{var}(s_{i + 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 desired output  $a_{d_i^{\mathrm{eq}},\mathrm{var}(s_{i + 2})} - a_{\hat{d}_i^{\mathrm{eq}},\mathrm{var}(s_{i + 2})} / a_{\hat{d}_i^{\mathrm{eq}},n_i^{\mathrm{cot}}}\cdot a_{d_i^{\mathrm{eq}},n_i^{\mathrm{cot}}}(\text{or} b_{d_i^{\mathrm{eq}}} - b_{\hat{d}_i^{\mathrm{eq}}} / a_{\hat{d}_i^{\mathrm{eq}},n_i^{\mathrm{cot}}}\cdot a_{d_i^{\mathrm{eq}},n_i^{\mathrm{cot}}})$ , which can be implemented as a look-up table (according to Lemma C.5). However, there are several special cases we have to consider:
-  $d_{i}^{\mathrm{eq}} = n_{i}^{\mathrm{cot}}$ . In this case, the coefficient is simply computed by normalizing the  $\hat{d}_{i}^{\mathrm{eq}}$ -th equation, which can also be implemented via a look-up table.
-  $d_i^{\mathrm{eq}} = \hat{d}_i^{\mathrm{eq}}$  and  $\hat{d}_i^{\mathrm{eq}} \neq n_i^{\mathrm{cot}}$ . In this case, the  $\hat{d}_i^{\mathrm{eq}}$ -th equation and the  $n_i^{\mathrm{cot}}$ -th equation are swapped according to Appendix B, and the coefficient should be instead computed by  $a_{n_i^{\mathrm{cot}},\mathrm{var}(s_{i + 2})} - a_{\hat{d}_i^{\mathrm{eq}},\mathrm{var}(s_{i + 2})} / a_{\hat{d}_i^{\mathrm{eq}},n_i^{\mathrm{cot}}} \cdot a_{n_i^{\mathrm{cot}},n_i^{\mathrm{cot}}}$ . Fortunately, the embeddings  $e_j^{\mathrm{cot\_num}}$  in the third and the fourth tasks contain exactly  $a_{n_i^{\mathrm{cot}},\mathrm{var}(s_{i + 2})}$  and  $a_{n_i^{\mathrm{cot}},n_i^{\mathrm{cot}}}$ .
Overall, the coefficient can be computed by a composition of look-up tables and (multivariate) conditional selection operations, which can be merged in a single two-layer MLP.
Now two more things remain to be done. The first is to obtain the 3-dimensional embedding  $l_{i+1}$  when  $s_{i+1}$  is a variable, while currently we have only obtained  $\operatorname{var}(s_{i+1})$ . However, we cannot compute the remaining two dimensions  $m^2 \sin \left( \frac{2 \operatorname{var}(s_{i+1}) \pi}{m} \right)$  and  $m^2 \cos \left( \frac{2 \operatorname{var}(s_{i+1}) \pi}{m} \right)$  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_j$  for any  $j$  such that  $s_j$  is the variable  $x_{\operatorname{var}(s_{i+1})}$  by using an attention head. Finally, the output is conditioned on the flag  $f_i^{\text{num}}$ : when  $f_i^{\text{num}}$  is true, this layer outputs the computed coefficient embedding; otherwise, it outputs  $e_i^{\text{next}}$  and  $l_{i+1}$ . We denote the output of this layer as  $x_i^{(4)} = (e_i^{\text{out}}, l_i^{\text{out}})$ .
Linear projection and softmax layer. Finally, we pass it through a softmax layer to predict the next token  $s_{i+1}$ . Unlike the proof of arithmetic expression, here the embedding  $l_i^{\mathrm{out}}$  is not one-hot (which contains  $\operatorname{var}(s_{i+1})$ ), so we need to additionally prove the following result: let the output logit corresponding to token  $t$  (before softmax) be  $z_t(\boldsymbol{e},\boldsymbol{l}) = \boldsymbol{w}_t \cdot [\boldsymbol{e}^\top,\boldsymbol{l}^\top]^\top + b_t$ , where  $\boldsymbol{w}_t$  and  $b_t$  are
parameters of the linear projection for logit  $t$ . Then, there exist parameters  $\{\pmb{w}_t, b_t\}_t$  such that for any two tokens  $t$  and  $\tilde{t}$  with  $t \neq \tilde{t}$
$$
\begin{array}{l} \operatorname {G a p} := z _ {t} \left(e _ {\mathrm {i d} (t)}, \operatorname {v a r} (t), m ^ {2} \sin \left(\frac {2 \operatorname {v a r} (t) \pi}{m}\right), m ^ {2} \cos \left(\frac {2 \operatorname {v a r} (t) \pi}{m}\right)\right) - \\ z _ {t} \left(\boldsymbol {e} _ {\text {i d} (\tilde {t})}, \operatorname {v a r} (\tilde {t}), m ^ {2} \sin \left(\frac {2 \operatorname {v a r} (\tilde {t}) \pi}{m}\right), m ^ {2} \cos \left(\frac {2 \operatorname {v a r} (\tilde {t}) \pi}{m}\right)\right) \geq \Theta (1). \\ \end{array}
$$
To prove the above result, simply set  $\pmb{w}_t = \left(\pmb{e}_{\mathrm{id}(t)}, \mathrm{var}(t), m^2 \sin \left(\frac{2 \mathrm{var}(t) \pi}{m}\right), m^2 \cos \left(\frac{2 \mathrm{var}(t) \pi}{m}\right)\right)$ . We have
$$
\begin{array}{l} \operatorname {G a p} = 1 + (\operatorname {v a r} (t)) ^ {2} + m ^ {4} - \mathbb {I} [ \operatorname {i d} (t) = \operatorname {i d} (\tilde {t}) ] - \operatorname {v a r} (t) \operatorname {v a r} (\tilde {t}) \\ - m ^ {4} \left(\sin \left(\frac {2 \operatorname {v a r} (t) \pi}{m}\right) \sin \left(\frac {2 \operatorname {v a r} (\tilde {t}) \pi}{m}\right) + \cos \left(\frac {2 \operatorname {v a r} (t) \pi}{m}\right) \cos \left(\frac {2 \operatorname {v a r} (\tilde {t}) \pi}{m}\right)\right) \\ = (1 - \mathbb {I} [ \mathrm {i d} (t) = \mathrm {i d} (\tilde {t}) ]) + \operatorname {v a r} (t) (\operatorname {v a r} (t) - \operatorname {v a r} (\tilde {t})) + m ^ {4} \left(1 - \cos \left(\frac {2 (\operatorname {v a r} (t) - \operatorname {v a r} (\tilde {t})) \pi}{m}\right)\right) \\ \end{array}
$$
When  $\operatorname{var}(t) = \operatorname{var}(\tilde{t})$ , we have  $\operatorname{id}(t) \neq \operatorname{id}(\tilde{t})$  and thus  $\mathrm{Gap} = 1$ . Otherwise,
$$
\begin{array}{l} \operatorname {G a p} \geq 1 - m ^ {2} + m ^ {4} (1 - \cos (2 \pi / m)) \\ = 1 - m ^ {2} + m ^ {4} \sin^ {2} (\pi / m) \geq 1, \\ \end{array}
$$
where we use the fact that  $\sin (x)\geq x / \pi$  whenever  $0 <   x\leq \pi /2$
Now it remains to conduct an error analysis and determine the scale of parameters. Similar to the proof of arithmetic expression, we can prove that all parameter values in the Transformer are bounded by  $O(\mathrm{poly}(n))$ .
E.2 Proof of Theorem 3.2
We will now prove that solving a system of linear equations without CoT is extremely difficult for bounded-depth autoregressive Transformers.
Theorem E.2. Assume  $\mathsf{TC}^0\neq \mathsf{NC}^1$ . For any prime number  $p$ , any integer  $L$ , and any polynomial  $Q$ , there exists a problem size  $m$  such that no log-precision autoregressive Transformer defined in Section 2 with depth  $L$  and hidden dimension  $d\leq Q(m)$  can solve the problem Equation  $(m,p)$ .
Proof. Our proof is based on leveraging the  $\mathsf{NC}^1$ -completeness of a classic problem: Unsolvable Automaton Membership Testing. According to Barrington's theorem [4, 5], given a fixed unsolvable automaton, judging whether the automaton accepts an input is complete in  $\mathsf{NC}^1$ . Below, we will prove that solving the system of linear equations is  $\mathsf{NC}^1$ -hard by demonstrating that the Unsolvable Automaton Membership Testing problem is  $\mathsf{NC}^0$  reducible to the problem of solving a system of linear equations. This will yield the conclusion since bounded-depth log-precision Transformers with polynomial size are in  $\mathsf{TC}^0$  [42].