Title: Dissecting Chain-of-Thought: Compositionality through In-Context Filtering and Learning
Venue: NeurIPS 2023
Accepted: yes
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-Thought: Compositionality through In-Context Filtering and Learning

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.

1 Introduction

The advent of transformers [Vaswani et al., 2017] has revolutionized natural language processing, paving the way for remarkable performance in a wide array of tasks. LLMs, such as GPTs [Brown et al., 2020], have demonstrated an unparalleled ability to capture and leverage vast amounts of data, thereby facilitating near human-level performance across a variety of language generation tasks. Despite this success, a deep understanding of their underlying mechanisms remains elusive.

Chain-of-thought prompting [Wei et al., 2022c] is an emergent ability of transformers where the model solves a complex problem [Wei et al., 2022b], by decomposing it into intermediate steps. Intuitively, this underlies the ability of general-purpose language models to accomplish previously-unseen complex tasks by leveraging more basic skills acquired during the pretraining phase. Compositional learning and CoT has enjoyed significant recent success in practical language modeling tasks spanning question answering, code generation, and mathematical reasoning [Perez et al., 2021, Imani et al., 2023, Yuan et al., 2023]. In this work, we attempt to demystify some of the mechanics underlying this success and the benefits of CoT in terms of sample complexity and approximation power. To

… …

…

…

…

…

# in-context sample

Error

ICL

CoT-I

CoT-I/O

Input:

x

1st hidden:

s

1

2nd hidden:

s

2

Output:

y

ICL

x

1

y

1

x

2

y

2

. . .

x

test

ˆ

y

CoT-I

x

1

s

1

1

s

2

1

y

1

x

2

s

1

2

. . .

x

test

ˆ

y

CoT-I/O

x

1

s

1

1

s

2

1

y

1

x

2

. . .

x

test

ˆ

s

1

. . .

x

test

ˆ

s

1

ˆ

s

2

. . .

x

test

ˆ

s

1

ˆ

s

2

ˆ

y

0

20

40

60

80

100

# in-context samples

0.0

0.1

0.2

0.3

0.4

Test risk

CoT-I/O

CoT-I

ICL

do this we explore the role of CoT in learning multi-layer perceptrons (MLPs) in-context, which we believe can lead to a first set of insightful observations. Throughout, we ask:

Does CoT improve in-context learning of MLPs, and what are the underlying mechanics?

In this work, we identify and thoroughly compare three schemes as illustrated in Figure 1. (a) ICL: In-context learning from input-output pairs provided in the prompt, (b) CoT-I: Examples in the prompt are augmented with intermediate steps, (c) CoT-I/O: The model also outputs intermediate steps during prediction. Our main contributions are:

Decomposing CoT into filtering and ICL: As our central contribution, we establish a rigorous and experimentally-supported abstraction that decouples CoT prompting into a filtering phase and an in-context learning (ICL) phase . In filtering , the model attends to the relevant tokens within the prompt based on an instruction. In ICL , the model runs inference on the filtered prompt to output a step and then moves onto the next step in the chain. Our Theorem 1 develops a theoretical understanding of this two-step procedure and formalizes how filtering and ICL phases of CoT can be implemented via the self-attention mechanism to learn MLPs.

Approximation and sample complexity: Through experiments and theory, we establish that intermediate steps in CoT-I improves the sample complexity of learning whereas step-by-step output improves the approximation ability through looping. Specifically, CoT-I/O can learn an MLP with input dimension d and k neurons using O (max( k, d )) in-context samples by filtering individual layers and solving them via linear regression - in contrast to the Ω( kd ) lower bound without step-augmented prompt. As predicted by our theory, our experiments (see Sec. 4) identify a striking universality phenomenon (as k varies) and also demonstrate clear approximation benefits of CoT compared to vanilla ICL.

Accelerated pretraining via learning shortcuts: We construct deep linear MLPs where each layer is chosen from a discrete set of matrices. This is in contrast to the above setting, where MLP weights can be arbitrary. We show that CoT can dramatically accelerate pretraining by memorizing these discrete matrices and can infer all layers correctly from a single demonstration. Notably the pretraining loss goes to zero step-by-step where each step 'learns to filter a layer' . Together, these showcase how CoT identifies composable shortcuts to avoid the need for solving linear regression. In contrast, we show that ICL (without CoT) collapses to linear regression performance as it fails to memorize exponentially many candidates (due to lack of composition).

The paper is organized as follows. In Section 2, we introduce the problem setup. Section 3 states our main theoretical results which decouple CoT into filtering and ICL. Section 4 provides empirical investigations of CoT with 2-layer MLPs, which validates our theoretical findings. Finally, we

elucidate the benefits of CoT during pretraining via experiments on deep linear MLPs in Section 5. Related work and discussion are provided in Sections 6 and 7, respectively.

2 Preliminaries and Setup

Notation. We denote the set { 1 , 2 , . . . , n } as [ n ] . Vectors and matrices are represented in bold text (e.g., x , A ), while scalars are denoted in plain text (e.g., y ). The input and output domains are symbolized as X and Y respectively (unless specified otherwise), and x ∈ X , y ∈ Y denote the input and output. Additionally, let F be a set of functions from X to Y . Consider a transition function f ∈ F where y = f ( x ) . In this section, we explore the properties of learning f , assuming that it can be decomposed into simpler functions ( f ℓ ) L ℓ =1 , and thus can be expressed as f = f L ◦ f L -1 ◦· · · ◦ f 1 .

2.1 In-context Learning

Following the study by Garg et al. [2022], the fundamental problem of vanilla in-context learning (ICL) involves constructing a prompt with input-output pairs in the following manner:

Here the transition function f ∈ F : 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 as p ( Spanish ) = (( apple, manzana ), ( ball, pelota ), . . . ) or 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 to learn a model that can accurately predict the output, given a prompt p and the test input x test , as shown in the following equation:

where ˜ f ∈ F is the test function which may differ from the functions used during training. Previous work [Zhou et al., 2022, Li et al., 2023c] has demonstrated that longer prompts (containing more examples n ) typically enhance the performance of the model.

2.2 Chain-of-thought Prompt and Prediction

As defined in (P-ICL), the prompt in vanilla ICL only contains input-output pairs of the target function. This demands that the model learns the function f ∈ 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-thought methods (e.g., [Wei et al., 2022c]) observed that prompts containing step-by-step instructions assist the model in decomposing the function and making better predictions. Specifically, consider a function composed of L subfunctions, represented as f := f L ◦ f L -1 ◦ . . . f 1 . Each intermediate output can be viewed as a step, enabling us to define a lengthn CoT prompt related to f with L steps (expressed with s ℓ , ℓ ∈ [ L ] ) as follows:

Here x i = s 0 i , y i = s L i and f ℓ ∈ F ℓ , which implies that f ∈ F L ×···F 1 := F .

Next we introduce two methodologies for making predictions within the CoT framework:

CoT over input only (CoT-I). Contrasted with ICL, CoT-I considers step-by-step instructions as inputs, nonetheless, the prediction for the last token is performed as a single entity. Our experiments indicate that this approach lowers the sample complexity for TF to comprehend the function ˜ f being learned (see the orange curves in Figures 5 and 6). The CoT-I prediction aligns with Eq. (2.1) as follows, while the prompt is determined by (P-CoT).

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 one shot. To mitigate this challenge, we consider

a scenario where in addition to implementing a CoT prompt, we also carry out CoT predictions. Specifically, for a composed problem with inputs formed via (P-CoT), the model recurrently makes L -step predictions as outlined below:

where at each step (step ℓ ), the model outputs an intermediate step ( ˆ s ℓ ) which is then fed back to the input sequence to facilitate the next-step prediction ( ˆ s ℓ +1 ). Following this strategy, the model only needs to learn the union of the sub-function sets, ⋃ L ℓ =1 F ℓ , whose complexity scales linearly with the number of steps L . Empirical evidence of the benefits of CoT-I/O over ICL and CoT-I in enhancing sample efficiency and model expressivity is reflected in the blue curves shown in Figures 5 and 6.

2.3 Model Training

