SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yingcong Li  
University of California, Riverside  
yli692@ucr.edu
Dimitris Papailiopoulos  
University of Wisconsin-Madison  
dimitris@papail.io
Kartik Sreenivasan Angeliki Giannou  
University of Wisconsin-Madison  
{ksreenivasa2, giannou}@wisc.edu
Samet Oymak  
University of Michigan & UC Riverside  
oymak@umich.edu
Chain-of-thought (CoT) is a method that enables language models to handle complex reasoning tasks by decomposing them into simpler steps. Despite its success, the underlying mechanics of CoT are not yet fully understood. In an attempt to shed light on this, our study investigates the impact of CoT on the ability of transformers to in-context learn a simple to study, yet general family of compositional functions: multi-layer perceptrons (MLPs). In this setting, we find that the success of CoT can be attributed to breaking down in-context learning of a compositional function into two distinct phases: focusing on and filtering data related to each step of the composition and in-context learning the single-step composition function. Through both experimental and theoretical evidence, we demonstrate how CoT significantly reduces the sample complexity of in-context learning (ICL) and facilitates the learning of complex functions that non-CoT methods struggle with. Furthermore, we illustrate how transformers can transition from vanilla in-context learning to mastering a compositional function with CoT by simply incorporating additional layers that perform the necessary data-filtering for CoT via the attention mechanism. In addition to these test-time benefits, we show CoT helps accelerate pretraining by learning shortcuts to represent complex functions and filtering plays an important role in this process. These findings collectively provide insights into the mechanics of CoT, inviting further investigation of its role in complex reasoning tasks.
The advent of transformers [Vaswani et al., 2017] has revolutionized natural language processing, paving the way for remarkable performance in a wide array of tasks. LLMs, such as GPTs [Brown et al., 2020], have demonstrated an unparalleled ability to capture and leverage vast amounts of data, thereby facilitating near human-level performance across a variety of language generation tasks. Despite this success, a deep understanding of their underlying mechanisms remains elusive.
Chain-of-thought prompting [Wei et al., 2022c] is an emergent ability of transformers where the model solves a complex problem [Wei et al., 2022b], by decomposing it into intermediate steps. Intuitively, this underlies the ability of general-purpose language models to accomplish previously-unseen complex tasks by leveraging more basic skills acquired during the pretraining phase. Compositional learning and CoT has enjoyed significant recent success in practical language modeling tasks spanning question answering, code generation, and mathematical reasoning [Perez et al., 2021, Imani et al., 2023, Yuan et al., 2023]. In this work, we attempt to demystify some of the mechanics underlying this success and the benefits of CoT in terms of sample complexity and approximation power. To
Figure 1: An illustration of ICL, CoT-I and CoT-I/O methods, using a 3-layer MLP as an example (top left, where  $x$ ,  $y$ ,  $s^1$ ,  $s^2$  denote input, output and hidden features respectively). The ICL method utilizes in-context examples in the form of  $(x, y)$  and makes predictions directly based on the provided  $x_{\mathrm{test}}$ . Both CoT-I and CoT-I/O methods admit prompts with samples formed by  $(x, s^1, s^2, y)$ . However, CoT-I/O uniquely makes recurrent predictions by re-inputting the intermediate output (as shown on the right). The performance of these methods is shown on the bottom left, with a more detailed discussion available in Section 4.
do this we explore the role of CoT in learning multi-layer perceptrons (MLPs) in-context, which we believe can lead to a first set of insightful observations. Throughout, we ask:
Does CoT improve in-context learning of MLPs, and what are the underlying mechanics?
In this work, we identify and thoroughly compare three schemes as illustrated in Figure 1. (a) ICL: In-context learning from input-output pairs provided in the prompt, (b) CoT-I: Examples in the prompt are augmented with intermediate steps, (c) CoT-I/O: The model also outputs intermediate steps during prediction. Our main contributions are:
- Decomposing CoT into filtering and ICL: As our central contribution, we establish a rigorous and experimentally-supported abstraction that decouples CoT prompting into a filtering phase and an in-context learning (ICL) phase. In filtering, the model attends to the relevant tokens within the prompt based on an instruction. In ICL, the model runs inference on the filtered prompt to output a step and then moves onto the next step in the chain. Our Theorem 1 develops a theoretical understanding of this two-step procedure and formalizes how filtering and ICL phases of CoT can be implemented via the self-attention mechanism to learn MLPs.
- Approximation and sample complexity: Through experiments and theory, we establish that intermediate steps in CoT-I improves the sample complexity of learning whereas step-by-step output improves the approximation ability through looping. Specifically, CoT-I/O can learn an MLP with input dimension  $d$  and  $k$  neurons using  $\mathcal{O}(\max(k, d))$  in-context samples by filtering individual layers and solving them via linear regression – in contrast to the  $\Omega(kd)$  lower bound without step-augmented prompt. As predicted by our theory, our experiments (see Sec. 4) identify a striking universality phenomenon (as  $k$  varies) and also demonstrate clear approximation benefits of CoT compared to vanilla ICL.
- Accelerated pretraining via learning shortcuts: We construct deep linear MLPs where each layer is chosen from a discrete set of matrices. This is in contrast to the above setting, where MLP weights can be arbitrary. We show that CoT can dramatically accelerate pretraining by memorizing these discrete matrices and can infer all layers correctly from a single demonstration. Notably the pretraining loss goes to zero step-by-step where each step "learns to filter a layer". Together, these showcase how CoT identifies composable shortcuts to avoid the need for solving linear regression. In contrast, we show that ICL (without CoT) collapses to linear regression performance as it fails to memorize exponentially many candidates (due to lack of composition).
The paper is organized as follows. In Section 2, we introduce the problem setup. Section 3 states our main theoretical results which decouple CoT into filtering and ICL. Section 4 provides empirical investigations of CoT with 2-layer MLPs, which validates our theoretical findings. Finally, we
elucidate the benefits of CoT during pretraining via experiments on deep linear MLPs in Section 5. Related work and discussion are provided in Sections 6 and 7, respectively.
2 Preliminaries and Setup
Notation. We denote the set  $\{1,2,\ldots ,n\}$  as  $[n]$ . Vectors and matrices are represented in bold text (e.g.,  $\pmb{x},\pmb{A}$ ), while scalars are denoted in plain text (e.g.,  $y$ ). The input and output domains are symbolized as  $\mathcal{X}$  and  $\mathcal{V}$  respectively (unless specified otherwise), and  $\pmb{x} \in \mathcal{X}$ ,  $\pmb{y} \in \mathcal{V}$  denote the input and output. Additionally, let  $\mathcal{F}$  be a set of functions from  $\mathcal{X}$  to  $\mathcal{V}$ . Consider a transition function  $f \in \mathcal{F}$  where  $\pmb{y} = f(\pmb{x})$ . In this section, we explore the properties of learning  $f$ , assuming that it can be decomposed into simpler functions  $(f_{\ell})_{\ell = 1}^{L}$ , and thus can be expressed as  $f = f_{L} \circ f_{L - 1} \circ \dots \circ f_{1}$ .
2.1 In-context Learning
Following the study by Garg et al. [2022], the fundamental problem of vanilla in-context learning (ICL) involves constructing a prompt with input-output pairs in the following manner:
$$
\boldsymbol {p} _ {n} (f) = \left(\boldsymbol {x} _ {i}, \boldsymbol {y} _ {i}\right) _ {i = 1} ^ {n} \quad \text {w h e r e} \quad \boldsymbol {y} _ {i} = f (\boldsymbol {x} _ {i}). \tag {P-ICL}
$$
Here the transition function  $f \in \mathcal{F} : \mathcal{X} \to \mathcal{Y}$  remains constant within a single prompt but can vary across prompts, and the subscript  $n$  signifies the number of in-context samples contained in the prompt. Considering language translation as an example,  $f$  is identified as the target language, and the prompt can be defined as  $p(\text{Spanish}) = ((apple, manzana), (ball, pelota), \ldots)$  or  $p(\text{French}) = ((cat, chat), (flower, fleur), \ldots)$ . Let TF denote any auto-regressive model (e.g., Decoder-only Transformer). The aim of in-context learning is to learn a model that can accurately predict the output, given a prompt  $p$  and the test input  $x_{\text{test}}$ , as shown in the following equation:
$$
\operatorname {T F} \left(\boldsymbol {p} _ {n} (\tilde {f}), \boldsymbol {x} _ {\text {t e s t}}\right) \approx \tilde {f} \left(\boldsymbol {x} _ {\text {t e s t}}\right) \tag {2.1}
$$
where  $\tilde{f} \in \mathcal{F}$  is the test function which may differ from the functions used during training. Previous work [Zhou et al., 2022, Li et al., 2023c] has demonstrated that longer prompts (containing more examples  $n$ ) typically enhance the performance of the model.
2.2 Chain-of-thought Prompt and Prediction
As defined in (P-ICL), the prompt in vanilla ICL only contains input-output pairs of the target function. This demands that the model learns the function  $f \in \mathcal{F}$  in one go, which becomes more challenging as  $\mathcal{F}$  grows more complex, since larger models and increased prompt length  $(n)$  are needed to make correct predictions (as depicted by the green curves in Figures 5 and 6). Existing studies on chain-of-thought methods (e.g., [Wei et al., 2022c]) observed that prompts containing step-by-step instructions assist the model in decomposing the function and making better predictions. Specifically, consider a function composed of  $L$  subfunctions, represented as  $f \coloneqq f_{L} \circ f_{L - 1} \circ \ldots \circ f_{1}$ . Each intermediate output can be viewed as a step, enabling us to define a length- $n$  CoT prompt related to  $f$  with  $L$  steps (expressed with  $s^{\ell}, \ell \in [L]$ ) as follows:
$$
\boldsymbol {p} _ {n} (f) = \left(\boldsymbol {x} _ {i}, \boldsymbol {s} _ {i} ^ {1}, \dots \boldsymbol {s} _ {i} ^ {L - 1}, \boldsymbol {s} _ {i} ^ {L}\right) _ {i = 1} ^ {n} \quad \text {w h e r e} \quad \boldsymbol {s} _ {i} ^ {\ell} = f _ {\ell} \left(\boldsymbol {s} _ {i} ^ {\ell - 1}\right), \ell \in [ L ]. \tag {P-CoT}
$$
Here  $\pmb{x}_i = s_i^0$ ,  $\pmb{y}_i = s_i^L$  and  $f_{\ell} \in \mathcal{F}_{\ell}$ , which implies that  $f \in \mathcal{F}_L \times \dots \mathcal{F}_1 \coloneqq \mathcal{F}$ .
Next we introduce two methodologies for making predictions within the CoT framework:
CoT over input only (CoT-I). Contrasted with ICL, CoT-I considers step-by-step instructions as inputs, nonetheless, the prediction for the last token is performed as a single entity. Our experiments indicate that this approach lowers the sample complexity for TF to comprehend the function  $\tilde{f}$  being learned (see the orange curves in Figures 5 and 6). The CoT-I prediction aligns with Eq. (2.1) as follows, while the prompt is determined by (P-CoT).
$$
\text {O n e - s h o t p r e d i c t i o n :} \mathrm {T F} \left(\boldsymbol {p} _ {n} (\tilde {f}), \boldsymbol {x} _ {\text {t e s t}}\right) \approx \tilde {f} \left(\boldsymbol {x} _ {\text {t e s t}}\right). \tag {2.2}
$$
CoT over both input and output (CoT-I/O). Despite the fact that CoT-I improves the sample complexity of learning  $\tilde{f}$ , the TF must still possess the capacity to approximate functions from the function class  $\mathcal{F}$ , given that the prediction is made in one shot. To mitigate this challenge, we consider
a scenario where in addition to implementing a CoT prompt, we also carry out CoT predictions. Specifically, for a composed problem with inputs formed via (P-CoT), the model recurrently makes  $L$ -step predictions as outlined below:
$$
\text {S t e p} 1: \operatorname {T F} \left(\boldsymbol {p} _ {n} (\tilde {f}), \boldsymbol {x} _ {\text {t e s t}}\right) := \hat {\boldsymbol {s}} ^ {1}
$$
$$
\text {S t e p 2 : T F} \left(\boldsymbol {p} _ {n} (\tilde {f}), \boldsymbol {x} _ {\text {t e s t}}, \hat {\boldsymbol {s}} ^ {1}\right) := \hat {\boldsymbol {s}} ^ {2}
$$
$$
\begin{array}{c} \vdots \\ \vdots \\ \vdots \end{array}
$$
$$
\operatorname {S e t p} L: \operatorname {T F} \left(\boldsymbol {p} _ {n} (\tilde {f}), \boldsymbol {x} _ {\text {t e s t}}, \hat {\boldsymbol {s}} ^ {1} \dots , \hat {\boldsymbol {s}} ^ {L - 1}\right) \approx \tilde {f} \left(\boldsymbol {x} _ {\text {t e s t}}\right), \tag {2.3}
$$
where at each step (step  $\ell$ ), the model outputs an intermediate step  $(\hat{s}^{\ell})$  which is then fed back to the input sequence to facilitate the next-step prediction  $(\hat{s}^{\ell +1})$ . Following this strategy, the model only needs to learn the union of the sub-function sets,  $\bigcup_{\ell = 1}^{L}\mathcal{F}_{\ell}$ , whose complexity scales linearly with the number of steps  $L$ . Empirical evidence of the benefits of CoT-I/O over ICL and CoT-I in enhancing sample efficiency and model expressivity is reflected in the blue curves shown in Figures 5 and 6.
2.3 Model Training
In Figure 1 and Section 2, we have discussed vanilla ICL, CoT-I and CoT-I/O methods. Intuitively, ICL can be viewed as a special case of CoT-I (or CoT-I/O) if we assume only one step is performed. Consequently, we will focus on implementing CoT-I and CoT-I/O for model training in the following.
Consider the CoT prompt as in (P-CoT), and assume that  $\pmb{x} \sim \mathcal{D}_{\mathcal{X}}$ , and  $f_{\ell} \sim \mathcal{D}_{\ell}, \ell \in [L]$ , where  $L$  denotes the number of compositions/steps, such that the final prediction should approximate  $f(\pmb{x}) = f_{L}(f_{L-1} \dots f_{1}(\pmb{x})) \coloneqq \pmb{y} \in \mathcal{V}$ . We define  $\ell(\hat{\pmb{y}}, \pmb{y}) : \mathcal{V} \times \mathcal{V} \to \mathbb{R}$  as a loss function. For simplicity, we assume  $f_{\ell}(\dots f_{1}(\pmb{x})) \in \mathcal{V}, \ell \in [L]$ . Let  $N$  represent the in-context window of TF, which implies that TF can only admit a prompt containing up to  $N$  in-context samples. Generally, our goal is to ensure high prediction performance given any length- $n$  prompt, where  $n \in [N]$ . To this end, we train the model using prompts with length from 1 to  $N$  equally and aim to minimize the averaged risk over different prompt size. Assuming the model TF is parameterized by  $\pmb{\theta}$  and considering meta learning problem, the objective functions for CoT-I and CoT-I/O are defined as follows.
$$
\hat {\pmb {\theta}} ^ {\mathrm {C o T - I}} = \arg \min _ {\pmb {\theta}} \mathbb {E} _ {(\pmb {x} _ {n}) _ {n = 1} ^ {N}, (f _ {\ell}) _ {\ell = 1} ^ {L}} \left[ \frac {1}{N} \sum_ {n = 1} ^ {N} \ell (\hat {\pmb {y}} _ {n}, f (\pmb {x} _ {n})) \right]
$$
$$
\hat {\boldsymbol {\theta}} ^ {\mathrm {C o T - I / O}} = \arg \min _ {\boldsymbol {\theta}} \mathbb {E} _ {(\boldsymbol {x} _ {n}) _ {n = 1} ^ {N}, (f _ {\ell}) _ {\ell = 1} ^ {L}} \left[ \frac {1}{N L} \sum_ {n = 1} ^ {N} \sum_ {\ell = 1} ^ {L} \ell (\hat {\boldsymbol {s}} _ {n} ^ {\ell}, \boldsymbol {s} _ {n} ^ {\ell}) \right]
$$
where  $\hat{\pmb{y}}_n = \mathrm{TF}(\pmb{p}_n(f),\pmb{x}_n)$  and  $\hat{s}_n^\ell = \mathrm{TF}(\pmb{p}_n(f),\pmb{x}_n\dots \pmb{s}_n^{\ell -1})$ $\pmb{p}_n(f)$  is given by (P-CoT), and as mentioned previously,  $s^0 = x$  and  $s^L = y$ . All  $\pmb{x}$  and  $f_{\ell}$  are independent, and we take the expectation of the risk over their respective distributions.
3 Provable Approximation of MLPs via Chain-of-Thought
In this section, we present our theoretical findings that demonstrate how CoT-I/O can execute filtering over the CoT prompt, thereby learning a 2-layer MLP with input dimension of  $d$  and hidden dimension of  $k$ , akin to resolving  $k$ $d$ -dimensional ReLU problems and  $1$ $k$ -dimensional linear regression problem. Subsequently, in Section 4.1, we examine the performance of CoT-I/O when learning 2-layer random MLPs. Our experiments indicate that CoT-I/O needs only  $O(\max(d, k))$  in-context samples to learn the corresponding MLP.
We state our main contribution of establishing a result that decouples CoT-based in-context learning (CoT-I/O) into two phases: (1) Filtering Phase: Given a prompt that contains features of multiple MLP layers, retrieve only the features related to a target layer to create an ICL prompt. (2) ICL Phase: Given filtered prompt, learn the target layer weights through gradient descent. Combining these two
Figure 2: Depiction of Theorem 1 that formalizes CoT as a Filtering + ICL procedure. In this illustration, we utilize in-context samples as depicted in Figure 1, characterizing the CoT prompt by  $(x, s^1, s^2, y)$ . We show a transformer TF, composed of  $\mathrm{TF}_{\mathrm{LR}} \circ \mathrm{TF}_{\mathrm{BE}}$ .  $\mathrm{TF}_{\mathrm{BE}}$  facilitates the implementation of filtering on the CoT prompt (refer to (P-CoT), illustrated on the left side of the figure), subsequently generating vanilla ICL prompts (refer to (P-ICL), illustrated on the right side of the figure). During each step of the ICL process,  $\mathrm{TF}_{\mathrm{LR}}$  employs gradient descent techniques to address and solve sub-problems, exemplified by linear regression in the context of solving MLPs.
phases, and looping over all layers, we will show that there exists a transformer architecture such that CoT-I/O can provably approximate a multilayer MLP up to a given resolution. An illustration is provided in Figure 2. To state our result, we assume access to an oracle that performs linear regression and consider the consider the condition number of the data matrix.
Definition 1 (MLP and condition number) Consider a multilayer MLP defined by the recursion  $\pmb{s}_i^\ell = \phi(\pmb{W}_\ell \pmb{s}_i^{\ell-1})$  for  $\ell \in [L]$ ,  $i \in [n]$  and  $\pmb{s}_i^0 = \pmb{x}_i$ . Here  $\phi(x) := \max(\alpha x, x)$  is a Leaky-ReLU activation with  $1 \geq \alpha > 0$ . Define the feature matrix  $\pmb{T}_\ell = [\pmb{s}_1^\ell \dots \pmb{s}_n^\ell]^\top$  and define its condition number  $\kappa_\ell = \sigma_{\max}(\pmb{T}_\ell) / \sigma_{\min}(\pmb{T}_\ell)$  (with  $\sigma_{\min} := 0$  for fat matrices) and  $\kappa_{\max} = \max_{0 \leq \ell < L} \kappa_\ell$ .
Assumption 1 (Oracle Model) We assume access to a transformer  $TF_{LR}$  which can run  $T$  steps of gradient descent on the quadratic loss  $\mathcal{L}(\boldsymbol{w}) = \sum_{i=1}^{n} (y_i - \boldsymbol{w}^\top \boldsymbol{x}_i)^2$  given a prompt of the form  $(\boldsymbol{x}_1, y_1, \ldots, \boldsymbol{x}_n, y_n)$ .
We remark that this assumption is realistic and has been formally established by earlier work [Giannou et al., 2023, Akyurek et al., 2022]. Our CoT abstraction builds on these to demonstrate that CoT-I/O can call a blackbox TF model to implement a compositional function when combined with filtering.
We now present our main theoretical contribution. Our result provides a transformer construction that first filters a particular MLP layer from the prompt through the attention mechanism, then applies in-context learning, and repeats this procedure to approximate the MLP output. The precise statement is deferred to the supplementary material.
Theorem 1 (CoT  $\Leftrightarrow$  Filtering+ICL) Consider a CoT prompt  $p_n(f)$  generated from an  $L$ -layer MLP  $f(\cdot)$  as described in Definition 1, and assume given test example  $(\boldsymbol{x}_{test},\boldsymbol{s}_{test}^{1},\dots \boldsymbol{s}_{test}^{L})$ . For any resolution  $\epsilon >0$ , there exists  $\delta = \delta (\epsilon)$ , iteration choice  $T = \mathcal{O}(\kappa_{\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{s}^i)_{i = 0}^{\ell -1}$  denote the first  $\ell -1$  CoT-I/O outputs of TF where  $\hat{s}^0 = \boldsymbol{x}_{test}$  and set  $\boldsymbol {p}[\ell ] = (\boldsymbol {p}_n(f),\boldsymbol{x}_{test},\hat{s}^1\dots \hat{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} \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).
$$
There exists a fixed projection matrix  $\Pi$  that applies individually on tokens such that the backend output obeys  $\| \Pi (TF_{BE}(\pmb {p}[\ell])) - \pmb{p}_{n}^{\text{filter}}\| \leq \delta$ .
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  $n$  and  $T$ . Consequently, after  $L$  rounds of CoT-I/O,  $TF$  outputs  $f(\boldsymbol{x}_{test})$  up to  $\epsilon$  accuracy.
(a) CoT-I/O:  $d = 10$  vs 20
(b) Alignment in Fig. 3(a)
(a) CoT-I/O: composed risk  $(d = 10)$
Figure 3: Solving 2-layer MLPs with varying input dimension  $d$  and hidden neuron size  $k$ .
(b) Risk of first layer
(c) Risk of second layer
Figure 4: We decouple the composed risk of predicting 2-layer MLPs into risks of individual layers.
Remark 1 Note that, this result effectively shows that, with a sufficiently good blackbox transformer  $TF_{LR}$  (per Assumption 1), CoT-I/O can learn an  $L$ -layer MLP using in-context sample size  $n > \max_{\ell \in [L]} d_{\ell}$  where  $d_{\ell}$  is the input dimension of  $\ell$ th layer. This is assuming condition number  $\kappa_{\max}$  of the problem is finite as soon as all layers become over-determined. Consequently, CoT-I/O needs  $\max(k, d)$  sample complexity to learn a two-layer MLP. This provides a formal justification for the observation that empirical CoT-I/O performance is agnostic to  $k$  as long as  $k \leq d$ .
We provide the concrete filtering statements based on the transformer architecture in Appendix A, and the key components of our construction are the following: (i) Inputs are projected through the embedding layer in which a set of encodings, an enumeration of the tokens  $(1,2,\dots ,N)$ , an enumeration of the layers  $(1,2,\ldots ,L)$  and an identifier for each layer already predicted are all attached. Notice that this "modification" to the input only depends on the sequence length and is agnostic to the token to-be-predicted. This allows for an automated looping over  $L$  predictions. (ii) We use this information to extract the sequence length  $N$  and the current layer  $\ell$  to-be-predicted. (iii) With these at hand, we construct an 'if-then' type of function using the ReLU layers to filter out the samples that are not needed for the prediction.
4 Experiments with 2-layer Random MLPs
For a clear exposition, we initially focus on two-layer MLPs, which represent 2-step tasks (e.g.,  $L = 2$ ). We begin by validating Theorem 1 using the CoT-I/O method, demonstrating that in-context learning for a 2-layer MLP with  $d$  input dimensions and  $k$  hidden neurons requires  $O(\max(d, k))$  samples. The results are presented in Section 4.1. Subsequently, in Section 4.2, we compare three different methods: ICL, CoT-I, and CoT-I/O. The empirical evidence highlights the advantages of CoT-I/O, showcasing its ability to reduce sample complexity and enhance model expressivity.