Chain-of-Thought Provably Enables Learning the (Otherwise) Unlearnable

About Theorem 4 (Error Lower Bound) and Question 5

1. This is a typo, we meant "when CoT is NOT used, we do not actually provide intermediate steps in the demonstration or ask the learner to learn the identity mapping" in the previous response.

2. When applying CoT ($k=2$), each demonstration is given in the format $x, z, y$ and the predictor is $h_{k=2} = h_2 \circ h_1$ where both $h_1 \in \mathcal H_1$ and $h_2 \in \mathcal H_2$ are linear function classes on non-linear feature maps (per equation (7)); when CoT is not applied ($k=1$), demonstration is $x,y$ and the predictor is $h_{k=1} \in \mathcal H'$, which is a single linear function class on non-linear feature map (per equation (7)). The key point of theorem 4 is simply that $\mathcal H'$ is a more restricted function class than $\mathcal H_2 \circ \mathcal H_1$, and consequently there exists complex target functions that $\mathcal H_2 \circ \mathcal H_1$ can learn well while $\mathcal H'$ can not. Neither writing $\mathcal H' = \mathcal H' \circ \{I\}$ or $\mathcal H_2 \circ \mathcal H_1 = \mathcal H_2 \circ \mathcal H_1 \circ \{I\}$ with change this result since they are essentially the same function class. However, theorem 4 only applies to two steps (note that theorem 4 states "for any decomposition operator T" rather than a sequence of decompositions $\{T_1, T_2, \ldots\}$) and does not deal with the composition of three function classes; in other words the remaining functions besides of the first step should be a linear function on non-linear feature map – both $\mathcal H' \circ \{I\}$ and $\mathcal H_2 \circ \mathcal H_1$ satisfy this while $\mathcal H_2 \circ \mathcal H_1 \circ \{I\}$ do not.

About the Relationship Between Performance and Reasoning Steps, and Question 3

It is true for the same column in Figure 2 more reasoning steps implies more complex target functions, and our theoretical prediction from lemma 2 is that if each step is equally easy to learn the performance will hurt with more reasoning steps; in other words, it does not draw negative conclusions about the case where there are more reasoning steps but each step is easier to learn, since in this case each individual error $\Delta(P_i, A_i )$ could be smaller, causing tradeoffs between error accumulation and smaller individual error. In such a case where there are more reasoning steps but individual steps become easier (and the target function remains the same), our theoretical prediction from theorem 3 is that if the hardest step becomes easier, CoT would help despite error accumulation. This is empirically verified by checking the same row of Figure 2, i.e. the target function is the same, increasing $H_{max}$ (decreasing $k$) would significantly hurt the performance.

Q1 & Weakness 1

As mentioned in Section 3.3, our general guidelines for forming an effective CoT are to decompose the hardest reasoning step into smaller steps that are easier to learn; parities and DNF serve as examples where we can provably show the task is unlearnable without CoT but learnable with CoT. While for more general problems it is typically infeasible to formally quantify the hardness, our guidelines still work by assessing the hardness of tasks in a qualitative or intuitive manner. For example, in the arithmetic reasoning task in Section 5.1, the hardness is qualitatively measured by the number of elementary functions required to construct the target function, and we found the performance is more significantly affected by the hardness of the hardest reasoning step in the CoT rather than the final task hardness.

Q2

The input and output formats of all experiments were provided in Appendix C.1 (for the arithmetic reasoning task in Section 5.1) and C.2 (for parities and DNF in Section 5.2). The CoT steps for the arithmetic reasoning task are $H-H_{max}+1$ since we hide $H_{max} -1$ steps in the original $H$ reasoning steps to control the hardness of the hardest step. The CoT steps for parities and DNF are $2$.

Q3

In Figure 2, as we increase $H$ (which is exactly the CoT steps when $H_{max} = 1$), the performance gradually drops. This is because, though each step is easy to learn, the error made at each step will accumulate and affect the final performance. This is consistent with our theory that increasing reasoning steps $k$ negatively impacts performance, despite the advantage of better expressiveness as pointed out in previous works such as [1].

[1] Chain of Thought Empowers Transformers to Solve Inherently Serial Problems, https://arxiv.org/abs/2402.12875

Q4 & Weakness 2 (Lemma 1)

The filtering process mentioned in lines 303-305 can be applied to reduce the depth requirement from $kt$ to $t$; however, unless we treat $t$ as a constant, reducing the depth to a constant seems difficult. More generally, it is an open research question what learning algorithms Transformers can implement with only constant depth since most ICL works mentioned in the related work Section 1.2 also require at least $t$ depth. Nevertheless, all subsequent results after Section 3.2 / Lemma 1 are adaptable for other Transformer constructions. Namely, if future works identify more efficient ways for expressing the learning algorithm (say, with constant depth and polynomial embedding width), our results apply as well.