In Figure 1 and Section 2, we have discussed vanilla ICL, CoT-I and CoT-I/O methods. Intuitively, ICL can be viewed as a special case of CoT-I (or CoT-I/O) if we assume only one step is performed. Consequently, we will focus on implementing CoT-I and CoT-I/O for model training in the following.

Consider the CoT prompt as in (P-CoT), and assume that x ∼ D X , and f ℓ ∼ D ℓ , ℓ ∈ [ L ] , where L denotes the number of compositions/steps, such that the final prediction should approximate f ( x ) = f L ( f L -1 . . . f 1 ( x )) := y ∈ Y . We define ℓ (ˆ y , y ) : Y × Y → R as a loss function. For simplicity, we assume f ℓ ( . . . f 1 ( x )) ∈ Y , ℓ ∈ [ L ] . Let N represent the in-context window of TF , which implies that TF can only admit a prompt containing up to N in-context samples. Generally, our goal is to ensure high prediction performance given any lengthn prompt, where n ∈ [ N ] . To this end, we train the model using prompts with length from 1 to N equally and aim to minimize the averaged risk over different prompt size. Assuming the model TF is parameterized by θ and considering meta learning problem, the objective functions for CoT-I and CoT-I/O are defined as follows.

and

where ˆ y n = TF ( p n ( f ) , x n ) and ˆ s ℓ n = TF ( p n ( f ) , x n · · · s ℓ -1 n ) . p n ( f ) is given by (P-CoT), and as mentioned previously, s 0 = x and s L = y . All x and f ℓ are independent, and we take the expectation of the risk over their respective distributions.

3 Provable Approximation of MLPs via Chain-of-Thought

In this section, we present our theoretical findings that demonstrate how CoT-I/O can execute filtering over the CoT prompt, thereby learning a 2-layer MLP with input dimension of d and hidden dimension of k , akin to resolving k d -dimensional ReLU problems and 1 k -dimensional linear regression problem. Subsequently, in Section 4.1, we examine the performance of CoT-I/O when learning 2-layer random MLPs. Our experiments indicate that CoT-I/O needs only O (max( d, k )) in-context samples to learn the corresponding MLP.

We state our main contribution of establishing a result that decouples CoT-based in-context learning (CoT-I/O) into two phases: (1) Filtering Phase: Given a prompt that contains features of multiple MLP layers, retrieve only the features related to a target layer to create an ICL prompt. (2) ICL Phase: Given filtered prompt, learn the target layer weights through gradient descent. Combining these two

CoT prompt

Filtering

Step

1

ICL:

Step

2

ICL:

Step

3

ICL:

Prediction

x

1

s

1

1

s

2

1

y

1

x

2

s

1

2

s

2

2

y

2

. . .

x

test

x

1

s

1

1

x

2

s

1

2

. . .

x

test

ˆ

s

1

s

1

1

s

2

1

s

1

2

s

2

2

. . .

ˆ

s

1

ˆ

s

2

s

2

1

y

1

s

2

2

y

2

. . .

ˆ

s

2

ˆ

y

phases, and looping over all layers, we will show that there exists a transformer architecture such that CoT-I/O can provably approximate a multilayer MLP up to a given resolution. An illustration is provided in Figure 2. To state our result, we assume access to an oracle that performs linear regression and consider the consider the condition number of the data matrix.

Definition 1 (MLP and condition number) Consider a multilayer MLP defined by the recursion s ℓ i = ϕ ( W ℓ s ℓ -1 i ) for ℓ ∈ [ L ] , i ∈ [ n ] and s 0 i = x i . Here ϕ ( x ) := max( αx, x ) is a Leaky-ReLU activation with 1 ≥ α > 0 . Define the feature matrix T ℓ = [ s ℓ 1 . . . s ℓ n ] ⊤ and define its condition number κ ℓ = σ max ( T ℓ ) /σ min ( T ℓ ) (with σ min := 0 for fat matrices) and κ max = max 0 ≤ ℓ<L κ ℓ .

Assumption 1 (Oracle Model) We assume access to a transformer TF LR which can run T steps of gradient descent on the quadratic loss L ( w ) = ∑ n i =1 ( y i -w ⊤ x i ) 2 given a prompt of the form ( x 1 , y 1 , . . . , x n , y n ) .

Weremark that this assumption is realistic and has been formally established by earlier work [Giannou et al., 2023, Akyürek et al., 2022]. Our CoT abstraction builds on these to demonstrate that CoT-I/O can call a blackbox TF model to implement a compositional function when combined with filtering.

We now present our main theoretical contribution. Our result provides a transformer construction that first filters a particular MLP layer from the prompt through the attention mechanism, then applies in-context learning, and repeats this procedure to approximate the MLP output. The precise statement is deferred to the supplementary material.

Theorem 1 (CoT ⇔ Filtering+ICL) Consider a CoT prompt p n ( f ) generated from an L -layer MLP f ( · ) as described in Definition 1, and assume given test example ( x test , s 1 test , . . . s L test ) . For any resolution ϵ > 0 , there exists δ = δ ( ϵ ) , iteration choice T = O ( κ 2 max log(1 /ϵ )) , and a backend transformer construction TF BE such that the concatenated transformer TF = TF LR ◦ TF BE implements the following: Let (ˆ s i ) ℓ -1 i =0 denote the first ℓ -1 CoT-I/O outputs of TF where ˆ s 0 = x test and set p [ ℓ ] = ( p n ( f ) , x test , ˆ s 1 . . . ˆ s ℓ -1 ) . At step ℓ , TF implements

Filtering. Define the filtered prompt with input/output features of layer ℓ ,

There exists a fixed projection matrix Π that applies individually on tokens such that the backend output obeys ∥ Π ( TF BE ( p [ ℓ ])) -p filter n ∥ ≤ δ .

Gradient descent. The combined model obeys ∥ TF ( p [ ℓ ]) -s ℓ test ∥ ≤ ℓ · ϵ/L .

TF BE has constant number of layers independent of n and T . Consequently, after L rounds of CoT-I/O, TF outputs f ( x test ) up to ϵ accuracy.

0

20

40

60

80

100

# in-context samples

0.0

0.1

0.2

0.3

0.4

Test risk

k

=4

k

=8

k

=16

d

=10

d

=20

(a) CoT-I/O:

d

= 10

vs

20

0

2

4

6

8

10

# in-context samples/

d

0.0

0.1

0.2

0.3

0.4

Test risk

d

=10,

k

=16

d

=20,

k

=16

d

=10,

k

=16

d

=20,

k

=16

(b) Alignment in Fig. 3(a)

0

20

40

60

80

100

# in-context samples

0.0

0.1

0.2

0.3

0.4

Test risk

k

=4

k

=16

k

=64

First layer risk

Second layer risk

First layer risk

Second layer risk

(a) CoT-I/O: composed risk (

d

= 10

)

0

20

40

60

80

100

# in-context samples

0.0

0.1

0.2

0.3

0.4

Test risk

k

=4

k

=16

k

=64

(b) Risk of first layer

0

20

40

60

80

100

# in-context samples

0.0

0.1

0.2

0.3

0.4

Test risk

k

=4

k

=16

k

=64

(c) Risk of second layer

Remark 1 Note that, this result effectively shows that, with a sufficiently good blackbox transformer TF LR (per Assumption 1), CoT-I/O can learn an L -layer MLP using in-context sample size n > max ℓ ∈ [ L ] d ℓ where d ℓ is the input dimension of ℓ th layer. This is assuming condition number κ max of the problem is finite as soon as all layers become over-determined. Consequently, CoT-I/O needs max( k, d ) sample complexity to learn a two layer MLP. This provides a formal justification for the observation that empirical CoT-I/O performance is agnostic to k as long as k ≤ d .

We provide the concrete filtering statements based on the transformer architecture in Appendix A, and the key components of our construction are the following: (i) Inputs are projected through the embedding layer in which a set of encodings, an enumeration of the tokens ( 1 , 2 , . . . , N ), an enumeration of the layers ( 1 , 2 , . . . , L ) and an identifier for each layer already predicted are all attached. Notice that this 'modification' to the input only depends on the sequence length and is agnostic to the token to-be-predicted. This allows for an automated looping over L predictions. (ii) We use this information to extract the sequence length N and the current layer ℓ to-be-predicted. (iii) With these at hand, we construct an 'if-then' type of function using the ReLU layers to filter out the samples that are not needed for the prediction.

