SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\left[ \begin{array}{c c c c c c c c c c c} \boldsymbol {x} _ {1} & \dots & \boldsymbol {s} _ {1} ^ {\ell - 1} & \boldsymbol {s} _ {1} ^ {\ell} & \dots & \boldsymbol {s} _ {2} ^ {\ell - 1} & \boldsymbol {s} _ {2} ^ {\ell} & \dots & \boldsymbol {s} _ {n} ^ {L} & \boldsymbol {x} _ {\text {t e s t}} & \dots & \hat {\boldsymbol {s}} ^ {\ell - 1} \\ \boldsymbol {0} & \dots & \boldsymbol {0} & \boldsymbol {0} & \dots & \boldsymbol {0} & \boldsymbol {0} & \dots & \boldsymbol {0} & \boldsymbol {0} & \dots & \boldsymbol {0} \\ \hline 1 & \dots & \ell & \ell + 1 & \dots & \ell & \ell + 1 & \dots & L + 1 & 1 & \dots & \ell \\ 1 & \dots & \ell & \ell + 1 & \dots & \ell & \ell + 1 & \dots & L + 1 & 1 & \dots & \ell \\ \hline \mathrm {N} & \dots & \mathrm {N} & \mathrm {N} & \dots & \mathrm {N} & \mathrm {N} & \dots & \mathrm {N} & \mathrm {N} & \dots & \mathrm {N} \\ \hline \ln (1) & \dots & \ln (\ell) & \ln (\ell + 1) & \dots & \ln (L + \ell + 1) & \ln (L + \ell + 2) & \dots & \ln (n (L + 1)) & \ln (n (L + 1) + 1) & \dots & \ln (\mathrm {N}) \\ 0 & \dots & 0 & 0 & \dots & 0 & 0 & \dots & 0 & 1 & \dots & 1 \\ 1 & \dots & 1 & 1 & \dots & 1 & 1 & \dots & 1 & 1 & \dots & 1 \\ \boldsymbol {m} _ {1} & \dots & \boldsymbol {m} _ {\ell} & \boldsymbol {m} _ {\ell + 1} & \dots & \boldsymbol {m} _ {L + \ell + 1} & \boldsymbol {m} _ {L + \ell + 2} & \dots & \boldsymbol {m} _ {n (L + 1)} & \boldsymbol {m} _ {n (L + 1) + 1} & \dots & \boldsymbol {m} _ {\mathrm {N}} \\ \end{array} \right]
$$
Notice that this step does not involve any error.
Step 2: Extract the identifier of the prediction to be made  $\ell$  (3 layers). Now, by setting the key and query weight matrices of the attention as to just keep the all ones row, we get a matrix that attends
equally to all tokens. Then the value weight matrix keeps the row that contains  $\ell$  ones and thus we get the number  $\frac{\ell}{\mathbf{N}}$ , propagated in all the sequence length. Then the attention output is as follows:
$$
\left[ \begin{array}{c c c c c c c c c c c} \boldsymbol {x} _ {1} & \dots & s _ {1} ^ {\ell - 1} & s _ {1} ^ {\ell} & \dots & s _ {2} ^ {\ell - 1} & s _ {2} ^ {\ell} & \dots & s _ {n} ^ {L} & \boldsymbol {x} _ {\text {t e s t}} & \dots & \hat {s} ^ {\ell - 1} \\ \boldsymbol {0} & \dots & \boldsymbol {0} & \boldsymbol {0} & \dots & \boldsymbol {0} & \boldsymbol {0} & \dots & \boldsymbol {0} & \boldsymbol {0} & \dots & \boldsymbol {0} \\ \hline 1 & \dots & \ell & \ell + 1 & \dots & \ell & \ell + 1 & \dots & L + 1 & 1 & \dots & \ell \\ 1 & \dots & \ell & \ell + 1 & \dots & \ell & \ell + 1 & \dots & L + 1 & 1 & \dots & \ell \\ \hline \mathrm {N} & \dots & \mathrm {N} & \mathrm {N} & \dots & \mathrm {N} & \mathrm {N} & \dots & \mathrm {N} & \mathrm {N} & \dots & \mathrm {N} \\ \frac {\ell}{\mathrm {N}} & \dots & \frac {\ell}{\mathrm {N}} & \frac {\ell}{\mathrm {N}} & \dots & \frac {\ell}{\mathrm {N}} & \frac {\ell}{\mathrm {N}} & \dots & \frac {\ell}{\mathrm {N}} & \frac {\ell}{\mathrm {N}} & \dots & \frac {\ell}{\mathrm {N}} \\ 0 & \dots & 0 & 0 & \dots & 0 & 0 & \dots & 0 & 1 & \dots & 1 \\ 1 & \dots & 1 & 1 & \dots & 1 & 1 & \dots & 1 & 1 & \dots & 1 \\ m _ {1} & \dots & m _ {\ell} & m _ {\ell + 1} & \dots & m _ {L + \ell + 1} & m _ {L + \ell + 2} & \dots & m _ {n (L + 1)} & m _ {n (L + 1) + 1} & \dots & m _ {\mathrm {N}} \end{array} \right].
$$
So far this step does not involve any error. As described in Lemma 3 to implement multiplication of two values up to any error  $\epsilon$ , we need 1) have the two numbers to be multiplied next to each other and an extra structure (some constants, which we consider that are encoded in the last rows of the matrix  $m_{i}$ ) and 2) the corresponding structure presented in Lemma 3. So we need to make the following transformation in the rows containing N and  $\frac{\ell}{N}$
$$
\begin{array}{l} \left[\begin{array}{c c c c}\mathrm {N}&\mathrm {N}&\dots&\mathrm {N}\\\frac {\ell}{\mathrm {N}}&\frac {\ell}{\mathrm {N}}&\dots&\frac {\ell}{\mathrm {N}}\end{array}\right]\rightarrow \left[\begin{array}{c c c c}\mathrm {N}&\frac {\ell}{\mathrm {N}}&\dots&\mathrm {N}\\\frac {\ell}{\mathrm {N}}&\frac {\ell}{\mathrm {N}}&\dots&\frac {\ell}{\mathrm {N}}\end{array}\right] (A.19) \\ \rightarrow \left[\begin{array}{c c c c}\ell&*&\dots&*\\*&*&\dots&*\end{array}\right] (A.20) \\ \rightarrow \left[\begin{array}{c c c c}\ell&\ell&\dots&\ell\\*&*&\dots&*\end{array}\right] (A.21) \\ \end{array}
$$
where  $*$  denotes inconsequential values. For the first step we assume that we have the necessary binary representations  $^5$  and then we use Lemma 1 which shows how we can perform the operation of copy/paste. This step involves an error and the current output would be  $\tilde{X}_1 = X^* +\varepsilon_1M_1$  where  $\varepsilon_{1}$  is controlled and we will determine it in the sequence,  $\| M_1\| \leq 1$  and  $X_{1}^{*}$  is the desired output. For the second step now we use Lemma 3 to perform the multiplication, this will affect some of the other values. To analyze the error of this step, we know that given  $X_{1}^{*}$  we will get the desired output  $X_{2}^{*} + \varepsilon_{2}M_{2}$ , where  $\| M_2\| \leq 1$  and  $\varepsilon_{2}$  is to be determined. However, are given  $\tilde{X}_{1}$  in the multiplication procedure which results in the following output  $\tilde{X}_2 = (\tilde{X}_1)^\top \tilde{X}_1 + \varepsilon_2M_2 =$ $X_{1}^{*\top}X_{1}^{*} + \varepsilon_{1}M_{1}^{\top}X_{1}^{*} + \varepsilon_{1}M_{1}^{\top}M_{1} + \varepsilon_{2}M_{2}$ . Thus by choosing  $\varepsilon_{1},\varepsilon_{2}$  to be of order  $\mathcal{O}(\delta)$  we get that  $\tilde{X}_2 = X_2^* +\frac{\delta}{C}\| M\|$  for some  $C$  chosen to be large enough. Then for the last step consider the follow (sub-)rows of the matrix  $X$
$$
\left[ \begin{array}{c c c c c} \ell & * & * & \dots & * \\ \mathbf {r} _ {1} & \mathbf {r} _ {1} & \mathbf {r} _ {1} & \dots & \mathbf {r} _ {1} \\ \mathbf {r} _ {1} & \mathbf {r} _ {2} & \mathbf {r} _ {3} & \dots & \mathbf {r} _ {\mathrm {N}} \end{array} \right] \tag {A.22}
$$
where  $\mathbf{r}_i$  is the binary representation of position  $i$ . By choosing  $W_{\mathbf{K}}, W_{\mathbf{Q}}$  as to
$$
\boldsymbol {W} _ {\mathbf {K}} \boldsymbol {X} = \left[ \begin{array}{l l l l} \mathbf {r} _ {1} & \mathbf {r} _ {2} & \dots & \mathbf {r} _ {\mathrm {N}} \end{array} \right], \quad \boldsymbol {W} _ {\mathbf {Q}} \boldsymbol {X} = \left[ \begin{array}{l l l l} \mathbf {r} _ {1} & \mathbf {r} _ {1} & \dots & \mathbf {r} _ {1} \end{array} \right] \tag {A.23}
$$
the size that we need will be  $2\log \mathrm{N}_{\mathrm{max}} + 1$ , where  $\mathrm{N}_{\mathrm{max}}$  is the maximum sequence length
Note that the true error is even smaller since the operation performed only affects one row.
and consider  $\pmb{W}_{\mathbf{V}}\pmb {X} = \left[ \begin{array}{cccc}\ell & * & \dots & *\\ \mathbf{0} & \mathbf{0} & \dots & \mathbf{0} \end{array} \right]$  we have that
$$
\begin{array}{l} \operatorname {a t t n} (X) = X + W _ {\mathrm {V}} X \operatorname {s o f t m a x} \left(\left(W _ {\mathrm {K}} X\right) ^ {\top} W _ {\mathrm {Q}} X\right) \\ = \boldsymbol {X} + \left[ \begin{array}{c c c c} \ell & * & \dots & * \\ \mathbf {0} & \mathbf {0} & \dots & \mathbf {0} \end{array} \right] s o f t m a x \left(\left[ \begin{array}{c c c c} \mathbf {r} _ {1} ^ {\top} \mathbf {r} _ {1} & \mathbf {r} _ {1} ^ {\top} \mathbf {r} _ {1} & \dots & \mathbf {r} _ {1} ^ {\top} \mathbf {r} _ {1} \\ \mathbf {r} _ {1} ^ {\top} \mathbf {r} _ {2} & \mathbf {r} _ {1} ^ {\top} \mathbf {r} _ {2} & \dots & \mathbf {r} _ {1} ^ {\top} \mathbf {r} _ {2} \\ \vdots & \vdots & \ddots & \vdots \\ \mathbf {r} _ {1} ^ {\top} \mathbf {r} _ {\mathrm {N}} & \mathbf {r} _ {1} ^ {\top} \mathbf {r} _ {\mathrm {N}} & \dots & \mathbf {r} _ {1} ^ {\top} \mathbf {r} _ {\mathrm {N}} \end{array} \right]\right) \\ = \boldsymbol {X} + \left[ \begin{array}{c c c c} \ell & * & \dots & * \\ \mathbf {0} & \mathbf {0} & \dots & \mathbf {0} \end{array} \right] \left[ \begin{array}{c c c c} 1 & 1 & \dots & 1 \\ 0 & 0 & \dots & 0 \\ \vdots & \vdots & \dots & \vdots \\ 0 & 0 & \dots & 0 \end{array} \right] + \varepsilon_ {3} \mathbf {M} _ {3} \\ \end{array}
$$
By subtracting one identity head for the first that we focus on as described in Lemma 1 we have that  $\operatorname{attn}(\mathbf{X})$  results in the desired matrix. We output this result and overwrite  $\ell / \mathrm{N}$ , thus we have
$$
\left[ \begin{array}{c c c c c c c c c c c} \boldsymbol {x} _ {1} & \dots & s _ {1} ^ {\ell - 1} & s _ {1} ^ {\ell} & \dots & s _ {2} ^ {\ell - 1} & s _ {2} ^ {\ell} & \dots & s _ {n} ^ {L} & \boldsymbol {x} _ {\text {t e s t}} & \dots & \hat {s} ^ {\ell - 1} \\ \boldsymbol {0} & \dots & \boldsymbol {0} & \boldsymbol {0} & \dots & \boldsymbol {0} & \boldsymbol {0} & \dots & \boldsymbol {0} & \boldsymbol {0} & \dots & \boldsymbol {0} \\ \hline 1 & \dots & \ell & \ell + 1 & \dots & \ell & \ell + 1 & \dots & L + 1 & 1 & \dots & \ell \\ 1 & \dots & \ell & \ell + 1 & \dots & \ell & \ell + 1 & \dots & L + 1 & 1 & \dots & \ell \\ \hline \mathrm {N} & \dots & \mathrm {N} & \mathrm {N} & \dots & \mathrm {N} & \mathrm {N} & \dots & \mathrm {N} & \mathrm {N} & \dots & \mathrm {N} \\ \ell & \dots & \ell & \ell & \dots & \ell & \ell & \dots & \ell & \ell & \dots & \ell \\ 0 & \dots & 0 & 0 & \dots & 0 & 0 & \dots & 0 & 1 & \dots & 1 \\ 1 & \dots & 1 & 1 & \dots & 1 & 1 & \dots & 1 & 1 & \dots & 1 \\ m _ {1} & \dots & m _ {\ell} & m _ {\ell + 1} & \dots & m _ {L + \ell + 1} & m _ {L + \ell + 2} & \dots & m _ {n (L + 1)} & m _ {n (L + 1) + 1} & \dots & m _ {\mathrm {N}} \end{array} \right] \tag {A.24}
$$
The new error introduced by this operation is again controlled as to be  $\varepsilon_3M_3 = \frac{\delta}{C}$ , with  $\| M_3\| \leq 1$ .
Step 3: Create  $\ell + 1$  (0 layer). We first copy the row with  $\ell$  to the row with  $N$ , this can trivially be done with a ReLU layer that outputs zero everywhere else except for the row of  $N$ s that output  $(\ell)_+ - (N)_+$  to account also for the residual. We now use one of the bias terms (notice that  $\ell$  is always positive) and set to one in one of the two rows that contain the  $\ell$ , again we account for the residual as before; everything else remains unchanged. Thus, we have
$$
\left[ \begin{array}{c c c c c c c c c c c} \boldsymbol {x} _ {1} & \dots & s _ {1} ^ {\ell - 1} & s _ {1} ^ {\ell} & \dots & s _ {2} ^ {\ell - 1} & s _ {2} ^ {\ell} & \dots & s _ {n} ^ {L} & \boldsymbol {x} _ {\text {t e s t}} & \dots & \hat {\boldsymbol {s}} ^ {\ell - 1} \\ \boldsymbol {0} & \dots & \boldsymbol {0} & \boldsymbol {0} & \dots & \boldsymbol {0} & \boldsymbol {0} & \dots & \boldsymbol {0} & \boldsymbol {0} & \dots & \boldsymbol {0} \\ \hline 1 & \dots & \ell & \ell + 1 & \dots & \ell & \ell + 1 & \dots & L + 1 & 1 & \dots & \ell \\ 1 & \dots & \ell & \ell + 1 & \dots & \ell & \ell + 1 & \dots & L + 1 & 1 & \dots & \ell \\ \hline \ell + 1 & \dots & \ell + 1 & \ell + 1 & \dots & \ell + 1 & \ell + 1 & \dots & \ell + 1 & \ell + 1 & \dots & \ell + 1 \\ \ell & \dots & \ell & \ell & \dots & \ell & \ell & \dots & \ell & \ell & \dots & \ell \\ 0 & \dots & 0 & 0 & \dots & 0 & 0 & \dots & 0 & 1 & \dots & 1 \\ 1 & \dots & 1 & 1 & \dots & 1 & 1 & \dots & 1 & 1 & \dots & 1 \\ m _ {1} & \dots & m _ {\ell} & m _ {\ell + 1} & \dots & m _ {L + \ell + 1} & m _ {L + \ell + 2} & \dots & m _ {n (L + 1)} & m _ {n (L + 1) + 1} & \dots & m _ {\mathrm {N}} \end{array} \right] \tag {A.25}
$$
This operation can collectively be implemented (add the bias + copy the row) in the ReLU layer of the previous transformer layer that was not used in the previous step.
Step 4: Create the binary bits (2 layers). We will now use the information extracted in the previous steps, to create the binary indicators/bit to identify which tokens we want to filter. This can be easily implemented with one layer of transformer and especially the ReLU part of it. Notice that if we subtract the row that contains the tokens that have already been predicted, i.e.,  $[\ell \ell \dots \ell]$  from the row that contains  $[12\ldots L12\ldots L\ldots L]$  we will get zero only in the positions that we want to filter out and some non-zero quantity in the rest. This is trivially implemented with one ReLU layer. So, we need to implement an if..then type of operation. Basically, if the quantity at hand is zero we want to set the bit to one, while if it non-zero to set it to be zero. This can be implemented with the following ReLU part of a transformer layer
$$
\operatorname {T F} \left(x _ {i}\right) = 1 - \left(x _ {i}\right) _ {+} - \left(- x _ {i}\right) _ {+} + \left(x _ {i} - 1\right) _ {+} + \left(- x _ {i} - 1\right) _ {+} - \left(\left(x _ {i}\right) _ {+} - \left(- x _ {i}\right) _ {+}\right) \tag {A.26}
$$
the last two terms are to account for the residual. Again the rest of the rows do not change and are zeroed-out.
Step 5: Implement the filtering (1 layer). We now apply Lemma 4 and our proof is completed. In this step as in Lemma 4, the error remains of the order of the error of the previous step. Thus, by fine-tuning the constants appropriately, depending on 1) the bound on the input as of Assumption 2, the targeted error  $\delta$  and the constant used in Lemma 4 we achieve that the error has the target upped bound.
Theorem 2 (Theorem 1 restated) Consider a prompt  $\pmb{p}_n(f)$  generated from an  $L$ -layer MLP  $f(\cdot)$  as described in Definition 1, and assume given test example  $(\pmb{x}_{test},\pmb{s}_{test}^{1},\dots,\pmb{s}_{test}^{L})$ . For any resolution  $\epsilon >0$ , there exists  $\delta = \delta (\epsilon)$ , iteration choice  $T = \mathcal{O}(\kappa_{\mathrm{max}}^2\log (1 / \epsilon))$ , and a backend transformer construction  $TF_{BE}$  such that the concatenated transformer  $TF = TF_{LR}\circ TF_{BE}$  implements the following: Let  $(\hat{\pmb{s}}^i)_{i = 1}^{\ell -1}$  denote the first  $\ell -1$  CoT-I/O outputs of TF and set  $\pmb {p}[\ell ] = (\pmb {p}_n(f),\pmb{x}_{test},\hat{\pmb{s}}^1\dots \hat{\pmb{s}}^{\ell -1})$ . At step  $\ell$ , TF implements
1. Filtering. Define the filtered prompt with input/output features of layer  $\ell$
$$
\boldsymbol {p} _ {n} ^ {f i l t e r} = \left( \begin{array}{c c c c c c c} \ldots \mathbf {0}, \boldsymbol {s} _ {1} ^ {\ell - 1}, & \mathbf {0} & \ldots \mathbf {0}, \boldsymbol {s} _ {n} ^ {\ell - 1}, & \mathbf {0} & \ldots \mathbf {0}, \hat {\boldsymbol {s}} ^ {\ell - 1} \\ \ldots \mathbf {0}, & \mathbf {0}, & \boldsymbol {s} _ {1} ^ {\ell} & \ldots \mathbf {0}, & \mathbf {0}, & \boldsymbol {s} _ {n} ^ {\ell} & \ldots \mathbf {0}, & \mathbf {0} \end{array} \right).
$$
2. Gradient descent. The combined model obeys  $\| TF(\pmb{p}[\ell]) - \pmb{s}_{test}^{\ell} \| \leq \ell \cdot \epsilon / L$ .
$TF_{BE}$  has constant number of layers independent of  $T$  and  $n$ . Consequently, after  $L$  rounds of CoT-I/O,  $TF$  outputs  $f(\boldsymbol{x}_{test})$  up to  $\epsilon$  accuracy.
Proof. We apply Lemma 5 from which it is clear that there exists a projection such that the result stated in (1. Filtering) holds, and it is independent to  $T$ ,  $n$  and  $\ell$ . Next we turn to prove (2. Gradient descent). In Definition 1, we assume that the network's activation function is leaky-ReLU, i.e.,
$$
\phi (x) = \left\{ \begin{array}{l l} x, & \text {i f} x \geq 0 \\ \alpha x, & \text {o t h e r w i s e .} \end{array} \right. \tag {A.27}
$$
Thus, as a first step we construct the inverse of leaky-ReLU and apply it in the second row of  $\pmb{p}_n^{\mathrm{filter}}$  where the inverse of leaky-ReLU is
$$
\phi^ {- 1} (y) = \left\{ \begin{array}{l l} y, & \text {i f} y \geq 0 \\ y / \alpha , & \text {o t h e r w i s e .} \end{array} \right. \tag {A.28}
$$
This can be implemented with the following activation function (denoted by  $\sigma (\cdot)$ ) using ReLUs:
$$
\sigma (x) = (x) _ {+} - 1 / \alpha (- x) _ {+}. \tag {A.29}
$$
After it, it remains  $\mathrm{TF}_{\mathrm{LR}}$  to solve linear regression problems. Taking  $\ell$ th layer, first neuron prediction as an example, and letting  $\pmb{x}_i^{\prime} := \pmb{s}_i^{\ell - 1}$ ,  $y_i^{\prime} := \phi^{-1}(\pmb{s}_i^\ell[0])$  and  $\pmb{w} = \pmb{W}_\ell[0]$ , linear regression has form of  $y_i^{\prime} = \pmb{w}^\top \pmb{x}_i^{\prime}$  for  $i \in [n]$ . Notice that the extra zeros do not contribute in the update performed by gradient descent, and after gradient descent has been performed, we apply back the leaky ReLU. Then following Assumption 1, since we assume  $\mathrm{TF}_{\mathrm{LR}}$  performs the same as gradient descent optimizer,
given matrix condition as described in Definition 1, after running  $T$  iterations of gradient descend on the linear regression problem and considering each layer prediction with resolution  $\epsilon /L$ , we can get that  $\| \mathrm{TF}_{\mathrm{LR}}(\pmb {p}[\ell ]) - \bar{\pmb{s}}^{\ell}\| \leq \epsilon /L$ , where  $\bar{s}^{\ell} = \phi (\pmb{W}_{\ell}\hat{s}^{\ell -1})$  is the correct prediction if taking  $\hat{s}^{\ell -1}$  as input. Then we have
$$
\left\| \mathrm {T F} _ {\mathrm {L R}} (\boldsymbol {p} [ \ell ] - \boldsymbol {s} _ {\text {t e s t}} ^ {\ell}) \right\| \leq \left\| \mathrm {T F} _ {\mathrm {L R}} (\boldsymbol {p} [ \ell ]) - \bar {\boldsymbol {s}} ^ {\ell} \right\| + \left\| \boldsymbol {W} _ {\ell} \left(\hat {\boldsymbol {s}} ^ {\ell - 1} - \boldsymbol {s} _ {\text {t e s t}} ^ {\ell - 1}\right) \right\| \lesssim \epsilon / L + \left\| \mathrm {T F} _ {\mathrm {L R}} (\boldsymbol {p} [ \ell - 1 ] - \boldsymbol {s} _ {\text {t e s t}} ^ {\ell - 1}) \right\|.
$$
Let  $\mathrm{TF}_{\mathrm{LR}}(\pmb{p}[0])$  returns  $\pmb{x}_{\mathrm{test}}$  and therefore  $\| \mathrm{TF}_{\mathrm{LR}}(\pmb{p}[0]) - \pmb{x}_{\mathrm{test}} \| = 0$ . Combining results in that  $\| \mathrm{TF}_{\mathrm{LR}}(\pmb{p}[\ell] - \pmb{s}_{\mathrm{test}}^{\ell}) \| \lesssim \ell \cdot \epsilon / L$ . Since from Lemma 5 we have that we can choose  $\delta$  to be arbitrary. Let  $\delta = \epsilon / L$ , where  $L$  is the total predictions we will make. Then it will result in  $\| \mathrm{TF}(\pmb{p}[\ell] - \pmb{s}_{\mathrm{test}}^{\ell}) \| \lesssim \ell \cdot \epsilon / L$ , which completes the proof.
In this section, we provide the implementation details of our experiments.
B.1 Model Evaluation
Recap the same setting as in Section 2.3 and assume we have pretrained models with parameters  $\hat{\theta}^{\mathrm{CoT-I}}$  and  $\hat{\theta}^{\mathrm{CoT-I/O}}$ . Next we make predictions following Section 2.2. Letting  $\ell(\cdot, \cdot): \mathcal{Y} \times \mathcal{Y} \to \mathbb{R}$  be loss function, we can define test risks as follows.
$$
\mathcal {L} ^ {\mathrm {C o T - I}} (n) = \mathbb {E} _ {(\boldsymbol {x} _ {i}) _ {i = 1} ^ {n}, (f _ {\ell}) _ {\ell = 1} ^ {L}} [ \ell (\hat {\boldsymbol {y}} _ {n}, f (\boldsymbol {x} _ {n})) ] \text {w h e r e} \hat {\boldsymbol {y}} _ {n} = \mathrm {T F} (\boldsymbol {p} _ {n} (f), \boldsymbol {x} _ {n}; \hat {\boldsymbol {\theta}} ^ {\mathrm {C o T - I}})
$$
$$
\mathcal {L} ^ {\mathrm {C o T - I / O}} (n) = \mathbb {E} _ {(\boldsymbol {x} _ {i}) _ {i = 1} ^ {n}, (f _ {\ell}) _ {\ell = 1} ^ {L}} [ \ell (\hat {\boldsymbol {y}} _ {n}, f (\boldsymbol {x} _ {n})) ] \text {w h e r e} \hat {\boldsymbol {y}} _ {n} = \mathrm {T F} (\boldsymbol {p} _ {n} (f), \boldsymbol {x} _ {n}, \hat {\boldsymbol {s}} ^ {1} \dots , \hat {\boldsymbol {s}} ^ {L - 1}; \hat {\boldsymbol {\theta}} ^ {\mathrm {C o T - I / O}}).
$$
Here, we use  $\mathcal{L}(n)$  to define the test risk when given prompt with  $n$  in-context samples. Then, results shown in Figures 3&4&5&6&7(a) are test risks  $\mathcal{L}(n)$  given  $n\in [N]$ . Following model training and evaluation, we can see that once loss function is the same for both training and predicting, CoT-I (as well as ICL) accepts training risk  $\mathcal{L}_{\mathrm{train}}^{\mathrm{CoT - I}} = \frac{1}{N}\sum_{n = 1}^{N}\mathcal{L}^{\mathrm{CoT - I}}(n)$ .
B.2 Implementation
All the transformer experiments use the GPT-2 model [Radford et al., 2019] and our codebase is based on prior works [Garg et al., 2022, Wolf et al., 2019]. Specifically, the model is trained using the Adam optimizer with learning rate 0.0001 and batch size 64, and we train 500k iterations in total for all ICL, CoT-I and CoT-I/O methods. Each iteration randomly samples inputs  $\boldsymbol{x}$ s and functions  $f_{s}$ . We also apply curriculum learning over the prompt  $n$  as Garg et al. [2022] did for 2-layer random MLPs (Sec. 4.2). For both training and testing, we use the squared error as the loss function, i.e.,  $\ell(\hat{y}, y) = \| \hat{y} - y \|^2$  (or  $\ell(\hat{y}, y) = (\hat{y} - y)^2$  for scalar).
C.1 Out-of-distribution Experiments
While our primary focus lies in in-distribution scenarios, wherein the test examples are presumed to follow the same distribution as the training data, this subsection extends our analysis to out-of-distribution (OOD) settings.
To provide a clearer understanding of our model's behavior under distribution shifts and to quantitatively assess the impact on test risk, we conduct experiments with varying levels of distribution shift. Specifically, we evaluate the test risk when the prompt is supplemented with 100 in-context examples. The results of these experiments are depicted in Figure 9, which serves to illustrate how our model's performance is affected under these OOD conditions. In Fig. 9(a), we analyze noisy in-context samples during testing. The solid and dashed curves represent the test risks, corresponding to the noisy in-context samples whose (input, output) takes the form of either  $(x, y + \text{noise})$  or
(a) Test with noisy data
(b) Test with different  $k$
(c) Test with different  $d$
Figure 9: We implement robustness experiments under different distribution shift levels. In Fig. 9(a) we add noise to the label  $y$  (solid) or input features  $x$  (dashed). In Fig. 9(b), we in-context learn an MLP with  $k \in [8]$  hidden nodes whereas transformer is trained for MLPs with 8 hidden nodes. In Fig. 9(b), we consider misspecification of input dimension: TF is trained with  $d = 10$  but we feed a neural net with  $d \leq 10$ .
(a) Small GPT-2,  $d = 10$ ,  $k = 4$
Figure 10: We train ICL with more in-context examples. The main conclusion is that: For small GPT, ICL can indeed not approximate the neural net even with many examples (unlike CoT) whereas, for large GPT, ICL can do so (although much less efficient). This is in line with our theoretical intuitions on the expressivity benefits of CoT.
(b) Standard GPT-2,  $d = 10$ ,  $k = 4$
$(x + \text{noise}, y)$ , respectively. The results indicate that CoT exhibits greater robustness compared to ICL, and the test risks increase linearly with the noise level, with attributed to the randomized MLPs setting. Additionally, in Figs. 9(b) & 9(c), we instead explore out-of-distribution test tasks where test MLPs differ in  $(d, k)$  from the training phase. For both subfigures, we firstly train small GPT-2 using 2-layer MLPs with  $d = 10$ ,  $k = 8$ . In Fig. 9(b), we fix  $d = 10$  and vary  $k$  from 1 to 8, whereas in Fig. 9(c), we fix  $k = 8$  and vary  $d$  from 1 to 10. In both instances, the findings reveal that CoT's performance remains almost consistent when  $k \geq 4$  or  $d \geq 6$ , and ICL is unable to surpass it. The improved performance of ICL with smaller values of  $d$  or  $k$  again reinforces our central assertion that ICL requires  $O(dk)$  samples for in-context learning of the 2-layer random MLP, and reducing either  $d$  or  $k$  helps in improving the performance. Given that we employ the ReLU activation function, smaller values of  $d$  or  $k$  can lead to significant bias in the intermediate feature. Consequently, CoT cannot derive substantial benefits from this scenario, resulting in a decline in performance.
C.2 Further Evidence of Model Expressivity