SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Therefore,  $|f(a,b) - ab| < \frac{10M^3}{3\lambda}$ . Set  $\lambda \geq \frac{10M^3}{3\epsilon}$ , and then we can obtain  $|f(a,b) - ab| < \epsilon$ . Moreover, each weight element in the MLP is upper bounded by  $O(\lambda^2)$ , which is clearly  $O(\mathrm{poly}(M,1 / \epsilon))$ .
Next, we will demonstrate that a two-layer MLP with GeLU activation can efficiently approximate a two-layer MLP with ReLU activation, with all weights upper bounded by  $O(\mathrm{poly}(1 / \epsilon))$ . This result is useful in proving subsequent lemmas.
Lemma C.2. Let  $g: \mathbb{R}^{d_1} \to \mathbb{R}^{d_2}$  be a two-layer MLP with ReLU activation, and all parameter values are upper bounded by  $M$ . Then, for any  $\epsilon > 0$ , there exists a two-layer MLP  $f$  of the same size with GeLU activation and parameters upper bounded by  $O(\mathrm{poly}(M, 1/\epsilon))$  in the  $\ell_{\infty}$  norm, such that for all  $x \in \mathbb{R}^{d_1}$ , we have  $\| f(x) - g(x) \|_{\infty} \leq \epsilon$ .
Proof. Let  $\pmb{g}(\pmb{x}) = \pmb{W}_2 \cdot \mathrm{ReLU}(\pmb{W}_1\pmb{x})$ . We construct  $\pmb{f}(x) = \frac{1}{\lambda}\pmb{W}_2 \cdot \mathrm{ReLU}(\lambda \pmb{W}_1\pmb{x})$  where  $\lambda > 0$  is a sufficiently large constant. To prove that  $\| \pmb{f}(\pmb{x}) - \pmb{g}(\pmb{x}) \|_{\infty} \leq \epsilon$  for all  $\pmb{x} \in \mathbb{R}^{d_1}$ , it suffices to prove that  $\| \pmb{W}_2(\mathrm{ReLU}(\pmb{z}) - \frac{1}{\lambda}\mathrm{ReLU}(\lambda \pmb{z})) \|_{\infty} \leq \epsilon$  for all  $\pmb{z} \in \mathbb{R}^d$  where  $d$  is the hidden size. Since
$$
\begin{array}{l} \left\| \boldsymbol {W} _ {2} \left(\operatorname {R e L U} (\boldsymbol {z}) - \frac {1}{\lambda} \operatorname {G e L U} (\lambda \boldsymbol {z})\right) \right\| _ {\infty} \leq \| \boldsymbol {W} _ {2} \| _ {\infty} \left\| \operatorname {R e L U} (\boldsymbol {z}) - \frac {1}{\lambda} \operatorname {G e L U} (\lambda \boldsymbol {z}) \right\| _ {\infty} \\ \leq M d \left\| \mathrm {R e L U} (\boldsymbol {z}) - \frac {1}{\lambda} \mathrm {G e L U} (\lambda \boldsymbol {z}) \right\| _ {\infty}, \\ \end{array}
$$
it suffices to consider the scalar setting and prove that  $\left|\frac{1}{\lambda}\mathrm{GeLU}(\lambda y) - \mathrm{ReLU}(y)\right| \leq \epsilon / Md$  for all  $y \in \mathbb{R}$ . By definition of ReLU and GeLU, we have
$$
\left| \frac {1}{\lambda} \operatorname {G e L U} (\lambda y) - \operatorname {R e L U} (y) \right| = \frac {1}{\lambda} \left| \operatorname {R e L U} (\lambda y) - \int_ {- \infty} ^ {\lambda y} \frac {\lambda y}{\sqrt {2 \pi}} \exp \left(- \frac {t ^ {2}}{2}\right) \mathrm {d} t \right|. \tag {8}
$$
When  $y \geq 0$ , (8) becomes
$$
\frac {1}{\lambda} \left| \int_ {- \infty} ^ {+ \infty} \frac {\lambda y}{\sqrt {2 \pi}} \exp \left(- \frac {t ^ {2}}{2}\right) d t - \int_ {- \infty} ^ {\lambda y} \frac {\lambda y}{\sqrt {2 \pi}} \exp \left(- \frac {t ^ {2}}{2}\right) d t \right| = \frac {1}{\lambda} \int_ {\lambda y} ^ {+ \infty} \frac {\lambda y}{\sqrt {2 \pi}} \exp \left(- \frac {t ^ {2}}{2}\right) d t.
$$
Combined with the case of  $y < 0$ , (8) can be consistently written as
$$
\begin{array}{l} \left| \frac {1}{\lambda} \mathrm {G e L U} (\lambda y) - \mathrm {R e L U} (y) \right| = \frac {1}{\lambda} \int_ {\lambda | y |} ^ {+ \infty} \frac {\lambda | y |}{\sqrt {2 \pi}} \exp \left(- \frac {t ^ {2}}{2}\right) \mathrm {d} t \\ \leq \frac {1}{\sqrt {2 \pi} \lambda} \int_ {\lambda | y |} ^ {+ \infty} t \exp \left(- \frac {t ^ {2}}{2}\right) \mathrm {d} t = \frac {1}{\sqrt {2 \pi} \lambda} \exp \left(- \frac {\lambda^ {2} y ^ {2}}{2}\right) \\ \leq \frac {1}{\sqrt {2 \pi} \lambda}. \\ \end{array}
$$
Picking  $\lambda = \frac{Md}{\sqrt{2\pi\epsilon}}$  yields the desired result and completes the proof.
Equipped with the above result, we now prove that a two-layer MLP with GeLU activation can perform linear transformation and conditional selection.
Proposition C.3. Let  $\pmb{f} : \mathbb{R}^{d_1} \to \mathbb{R}^{d_2}$  be a two-layer MLP with GeLU activation, and the hidden dimension is  $2d_2$ . Let  $\pmb{W} \in \mathbb{R}^{d_2 \times d_1}$  be any matrix and denote  $M = \max_{ij} |W_{ij}|$ . Then, for any  $\epsilon > 0$ , there exist MLP parameters with  $\ell_{\infty}$  norm bounded by  $O(\mathrm{poly}(M, 1/\epsilon))$ , such that for any  $\pmb{x} \in \mathbb{R}^{d_1}$ , we have  $\| \pmb{f}(\pmb{x}) - \pmb{W}\pmb{x} \|_{\infty} \leq \epsilon$ .
Proof. We can use a two-layer MLP with ReLU activation to implement  $g(x) = Wx$  by the following construction:
$$
\boldsymbol {W} \boldsymbol {x} = \operatorname {R e L U} (\boldsymbol {W} \boldsymbol {x}) + \operatorname {R e L U} (- \boldsymbol {W} \boldsymbol {x})
$$
Combined with Lemma C.2, we can also implement  $g(x)$  by a two-layer MLP with GeLU activation.
Lemma C.4. Define the selection function  $g: \mathbb{R}^d \times \mathbb{R}^d \times \mathbb{R} \to \mathbb{R}^d$  as follows:
$$
\boldsymbol {g} (\boldsymbol {x}, \boldsymbol {y}, t) = \left\{ \begin{array}{l l} \boldsymbol {x} & i f t \geq 0, \\ \boldsymbol {y} & i f t <   0. \end{array} \right. \tag {9}
$$
Let  $\pmb{f}:\mathbb{R}^d\times \mathbb{R}^d\times \mathbb{R}\to \mathbb{R}^d$  be a two-layer MLP with GeLU activation, and the hidden dimension is  $2d + 2$ . Then, for any  $\epsilon >0$ ,  $\alpha >0$ , and  $M > 0$ , there exist MLP parameters with  $\ell_{\infty}$  norm bounded by  $O(\mathrm{poly}(M,1 / \alpha ,1 / \epsilon))$ , such that for all  $\pmb {x}\in [-M,M]^d$ ,  $\pmb {y}\in [-M,M]^d$ , and  $t\in [-\infty , - \alpha ]\cup [\alpha , + \infty ]$ , we have  $\| f(\pmb {x},\pmb {y},t) - g(\pmb {x},\pmb {y},t)\|_{\infty}\leq \epsilon$
Proof. We can simply use a two-layer MLP with ReLU activation to implement  $g$  by the following construction:
$$
\boldsymbol {h} (\boldsymbol {x}, \boldsymbol {y}, t) = \left(\boldsymbol {h} _ {1}, \boldsymbol {h} _ {2}, h _ {3}, h _ {4}\right) := \left(\boldsymbol {x} + \alpha^ {- 1} M t \mathbf {1} _ {d}, \boldsymbol {y} - \alpha^ {- 1} M t \mathbf {1} _ {d}, \alpha^ {- 1} M t, - \alpha^ {- 1} M t\right) \in \mathbb {R} ^ {2 d + 2},
$$
$$
\boldsymbol {f} (\boldsymbol {x}, \boldsymbol {y}, t) = \operatorname {R e L U} \left(\boldsymbol {h} _ {1}\right) - \operatorname {R e L U} \left(h _ {3}\right) \mathbf {1} _ {d} + \operatorname {R e L U} \left(\boldsymbol {h} _ {2}\right) - \operatorname {R e L U} \left(h _ {4}\right) \mathbf {1} _ {d},
$$
where  $\mathbf{1}_d$  is the all-one vector of  $d$  dimension. It is easy to check that, for all  $\pmb{x} \in [-M, M]^{d_1}$ ,  $\pmb{y} \in [-M, M]^{d_2}$ , and  $t \in [-\infty, -\alpha] \cup [\alpha, +\infty]$ , we have  $\pmb{f}(\pmb{x}, \pmb{y}, t) = \pmb{g}(\pmb{x}, \pmb{y}, t)$ . Moreover, all parameters are bounded by  $O(M / \alpha)$ . Therefore, by using Lemma C.2, we can also implement  $\pmb{g}(\pmb{x})$  by a two-layer MLP with GeLU activation and all parameters are bounded by  $O(\mathrm{poly}(M, 1 / \alpha, 1 / \epsilon))$ .
We final show that a two-layer MLP can efficiently represent a look-up table. Consider a  $k$ -dimensional table of size  $d^k$ , where each element in the table is an integer ranging from 1 to  $d$ . Denote the set  $\mathcal{D} = \{e_i : i \in [d]\}$ , where  $e_i$  is a  $d$ -dimensional one-hot vector with the  $i$ -th element being 1. The above look-up table can thus be represented as a discrete function  $g : \mathcal{D}^k \to \mathcal{D}$ . The following lemma shows that  $g$  can be implemented by a two-layer MLP with GeLU activation.
Lemma C.5. Let  $g: \mathcal{D}^k \to \mathcal{D}$  be any function defined above, and let  $f: \mathbb{R}^{k \times d} \to \mathbb{R}^d$  be a two-layer MLP with GeLU activation and bias, and the hidden dimension is  $d^k$ . Then, for any  $\epsilon > 0$ , there exist MLP parameters with  $\ell_{\infty}$  norm bounded by  $O(\mathrm{poly}(k,1 / \epsilon))$ , such that for all  $x \in \mathcal{D}^k \subset \mathbb{R}^{k \times d}$  and all perturbation  $\delta \in [-1/2k,1/2k]^{k \times d}$ , we have  $\| f(x + \delta) - g(x) \|_{\infty} \leq \epsilon + 2k \| \delta \|_{\infty}$ , where  $\| \delta \|_{\infty}$  is the vector  $\ell_{\infty}$ -norm applied to the flattened matrix  $\delta$ .
Proof. We can simply use a two-layer MLP with ReLU activation to implement  $\pmb{g}$  by the following construction. Denote the index of the MLP hidden layer as  $(i_1, \dots, i_k) \in [d]^k$ . We can construct the weights of the first MLP layer such that
$$
h _ {\left(i _ {1}, \dots , i _ {k}\right)} (\boldsymbol {x}) = 2 \left(x _ {i _ {1}} + x _ {d + i _ {2}} \dots + x _ {(k - 1) d + i _ {k}}\right) - 2 k + 1.
$$
We can then construct the weights of the second layer such that the final output of the MLP is
$$
f_{j}(\boldsymbol {x}) = \sum_{g_{j}(e_{i_{1}},\dots ,e_{i_{k}}) = 1}\mathrm{ReLU}(h_{(i_{1},\dots ,i_{k})}(\boldsymbol {x})).
$$
One can check that  $\pmb{f}(\pmb{x}) = \pmb{g}(\pmb{x})$  holds for all  $\pmb{x} \in \mathcal{D}^k \subset \mathbb{R}^{d \times k}$ . Furthermore, we have
$$
\begin{array}{l} \| \boldsymbol {f} (\boldsymbol {x} + \delta) - \boldsymbol {g} (\boldsymbol {x}) \| _ {\infty} = \| \boldsymbol {f} (\boldsymbol {x} + \delta) - \boldsymbol {f} (\boldsymbol {x}) \| _ {\infty} \\ = \max  _ {j \in [ d ]} \left| \sum_ {g _ {j} \left(\boldsymbol {e} _ {i _ {1}}, \dots , \boldsymbol {e} _ {i _ {k}}\right) = 1} \left(\operatorname {R e L U} \left(h _ {\left(i _ {1}, \dots , i _ {k}\right)} (\boldsymbol {x} + \boldsymbol {\delta})\right) - \operatorname {R e L U} \left(h _ {\left(i _ {1}, \dots , i _ {k}\right)} (\boldsymbol {x})\right) \right. \right| \\ \leq \max  _ {(i _ {1}, \dots , i _ {k}) \in [ d ] ^ {k}} \left| h _ {(i _ {1}, \dots , i _ {k})} (\boldsymbol {x} + \boldsymbol {\delta}) - h _ {(i _ {1}, \dots , i _ {k})} (\boldsymbol {x}) \right| \leq 2 k \| \boldsymbol {\delta} \| _ {\infty} \\ \end{array}
$$
for all perturbations  $\delta \in [-1/2k, 1/2k]^{k \times d}$ . Thus by using Lemma C.2, we can also implement  $\pmb{g}(\pmb{x})$  by a two-layer MLP with GeLU activation and all parameters are bounded by  $O(\mathrm{poly}(k, 1/\epsilon))$ .
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 \in \mathbb{N}$  be an integer and let  $\boldsymbol{x}_1, \boldsymbol{x}_2, \dots, \boldsymbol{x}_n$  be a sequence of vectors where  $\boldsymbol{x}_i = (\tilde{\boldsymbol{x}}_i, r_i, 1) \in [-M, M]^{d+2}$ ,  $\tilde{\boldsymbol{x}}_i \in \mathbb{R}^d$ ,  $r_i \in \mathbb{R}$ , and  $M$  is a large constant. Let  $K, Q, V \in \mathbb{R}^{d' \times (d+2)}$  be any matrices with  $\| V \|_{\infty} \leq 1$ , and let  $0 < \rho, \delta < M$  be any real numbers. Denote  $q_i = Q x_i$ ,  $k_j = K x_j$ ,  $v_j = V x_j$ , and define the matching set  $S_i = \{j \leq i : |q_i \cdot k_j| \leq \rho\}$ . Equipped with these notations, we define two basic operations as follows:
- COPY: The output is a sequence of vectors  $\mathbf{u}_1, \dots, \mathbf{u}_n$  with  $\mathbf{u}_i = \mathbf{v}_{\mathrm{pos}(i)}$ , where  $\mathrm{pos}(i) = \operatorname*{argmax}_{j \in S_i} r_j$ .
- MEAN: The output is a sequence of vectors  $\mathbf{u}_1, \dots, \mathbf{u}_n$  with  $\mathbf{u}_i = \mathrm{mean}_{j \in S_i} \mathbf{v}_j$ .
The output  $\pmb{u}_i$  is undefined when  $S_i = \emptyset$ . We next make the following regularity assumption:
Assumption C.6. The matrices  $Q, K, V$  and scalars  $\rho, \delta$  satisfy that for all considered sequences  $x_1, x_2, \dots, x_n$ , the following hold:
- For any  $i, j \in [n]$ , either  $|\pmb{q}_i \cdot \pmb{k}_j| \leq \rho$  or  $\pmb{q}_i \cdot \pmb{k}_j \leq -\delta$ .
- For any  $i, j \in [n]$ , either  $i = j$  or  $|r_i - r_j| \geq \delta$ .
Assumption C.6 says that there are sufficient gaps between the attended position (e.g.,  $\mathrm{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  $\rho \leq \frac{\delta^2}{8M}$ . For any  $\epsilon > 0$ , there exists an attention layer with embedding size  $O(d)$  and one causal attention head that can approximate the COPY operation defined above. Formally, for any considered sequence of vectors  $\mathbf{x}_1, \mathbf{x}_2, \ldots, \mathbf{x}_n$ , denote the corresponding attention output as  $\mathbf{o}_1, \mathbf{o}_2, \ldots, \mathbf{o}_n$ . Then, we have  $\| \mathbf{o}_i - \mathbf{u}_i \|_{\infty} \leq \epsilon$  for all  $i \in [n]$  with  $S_i \neq \emptyset$ . Moreover, the  $\ell_{\infty}$  norm of attention parameters is bounded by  $O(\mathrm{poly}(M, 1/\delta, \log(n), \log(1/\epsilon)))$ .
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:  $(\lambda q_i, \mu) \in \mathbb{R}^{d+1}$
Key:  $(\pmb {k}_i,r_i)\in \mathbb{R}^{d + 1}$
- Value:  $\mathbf{v}_i \in \mathbb{R}^d$
where  $\lambda$  and  $\mu$  are constants that will be defined later. Denote  $a_{ij}$  as the attention score, then
$$
a _ {i, j} = \frac {\exp (\lambda (\boldsymbol {q} _ {i} \cdot \boldsymbol {k} _ {j}) + \mu r _ {j})}{\sum_ {j ^ {\prime}} \exp (\lambda (\boldsymbol {q} _ {i} \cdot \boldsymbol {k} _ {j ^ {\prime}}) + \mu r _ {j ^ {\prime}})} = \frac {\exp (\lambda (\boldsymbol {q} _ {i} \cdot \boldsymbol {k} _ {j}))}{\sum_ {j ^ {\prime}} \exp (\lambda (\boldsymbol {q} _ {i} \cdot \boldsymbol {k} _ {j ^ {\prime}}) + \mu (r _ {j ^ {\prime}} - r _ {j}))}.
$$
Since  $\rho \leq \frac{\delta^2}{8M}$  and  $M \geq \delta$ , we have  $\delta - \rho \geq \frac{7}{8} \delta$ . By setting  $\lambda = \frac{8M \ln \left( \frac{2nM}{\epsilon} \right)}{\delta^2}$  and  $\mu = \frac{3 \ln \left( \frac{2nM}{\epsilon} \right)}{\delta}$  (which are bounded by  $O(\mathrm{poly}(M, 1 / \delta, \log(n), \log(1 / \epsilon))))$ ), we have
$$
\begin{array}{l} a _ {i, \operatorname {p o s} (i)} \geq \frac {\exp (- \lambda \rho)}{\exp (- \lambda \rho) + (n - 1) \exp (\max (- \lambda \delta + 2 M \mu , \lambda \rho - \mu \delta))} (10) \\ = \frac {1}{1 + (n - 1) \exp (\max (- \lambda (\delta - \rho) + 2 M \mu , 2 \lambda \rho - \mu \delta))} \\ \geq 1 - n \exp \left(\max  (- \lambda (\delta - \rho) + 2 M \mu , 2 \lambda \rho - \mu \delta)\right) (11) \\ \geq 1 - n \exp \left(\max  \left(- \frac {M}{\delta} \ln \left(\frac {2 n M}{\epsilon}\right), - \ln \left(\frac {2 n M}{\epsilon}\right)\right)\right) \\ \geq 1 - n \exp \left(- \ln \left(\frac {2 n M}{\epsilon}\right)\right) (12) \\ = 1 - \frac {\epsilon}{2 M}, \\ \end{array}
$$
where in (10) we use Assumption C.6, which implies that whenever  $j' \neq \mathrm{pos}(i)$ , either  $\pmb{q}_i \cdot \pmb{k}_{j'} \leq -\delta$  or  $(\pmb{q}_i \cdot \pmb{k}_{j'} \leq \rho$  and  $r_{j'} - r_j \leq -\delta)$ ; in (11) we use the inequality  $\frac{1}{1 + x} \geq 1 - x$  for all  $x \geq 0$ ; in (12) we use the fact that  $M \geq \delta$ . We thus have
$$
\begin{array}{l} \left\| \boldsymbol {o} _ {i} - \boldsymbol {u} _ {i} \right\| _ {\infty} = \left\| \sum_ {j} a _ {i j} \boldsymbol {v} _ {j} - \boldsymbol {v} _ {\operatorname {p o s} (i)} \right\| _ {\infty} \leq M \| \boldsymbol {V} \| _ {\infty} \cdot \left(1 - a _ {i, \operatorname {p o s} (i)} + \sum_ {j \neq \operatorname {p o s} (i)} a _ {i, j}\right) \\ = M \| \boldsymbol {V} \| _ {\infty} (2 - 2 a _ {i, \mathrm {p o s} (i)}) \leq \epsilon , \\ \end{array}
$$
Lemma C.8. Assume Assumption C.6 holds with  $\rho \leq \frac{\delta\epsilon}{16M\ln\left(\frac{4Mn}{\epsilon}\right)}$ . For any  $0 < \epsilon \leq M$ , there exists an attention layer with embedding size  $O(d)$  and one causal attention head that can approximate the MEAN operation defined above. Formally, for any considered sequence of vectors  $\mathbf{x}_1,\mathbf{x}_2,\ldots ,\mathbf{x}_n$  denote the attention output as  $\pmb{o}_1,\pmb{o}_2,\dots ,\pmb{o}_n$ . Then, we have  $\| \pmb {o}_i - \pmb {u}_i\|_{\infty}\leq \epsilon$  for all  $i\in [n]$  with  $S_{i}\neq \emptyset$ . Moreover, the  $\ell_{\infty}$  norm of attention parameters is bounded by  $O(\mathrm{poly}(M,1 / \delta ,\log (n),\log (1 / \epsilon)))$ .
Proof. The purpose of the attention head is to average across all tokens that satisfy the condition  $\pmb{q}_i \cdot \pmb{k}_j \approx 0$ . To achieve this, we construct the key, query, and value vectors as follows:
- Query:  $\lambda \pmb{q}_i\in \mathbb{R}^d$
Key:  $\pmb{k}_i\in \mathbb{R}^d$
- Value:  $\mathbf{v}_i \in \mathbb{R}^d$
where  $\lambda$  is a constant which will be defined later. Denote  $a_{ij}$  as the attention score, then
$$
a _ {i, j} = \frac {\exp (\lambda \boldsymbol {q} _ {i} \cdot \boldsymbol {k} _ {j})}{\sum_ {j ^ {\prime}} \exp (\lambda \boldsymbol {q} _ {i} \cdot \boldsymbol {k} _ {j ^ {\prime}})}.
$$
By setting  $\lambda = \frac{1}{\delta}\ln \left(\frac{4Mn}{\epsilon}\right)$  (which is bounded by  $O(\mathrm{poly}(M,1 / \delta ,\log (n),\log (1 / \epsilon))))$  we have:
$$
\begin{array}{l} \sum_ {j \notin \mathcal {S} _ {i}} a _ {i j} \leq \frac {(n - | \mathcal {S} _ {i} |) \exp (- \lambda \delta)}{(n - | \mathcal {S} _ {i} |) \exp (- \lambda \delta) + | \mathcal {S} _ {i} | \exp (- \lambda \rho)} (13) \\ = \frac {1}{1 + \frac {| \mathcal {S} _ {i} |}{n - | \mathcal {S} _ {i} |} \exp (- \lambda (\rho - \delta))} \\ <   n \exp (\lambda \rho) \exp (- \lambda \delta)) (14) \\ \leq n \exp \left(\frac {\epsilon}{1 6 M}\right) \exp \left(- \ln \left(\frac {4 M n}{\epsilon}\right)\right) \\ <   \frac {\epsilon}{3 M}, (15) \\ \end{array}
$$
where in (13) we use Assumption C.6, which implies that  $\pmb{q}_i \cdot \pmb{k}_j \leq -\delta$  for all  $j \notin S_i$  and  $\pmb{q}_i \cdot \pmb{k}_j \geq -\rho$  for all  $j \in S_i$ ; in (14) we use the inequality  $\frac{1}{1 + x} < \frac{1}{x}$  for all  $x > 0$ ; in (15) we use that the assumption that  $\epsilon \leq M$  and the fact that  $\exp(1/16) < 4/3$ .
Similarly, for any  $j\in \mathcal{S}_i$  , we have
$$
\begin{array}{l} \left| a _ {i j} - \frac {1}{| \mathcal {S} _ {i} |} \right| \leq \max  \left(\frac {1}{| \mathcal {S} _ {i} |} - \frac {\exp (- \rho \lambda)}{| \mathcal {S} _ {i} | \exp (\rho \lambda) + (n - | \mathcal {S} _ {i} |) \exp (- \lambda \delta)}, \frac {\exp (\rho \lambda)}{| \mathcal {S} _ {i} | \exp (- \rho \lambda)} - \frac {1}{| \mathcal {S} _ {i} |}\right) (16) \\ \leq \frac {1}{| \mathcal {S} _ {i} |} \max \left(1 - \frac {1}{\exp (2 \rho \lambda) + n \exp (- \lambda (\delta - \rho))}, \exp (2 \rho \lambda) - 1\right) \\ \leq \frac {1}{| \mathcal {S} _ {i} |} \max  \left(\exp (2 \rho \lambda) - 1 + n \exp (- \lambda (\delta - \rho)), \exp (2 \rho \lambda) - 1\right) (17) \\ = \frac {1}{\left| \mathcal {S} _ {i} \right|} \left(\exp (2 \rho \lambda) - 1 + n \exp (- \lambda (\delta - \rho))\right) \\ \leq \frac {1}{| \mathcal {S} _ {i} |} \left(\exp \left(\frac {\epsilon}{8 M}\right) - 1 + \frac {\epsilon}{3 M}\right) (18) \\ \leq \frac {2 \epsilon}{3 M | \mathcal {S} _ {i} |} (19) \\ \end{array}
$$
where in (16) we use Assumption C.6 similarly as before; in (17) we use the inequality  $1 - \frac{1}{x} \leq x - 1$  for all  $x > 0$ ; in (18) we use the inequality previously derived in (15); in (19) we use the inequality  $\exp(x) \leq 1 + 2x$  for all  $0 \leq x \leq 1$ . We thus obtain
$$
\| \boldsymbol {o} _ {i} - \boldsymbol {u} _ {i} \| _ {\infty} = \left\| \sum_ {j} a _ {i j} \boldsymbol {v} _ {j} - \frac {1}{| \mathcal {S} _ {i} |} \sum_ {j \in \mathcal {S} _ {i}} \boldsymbol {v} _ {j} \right\| _ {\infty} \leq M \| \boldsymbol {V} \| _ {\infty} \cdot \left(\sum_ {j \notin \mathcal {S} _ {i}} a _ {i j} + \sum_ {j \in \mathcal {S} _ {i}} \left| a _ {i j} - \frac {1}{| \mathcal {S} _ {i} |} \right|\right) \leq \epsilon ,
$$
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/MLP 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  $\pmb{f}:\mathbb{R}^d\to \mathbb{R}^d$ , and let  $\pmb {y} = \pmb {f}(\pmb {x})$ . It is easy to see that we can construct another MLP (or attention layer) denoted as  $g:\mathbb{R}^{2d}\rightarrow \mathbb{R}^{2d}$  such that  $g(\pmb {x},\pmb {0}) + (\pmb {x},\pmb {0}) = (\pmb {0},\pmb {y}) + (\pmb {x},\pmb {0}) = (\pmb {x},\pmb {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 proofs presented in Appendices D, E and G. 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  $W_{O}^{(l,h)}$  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.1. For any prime  $p$  and integer  $n > 0$ , there exists an autoregressive Transformer defined in Section 2 with hidden size  $d = O(\mathrm{poly}(p))$  (independent of  $n$ ), depth  $L = 5$ , and 5 heads in each layer that can generate the CoT solution defined in Appendix B for all inputs in Arithmetic  $(n, p)$ . Moreover, all parameter values in the Transformer are bounded by  $O(\mathrm{poly}(n))$ .
Proof sketch. The intuition behind our construction is that when the CoT output proceeds to a certain position, the Transformer can read the context related to this position and determine whether it should copy a token or perform a 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 a 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 arithmetic rules  $(+, -, \times, \div)$ .
Proof. We construct each layer as follows.
Token Embeddings. Assume that we have a sequence of tokens  $s_1, \ldots, s_i$  and we want to generate the next token  $s_{i+1}$ . For any  $j \in [n]$ , let  $\mathrm{id}(s_j)$  be the index of token  $s_j$  in the embedding dictionary, with values ranging from 1 to the number of tokens. We can embed the token  $s_j$  by  $x_j^{(0)} = (e_{\mathrm{id}(s_j)}, j, 1) \in \mathbb{R}^{\mathrm{num\_tokens} + 2}$ , where  $e_j$  is a one-hot vector with the  $j$ -th element being 1,  $j \in \mathbb{N}_+$  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  $(\mathbf{\Phi} = \mathbf{\Phi})$  in previous tokens, denoted as  $n_i^{\overline{=}}$ , i.e.,  $n_i^{\overline{=}} = |\{j \leq i : s_j = \mathbf{\Phi}' = \mathbf{\Phi}\}|.$
2. Copy the position of the last equal sign, denoted as  $p_i^{\equiv}$ , i.e.,  $p_i^{\equiv} = \max \{j : j \leq i, s_j = ' = '\}$ . If the set  $\{j : j \leq i, s_j = ' = '\}$  is empty, define  $p_i^{\equiv} = 0$ .
3. Compute  $i^2$ .