4 Experiments with 2-layer Random MLPs

For a clear exposition, we initially focus on two-layer MLPs, which represent 2-step tasks (e.g., L = 2 ). We begin by validating Theorem 1 using the CoT-I/O method, demonstrating that in-context learning for a 2-layer MLP with d input dimensions and k hidden neurons requires O (max( d, k )) samples. The results are presented in Section 4.1. Subsequently, in Section 4.2, we compare three different methods: ICL, CoT-I, and CoT-I/O. The empirical evidence highlights the advantages of CoT-I/O, showcasing its ability to reduce sample complexity and enhance model expressivity.

Dataset. Consider 2-layer MLPs with input x ∈ R d , hidden feature (step-1 output) s ∈ R k , and output y ∈ R . Here, s = f 1 ( x ) := ( Wx ) + and y = f 2 ( s ) := v ⊤ s , with W ∈ R k × d and v ∈ R k being the parameters of the first and second layer, and ( x ) + = max( x, 0) being ReLU activation. The function is composed as y = v ⊤ ( Wx ) + . We define the function distributions as follows: each entry of W is sampled via W ij ∼ N (0 , 2 k ) , and v ∼ N (0 , I k ) , with inputs being randomly sampled through x ∼ N (0 , I d ) 1 . We apply the quadratic loss in our experiments. To avoid the implicit bias due to distribution shift, both training and test datasets are generated following the same strategy.

CoT-I/O

CoT-I

ICL

0.0

0.1

0.2

0.3

Averaged risk

Standard

Small

Tiny

(a) Averaged risk

0

20

40

60

80

100

# in-context samples

0.0

0.1

0.2

0.3

0.4

Test risk

CoT-I/O

CoT-I

ICL

(b) Standard GPT-2

0

20

40

60

80

100

# in-context samples

0.0

0.1

0.2

0.3

0.4

Test risk

CoT-I/O

CoT-I

ICL

(c) Small GPT-2

0

20

40

60

80

100

# in-context samples

0.0

0.1

0.2

0.3

0.4

Test risk

CoT-I/O

CoT-I

ICL

(d) Tiny GPT-2

CoT-I/O

CoT-I

ICL

0.00

0.05

0.10

0.15

0.20

Averaged risk

k

=4

k

=8

k

=16

(a) Averaged risk

0

20

40

60

80

100

# in-context samples

0.0

0.1

0.2

0.3

0.4

Test risk

CoT-I/O

CoT-I

ICL

(b)

k

= 4

0

20

40

60

80

100

# in-context samples

0.0

0.1

0.2

0.3

0.4

Test risk

CoT-I/O

CoT-I

ICL

(c)

k

= 8

0

20

40

60

80

100

# in-context samples

0.0

0.1

0.2

0.3

0.4

Test risk

CoT-I/O

CoT-I

ICL

(d)

k

= 16

4.1 Empirical Evaluation of CoT-I/O Performance

To investigate how MLP task impacts CoT-I/O performance, we train 2-layer MLPs with varying input dimensions ( d ) and hidden layer sizes ( k ). The results are presented in Figures 3 and 4, and all experiments utilize small GPT-2 models for training 2 .

CoT-I/O performance is agnostic to k when k ≤ d (Figure 3). In Fig. 3(a), we train MLPs with d = 10 , 20 and k = 4 , 8 , 16 . Solid and dashed curves represent the CoT-I/O test risk of d = 10 and 20 respectively for varying in-context samples. The results indicate that an increase in d will amplifies the number of samples needed for in in-context learning, while the performance remains unaffected by changes in k ∈ { 4 , 8 , 16 } . To further scrutinize the impact of d on CoT-I/O accuracy, in Fig. 3(b), we adjust the horizontal axis by dividing it by the input dimension d , and superimpose both d = 10 , k = 16 (blue solid) and d = 20 , k = 16 (orange dashed) results. This alignment of the two curves implies that the in-context sample complexity of CoT-I/O is linearly dependent on d .

Large k dominates CoT-I/O performance (Figure 4). We further investigate the circumstances under which k begins to govern the CoT-I/O performance. In Figure 4(a), we replicate the same experiments with d = 10 , but train with wider MLPs ( k = 64 ). Blue, orange and green curves represent results for k = 4 , 16 , 64 respectively. Since the hidden dimension k = 64 is larger, learning the second layer requires more hidden features ( s ), thus N = 100 in-context samples (providing 100 s s) are insufficient to fully restore the second layer, leading to performance gaps between k = 4 , 16 and k = 64 . To quantify the existing gaps, we conduct single-step evaluations for both the first and the second layers, with the results shown in Figures 4(b) and 4(c). Specifically, let p n ( ˜ f ) be a test prompt containing n in-context samples where ˜ f represents any arbitrary 2-layer MLP. Given a test sample ( x test , s test , y test ) , the layer predictions are performed as follows.

The test risks are calculated by ∥ ˆ s -s test ∥ 2 and (ˆ y -y test ) 2 . The risks illustrated in the figures are normalized for comparability (refer to the appendix for more details). Evidence from Fig. 4(b) and 4(c) shows that while increasing k does not affect the first layer's prediction, it does augment the number of samples required to learn the second layer. Moreover, by plotting the first layer risks of k = 4 , 16 (blue/orange dotted) and second layer risk of k = 64 (green dashed) in Fig. 4(a), we can see that they align with the CoT-I/O composed risks. This substantiates the hypothesis that CoT-I/O learns 2-layer MLP through compositional learning of separate layers.

0

2

4

6

8

10

# in-context samples

0.0

0.2

0.4

0.6

0.8

1.0

Test risk

ICL

CoT-2

CoT-3

CoT-6

Least squares

n

=

d

(a) Point-to-point meta prediction

0

100k

200k

300k

400k

500k

# iteration

0.0

0.2

0.4

0.6

0.8

1.0

Test risk

ICL

CoT-2

CoT-3

CoT-6

(b) One-shot prediction over time

4.2 Comparative Analysis of ICL, CoT-I and CoT-I/O

Varying model sizes (Figure 5). We initially assess the benefits of CoT-I/O over ICL and CoT-I across different TF models. With d = 10 and k = 8 fixed, we train three different GPT-2 models: standard, small and tiny GPT-2. The small GPT-2 has 6 layers, 4 attention heads per layer and 128 dimensional embeddings. The standard GPT-2 consists of twice the number of layers, attention heads and embedding dimensionality compared to the small GPT-2, and tiny GPT-2, on the other hand, possesses only half of these hyperparameters compared to the small GPT-2. We evaluate the performance using prompts containing n in-context samples, where n ranges from 1 to N ( N = 100 ). The associated test risks are displayed in Figs. 5(b), 5(c) and 5(d). The blue, orange and green curves correspond to CoT-I/O, CoT-I and ICL, respectively. In Fig. 5(a), we present the averaged risks. The results show that using CoT-I/O, the small GPT-2 can solve 2-layer MLPs with approximately 60 samples, while CoT-I requires the standard GPT-2. Conversely, ICL is unable to achieve zero test risk even with the standard GPT-2 model and up to 100 samples. This indicates that to learn 2-layer MLPs in a single shot, ICL requires at least O ( dk + d ) samples to restore all function parameters. Conversely, CoT-I and CoT-I/O can leverage implicit samples contained in the CoT prompt. Let f 1 ∈ F 1 (first layer) and f 2 ∈ F 2 (second layer). By comparing the performance of CoT-I and CoT-I/O, it becomes evident that the standard GPT-2 is capable of learning the composed function f = f 2 ◦ f 1 ∈ F , which the small GPT-2 cannot express.