Q5

Note that in our setup, we do not consider learning only one target function $f$ (since being able to learn one function is not practically meaningful) but a family of target functions or a distribution family $\mathcal P = {(f_1, \mathcal D_1), (f_2, \mathcal D_2), \cdots }$. With CoT, the first-step learner is expected to output a different first-layer predictor $h_1$ for each unique first-layer target function $f_1$ in $\mathcal P_1$ (and of course for the same target function, the learner could also output a different predictor on random samples), and thus the possible number of first-step predictors $|\mathcal H_1’|$ is expected to be large in general. Without CoT, the learner also outputs a different predictor $h$ for each unique target function $f$ in $\mathcal P$; however, for each predictor (which outputs the prediction in just one step), it could be equivalently viewed as $h = h \circ \operatorname{Id}$ where the first step is the identity mapping and the second step is $h$ itself, and thus $|\mathcal H_1’| = 1$. Note that this is only for theoretical convenience for incorporating both cases in a single analytic framework; when CoT is not used, we do not actually provide intermediate steps in the demonstration or ask the learner to learn the identity mapping.

Weakness 2 (Lemma 2)

We treat $k$ as a constant since the error will accumulate and the upper bound does not hold asymptotically in $k$. This reveals a practical limitation of CoT when increasing the reasoning steps, which has not been revealed by previous theoretical works on CoT that purely focus on the expressiveness of Transformers. In experiments, we also find that when each step is equally easy to learn, increasing the reasoning steps often hurts the performance, which is consistent with our theoretical prediction.

Q1

It is standard in statistical learning to distinguish learner and predictor (e.g., see [1], Chapter 2.1): the learner is defined as a mapping from a set of samples to a function, i.e., $(\mathcal{X} \times \mathcal{Y})^N \rightarrow \mathcal{Y}^{\mathcal X}$, and the predictor (a.k.a. hypothesis) is defined as a function from $\mathcal X$ to $\mathcal Y$. In $\mathcal M_\theta(D,x) = \mathcal A(D)(x) = h_D(x)$, the first equality holds since mathematically any function $f(x_1, x_2)$ can be uniquely expressed in curried form as $g(x_1)(x_2)$, where $g$ is a function that maps an input $x_1$ to another function of $x_2$, i.e., $f(x_1, \cdot)$. The second equality is simply defining a notation (i.e., writing $g(x_1, x_2)$ as $h_{x_1}(x_2)$), which is also standard in learning theory.

PAC learning is chosen since it provides a formal way to quantify the learnability of tasks in terms of the required number of samples (i.e., sample complexity) to successfully learn the task (i.e., achieve a desired level of accuracy and confidence). We show that with CoT, the sample complexity can be reduced to that of the hardest task, and without CoT, there exist tasks that are not PAC learnable regardless of how many samples are given and how many GD steps are used. These results formally quantify how CoT helps in learning initially complex tasks.

[1] Understanding Machine Learning: From Theory to Algorithms

Q2

Definition 2 uses multisets since there might exist repeated elements. For instance, when the first target function is $f = f_2 \circ f_1$ and the second target function is $g = g_2 \circ f_1$, the subtask target function $f_1$ is repeated.

Q3

Thanks for the suggestion. We comprehensively discussed the connection in Appendix B and will find ways to incorporate some of them into the main text in the next version.

Q4

At a high level, to extend to $k > 1$, instead of using self-attention layers to simulate the training dynamics of the output $y$, as is done when $k = 1$, we construct self-attention layers to simulate the training dynamics of all intermediate results ${z_i}{i=1}^k$, including the output, as described in equations (20)-(23). To do so, we use a projection matrix to expand the dimension of embeddings from $d$ to $2d$, which allows it to store both the initial intermediate result $z_i$ and the prediction error (a.k.a. residual) $z_i - h_i(z{i-1})$. The former is used for computing the kernel for layer $i$ (i.e., the blue part in equation (9)), and the latter is used for the training dynamics of layer $i-1$ (i.e., the red part in equation (9)). $W_K, W_Q, W_V$ in equations (28) and (29) are constructed in such a way that they attend only to the part of embeddings that correspond to the current layer. The complete proof is given in Appendix A.1.

Q5

The formatting issue has been fixed.

Q6

We have updated the “Extension to $k > 1$” paragraph below Lemma 1 to make the connection between Transformer layers and GD steps clearer and more explicit.

