Chain-of-Thought Provably Enables Learning the (Otherwise) Unlearnable
Authors: Chenxiao Yang
Zhiyuan Li
David Wipf
Venue: ICLR 2025 Poster (2025)
Keywords: large language models
learning theory
chain-of-thought
Primary area: foundation or frontier models, including LLMs

Abstract: Modern language models have demonstrated remarkable reasoning capabilities by using chain-of-thought (CoT). One hypothesis about the inner workings of CoT is that it breaks down originally complex tasks into smaller subtasks that are more amenable to learning. We formalize this notion by showing possibility and impossibility results of learning from in-context demonstrations with and without CoT. In particular, with CoT, we examine a family of learning algorithms that learn a task step-by-step, capable of composing simpler functions from individual reasoning steps to form an overall complex function. This process reduces the difficulty of learning a task to that of the hardest reasoning step in the chain. Moreover, we prove Transformers can express this algorithm and thus they can efficiently in-context learn arbitrary tasks as long as these tasks can be decomposed into a finite number of subtasks, each of which are efficiently learnable. In contrast, without CoT, we demonstrate that there exist tasks that are inherently unlearnable by the same algorithm. Overall, our results suggest several provably effective ways for decomposing target problems to instantiate CoT. Empirically, we demonstrate our proposed CoT construction significantly enhances the reasoning capabilities of real-world LLMs in solving challenging arithmetic reasoning tasks, including learning polynomials and Boolean formulas.