Varying MLP widths (Figure 6). Next, we explore how different MLP widths impact the performance (by varying the hidden neuron size k ∈ { 4 , 8 , 16 } ). The corresponding results are depicted in Figure 6. The blue, orange and green curves in Fig. 6(b), 6(c) and 6(d) correspond to hidden layer sizes of k = 4 , 8 , and 16 , respectively. Fig. 6(a) displays the averaged risks. We keep d = 10 , N = 100 fixed and train with the small GPT-2 model. As discussed in Section 4.1, CoT-I/O can learn a 2-layer MLP using around 60 samples for all k = 4 , 8 , 16 due to its capability to deconstruct composed functions. However, CoT-I can only learn the narrow MLPs with k = 4 , and ICL is unable to learn any of them. Moreover, we observe a substantial difference in the performance of ICL and CoT-I with varying k (e.g., see averaged risks in Fig. 6(a)). This can be explained by the fact that enlarging k results in more complex F 1 and F 2 , thus making the learning of F = F 2 ×F 1 more challenging for ICL and CoT-I.

5 Further Investigation on Deep Linear MLPs

In Section 4, we have discussed the approximation benefits of CoT-I/O and how it in-context learns 2-layer random MLPs by parallel learning of k d -dimensional ReLU and 1 k -dimensional linear regression. In this section, we investigate the capability of CoT-I/O in learning longer compositions. For brevity, we will use CoT to refer to CoT-I/O in the rest of the discussion.

Dataset. Consider L -layer linear MLPs with input x ∈ R d ∼ N (0 , I d ) , and output generated by y = W L W L -1 · · · W 1 x , where the ℓ th layer is parameterized by W ℓ ∈ R d × d , ℓ ∈ [ 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 k : ¯ W ⊤ k ¯ W k = I , k ∈ [ K ] } 3 . Therefore, to learn the L -layer neural net, CoT only needs to learn ¯ F with | ¯ F| = K , whereas ICL needs to learn the function set ¯ F L , which contains K L random matrices.

0

10k

20k

30k

40k

50k

# iteration

0.0

0.2

0.4

0.6

0.8

1.0

Test risk

ICL

CoT-2

CoT-3

CoT-6

(a) Composed risk

0

10k

20k

30k

40k

50k

0.0

0.2

0.4

0.6

0.8

1.0

Layer 6

0

10k

20k

30k

40k

50k

0.0

0.2

0.4

0.6

0.8

1.0

Layer 3

0

10k

20k

30k

40k

50k

0.0

0.2

0.4

0.6

0.8

1.0

Layer 2

0

10k

20k

30k

40k

50k

0.0

0.2

0.4

0.6

0.8

1.0

Test risk

Layer 4

0

10k

20k

30k

40k

50k

0.0

0.2

0.4

0.6

0.8

1.0

Test risk

Layer 1

0

10k

20k

30k

40k

50k

# iteration

0.0

0.2

0.4

0.6

0.8

1.0

Layer 5

(b) Risk of each layer output

Composition ability of CoT (Figures 7). Set d = 5 , L = 6 and K = 4 . At each round, we randomly select L matrices W ℓ , ℓ ∈ [ L ] from ¯ F so that for any input x , we can form a chain

where s ℓ = W ℓ s ℓ -1 , ℓ ∈ [ L ] and s 0 := x . Let CoTX denote X -step CoT-I/O method. For example, the in-context sample of CoT-6 has form of ( x , s 1 , s 2 , . . . s 5 , y ) , which contains all the intermediate outputs from each layer; while CoT-3, CoT-2 have prompt samples formed as ( x , s 2 , s 4 , y ) and ( x , s 3 , y ) respectively. In this setting, ICL is also termed as CoT-1, as its prompt contains only ( x , y ) pairs. To solve the length-6 chain, CoTX needs to learn a model that can remember 4 6 /X matrices. Therefore, ICL is face a significantly challenge sincd it needs to remember 4 6 = 4 , 096 matrices (all combinations of the 4 matrices used for training and testing) compared to just 4 for CoT-6.

We train small GPT-2 models using the CoT-2/-3/-6 and ICL methods, and present the results in Fig. 7(a). As evident from the figure, the performance curves of CoT-2 (orange), CoT-3 (green) and CoT-6 (red) overlap, and they can all make precise predictions in one shot (given an in-context example n = 1 ). It seems that TF has effectively learned to remember up to 64 matrices (for CoT-2) and compose up to 6 layers (for CoT-6). However, ICL (blue) struggles to learn the 6-layer MLPs in one shot. The black dashed curve shows the solution for linear regression y = β ⊤ x computed directly via least squares given n random training samples, where x is the input and y is from the output of the 6-layer MLPs (e.g., y [0] ). The test risks for n = 1 , . . . 10 are plotted (in Fig. 7(a)), which show that the ICL curve aligns with the least squares performance. This implies that, instead of remembering all 4 , 096 matrices, ICL solves the problem from the linear regression phase.

In addition to the meta-learning results which highlight the approximation benefits of multi-step CoT, we also investigate the convergence rate of CoT-2/-3/-6 and ICL, with results displayed in Fig. 7(b). We test the one-shot performance during training and find that CoT-6 converges fastest. This is because it has the smallest sub-function set, and given the same tasks (e.g., deep neural nets), shortening the chain leads to slower convergence. This supports the evidence that taking more steps facilitates faster and more effective learning of complex problems.

Evidence of Filtering (Figure 8). As per Theorem 1 and the appendix, transformers can perform filtering over CoT prompts, and the results from 2-layer MLPs align with our theoretical findings. However, can we explicitly observe filtering behaviors? In Fig. 8(a), we display the results of the first 50k iterations from Fig. 7(b), and observe risk drops in CoT-6 (red) at the 15k and 25k iteration (shown as grey dotted and dashed lines). Subsequently, in Fig. 8(b), we plot the test risk of each layer prediction (by feeding the model with correct intermediate features not the predicted ones), where CoT-6 (red) predicts the outputs from all 6 layers ( s 1 , · · · , s L ). From these figures, we can identify risk drops when predicting different layers, which appear at either 15k (for layer 2, 3, 4, 5, 6) or 25k (for layer 1) iteration. This implies that the model learns to predict each step/layer function independently. Further, we test the filtering evidence of the ℓ th layer by filling irrelevant positions with random features. Specifically, an in-context example is formed by

The test risks are represented by black dotted curves in Fig. 8(b), which aligned precisely with the CoT-6 curves (red). This signifies that each layer's prediction concentrate solely on the corresponding intermediate steps in the prefix, while disregarding irrelevant features. This observation provides evidence that the process of filtering is indeed performed.

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, chain-of-thought (CoT) prompts have made impressive improvements in performing complex reasoning by decomposing it into step-by-step intermediate solutions [Wei et al., 2022c, Narang et al., 2020, Lampinen et al., 2022, Wei et al., 2022b, Zhou et al., 2022, Nye et al., 2021, Veliˇ ckovi´ c and Blundell, 2021, Lanchantin et al., 2023], which in general, shows the ability of transformer in solving compositional functions. Lee et al. [2023], Dziri et al. [2023] study the problem of teaching arithmetic to small transformers and show that breaking the task down into small reasoning steps allows the model to learn faster. Li et al. [2023b] show that small student transformer models can learn from rationalizations sampled from significantly larger teacher models. The idea of learning how to compose skills has been well studied in other literatures [Sahni et al., 2017, Liška et al., 2018]. More specifically, for the problem of learning shallow networks, there are several well known hardness results Goel et al. [2017, 2020], Zhang et al. [2019]. In particular, Hahn and Goyal [2023] shows a formal learnability bound which implies that compositional structure can benefit ICL. However, most of the work focuses on investigating empirical benefits and algorithmic designing of CoT, and there exists little effort studying the underlying mechanisms of CoT.

Considering the expressivity of the transformer architecture itself, Yun et al. [2019] showed that TFs are universal sequence to sequence approximators. More recently, Giannou et al. [2023] use an explicit construction to show that shallow TF s can be used to run general purpose programs as long as we loop them. Other works have also shown the turing-completeness of the TF architecture but these typically require infinite/high precision and recursion around attention layers [Wei et al., 2022a, Pérez et al., 2019, 2021, Liu et al., 2022]. Closer to our work, Akyürek et al. [2022], Von Oswald et al. [2023], von Oswald et al. [2023] prove that a transformer with constant number of layers can implement gradient descent in solving linear regression, and Giannou et al. [2023] introduce similar results by looping outputs back into inputs. Ahn et al. [2023] prove this result from an optimization perspective and show that the global minimum of the training objective implements a single iteration of preconditioned gradient descent for transformers with a single layer of linear attention. Zhou et al. [2023] introduce the RASP-Generalization Conjecture which says that Transformers tend to length generalize on a task if the task can be solved by a short RASP program which works for all input lengths. In this work, we prove CoT can be treated as: first apply filtering on the CoT prompts using special construction, and then in-context learn the filtered prompt.