Q7

Theorem 4 is an analysis of Algorithm 1, which holds for all numbers of GD steps $t$, all decompositions $T$, and compares $k=1$ (no CoT) and $k=2$ (one step of CoT). Particularly, the result in Theorem 4 is derived based on the representational limit of the linear hypothesis class used by the ICL algorithm. It states that regardless of how many GD steps are used and how many in-context samples are provided, the learning algorithm cannot successfully learn certain tasks (such as the parity example) without CoT, since the expected error always has a non-trivial lower bound. This forms a separation from using CoT, in which case we have shown that the task becomes learnable. With sufficient GD steps (which can be implemented by Transformers with a sufficient number of layers) and a sufficient number of in-context examples, the prediction error could be made arbitrarily small; this is otherwise impossible without CoT.

Some other clarifications: The operator $T$ is predetermined and used to form the CoT examples (the learning algorithm itself does not construct CoT samples). Given a set of CoT examples generated based on the decomposition $T$, $\mathcal H’_1$ is the set of predictors the learning algorithm can return. In this case, without CoT, we could equivalently view the learning algorithm as always returning the identity mapping for the first step, i.e., $|\mathcal H’_1| = 1$.

We will improve the presentation to make the results clearer, and please let us know if you have more questions.

Thank you for your response.

Q1

Thank you for the explanation. As I understand it now, you use the standard PAC learning definition with partial eval (curry/uncurry). My confusion stems from the phrasing "To define a notion of learnability in the context of language modeling, we extend the classic PAC learnability [42] to the setting in which the learner is associated with a parametric model" which made it seem like the standard definition was actually modified.

Q2

Thank you for the response, this addresses my question.

Q3

Thank you.

Q4

Thank you for this explanation. Perhaps it would help the reader if the "Extension to k > 1" header was placed before the introduction of Lemma 1, or Lemma 1 was modified to explicitly talk about the extension to k > 1.

Q5 and Q6

Thank you.

Q7

Thank you for the response. I still don't quite understand how the size of the transformer is controlled in the baseline (no CoT) vs. CoT case in theorem 4. Put another way, assuming that the baseline and CoT transformers have the same number of layers and are run for the same number of gradient steps (although they could be used differently), can you explain at a high-level why CoT overcomes these barriers? Is the "main insight" that the task decomposition (given to us by an oracle) means that we can target each reasoning step at a different loss function aligned with the appropriate task? I ask again because I believe this work addresses an important problem. However, I would like to understand in simple terms where all the heaving lifting is done, which as I understand it now, is given to you via the task decomposition oracle and the ability to choose a different loss function for each intermediate task.

Q1

We have modified the sentence to avoid ambiguity.

Q4

We have modified lemma 1 to explicitly state k>1.

Q7

In a high-level and informally, the results in section 4 state that: without CoT, the learning algorithm 1 (regardless of how many GD steps are used) is unable to output predictors that can accurately solve the task $P$ if the task itself is inherently complex; in contrast, with CoT, instead of directly learning the task, the algorithm can learn two subtasks $P_1$ and $P_2$ that are easier than the original task $P$ and compose predictors obtained from each subtask to solve the initial task $P$. More specifically, directly learning the task is hard because the hypothesis class $\mathcal H$ is not sufficient (i.e. not expressive enough) to represent all target functions in $P$; decomposition makes learning easier because it enables learning $P$ step by step, that is we can first use $\mathcal H_1$ to learn $P_1$ and $\mathcal H_2$ to learn $P_2$ and then compose the learned predictors.

It is important to note that theorem 4 itself does not directly link with Transformers — this is a result about algorithm 1. The way to link the results in section 4 to Transformers is by combining them with lemma 1, that is: without CoT, there does not exist such a Transformer (even it has infinite depth), under the construction we provided in lemma 1, that can learn the task; but with CoT, there indeed exists Transformers (with finite depth) that can simulate algorithm 1 and successfully learn the task. Overall, the insight here is not that “we can target each reasoning step at a different loss function aligned with the appropriate task”, but CoT makes learning a task easier by decomposing it into subtasks that are learnable by a learning algorithm (i.e. algorithms) Transformers can actually implement (otherwise without CoT more powerful algorithms are needed to learn the task but it’s an open question whether ICL can go beyond linear function classes considered in this paper).

Weakness

Thank you for the suggestion. The analysis of CoT in this paper is indeed conceptually related to the idea of learning from easier to more difficult tasks. We will discuss the connection in the related work section.

Q1

