Dissecting Chain-of-Thought: Compositionality through In-Context Filtering and Learning
Authors: Yingcong Li
Kartik Sreenivasan
Angeliki Giannou
Dimitris Papailiopoulos
Samet Oymak
Venue: NeurIPS 2023 poster (2023)
Keywords: chain-of-thought
in-context learning
attention
compositional learning
approximation
length generalization

Abstract: 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.

Dissecting Chain-of-Thoughe Compositionality through In-Context Filtering and Learning
Yingcong Li University of Californa, Riverside yli6920ucr. edu
Kartik Sreenivasan
Angeliki Giannou
University of Miscoonsin-Madison edu
Dimitris Papailiopoulos University of Wiicoosinn-Madioon dimitris@papail..i
Samet Oymak
University of Michigan & UC Riverside oymakšumi ch. edu
Abstract
Chain-of-thought (CoT) is method that enables language models handle com- plex reasoning tasks by decomposing them into simpler steps. Despite its success, the underlying mechanics of CoT are not yet fully understood. In attempt to shed light on this, our study investigates the impact of CoT on the ability of transformers to in-context learn 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 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, 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.
Introduction
The advent of transformers [Vaswani et al., 2017] has revolutionized natural language processing, paving the way for remarkable performance in wide array of tasks. LLMs, such as GPTss [Brown et al., 2020], have demonstrated an unparalleled ability to capture and leverage vast amounts data, thereby facilitating near human-level performance across variety of language generation tasks. Despite this success, deep understanding of their underlying mechanisms remains elusive.
Chain-of-thoughe prompting [Wei et al., 2022c] is 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 seneral-purposs 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 demystify some of the mechanics underlying this success and the benefits of CoT in terms sample complexity and approximation power. To
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Input:
O...O...
ICL
CoT-1/O
hidden:
...
12 Ples
o o 11
hidden
Chipul: y
0.4
Co
0.3
CoTi
CoT-I
Test risk
#1s U1 21
Rest
4
x²
0.2
0.1
0.0
60 80 100
in-context samples
Figure 1: An illustration of ICL, CoT-II and CoT-I/O methods, using 3-layer MLP as an example (top left, where æ, y, s', s² denote input, output and hidden features respectivaly). The ICL method utilizes in-context examples in the form of (z, y) and makes predictions directly based the provided Test- Both CoT-II and CoTT-I/O methods admit prompts with samples formed by (æ, s s² 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 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 experimentaall-supeen suppoored abstraction that decouples CoT prompting into a filtering phase and an in-context learning (ICL) phase. In filtering, the model attends the relevant tokens within the prompt based on an instruction. In ICL, the model runs inference on the filtered prompt to output step and then moves onto the next step the chain. Our Theorem develops theoretical understanding of this two-step procedure and formalizes how filtering and ICL phases of CoT can implemented via the self-attention mechanism to learn MIPs.
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 & neurons using (max(k:: d)) in-context samples by filtering individual layers and solving them via linear regression in contrast to the (2kd) lower bound without -uggeetedd prompt. As predicted by theory, our experiments (see Sec. 4) identify a striking universality phenomenon (as 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 single demonstration. Notably the pretraining loss goes to zero stpteb where each step "learns filter a layer". Together, these showcase how CoT identifies composable shortcuts to avoid the need for solving linear regression. In contrast, show that ICL (without CoTT collapses to linear regression performance 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 ICI. Section 4 provides empirical investigations of CoT with 2-layer MLPs, which validates our theoretical findings. Finally,

---

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. ... as [n). Vectors and matrices are represented in bold text (e.g., z, A), while scalars are denoted in plain text ge.g., y). The input and output domains are symbolized as X' and Y respectively (unless specified otherwise), and a E X, y E 3 denote the input and output. Additionally, let F be a set of functions from X' to V. Consider a transition function F where y f(x). In this section, we explore the properties of learning f. assuming that it can be decomposed into simpler functions (fe) 1. and thus can be expressed as f = fL o ooof fL-10-... ft.
2.1 In-context Learning
Following the study by Garg et al. [2022], the fundamental problem of vanilla in in-context learning (ICL. involves constructing prompt with input-output pairs the following manner.
Pr (f) (z; y:)i where y: f(x:).
(P-ICL)
Here the transition function f E X + 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 p (S)annish) (lapple, manzana), (ball. pelota). ...) p (French)=((cat, chat), (flower, fleur).... Let TF denote any auto-regressive model (e.g., Decoder-only Transformer). The aim of in-context learning is learn model that can accurately predict the output, given prompt and the test input Teas, as shown in the following equation:
TF (f), Drea) R3 (
(2.1)
where f E F is the test function which may differ from the functions used during training. Previous work [Zhou et al., 2022, Li al., 2023c] has demonstrated that longer prompts (containing more examples n)) typically enhance the performance of the model.
2.2
Chain-of-thoughh 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 E F in one go, which becomes more challenging as 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 offin-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 T. subfunctions, represented as f f= fr. Jt-1 o... fi- Fach intermediate output can be viewed as step, enabling us to define a length-n CoT prompt related to with L, steps (expressed with s l (L.) as follows:
Pn (f) = =xx,, 8,, s
J.-1 Lyn s 1= where of = = le [L].
(P-CoT)
Here I; = s,. y, = s and fe E Fe. which implies that f FL X :
Next we introduce two methodologies for making predictions within the CoT framework:
CoT over input only (CoT-I). Contrasted with ICI., CoT-I considers step-by-step instructions as inputs. nonetheless, the prediction for the last token is performed as single entity. Our experiments indicate that this approach lowers the sample complexity for TF to comprehend the function f being learned (see the orange curves Figures and 6). The CoT-I prediction aligns with Eq. (2.1) as follows, while the prompt is determined by (P-CoT).
One-shot prediction: TF(p.(f). Drea) Ri f(
(2.2)
CoT over both input and output (CoT-I/O). Despite the fact that CoT-I improves the sample complexity of learning f. the TF must still possess the capacity to approximate functions from the function class F. given that the prediction is made in shot. To mitigate this challenge, we consider
3

---

scenario where in addition to implementing a CoT prompt, we also carry out CoT predictions. Specifically, for composed problem with inputs formed via (P-CoT), the model recurrently makes T-step predictions as outlined below:
Step 1: TF(Pn Drea) ==
Step 2: TF (F (f), Tress 8+) )+ = 8²
Setp L: (F(,,. Tuea Rusa, 8² .8 22(2f (I
(2.3)
where at each step (step E), the model outputs intermediate step (&*) which is then fed back the input sequence to facilitate the next-step prediction ($*+1). Following this strategy, the model only needs to learn the union of the sub-function sets. U Fr. whose complexity scales linearly with the number of steps Empirical evidence of the benefits of CoT-I/O over ICL and CoT-II 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 and Section 2, we have discussed vanilla ICL, CoT-I and CoT-I/O methods. Intuitively, ICL can be viewed as special case of CoTTII (or CoT--//O) if we assume only one step is performed. Consequently, we will focus on implementing CoT-I and CoT-I/O for model training the following.
Consider the CoT prompt as in (P-CoTT), and assume that x ~ Dx, and ft ~ Dr, l E [4]. where I. denotes the number of compositions/stepp, such that the final prediction should approximate f(x)) = (J (x)) == y E y. We define {(j y) y x y - as a loss function. For simplicity, we assume fe(... fi (xx) E V. l E [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 ensure high prediction performance given any length-n prompt, where n [.V. To this end, we train the model using prompts with length from to N equally and aim to minimize the averaged risk over different prompt size. Assuming the model TF parameterized by 0 and considering meta learning problem, the objective functions for CoTT-I and CoT-I/O are defined as follows.
N E/nn f(xn)) ((Zn))
@COT- = arg min ((o.)
N n=1
and
N
L
@C = arg min E (
Σ
((s n. )
NL
where dn (F(p,.)). vn) and 8 TF(P:) In ..... sf
(f) is given by (P-CoT), and as mentioned previously, d x and s y. All x and fe 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 exccute filtering over the CoT prompt, thereby learning a 2-layer MLP with input dimension of d and hidden dimension of k. akin to resolving & d-dimensional ReLU problems and 1 k-dimensional linear regression problem. Subsequently, in Section 4. we examine the performance of CoT-I/O when learning 2-layer random MLPs. Our experiments indicate that CoT-I/O needs only O (mmax(d, k)) in-context samples to leam the corresponding MLP.
We state our main contribution of establishing a result that decouples CoT-based in-context learning (Co'T-I/O) into two phases: (1) Filtering Phase: Given prompt that contains features of multiple MIP layers, retrieve only the features related to a target layer to create an ICL. prompt. (2) ICI. Phase: Given filtered prompt, learn the target layer weights through gradient descent. Combining these two

---

31
122
Peest
Step ICL:
m1 #1 of 21 a2
Filtering
Step ICL:
CoT prompt
Step 3 ICL:
Prediction
Figure 2: Depiction of Theorem 1 that formalizes CoT Filtering ICL procedure. In this illustration, we utilize in-context samples as depicted in Figure 1, characterizing the CoT prompt by (æ, 82 s² y). We show transformer TF, composed of TF La TFaE. TF facilitates the imple- mentation of filtering on the CoT prompt (refer to (P-CoT), illustrated on the left side of the figure), subsequently generating vanilla ICL prompts (refer (P-ICL), illustrated on the right side of the figure). During each step of the ICL process, LF employs gradient descent techniques to address and solve sub-pu 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 multilayer MLP up given resolution. An illustration is provided Figure 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 (MLP and condition number) Consider a multilayer MLP defined by the recursion for activation with 1 2 a O. Define the feature matrix T, = [ ... " and define condition number Tmax (onax(Te))) (Te)/oma T (Ti) (with min == for fat matrices) and Kmax maxogi<L Ke.
of = @( -/11 ) & E [L], i E [n] and sf = D; Here Φ(x) == max( x) is a Leaky-RelL
Assumption (Oracle Model) We assume access t0 transformer TF which can run I steps of gradient descent on the quadratic loss L(w) Em= (1: w xx)² given prompt of the form (æ1. V1,
Rn: ).
We remark that this assumption is realistic and has been formally established by earlier work [Giannou al., 2023, Akyürek et al., 2022]. Our CoT abstraction builds these 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 transformer construction that first filters particular MLP layer from the prompt through the attention mechanism, then applics in-context learning, and repeats this procedure to approximate the MLP output. The precise statement deferred to the supplementary material.
Consider CoT prompt Pn. (f) generated from L-layer MLP f(-) as described Definition 1, and assume given test example (Resrs 5 B(.)). ) For any resolution O, there exists 8 ô(e), iteration choice T = ( maax lo)(1/e)), and backend transformer construction TF such that the concatenated transformer TF TF TF implements the following: Let (8)) denote the first f. - CoT-1/O outputs of where 8
Theorem (CoT=Fitterig==--ee
Brss and set
pl(] (Pn(f), Frer 8). &f. At step f, TF implements
1. Filtering. Define the filtered prompt with input/output features of layer l,
p
fiter
0,
0,
5 O,
0
0
8
0,
0,
8
0,
8
O,
0
There exists a fixed projection matrix II that applies individually tokens such that the backend output obeys (IL(FF=e(P())) - pilter < 8.
2. Gradient descent. The combined model obeys || [F(p(() s S Se. e/L.
TFas has constant number of layers independent of n and T. Consequently, after rounds of CoT-1/O.
TF outputs f (x up accuracy.
5

---

d= 10
d= 10
0.3
Test risk
4=4
Test risk
0.2
0.2
16
0.1
0.1
0.0
100
0.0
in-context samples
in-context samples/d
(a) CoT-I/O us 20
(b) Alignment in Fig. 3(a)
Figure Solving 2-layer MLPs with varying input dimension d and hidden neuron size k.
0.4
0.4
0.4
layer
0.3
Second layer risk
0.3
0.3
Test risk
0.2
k 16
16
16
k 64
k 64
K 64
0.1
0.1
0.1
0.0
20
40
60
80
100
0.01
20
40
60
80
100
0.0
20
40
60
BD
100
in-context samples
in-context samples
in-context samples
(a) CoT-I/O: composed risk (d
10)
(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 Note that, this result effectively shows that, with a sufficiently good blackbox transformer TFLA (per Assumption 1), CoT-I/O can learn an L-layer MLP using in-context sample size >> maX de where de is the input dimension of Eth layer. This is assuming condition number Emax of the problem is finite as soon as all layers become ver-ddeteminned. Consequently, CoT-1/O needs max(k, d) sample complexity to learn a two layer MLP. This provides formal justification for the observation that empirical CoT-1/O performance is agnostic to k as long as k S d.
We provide the conerete 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 21, N), an enumeration of the layers (1,2, L) and 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 automated looping over L predictions. (ii) We use this information to extract the sequence length N and the current layer to-be-pree bo-be-predicted. (iii) With these at hand, construct 'if-then' type of function using the ReLU layers to filter out the samples that are not needed for the prediction.
Experiments with 2-layer Random MIPs
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 using the CoT-I/O method, demonstrating that in-context learning for 2-layer MLP with d input dimensions and l hidden neurons requires O Ommxx(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.
Dataset. Consider 2-layer MLPs with input æ E Rd, hidden feature (step-1 I output) s E Rt and output y R. Here, fi (x: == (Wx)+ and y = f2(s) == v s, with W E Rkxd and R* being the parameters of the first and second layer, and (x)+ max(x; 0) being ReLU activation. The function is composed as y = v (W>).. We define the function distributions as follows: each entry of W is sampled via Wij (0, )) and AO0 4x), with inputs being randomly sampled through (0, I.)'. - We apply the quadratic loss our experiments. To avoid the implicit bias due to distribution shift, both training and test datasets are generated following the same strategy.
FFllowing this strategy for data generation, the expected norms of æ, s and are equivalent, and the risk curves displayed in the figures are normalized for comparison.

---

0.3
Standaro Small Tiny
Averaged risk
KL
ICL
JCL
0.2
Co
contaat sampless
Ininontext samples
inicontext samples
(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.
0.4
Averaged risk
0..1
KL
0.3
ICL
ICL
0.2
0.30
0.05
80
In-oontext samples
In-contoxt samples
(a) Averaged risk
(h)
(c) A
(d) = 16
Figure 6: Comparison of the three methods for solving 2-layer MLPs with different hidden sizes.
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 and 4, and all experiments utilize small GPT-2 models for training".
CoT-I/O performance is agnostic to k when k v d (Figure 3). In Fig. 3(a), train MLPs with = 10, 20 and = 4, 8, 16. Solid and dashed curves represent the CoT-I/O test risk of = 10 and 20 respectively for varying in-context samples. The results indicate that increase in will amplifies the number of samples needed for in in-context learning, while the performance remains unaffected by changes in (4, 8, 16). To further scrutinize the impact of 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, 16 (blue solid) and 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 dominates CoT-I/O performance (Figure 4). We further investigate the cireumstances under which & 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 the second layer requires more hidden features (s), thus N = 100 in-context samples (providing 100 8s) are insufficient to fully restore the second layer, leading to performance gaps between & = 4, 16 and & 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., (f) be test prompt containing n in-context samples where represents any arbitrary 2-layer MLP. Given test sample (Ereas Seas Musa). ) the layer predictions are performed as follows.
64 larger, learning
layer prediction: TF(p, (f), B.n)) == 8,
2nd layer prediction: TF TF (.f) Brest: Steat :
The risks are calculated by ||s - Sua and (i 3 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 & does not affect the first layer's prediction, does augment the number of samples required to learn the second layer. Moreover, by plotting the first layer risks of 4, 16 (blue/orange dotted) and second layer risk of k: = (green dashed) in Fig. 4(a), we can see that they align with the CoT--1/O composed risks. This substantiates the hypothesis that CoT--I/O learns layer ML. through compositional learning of separate layers.
Our code is available https: http::///tttttooooooo/oggggggttt htwpa//oooon--/iiiittiiiiiicciai
7

---

1.0
n=d
Test nsk
CaT-2
risk
Test Tiskrisk
ICL
CoT-3
3a
CoT-
Least squares
100k 200k 300k
400k 500k
#f in-context samples
# iteration
(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 Xtstep CoT-I/O. Fig. 7(a) illustrates point-to-poiee 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 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 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-II and ICI., respectively. results show that using CoTT-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. unable achieve zero test risk even with the standard GPT-2 model and up to 100 samples. This indicates that to learn 2-layer MLPs single shot, ICL requires at least (dk + d) samples to restore all function parameters. Conversely, CoT-I and CoT-1/O can leverage implicit samples contained in the CoT prompt. Let fi E Fi (first layer) and f2 E F2 (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 fi E F, which the small GPT-2 cannot express.
Fig. 5(a), we present the averaged risks. The
Varying MLP widths (Figure 6). Next, we explore how different MLP widths impact the performance (by varying the hidden neuron size k E {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 i- 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-1/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 them. Moreover, we observe a substantial difference in the performance of ICL and CoT-I with varying (e.g.. averaged risks in Fig. 6(a)). This be explained by the fact that enlarging results in more complex Fi and F2, thus making the learning of F2 x Fi more challenging for ICL and CoT-I.
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 refer CoT-I/O in the rest of the discussion.
Dataset. Consider L-layer linear MLPs with input æ E Rd ~ N(0 Ia), and output generated by y W- Wi where the /th layer is parameterized by W. Rdxd E [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. f (W W,T W I, k [K]].. Therefore, to learn the T.-layer neural net, CoT only needs to leam FF with |F K, whereas ICL needs learn the function set F4, which contains K- random matrices.
'This assumption ensures that the norm of the feature remains constant across layers ||>||| = |g enabling fair evaluation across different layers.
||s

---

1.0
Leyer
Lawe2
2.
6.3
0.8
3.6
ICL
CoT-2
Test
CoT-3
236
23c
264
22.
160
Laper
Lgw
Layors
CoT-6
6.5
0.0
20k
30k
50k
# iteration
leerstinn
(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 decouple the composed risks from predicting 6-1ayer 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). randomly select L matrices We, E [L] from I so that for any input æ, we can form chain
Set d = 5, L 6 and K = 4. At each round, we
x + s + 5²
where sf = Wisf-1, [L] and s == x. Let CoT-X denote X Xteppp CoT-I/O method. For example, the in-context sample of CoT-6 has form of (ax, s s? ...... S², y), which contains all the intermediate outputs from each layer; while CoT-3, CoT-2 have prompt samples formed as (æ, s² s y) and (xr, 8", 3) respectively. pairs. To solve the length-6 chain, CoT-X needs to learn a model that can remember 46/X matrices. Therefore, ICI. is face significantly challenge sined it needs to remember 4 4, 096 matrices (all combinations of the matrices used for training and testing) compared to just 4 for CoT-6.
this setting, ICL is also termed as CoT-1, as its prompt contains only (æ. y)
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 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 = B computed directly via least squares given random training samples, where æ is the input and is from the output of the 6-layer MI.ss (e.g., y (0). The test risks for = 1.... 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 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) 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 identify risk drops when predicting different layers, which appear at either 15k (for layer 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 tth layer by filling irrelevant positions with random features. Specifically, in-context example is formed by
s²). From these figures, we can
(z
where s We(s*=1 and Ia).
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 the prefix, while disregarding irrelevant features. This observation provides evidence that the process of filtering is indeed performed.

---

6 Related Work
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, complex reasoning by decomposing it into step-by -t-pppp 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, Veličković 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 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 Icarnability 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.
thain-off-thughht (CoT) prompts have made impressive improvements in performing
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 corpleeeomesssness 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 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 -Gne-Genrrliization Conjecture which says that Transformers tend to length generalize 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 thaingtf-thought prompting and shed light on how it enables composi- tional 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 theoreti- cally demonstrated that to learn 2-layer MLP with d-dimensional input and k neurons, CoT-I/O requires O O(max(d, (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, 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.I, 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 practical problems such as code generation and mathematical reasoning? We have shown that CoT-1/O can rely on linear regression oracle to learn an MI.. To what extent can transformers approximate MLPs without CoT-I/O (c.g. with C-T-I) and what are the corresponding lower/upper bounds?
10

---

8 Acknowledgements
This work was supported by the NSF CAREER awards #2046816 and #1844951, NSF grant #2212426. AFOSR & AFRL Center of Excellence Award 199500-18-10166 a Google Research Scholar award, an Adobe Data Science Research award, an Army Research Office grant W911NF21 10312, and an ONR Grant No. N00014-21-1-2806.
References
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. arXiv preprint arXiv:2306. N.XY::2306.00297,, 2023.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algo- rithm is in-context learning? investigations with linear models. arXiv preprint arXiv: 156,1 2022.
Haldun Balim, Zhe Du, Samet Oymak, and Necmiye Ozay. Can transformers learn optimal filtering for unknown systems? arXiv preprint wrXi:2238.0855.6 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 meta optimizers. arXiv preprint arXiv:2212. 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, al. Faith and fate: Limits of transformers on compositionality. arXiv preprint 2..iv:2335..86665 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 Pa- pailiopoulos. Looped transformers as programmable computers. arXiv preprint arXiv:2301.13190 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 2707:7330..777,,, 2023.
Shima Imani, Liang Du, and Harsh Shrivastava. Mathprompter: Mathematical reasoning using large language models. arXiv preprint arXiv:2303.05330 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 preprini 4rXiv: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.008: 2023.
Nayoung Lee, Kartik Sreenivasan, Jason D Lee, Kangwook Lee, and Dimitris Papailiopoulos. Teaching arithmetic to small transformers. arXiv preprint wxXiv2207.00388. 2023.
Brian Lester, Rami AI-Rfou, and Noah Constant. The power of scale for parameter-efficient prompt tuning, arXiv preprint (tt.::104400991) 2021.

---

Hongkang Li, Meng Wang, Sijia Liu, and Pin- Yu Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. arXiv preprint arXiv:2302.06015. 2023a.
Liunian Harold Li, Jack Hessel, Youngjae Yu, Xiang Ren, Kai- Chang, and Yejin Choi. Sym- bolic chain-of- thought distillation: Small models can also" think" step-by-step arXiv preprint arXiv:2306. 14050, 2023b.
Yingcong Li, M Emrullah Iidiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in in-context learning. International Conference on Machine Learning, 2023c.
Adam Liška, Germán Kruszewski, and Marco Baroni. Memorize or generalize? searching for compositional mn in a haystack. arXiv preprint /rXiv:/802.066 2018.
Bingbin Liu, Jordan Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. arXiv preprini arXiv:2210. 10749, 2022.
Xinxi Lyu, Sewon Min, Iz Beltagy, Luke Zettlemoyer, and Hannaneh Hajishirzi. Z-icl: Zero-shot in-context learning with pseudo-demoosstaatonn... arXiv preprint arXiv:2212.0986, 2022.
Sewon Min, Mike Lewis, Luke Zettlemoyer, and Hannaneh Hajishirzi. Metaicl: Learning to learn in context. arXiv preprint arXiv::X 10. 2021.
Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. Rethinking the role of demonstrations: What makes in-context learning work? arXiv preprint arXiv:2202. 2022.
Sharan Narang, Colin Raffel, Katherine Lee, Adam Roberts, Noah Fiedel, and Karishma Malkan. Wt5?! training text-to-text models to explain their predictions. arXiv preprint arXiv:2004. 14546, 2020.
Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, et al. Show your work: Scratchpads for intermediate computation with language models. arXiv preprint arXiv:2112.00114 2021.
Jorge Pérez, Javier Marinković, and Pablo Barceló. On the turing completeness of modern neural network architectures. arXiv preprint urXiv:1901.033 2019.
Jorge Pérez, Pablo Barceló, and Javier Marinkovic. Attention is turing complete. The Journal of Machine Learning Research, 22(1):3463-3497 2021.
Luis Perez, Lizi Ottens, and Sudharshan Viswanathan. Automatic code generation using pre-trained language models. arXiv preprint 10Xiv:2102.1055:: 2021.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Himanshu Sahni, Saurabh Kumar, Farhan Tejani, and Charles Ishell. Learning to compose skills. arXiv preprint arXiv: /7.1. 1289, 2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
Petar Veličković and Charles Blundell. Neural algorithmic reasoning. Patterns, 2(7): 100273, 2021.
Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. arXiv preprint arXiv:221..0767., 2022.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference Machine Learning, pages 35151-35174, PMLR, 2023.
12

---

Johannes von Oswald. Eyvind Niklasson, Maximilian Schlegel, Seijin Kobayashi, Nicolas Zucchet, Nino Scherrer, Nolan Miller, Mark Sandler, Max Vladymyrov, Razvan Pascanu, al. Uncovering mesa-optimization algorithms in transformers. arXiv preprint wrxiv:2309005858, 2023.
Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: case study on approximating turing machines with transformers. Advances in Neural Information Processing Systems, 15:12071-1208 2022a.
Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al. Emergent abilities of large language models. 2022b.
arXiv preprini
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Ed Chi, Quoc Le, and Denny Zhou, Chain of thought prompting elicits reasoning in large language models. arXiv preprint arXiv:2201.11903, 2022c.
Jerry Wei, Jason Wei, Yi Tay, Dustin Tran, Albert Webson, Yifeng Lu, Xinyun Chen, Hanxiao Liu, Da Huang, Denny Zhou, et al. Larger language models do in-context learning differently, arXiv preprint arXiv:2303.03: 03846, 2023,
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi,
Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Huggingface's transformers: State-of-the-att natural language processing. arXiv preprint 1rXiv:1910.03771, 2019.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111. 0200000 2021.
Zheng Yuan, Hongyi Yuan, Chuanqi Tan, Wei Wang, and Songfang Huang. How well do large language models perform in arithmetic tasks? arXiv preprint arXiv:2304. 2xXi:22344002055, 2023.
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequenee functions? arXiv preprint arXiv:1912.10077, 2019.
Xiao Zhang, Yaodong Yu, Lingxiao Wang, and Quanquan Gu. Learning one-hidden-layer relu networks via gradient descent. In The 22nd international conference on artificial intelligence and statistics, pages 1524-1534, PMLR, 2019.
Hattie Zhou, Azade Nova, Hugo Larochelle, Aaron Courville, Behnam Neyshabur, and Hanie Sedghi. Teaching algorithmic reasoning via in in-context learning. arXiv preprint 2eXiv:221.09006 2022.
Hattie Zhou, Arwen Bradley, Etai Littwin, Noam Razin, Omid Saremi, Josh Susskind, Samy Bengio, and Preetum Nakkiran. What algorithms can transformers learn? study in length generalization. arXiv preprint arXiv:2310. /6028, 2023.
13

---

Appendix
Table of Contents
A Construction
14
A.1 The Transformer Architecture
14
A.2 Positional Encodings
14
A.3 Constructing Some Useful "Black-box" Functions
A.4 Results on Filtering
B Experimental Details
22
B.I Model Evaluation
22
B.2 Implementation
22
Additional Experimental Results
22
C.I Out-of-distribation Experiments
22
C.2 Further Evidence of Model Expressivity
................................
23
C,3 Filtering Evidence in 2-layer MLPs
24
C.4 Comparison of Filtered CoT with ICL
25
C.5 CoT across Different Sizes of GPT-2
25
C.6 Comparison Transforner Prediction and Linear Regression
26
A Construction
A.1 The Transformer Architecture
For the purpose of this proof, we consider encoder-based transformer architectures and assume that the positional encodings are appended to the input". We also consider that the heads are added and each one has each own key, query and value weight matrices. Formally, we have
attn (x) X
wh X softmax W
(A.1)
TF (F(X)= attn(X) (X) (2WWiarte jataa( X) b, 1,) b21n.
(A.2)
where X E Rdxn, is the number of heads used and f(x) (x)+ is the ReLU activation. We also make use of the temperature X, which is parameter of the softmax. Specifically, softmax(x:) (e>e./ E, eAa J.- Notice that as À + co we have softmax(x) - max; V4. We also assume that the inputs are bounded; we denote with Nmar the maximum sequence length of the model.
Assumption Each entry is bounded by some constant which implies that |X|| S fo some large e' that depends on the maximum sequence length and the width of the transformer.
A.2 Positional Encodings
In the constructions below we use combination of the binary representation of each position, as well as some additional information (e.g. 0-1 bits) as described in the following sections. The binary representations and in general encodings construct, require only logarithmic space with respect
*We note here that in terms of construction adding the encodings appending them the input can be viewed a similar manner. Specifically, we can consider that the up-projecties step projects the input to zero padded vectors, while the encodings are orthogonal to that projection in the sense that the have zero in the non-zen zero coordinates of the input. In that case adding the positional encodings corresponds to appending them to the input.
14

---

to the sequence length. Notice that binary representation of the positions satisfy the following two conditions:
1. Let be the binary representation of the i-th position, then there exists 0 such that r r, r. r, + for all j *
2. There exists layer transformer that can implement the addition of two pointers (see Lemma 2
A.3 Constructing Some Useful "Black-box" Functions
We follow the construction of previous work [Giannou et al., 2023] and use the following individual implementations as they do. We repeat the statements here for convenience. The first lemma is also similar to Akyürek et al. [2022], we however follow slightly different proof.
Lemma A transformer with one layer, two heads, and embedding dimension of O n + d). where d is the input dimension and n is the sequence length, can copy any block of the input to any other block of the input with error arbitrarily small.
Proof. Consider an input of the following form
32
En
X
yn
(A.3)
In
Tn.
where I;, y: E Rd are column vectors for all = 1....,, We present how we can move block of data from the (d 1: 2d, k) position block to the (1 d, 1) position block, meaning move the point to the position of the point x1. is straightforward to generalize the proof that follows to move blocks from (i k.
:' + j') for arbitrary i, j.J". J',
Let in Eq. A.1 W = 0
0] and We = 10
I and W
I and
K
w Q 0
I then we have
r r
r r2
Fk
Fn ry Fn
r2 Pk
12 12
rg Fa
(W =
(A.4)
rg rk
rg F2
r Fk
rg Fn
r. r,
Fn
As A + oo and after the application of the softmax operator the above matrix becomes equal ek ...e,, + cM, where e; is the one-hot vector with 1 in the k-th position, and is controllable by the temperature parameter and can be arbitrary small. Let finally
lek e2
|M <
0
W
we get that
0
0
W X softmax (WxX)((W(x)) =
(W!
22
En 0
+
+cM
(A.5)
0
I
0
0
0
By repeating the exact same steps for the second head and letting W adding back the residual get the desired result.
and
0
0
0
slightly different implementation be found in Giannou al. [2023], in which the ReLU layers are also used, together with indicator vectors that define whether a position should be updated or not.
Lemma There exists a 1-hidden layer feedforward, ReLU network, with 8d activations in the hidden layer and d neurons in the outpul layer that when given two d-dimensional binary vectors representing two non-negun non-negative integers, can output the binary vector representation of their sum, as long the sum is less than 2d+1
15

---

Lemma 3 Let A E R and B E Raxx. Then for any 0 there exists transformer-basee function block with 2 layers, I head and width O(d) that outputs the multiplication A A'B. cM, for some |M]] 1
Remark Notice that based on the proof of this lemma, the mutrices/scalars/vettoo need to be in the same rows, i.c., Q [A B). By also appending the appropriate binary encodings we can move the output at any specific place we choose as in Lemma 1. Also, following the proof of the paper the imput matrix is
Q
0
0
0
X
I
0
0
r(1)
r2)
where r(),r') are chosen as to specify the position of the result.
A.4 Results Filtering
Lemma Assume that the input to a transformer layer is of the following form
D2
nn
0
0
0
X
(A.6)
b
b2
b
bn.
b/2
by
b
where by, B, E {0, with zero indicating that the corresponding point should be ignored. Then there exists transformer TF consisting only of ReLU layer that performs this filtering, i.e.,
1h * 0)
1 + 0)æ2
1 bn * xn
1 1(b, + 0n
1 1K * 0) b
[b
+ 0)
1
16
#
xn
1 16 7 0)
X)
(A.7)
b2
n
Un
b2
bn
bn
Proof.
The layer is the following:
TF(x;) I; (-Cb; x1)+ - (-Cb, + xz)+
(A.8)
for some large constant C. Notice that if b; 1 the output is just æ, But if b, = 0 then the output is zero. For the second set instead of using the bits b,, we use the b/.
Remark 3 Notice that if some b, is instead of 1 c/C. Then the output of the above layer would be
TF(a..) n4 + (-C (
(A.9)
= 27 while some b, is instead of zero, and assuming that x; > c 0 - the output would be
(A.10)
TF(a.)) = x; + (-Ce -x;)+ (-Cb; + xi.
(A.11)
+ Ce -
(A.12)
Sc
(A.13)
If [xy] S again the output would be less than equal to c, where can be arbitrarily small.
Our target is to create these binary tokens, as to perform the filtering. We now describe for clarity how next word prediction performed. An d x N-dimensional input given to the transformer architecture, which is up-pn up-projected using the embedding layer; the positional encodings are also added/appended to the input. At the last layer, the last token predicted is appended the initial input to the transformer architecture. The only difference of the new input including the positional encodings ( the input for the next iteration is the n + 1-th token. This is a property that our construction maintains, i.e., the positional encodings used are oblivious to the prediction step performed and are
16

---

always the same for each individual token. We consider encoder-based architectures in all of our lemmas below.
In the subsequent lemma we construct an automated process that works along those guidelines. To do assume that the input to the transformer contains the following information:
1. An enumeration of the tokens from 1 to N.
2. The In of the above enumeration.
3. Zeros for the tokens that correspond to the data points, 1 for each token that the Biest or is prediction based on it.
4. An enumeration 1 to L for each one of the data points provided. For example, if given three sets of data we would have : 1... L 1... 1...
5. Some extra information that is needed to implement a multiplication step as described in Lemma and to move things to the correct place.
The above information can be viewed as part of the encodings that are appended the input of the transformer. Formally, we have
Lemma Consider that a prompt with n in-context samples is given and the l 1-th prediction has been made, and the transformer predict the l-th one. Assume the input the transformer is:
e-1
s
82
E-1
s
8/
Eiezt
0
0
0
0
0
0
2+1
++1
L
e
(+1
E+1
L+1
f.
X
2+1
++++11
L+ 2+2
11
n(I 11)
N
In(1)
(n(() In(++ 11 ....h++++))
In ++ ... In(n(I 1))
In(n(I+ +1) In(N)
0
0
0
1
1
1
1
m2
mr
mi+1
mL+0+1
MLL+122
mm(1++)441
my
where ( denote the first l 1 recurrent outputs of TF and for simplicity, let N == n(L ( + denote the total number of tokens. Then there exists a transformer TF consisting of 7 layers that has output
8-1
0
8-1
0
0
8
0
82
0
0
0
X)
{+1 1
L+E+1
L+l+2
n(L+1)
n(L + +1
N
In(1)
In(() In(2 +
In(T+++1) In( +
In(n( (L +
In(n +
In(N)
0
0
0
mi
m.
Mc+1
Mi+E+1
ML+1-2
Mm(L+1)
mn(L+11441
my
with error up 5M, where [M] and 8 > 0 is a constant that controlled and can be arbitrarily small.
Proof.
Step 1: Extract the sequence length (1 layer). Let
WK
0 000010000
W
00000000000000000000000000000000 (A.14)
17

---

and thus
ln(1)
ln
ln(2)
In(2)
(2(2)
(WkX) (Wxx))())))
(A.15)
[n(N)
In(N)
[N)
So after the softmax is applied we have
2
@ffmax((W<F) (WoX)) =
Di=i
Σ
i=
Σ
2i=1"
(A.16)
iml
We then set the weight value matrix as to zero-out all lines except for one line as follows
0
0
W X =
1 2
(A.17)
0
After adding the residual and using extra head where the softmax returns identity matrix and the value weight matrix is minus the identity, we get attention output
s/
A
Fiest
R-1
0
8
8+1
E.
8+1
L+1
E
{+1 1
8+1
L+1
Σ
1
ΣN 1
ΣN
1
Σ
ΣN 1
Σ 1
12
Σ 12
E
Σ
Σ
ΣN
E
Σ i
Σ
I1(1)
In(e)
In(E + 1)
In( + & 11)
In +++2)
In(n(L + +1)) In(n (L +
+11
In(N)
0
0
0
1
1
1
mi
me
mi+1
m
mL+442
(LL+1
Ma(L+1)44
my
2N(N + 12N + 1)
2N +
Notice that
We then use the ReLU layer as to multiply
6N(N + +1)
3
with 3/2 and subtract from this column. This results attention output
$ --1
s/
211
Beest
0
0
0
0
0
0
0
1
8
8 2+1 1
e 1
L 1
1
8
2 2 1
L+1
1
N
N
N
N
N
N
N
In(1) ... In(2) In(E+ 1) ... In(L + & 1) In + E+2) ... In(n(L + 1)) In(a (L + 1) + 1) ... IN(N)
1
1 mi+1
1
m
mL+0+2
mo((+41)41
Notice that this step does not involve any error.
Step 2: Extract the identifier the prediction to be made layers). Now, by setting the key and query weight matrices of the attention to just keep the all ones row, get matrix that attends
18

---

equally to all tokens. Then the value weight matrix keeps the row that contains ones and thus we
l N
get the number
propagated all the sequence length. Then the attention output is follows:
31
2
f-1
82
8
Fesst
0
0
0
0
0
L
L
8
L +
1
1
L+]
&
8+1
L
L+1
N
N
N
N
N
N
N
N
(A.18)
E
e
E
i
N
N
N
N
N
0
0
0
1
1
1
1
m1 my me+1 ML+E-1 ML++42 (L++1) mL(L++)+11
mx
So far this step does not involve any eπrоr. As described in Lemma to implement multiplication of two values up to any error e, we need 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;) and 2) the corresponding structure presented in Lemma 3. So we need to make the following
L
transformation in the rows containing N and
N
N
N
N
N
N
(A.19)
N
N
N
N
N
N
l
(A.20)
(A.21)
where denotes inconsequential values. For the first step we assume that we have the necessary binary representations and then we use Lemma 1 which shows how we can perform the operation of copy/paste. This step involves error and the current output would be X X E M,, where E is controlled and we will determine in the sequence, || S and X is the desired output. For the second step now we use Lemma 3 to perform the multiplication, this will affect some the other values. To analyze the error this step, we know that given X we will get the desired output X; + E2 where M2 I and E2 is to be determined. However, are given X1 the multiplication procedure which results the following output X2 = (X,) X + M2 = xi Xi + é M, Xi + M Mi + M2M Thus by choosing 11, 82 to be of order Ol5)we get that X, X; + 8 & M|] for some C' chosen to be large enough. Then for the last step consider the follow (sub-)rows of the matrix X
F1 F1 F1
F1
(A.22)
ri F3
IN
where is the binary representation of position i. By choosing WK. Wq as to
ra rx]. WaX
Fi
ri
(A.23)
5 'the size that we need will he 2 log Nmxx + 1, where Nmux the maximum sequence length "Note that the true error is even smaller since the operation performed only affects row.
19

---

i
and consider WvX =
we have that
0 0 0
attm x = X + Wv X ssofmaxx (W
WQ
r
r
r r
rir
r2
r/
ri
r/ra
X
softmax
0 0
rJ
T/TN
r.rs]
0
X
0
By subtracting one identity head for the first that we focus on as described in Lemma we have that atm((X (X) results the desired matrix. We output this result and overwrite E/N, thus we have
82
nt
Frest
g--1
0
0
0
0
0
0
0
0
1
+
+
L
l
1
8+1
{+1
L+1
8
N
N
N
N
N
N
N
N
(A.24)
L
i
E
0
0
0
0
0
1
1
1
1
1
m1
Mc
m
ML44+1
M ... Mm(L+1) Mm(L+1)+1
mN
The new error introduced by this operation again controlled as to be M3 C
with || S 1.
Step 3: Create f+ (0 layer). We first copy the with é to the row with N, this trivially be done with a ReLU layer that outputs zero everywhere else except for the row of Ns that output (f)+ (N)+ to account also for the residual. We now use of the bias terms (notice that is always positive) and set to one in one of the two rows that contain the L, again we account for the residual as before; everything else remains unchanged. Thus, we have
8
L
Frsst
-
0
0
0
0
0
0
L
1
8
+ 1 ...
L
1
E
l
8
{+1
L+1
{+1
6+1
6+1
6+1
{+1
l 1
(A.25)
E
8
E
0
1
1
1
1
1
mi
M1.441 MI.+142 ... Mm(L+1) Mm(L+1)+1 ...
This operation can collectively implemented (add the bias + copy the row) in the ReLU layer of the previous transformer layer that was not used in the previous step.
20

---

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 identify which tokens we want to filter. This can be easily implemented with one layer of transformer and especially the ReL.U part of it. Notice that if we subtract the row that contains the tokens that have already been predicted, iee, [e (... e from the row that contains [1 2... L 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 non-zero to set it to be zero. This can be implemented with the following ReLU part of transformer layer
2... L... L. we will get zero only in the positions that we want to
TF (+:) = - (xt)+ + - (-x:)+ + + (et 1)+ + (-x; - 1)+ - ((x+)) (-xx)+)
(A.26)
the last two terms are 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 the input of Assumption the targeted error å and the constant used in I Lemma 4 we achieve that the error has the target upped bound.
Theorem 2 (Theorem 1 restated) Consider a prompt (f) generated from an L-layer MLP f()) as described in Definition I, and assume given test example (Frest Slyg-... s For any resolution e > O, there exists 8 = 8(e), iteration choice T = O(Kmax log(1 (1/e)), and backend transformer construction TF such that the concatenated transformer IF = IF o TF implements the following:
Let (& } denote the first &. CoT-I/O outputs of TF and set pl6] = (fn( Fwas 82 11 At
step e, TF implements
1. Filtering. Define the filtered prompt with input/output features of layer l
filter ... (...0.
0,
8 0, 8-
0, 81
0,
8 0,
There exists fixed projection matrix II that applies individually on tokens such that the backend output obeys (Fss(((()) - p fitter
S 8.
2. Gradient descent. The combined model obeys || TF(p(() s S e. ele/TT
TF has constant number of layers independent Tf and Consequently, after I. rounds of CoT-I/O. TF outputs ((x)s) up to 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 T, and f. Next we turn to prove (2. Gradient descent). In Definition 1, we assume that the network's s activation function is leaky-ReLU,, iee.
i 2 ax, otherwise.
((x)
=
(A.27)
Thus, as a first step we construct the inverse of leaky-ReLU and apply it in the second row of p, filter where the inverse of leaky- Reek- is
(v) = y,
y > 0
(A.28)
Ly/a,
otherwise.
This can be implemented with the following activation function (denoted by o-)) using ReLUs:
(x)) = (x)+ - (11))-1//(())) 1/a(-x)+.
(A.29)
After it, it remains TF to solve linear regression problems. Taking 4th layer, first neuron prediction as example, and letting x : 8 f-1 v/ . - (0) and w Wi[0]. linear regression has form of v/ w a/ for i E [m]. 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 ReL.U. Then following Assumption 1, since we assume TFLa performs the same as gradient descent optimizer,
21

---

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 c/ L, we can get that [P(() s'1 S e/ L. where s = =(W is the correct prediction if taking sf as input. Then we have
[TF (p(() (lml) & [TFLx(D((]) 1'11 + 1 s(on))) S e/ | -(==)
Let La (0) returns Fuss and therefore (p()) T
0. Combing results in that [P[( ا(m)) S L e.cLL. Since from Lemma 5 we have that we can choose 5 be ar- bitrary. Let 5 = L/L, where L is the total predictions we will make. Then will result in [TF()(() (l)ll Il S e.e/L, L, which completes the proof.
B Experimental Details
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
@CoT and @CoT Next we make predictions following Section 2.2. Letting ((:-)) y x y - be loss function, can define test risks follows.
(n) = [(m:)2 [0(2n (an)) where Ú. - TFF(P(()) (
and (CoTH1O(1)) [[(r)) [(( ((x,)) where ÿ.. = TF(P, xn, 8 ....,,8--14
8-,-1; (CoT-1O).
Here, we use [(n) to define the test risk when given prompt with n in-context samples. Then, results shown in Figures 3&&&&667776 (()) are test risks [(n) given n E [N]. Following model training and evaluation, we can see that once loss function is the same for both training and predicting, CoT-II (as
well as ICL) accepts training risk CCOTI N Σ (n.T)(n).
B.2 Implementation
All the transformer experiments use the GPT-2 model [Radford 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 total for all ICI., CoTT-I and CoT-I/O methods. Each iteration randomly samples inputs ass and functions fs. 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., ((y. y) = 9 y (or E(j. = (g- - yy for scalar).
C
Additional Experimental Results
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 quanti- tatively 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, + noise) or
22

---

0.10
0.07
CoT-110
CoT-I
0.06
CoT-I
0.08
COT-1
Test risk
0.3
ICL
ICL
ICL
0.06
0.04
0.2
0.03
0.04
0.02
0.0
0.00
0.0
0.1
0.2
0.3
0.4
0.5
2
5
7
10
Noise level (scaled)
k
d
(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. (()) we add noise the label y (solid) input features æ (dashed), In Fig. 9()), we in-context learn an MLP with & [8] hidden nodes whereas transformer is trained for MLPs with 8 hidden nodes, In
Fig. 9(b), consider misspecification of input dimension: TF is trained with d 10 but we feed neural net with S 10.
0.4
0.4
CoT-I/O
CoT-I/0
0.3
CoT-I
0.3
CoT-I
Test risk
[CL
Test risk
ICL
0.2
0.2
0.1
200 0.0
25
50
100 125 150 175
in-context samples
# in-context samples
(a) Small GPT-2, d 10, A =
(b) Standard GPT-2, d = 10, k =
Figure 10: We train ICL with more in-context examples. The main conclusion is that: For small GPT, ICL indeed not approximate the neural 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.
(a + 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. (())&(cc)
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 MIPs with d 10, k 8. In Fig. 9(),, fix
instead explore
10 and vary l from 1 to 8, whereas in Fig. 9(c), we fix = 8 and vary d from to 10. In both instances, the findings reveal that CoT's performance remains almost consistent when A: > 4 or d 2 6, and ICI. is unable surpass it. The improved performance of ICL with smaller values of d or A again reinforces our central assertion that ICL requires ((dk) samples for in-context learning of the 2-layer random MLP, and reducing either d or & helps in improving the performance. Given that we employ the ReLU activation function, smaller values of d can lead to significant bias 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
There are two determinants of model learnability in-context tasks: sample complexity and model expressivity. Sample complexity pertains to the number of samples needed to precisely solve problem. However, when the transformer model small, even with 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 too small to express even a single layer of 2-layer MI.Ps. Additionally, we have run more experiments, and the results are shown in Figure 10. Both Figures 10(a) and 10(b)
23

---

0.8
0.4
CaT
d= 10
0.6
CoT-zero
0.3
d= 20
Test risk
Co T-random
Test risk
0.4
0.2
k
16
0.0
20
40
60
80
100
0.0
20
40
60
80
100
in-context samples
in-context samples
(a) Test filtering first layer prediction
(b) CoT-1/O: composed risk (same as Fig. 3(a))
Figure 11: Fig. (1(a) presents a filtering evidence of 2-layer MLPs. Given a 2-layer MLP in-context example (x, 8, 8,y), CoT admits (æz 8, y) as test sample; while test samples of CoT-zero and CoT- random are formed by (æ, 8, 0) and (æt, 8, 2) where z ~ (O, d). Fig. 11(b) is directly cloned form Fig. ((a) with error bar for better comparison with results in Fig. 12.
= 10
d= 10
d0
Test risk
Test risk
0.2
0.2
0.1
0.1
0.0
# in-context samples
# in-context samples
(a) CoT-I/O: layer prediction
(b) Train ICL with (ar, s) pairs
0.4
0.4
10
d= 10
Test risk
Test
0.2
0.2
16
16
0.1
0.1
0.0
20
60
80
100
0.0
60 80
100
#
in-context samples
# in-context samples
(c) CoT-I/O: 2nd layer prediction
(d) Train ICL with (s.w) pairs
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. I), which shows the filtered CoT results. In Fig. 12(b)&12(0) method taking (æ, 8) and (s, y) as inputs.
train two additional models using ICL
detail training models with MLP tasks of dimensions = 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. Fig. 10(b), we train a larger model, and the results demonstrate that the standard GPT-2 sufficient to express 2-layer MLP with d 10 and ls 4.
C.3 Filtering Evidence in 2-layer MIPs
Section has demonstrated the occurrence filtering in the linear deep MLPs setting (black dotted curves in Fig. 8(b)). In this section, 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 h: 8. Assume we have model pretrained using CoT-I/O method. As described Section 4.2, during training, the prompt consists of in-context samples in the form of (x, s, y) where (Wx) x)4 and y v s. To investigate filtering,
24

---

0.4
0.4
0.4
0.3
0.3
0.3
Test risk
layer=3
head=2
embedding -64
0.2
layer-6
0.2
# head-4
0.2
# embedding -128
layer=12
head-8
embedding -55
0.0
20
40
60
80
100
0.0
20
40
60
80
100
0.0
20
40
60
80
100
in-context samples
in-context samples
in-context samples
(a) Varying GPT-2 layers
(b) Varying GPT-2 heads
(c) Varying embedding dimensions
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, 2563).
we make three different predictions evaluate the intermediate output, whose test prompts have in-context examples with the following forms:
CoT: (x, s,y), CoT-zero: (æ, s, O). CoT-random: (æ, s,2),
where 2 ~ (O, d). The results are displayed in Figure 1 l1 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 prediction for input æ, will attend only to (æ, 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 interpretation, 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 ICL.
For convenience and easier comparison, we repeat the same results as Fig. 3(a) in Fig. 11(b), where d E {10, 20), & {4,8, 16), and train with 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 (æ, s, y). Given that filtering happens,
make first and second layer predictions following Section 4.1 and results presented in Fig. 12(a) and Fig. 12(c), respectively. These results show the performances of the filtered CoT prompts. Next, need compare the performance with separate ICL. training. To achieve this goal, we train small GPT-2 model using ICL. method with prompt containing (æt, 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), well Fig. 12(c) and 12(d), we observe that after filtering, CoT-i/O achieves similar performance as individually training 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 performance. Therefore in this section, investigate how different components of transformer model affect the resulting performance by run CoT-1/O various GPT-2 architectures.
We maintain the same setting as in Section 4.2, fix and 8, and consider base GPT-2 model (small GPT-2) with 6 attention layers, 4 heads in each layer and 22-dimensional embeddings.
25

---

0.8
1.0
Gradiem descent
GPT-2
0.6
0.8
Least squares
Test risk
0.6
0.4
16
16
Test
64
64
0.2
0.0
0.0
0
20
40
60
80
100
0
20
40
60
80
100
# in-context samples
# in-context samples
(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()), we fix the number 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 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 in this section. Previous work [Giannou et al., 2023, Akyürek 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
T
n arg min W
Σ
(1 (.)))))4 will where n, ~ NO, Ia), V1 = (un* mx)+
for some w² NNO,2I, 2Ia) and is the training sample size. Then the normalized test risks are computed by [(n) Em an (l), x)+ //2 and show point-to-point results for [N] in black dashed curve Fig. 14(a)7. As for the second layer, solve least squares problems as follows
in = Siy where S E paxk S = (W**))+ yli] = x* S[], æ, E NO Ia)
for some W E Rixd ~ // 2/k) and o' ~ NNO, Iz). Here, represents the pseudo-invers operator. Then calculate the normalized test risk of least square solution (given training samples) C(n) = w- in Fig. 14(b) where blue, orange and green dashed curves correspond to solving the problem using different values of & E {4, 16, 64). In this figure, the curves for A 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 the increased complexity of the function set with higher dimensionality (k 64). Learning such complex functions requires larger TF model.
s 2 where
(W*x)+ and y 0. The results presented
"To mitigate the bias introduced by ReLU activation. subtract the mean value during prediction, ie..
[(n) Lw [ll[ x)+ y (La ((ú,, x)+
///d
26


Reviews:
[{"review_id": "xj11eXiZ3R", "reviewer": "Reviewer_mo8c", "summary": "In this paper, the authors explore the mechanics of Chain-of-Thought (CoT), a method that has successfully enabled language models to handle complex reasoning tasks by decomposing them into simpler steps.\nThe study aims to understand the underlying mechanics of CoT by investigating its impact on the ability of transformers to in-context learn a simple yet general family of compositional functions: multi-layer perceptrons (MLPs).\nThe authors reveal that the success of CoT can be attributed to two distinct phases: focusing on data related to each step of the composition and in-context learning the single-step composition function. \nThey provide experimental and theoretical evidence that demonstrates 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.", "questions": "1. Do the conclusions drawn in this study remain consistent in an out-of-distribution setting? For example, what might occur if the MLPs in the test and training sets exhibit notable differences, such as having different widths?", "limitations": "This paper does not explicitly address its limitations. The most significant limitation of the study is the unclear contribution towards answering the question of how the Chain-of-Thought (CoT) method can assist in solving complex reasoning tasks.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "1. This work explores CoT in the learning of MLPs thatis a novel perspective. Such a simplicity can also help to dig the into inner machenism of CoT and ICL.\n2. This work experimentally compare three schemes (ICL, CoT-I, and CoT-I/O), providing valuable insights into their differences and the benefits of CoT prompting.\n3. The paper provides a formalized theorem that explains how a transformer architecture can realize the CoT process for MLPs.", "weaknesses": "1. The Chain-of-Thought (CoT) implementation presented in this study appears to be limited in terms of reasoning steps, which may not be fully consistent with the original motivation behind proposing CoT as a method for complex reasoning tasks.\n\nIn most existing work, CoT has the potential to assist in various complex reasoning tasks as it provides the model with step-by-step guidance on solving the input. \nEach reasoning step should include both the intermediate state and the process through which this state is generated from the previous one. \nThis is typically implemented using natural language descriptions or formal language expressions.\n\nHowever, the CoT implementation in this paper focuses on the intermediate states in MLPs and does not provide information on how each intermediate state is produced. \nBy only presenting intermediate states rather than complete reasoning steps, the CoT explored in this study primarily reflects the capability of state transition rather than reasoning.\n\n2. The study appears to have limited analysis on generalization capabilities, which is a fundamental aspect of in-context learning (ICL).\n\nThis study assumes that there is no distribution shift between the training and test datasets. \nHowever, a fundamental aspect of in-context learning (ICL) is the ability to learn new tasks from in-context examples. \nConsequently, it is crucial to examine the performance in a generalization setting. \nFor example, [1], which shares a similar philosophy with this work, extensively discussed the behavior of ICL on out-of-distribution prompts.\n\n[1] Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed feedback and for recognizing the novelty and value our results contribute, particularly in understanding the inner mechanisms of CoT and ICL.\n\n>**Reasoning steps vs intermediate state:** We appreciate the reviewer's insightful comment on this matter. In the realm of NLP, the reasoning process targeting the intermediate state is intricately interwoven within the language sentences and pretrained LLMs. Our work, however, highlights that the advantage of CoT over ICL lies in the decomposition of complex problems into intermediate steps, rather than merely its competence in semantic comprehension. Though our MLP setting may seem to lack a distinct \"reasoning process\", it does implicitly include reasoning within positional embedding. Inspired by prior work [Garg et al. 2022, Wei et al. 2022], we employ a CoT-based MLP setting, enabling us to precisely and distinctly examine the underlying CoT mechanism.\n> \n>We respectfully disagree with the claim that \"_the CoT explored in this study primarily reflects the capability of state transition rather than reasoning_\". The process of ICL/CoT involves acquiring the ability to perform a novel task through demonstrations. When provided with too few in-context examples, accurate predictions become unattainable, as evidenced by the figures in our paper where fewer samples (smaller x-axis) correspond to larger errors. Thus, CoT method involves not merely state transition but an in-context learning phase to discern and restore the feature generation function.\n\n> **Out-of-distribution evaluations (Fig 2 in the attached file):**  Great point! Our investigation has focused on the standard setting where train and test samples have identical distribution. Regrettably, we did not explore the out-of-distribution aspect of CoT in our original manuscript. Following the reviewer's suggestion, we have commenced new experiments, with preliminary findings shown in the attached Figure 2. In short, **our observations on CoT translate well to several out-of-distribution setting such as incorporating label or feature noise and misspecified input or hidden dimensions**. To elucidate our results and examine the test risk under varying distribution shift levels, we plot test risk evaluated when prompt has 100 in-context examples. In Fig 2(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 $(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 Fig 2(b)(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 2(b), we fix $d=10$ and vary $k$ from $1$ to $8$, whereas in Fig 2(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\\geq6$, 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(d*k)$ samples for in-context learning of the 2-layer random MLP, and reducing either $k$ or $d$ helps in improving the performance. Given that we employ the ReLU activation function, smaller values of $d$ or $f$ can lead to significant bias in the intermediate feature. Consequently, CoT cannot derive substantial benefits from this scenario, resulting in a decline in performance. To sum up, we thank the reviewer for raising this issue and will integrate OOD experiments in our revision.\n\n> **Robustness implications of our work on broader in-context learning** We also briefly discuss how our formalism that decouples CoT into filtering and in-context learning stages has adversarial/distributional robustness implications for in-context learning. The related works on in-context learning such as [Garg et al. NeurIPS'21, Akyurek et al. ICLR'23] focus on the scenario where input prompt contains IID (input, label) pairs. A natural question is: **What if the prompt contains non-IID data, for instance, is in-context learning robust to outlier features?** In this work, we show that transformer can provably *filter a heterogeneous prompt* to obtain an *purified prompt containing IID features* amenable to in-context learning (our Theorem 1). Thus, beyond CoT, our work has implications for outlier/adversarially-robust in-context learning. Our Filtering+ICL formalism has the following interpretation: *Train the transformer with heterogeneous prompts so that it learns how-to-filter. Then, it can implement outlier-robust ICL during inference time*. We will add a discussion on this in the last section.\n\n> **Lack of limitations:** We acknowledge that we should have provided an explicit discussion of limitations. We've now addressed potential limitations in our general response and intend to include a designated subsection in our revised version.\n\nWe extend our gratitude to the reviewer once more for their insightful comments. In response to their feedback, we are actively engaged in enhancing our discussion to elucidate our problem more comprehensively. Additionally, we are delving deeper into the out-of-distribution aspect by conducting further experiments with varying parameters, such as different $(d,k)$ MLPs pretrained models, alternative GPT-2 architectures, diverse types of distribution, and more."}, {"review_id": "lBSIrHpcXb", "reviewer": "Reviewer_8qtM", "summary": "This paper aims to demystify the mechanism lying in the in-context learning (ICL) and chain-of-thought (CoT). It reveals how CoT significantly reduces the sample complexity of ICL. It uses a two-layer MLP, and a backbone GPT-2 model for exploration.   The experimental results reveal some interesting findings, e.g., the in-context samples needed is linearly dependent on the input dimension. The paper also provides theoretical analysis of probable approximation of MLPs via chain-of-thought.", "questions": "- How do you use the intermediate MLPs features as the CoT prompts for the GPT-2? It seems that the features are not tokens to be prompts.\n- Why do you use the input dimension d and hidden size k of a MLP layer as the measures for experiments? Are there any specific reasons to use them as the proxy for some properties of the target task?", "limitations": "No limitation section is found.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper provides both experimental and theoretical evidence to support its claims about how CoT works.\n- The findings about the relations between the in-context samples needed and the input MLP dimension are interesting and insightful.", "weaknesses": "- Lack of clarity: Some parts of the paper may be unclear or difficult to follow, which could make it challenging for readers to understand the key findings and contributions of the study. Some of the concepts are not explained very well in the beginning (e.g, what is compositional learning? What is MLPs in-context), thus making the paper abstruse.\n- The main contribution of the paper is not clear. The paper claims to dissect the mechanism of CoT", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their careful reading and for recognizing the interest and insights in our work. We hope that we have addressed your questions and concerns adequately below.\n\n> **Lack of clarity:** We apologize to the reviewer for any confusion stemming from the unclear aspects of our work. Our commitment to enhancing clarity will be reflected in our planned revisions. While we're unable to make changes during this review stage, our general response outlines our reorganization strategy, and we're eager to consider any additional recommendations. We'll also refine the introduction to better articulate our motivations, main contributions, and notations.\n> * **_Compositional learning_:** It targets a learning method that decomposes a complex problem into several intermediate stages, addressing each sequentially and combining the solutions to render a final prediction. CoT illustrates this method well. Unlike ICL, which requires recovering a function $f$ from data such as $(x_1,y_1,x_2,y_2,...)$ where $y=f(x)$, CoT utilizes prompts of the form $(x_1,s_1,y_1,x_2,s_2,y_2,...)$ where $s=g_1(x)$ and $y=g_2(s)$, and solves the functions $g_1,g_2$ separately. By composing them to obtain $f:=g_2\\circ g_1$, CoT leverages a compositional advantage. This makes the process of solving subfunctions more feasible and sample efficient.\n> * **_MLPs in-context_:** We regret the ambiguity. Our statement that \"_transformers can in-context learn MLPs_\" means that transformers are capable of learning MLP tasks from in-context samples without altering the model weights, which diverges from traditional learning scenarios where model weights are tuned to optimize performance on specific tasks.\n\n> **Contribution is not clear:** We apologize that this does not come through. In this work, our primary aim is to dissect and elucidate the mechanism of CoT in a more accessible and simplified setting. In the context of LLMs, the complexity arising from the vast amount of pretraining data and intricate semantic information makes it challenging to discern why concepts such as \"let's think step by step\" are effective. Our work confronts this challenge by modeling the CoT prompt using random MLPs. Through both empirical and theoretical results, we demonstrate that the transformer's strength in CoT lies in its ability to learn compositional functions and separate out filtering and learning processes. This novel understanding of CoT's underlying mechanism results in significantly improved sample efficiency in both the training and in-context inference stages. We have included an attached Figure 1 to offer a visual representation of these findings, enhancing understanding, and have also provided a more comprehensive explanation in the general response.\n\n> **Tokens vs prompts:** We concur that GPT-2 contains embedding layer that embeds language words into continuous vector. For our MLP-based prompting study, we clarify that we've substituted the GPT-2 embedding layer with up and down linear projections, aligning with Garg et al. This allows tokens to map into the same dimension, and we'll add further implementation details in our experiment section.\n\n> **Why use $d$ and $k$ as proxy for target task difficulty?** We thank the reviewer for this insightful question. To study the compositional learning ability and sample complexity of CoT, in this work, we explore the learning of complex functions $f=g_2\\circ g_1$, which can be decomposed into two less intricate tasks. Using 2-layer MLPs as examples, the intermediate feature is generated via $s=g_1(x):=\\phi(Wx)$ with $W\\in\\mathbb{R}^{k\\times d}$ presenting the weights of the first layer and $\\phi$ as the activation. The second/final layer is functionalized by $y=g_2(s):=v^\\top s$ where $v\\in\\mathbb{R}^k$. CoT prompts are thus framed as $(x,s,y)$. Based on this setting, $(d,k)$ are the hyperparameters that control the complexity of the MLP, and increasing $d$ and $k$ correspondingly heightens the difficulty of recovering the MLP. To solve 2-layer MLP with dimensions $d,k$, ICL indeed necessitates $O(d*k)$ in-context samples, whereas CoT only requires $O(\\max(d,k))$ samples, thanks to its innate compositional learning capacity.\n\nWe sincerely apologize for failing to include the limitations in the current version of our paper. We acknowledge that addressing the potential limitations is vital for a complete understanding of our work, and we have taken steps to include them in our general response. We are also committed to incorporating the reviewer's insightful comments to enhance our paper by adding further clarification and details to highlight and motivate our primary contributions. Thank you for bringing these vital aspects to our attention."}, {"review_id": "Ry1gaJo7gA", "reviewer": "Reviewer_zLZS", "summary": "The paper proposes to study chain-of-thought (prompting) in the setting of learning MLPs. The authors build on top of recent work studying in-context learning linear regression tasks in the light of gradient descent and extent their setting to learning non-linear functions. In order to study chain-of-thought prompting they either provide features / hidden activations from the \"teacher\" MLP the student (CoT-I setting) or make the Transformer produce its own input by \"looping\" (CoT-I/O  setting). \nThey provide theoretical results and empirical evidence that the Transformer is more sample efficient due to CoT by leveraging the given inputs and/or can remember weights coming from the family of teachers to allow for the self-production of features due to looping. \n", "questions": "My main concern is the following: Given that CoT is given the activations of the teacher, the sequence length of these experiments is larger. I think it is quite crucial to contrast the performance of your trained models when given different input data. If I understand correctly, you are hypothesizing that CoT works in LLMs is because the extra data is representative of features of the input data.\n\n3 Naive ablations come to mind: Just train a model on your setup where the extra data that is provided during CoT (s_0, ...,) is either constant, random or actually the same data given multiple times. That would provide more evidence that in your setting indeed the teachers activations are crucial. It could (I doubt it myself) be that just more naive prompts do the trick. ", "limitations": "I think the limitations that the setting is quite constructed and might not have anything to do how CoT works in LLMs could be a bit more strong. ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "I like the abstraction of CoT presented in the paper. The experiments are convincing and supported by the theoretical results. Although this setting is quite simple, it nicely extends the recent studies of gradient descent on linear regression. The results and the line of thinking is very intuitive, I like the ideas and the empirical execution of the paper. ", "weaknesses": "Although I am very familiar with the setup of the paper, I still had problems understanding it. The presentation can be made much clearer. \nPlease work on your presentation and think about how to structure the paper in a clearer and structured way. It is a bit confusing that you have 1.5 empirical sections and the paper is very dense. I think it would benefit the paper if the authors work on restructuring the presentation given the (hopefully available) extra page if accepted. \n\nIt would be nice to clearly explain what is meant with \"learning an MLP\". This differs from classic student teacher frameworks since you are not learning and also don't have to learn (in the CoT-I setting as far as I understand) the weight of the MLPs but solely rely on gradient descent on a linear regression task which acts on the given features in-context. Given these features in-context, it feels quite obvious that the CoT setting is indeed outperforming classic in-context learning. Nevertheless, I think it is still quite interesting. \n\nAlso it would be helpful to clearly state your abstraction/hypotheses that in LLMs the additional data needed to outperform plain ICL are \"features\" of the data motivating your MLP abstraction. This only became clear to me after reading a few times. \n\nIn the CoT-I/O setting, please explain why you are not using MLPs but linear deep nets as the teacher.  Although the presented experiments make sense, the training of Transformers is usually done differently as TFs are not trained to do CoT-I or CoT-I/O. Please comment on this.\n\nThere are a couple of ablations / interesting experiments I would like to see.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed review and encouraging feedback! We hope that we have addressed your questions and concerns adequately below.\n\n> **Lack of clarity:** We are sorry that the reviewer feels this way. Taking the reviewer's suggestion into sincere consideration, we are working on restructuring the paper to improve readability and hope that the new manuscript will be much better structured and easier to read. We have discussed our plan for reorganization in the general response. Basically, we will keep Section 3 pure theory, pack all the 2-layer MLP experiments (current Sections 3.1 and 4.2) together in Section 4, and Section 5 will further introduce our deep linear MLP results. We are glad to hear back if you have any further suggestions and comments.\n\n>**Explain \"Learning an MLP\":** We are encouraged that the reviewer finds our work interesting! In our work, we define \"learning an MLP\" as in-context learning a function whose performance is close to the target MLP. This is indeed as the reviewer comments, that transformers have the ability to implicitly learn the target MLP from the demonstrations without altering the model parameters. We will clarify this further in our updated manuscript. We also appreciate the reviewer's comment on the implicit gradient descent ability of ICL. Although the result that CoT > ICL is intuitive and not surprising, to the best of our knowledge, seldom do works look into its mechanism. Our work introduces a novel framework to explain it as CoT = filtering + ICL. We also provide more discussion regarding this in the general response and visualize it in attached Figure 1. Specifically, CoT prompts are filtered out into different steps of the ICL process, and the model performs gradient descent for each ICL and links the output sequentially. Theoretical and empirical evidence in this paper has been provided to prove it. To reiterate, our goal is to better understand an existing phenomenon rather than propose a new one.\n\n> **Clarify the connection to LLMs:** Thank you for the suggestion. We will definitely add more discussion regarding how our CoT setup can be related to real NLP settings in the revised manuscript.\n\n> **Why using deep linear MLPs instead of standard MLPs?** Great question! We could have indeed used standard MLPs. However, a drawback of ReLU activation is that, even with proper normalization, the feature distribution becomes more and more heavy tailed as we get deeper in the network. This heavy tail impedes the training process of the transformer resulting in very slow and brittle experimentation when depth is large. We opted to use deep linear MLPs as they better preserve the input distribution (especially with random unitary matrices). A middle ground could have been ResNet :) Let us know if you have further questions.\n\n> **TFs are not trained to do CoT:** We would like to remark that the next token prediction loss exactly mimics the CoT-I structure. What is unclear is whether the text in the pretraining data has a compositional structure. We believe that Nye et al. (check ref)'s work on \"take it step by step\" actually indicates that this is true since the pretrained model responds to the \"step by step\" prompt.\n\n> **Ablation experiments (attached Fig 4):** In attached Figure 4, we run the ICL experiment based on the settings the reviewer suggests. Due to time limitations, we are still working on the scenario where intermediate space is filled with repeatting data. Here, we train a small GPT-2 with $d=10$ and $k=8$. Blue, orange, and green curves show the evaluation results using the methods introduced in the paper. Here, we also try a variant setting where the prompt is provided in the form $(x,[?],y)$, and solid, dashed black curves show the results where the padded intermediate feature [?] is random and constant $-1$, respectively. We choose $-1$ as the constant number since the correct feature is outputed after ReLU and is non-negative. Results have shown that without given meaningfully intermediate information, ICL finds it hard to target the correct task functions.\n\nFinally, we apologize for missing the limitations discussion, and we agree with the reviewer that this work lacks an explanation of the connection between our MLP-based setting and LLM applications. We will further modify our paper to articulate our work more deeply and clearly."}, {"review_id": "LEpxRN9uhI", "reviewer": "Reviewer_wZge", "summary": "This paper investigates how the transformer-based model can compositionally learn some complex functions (e.g., an MLP) by breaking them down to some atom problems (e.g., linear mapping). Such an ability is also the crux of the success of Chain-of-thought (CoT) in-context learning (ICL) methods in large language models. Hence this paper bridges the gap between these two fields by theoretically analyzing the sample complexity of different methods and the accelerating effect of CoT in pertaining. Although I believe the contribution of the paper is solid, novel, and helpful to the field, there are plenty of problems that make the paper hard to follow. I wish the authors could polish the paper and tackle some of my concerns to make the paper stronger. Hence at this stage, I would only give a borderline rejection. I would be very happy to increase my score during the rebuttal phase.", "questions": "1. For the experiment in Figure 1, will ICL and CoT-I converge to a lower value if the number of samples keeps increasing? Because the sample complexity of CoT-I/O is lower than the other two and it converges at around 100 samples, would the other two methods converge if we use 500 samples?\n\n2. In Figure 2,3, it is hard to distinguish different methods using dashed lines, using different markers might be good.\n\n3. In Figure 3, if we want to see the influence of k, why not draw a figure using k as the x-axis?\n\n4. Figure 3 is on page 2, but it is first referred at page 4 — it is hard to locate the figure when reading the paper.\n\n5. Some typos and imprecise expressions:\n    \n    a.) In section 3.1, the first sentence, ‘we train 2-layer MLPs’. Do we train this MLP? IIUC, this MLP is used to generate the training samples and is fixed all the time.\n    \n    b.) In section 3.2, the last paragraph, ‘consider the consider the condition’.", "limitations": "I’m not sure whether this is an unavoidable problem of the CoT-style ICL, but it would be nice if the author can answer the following problem somewhere in the paper, which might make the paper more solid. I think comparing the sample complexity (i.e., the number of examples in prompting) between these 3 algorithms is unfair. Imagine we have $k$ examples and the problem has L steps. Then in ICL, there are only $k$ input vectors and 1 supervisory signal. In CoT-I, there are $Lk$ input vectors and 1 supervisory signal. For CoT-I/O, there are $Lk$ input vectors and $(L-1)k$ supervisory signals (I’m not sure in this paper's setting, whether ground truth $s_n^l$ is accessible in CoT-I/O, but if that is true, the comparison would be more unfair.) Given the differences of the input and supervisory signals, it might be straightforward to conclude that CoT has smaller sample complexity than ICL.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 3, "strengths": "See the summary part.", "weaknesses": "1. I find the paper hard to follow, maybe because the theoretical part, which I believe is the most important contribution, is too abstract. It is not easy to get intuition from current section 3.2, so maybe explaining how theorem 1 is formulated is helpful. \n\n2. The experimental parts are not well organized. First, it appears almost everywhere (in section 2,3,4), which breaks the flow of the paper significantly. Also, some of the subfigures in Fig3,4,5 contain similar information. So it would be good to re-organize these results to make the paper easier to follow.\n\n3. In section 4.3, the paper switches from the ICL setting to the pertaining setting. It is beneficial to clearly explain what’s the difference between them and why we need experiments in this setting.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed feedback and encouraging feedback. We hope that we have addressed your questions and concerns adequately below.\n\n> **W1. Theory is too abstract:** We provide below some more detailed description of the construction that implements this filtering, which we will ensure to include in our revised version. \n>\n> First off, we include a figure in the uploaded pdf which shows the process in a high level depiction. The idea is to filter out pairs of data that are relevant to the current prediction. So, assuming that we currently want to predict the $l$-th layer's output. In that case we want to filter out the input data $x_{l-1},x_{l}$, $x_{L+l-},x_{L+l}, ...$, where $L$ is the total number of layers and $x_i$ are the input data. To do so, we implement the following steps:\n>\n>1. Given $n$ bits $b_i$, which take values $0$ or $1$ we zero out any data point $x_i$ that its corresponding bit is zero. \n2. The second step is to construct these bits, which are the indicators of what needs to be filtered. As we also mention in our main paper, this procedure of filtering is agnostic to the token-to-be predicted and it is implemented in an automated way. \n> \n> A detailed proof is given in section A of the appendix. Finally, we will highlight that the *self-attention mechanism plays a critical role in both filtering and ICL stages highlighting the transformer-specific nature of our theorem*. For instance, attention layer is crucial for selecting which tokens should be processed next (see the attached Figure 1 for visualization). We will provide a discussion on this and also (experimentally) visualize the theorem's message on \"Chain-of-Thought <=> Filtering + ICL + Looping\" via attention maps.\n\n> **W2. Reorganize experimental section:** We recognize the lack of organization in our work and have made plans to revise and restructure it. Specifically, we will combine the experimental results and focus Section 3 purely on theory. Details can be found in the general response.\n\n> **W3. ICL vs Pretraining setting** Thanks for raising this question: The initial part of the paper focuses on the inference phase of in-context learning. That is, we are interested in the number of examples in the prompt to correctly predict the test query. By \"pretraining\" we mean the training phase of the in-context learning and we study the number of prompts the transformer needs to be trained with so that it can successfully in-context learn during inference. Our conclusion (based on deep linear MLPs) is that CoT helps improve the sample complexity of this training phase by learning shortcuts to represent complex functions. To avoid confusion, we will replace all \"pretraining\" phrases with \"ICL training\"\n\n> **Q1. Will ICL converge given more in-context samples? (attached Figure 5)** We appreciate the reviewer's query and the short answer is: NO (unless we enlarge the model size). 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 4 in the paper illustrates 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 5 in the attached file. Both Figures 5(a) and 5(b) detail training models with MLP tasks of dimensions $d=10$ and $k=4$. In Figure 5(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 Figure 5(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$.\n\n> **Q2.** Thank you for the constructive point and suggestion. We will integrate the changes in the updated version.\n\n> **Q3.** We appreciate this perspective. We did consider this during our research. However, since test risks fluctuate significantly with in-context sample sizes, plotting the x-axis as $k$ presents challenges. We would have to either focus on a specific in-context sample size or display the averaged risks, making it difficult to determine whether CoT is universally superior to ICL or only in particular cases of in-context samples. What's more, current plots could clearly show the performance alignments over different settings. Nevertheless, we still wanted to illustrate how the test risks vary with $k$, and for this purpose, we have included a bar figure in Figure 5(a) in the paper, showing the averaged risks.\n\n> **Q4.** We apologize for the lack of organization and appreciate your patience. We have devised a plan to reorganize our work, as outlined in the general response, and we hope this addresses your concerns.\n\n> **Q5.** Thank you for pointing out the errors, and we apologize for any confusion. The sentences and typos have been corrected. \n\n> **Limitations:** We apprepriate the reviewer's thoughtful feedback. While we agree ICL has a smaller in-context window size compared to CoT, we want to emphasize that our setting is aligned with the common practice where reasoning prompts are typically denser than merely providing final answers, and explanations prove helpful. Moreover, as Reviewer zLZS syggests, we have conducted varant ICL experiments by filling the explanation with some random data and the results are shown in attached Figures 4. Fig 5(a) further illustrates that even with more in-context samples, ICL cannot surpass CoT due to task complexity and model expressivity limitations. "}]

Author rebuttal:
We thank the reviewers for their constructive comments and insightful questions. We are gratified that many of the reviewers found our work insightful and interesting. In the following sections, we will summarize our key contributions, respond to shared concerns raised by the reviewers, and provide explanations for the new figures presented.

We briefly recap our **main contribution and message**: There are several major works on in-context learning (ICL) theory (e.g. [Xie et al. ICLR'22, Garg et al. NeurIPS'22]). These focus on learning from the examples provided in the prompt. Compared to these, we consider a novel setting where prompt contains chain-of-features obtained from a multistep problem-solving (or reasoning) process. We show that when solving a new problem, **chain-of-thought operates by (1) filtering the heterogeneous prompt and retrieving the relevant features, (2) running in-context learning on the filtered prompt, and (3) Loop back to (1).** This process is illustrated in attached Figure 1. Our key contributions are establishing how transformers can provably implement this process and demonstrating the empirical and theoretical value of CoT in terms of approximation and sample complexity. In our Theorem 1, the attention mechanism provides a critical role in both filtering and ICL stages highlighting the transformer-specific nature of our results. We will further elaborate on this in the final manuscript. 

To proceed, we address the shared concerns by the reviewers:

- **Paper Organization.** Reviewers 8qtM, zLZS, and wZge noted concerns with the paper's clarity, citing confusion over the existence of both empirical and experimental results sections. We appreciate their feedback and have laid out a plan to enhance the paper's readability: 
    1. Section 3 will be focused solely on theory, and additional discussion will be included to better motivate and clarify our theoretical results. For instance, we will include an illustration to elucidate how the theory functions (see Fig. 1 in the .pdf and also the response to Reviewer wZge for further details).
    2. All experiments involving 2-layer MLPs will be consolidated in Section 4.
    3. Section 5 will address the training aspects of CoT by exploring deep linear MLPs.

- **Lack of Limitations Section.** We recognize our oversight in neglecting to include an explicit discussion of the paper's limitations and sincerely apologize for this omission. Rest assured, we will incorporate a limitations section in our revised manuscript. Our notable limitations include: 
    1. We should have provided a deeper study of out-of-distribution scenarios. In response to Reviewer mo8c, we have partly addressed this limitation (see Fig 2 in the pdf).
    2. Our key contributions are theoretical and we work with a synthetic problem setting which extends [Garg et al. NeurIPS'22, Akyurek et al. ICLR'23] to compositional functions (namely MLPs) and CoT.

**Supporting experiments:** We provide the following figures to address reviewer inquiries.

- Figure 1 aims to clarify the chain-of-thought prompting formalized by our work by visualizing how CoT can be decoupled into filtering and in-context-learning stages.
- Figure 2 provides new robustness experiments in response to Reviewer mo8c. In Fig 2(a), we add noise to the label $y$ (solid) or input features $x$ (dashed). This shows that CoT methods (trained on noiseless data) are fairly robust to noisy data at test-time. In Fig 2(b), we in-context learn an MLP with $k$ hidden nodes whereas transformer is trained for MLPs with $8$ hidden nodes. This shows that CoT-I/O is robust to misspecification as long as it is small. However, if $k$ is very small, then CoT-I/O suffers more distribution shift. This makes sense because CoT-I/O relies more on hidden features compared to others. In Fig 2(c), we consider misspecification of input dimension: TF is trained with $d=10$ but we feed a neural net with $d<10$. This reveals that CoT is pretty robust and, in fact, CoT-I is the most robust to misspecification. We speculate this is because CoT-I makes similar use of input and hidden features.
- In Figure 3, we investigate deep nonlinear MLPs. The main takeaway is that, more CoT steps help but not that much. We believe this might be because of the random generation of neural nets. Specifically, it is possible that when neural net weights are fully random, a two layer neural net might be able to accurately approximate a 4 layer network. This would mean only 2 CoT steps are needed to do a good job. 
- In Figure 4, we put vanilla ICL in CoT format and feed $(x,[?],y)$ in response to Reviewer zLZS. We set [?] mark as random gaussian features or all [-1] vector. In both cases, the performance coincide with vanilla ICL $(x,y)$ which we found to be intuitive.
- In Figure 5, we train the models with more in-context examples as requested by Reviewer wZge. 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.


Finally, as for the concerns/questions raised, we believe that we successfully addressed all of them sufficiently and reply in line to each review. We would be grateful to respond further reviewer inquiries during the discussion week.

Meta-review: The authors make a novel contribution towards understanding the phenomenon of CoT. Reviewers converged on recommending acceptance after a lively rebuttal phase. One lingering concern is about the clarity of presentation -- but much of this is likely due to the abstract nature of the contribution and the authors have made a strong effort to respond to reviewer feedback during rebuttal + laid out a plan for improving it. The contribution is valuable and novel, and worth discussion at NeurIPS
Decision: Accept (poster)