7 Conclusion, Limitations, and Discussion

In this work, we investigate chain-of-thought prompting and shed light on how it enables compositional learning of multilayer perceptrons step-by-step. Specially, we have explored and contrasted three methods: ICL, CoT-I and CoT-I/O, and found that CoT-I/O facilitates better approximation and faster convergence through looping and sample efficiency. Additionally, we empirically and theoretically demonstrated that to learn a 2-layer MLP with d -dimensional input and k neurons, CoT-I/O requires O (max( d, k )) in-context samples whereas ICL runs into approximation error bottlenecks.

While we have provided both experimental and theoretical results to validate the advantages of CoT, it is important to note that our analysis in the main text pertains to in-distribution scenarios. In an effort to address this limitation and demonstrate the robustness of CoT, we have conducted additional simulations, as detailed in Appendix C.1, where the test samples follow a different distribution than the training examples. Also, we note that our focus has been primarily on MLP-based tasks, where the subproblems are essentially instances of simple linear regression. It would be valuable to explore how CoT might influence the training of tasks characterized by more complex structures, longer compositional chains, and a broader variety of subproblems.

There are several interesting avenues for future research to build on our findings. To what extent does our decoupling of CoT (filtering followed by ICL) align with the empirical evidence in practical problems such as code generation and mathematical reasoning? We have shown that CoT-I/O can rely on a linear regression oracle to learn an MLP. To what extent can transformers approximate MLPs without CoT-I/O (e.g. with CoT-I) and what are the corresponding lower/upper bounds?

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.00297 , 2023.

Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. arXiv preprint arXiv:2211.15661 , 2022.

Haldun Balim, Zhe Du, Samet Oymak, and Necmiye Ozay. Can transformers learn optimal filtering for unknown systems? arXiv preprint arXiv:2308.08536 , 2023.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems , 33:1877-1901, 2020.

Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Zhifang Sui, and Furu Wei. Why can gpt learn incontext? language models secretly perform gradient descent as meta optimizers. arXiv preprint arXiv:2212.10559 , 2022.

Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jian, Bill Yuchen Lin, Peter West, Chandra Bhagavatula, Ronan Le Bras, Jena D Hwang, et al. Faith and fate: Limits of transformers on compositionality. arXiv preprint arXiv:2305.18654 , 2023.

Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems , 35:30583-30598, 2022.

Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D Lee, and Dimitris Papailiopoulos. Looped transformers as programmable computers. arXiv preprint arXiv:2301.13196 , 2023.

Surbhi Goel, Varun Kanade, Adam Klivans, and Justin Thaler. Reliably learning the relu in polynomial time. In Conference on Learning Theory , pages 1004-1042. PMLR, 2017.

Surbhi Goel, Adam Klivans, Pasin Manurangsi, and Daniel Reichman. Tight hardness results for training depth-2 relu networks. arXiv preprint arXiv:2011.13550 , 2020.

Michael Hahn and Navin Goyal. A theory of emergent in-context learning as implicit structure induction. arXiv preprint arXiv:2303.07971 , 2023.

Shima Imani, Liang Du, and Harsh Shrivastava. Mathprompter: Mathematical reasoning using large language models. arXiv preprint arXiv:2303.05398 , 2023.

Andrew K Lampinen, Ishita Dasgupta, Stephanie CY Chan, Kory Matthewson, Michael Henry Tessler, Antonia Creswell, James L McClelland, Jane X Wang, and Felix Hill. Can language models learn from explanations in context? arXiv preprint arXiv:2204.02329 , 2022.

Jack Lanchantin, Shubham Toshniwal, Jason Weston, Arthur Szlam, and Sainbayar Sukhbaatar. Learning to reason and memorize with self-notes. arXiv preprint arXiv:2305.00833 , 2023.

Nayoung Lee, Kartik Sreenivasan, Jason D Lee, Kangwook Lee, and Dimitris Papailiopoulos. Teaching arithmetic to small transformers. arXiv preprint arXiv:2307.03381 , 2023.

Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for parameter-efficient prompt tuning. arXiv preprint arXiv:2104.08691 , 2021.

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-Wei Chang, and Yejin Choi. Symbolic chain-of-thought distillation: Small models can also" think" step-by-step. arXiv preprint arXiv:2306.14050 , 2023b.

Yingcong Li, M Emrullah Ildiz, 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 a compositional rnn in a haystack. arXiv preprint arXiv:1802.06467 , 2018.

Bingbin Liu, Jordan T Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. arXiv preprint arXiv:2210.10749 , 2022.

Xinxi Lyu, Sewon Min, Iz Beltagy, Luke Zettlemoyer, and Hannaneh Hajishirzi. Z-icl: Zero-shot in-context learning with pseudo-demonstrations. arXiv preprint arXiv:2212.09865 , 2022.

Sewon Min, Mike Lewis, Luke Zettlemoyer, and Hannaneh Hajishirzi. Metaicl: Learning to learn in context. arXiv preprint arXiv:2110.15943 , 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.12837 , 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´ c, and Pablo Barceló. On the turing completeness of modern neural network architectures. arXiv preprint arXiv:1901.03429 , 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 arXiv:2102.10535 , 2021.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

Himanshu Sahni, Saurabh Kumar, Farhan Tejani, and Charles Isbell. Learning to compose skills. arXiv preprint arXiv:1711.11289 , 2017.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Petar Veliˇ ckovi´ c 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:2212.07677 , 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 on Machine Learning , pages 35151-35174. PMLR, 2023.

Johannes von Oswald, Eyvind Niklasson, Maximilian Schlegel, Seijin Kobayashi, Nicolas Zucchet, Nino Scherrer, Nolan Miller, Mark Sandler, Max Vladymyrov, Razvan Pascanu, et al. Uncovering mesa-optimization algorithms in transformers. arXiv preprint arXiv:2309.05858 , 2023.

Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: a case study on approximating turing machines with transformers. Advances in Neural Information Processing Systems , 35:12071-12083, 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. arXiv preprint arXiv:2206.07682 , 2022b.

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.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-art natural language processing. arXiv preprint arXiv: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.02080 , 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.02015 , 2023.

Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequence 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-context learning. arXiv preprint arXiv:2211.09066 , 2022.

Hattie Zhou, Arwen Bradley, Etai Littwin, Noam Razin, Omid Saremi, Josh Susskind, Samy Bengio, and Preetum Nakkiran. What algorithms can transformers learn? a study in length generalization. arXiv preprint arXiv:2310.16028 , 2023.

Appendix

Table of Contents

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 4 . We also consider that the heads are added and each one has each own key, query and value weight matrices. Formally, we have

where X ∈ R d × n , H is the number of heads used and f ( x ) = ( x ) + is the ReLU activation. We also make use of the temperature λ , which is a parameter of the softmax. Specifically, softmax ( x ) = { e λx i / ∑ j e λx j } i . Notice that as λ →∞ we have softmax ( x ) → max i x i . We also assume that the inputs are bounded; we denote with N max the maximum sequence length of the model.

Assumption 2 Each entry is bounded by some constant c , which implies that ∥ X ∥ ≤ c ′ , for some large c ′ that depends on the maximum sequence length and the width of the transformer.

A.2 Positional Encodings

In the constructions below we use a 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 we construct, require only logarithmic space with respect

to the sequence length. Notice that binary representation of the positions satisfy the following two conditions:

̸

Let r i be the binary representation of the i -th position, then there exists ε > 0 such that r ⊤ i r i > r ⊤ i r j + ε for all j = i .

There exists a one 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 a slightly different proof.

