SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Dataset. Consider 2-layer MLPs with input  $\pmb{x} \in \mathbb{R}^d$ , hidden feature (step-1 output)  $\pmb{s} \in \mathbb{R}^k$ , and output  $y \in \mathbb{R}$ . Here,  $s = f_1(\pmb{x}) \coloneqq (\pmb{W}\pmb{x})_+$  and  $y = f_2(\pmb{s}) \coloneqq \pmb{v}^\top \pmb{s}$ , with  $\pmb{W} \in \mathbb{R}^{k\times d}$  and  $\pmb{v} \in \mathbb{R}^k$  being the parameters of the first and second layer, and  $(x)_+ = \max(x,0)$  being ReLU activation. The function is composed as  $y = v^\top(\pmb{W}\pmb{x})_+$ . We define the function distributions as follows: each entry of  $\pmb{W}$  is sampled via  $W_{ij} \sim \mathcal{N}(0,\frac{2}{k})$ , and  $\pmb{v} \sim \mathcal{N}(0,\pmb{I}_k)$ , with inputs being randomly sampled through  $x \sim \mathcal{N}(0,\pmb{I}_d)^1$ . We apply the quadratic loss in our experiments. To avoid the implicit bias due to distribution shift, both training and test datasets are generated following the same strategy.
$^1$ Following this strategy for data generation, the expected norms of  $x$ ,  $s$  and  $y$  are equivalent, and the risk curves displayed in the figures are normalized for comparison.
(a) Averaged risk
(b) Standard GPT-2
(c) Small GPT-2
(d) Tiny GPT-2
Figure 5: Comparison of the three methods for solving 2-layer MLPs using different GPT-2 models.
(a) Averaged risk
Figure 6: Comparison of the three methods for solving 2-layer MLPs with different hidden sizes.
(b)  $k = 4$
(c)  $k = 8$
(d)  $k = 16$
4.1 Empirical Evaluation of CoT-I/O Performance
To investigate how MLP task impacts CoT-I/O performance, we train 2-layer MLPs with varying input dimensions  $(d)$  and hidden layer sizes  $(k)$ . The results are presented in Figures 3 and 4, and all experiments utilize small GPT-2 models for training $^{2}$ .
CoT-I/O performance is agnostic to  $k$  when  $k \leq d$  (Figure 3). In Fig. 3(a), we train MLPs with  $d = 10, 20$  and  $k = 4, 8, 16$ . Solid and dashed curves represent the CoT-I/O test risk of  $d = 10$  and 20 respectively for varying in-context samples. The results indicate that an increase in  $d$  will amplify the number of samples needed for in-context learning, while the performance remains unaffected by changes in  $k \in \{4, 8, 16\}$ . To further scrutinize the impact of  $d$  on CoT-I/O accuracy, in Fig. 3(b), we adjust the horizontal axis by dividing it by the input dimension  $d$ , and superimpose both  $d = 10$ ,  $k = 16$  (blue solid) and  $d = 20$ ,  $k = 16$  (orange dashed) results. This alignment of the two curves implies that the in-context sample complexity of CoT-I/O is linearly dependent on  $d$ .
Large  $k$  dominates CoT-I/O performance (Figure 4). We further investigate the circumstances under which  $k$  begins to govern the CoT-I/O performance. In Figure 4(a), we replicate the same experiments with  $d = 10$ , but train with wider MLPs ( $k = 64$ ). Blue, orange and green curves represent results for  $k = 4, 16, 64$  respectively. Since the hidden dimension  $k = 64$  is larger, learning the second layer requires more hidden features ( $s$ ), thus  $N = 100$  in-context samples (providing 100  $ss$ ) are insufficient to fully restore the second layer, leading to performance gaps between  $k = 4, 16$  and  $k = 64$ . To quantify the existing gaps, we conduct single-step evaluations for both the first and the second layers, with the results shown in Figures 4(b) and 4(c). Specifically, let  $p_n(\tilde{f})$  be a test prompt containing  $n$  in-context samples where  $\tilde{f}$  represents any arbitrary 2-layer MLP. Given a test sample ( $x_{\text{test}}, s_{\text{test}}, y_{\text{test}}$ ), the layer predictions are performed as follows.
1st layer prediction:  $\mathrm{TF}(\pmb {p}_n(\tilde{f}),\pmb{x}_{\mathrm{test}})\coloneqq \hat{\pmb{s}}$
2nd layer prediction:  $\mathrm{TF}(\pmb {p}_n(\tilde{f}),\pmb{x}_{\mathrm{test}},\pmb{s}_{\mathrm{test}})\coloneqq \hat{y}$
The test risks are calculated by  $\| \hat{s} - s_{\text{test}} \|^2$  and  $(\hat{y} - y_{\text{test}})^2$ . The risks illustrated in the figures are normalized for comparability (refer to the appendix for more details). Evidence from Fig. 4(b) and 4(c) shows that while increasing  $k$  does not affect the first layer's prediction, it does augment the number of samples required to learn the second layer. Moreover, by plotting the first layer risks of  $k = 4, 16$  (blue/orange dotted) and second layer risk of  $k = 64$  (green dashed) in Fig. 4(a), we can see that they align with the CoT-I/O composed risks. This substantiates the hypothesis that CoT-I/O learns 2-layer MLP through compositional learning of separate layers.
2Our code is available at https://github.com/yingcong-li/Dissecting-CoT.
(a) Point-to-point meta prediction
(b) One-shot prediction over time
Figure 7: Evaluations over deep linear MLPs using CoT-I/O and ICL where CoT-  $X$  represents the  $X$ -step CoT-I/O. Fig. 7(a) illustrates point-to-point meta results where the model is trained with substantial number of samples. In contrast, Fig. 7(b) displays the one-shot performance (with only one in-context sample provided) when making predictions during training.
4.2 Comparative Analysis of ICL, CoT-I and CoT-I/O
Varying model sizes (Figure 5). We initially assess the benefits of CoT-I/O over ICL and CoT-I across different TF models. With  $d = 10$  and  $k = 8$  fixed, we train three different GPT-2 models: standard, small and tiny GPT-2. The small GPT-2 has 6 layers, 4 attention heads per layer and 128 dimensional embeddings. The standard GPT-2 consists of twice the number of layers, attention heads and embedding dimensionality compared to the small GPT-2, and tiny GPT-2, on the other hand, possesses only half of these hyperparameters compared to the small GPT-2. We evaluate the performance using prompts containing  $n$  in-context samples, where  $n$  ranges from 1 to  $N$  ( $N = 100$ ). The associated test risks are displayed in Figs. 5(b), 5(c) and 5(d). The blue, orange and green curves correspond to CoT-I/O, CoT-I and ICL, respectively. In Fig. 5(a), we present the averaged risks. The results show that using CoT-I/O, the small GPT-2 can solve 2-layer MLPs with approximately 60 samples, while CoT-I requires the standard GPT-2. Conversely, ICL is unable to achieve zero test risk even with the standard GPT-2 model and up to 100 samples. This indicates that to learn 2-layer MLPs in a single shot, ICL requires at least  $\mathcal{O}(dk + d)$  samples to restore all function parameters. Conversely, CoT-I and CoT-I/O can leverage implicit samples contained in the CoT prompt. Let  $f_{1} \in \mathcal{F}_{1}$  (first layer) and  $f_{2} \in \mathcal{F}_{2}$  (second layer). By comparing the performance of CoT-I and CoT-I/O, it becomes evident that the standard GPT-2 is capable of learning the composed function  $f = f_{2} \circ f_{1} \in \mathcal{F}$ , which the small GPT-2 cannot express.
Varying MLP widths (Figure 6). Next, we explore how different MLP widths impact the performance (by varying the hidden neuron size  $k \in \{4,8,16\}$ ). The corresponding results are depicted in Figure 6. The blue, orange and green curves in Fig. 6(b), 6(c) and 6(d) correspond to hidden layer sizes of  $k = 4,8$ , and 16, respectively. Fig. 6(a) displays the averaged risks. We keep  $d = 10$ ,  $N = 100$  fixed and train with the small GPT-2 model. As discussed in Section 4.1, CoT-I/O can learn a 2-layer MLP using around 60 samples for all  $k = 4,8,16$  due to its capability to deconstruct composed functions. However, CoT-I can only learn the narrow MLPs with  $k = 4$ , and ICL is unable to learn any of them. Moreover, we observe a substantial difference in the performance of ICL and CoT-I with varying  $k$  (e.g., see averaged risks in Fig. 6(a)). This can be explained by the fact that enlarging  $k$  results in more complex  $\mathcal{F}_1$  and  $\mathcal{F}_2$ , thus making the learning of  $\mathcal{F} = \mathcal{F}_2 \times \mathcal{F}_1$  more challenging for ICL and CoT-I.
5 Further Investigation on Deep Linear MLPs
In Section 4, we have discussed the approximation benefits of CoT-I/O and how it in-context learns 2-layer random MLPs by parallel learning of  $k$ $d$ -dimensional ReLU and  $1$ $k$ -dimensional linear regression. In this section, we investigate the capability of CoT-I/O in learning longer compositions. For brevity, we will use CoT to refer to CoT-I/O in the rest of the discussion.
Dataset. Consider  $L$ -layer linear MLPs with input  $\pmb{x} \in \mathbb{R}^d \sim \mathcal{N}(0, I_d)$ , and output generated by  $\pmb{y} = \pmb{W}_L \pmb{W}_{L-1} \cdots \pmb{W}_1 \pmb{x}$ , where the  $\ell$ th layer is parameterized by  $\pmb{W}_{\ell} \in \mathbb{R}^{d \times d}$ ,  $\ell \in [L]$ . In this work, to better understand the emerging ability of CoT, we assume that each layer draws from the same discrete sub-function set  $\bar{\mathcal{F}} = \{\bar{W}_k : \bar{W}_k^\top \bar{W}_k = I, k \in [K]\}^3$ . Therefore, to learn the  $L$ -layer neural net, CoT only needs to learn  $\bar{\mathcal{F}}$  with  $|\bar{\mathcal{F}}| = K$ , whereas ICL needs to learn the function set  $\bar{\mathcal{F}}^L$ , which contains  $K^L$  random matrices.
3This assumption ensures that the norm of the feature remains constant across layers ( $\| \pmb{x}\| = \| \pmb{y}\| = \| \pmb{s}^{\ell}\|$ ), enabling fair evaluation across different layers.
(a) Composed risk
(b) Risk of each layer output
Figure 8: Fig. 8(a) is generated by magnifying the initial 50k iterations of Fig. 7(b), and we decouple the composed risks from predicting 6-layer linear MLPs into predictions for each layer, and the results are depicted in Fig. 8(b). Additional implementation details can be found in Section 5.
Composition ability of CoT (Figures 7). Set  $d = 5$ ,  $L = 6$  and  $K = 4$ . At each round, we randomly select  $L$  matrices  $W_{\ell}$ ,  $\ell \in [L]$  from  $\mathcal{F}$  so that for any input  $x$ , we can form a chain
$$
\boldsymbol {x} \rightarrow s ^ {1} \rightarrow s ^ {2} \dots \rightarrow s ^ {6} := \boldsymbol {y},
$$
where  $s^{\ell} = W_{\ell}s^{\ell -1}, \ell \in [L]$  and  $s^0 \coloneqq x$ . Let CoT-  $X$  denote  $X$ -step CoT-I/O method. For example, the in-context sample of CoT-6 has form of  $(x,s^{1},s^{2},\ldots s^{5},y)$ , which contains all the intermediate outputs from each layer; while CoT-3, CoT-2 have prompt samples formed as  $(x,s^{2},s^{4},y)$  and  $(x,s^{3},y)$  respectively. In this setting, ICL is also termed as CoT-1, as its prompt contains only  $(x,y)$  pairs. To solve the length-6 chain, CoT-  $X$  needs to learn a model that can remember  $4^{6 / X}$  matrices. Therefore, ICL is face a significantly challenge since it needs to remember  $4^{6} = 4$ , 096 matrices (all combinations of the 4 matrices used for training and testing) compared to just 4 for CoT-6.
We train small GPT-2 models using the CoT-2/-3/-6 and ICL methods, and present the results in Fig. 7(a). As evident from the figure, the performance curves of CoT-2 (orange), CoT-3 (green) and CoT-6 (red) overlap, and they can all make precise predictions in one shot (given an in-context example  $n = 1$ ). It seems that TF has effectively learned to remember up to 64 matrices (for CoT-2) and compose up to 6 layers (for CoT-6). However, ICL (blue) struggles to learn the 6-layer MLPs in one shot. The black dashed curve shows the solution for linear regression  $y = \beta^{\top}x$  computed directly via least squares given  $n$  random training samples, where  $x$  is the input and  $y$  is from the output of the 6-layer MLPs (e.g.,  $y[0]$ ). The test risks for  $n = 1, \ldots, 10$  are plotted (in Fig. 7(a)), which show that the ICL curve aligns with the least squares performance. This implies that, instead of remembering all 4,096 matrices, ICL solves the problem from the linear regression phase.
In addition to the meta-learning results which highlight the approximation benefits of multi-step CoT, we also investigate the convergence rate of CoT-2/-3/-6 and ICL, with results displayed in Fig. 7(b). We test the one-shot performance during training and find that CoT-6 converges fastest. This is because it has the smallest sub-function set, and given the same tasks (e.g., deep neural nets), shortening the chain leads to slower convergence. This supports the evidence that taking more steps facilitates faster and more effective learning of complex problems.
Evidence of Filtering (Figure 8). As per Theorem 1 and the appendix, transformers can perform filtering over CoT prompts, and the results from 2-layer MLPs align with our theoretical findings. However, can we explicitly observe filtering behaviors? In Fig. 8(a), we display the results of the first 50k iterations from Fig. 7(b), and observe risk drops in CoT-6 (red) at the 15k and 25k iteration (shown as grey dotted and dashed lines). Subsequently, in Fig. 8(b), we plot the test risk of each layer prediction (by feeding the model with correct intermediate features not the predicted ones), where CoT-6 (red) predicts the outputs from all 6 layers  $(s^1,\dots ,s^L)$ . From these figures, we can identify risk drops when predicting different layers, which appear at either 15k (for layer 2, 3, 4, 5, 6) or 25k (for layer 1) iteration. This implies that the model learns to predict each step/layer function independently. Further, we test the filtering evidence of the  $\ell$ th layer by filling irrelevant positions with random features. Specifically, an in-context example is formed by
$$
\left(\boldsymbol {z} ^ {0}, \dots , \boldsymbol {s} ^ {\ell - 1}, \boldsymbol {s} ^ {\ell}, \boldsymbol {z} ^ {\ell + 1}, \dots \boldsymbol {z} ^ {L}\right), \text {w h e r e} \boldsymbol {s} ^ {\ell} = \boldsymbol {W} _ {\ell} \left(\boldsymbol {s} ^ {\ell - 1}\right) \text {a n d} \boldsymbol {z} \sim \mathcal {N} (0, \boldsymbol {I} _ {d}).
$$
The test risks are represented by black dotted curves in Fig. 8(b), which aligned precisely with the CoT-6 curves (red). This signifies that each layer's prediction concentrate solely on the corresponding intermediate steps in the prefix, while disregarding irrelevant features. This observation provides evidence that the process of filtering is indeed performed.
With the success of LLMs and prompt structure [Lester et al., 2021], there is growing interest in in-context learning (ICL) from both theoretical and experimental lens [Garg et al., 2022, Brown et al., 2020, von Oswald et al., 2022, Dai et al., 2022, Min et al., 2022, Lyu et al., 2022, Li et al., 2023c, Balim et al., 2023, Xie et al., 2021, Min et al., 2021, Wei et al., 2023, Li et al., 2023a]. As an extension, chain-of-thought (CoT) prompts have made impressive improvements in performing complex reasoning by decomposing it into step-by-step intermediate solutions [Wei et al., 2022c, Narang et al., 2020, Lampinen et al., 2022, Wei et al., 2022b, Zhou et al., 2022, Nye et al., 2021, Velicković and Blundell, 2021, Lanchantin et al., 2023], which in general, shows the ability of transformer in solving compositional functions. Lee et al. [2023], Dziri et al. [2023] study the problem of teaching arithmetic to small transformers and show that breaking the task down into small reasoning steps allows the model to learn faster. Li et al. [2023b] show that small student transformer models can learn from rationalizations sampled from significantly larger teacher models. The idea of learning how to compose skills has been well studied in other literatures [Sahni et al., 2017, Liška et al., 2018]. More specifically, for the problem of learning shallow networks, there are several well known hardness results Goel et al. [2017, 2020], Zhang et al. [2019]. In particular, Hahn and Goyal [2023] shows a formal learnability bound which implies that compositional structure can benefit ICL. However, most of the work focuses on investigating empirical benefits and algorithmic designing of CoT, and there exists little effort studying the underlying mechanisms of CoT.
Considering the expressivity of the transformer architecture itself, Yun et al. [2019] showed that TFs are universal sequence to sequence approximators. More recently, Giannou et al. [2023] use an explicit construction to show that shallow TFs can be used to run general purpose programs as long as we loop them. Other works have also shown the tuning-completeness of the TF architecture but these typically require infinite/high precision and recursion around attention layers [Wei et al., 2022a, Pérez et al., 2019, 2021, Liu et al., 2022]. Closer to our work, Akyürek et al. [2022], Von Oswald et al. [2023], von Oswald et al. [2023] prove that a transformer with constant number of layers can implement gradient descent in solving linear regression, and Giannou et al. [2023] introduce similar results by looping outputs back into inputs. Ahn et al. [2023] prove this result from an optimization perspective and show that the global minimum of the training objective implements a single iteration of preconditioned gradient descent for transformers with a single layer of linear attention. Zhou et al. [2023] introduce the RASP-Generalization Conjecture which says that Transformers tend to length generalize on a task if the task can be solved by a short RASP program which works for all input lengths. In this work, we prove CoT can be treated as: first apply filtering on the CoT prompts using special construction, and then in-context learn the filtered prompt.
7 Conclusion, Limitations, and Discussion
In this work, we investigate chain-of-thought prompting and shed light on how it enables compositional learning of multilayer perceptrons step-by-step. Specially, we have explored and contrasted three methods: ICL, CoT-I and CoT-I/O, and found that CoT-I/O facilitates better approximation and faster convergence through looping and sample efficiency. Additionally, we empirically and theoretically demonstrated that to learn a 2-layer MLP with  $d$ -dimensional input and  $k$  neurons, CoT-I/O requires  $\mathcal{O}(\max(d, k))$  in-context samples whereas ICL runs into approximation error bottlenecks.
While we have provided both experimental and theoretical results to validate the advantages of CoT, it is important to note that our analysis in the main text pertains to in-distribution scenarios. In an effort to address this limitation and demonstrate the robustness of CoT, we have conducted additional simulations, as detailed in Appendix C.1, where the test samples follow a different distribution than the training examples. Also, we note that our focus has been primarily on MLP-based tasks, where the subproblems are essentially instances of simple linear regression. It would be valuable to explore how CoT might influence the training of tasks characterized by more complex structures, longer compositional chains, and a broader variety of subproblems.
There are several interesting avenues for future research to build on our findings. To what extent does our decoupling of CoT (filtering followed by ICL) align with the empirical evidence in practical problems such as code generation and mathematical reasoning? We have shown that CoT-I/O can rely on a linear regression oracle to learn an MLP. To what extent can transformers approximate MLPs without CoT-I/O (e.g. with CoT-I) and what are the corresponding lower/upper bounds?
This work was supported by the NSF CAREER awards #2046816 and #1844951, NSF grant #2212426, AFOSR & AFRL Center of Excellence Award FA9550-18-1-0166, a Google Research Scholar award, an Adobe Data Science Research award, an Army Research Office grant W911NF2110312, and an ONR Grant No. N00014-21-1-2806.
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. arXiv preprint arXiv:2306.00297, 2023.
Ekin Akyurek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. arXiv preprint arXiv:2211.15661, 2022.
Haldun Balim, Zhe Du, Samet Oymak, and Necmiye Ozay. Can transformers learn optimal filtering for unknown systems? arXiv preprint arXiv:2308.08536, 2023.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Zhifang Sui, and Furu Wei. Why can gpt learn in-context? language models secretly perform gradient descent as meta optimizers. arXiv preprint arXiv:2212.10559, 2022.
Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jian, Bill Yuchen Lin, Peter West, Chandra Bhagavatula, Ronan Le Bras, Jena D Hwang, et al. Faith and fate: Limits of transformers on compositionality. arXiv preprint arXiv:2305.18654, 2023.
Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30598, 2022.
Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D Lee, and Dimitris Papailiopoulos. Looped transformers as programmable computers. arXiv preprint arXiv:2301.13196, 2023.
Surbhi Goel, Varun Kanade, Adam Klivans, and Justin Thaler. Reliably learning the relu in polynomial time. In Conference on Learning Theory, pages 1004-1042. PMLR, 2017.
Surbhi Goel, Adam Klivans, Pasin Manurangsi, and Daniel Reichman. Tight hardness results for training depth-2 relu networks. arXiv preprint arXiv:2011.13550, 2020.
Michael Hahn and Navin Goyal. A theory of emergent in-context learning as implicit structure induction. arXiv preprint arXiv:2303.07971, 2023.
Shima Imani, Liang Du, and Harsh Shrivastava. Mathprompter: Mathematical reasoning using large language models. arXiv preprint arXiv:2303.05398, 2023.
Andrew K Lampinen, Ishita Dasgupta, Stephanie CY Chan, Kory Matthewson, Michael Henry Tessler, Antonia Creswell, James L McClelland, Jane X Wang, and Felix Hill. Can language models learn from explanations in context? arXiv preprint arXiv:2204.02329, 2022.
Jack Lanchantin, Shubham Toshniwal, Jason Weston, Arthur Szlam, and Sainbayar Sukhbaatar. Learning to reason and memorize with self-notes. arXiv preprint arXiv:2305.00833, 2023.
Nayoung Lee, Kartik Sreenivasan, Jason D Lee, Kangwook Lee, and Dimitris Papailiopoulos. Teaching arithmetic to small transformers. arXiv preprint arXiv:2307.03381, 2023.
Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for parameter-efficient prompt tuning. arXiv preprint arXiv:2104.08691, 2021.