SlowGuess/NeurIPS_2023_OCR huggingface dataset:
There are two determinants of model learnability in in-context tasks: sample complexity and model expressivity. Sample complexity pertains to the number of samples needed to precisely solve a problem. However, when the transformer model is small, even with a sufficiently large number of samples, due to its lack of expressivity, ICL cannot achieve zero test risk. This contrasts with CoT, which decomposes complex tasks into simpler sub-tasks, thereby requiring smaller models for expression. Figure 5 has illustrated the expressivity of different GPT-2 architectures, showing that the tiny GPT-2 model is too small to express even a single layer of 2-layer MLPs. Additionally, we have run more experiments, and the results are shown in Figure 10. Both Figures 10(a) and 10(b)
(a) Test filtering on first layer prediction
(b) CoT-I/O: composed risk (same as Fig. 3(a))
Figure 11: Fig. 11(a) presents a filtering evidence of 2-layer MLPs. Given a 2-layer MLP in-context example  $(\boldsymbol{x},\boldsymbol{s},\boldsymbol{y})$ , CoT admits  $(\boldsymbol{x},\boldsymbol{s},\boldsymbol{y})$  as test sample; while test samples of CoT-zero and CoT-random are formed by  $(\boldsymbol{x},\boldsymbol{s},0)$  and  $(\boldsymbol{x},\boldsymbol{s},z)$  where  $z\sim \mathcal{N}(0,d)$ . Fig. 11(b) is directly cloned from Fig. 3(a) with error bar for better comparison with results in Fig. 12.
(b) Train ICL with  $(\pmb {x},\pmb {s})$  pairs
(a) CoT-I/O: 1st layer prediction
(c) CoT-I/O: 2nd layer prediction
Figure 12: We compare the performance of filtered CoT and ICL. In Fig. 12(a)&12(c), we decouple the composed risk of predicting 2-layer MLPs into risks of individual layers (following Section 4.1), which shows the filtered CoT results. In Fig. 12(b)&12(d), we train two additional models using ICL method taking  $(\boldsymbol{x},\boldsymbol{s})$  and  $(s,y)$  as inputs.
(d) Train ICL with  $(s,y)$  pairs
detail training models with MLP tasks of dimensions  $d = 10$  and  $k = 4$ . In Fig. 10(a), we use a small GPT-2 model, and the results show that the test risk stops decreasing even with more in-context examples. In Fig. 10(b), we train a larger model, and the results demonstrate that the standard GPT-2 is sufficient to express a 2-layer MLP with  $d = 10$  and  $k = 4$ .
C.3 Filtering Evidence in 2-layer MLPs
Section 5 has demonstrated the occurrence of filtering in the linear deep MLPs setting (black dotted curves in Fig. 8(b)). In this section, we present further empirical evidence based on the 2-layer MLPs setting discussed in Section 4.2.
Follow the same setting as Figure 3(a) and choose  $d = 10$  and  $k = 8$ . Assume we have a model pretrained using CoT-I/O method. As described in Section 4.2, during training, the prompt consists of in-context samples in the form of  $(\pmb{x},\pmb{s},y)$  where  $s = (W\pmb{x})_{+}$  and  $y = \pmb{v}^{\top}\pmb{s}$ . To investigate filtering,
(a) Varying GPT-2 layers
(b) Varying GPT-2 heads
Figure 13: To further investigate how model architectures impact the prediction performance, we fix the number of heads and embedding dimension in Fig. 13(a) and change the layer number in  $\{3,6,12\}$ . Similarly for Fig. 13(b)&13(c) but instead, change number of heads (in  $\{2,4,8\}$ ) and embedding dimensions (in  $\{64,128,256\}$ ).
(c) Varying embedding dimensions
we make three different predictions to evaluate the intermediate output, whose test prompts have in-context examples with the following forms:
$$
\operatorname {C o T}: (x, s, y), \quad \operatorname {C o T - z e r o}: (x, s, 0), \quad \operatorname {C o T - r a n d o m}: (x, s, z),
$$
where  $z \sim \mathcal{N}(0, d)$ . The results are displayed in Figure 11(a) where blue, orange and green curves represent first layer prediction results using CoT, CoT-zero and CoT-random prompts, respectively. From this figure, we observe that the three curves are well aligned, indicating that when making a prediction for input  $x$ , TF will attend only to  $(x, s)$  and ignore  $y$ . Therefore filling the positions of  $y$  with any random values (or zero) will not change the performance of first layer prediction.
C.4 Comparison of Filtered CoT with ICL
Until now, many experimental results have shown that CoT-I/O provides benefits in terms of sample complexity and model expressivity compared to ICL. As an interpretation, we state that CoT can be decoupled into two phases: Filtering and ICL, and theoretical results have been provided to prove this statement. As for the empirical evidence, Sections 5 and C.3 precisely show that filtering does occur in practice. In this section, we provide additional experiments to demonstrate that, after filtering, CoT performs similarly to ICL.
For convenience and easier comparison, we repeat the same results as Fig. 3(a) in Fig. 11(b), where  $d \in \{10,20\}$ ,  $k \in \{4,8,16\}$ , and train with a small GPT-2. We again recap the data setting for the 2-layer MLP, where the in-context examples of CoT prompt are in the form of  $(x,s,y)$ . Given that filtering happens, we make first and second layer predictions following Section 4.1 and results are presented in Fig. 12(a) and Fig. 12(c), respectively. These results show the performances of the filtered CoT prompts. Next, we need to compare the performance with separate ICL training. To achieve this goal, we train a small GPT-2 model using ICL method with prompt containing  $(x,s)$  pairs (first layer). The test results are shown in Fig. 12(b). Additionally, in Fig. 12(d), we train another small GPT-2 model using ICL but with prompts containing  $(s,y)$  pairs (second layer). By comparing Fig. 12(a) and 12(b), as well as Fig. 12(c) and 12(d), we observe that after filtering, CoT-I/O achieves similar performance as individually training a single-step problem through the ICL phase.
C.5 CoT across Different Sizes of GPT-2
In Figure 5, we have demonstrated that larger models help in improving performance due to their ability of solving more complex function sets. However, since tiny, small and standard GPT-2 models scale the layer number, head number and embedding dimension simultaneously, it is difficult to determine which component has the greatest impact on performance. Therefore in this section, we investigate how different components of transformer model affect the resulting performance by run CoT-I/O on various GPT-2 architectures.
We maintain the same setting as in Section 4.2, fix  $d = 10$  and  $k = 8$ , and consider a base GPT-2 model (small GPT-2) with 6 attention layers, 4 heads in each layer and 128-dimensional embeddings.
(a) 1st layer: compare with GD solving ReLU
(b) 2nd layer: compare with least squares
Figure 14: Fig. 14(a): compare transformer results (solid) with gradient descent optimizer (dashed) when solving first layer of 2-layer MLPs; Fig. 14(b): compare transformer result (solid) with least squares optimizer (dashed) when solving the second layer of 2-layer MLPs.
In Fig. 13(a), we fix the number of heads at 4 and the embedding dimension at 128, while varying the number of layers in  $\{3,6,12\}$ . Similarly, we explore different models with different numbers of heads and embedding dimensions, and the results are respectively presented in Fig. 13(b) and 13(c). Comparing them, we can observe the following: 1) once the problem is sufficiently solved, increasing the model size does not significantly improve the prediction performance (see Fig. 13(b)&13(c); 2) the number of layers influences model expressivity, particularly for small GPT-2 architecture (see Fig. 13(a)).
C.6 Comparison of Transformer Prediction and Linear Regression
We also provide experimental findings to verify Assumption 1 in this section. Previous work [Giannou et al., 2023, Akyurek et al., 2022] has theoretically proven that TF can perform similar to gradient descent, and empirical evidence from [Dai et al., 2022, Garg et al., 2022, Li et al., 2023c] suggests that TF can even be competitive with Bayes optimizer in certain scenario. To this end, we first repeat the same first/layer predictions from Figure 4 in Figure 14, where  $d = 10$  and blue, orange and green solid curves represent the performances of  $k = 4, 16, 64$  using pretrained small GPT-2 models. We also display the evaluations of gradient descent/least square solutions in dashed curves. Specifically, in Fig. 14(a), we solve problem
$$
\hat {\boldsymbol {w}} _ {n} = \arg \min  _ {\boldsymbol {w}} \frac {1}{n} \sum_ {i = 1} ^ {n} \| (\boldsymbol {w} ^ {\top} \boldsymbol {x} _ {i}) _ {+} - y _ {i} \| ^ {2} \text {w h e r e} \boldsymbol {x} _ {i} \sim \mathcal {N} (0, \boldsymbol {I} _ {d}), y _ {i} = (\boldsymbol {w} ^ {\star^ {\top}} \boldsymbol {x} _ {i}) _ {+}
$$
for some  $\boldsymbol{w}^{\star} \sim \mathcal{N}(0, 2\boldsymbol{I}_d)$  and  $n$  is the training sample size. Then the normalized test risks are computed by  $\mathcal{L}(n) = \mathbb{E}_{\boldsymbol{w}^{\star}, \boldsymbol{x}}[\| (\hat{\boldsymbol{w}}_n^\top \boldsymbol{x})_+ - y\|^2] / d$ , and we show point-to-point results for  $n \in [N]$  in black dashed curve in Fig. 14(a). As for the second layer, we solve least squares problems as follows
$$
\hat {\boldsymbol {v}} _ {n} = \boldsymbol {S} ^ {\dagger} \boldsymbol {y} \text {w h e r e} \boldsymbol {S} \in \mathbb {R} ^ {n \times k}, \boldsymbol {S} [ i ] = (\boldsymbol {W} ^ {\star} \boldsymbol {x} _ {i}) _ {+}, \boldsymbol {y} [ i ] = \boldsymbol {v} ^ {\star^ {\top}} \boldsymbol {S} [ i ], \boldsymbol {x} _ {i} \in \mathcal {N} (0, \boldsymbol {I} _ {d})
$$
for some  $W^{\star} \in \mathbb{R}^{k \times d} \sim \mathcal{N}(0, 2 / k)$  and  $\pmb{v}^{\star} \sim \mathcal{N}(0, I_k)$ . Here,  $\dagger$  represents the pseudo-inverse operator. Then we calculate the normalized test risk of least square solution (given  $n$  training samples) as  $\mathcal{L}(n) = \mathbb{E}_{\pmb{W}^{\star}, \pmb{v}^{\star}, \pmb{x}}[\|\hat{\pmb{v}}_n^\top \pmb{s} - y\|^2] / d$  where  $\pmb{s} = (\pmb{W}^{\star} \pmb{x})_{+}$  and  $y = \pmb{v}^{\star} \top \pmb{s}$ . The results are presented in Fig. 14(b) where blue, orange and green dashed curves correspond to solving the problem using different values of  $k \in \{4, 16, 64\}$ . In this figure, the curves for  $k = 4, 16$  are aligned with GPT-2 risk curves, which indicates that TF can efficiently solve linear regression as a least squares optimizer. However, the curve for  $k = 64$  does not align, which can be attributed to the increased complexity of the function set with higher dimensionality ( $k = 64$ ). Learning such complex functions requires a larger TF model.
To mitigate the bias introduced by ReLU activation, we subtract the mean value during prediction, i.e.,  $\mathcal{L}(n) = \mathbb{E}_{\boldsymbol{w}^{\star},\boldsymbol{x}}[||(\hat{\boldsymbol{w}}_n^\top \boldsymbol{x})_+ - y - (\mathbb{E}_x[(\hat{\boldsymbol{w}}_n^\top \boldsymbol{x})_+ - y])\| ^2 ] / d.$