Lemma 1 A transformer with one layer, two heads, and embedding dimension of O (log 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

where x i , y i ∈ R d are column vectors for all i = 1 , . . . , n . We present how we can move a block of data from the ( d +1 : 2 d, k ) position block to the (1 : d, 1) position block, meaning to move the point y k to the position of the point x 1 . It is straightforward to generalize the proof that follows to move blocks from ( i : i + k, j ) to ( i ′ : i ′ + k, j ′ ) for arbitrary i, i ′ , j, j ′ , k .

Let in Eq. A.1 W 1 K = [ 0 0 I 0 ] and W 1 Q = [ 0 0 0 I ] and W 2 K = [ 0 0 0 I ] and W 2 Q = [ 0 0 0 I ] then we have

As λ → ∞ and after the application of the softmax operator the above matrix becomes equal to [e k e 2 . . . e k . . . e n ] + ϵ M , where e i is the one-hot vector with 1 in the k -th position, ∥ M ∥ ≤ 1 and ϵ is controllable by the temperature parameter and can be arbitrary small. Let finally

By repeating the exact same steps for the second head and letting W 2 V = [ -I 0 0 0 0 0 0 0 ] and adding back the residual we get the desired result.

A slightly different implementation can be found in Giannou et 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 2 There exists a 1-hidden layer feedforward, ReLU network, with 8 d activations in the hidden layer and d neurons in the output layer that when given two d -dimensional binary vectors representing two non-negative integers, can output the binary vector representation of their sum, as long as the sum is less than 2 d +1 .

Lemma 3 Let A ∈ R d × m and B ∈ R d × n . Then for any ϵ > 0 there exists a transformer-based function block with 2 layers, 1 head and width r = O ( d ) that outputs the multiplication A ⊤ B + ϵ M , for some ∥ M ∥ ≤ 1 .

Remark 2 Notice that based on the proof of this lemma, the matrices/scalars/vectors need to be in the same rows, i.e. , 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 input matrix is

where r (1) , r (2) are chosen as to specify the position of the result.

A.4 Results on Filtering

Lemma 4 Assume that the input to a transformer layer is of the following form

where b i , b ′ i ∈ { 0 , 1 } , with zero indicating that the corresponding point should be ignored. Then there exists a transformer TF consisting only of a ReLU layer that performs this filtering, i.e.,

̸

̸

Proof. The layer is the following:

for some large constant C . Notice that if b i = 1 the output is just x i . But if b i = 0 then the output is zero. For the second set instead of using the bits b i , we use the b ′ i .

Remark 3 Notice that if some b i is instead of 1 , 1 ± ε , ε < c/C . Then the output of the above layer would be

while if some b i = ± ε instead of zero, and assuming that x i > c > 0 or x i < -c < 0 the output would be

If | x i | ≤ c , again the output would be less than or equal to c , where c 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 is performed. An d × N -dimensional input is given to the transformer architecture, which is 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 to 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

̸

̸

̸

̸

̸

̸

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 so we assume that the input to the transformer contains the following information:

An enumeration of the tokens from 1 to N .

The ln of the above enumeration.

Zeros for the tokens that correspond to the data points, 1 for each token that is the x test or it is a prediction based on it.

An enumeration 1 to L for each one of the data points provided. For example, if we are given three sets of data we would have : 1 . . . L 1 . . . L 1 . . . L .

Some extra information that is needed to implement a multiplication step as described in Lemma 3 and to move things to the correct place.

The above information can be viewed as part of the encodings that are appended to the input of the transformer. Formally, we have

with error up to δ M , where ∥ M ∥ ≤ 1 and δ > 0 is a constant that is controlled and can be arbitrarily small.

Proof.

Step 1: Extract the sequence length (1 layer). Let

and thus

So after the softmax is applied we have

We then set the weight value matrix as to zero-out all lines except for one line as follows

After adding the residual and using an extra head where the softmax returns identity matrix and the value weight matrix is minus the identity, we get attention output

Notice that ∑ N i =1 i 2 ∑ N i =1 i = 2 N ( N +1)(2 N +1) 6 N ( N +1) = 2 N +1 3 . We then use the ReLU layer as to multiply with 3 / 2 and subtract 1 from this column. This results to attention output

Notice that this step does not involve any error.

Step 2: Extract the identifier of the prediction to be made ℓ (3 layers). Now, by setting the key and query weight matrices of the attention as to just keep the all ones row, we get a matrix that attends

.

equally to all tokens. Then the value weight matrix keeps the row that contains ℓ ones and thus we get the number ℓ N , propagated in all the sequence length. Then the attention output is as follows:

So far this step does not involve any error. As described in Lemma 3 to implement multiplication of two values up to any error ϵ , we need 1) have the two numbers to be multiplied next to each other and an extra structure (some constants, which we consider that are encoded in the last rows of the matrix m i ) and 2) the corresponding structure presented in Lemma 3 . So we need to make the following transformation in the rows containing N and ℓ N

where ∗ denotes inconsequential values. For the first step we assume that we have the necessary binary representations 5 and then we use Lemma 1 which shows how we can perform the operation of copy/paste. This step involves an error and the current output would be ˜ X 1 = X ∗ + ε 1 M 1 , where ε 1 is controlled and we will determine it in the sequence, ∥ M 1 ∥ ≤ 1 and X ∗ 1 is the desired output. For the second step now we use Lemma 3 to perform the multiplication, this will affect some of the other values. To analyze the error of this step, we know that given X ∗ 1 we will get the desired output X ∗ 2 + ε 2 M 2 , where ∥ M 2 ∥ ≤ 1 and ε 2 is to be determined. However, are given ˜ X 1 in the multiplication procedure 6 which results in the following output ˜ X 2 = ( ˜ X 1 ) ⊤ ˜ X 1 + ε 2 M 2 = X ∗ 1 ⊤ X ∗ 1 + ε 1 M ⊤ 1 X ∗ 1 + ε 1 M ⊤ 1 M 1 + ε 2 M 2 . Thus by choosing ε 1 , ε 2 to be of order O ( δ ) we get that ˜ X 2 = X ∗ 2 + δ C ∥ M ∥ for some C chosen to be large enough. Then for the last step consider the follow (sub-)rows of the matrix X

where r i is the binary representation of position i . By choosing W K , W Q as to

By subtracting one identity head for the first that we focus on as described in Lemma 1 we have that attn ( X ) results in the desired matrix. We output this result and overwrite ℓ/ N, thus we have

The new error introduced by this operation is again controlled as to be ε 3 M 3 = δ C , with ∥ M 3 ∥ ≤ 1 .

Step 3: Create ℓ +1 (0 layer). We first copy the row with ℓ to the row with N , this can trivially be done with a ReLU layer that outputs zero everywhere else except for the row of N s that output ( ℓ ) + -( N ) + to account also for the residual. We now use one of the bias terms (notice that ℓ is always positive) and set to one in one of the two rows that contain the ℓ , again we account for the residual as before; everything else remains unchanged. Thus, we have

This operation can collectively be implemented (add the bias + copy the row) in the ReLU layer of the previous transformer layer that was not used in the previous step.

Step 4: Create the binary bits (2 layers). We will now use the information extracted in the previous steps, to create the binary indicators/bit to identify which tokens we want to filter. This can be easily implemented with one layer of transformer and especially the ReLU part of it. Notice that if we subtract the row that contains the tokens that have already been predicted, i.e., [ ℓ ℓ . . . ℓ ] from the row that contains [1 2 . . . L 1 2 . . . L . . . L ] we will get zero only in the positions that we want to filter out and some non-zero quantity in the rest. This is trivially implemented with one ReLU layer. So, we need to implement an if..then type of operation. Basically, if the quantity at hand is zero we want to set the bit to one, while if it non-zero to set it to be zero. This can be implemented with the following ReLU part of a transformer layer

the last two terms are to account for the residual. Again the rest of the rows do not change and are zeroed-out.

Step 5: Implement the filtering (1 layer). We now apply Lemma 4 and our proof is completed. In this step as in Lemma 4, the error remains of the order of the error of the previous step. Thus, by fine-tuning the constants appropriately, depending on 1) the bound on the input as of Assumption 2, the targeted error δ and the constant used in Lemma 4 we achieve that the error has the target upped bound.