Published as a conference paper at ICLR 2025
CHAIN-OF-THOUGH PROVABLY ENABLES LEARNING THE (OTHERWISE) UNLEARNABLE
Chenxiao Yangl; Zhiyuan Li', David Wipf i Toyota Technological Institute at Chicago, & Amazon Web Services
(chenxiac, zhiyuanlsj@ttic.et davicwipf@gmail.
ABSTRACT
Modern language models have demonstrated remarkable reasoning capabilities by using chain-of-thought (CoT). One hypothesis about the inner workings of CoT is that it breaks down originally complex tasks into smaller subtasks that are more amenable to learning. We formalize this by showing possibility and impossibility results of learning from in-context demonstrations with and without COT. In partieular, with COTT, we examine a family of learning algorithms that learn a task step-by-step, capable of composing simpler functions from individual reasoning steps to form an overall complex function. This process reduces the difficulty of learning task that of the hardest reasoning step in the chain. Moreover, we prove Transformers can express this algorithm and thus they can efficiently in-context learn arbitrary tasks as long these tasks can be decomposed into finite number of subtasks, each of which are efficiently learnable. In contrast, without CoT, demonstrate that there exist tasks that inherently unlearnable by the same algorithm. Overall, our results suggest several provably effective ways for decomposing target problems to instantiate COT. Empirically, we demonstrate our proposed CoT construction significantly enhances the reasoning capabilities of real-world LLMs in solving challenging arithmetic reasoning tasks, including learning polynomials and Boolean formulas.
INTRODUCTION
Complex problem solving often involves breaking down an originally challenging task into smaller, more manageable subtasks, learning from these subtasks, and then composing the acquired skills to address the overall task - a strategy that reflects how humans naturally solve problems. One method that mimics this process is called al., 2022, Reynolds & McDonell, [02T Nye et al.) (021), whereby a model is provided with demonstrations involving detailed reasoning steps and subsequently instructed to generate thoughts step-by-step before yielding the final answer. Modern language models rely on CoT or variants thereof Yao et al. 2024 Besta et al. 2024) to tackle complex tasks ranging from commonsense reasoning to mathematical proofs (Cobbe et al. 2021; Rae et al., 2021: Srivastava et al. 2022). sometimes even exceeding the capabilities of human experts.
(CoT) (Wei)
Despite the empirical success of CoT, theoretical investigation thus far has remained relatively sparse. Some previous efforts have made strides from the perspective of Transformer expressiveness (Li et al. 2024 Feng et al. (2024) or through case studies of in-context learning MLPs (Li et al. 2023b). However, these works either focus on mostly generic cases without specifying conerete/actionable ways of actually decomposing a task, or involve quite restricted scenarios where both the task and the CoT are predetermined. Since in practice not all intermediate steps in CoT are equally useful, e.g. Zhang et al. (2022); Press et al. (2022), and the target tasks are typically diverse, there still exists a considerable gap berween current theory and practical scenarios. targeting this gap, we ask the following core question:
How do task decompositions affect the ability of a model to learn complex reasoning tasks, and what principles can guide the corresponding optimal CoT design?
"Work was done during author's internship at Amazon Web Services.

---

Published as a conference paper at ICLR 2025
Learn Step-by-Siep (Alg. 1) Expressible by Transformer? (v Lem. 1)
Hurdness of learning (Thm. 4 & Cor. 5)
Learner
A
Input
Outpat
w/ CoT
Input
Step
Step
Output
w/o CoT
(Bottleneck)
Step
The hardest subtask is efficiently learnable
How to form effective CoT? (Experiments) Figure I: An overview of our analysis. Arrow length indicates the difficulty of learning step.
+
The overall tack is efficiently leornable
(Lem. Thm. 3)
One possible explanation for the efficacy of CoT is that it reduces the difficulty of learning a complex task the level of learning a series of subtasks, The implication here is simply that, assuming these subtasks are suitably designed and orchestrated, then the necessary skills acquired at the subtask level through CoT will suffice to outperform attempts at learning directly from the overall task. In this work, we formalize these intuitions in mathematically rigorous way, and take initial steps towards quantifying the benefits of different task decomposition schemes.
Problem Setup. To define notion of learnability in the context of language modeling, we adopt the classic PAC learnability (Valiant. 1984) in the setting where the learner is associated with a parametric model - a setting known as in-context learning (Brown et al. 2020; Garg et al. (022). Unlike conventional supervised learning, the learning considered here occurs during test time, which highly useful feature enabling the model to adapt to new tasks that were not explicitly seen in pre-training. Concretely, we say a task is in-context learnable by a parametric model M if there exists parameter configuration 0 such that, for arbitrary distribution from the task, the model can, upon receiving set of i.i.d. demonstrationnssamplee D) and a query/test sample output a prediction for the query, i.e. Ma( Mg(Dxx), x),, that is close the ground-truth y with high probability. learning algorithm A is implicit in this process, which takes D as input and outputs predictor or hypothesis hp : X + y. The relation between the algorithm and the model Mg can be written as
(D, = (x) = hn())
To account for the effects of CoT,, we also formalize the notion of 'task decomposition" as transforming the distribution from which 1) generated into a sequence of distributions that generate D with detailed intermediate steps. (Section
1.1 CONTRIBUTIONS
Main Results. To formally answer whether or not CoT can help a model learn complex reasoning tasks, investigate the in-context learnability of tasks w.r.t. different decomposition schemes. We find the answer is often affirmative depending on the task decomposition and summarize our findings via the following two informal statements:
1. Regardless of how complex reasoning task is, it can be efficiently learned by Transformers in context as long as it can be decomposed into finite number of reasoning steps, each of which is efficiently learnable by learner Transformers can perform. (Section 3
2. There exist inherently hard tasks that are not learnable without CoT regardless of the sample size, and yet nonetheless become learnable by using CoT with specific decomposition schemes we introduce. (Section
In aggregate, our results suggest a broad range of scenarios where CoT can indeed effectively reduce the hardness of learning, from that of the overall task to that of the hardest reasoning step in the chain, or even from an unlearnable level to the learnable level. These results also suggest several actionable ways to form effective intermediate steps of CoT.
To obtain the above results, we introduce a class of learning algorithms ACoT enabled by CoT, dubbed step-by-step learning. This class of algorithms takes as input CoT examples, and outputs a complex predictor hp by composing predictors (h:) [he]eek]) obtained from individual algorithms, where is the number of reasoning steps (Algorithm []. Given fixed k. the expected overall prediction error made by hD can be upper bounded by the individual errors made by predictors {h. on their respective reasoning steps (Lemma 2. Leveraging this fact, prove that the difficulty of learning the overall task can be reduced to that of the hardest constituent step of CoT, since the sample complexity of this step determines the sample complexity of learning the overall task (Theorem 3).
Furthermore, we establish that the capabilities of AçOT described above can be achieved by Trans- formers the de facto parametric model M used in language modeling. Specifically, we show
2

---

Published as a conference paper at ICLR 2025
that linear depth (w.rt. k) and constant embedding size Transformer is sufficient to express Acor in an end-to-end manner, if the individual algorithms for learning each step are instantiated as kernel gradient descent that produces predictor h, : + W Wo(z) with non-linear feature map o(+) (Lemmal) D. Combining these results suggests that a reasoning task is efficiently learnable by Transformers in-context as long each subtask is efficiently learnable by these individual algorithms.
Within the same analytical framework, we further present impossibility results illustrating the hardness of learning without COT. Particularly, we consider the task of learning a Boolean function class called sparse parity. Without CoT, this task is impossible to learn by the variant of algorithm ACOT with k 1 due limited approximation power of the hypothesis class that is used to learn a single step. whereas if we enable CoT with specific intermediate reasoning steps, the task becomes learnable with guarantees of smaller errors (Theorem & Corallary 5
Experiments. For empirical verification, we consider new arithmetic reasoning tasks and test if the decomposition schemes from our analysis are practically effective (Section [5]. Specifically, we propose to construct complex reasoning tasks with varying overall hardness and hardness of subtasks. We observe that the performance of real-world LLMs improve significantly as the difficulty of the hardest step in the CoT reduces, regardless of the overall task complexity. Additionally, on learning two inherently hard Boolean functions, we demonstrate that our introduced CoT significantly enhances reasoning performance, sometimes improving accuracy from nearly random guessing to nearly perfect.
1.2 RELATED WORK
CoT has been analyzed through the lens of the expressiveness of finite- seitee-reeision Transformers Feng et al. (2024); ;; Li et al. (2024). In particular, Feng et al. (2024) proves CoT can enable Transformers to solve some specific tasks such as basic arithmetic and linear equations, while Li et al. (2024) proves that increasing the step number in the CoT allows them emulate circuits of increasing depth. More closely related to ours, Li et al. (203b) studies a scenario in which the intermediate steps of CoT are intermediate layers of MLP, showing that Transformers can in-context learn the MLP. We also note that there are many recent papers that empirically analyze CoT, e.g. (Wang et al. 2022 Fu et al.
2023b Madaan & Yazdanbakhshi [022] (Turpin et al.] 2024] Prabhakar et al.] (2224) among others.
However, previous works are generally limited in one of two ways: 1) either the problem setting overly restricted, focusing on fixed tasks and specific forms of decomposition, such as solving arithmetic equations (Feng al.) 2024) or learning MLPs (Li et al. 2023b), 2) in more complex regime (Li ct al. 2024), where though the possibility for improvement is shown, the effects of specific intermediate steps are not quantified. In contrast, our work introduces new setting, one which is general enough to encompass learning all distribution families, while also explicitly accounting for the effects of specific decomposition schemes. Our contribution is also related to recent line of work on
ICL (Akyürek et al.
20233 Von Oswald et al., 2023 Dai et al. 2023; von Oswald et al.) 2023] Cheng al. 2024] Li al. (202a)), partieularly demonstrations of how Transformers with certain weight
constructions perform ICL similarly to optimization algorithms like gradient descent (Von Oswald
al.] (202)) or its variants (Giannou et al.) 2024] Fu et al. 2023a). Notably, [henng al. (2024)
proves that Transformers could implement kernel gradient descent, which is closely related to ours. The idea of task decomposition and generalizing from easier to harder tasks has also been explored in
other works of learning theory, e.g. Natarajan 1989): Tadepalli (2008): Wies et al. (2022).
2 PRELIMINARIES
Chain-of-Thoupg (CoT). We follow the commonly-adopted few-shot CoT setting, where demon- strations are augmented with intermediate steps and the prediction is also in the format of COT. In particular, for k reasoning steps, we denote each demonstration as
e (I, 21, 9))
v) E X Zi... Ze-1 y
(2)
where 2 Z, represents the t-th intermediate reasoning step. For convenience, also let 20 and 24 = y. Let d(Z,) be the dimension of Z,, and d Σ d(Z,).
In-Context Learning (ICL), In ICL, the base model is provided with N demonstrations or in- context examples et = ()) for E [N] where x E X and y = f(x) E y. We denote learning algorithm learner A (X' x ))N H. which takes set of demonstrations
3

---

Published as a conference paper at ICLR 2025
as imput and outputs predictor hypothesis h X' + y from hypothesis class H. Given set of demonstrations and a query r((+1) (which we assume is from the same distribution as the demonstrations), the goal of the base model is to learn predictor and use this predictor to make predictions on the query x((+1), For base model Transformer TFo(-) with parameters 8, we have
(F ())) :iE [N]], xN+1)) A(f(=")) ()) : ((())) h(x(N+1)). (3)
Transformers. Let E = {e(() : E [N]} E R2 be the concatenation of samples, and let e(N++) = (1), 0) E Rd whose dimension aligns with other examples. Following prior
work (Von Oswald et al. 2023) Ahn et al.. 2023 Cheng et al. 2024, Zhang al. 2023), single-head self-attention layer with weights WK, Wo, Wv E A updates e(M+1) as
e + ((N+1) + WyEooEE W W Woe(N+1))
where a is non-linearity that could be specified as Softmax, ReLU or some kemel functions, e.g. (Choromanski et al. | Katharopoulos et al., 2020 Wang et al.) 2020 Peng et al. [021T) Stacking multiple self-attentin layers (with or without an MLP module applied between each self- attention layer) gives us the Transformer considered in this paper. Following (Von Oswald et al.)
2023), and for subtle technical reasons related to the construction in Sec 3.2 we exclude the query token when computing the attention.
Our Setup. We are now ready to formally define the problem setup. Let D be distribution over X. and f: + V/ target function. An input distribution and target function pair (f, D) defines the generating process of in-context examples, namely examples are drawn based Let P be a family of distributions defined as a set of (f. D) pairs, representing a certain task a model aims to solve. For instance, the target functions in P could be defined all polynomials, Boolean functions, etc. The following error quantifies how successfully an algorithm can learn the task:
D, y f()).
A)
mAx (J,D)E
EEDD (1 (h(x), ((:))
where is the squared loss (which could be extended to other convex loss functions). h is the predictor given by a learner A on N i.i.d. examples from D). Minimizing this error guarantees successful learning of all target functions within the family P.
Definition 1 (In-Context Learnability). We say parametric model M 0 + Ma (i.e. a functional mapping from parameter space to function space) can learn task P there exists a learning algorithm A and function NA (0, 1) + N. such that for any confidence and accuracy parameters 5, e. E (0, 1):
1. Under a certain parameter choice O, we have Mo(;,z) = A())(x) for any query x.
2. Given (6 E) i.idd. examples, the algorithm A returns predictor h such that with probability of at least 1 1 6, A(P,A) S.
Moreover, we say P can be efficiently in-context learned by .M if both the running time of A and the sample size (ā e) are polynomial in 8 and -1 Note that the notion of learnability here is different from the classic PAC learnability (Shalev-Shware per in our definition, the model itself acts not as a predictor, but a learner. This is a highly desirable property for modern language models as it allows them to meta-learn nut-of-distribution tasks that were not available during pre-training (Brown et al. 2020). In the rest of this paper, we stipulate the model Transformer TF, & as defined via (4). unless otherwise stated.
& Ben- BaviD 2014) in the sense that,
Moving forward, we focus on whether or not CoT can improve the in-context learnability of different reasoning tasks. In principle this can be studied from multiple vantage points, such the the effect of CoT on the sample efficiency (8, e). or the existence of cases where a task is initially not learnable but becomes so once CoT enabled, etc. Notably, addressing these issues depends critically on the specific forms of intermediate CoT steps involved. To accommodate this aspect, we next formalize the notion of task decomposition.
Definition 2 (Task Decomposition). A decomposition operator T is such that, for every (J,D) E P, a target function can be decomposed as T(J) = (J2, fi) subject to f= =2 Sú, where fi X + Z and fa Z + for another space Z. This operation induces two new distribution families
P ((fi.D): (f2 fr. D) E PB and P2 = (((,f(())) (f2 fi, D) E PB
(6)
where {-B is multiset allowing repeating elements, and D' 2 + R is determined by S1 and D.
4

---

Published as a conference paper at ICLR 2025
The decomposition always exists and is not unique (e.g. fi can be arbitrary bijection). Particularly when fi is specified as the identity mapping, we have P2 = P, in which case CoT is unlikely to work. For each decomposition,
associate it with the generating process of demonstrations (denoted as 2), i.e. x ~ D, 2 = (i)).. y = f2(z). This affects how examples are generated in Definition) 1 and thus could impact the learnability of task. Note that this definition be generalized CoT with & > 2 steps by sequentially applying the decomposition.
3 IMPROVED LEARNABILITY BY TASK DECOMPOSITION
To demonstrate how CoT improves learnability by task decomposition, will first define learning algorithm Acor (Sec 3.1 and prove that it can be expressed by Transformers with linear depth (Sec (S2c Then, derive upper bound for the overall error, which allows us show CoT can adjust the learnability of a complex task to the learnability of simpler subtasks (Sec 3.3
3.1 LEARNING ALGORITHMS ENABLED BY CHAIN-OF-THOUGHTS
Consider a class of learning algorithms AC.T enabled by CoT, which involves several individual algorithms {A, 1 and learns increasingly complex compositional functions with more CoT steps:
Algorithm 1: Step-by-Step Learning with CoT (ACot)
( Input: Demonstrations {(=0) 2(() ,. = } Output: Predictor h X' -
Jij= individual learning algorithms (A.)
for 1,... do h4 + A.((2 )
hy ...... h2 h1
This class of algorithms take input of demonstrations, where each demonstration contains reasoning steps, and outputs predictor X' ~ y. The learning is performed in a step-by-step manner; that is, for each reasoning step i [kt], an individual algorithm A, is used to learn predictor h; : Z; + - E Hi The learned predictors h1, h22 -.2.-.. , are then composed to obtain the desired overall predictor. Note that the algorithm can be naturally extended to scenarios where each step is function of all preceding steps, by redefining 2, in the algorithm as concatenation of {2)bis in the initial CoT. Therefore, without loss of generality, we assume that the CoT satisfies the Markov property, meaning that each step is conditionally dependent only on the last step.
3.2 EXPRESSIVENESS OF TRANSFORMERS
Next, we demonstrate parameter choices 0 that connect Transformers TF, and algorithm AçoT instantiated a certain way.
Instantiation of ACoT. While the algorithm could have many different instantiations, in this paper, we define A, as empirical risk minimization: using gradient descent to minimize a squared loss L; over in-context examples to learn predictor from hypothesis class Hi, which is defined as linear model on fixed non-linear features. Specifically,
N
C,
1
h ( W1 [Wl]22 S B},
(7)
2
where 6; - K is non-linear feature map K K-dimenssioall space, and W, { are learnable weights initialized to zero and subsequently with norm bounded by B. Therefore, the overall predictor obtained from this composition be written stacked sequence of multiple non-linearities and linear transformations, i.e.
WLou(. (W2g
E H. = H4 .... H2 Hi.
(8)
For example, if d1 is specified as the identity mapping, whereas de for i * are conventional activation functions. 8 could represent k-layer deep neural network. And beyond this, for generic feature maps, could represent more powerful functions. As the step number increases, the predictor h also becomes more complex. Below, connect the algorithm with Transformers.
5

---

Published as a conference paper at ICLR 2025
Lemma 1 (Transformers Learn Step-by-Step). For any k >> 1, given set of CoT demonstrations with k reasoning steps as per (2) and query (N(+1), Transformers with linear depth kt and constant embedding size 2d can express ACOT, where As is steps of GD on squared loss and Hi the hypothesis class defined in whose feature map aligns with the attention as is specified in the
Appendix(A.I
Proof Sketch. Consider simplified case k = 1. The loss is L EigIN IIW (Wo((()) - 13/2
and GD with fixed step size updates the weights as W + W - W. This process also induces dynamies in function space, i.e. the evolution of the learned predictor as the weights update. For query x (1)1), the function-space dynamics could be written as (see derivation in the proof)
(GD Dynamics)
h(a(N+1)) h(x(N+4)
7 (Y P)
(XX) o(z(N+1)),
(9)
Predictions
Residuals
Kernel Function
where Y y= - E [R()))x, y = [1(2 ((X) = [=(=()) E RR The residuals
Y --- Ý are equivalent to Y at initialization the weights are initialized to zero. The last term represents a kernel function w.r.t. the feature map o, quantifying the similarity between the test (i.e.
query) and training examples (i.e. demonstrations).
For comparison, we also rewrite the self-attention layer, where e(0 is reinterpreted as concatenation of input and the residual (z)(*,, ( - ))((())), which at initialization is equivalent to (x((), yy()))
(Transformer Layer) c(N+1) <- ((++1) + Wy a(E W Wowl
(10)
Skip Connection
Embeddinn
Attention Module
where the last term is the attention module. As is specified in Appendix A. with simple choices
of Wv, WK, Wo E 0, one can show (1O) subsumes (9): that is. Transformers can perform kernel regression in their forward pass. The key here is the connection between the kernel function and the attention matrix, which has been widely studied in previous literature, e.g. (T'ai et al. 2019 Wright
here
& Gonzalez |02T] Chen et al.] (2024; (Choromanski al., 2021 Katharopoulos et al. [2020; Wang
et al. 2020 Peng et al.)) 2021) and also discussed in the setting of ICI. (Von Oswald et al. 2023
Cheng et al. 2024 Guo al. (024). Appendix B we provide a comprehensive discussion their connections and how the feature map is related to Transformer parameter and architectural choices.
Extension to 1. To extend this result to COT, we define k loss functions (eC) associated with k reasoning steps. Each loss function is convex w.r.t. the weights of the corresponding predictor. In the forward pass, similar with (9) Transformers implement (kernel) GD dynamics in function space to minimize these loss functions. One challenge of retaining the connection between (9) and (10) while further incorporating CoT is that, for compositional non-linear predictors h in (). updating weights in prior-step predictors (e.g. Wi in h2) could introduce non-linear dynamics in the final prediction h(x) from (8). We show that this issue can be circumvented if the learning is done in a step-by-step manner as in ACOT, namely Transformers first learn a preceding reasoning step using t layers, then proceed to learn the next step using & layers. This results in a total of kt layers for k steps. Note that the construction here not unique and similar conclusions could be drawn from other setups, such as recurrently making k predictions (Li et al. 2023b) in k forward passes, which would additionally require the Transformer to perform the so-called filtering process but could potentially reduce the depth requirement to a constant. Our construction differs from Von Oswald et al. (2033) Cheng et al. (2024) by extending their proof to the case where CoT is used.
3.3 EFFECTS OF TASK DECOMPOSITION
We now proceed to answer when and how CoT improves the learnability of task by studying the in-context learnability of a distribution family P w.r.t. different task decomposition schemes. Before presenting the main result, we analyze how well can the predictor h generalizes to unseen queries. This is accomplished by studying the final error A (P,,A) ) made by the learning algorithm AçTT and its relation with individual errors made by individual algorithms at each step.
Consider a fixed number of steps k. According to Definition 2 applying a sequence of decomposition operators [Ti) on P produces k new distribution families {P:) (E where P; is the induced distribution family that generates the i-th reasoning step. As per equation (5), A ((P..A)) refers the error of learning P; using the individual algorithm A.. The following lemma upper bounds the final error (PP, A) by the individual errors.

---

Published as a conference paper at ICLR 2025
Lemma 2 (Error Upper Bound). Fix 1. Thenfor any distribution faminy P and any decomposition operators (Ti)ie(k-1) applied on P, the predictor returned by ACoT has an error upper bound
A(P, S cs E A(P,,AA) CK 2 max, Ii= 2B Lip(o))²
(11)
where Cj is a constant depending on the hypothesis class Hi and the step number k.
Note that this result does not rely on the specific individual algorithm we choose in Acer as long as the predictor for each step has a Lipschitz constant and the loss function is convex. The upper bound suggests that, to minimize A (P, AA). suffices for each individual algorithm to minimize the error made at its corresponding step. In fact, it actually suffices to minimize the largest error made at the hardest reasoning step argmax; (A(P;,A:); as is revealed by the following result.
Theorem 3 (Improved Learnability with CoT). With CoT, any distribution family P is efficiently in-context learnable by the Transformer in Lemma [ if there exists a finite sequence of deconposition operators {Ti [E())) such that each induced P. if efficiently PAC learnable by the individual algorithm A,, Particularly, the sample complexity is Na(e, max NA, (6/k, e/ck) where NA; is the sample complexity for learning P, by A..
Theoretical Practical Implications. This result indicates that, regardless of how complex the original reasoning task P is, in order for Transformer to efficiently learn this task, it is sufficient to make each subtask P; efficiently learnable. In other words, CoT can reduce the difficulty of learning task to the difficulty of learning the hardest subtask in the chain; here, the hardest step refers the least sample-efficien one. This result also leads to a simple practical lesson for designing CoT: an effective way to form CoT is by decomposing the hardest reasoning step into smaller steps that are easier to learn. Such result aligns with existing empirical practices of decomposing challenging
tasks, e.g. (Zhou et al.) 2022; Khot et al. 2022 Zhang al., 2022). and will be validated in greater depth by new experiments in Section 5.1
To understand this result, notice that in ACTT each subtask shares the same sample size, and thus one has to choose this size based on the hardest step to ensure each individual step can be successfully learned so that the overall task can be successful learned. Another way view this through the lens of error (Lemma 2): in the limit of large sample size N, the individual error at the step with the worst rate will dominate the overall error, regardless the constant coefficient associated with each A (Pr,s, h:); therefore, the hardest step becomes the bottleneck.
Compounding Error Issue. Nevertheless, we note one caveat of Lemma 2 is that it is not asymp- totic in & which has been treated as constant. Therefore, the result in this section does not hold if k scales up, in which case errors could accumulate over reasoning steps and grow exponentially in the horizon. This implies trade-off between the step number and hardness of subtasks: decomposing the hardest subtask improves the learnability, but also introduces the risk of compounding error, which renders scaling up a practical challenge despite the theoretical merit. Aligned with our theory, the compounding error issue has indeed been widely observed in practice, and approaches such as self-correction or refinement have been proposed mitigate this issue, e.g. (Wang al.. [023] Yao et al. 2024 Madaan et al. 2024). From our experiments in Section 5.1 we also find when the hardness of each reasoning step is approximately the same, increasing & increases the overall task hardness and could hurt performance; however, for fixed overall task hardness, reducing the hardness of the hardest step can significantly improve the performance, which aligns with Theorem 3 We also note that the required GD steps for each step depends on parameters and 6 and depth of the constructed Transformer in Lemma [ould increase for harder tasks.
4 HARDNESS OF LEARNING WITHOUT CHAIN-OF- THOUGHT
In this section, will present impossibility results demonstrating that there exist inherently hard tasks that are not learnable without CoT but learnable after being decomposed.
4.1 LOWER BOUND
We begin by presenting a general lower bound on the error A(P,A). This bound applies to both the case where there is no COT (k- 1), and the case of CoT with one intermediate step (k 2) each demonstration is in the form of (z, 2, y). The proof is based on the limited power of the hypothesis class H to approximate the target functions in P (see Appendix A.4)
7

---

Published as a conference paper at ICLR 2025
Theorem 4 (Error Lower Bound). For any distribution famity P and decomposition operator T, suppose ACoT returns a first-step predictor h1 from a finite set Hi S H1. Then the overall ervor has lower bound given by
A(P.A) BVHKHH ((r)))
(12)
where B and K are constants depending on the definition of the hypothesis class as per 7 is certain variable depending on P. we defer its definition Appendix
Var(P)
We discuss the implications of this theorem in two separate cases: when CoT is used and not used.
w/ CoT. In this case, the bound is generally loose or even uninformative, because the term [H|]] is typically expected to be large. This makes sense since when CoT is used, the hypothesis class H H2 o H1 per equation (8) can potentially approximate a wide range of functions analogous to what a two-layer neural network can approximate. Despite this, the bound is still useful to identify cases where the CoT is suboptimal. For instance, consider a dummy CoT where the first step is the identity mapping (i.e., x = 2), and assume the learner has successfully learned this step, which indicates that h1 is also an identity mapping, and hence [7[] = 1. In this case, the lower bound would increase and potentially become positive depending Var(P) which will be discussed below, suggesting this dummy CoT is suboptimal.
w/o COT. A more interesting scenario is when there is COT. In this case the lower bound from 2022), quantifying how hard it is to approximate P using the hypothesis class H. = { + W Wo(z) [[ S B} per definition in 77 The key quantity in the lower bound is Var(P), which indicates the intrinsic complexity of the
12 reduces to 1/2 - B K Var(P) (Malach & Shalev-Shware
task; the more complex P is, the smaller Var(P) is. In the next subsection, we will discuss learning specific family of Boolean functions called parities; these functions underpin a concrete scenario
whereby the task is unlearnable without CoT but becomes learnable when CoT is used.
4.2 ILLUSTRATIVE EXAMPLE: LEARNING PARITIES
Boolean functions are mappings from an input space X' = (=1)" of n binary bits to an output space y = {=1]. In particular, parities are family of functions that compute the exclusive-or (XOR) of bits at some predefined positions in the input, which are notoriously hard to learn (Keamn) 1998; Shhlev-Shwart is determined by subset S C [n]. For each S, the corresponding parity function is defined as
al. 2017 Daniely & Malach] 2020). The specific form of parity function
Xs xs(xa) (ar) llles x[i]] where x|i] is the i-th bit of the input. The distribution family defined as
PXOR(() A ((ss D) SS [n]}
(13)
where D is a fixed input distribution uniform over {=1}". Particularly, applying the lower bound to learning parities, we have the following result:
Corollary 5 (Learning Parities). For any hypothesis class defined in 7 there exists some sufficiently large n E N such that PxoR(n) is not learnable by the corresponding algorithm ACOT when 1, but learnable after using the following decomposition to form the CoT (k 2):
for i S x/11]
1st Step: 2|1] = xx.s(((())
2nd Step: y X2,s(2)
II 2/1]. (14)
for i &
where the first step (1.s(x) learns to select relevant features from while masking irrelevant ones, and the second step (2)s(2) computes XOR of all bits in z.
The proof is deferred to Appendix A.5
Intuitively, parities are hard to learn without CoT because the 2" target functions in the family form an orthogonal basis for the space of all Boolean functions. Consequently, a linear function class with a fixed feature space dimension K can not approxi- mate all parities, as this is as hard as approximating all Boolean functions. Concretely, we have Var(PXoR(m)) = 2-" and thus the lower bound will become positive for sufficiently large n; this implies there always exists least one target function in Pxon(n) and some Eo 0 such that eo holds true, and thus the task is not learnable per definition in Section 2 regardless the sample size. CoT resolves this issue, since the first-step predictor with learnable weights can adapt to the relevant
Note that in the case of CoT, by default | and the bound still applies.
8

---

Published as a conference paper at ICLR 2025
Hoar
Hog
Hous
Hage
Hage
100
100
H-1
97
H-1
65
H=2
92
35
90
H=2
43
27
(%) soru SECCES
34
(N) 809 s9930ns
45
24
H-5
-00
40
H-6
H-6
30
H=7
H=7
H=8
16
H=8
13
0
GPT-4o
(b) GPT-3.5-turbo
Figure 2: Success rate of GPT-4o and GPT-3.5-turis for learning compositional functions. H denotes the number of elementary functions used construct the target function; Hmaa denotes the maximal number of elementary functions to construct a reasoning step.
bits in the input x (determined by S), such that the second-step only needs to learn fixed XOR function. In particular, with CoT, both steps or subtasks become learnable by a linear function class with fixed feature space dimension. Therefore, we have illustrated such a setting where successful learning is impossible without COT. The effectiveness of the designed CoT is verified in Section 5.2
5 EXPERIMENTS
As empirical verification, we consider new arithmetic reasoning tasks and evaluate the performance of real-world LLMss, including GPT-4o and GPT-3.5-turbo. Sec 5.1 studies the connection between the hardest step and the overall reasoning performance. Sec[5.2 [222ttsss tests whether the specific forms of CoT we introduced can improve the performance. See detailed experimental setups in Appendix Codes are available httttt:s///
5.1 INCREASINGLY COMPLEX FUNCTIONS
Since benchmarks are lacking where one precisely control the hardness of tasks and CoT steps, we first present a method to construct such tasks by incrementally building upon challenging subtasks.
Constructing Highly Challenging Tasks. Let consider class of elementary functions F. where each function maps from the input space X' to itself. In general, these elementary functions should be considered equally easy learn. Then, we sample a sequence of these functions [1, ./,y...,fy Fa; composing them gives us a target function f = fr o ʃ2 o f1 : X' - X' whose complexity increases as H increases. We consider an instantiation by defining the input space as the space of two integers x e 2², The elementary functions are defined as choosing one integer and using it to perform basic arithmetic operation (drawn from +, or x) with another number. Therefore, F. consists of
2 2|0 + = + =[1],
2 - 2 =[1], z + 200 x =[1] 2 2|1] + 2 + =(0), z[1] 2[1]=1110 z(0],
(15)
2[1] + = x 2(0].
While each elementary function (15) is simple, the overall target function f can become highly complex, possibly representing polynomial functions on =|0] and x[1] up to arbitrary order and number of terms. Moreover, to quantify the hardest step, we do not reveal all intermediate steps of f in the demonstrations provided to LLMs. Instead, we stipulate that there exists least one step E [k] where the function from 2;-1 to 2; is constructed from Hmar elementary functions, whereas all other steps use fewer of them. For example, given H elementary functions f1 210) 2 + =[1], fa 2 < 211 x z(0) and f3 : x(1] + 2[1] 2(0]. the hardest step can be expressed as fa o fa a fi
{
2 - 0 + = 1
(16)
1 2; 2 (1) (2-1(11-1) -1)
Results. We test the performance of real-world LLMs on the reasoning task described above with respect to different overall hardness H and the hardest step Hmar- We report their success rates

---

Published as a conference paper at ICLR 2025
(n, k)-Parities
(10,1)
(10,2)
(10,3)
(10,4)
(10,5)
(10,6)
(10,7)
(10,8)
(10,9)
(10,10)
GPT-4o w/o CoT
87
51
48
50
52
$1
GPT-4o w CoT
73
66
58
62
GPT-3.5-turbo w/o CoT
75
62
60
47
59
58
51
56
55
54
GPT-3.5-turbo CoT
80
76
72
74
75
69
57
63
57
57
Table I: Success rate (%) of GPT-4o and GPT-3.5-turbo of learning (n.k)-parities.
3-Term DNF
Width 3
Width 4
Width
Width 6
Width 7
Width
Width 9
Width 10
GPT-4o w/o CoT
85
81
77
73
68
66
62
74
GPT-4o w CoT
96
87
86
88
81
80
80
84
GPT-3.5-turbo CuT
74
62
58
GPT-3.5-turbo CoT
90
81
73
Table 2: Success rate (%) of GPT-4o and GPT-3,5-turbo of learning 3-term DNF.
across 100 i.i.d. sampled target functions for each H and Haax- And for each target function, the LLMs are provided with demonstrations and asked to infer the computation process and apply derive the output for an unseen input. As shown in Fig. (and more results in Appendix C.I). the success rate of LLMs quickly drops as Hmar increases. In particular, GPT-40 can successfully learn the target function in most cases when H 1; however, performs significantly worse as Hmoz increases from 1 to 4, then fails as Haaz becomes even larger. These phenomena corroborate our result that reducing the complexity the hardest step critical successfully handle the task.
5.2 CANONICAL BOOLEAN FUNCTIONS
We further evaluate LLMs two families of Boolean functions: parities and disjunctive normal form (DNF), which are known hard to learn (Daniely & Vardi 2021 Malach & Shalev-Shwer Shwart)S 2022).
Task Descriptions. An (n,, k) -ppriityyy function computes the XOR (+) of subset of k variables from a total of n input binary bits (n 10 in our experiments). It outputs if an odd number of the relevant variables are 1, and 0 otherwise. Meanwhile, DNF function is disjunction (logical OR) of conjunctions (logical ANDs) of literals; and in the experiments, we consider a family of 3-term DNFs the width and m {+1) 3w is latent variable whose value determines the target function (i.e. m.j 1 invalidates 2.
[(x) = v 1=1 A== (mig V may) where
Results. For each & in parities and w in DNFs, we similarly i.i.d. sample 100 target functions For each function, we provide LLMs with 100 in-context examples, ask them to find patterns these examples, and return the output for each query. Tables and report the success rates. Particularly terms of parities, we find even GPT-4o generally performs no better than random guessing (with an expected accuracy of 50%) when k >> 3. Then, we provide LLMs CoT examples with intermediate steps that are provably effective by applying our results derived Section 4 for parities, the intermediate step is defined as =[i] = x if E S otherwise 0; for DNFs, 2[i,
x if
m[i,j] otherwise 1. Results Tables aad 2 demonstrate the designed CoT significantly improves the performance, e.g. GPT-4o achieves an almost perfect success rate of 94 on (10, 4)-parity, while without CoT the success rate is 54, which is close to random guessing.
6 CONCLUSION AND DISCUSSION
In this paper, we quantify the benefits of task decompositions within the setting of learning tasks in- context by Transformers. We note that a limitation of our work thus far is that, despite quantifying the potential for Transformer to efficiently learn a complex task by CoT, it nonetheless remains unclear on case-by- case-by-case by-case basis whether real-world L.MSs will actually achieve success in practice. This is because of confounding issues related to model training procedures, including dataset properties, the actual optimization process, and fine-tuning. Hence an interesting future direction is to delve more deeply into these issues.
For parities, sample from uniform distrihution; for DNF, we sample from non-uniform distribution to ensure the label (0/1) is balanced for 3.
10

---

Published as a conference paper at ICLR 2025
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.0x::23006000227,, 2023.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in in-context learning? investigations with linear models. International Conference on Learning Representations, 2023.
Maciej Besta, Nils Blach, Ales Kubicek, Robert Gerstenberger, Michal Podstawski, Lukas Gianinazzi, Joanna Gajda, Tomasz Lehmann, Hubert Niewiadomski, Piotr Nyczyk, et al. Graph of thoughts: Solving elaborate problems with large language models. In Proceedings of the AAAI Conference on Artificial Intelligence, 2024.
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.
Yingyi Chen, Qinghua Tao, Francesco Tonin, and Johan Suykens. Primal-attention. Self-attention through asymmetric kernel svd in primal representation. Advances in Neural Information Process- ing Systems, 36, 2024.
Xiang Cheng, Yuxin Chen, and Suvrit Sra. Transformers implement functional gradient descent to learn non-linear functions in context. International Conference Machine Learning. 2024.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. International Conference on Learning Representations, 2021.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reichhro Nakano, et al. Training verifiers to solve 2021.
math word problems. arXiv preprint arXiv:2110.1444
Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Shaming Ma, Zhifang Sui, and Furu Wei. Why can gpt learn in-context? language models secretly perform gradient descent as meta-optimizers. In Findings of the Association for Computational Linguistics: ACL 2023, pp. 4005-4019, 2023.
Amit Daniely and Eran Malach. Learning parities with neural networks. Advances in Neural Information Processing Systems, 33:20356-20365, 2020.
Amit Daniely and Gal Vardi. From local pseudorandom generators to hardness of learning. In Conference on Learning Theory, pp. 1358 1394. PMLR, 2021.
Guhao Feng. Bohang Zhang, Yuntian Gu, Haotian Ye, Di He, and Liwei Wang. Towards revealing the mystery behind chain of thought: theoretical perspective. Advances in Neural Information Processing Systems, 36, 2024.
Deqing Fu, Tian-Qi Chen, Robin Jia, and Vatsal Sharan. Transformers learn higher-order optimization methods for in-context learning: A study with linear models. arXiv preprint arXiv:2310.17086.
2023a.
Yao Fu, Hao Peng, Ashish Sabharwal, Peter Clark, and Tushar Khot. Complexity-based prompting for multi-step reasoning. In The Eleventh International Conference on Learning Representations, 2023b.
Shivam Garg, Dimitris Tsipras, Percy 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, Liu Yang, Tianhao Wang, Dimitris Papailiopoulos, and Jason D Lee. How well can transformers emulate in-context newton's method? arXiv preprint arXiv:2403.033 088.,333,,, 2024.
Tianyu Guo, Wei Hu, Song Mei, Huan Wang, Caiming Xiong, Silvio Savarese. and Yu Bai. How do transformers learn in-context beyond simple functions? a case study on learning with representa- tions. International Conference Learning Representations, 2024.
11

---

Published as a conference paper at ICLR 2025
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rans: Fast autoregressive transformers with linear attention. In International conference on machine learning. pp. 5156-5165. PMLR, 2020.
Michael Kearns. Efficient noise-tolerant learning from statistical queries. Journal of the ACM (JACM), 45(6):983-1006, 1998.
Tushar Khot, Harsh Trivedi, Matthew Finlayson, Yao Fu, Kyle Richardson, Peter Clark, and Ashish Sabharwal. Decomposed prompting: A modular approach for solving complex tasks. arXiv preprint arXiv:22
2022.
Juno Kim and Taiji Suzuki. Transformers learn nonlinear features in context: Nonconvex mean-field dynamics on the attention landscape. International conference on machine learning, 2024.
Yingcong Li, Muhammed Emrullah lidiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in in-context learning. In International Conference on Machine Learning, pp. 19565-19594. PMLR, 2023a.
Yingcong Li, Kartik Sreenivasan, Angeliki Giannou, Dimitris Papailiopoulos, and Samet Oymak. Dissecting chain-of thothtt A study on compositional in-context learning of mlps. arXiv preprint arXiv:2305. 18869, 2023b.
Zhiyuan Li, Hong Liu, Denny Zhou, and Tengyu Ma. Chain of thought empowers transformers to solve inherently serial problems. International Conference on Learning Representations, 2024.
Aman Madaan and Amir Yazdanbakhsh, Text and patterns: For effective chain of thought, it takes two tango. arXiv preprint arXiv:2209.07686. 2022.
Aman Madaan, Niket Tandon, Prakhar Gupta, Skyler Hallinan, Luyu Gao, Sarah Wiegreffe, Uri Alon, Nouha Dziri, Shrimai Prabhumoye, Yiming Yang, et al. Self-refine: Iterative refinement with self-feedback. Advances in Neural Information Processing Systems, 36, 2024.
Eran Malach and Shai Shalev-Shwartz. When hardness of approximation meets hardness of learning. Journal of Machine Learning Research, 23(91):1-24, 2022.
BK Natarajan. Learning from exercises. In Proceedings of the 1989 Workshop on Computational Learning Theory, pp. 72-86, 1989.
Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, al. Show your work: Scratchpads for intermediate computation with language models. arXiv preprint 1 14, 2021.
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. arXiv preprint arXiv:2103.02143: 2021.
Akshara Prabhakar, Thomas L Griffiths, and R Thomas McCoy. Deciphering the factors influencing the efficacy of chain-of-thought: Probability, memorization, and noisy reasoning. arXiv preprint 2oXiv:2407.01687 2024.
Ofrr Press, Muru Zhang, Sewon Min, Ludwig Schmidt, Noah A Smith, and Mike Lewis. Measuring and narrowing the compositionality gap in language models. arXiv preprint 20X0v:2210.03350, 2022.
Jack Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring. Susannah Young, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint 12. 2021.
Laria Reynolds and Kyle McDonell. Prompt programming for large language models: Beyond the
few-shot paradigm. In Extended Abstracts of the 2021 CHI Conference Human Factors in Computing Systems, pp. 1-7, 2021.
Shai Shalev-Shwartz and Shai Ben-David. Understanding machine learning: From theory to algorithms. Cambridge university press, 2014.
12

---

Published as a conference paper at ICLR 2025
Shai Chalev-Shwartz Ohad Shamir, and Shaked Shammah. Failures of gradient-basee deep learning. In Interaational Conference on Machine Learning, pp. 3067-3075. PMLR, 2017.
Aarohi Srivastava, Abhinav Rastogi, Abhishek Rao, Abu Awal Md Shoeb, Abubakar Abid, Adam Fisch, Adam R Brown, Adam Santoro, Aditya Gupta, Adrià Garriga-Alonse et al. Beyond the imitation game: Quantifying and extrapolating the capabilities of language models. arXiv preprint arXiv:2206.046135 2022.
Prasad Tadepalli. Learning to solve problems from exercises. Computational Intelligence, 24(4): 257-291, 2008.
Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis-Philippe Philippe Morency, and Ruslan Salakhut- dinov. Transformer dissection: An unified understanding for transformer's s attention via the lens of kernel. In Proceedings of the Conference Empirical Methods in Natural Language Processing, 2019.
Miles Turpin, Julian Michael, Ethan Perez, and Samuel Bowman. Language models don't t always say what they think: unfaithful explanations in chain-of-thoug prompting. Advances in Neural Information Processing Systems, 36, 2024.
Leslie G Valiant. A theory of the learnable. Communications of the ACM, 27(11):1134-11 1984.
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, pp. 35151-35174. PMLR, 2023.
Johannes Oswald, Eyvind Niklasson, Maximilian Schlegel, Seijin Kobayashi, Nicolas Zucchet, Nino Scherrer, Nolan Miller, Mark Sandler, Max Vladymyrov, Razvan Pascanu, et al. Uncovering mesa-optimizztioon algorithms in transformers. arXiv preprint arXiv:2309.05858 2023.
Boshi Wang, Sewon Min, Xiang Deng, Jiaming Shen, You Wu, Luke Zettlemoeer and Huan Sun. Towards understanding chain-of ofciin--fffnn -ii--ffftthooo thought prompting: An empirical study of what matters. arXiv preprint arXiv:2212. 10001, 2022.
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006. 200.60000044 2020.
Xuezhi Wang. Jason Wei, Dale Schuurmans, Quoe V Le, Ed H Chi, Sharan Narang, Aakanksha Chowdhery, and Denny Zhou. Self-consistency improves chain of thought reasoning in language models. In The Eleventh International Conference on Learning Representations, 2023.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in Neural Information Processing Systems, 35:24824-24837, 2022.
Noam Wies, Yoav Levine, and Amnon Shashua. Sub-task decomposition enables learning in sequence to sequence tasks. arXiv preprint arXiv:2204.02892, 2022.
Matthew A Wright and Joseph E Gonzalez. Transformers deep binary kernel machines. arXiv preprint arXiv:2106.01506, 2021.
non-mercer
Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Tom Griffiths, Yuan Cao, and Karthik Narasimhan. Tree of thoughts: Deliberate problem solving with large language models. Advances in Neural Information Processing Systems, 36, 2024.
Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. Advances in Neural Information Processing Systems, 2023.
Zhuosheng Zhang, Aston Zhang, Mu Li, and Alex Smola. Automatic chain of thought prompting in large language models. arXiv preprint arXiv:2210.034933 2022.
Denny Zhou, Nathanael Schärli, Le Hou, Jason Wei, Nathan Scales, Xuezhi Wang, Dale Schuurmans, Claire Cui, Olivier Bousquet, Quoc Le, et al. Least-to-most prompting enables complex reasoning in large language models. arXiv preprint arXiv:2205. 10625, 2022.
13

---

Published as a conference paper at ICLR 2025
CONTENTS
Introduction
1.1
Contributions
Related Work
2
Preliminaries
3
Improved Learnability by Task Decomposition
3.1
Learning Algorithms Enabled by
3.2
Expressiveness of Transformers
3.3
Effects of Task Decomposition
Hardness of Learning without Chain-of- Thoought
Lower Bound
4.2
Illustrative Example: Learning Parities
5
Experiments
5.1
Increasingly Complex Functions
5.2
Canonical Boolean Functions
10
6
Conclusion and Discussion
10
Proofs
15
A.1 Lemma
Expressiveness of Transformers
15
A.2 Lemma
2
Upper Bound
17
A.3
Theorem
3
CoT Improves Learnability
18
A.4
Theorem 4
Lower Bound
19
A.5
Corollary
Learning Parities
Background: Connection between Attention and Kernel
22
Experimental Details and Additional Results
24
Increasingly Complex Functions
24
Canonical Boolean Functions
27
14

---

Published as a conference paper at ICLR 2025
A PROOFS
A. LEMMA T EXPRESSIVENESS OF TRANSFORMERS
Given demonstrations {(z8)) 200 which defines loss function quantifying the error of a particular predictor h4. These loss functions are
( د we could create k training sets, each of
N
- - Σ
(((221)
[ :iE
: [+]}
(17)
Note that while the overall predictor
h = h4 ..... o h2 h1 Wiob(....
(18)
a non-linear function, loss functions in (17) are convex with respective to the weights of their
corresponding linear predictors {W; : [] J. Let also denote hgw ha o... h2 o hi for
k' & k. Using gradient descent to minimize L; with fixed step size v) induces the following training dynamics in weight space
+ W. W; (2, (2;;1,) 0z
(19)
where h,,77-1) (= [EENN Rd(2)) (7,-1) (cs: JIN E RRXN Note that difference between our setup and the conventional supervised learning setup is that, the latter is interested in the variation of output with different weights, while in this paper, we are also interested
in the dynamics of intermediate steps. Particularly:
For i' < i, the dynamics of intermediate steps induced by GD is
hige hg (2)+))(
(20)
namely the variation of upper layer weights does not affect lower layer representations (i.e. interme- diate steps).
For i' = i, the dynamics is
hgt()(N+1))+ (oss
+ (W, VW oi (
(21)
-Ng)(x((+1)) 1/ (Z: hz th (z( (6
Let Ki be the kernel function defined by the feature map Φie we have
(N+()( + (a (N+1) + ( - ()) k2 (( (x(
(22)
For i' > i, the dynamics is
h; o ... ... h141 o hev (ev(a((+)) (2(
+ h o h041 (N+1() + (Z: -x (x(N)))))
(23)
which in general intractable since h2 o... hr+1 is non-linear. However, if upper layer weights in h4 h141 are 0, hgg (N+1)) will become 0 as well and thus circumvent (23)
Recall also that based on the definition in Section 2 the self-attention layer can be written as
(N+1) - (x+1) Wy E W W WoelN+1)..
(Self-Attention)
(24)
where E = (z, ž1, ... -1 2x) E Rdd at the input layer, E d(z.).
In the following construction, show that in the forward pass of Transformer, 24 could express dynamics of all intermediate steps, including (20). (2 and (23) based on a certain order in which minimization of losses in (1) is performed. Particularly,
construction, Transformer will sequentially minimize loss functions in (17). In other words, the lower layers of the Transformer learn prior reasoning steps, while upper layers of the Transformer learn later reasoning steps.
15

---

Published as a conference paper at ICLR 2025
Expanded Embedding Construction. Let d Er= d( Z3) be the total dimension of the original Zn and 2 Zu. Let d' 2d d( d( Zx) which will be the dimension of our expanded embedding. We now construct a projection matrix P E Rd xd to map each demonstration into an expanded embeddings that keeps track of inputs and residuals in single vector. Concretely, we partition perform the following transformation:
demonstration (20, (10.21 .......,2... 2k) where 20
La(Zo)
Ia(z.)
0
T Ta(zz)
0
-l(()))
P =
(25)
0
0
Ia(z))
0
:
Ta(zz)/
where each (2) denotes d[ (24) ) d(22) identity block (or negative identity block). Applying this P to vector (20.21. 2..... 2x. E Rd yields an expanded embedding
= (zo: 21, 2k,
R
(26)
At initialization, each hu (z;-1) is zero, the rows with a(2,) produce residuals (h:(zi-1) 2z. which imitially reduce to (0 2;) 24 Hence, for a demonstration (x, 21. 22, ,.,,..,, 2(6 with x 20. we explicitly obtain
e x h1(x) 21, 21, .... hy- (2 - 2k-1, 2k-1, hy (zk-1) - zx E Rdd (27)
Similarly, for the query x (N++1) Zo. we have the expanded embedding
(N+1) (N++)) h1 (N+()) 0, lggk 1 (N) O, leg( (2.(11())) E Rd (28)
reflecting that for the query we store prediction hgg (N each layer in own coordinate, with zeros in place of 2; or differences as needed to run gradient- based updates in the Transformer. Here we leverage fact that 21, 1 are all-0 vectors since they are unknown.
Weight Matrices Construction. For layers that minimize the i-th step's loss function Li, we construct Transformer weights in the corresponding self-attention layer as follows.
0 Od
0
0
Wy =
n (
0
(29)
0
0 Od-(i)
where
1-1 Σ d(zz) d(x), d+(1) 5=0
d((t) = 2
d( d(Zz) d(Zx) d(Zi). Σ j=i
selecting the residual coordinates the expanded embedding (i.e. 2; - (z)-1))) This extracts those
residuals, multiplies them by -7) and writes them into the sublayer prediction coordinate.
We let Wx and We select the appropriate coordinates for = (from demonstrations) and 1 (N+1) (from the query):
0 4p
0
Wx
0
Id(z)-1)
0
(30)
0
0 (g)
0 (()
0
We =
I L
0
(31)
0 Og
16

---

Published as a conference paper at ICLR 2025
where
di(i) 2
> d(z,) d(x) d(Z;-1), ==
d(1) 2
_ d(zz) d(Zx), (32)
and similarly (i), (i) define the row/column slices for the query coordinate. Stacking self- attention layers (each with the above Wv, WK. Wo for subtask i) carries out gradient steps on C. Finally, iterating over I,..., k in ascending order completes minimization of all losses in equation 17 yielding the final predictor that includes all intermediate and final predictions.
Extension kernel regression follows from Proposition in Von Oswald et al. (2023) and Proposi- tion 1 in Cheng et al. (2024). Refer also to Appendix B for a discussion the connection between kernel and attention.
A.2 LEMMA 2 UPPER BOUND
Given a target function f and an input distribution (zx). the algorithm returns a predictor based on N ii.d. samples, whose expected error is defined as f(x))) Let us first consider single step of task decomposition.
Lemma 6. For any distribution family P and decomposition operator T, the predictor returned by ACOT demonstrations sampled from the corresponding distributions in Pi and P2 has error upper bound
AP) s 2 max (1, ((ax(() A(P:.
(33)
where CB = B Lip())² is a constant determined by the hypothesis class H. in 7
Proof. For a family of distributions P, the error is defined as
AP,, A) A max (J,D)P
ENED ( [()(),
(34)
where we slightly abuse notation here as h is also dependent on the distribution (J: D) and the learning algorithm. For a certain decomposition operator T. the target function can be expressed as == f2 fi and the predictor h fı2 o h1. We have
A)
max ERND (F,D)E
((h2 h))(s) - (h2 o fz)(z) + (he o fi)(z) f(x)))
(35)
Suppose the feature map (xz) for h2 has Lipschitz constant
[kx(x) o (xx))l2
Lip(ø) sup at2' || z'll2
(36)
and by Jensen's 's inequality, we have
A(P,A) S max EAND (F.D)P
[(h2 (1120)1)(2)) (h2 f1)(z)) + ((h2 f1)(x) - [(x))
(37)
< max EAND (J.))E
B² Lip(())) - + (h2(z) ff(=))²
(38)
< B² Lip( max ENND [1/b(xx) - (z) + max ENND (12()))fe(2)) (N.D)EP
(T.D)EP
(39)
where B² Lip(0)2 a constant determined by the definition of hypothesis class. Given decomposition operator T,, the distribution family can be decomposed into Pi and P2. It follows that
APP, A) < 2 Lip() + 22 A2)
(40)
Let CB max(1, B² ))) get the desired upper bound.
17

---

Published as a conference paper at ICLR 2025
To extend this lemma to the case where there are more reasoning steps, i.e. 2, let
Individual Error:
(P,.. = A max BELINDD (h4(2+-1) - 2;]] (S.D)EP;
Accumulated Error: Á(P, (P5 Á. max ENND 11 ((hg o.... hz)(z) -x]] (S.D)EP
(41)
A; is the error of learning an individual reasoning step i, which is consistent with its definition in Section 2 Å is the accumulated error from the first step to step i. Note that we slightly abuse the notation here since A(P, A) is actually equivalent Åg. We also have A A1. Therefore
Á = max E
2 [1(,; ....... h1) h -2112
(42)
(r,D)EP
= max EUND (1,))P
2 [1(0g ol1..o....h))()))) 1=(21-1) 1) + h=(21-1) - =:12
(43)
< 2B2 Lip(Ør)2
- 2
(44)
< ...
s 2 Σ j=1
[ 2B Lip(o,)² 1++
A
(45)
We use the lemma from
to Therefore, we have the following upper bound
A(P, = Å S c Σ A(P..
(46)
where Ca = 2 max, M 2B Lip(d;) 2 is constant depending on the hypothesis class Hi and k. Note that in the derivation, we could remove () and directly analyze the expected error EzeD (eowD[((((x)) [[((()), which will give us a similar result: the final error is upper bounded by the sum of individual errors up constant.
A.3 THEOREM COT IMPROVES LEARNABILITY
Recall the definition of in-context learnability: we say parametric model M can learn task P if there exists learning algorithm A and function N : (0, 1)² + N, such that for any confidence and accuracy parameters d, e (0, 1):
1. Under certain parameter choice f, we have Mo(., (x)) for any query x.
2. Given (6, e) i.i.d. examples, the learning algorithm A returns predictor h such that with probability of at least - 5, A(P,A) <
Moreover, we say P can be efficiently in-context learned by . both the running time of A and the sample size N (ō, e) polynomial in 5-1 and -1
To see how it can be reduced to the learnability of subtasks after using CoT, recall the definition of PAC learnability:
Definition 3. We say subtask P; can be efficiently learned by an algorithm A; if its sample complexity and time complexity scale as poly (,,)) where ò; and e are confidence and accuracy parameter for the subtask.
Since by Lemma] 1 know there exists such a parameter choice such that TF)(; (-or(-) our goal is then to show: if every subtask P, is efficiently learnable by its corresponding individual algorithm A,, the overall task is also efficiently learnable by ACoT-
By Lemma 2 we know that for fixed k, the accuracy parameter e of the overall task is upper bounded by the sum of the accuracy parameters of subtasks { up constant coefficient. Additionally,
18

---

Published as a conference paper at ICLR 2025
notice that, given n (not necessarily independent) events A2, An. each occurring with probability P ((A,) = 8,, we have
n n
12
p
A
max
0,
64
(47)
This means the confidence parameter of the overall task ō is also upper bounded by the sum of confidence parameters all subtasks {6.) Namely
S Ck
611
S
(48)
In terms of the time complexity of AcoT. since each individual algorithm A, runs in time polynomial
in 0,- and e, we have (Am) = poly(6,1 (.,)(6,,.._)) Moreover, one can choose ón = 62 = 5g- 8', 61 = 62 a simple combination of all individual algorithms (Ai:)= we have
CR = d', correspondingly have 8' 2 5/k and e' 2 c/ck- Since Acor is
A
TiCo(ACoT)
Ti)) < poly(o (0-0001111 ,)
(49)
meaning the learning algorithm is also computationally efficient.
In terms of the sample complexity of ACsT, we also let d1 = =
8/k and
El = E2
Eg e' 2 c/C2. Note that the case of CoT, all individual algorithms use the same number of samples as the algorithm Acor. To show the algorithm is sample-efficient, we can simply choose N c) be the largest among { AN c i.e.
( é) max N.. (8',
(50)
KE(
which ensures that ACTT can achieve high accuracy with high probability. Since any N MA,(0, (6i, e;) is polynomial in 8' and e'-1 holds that N6 e) is also polynomial in 8-1 and - However, Na(ō,e) c) < | NA, (8', E'), we have no guarantee that each step can successfully learn their corresponding reasoning step. From here, we also proved that the sample efficiency of the overall step is that of the hardest step in the CoT.
Since ACoT can efficiently learn P, and Transformer with linear depth and constant embedding size can express Acor, we complete the proof.
A.4 THEOREM LOWER BOUND
The in-context learning error has approximation error lower bound, that the minimum error achievable by predictor in the hypothesis class H. = {he h1 h1 H2, ha E H2)
max EZND (h(x). f())) Z max min ENED [1 [())). (/D)EP
(51)
(r.))P (h, Ma) X
Thus suffices to lower bound the approximation error.
To do so, notice that the learned first-step predictor f2 is from finite function class Hi, which is subset of the initial hypothesis class H4. We show the approximation power of h(x) with finite-sized H' is lower bounded by a linear class whose size depends on Mi, In particular, suppose the hypothesis class
H1 = (h1,1, 11,2.... A An,pag...
(52)
and based on the index of h2 in H', the predictor h(x) can be re- re-written
h(x,j) A W2 hz))(z)
(53)
U
(54)
19

---

Published as a conference paper at ICLR 2025
where Pi.j = = his5 and U W2,1 if otherwise 0. As (4 is an inner product of weight vector RKIR and feature vector ((x) [+1]K(H] in an expanded space, f(xz,) j) reduces to a linear function.
Since the squared loss [(, is convex w.r.t. U for arbitrary and j. its expectation (h) - EgwD P (h(x, j): [xx)) is also convex w.r.t. U. Moreover, h(x,j when W2 or U
goes to 0. Therefore, given any (f, D) E P and any predictor h with fixed W2 and j (and thus fixed
U), by first-order condition, we have
[h) 2 (0) + (U - O, (b)
(55)
2 |1|| (b)
1
(56)
2
where the last equation uses the fact L (0) Ez~D |1 [())] 1/2 and inequality (v, u) -1lul2/|/o|l2-
2
Notice U has the same norm W2 and thus || < B. Moreover, have
D(f A (NL))(())) Vy | (h ((e) lu=ol
(57)
DRnD Vol (h(x, (), f(x)) lu=ol
(58)
Vu 1 2 (() ((=)) (()))))))))
EzND
(59)
- |EanD [
(60)
K |M
E (x)
(61)
Subjecting to (56 gives us that, for any (f, E P and any predictor h(x) obtained from in-context learning, i.c.
1
min (2), x
(,)
B- ((S
(62)
2
Following from 51. the in-context learning error has lower bound
A(P, A) Z max (JD)EP (HH, x
min
(h)
(63)
2 max (AD)EP
B ())
(64)
2
> B(J.D)E
B. ())
(65)
1
((J.D))
B B- -EE,,D)E
(66)
2
By noting | < 3 we have
A) 2
B B.DDer (f.
(67)
K Pg
B. B.E
ENnD (2))
(68)
Let
Var(P) 2 supp EEND [o())
(69)
which determined by target functions and distributions in P, and could be understood as the intrinsic complexity of the distribution family (i.e. the more complex P is, the smaller Var(P) is).
20

---

Published as a conference paper at ICLR 2025
follows that.
K P21
AP,, A)
B
BUID)EP BOWD ( (x))
(70)
2
BV K|HL Var(P),
(71)
completing the proof.
A.5
COROLLARY 5 LEARNING PARITIES
w/o COT.
The distribution family of parities with input size n is defined as
PxoR(a) - {(xs D): S S [m]} where (s(x)
II 2|1]
(72)
iES
Based on 669 we can derive an upper bound on Var(P) for parities, i.e.
Var(PXOR(m) = sup E D)
(73)
sup
Be-Dlo(r)xs(T))
(74)
[XXOR(]]
SC(M)
where (PxoR(n) = 2" and 2 is the power set of [n]. Note that {xs: S S [ni]] forms Fourier
basis of Boolean functions, meaning that for any pair of different subsets S1, S2 S [n]. we have
Ez~D(Xs; (s)xs, (z)] = 0.
(75)
Therefore Eg~D(o(<)) Xs) (z)] exactly the Fourier coefficient of Boolean function ((x) correspond- ing to S, which we denote as @(S). Therefore
S( 1
Var(PXOR(m))
(76)
2n
8c(M]
Since, when there is no CoT, the error lower bound is
K
A A(Vor((x))(()) )) BV K Var(Pxo(()) By
(77)
2n
we can choose n > + log, B²K K such that this lower bound is always positive (regardless of the sample size).
w/ CoT.
Consider the following CoT
1st Step:
z|i] = Xi,s(
{{
x||]] for i E S 1 for
2nd Step: y y=x2,5(;) = II =|1], (78)
we show each step can be approximated by the hypothesis class defined
H4 {--1+ Widr(zz-1) ||all2 S B}
(79)
where di 2,-1 R is non-linear feature map to a K-dimensional space, and W; dRd().)
are learnable weights. We do this by construction. For the first-step predictor (1.s(z), we let o (z)
(x, 1) be the connection of x and for the second-step predictor (2)s we Φ2(x) = II. =
With these features, both steps can be learned by simple linear model W ar which also means the expected error could be arbitrarily small with sufficiently large sample size. The sample complexity follows the standard analysis for linear regression.
In contrast, when there no CoT, APP.A)
A) has a positive lower bound, meaning there always exists least one (f, D) E P such that the expected emror is greater than the lower bound.
21

---

Published as a conference paper at ICLR 2025
B BACKGROUND: CONNECTION BETWEEN ATTENTION AND KERNEL
As the background knowledge for understanding the construction of Transformers in Lemma 1 here provide non-exhaustive summary of the connections between the attention matrix
Transformers and the kernel method from previous works {Von Oswald et al. 2023 Cheng et al.
2024 Guo al. 2024 Tsai al. 2019; Wright & Gonzalez 2021] Chen et 2024). We rewrite
the kernel gradient descent dynamics and the Transformer layer here for reference
(GD Dynamics)
M(a(N+1) ((r(N+1)) +
n ( Y) ((X) o(r(N+1)),
(80)
Predictions
Residuals
Kernel Function
(Transformer Layer)
(N 11)
e(N
WvE W Woe
(81)
Skip Connection
Embedding
Atteention Module
Here, the kernel is induced by s(x, y) <(x) o(y). The definition of could be flexibly chosen depending on the practical implementation of Transformers.
In-Context Learning. Von Oswald et (223) (Proposition 1) demonstrates in the most simple case where both the non-linearity in Transformer and the feature map are identity mappings, the yields
(a(x)) 0
and Wx WQ = [u(x) Dacy))
weight constructions Wv
E W W Woe
X ZN+1: Wv WV=E=nn00uar Y Y).
(82)
In this case, n(x; v/ x: y is the inner product kernel. Their Proposition 2 further discusses the case which accommodates the role of the MLP module in the Transformer architecture. Specifically the MLP module transforms the token MLP (e), and thus with the same Wy, WK. Wq and identty mapping a, the kernel is
<(x; v) = (L) MLPo(y).
(83)
Guo al. (2024) also leverages the representation power of MLP and considers the case where Ø: R"
R is a fixed representation function, which can in principle be chosen arbitrarily as long as the features can represented by a MLP. Theorem in this paper provides construction where is chosen to be the normalized ReLU and proves that Transformers can perform in-context ridge regression. Following this work, more recently Kim & Suzuki (2024) similarly concluded that MLP layer can extend the class of learnable functions of ICL to the Barron space.
Cheng et al. (2024) considers the case where a is non-linear, and their Proposition proves the non-linearity matches the kernel the Transformers can perform functional (kernel) gradient descent regression w.r.t. the reproducing kernel Hilbert space metric of this kernel; the proof is exactly based on the relation between (80) and (8I) For example, they demonstrate that the Softmax attention corresponds to the exponential kernel
x(x, v) exp (-
(84)
One caveat here is that the attention matrix in this case is asymmetrie; to address this, one could treat the normalization term as pre-conditioner of the optimization algorithm, and the construction of Transformers should be modified accordingly. We refer the readers to more discussions in their paper.
General Discussions. Even before ICI., the connection between kernel and attention a topic that has been widely discussed (e.g. (Tsai et al.) 2019 Wright & Gonzalez 2021 Chen et al. (2024) etc.) and references therein). For instance, Wright & Gonzalez (2021) proves that the standard attention matrix is reproducing kernel for reproducing kernel Banach space and gives explicit formulation of the feature maps (Proposition 1). Their Theorem 2 further demonstrates that Transformers can learn any binary non-Mercer reproducing kernel Banach space pair.
In practice, variants of Transformers with kernelized attentions (e.g. (Choromanski et al., 2021
Katharopoulos et al. 2020 Wang al. 2020 Peng et al. 202T) etc.) such those relying random
22

---

Published as a conference paper at ICLR 2025
Fourier features very popular and effective. These implementations often considered for efficiency considerations, since once the attention is kernelized, one can switch the order of matrix multiplication to accelerate the computation and reduce the time complexity from quadratic to linear w.r.t. the input length.
23

---

Published as a conference paper at ICLR 2025
C EXPERIMENTAL DETAILS AND ADDITIONAL RESULTS
C.I INCREASINGLY COMPLEX FUNCTIONS
Sampling. For this task, we sample a sequence of functions from elementary functions with the following probabilities:
1
1
1
P (z < =|0] 2[1])
( + - =[1])
[ 2 2[1]) 4
g'
8'
1 ( = =00))
1 P(+ 2 x 2 4
1
(x][1] - = =|0))
8*
8'
Here, the multiplication operation is more likely to be sampled than addition subtraction. For the input values, we uniformly select two unique integers from the set (2,3..... 10). This range is chosen avoid excessively large numbers, which difficult to handle, and trivial cases, such as when z00] = x rr]1, which could result zero values that are easily predictable by LLMss For this task, explicitly ensure that each sample is unique prevent repeated samples in training and testing. This is done to avoid scenarios where LLMs could simply memorize the results from the demonstrations and use them to answer the query.
CoT Prompting. Given certain H namely the maximal number of elementary functions to construct a reasoning step, we implement it by randomly masking H consecutive intermediate steps. For example, when T and H 3, example prompt is given follows:
Given two numbers, sequentially apply predefined arithmetic operations (addition, subtraction, multiplication) to transform them. Each step involves specific predefined operation on one of the numbers. If any operations or intermediate results are missing, deduce these to complete the transformation and arrive at the final output.
Input: 7, 5
Stepl: 7, -2
Step2: 5, -2
Step3: missing
Step4: missing
Step5: 5, 75
Output: -70, 75
Input: 2,
Stepl: 2, 1
Step2: 3,
Step3: missing
Step4: missing
Step5: 3, 36
Output: -33, 36
Input: 5,
What is the output? Your answer should end in the format 'Stepl: ?, Step2: ?, Output:?'.
Note that the missing steps are consistent for one trial. We test 100 times to compute the success rate.
addition to reporting the success rate of LLMs for predicting the final output as presented in Section 5| we also evaluate their success rate for predicting intermediate steps. This provides a more comprehensive assessment of the LLMS' performance, as even they might fail to predict the final step but could still succeed in predicting the intermediate steps.
Additional Results.
24

---

Published as a conference paper at ICLR 2025
H Hooe
H
Mma
Hose
1
11
H=1
H=1
H=2
92
35
80
H=2
43
27
H
(1%) 55803Sn
H=4
88
22
16
60
4=4
24
7
82
21
12
3
35
13
4
40
H-6
H-6
2
H=
69
23
D
20
H=7
15
12
4
H-8
H=8
(a) GPT-4o
(b) GPT-3.5-turbo
Figure Success rate for predicting the last step (i.e. namely the output).
Hine B
H
Minsa
100
100
H=1
H=2
98
80
H=2
14
H=3
91
15
success Rate (%)
H=3
25
16
60
H=55
89
17
10
0
H=5
36
10
5
H-7
73
20
D
H-7
20
9
2
H=8
78
14
H=8
16
(a) GPT-40
(b) GPT-3.5-urro
Figure 4: Success rate for predicting the second to last step.
Hoz
Higg 33
Ho
Moas 6
21
22 33
Moar
24
100
100
H=1
H=1
80
80
Hm3
94
56
H-3
33
24
2993 10332210
60
60
H=4
92
46
10
35
26
H=55
H=5
-40
440
H-6
87
29
13
2
H-6
37
12
0
0
20
37
20
H=88
74
21
D
#=8
21
10
(a) GPT-40
(b) GPT-3.5-turbo
Figure 5: Success rate for predicting the third to last step.
25

---

Published as a conference paper at ICLR 2025
H Hoo
H
Hoae
1
1
H=1
H=1
H=2
80
H
H-3
(06) 55803Sn
H=44
100
75
60
H=4
43
23
3
54
32
20
1
40
H-6
41
H-6
18
1
H=
90
49
31
D
20
H=7
34
17
10
2
0
H-8
H=8
(a) GPT-4o
(b) GPT-3.5-turbo
Figure 6: Success rate for predicting the fourth to last step.
Hine B
Honz
M
100
100
H=1
H=2
80
H=2
H=3
Success Rate (%)
H=3
60
Hm4
H=5
97
76
66
48
2
H=5
40
32
31
19
12
H-7
93
57
42
20
0
D
H-7
30
23
9
H=8
86
32
0
H=8
17
4
(a) GPT-40
GPT-3.5-turbo
Figure 7: Success rate for predicting the fifth to last step.
Haz
Hong Hig 3
Ho m4
6
100
100
1=11
H=1
80
80
Hm3
H-3
2993 10332200
60
60
H=4
H=44
H=55
H=5
-40
400
H-6
98
77
76
58
62
H-6
41
29
23
33
24
0
(%6
20
21
20
H=88
92
56
50
36
13
#=8
24
11
(a) GPT-40
(b) GPT-3.5-turbo
Figure 8: Success rate for predicting the sixth to last step.
26

---

Published as a conference paper at ICLR 2025
C.2 CANONICAL BOOLEAN FUNCTIONS
Examples of standard and CoT prompts for (10, -ararity and DNF with width 6 are given as follows:
Standard prompt for parity:
Predict the output based on pattern in the input binary string.
Input: 1010101100 Output: 0
Input: 0100000011 Output: 0
Input: 0110010000
Output: 1
Input: 1010100001
Output: 0
Input: 1010100001 Output:
Input: 1011000111
What is the output? Directly answer the question in the format 'Output:'.
CoT prompt for parity:
Replace some bits located at specific predefined positions in the binary string with 0 to form new string. Then, based some 4 patterns in the new string to predict the output.
Input: 1000010001
New string: 0000010001
Output: 0
Input: 0100110111
New string: 0000110001 Output: 1
Input: 0101001000
New string: 0001000000 Output: 1
Input: 0100000010
New string: 0000000000
Output:
Input: 1011010000
New string: 0001010000
Output: 0
Input: 010010011
What is the output? Di Directly the question in the format 'New string:, Output:'.
27

---

Published as a conference paper at ICLR 2025
Standard prompt for DNF:
Predict the output based pattern in the input binary string.
Input: 00101 000001 010011
Output: 1
Input:
100111 011001 010111
Output:
Input: 010010 011010 101011 Output: 1
Input: 010101 001101 001001 Output: 1
Input: 111110 011001 010111
Output: 1
Input: 010001 110101 011011 what is the output? Directly answer the question in the format 'Output:'.
CoT prompt for DNF:
Replace some bits located at specific predefined positions in the binary string with 1 to form string. Then, based some 4 patterns in the new string to predict the output.
Input: 000101 011111 011010
New string: 100111 111111 011111 Output: 1
Input: 001101 100111 000101
New string: 101111 110111 000101
Output: 0
Input: 100001 011001 001010
New string: 100011 111011 001111 Output: 0
Input: 001100 010100 101011
New string: 101111 110110 101111
Output:
Input: 010101 111011 100101 New string: 110111 111011 100101 Output: 0
Input:
000100 011110 100000
What is the output? Directly answer the question in the format 'New string:, Output:'.
28


Reviews:
[{"review_id": "qbLmcyZRL8", "reviewer": "Reviewer_S1ny", "summary": "This paper studies the effectiveness of chain of thought (CoT) in the in-context learning setting. Particularly, it is shown that there is a parametrization of the Transformer architecture that allows in-context learning of certain functions with CoT. It is further claimed that there are functions that are only learnable with CoT and not otherwise in this setting. The paper concludes with experiments with GPT models and their in-context learning abilities based on the granularity of the CoT provided to them.", "questions": "- The depth requirement in Theorem 1, $kt$, seems rather impractical and quite large. Further, I think there are several clarifications required. For start, I think $t$ itself is not a parameter of learning setting but part of the learning algorithm. We are in the PAC setting in this paper. So it would indeed make sense to give the depth of Transformer based on $k, \\delta, \\epsilon$ instead. In particular, for me it’s not clear if the complexity in $k$ is really linear. Based on eq. (11) we know that as $k$ increases, we have a compounding and increasing error. Thus, probably more GD iterations are also needed. I.e., for a given $\\epsilon, \\delta$, parameter $t$ is not proven to be independent of $k$. So effectively, I’m not sure if the depth is linear in $k$ as claimed. I think giving the depth in $k, \\delta, \\epsilon$ would fix all these issues. \n- The negative result for the in-context learning is given for the hypothesis class in eq. (7) and not Transformers. Can we say something about Transformers not being able to learn the functions in question in an in-context fashion?\n- In definition 1, part 2, how is $\\mathbb{E}_x[l(h(x),f(x))] \\leq \\epsilon$ with prob. $1-\\delta$ equivalent with $\\Delta(\\mathcal{P}, h) \\leq \\epsilon$?\n- For Theorem 3, it is said that the number of samples we need is equal to the number of samples needed for the hardest subtask. However, the samples used by different subtasks are not independent. Can you further clarify how is this taken into account?\n- Given a fixed architecture (with fixed depth), can you please further elaborate on the effect of making CoT more granular and also the effect of increasing the number of in-context samples?\n- In the in-context setting, when we go the in-context setting we provide more info by giving each sample. In other words, comparing the sample complexities with and without CoT doesn’t seem to be totally fair. Maybe, can you give a comparison based on the total number of tokens given? \n\n\n\nSuggestions/Minor remarks:\n- For me, the notation $\\Delta(\\mathcal{P}, h)$ used in equation 5 is not clear. $h$ on the right hand side depends on the function being learned however the $h$ is on the left hand side should be independent of it and represent the learning algorithm. Maybe replace it with $\\mathcal{A}$?\n- I would suggest that the authors include the existence of a hypothesis class in page 2 to increase the clarity.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is generally well-written and theoretically studies an important problem: ‘effectiveness of CoT in the in-context learning setting’. Moreover, the paper generally proposes rather clear definitions of in-context learning with and without CoT based on the PAC framework. \n- A particularly interesting theoretical prediction is that making CoT more granular and longer does not necessarily improve the performance.", "weaknesses": "- One problem with in-context learning works as done in this paper is that it is only proven that there is a Transformer parametrization suitable for in-context learning (similar to an expressivity result). However, it is not clear whether this parameterization is realistic or not. I.e., if a Transformer pre-training is done in a reasonable manner, it’s not clear if it's possible/likely that the Transformer achieves the needed parameterization. (E.g., the claimed parametrization may not be reachable by GD in practice and there is generally no insight on the pre-training process.)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QuDAcrRsPg", "reviewer": "Reviewer_dJXS", "summary": "This paper analyzes Chain of Thought (CoT) learning in LLMs. The key idea behind CoT is to learn hard tasks by being shown intermediate steps involved in learning. The paper formulates this as in-context learning (ICL) which means that a parametric language model can be trained to generalize a set of demonstrations (with k steps) to predict the final output for a new input. The paper shows that CoT can be implemented by a step-by-step algorithm, where each step is generalized by a learning algorithm A_i from the examples of the i'th stage. The key lemma is to show that a transformer of depth linear in k can express the learning algorithm A_{CoT} = A_1...A_k as a sequence of gradient descent steps. \nThe main theorem then shows that CoT can efficiently learn a set of tasks which can be decomposed into subtasks with a sample complexity bounded by the sample complexity of the hardest subtask. Negatively, the paper gives a lower bound on the overall error as a function of hypothesis space and the problem class. \nThe concrete results show that the class of parity functions is not polynomially learnable in 1 stage, but can be learned in the CoT framework in 2 stages. \nThe paper introduces a new benchmark for CoT with controllable hardness. Each stage of each function manipulates two integers through addition and multiplication. The hardness is controlled by skipping some of the intermediate steps of the function. The success rates on GPT-40 and GPT-3.5 illustrate the predicted result - they degrade gradually with increasing the maximum number of contiguous missing steps.  They also show positive results on an (n,k)-parity task and on 3-term DNF learning using the same LLMs.", "questions": "Since this work builds upon a number of previous works in CoT and transformers, please state precisely what was previously known and what are the novel (theoretical) contributions.  \n\nIt is not always obvious how to decompose a given task. Do you have any advice on how to approach this on a practical problem?\n\nThe importance of the work is diminished by having to manually decompose each new task to generate demonstrations. Do you have any thoughts on how to reduce this burden?", "limitations": "", "rating": 10, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The paper gives a strong theoretical justification and analysis of CoT in the context of LLMs, building on prior work on understanding the prior work on viewing LLMs as few shot learners. The analysis is sound and well-presented. \n\nThe paper shows concretely how CoT makes otherwise unlearnable classes learnable through decomposition.   \n\nThe empirical results support the theory and clarify the paper. \n\nThe paper presents some interesting benchmarks including a challenging task that can express arbitrary polynomials over integers, which can be controlled for hardness.", "weaknesses": "It appears that this work is somewhat related to `learning from exercises' in the problem solving setting (eg, https://dl.acm.org/doi/10.5555/93335.93344, and https://onlinelibrary.wiley.com/doi/10.1111/j.1467-8640.2008.00330.x) and curriculum learning (http://arxiv.org/abs/2101.10382). It would be good to see a discussion of the relationship.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SjPS4TnjV8", "reviewer": "Reviewer_XoWP", "summary": "This paper studies how chain-of-thought (CoT) impacts in-context learning (ICL) for transformer models. There are two main results. First, the paper shows that a linear-depth transformer model can implement a learning algorithm A_CoT that chains together k gradient descent learning algorithms A_i that run for t steps each for 1 <= i <= k (Lemma 1). An upper bound on error is given for A_CoT, thereby showing that a composite task is efficiently learnable by a transformer model if each subtask is efficiently learnable by each of the k individual gradient learning algorithms (Lemma 2), where the hardness of learning is bounded by the difficulty of the most difficult subtask (Theorem 3). Second, the authors lower-bound the error of A_CoT in the case of finite hypothesis class for A_1 as a function of the size of that hypothesis class (Theorem 4), which suggests that there may be functions that are learnable with CoT that are not learnable without CoT. A concrete witness is identified for learning parities (Corollary 5).", "questions": "Please see weaknesses for questions.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper studies an important problem in trying to theoretically elucidate CoT and ICL for transformer models. There is also preliminary experimental results that corroborate the theoretical results.", "weaknesses": "The presentation is quite dense could be improved to explain the motivation and intuitions behind definitions/results.\n- For instance, the definition of in-context learning (Definition 1) makes a distinction between predictors and a learner (around line 195-196). From my understanding, the definition of Model(Data, query) = LearningAlg(Data)(query) = predictor_D(query) is just partial evaluation, and so the distinction that is drawn between predictors and learners is somewhat artificial. Rather, the main idea seems to me to formulate ICL, and so we just define a function that takes a query that is a tuple (demonstrations, x) as opposed to x. Moreover, it would be helpful to explain why PAC learning is chosen and to give pointers (presumably to Theorem 4) on how exactly PAC learning will be used to help you demonstrate your results and why it is sensible.\n- Why does Definition 2 use multisets?\n- Some of the material in the appendix on how your Transformers implement gradient descent result differs from Cheng et al. and Von Oswald et al.'s would be helpful in the main text to give more context.\n- Could you offer more clarification on how to extend the results to handle k > 1 in lines 294-305 which is from what I understand is the main point of innovation of the paper. However, I can't seem to find more details in the main text or the appendix.\n- The references for this paper do not use ICLR's reference style ...\n- The statement of Lemma 1 might be helped by stating that t steps of gradient descent map to t layers of self-attention which is mentioned in the appendix.\n- I find the statement of Theorem 4 to be somewhat strange. If I understood correctly, without CoT, we can essentially run for t gradient steps on a neural network that predicts from a representation that has t self-attention layers. With CoT for k steps, we can essentially run for kt gradient descent steps with kt self-attention layers and k prediction layers. How does Theorem 4 control for the same number of layers across reasoning steps k? Intuitively, I expect a theorem showing that there functions that are learnable with CoT that are not learnable without CoT to (1) show that a kt layer network run for kt gradient descent steps can not learn the task purely with an identity (trivial) task decomposition and (2) that a task decomposition (T in theorem 4) used by CoT can be constructed in polynomial time *uniformly* in the task with a Hypothesis class H_1'. However, there are no limits on the task decomposition T in the theorem and it's currently not clear to me how the size of the Transformer is controlled in the current statement of Theorem 4. Perhaps I misunderstood and would be curious to get more insight into what is actually going on.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6KXeFPblxr", "reviewer": "Reviewer_PsRB", "summary": "This paper investigates the impact of subtask complexity on the reasoning capabilities of Transformers when employing Chain-of-Thought (CoT). Theoretical results demonstrate the expressive power of Transformers to learn from in-context CoT samples through in-context gradient descent. Furthermore, the paper establishes upper and lower bounds for the error bound $\\Delta(\\mathcal{P}, h)$, highlighting the significance of task decomposition. Experimental results illustrate the effectiveness of CoT in parity and DNF tasks.", "questions": "1. Can you describe how to design effective CoT formats for more general problems using the strategies proposed in this paper?\n\n2. What are the input and output formats (including CoT steps) for all experiments? What is the length of CoT reasoning steps for each task?\n\n3. Lines 355-357 indicate that increasing reasoning steps $k$ negatively impacts performance. Can you elaborate on how this conclusion is drawn from the experiments in Section 5?\n\n4. Is it feasible to develop a Transformer with constant depth as proposed in Lemma 1? For instance, could the filtering process mentioned in lines 303-305 be applied? A polynomial embedding size is acceptable in the context of constant layers, which aligns more closely with practical applications.\n\n5. In lines 381-397, why is $|\\mathcal{H}_1'|$ expected to be large when applying CoT, while $|\\mathcal{H}_1'|=1$ without CoT? Given the problem setup, $\\mathcal{A}$_CoT is an algorithm that learns from demonstrations, meaning the output predictor $h$ is influenced by the input demonstrations and is not necessarily unique. In this context, the assertion that $|\\mathcal{H}_1'| = 1$ without CoT seems questionable.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "This work emphasizes the importance of understanding principles for designing CoT steps, providing several insights for creating more effective CoT formats when presenting in-context samples. The theoretical findings regarding error decomposition at each CoT intermediate step reveal the hardest subtask as a key bottleneck in reasoning.", "weaknesses": "1. While the paper proposes effective strategies for decomposing target problems as claimed in lines 23-25, these strategies appear limited to specific problems like parity or DNF. It remains uncertain whether these approaches can be generalized to CoT format design for a broader range of reasoning problems.\n\n2. The theoretical results are somewhat unconvincing:\n   - Lemma 1 indicates that a Transformer with depth linear to the number of reasoning steps $k$ and the number of gradient descent (GD) steps $t$ diverges from practical applications. Current large language models (LLMs) typically possess at most hundreds of layers, yet they can handle thousands of context lengths. Consequently, a construction involving constant-depth Transformers with polynomial hidden dimensions may yield more practical insights.\n   - In Lemma 2, the bound becomes less meaningful when $k$ is not treated as a constant. However, reasoning lengths often are not constant; the expressive power is strictly constrained. [1] illustrates that a constant-depth, log-precision Transformer with polynomial hidden dimensions and either no CoT reasoning or $O(\\log n)$ CoT reasoning steps cannot solve problems outside the $\\mathsf{TC}^0$ class, where $n$ is the problem size. To strengthen the effectiveness of CoT, the reasoning steps $k$ should generally be $\\Omega(\\log n)$, and in many instances, $k = \\text{poly}(n)$. Under these conditions, the bounds become quite loose.\n\n[1] Chain of Thought Empowers Transformers to Solve Inherently Serial Problems, https://arxiv.org/abs/2402.12875", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)