As discussed in Section 1.2, previous theoretical works on CoT either focused on overly restricted tasks or did not quantify the effects of different ways to decompose the task. This paper formalizes in-context learnability and first explores how different ways of task decomposition affect learnability.

Q2

Our results suggest that to improve the overall learnability of a complex task, it is desirable to decompose the hardest step into smaller steps, where the hardness here is quantified by the sample complexity. Unfortunately, in practice, except for certain well-defined problems such as parities and DNF considered in this work, there lacks a unified way to quantify the hardness, and the specific way to decompose the task should be studied on a case-by-case basis (e.g., in our experiments, learning higher-order polynomials is generally considered harder).

Q3

One potential way to eliminate manual efforts for designing CoT is to automate this process by using ‘find the hardest step and decompose it into subproblems’ or something similar as a prompt and instruct the LLMs to construct demonstrations under our principle, following existing works such as zero-shot CoT or auto CoT [1,2].

[1] Large language models are zero-shot reasoners, NeurIPS 2022

[2] Automatic chain of thought prompting in large language models, ICLR 2023

Q1

The number of the required steps $t$ might indeed increase with $k$ given fixed $\epsilon$ and $\delta$ due to compounding errors, and thus the depth dependence on $k$ might not be linear (if the goal is to achieve a fixed small error). While we could explicitly account for the dependency by, e.g., writing $t$ in terms of $\epsilon$ and $\delta$, it introduces additional technical difficulties since their relations depend also on the learning task, which is not specified in both theory and practice. For some tasks, such as parities in Section 4.2, it is infeasible to write $t$ in terms of $\epsilon$ and $\delta$ since the error is lower bounded and we could not achieve small error regardless of how large $t$ is. Therefore, in this paper, we do not restrict the depth of the Transformer and allow $t$ to be chosen arbitrarily (i.e. one that is optimal for the task). While in practice the depth of Transformers is restricted (which is a limitation shared by most in-context learning works), we nonetheless have made real-world predictions on how CoT affects the in-context learning performance, verified by our experiments.

Q2

It is prohibitively hard (and an open research question) to discuss in general what tasks Transformers cannot learn, without fixing a specific (or a class of) learning algorithms, as we do in this work. While we have thought about deriving information-theoretic lower bounds for certain tasks to show Transformers cannot in-context learn these tasks, we have not found separation results showing how CoT could help in this case. We hope to cover this in future work.

Q3

This is a typo. We meant $\forall (f, \mathcal{D}) \in \mathcal{P}, \mathbb{E}_{x \sim \mathcal{D}(x)}[l(h(x), f(x))] \leq \epsilon$ is equivalent to $\Delta(\mathcal{P}, h) \leq \epsilon$, and you are correct they are not equivalent with an additional probability constraint. Since all theoretical results are derived based on $\Delta(\mathcal{P}, h)$, we modified Definition 1 accordingly to avoid ambiguity.

Q4

It does not matter whether or not the samples for different subtasks are independent, since the probability (w.r.t. a random draw of samples) that the learner fails to learn the overall task is upper bounded by the probability that each individual subtask fails, even if a sample from one subtask is dependent on a sample from another task (i.e., by Bonferroni inequality (44)). Therefore, to successfully learn the overall task, it suffices to choose the sample number based on the hardest subtask. See the complete proof in Appendix A.3.

Q5

For a fixed architecture and assuming the depth is sufficiently large to express the learning algorithm where $t$ is set to be optimal (which is considered in our analysis), increasing the number of CoT steps would reduce the sample complexity to that of the hardest subtask. However, suppose the hardest subtask of two different CoTs has the same sample complexity; the one with the longer chain would suffer from the compounding error issue and thus overall perform worse. Increasing the number of in-context examples is always beneficial in our theory. These results have been verified in our experiments. For the case where the depth is insufficient to express the learning algorithm where $t$ is chosen to be optimal (which we have not discussed), there will exist more complicated tradeoffs between $k$ and $t$, as you mentioned, and these are generally prohibitive to quantify (which seems to be an open research question even beyond the scope of ICL/CoT/LLM).

Q6

If we restrict the number of total tokens, there will be a tradeoff between the number of samples and the number of CoT steps, which is generally hard to quantify if we do not consider a specific task and a specific form of CoT. However, for the parity example discussed in Section 4.2, since without CoT the lower bound in Equation (12) holds regardless of the sample size (and the number of tokens), whereas with CoT the problem becomes learnable, CoT provably has benefits for this particular example.

Minors

Thank you for the suggestion. We have modified the notation to $\Delta(\mathcal{P}, \mathcal{A})$ and will revise the introduction accordingly.