Theorem 2 (Theorem 1 restated) Consider a prompt p n ( f ) generated from an L -layer MLP f ( · ) as described in Definition 1, and assume given test example ( x test , s 1 test , . . . s L test ) . For any resolution ϵ > 0 , there exists δ = δ ( ϵ ) , iteration choice T = O ( κ 2 max log(1 /ϵ )) , and a backend transformer construction TF BE such that the concatenated transformer TF = TF LR ◦ TF BE implements the following: Let (ˆ s i ) ℓ -1 i =1 denote the first ℓ -1 CoT-I/O outputs of TF and set p [ ℓ ] = ( p n ( f ) , x test , ˆ s 1 . . . ˆ s ℓ -1 ) . At step ℓ , TF implements

Filtering. Define the filtered prompt with input/output features of layer ℓ ,

There exists a fixed projection matrix Π that applies individually on tokens such that the backend output obeys ∥ Π ( TF BE ( p [ ℓ ])) -p filter n ∥ ≤ δ .

Gradient descent. The combined model obeys ∥ TF ( p [ ℓ ]) -s ℓ test ∥ ≤ ℓ · ϵ/L .

TF BE has constant number of layers independent of T and n . Consequently, after L rounds of CoT-I/O, TF outputs f ( x test ) 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 to T , n and ℓ . Next we turn to prove (2. Gradient descent). In Definition 1, we assume that the network's activation function is leaky-ReLU, i.e.,

Thus, as a first step we construct the inverse of leaky-ReLU and apply it in the second row of p filter n where the inverse of leaky-ReLU is

This can be implemented with the following activation function (denoted by σ ( · ) ) using ReLUs:

After it, it remains TF LR to solve linear regression problems. Taking ℓ th layer, first neuron prediction as an example, and letting x ′ i := s ℓ -1 i , y ′ i := ϕ -1 ( s ℓ i [0]) and w = W ℓ [0] , linear regression has form of y ′ i = w ⊤ x ′ i for i ∈ [ n ] . Notice that the extra zeros do not contribute in the update performed by gradient descent, and after gradient descent has been performed, we apply back the leaky ReLU. Then following Assumption 1, since we assume TF LR performs the same as gradient descent optimizer,

given matrix condition as described in Definition 1, after running T iterations of gradient descend on the linear regression problem and considering each layer prediction with resolution ϵ/L , we can get that ∥ TF LR ( p [ ℓ ]) -¯ s ℓ ∥ ≤ ϵ/L , where ¯ s ℓ = ϕ ( W ℓ ˆ s ℓ -1 ) is the correct prediction if taking ˆ s ℓ -1 as input. Then we have

Let TF LR ( p [0]) returns x test and therefore ∥ TF LR ( p [0]) -x test ∥ = 0 . Combing results in that ∥ TF LR ( p [ ℓ ] -s ℓ test ) ∥ ≲ ℓ · ϵ/L . Since from Lemma 5 we have that we can choose δ to be arbitrary. Let δ = ϵ/L , where L is the total predictions we will make. Then it will result in ∥ TF ( p [ ℓ ] -s ℓ test ) ∥ ≲ ℓ · ϵ/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-I and ˆ θ CoT-I/O . Next we make predictions following Section 2.2. Letting ℓ ( · , · ) : Y × Y → R be loss function, we can define test risks as follows.

and

Here, we use L ( n ) to define the test risk when given prompt with n in-context samples. Then, results shown in Figures 3&4&5&6&7(a) are test risks L ( n ) given n ∈ [ N ] . Following model training and evaluation, we can see that once loss function is the same for both training and predicting, CoT-I (as well as ICL) accepts training risk L CoT-I train = 1 N ∑ N n =1 L CoT-I ( n ) .

B.2 Implementation

All the transformer experiments use the GPT-2 model [Radford et al., 2019] and our codebase is based on prior works [Garg et al., 2022, Wolf et al., 2019]. Specifically, the model is trained using the Adam optimizer with learning rate 0 . 0001 and batch size 64 , and we train 500 k iterations in total for all ICL, CoT-I and CoT-I/O methods. Each iteration randomly samples inputs x s and functions f s. We also apply curriculum learning over the prompt n as Garg et al. [2022] did for 2-layer random MLPs (Sec. 4.2). For both training and testing, we use the squared error as the loss function, i.e. , ℓ (ˆ y , y ) = ∥ ˆ y -y ∥ 2 (or ℓ (ˆ y, y ) = (ˆ y -y ) 2 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-ofdistribution (OOD) settings.

To provide a clearer understanding of our model's behavior under distribution shifts and to quantitatively assess the impact on test risk, we conduct experiments with varying levels of distribution shift. Specifically, we evaluate the test risk when the prompt is supplemented with 100 in-context examples. The results of these experiments are depicted in Figure 9, which serves to illustrate how our model's performance is affected under these OOD conditions. In Fig. 9(a), we analyze noisy in-context samples during testing. The solid and dashed curves represent the test risks, corresponding to the noisy in-context samples whose (input, output) takes the form of either ( x , y + noise ) or

0.0

0.1

0.2

0.3

0.4

0.5

Noise level (scaled)

0.0

0.1

0.2

0.3

0.4

Test risk

CoT-I/O

CoT-I

ICL

(a) Test with noisy data

1

2

3

4

5

6

7

8

k

0.01

0.02

0.03

0.04

0.05

0.06

0.07

Test risk

CoT-I/O

CoT-I

ICL

(b) Test with different

k

2

4

6

8

10

d

0.00

0.02

0.04

0.06

0.08

0.10

Test risk

CoT-I/O

CoT-I

ICL

(c) Test with different

d

0

25

50

75

100

125

150

175

200

# in-context samples

0.0

0.1

0.2

0.3

0.4

Test risk

CoT-I/O

CoT-I

ICL

CoT-I/O

CoT-I

ICL

(a) Small GPT-2,

d

= 10

,

k

= 4

0

25

50

75

100

125

150

175

200

# in-context samples

0.0

0.1

0.2

0.3

0.4

Test risk

CoT-I/O

CoT-I

ICL

CoT-I/O

CoT-I

ICL

(b) Standard GPT-2,

d

= 10

,

k

= 4

( x + noise , y ) , respectively. The results indicate that CoT exhibits greater robustness compared to ICL, and the test risks increase linearly with the noise level, with attributed to the randomized MLPs setting. Additionally, in Figs. 9(b)&9(c), we instead explore out-of-distribution test tasks where test MLPs differ in ( d, k ) from the training phase. For both subfigures, we firstly train small GPT-2 using 2-layer MLPs with d = 10 , k = 8 . In Fig. 9(b), we fix d = 10 and vary k from 1 to 8 , whereas in Fig. 9(c), we fix k = 8 and vary d from 1 to 10 . In both instances, the findings reveal that CoT's performance remains almost consistent when k ≥ 4 or d ≥ 6 , and ICL is unable to surpass it. The improved performance of ICL with smaller values of d or k again reinforces our central assertion that ICL requires O ( dk ) samples for in-context learning of the 2-layer random MLP, and reducing either d or k helps in improving the performance. Given that we employ the ReLU activation function, smaller values of d or k can lead to significant bias in the intermediate feature. Consequently, CoT cannot derive substantial benefits from this scenario, resulting in a decline in performance.

C.2 Further Evidence of Model Expressivity

There are two determinants of model learnability in in-context tasks: sample complexity and model expressivity. Sample complexity pertains to the number of samples needed to precisely solve a problem. However, when the transformer model is small, even with a sufficiently large number of samples, due to its lack of expressivity, ICL cannot achieve zero test risk. This contrasts with CoT, which decomposes complex tasks into simpler sub-tasks, thereby requiring smaller models for expression. Figure 5 has illustrated the expressivity of different GPT-2 architectures, showing that the tiny GPT-2 model is too small to express even a single layer of 2-layer MLPs. Additionally, we have run more experiments, and the results are shown in Figure 10. Both Figures 10(a) and 10(b)

0

20

40

60

80

100

# in-context samples

0.0

0.2

0.4

0.6

0.8

Test risk

CoT

CoT-zero

CoT-random

(a) Test filtering on first layer prediction

0

20

40

60

80

100

# in-context samples

0.0

0.1

0.2

0.3

0.4

Test risk

k

=4

k

=8

k

=16

d

=10

d

=20

(b) CoT-I/O: composed risk (same as Fig. 3(a))

0

20

40

60

80

100

# in-context samples

0.0

0.1

0.2

0.3

0.4

Test risk

k

=4

k

=8

k

=16

d

=10

d

=20

(a) CoT-I/O: 1st layer prediction

0

20

40

60

80

100

# in-context samples

0.0

0.1

0.2

0.3

0.4

Test risk

k

=4

k

=8

k

=16

d

=10

d

=20

(b) Train ICL with

(

x

,

s

)

pairs

0

20

40

60

80

100

# in-context samples

0.0

0.1

0.2

0.3

0.4

Test risk

k

=4

k

=8

k

=16

d

=10

d

=20

(c) CoT-I/O: 2nd layer prediction

0

20

40

60

80

100

# in-context samples

0.0

0.1

0.2

0.3

0.4

Test risk

k

=4

k

=8

k

=16

d

=10

d

=20

(d) Train ICL with

(

s

, y

)

pairs

detail training models with MLP tasks of dimensions d = 10 and k = 4 . In Fig. 10(a), we use a small GPT-2 model, and the results show that the test risk stops decreasing even with more in-context examples. In Fig. 10(b), we train a larger model, and the results demonstrate that the standard GPT-2 is sufficient to express a 2-layer MLP with d = 10 and k = 4 .

C.3 Filtering Evidence in 2-layer MLPs

Section 5 has demonstrated the occurrence of filtering in the linear deep MLPs setting (black dotted curves in Fig. 8(b)). In this section, we present further empirical evidence based on the 2-layer MLPs setting discussed in Section 4.2.

Follow the same setting as Figure 3(a) and choose d = 10 and k = 8 . Assume we have a model pretrained using CoT-I/O method. As described in Section 4.2, during training, the prompt consists of in-context samples in the form of ( x , s , y ) where s = ( Wx ) + and y = v ⊤ s . To investigate filtering,

0

20

40

60

80

100

# in-context samples

0.0

0.1

0.2

0.3

0.4

Test risk

# layer=3

# layer=6

# layer=3

# layer=6

(a) Varying GPT-2 layers

0

20

40

60

80

100

# in-context samples

0.0

0.1

0.2

0.3

0.4

Test risk

# head=2

# head=4

# head=8

# head=2

# head=4

# head=8

(b) Varying GPT-2 heads

0

20

40

60

80

100

# in-context samples

0.0

0.1

0.2

0.3

0.4

Test risk

# embedding=64

# embedding=256

# embedding=64

# embedding=256

(c) Varying embedding dimensions

we make three different predictions to evaluate the intermediate output, whose test prompts have in-context examples with the following forms:

where z ∼ N (0 , d ) . The results are displayed in Figure 11(a) where blue, orange and green curves represent first layer prediction results using CoT, CoT-zero and CoT-random prompts, respectively. From this figure, we observe that the three curves are well aligned, indicating that when making a prediction for input x , TF will attend only to ( x , s ) and ignore y . Therefore filling the positions of y with any random values (or zero) will not change the performance of first layer prediction.

C.4 Comparison of Filtered CoT with ICL

Until now, many experimental results have shown that CoT-I/O provides benefits in terms of sample complexity and model expressivity compared to ICL. As an interpretation, we state that CoT can be decoupled into two phases: Filtering and ICL , and theoretical results have been provided to prove this statement. As for the empirical evidence, Sections 5 and C.3 precisely show that filtering does occur in practice. In this section, we provide additional experiments to demonstrate that, after filtering, CoT performs similarly to ICL.

For convenience and easier comparison, we repeat the same results as Fig. 3(a) in Fig. 11(b), where d ∈ { 10 , 20 } , k ∈ { 4 , 8 , 16 } , and train with a small GPT-2. We again recap the data setting for the 2-layer MLP, where the in-context examples of CoT prompt are in the form of ( x , s , y ) . Given that filtering happens, we make first and second layer predictions following Section 4.1 and results are presented in Fig. 12(a) and Fig. 12(c), respectively. These results show the performances of the filtered CoT prompts. Next, we need to compare the performance with separate ICL training. To achieve this goal, we train a small GPT-2 model using ICL method with prompt containing ( x , s ) pairs (first layer). The test results are shown in Fig. 12(b). Additionally, in Fig. 12(d), we train another small GPT-2 model using ICL but with prompts containing ( s , y ) pairs (second layer). By comparing Fig. 12(a) and 12(b), as well as Fig. 12(c) and 12(d), we observe that after filtering, CoT-I/O achieves similar performance as individually training a single-step problem through the ICL phase.

C.5 CoT across Different Sizes of GPT-2

In Figure 5, we have demonstrated that larger models help in improving performance due to their ability of solving more complex function sets. However, since tiny, small and standard GPT-2 models scale the layer number, head number and embedding dimension simultaneously, it is difficult to determine which component has the greatest impact on performance. Therefore in this section, we investigate how different components of transformer model affect the resulting performance by run CoT-I/O on various GPT-2 architectures.

We maintain the same setting as in Section 4.2, fix d = 10 and k = 8 , and consider a base GPT-2 model (small GPT-2) with 6 attention layers, 4 heads in each layer and 128 -dimensional embeddings.

0

20

40

60

80

100

# in-context samples

0.0

0.2

0.4

0.6

0.8

Test risk

k

=4

k

=16

k

=64

Gradient descent

0

20

40

60

80

100

# in-context samples

0.0

0.2

0.4

0.6

0.8

1.0

Test risk

k

=4

k

=16

k

=64

GPT-2

Least squares

(a) 1st layer: compare with GD solving ReLU

(b) 2nd layer: compare with least squares

In Fig. 13(a), we fix the number of heads at 4 and the embedding dimension at 128 , while varying the number of layers in { 3 , 6 , 12 } . Similarly, we explore different models with different numbers of heads and embedding dimensions, and the results are respectively presented in Fig. 13(b) and 13(c). Comparing them, we can observe the following: 1) once the problem is sufficiently solved, increasing the model size does not significantly improve the prediction performance (see Fig. 13(b)&13(c)); 2) the number of layers influences model expressivity, particularly for small GPT-2 architecture (see Fig. 13(a)).

C.6 Comparison of Transformer Prediction and Linear Regression

We also provide experimental findings to verify Assumption 1 in this section. Previous work [Giannou et al., 2023, 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

for some w ⋆ ∼ N (0 , 2 I d ) and n is the training sample size. Then the normalized test risks are computed by L ( n ) = E w ⋆ , x [ ∥ ( ˆ w ⊤ n x ) + -y ∥ 2 ] /d , and we show point-to-point results for n ∈ [ N ] in black dashed curve in Fig. 14(a) 7 . As for the second layer, we solve least squares problems as follows

for some W ⋆ ∈ R k × d ∼ N (0 , 2 /k ) and v ⋆ ∼ N (0 , I k ) . Here, † represents the pseudo-inverse operator. Then we calculate the normalized test risk of least square solution (given n training samples) as L ( n ) = E W ⋆ , v ⋆ , x [ ∥ ˆ v ⊤ n s -y ∥ 2 ] /d where s = ( W ⋆ x ) + and y = v ⋆ ⊤ s . The results are presented in Fig. 14(b) where blue, orange and green dashed curves correspond to solving the problem using different values of k ∈ { 4 , 16 , 64 } . In this figure, the curves for k = 4 , 16 are aligned with GPT-2 risk curves, which indicates that TF can efficiently solve linear regression as a least squares optimizer. However, the curve for k = 64 does not align, which can be attributed to the increased complexity of the function set with higher dimensionality ( k = 64 ). Learning such complex functions requires a larger TF model.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nSee the summary part.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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’.\n\nLIMITATIONS:\nI’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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-22"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\nSTRENGTHS:\nI 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.\n\nWEAKNESSES:\nAlthough 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.\n\nQUESTIONS:\nMy 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.\n\nLIMITATIONS:\nI 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-29"}

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

STRENGTHS:
- The paper provides both experimental and theoretical evidence to support its claims about how CoT works.
- 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.
- The main contribution of the paper is not clear. The paper claims to dissect the mechanism of CoT

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.
- 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThis 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-21"}