Dissecting Chain-of-Thought: Compositionality through In-Context Filtering and Learning

Summary

This paper investigates how the transformer-based model can compositionally learn some complex functions (e.g., an MLP) by breaking them down to some atom problems (e.g., linear mapping). Such an ability is also the crux of the success of Chain-of-thought (CoT) in-context learning (ICL) methods in large language models. Hence this paper bridges the gap between these two fields by theoretically analyzing the sample complexity of different methods and the accelerating effect of CoT in pertaining. Although I believe the contribution of the paper is solid, novel, and helpful to the field, there are plenty of problems that make the paper hard to follow. I wish the authors could polish the paper and tackle some of my concerns to make the paper stronger. Hence at this stage, I would only give a borderline rejection. I would be very happy to increase my score during the rebuttal phase.

Strengths

See the summary part.

Weaknesses

1. I find the paper hard to follow, maybe because the theoretical part, which I believe is the most important contribution, is too abstract. It is not easy to get intuition from current section 3.2, so maybe explaining how theorem 1 is formulated is helpful. 2. The experimental parts are not well organized. First, it appears almost everywhere (in section 2,3,4), which breaks the flow of the paper significantly. Also, some of the subfigures in Fig3,4,5 contain similar information. So it would be good to re-organize these results to make the paper easier to follow. 3. In section 4.3, the paper switches from the ICL setting to the pertaining setting. It is beneficial to clearly explain what’s the difference between them and why we need experiments in this setting.

Questions

1. For the experiment in Figure 1, will ICL and CoT-I converge to a lower value if the number of samples keeps increasing? Because the sample complexity of CoT-I/O is lower than the other two and it converges at around 100 samples, would the other two methods converge if we use 500 samples? 2. In Figure 2,3, it is hard to distinguish different methods using dashed lines, using different markers might be good. 3. In Figure 3, if we want to see the influence of k, why not draw a figure using k as the x-axis? 4. Figure 3 is on page 2, but it is first referred at page 4 — it is hard to locate the figure when reading the paper. 5. Some typos and imprecise expressions: a.) In section 3.1, the first sentence, ‘we train 2-layer MLPs’. Do we train this MLP? IIUC, this MLP is used to generate the training samples and is fixed all the time. b.) In section 3.2, the last paragraph, ‘consider the consider the condition’.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

1 poor

Contribution

3 good

Limitations

I’m not sure whether this is an unavoidable problem of the CoT-style ICL, but it would be nice if the author can answer the following problem somewhere in the paper, which might make the paper more solid. I think comparing the sample complexity (i.e., the number of examples in prompting) between these 3 algorithms is unfair. Imagine we have $k$ examples and the problem has L steps. Then in ICL, there are only $k$ input vectors and 1 supervisory signal. In CoT-I, there are $Lk$ input vectors and 1 supervisory signal. For CoT-I/O, there are $Lk$ input vectors and $(L-1)k$ supervisory signals (I’m not sure in this paper's setting, whether ground truth $s_n^l$ is accessible in CoT-I/O, but if that is true, the comparison would be more unfair.) Given the differences of the input and supervisory signals, it might be straightforward to conclude that CoT has smaller sample complexity than ICL.

Summary

The paper proposes to study chain-of-thought (prompting) in the setting of learning MLPs. The authors build on top of recent work studying in-context learning linear regression tasks in the light of gradient descent and extent their setting to learning non-linear functions. In order to study chain-of-thought prompting they either provide features / hidden activations from the "teacher" MLP the student (CoT-I setting) or make the Transformer produce its own input by "looping" (CoT-I/O setting). They provide theoretical results and empirical evidence that the Transformer is more sample efficient due to CoT by leveraging the given inputs and/or can remember weights coming from the family of teachers to allow for the self-production of features due to looping.

Strengths

I like the abstraction of CoT presented in the paper. The experiments are convincing and supported by the theoretical results. Although this setting is quite simple, it nicely extends the recent studies of gradient descent on linear regression. The results and the line of thinking is very intuitive, I like the ideas and the empirical execution of the paper.

Weaknesses

Although I am very familiar with the setup of the paper, I still had problems understanding it. The presentation can be made much clearer. Please work on your presentation and think about how to structure the paper in a clearer and structured way. It is a bit confusing that you have 1.5 empirical sections and the paper is very dense. I think it would benefit the paper if the authors work on restructuring the presentation given the (hopefully available) extra page if accepted. It would be nice to clearly explain what is meant with "learning an MLP". This differs from classic student teacher frameworks since you are not learning and also don't have to learn (in the CoT-I setting as far as I understand) the weight of the MLPs but solely rely on gradient descent on a linear regression task which acts on the given features in-context. Given these features in-context, it feels quite obvious that the CoT setting is indeed outperforming classic in-context learning. Nevertheless, I think it is still quite interesting. Also it would be helpful to clearly state your abstraction/hypotheses that in LLMs the additional data needed to outperform plain ICL are "features" of the data motivating your MLP abstraction. This only became clear to me after reading a few times. In the CoT-I/O setting, please explain why you are not using MLPs but linear deep nets as the teacher. Although the presented experiments make sense, the training of Transformers is usually done differently as TFs are not trained to do CoT-I or CoT-I/O. Please comment on this. There are a couple of ablations / interesting experiments I would like to see.

Questions

My main concern is the following: Given that CoT is given the activations of the teacher, the sequence length of these experiments is larger. I think it is quite crucial to contrast the performance of your trained models when given different input data. If I understand correctly, you are hypothesizing that CoT works in LLMs is because the extra data is representative of features of the input data. 3 Naive ablations come to mind: Just train a model on your setup where the extra data that is provided during CoT (s_0, ...,) is either constant, random or actually the same data given multiple times. That would provide more evidence that in your setting indeed the teachers activations are crucial. It could (I doubt it myself) be that just more naive prompts do the trick.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

I think the limitations that the setting is quite constructed and might not have anything to do how CoT works in LLMs could be a bit more strong.

Summary

This paper aims to demystify the mechanism lying in the in-context learning (ICL) and chain-of-thought (CoT). It reveals how CoT significantly reduces the sample complexity of ICL. It uses a two-layer MLP, and a backbone GPT-2 model for exploration. The experimental results reveal some interesting findings, e.g., the in-context samples needed is linearly dependent on the input dimension. The paper also provides theoretical analysis of probable approximation of MLPs via chain-of-thought.

Strengths

- The paper provides both experimental and theoretical evidence to support its claims about how CoT works. - The findings about the relations between the in-context samples needed and the input MLP dimension are interesting and insightful.

Weaknesses

- Lack of clarity: Some parts of the paper may be unclear or difficult to follow, which could make it challenging for readers to understand the key findings and contributions of the study. Some of the concepts are not explained very well in the beginning (e.g, what is compositional learning? What is MLPs in-context), thus making the paper abstruse. - The main contribution of the paper is not clear. The paper claims to dissect the mechanism of CoT

Questions

- How do you use the intermediate MLPs features as the CoT prompts for the GPT-2? It seems that the features are not tokens to be prompts. - Why do you use the input dimension d and hidden size k of a MLP layer as the measures for experiments? Are there any specific reasons to use them as the proxy for some properties of the target task?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

No limitation section is found.

Summary

In this paper, the authors explore the mechanics of Chain-of-Thought (CoT), a method that has successfully enabled language models to handle complex reasoning tasks by decomposing them into simpler steps. The study aims to understand the underlying mechanics of CoT by investigating its impact on the ability of transformers to in-context learn a simple yet general family of compositional functions: multi-layer perceptrons (MLPs). The authors reveal that the success of CoT can be attributed to two distinct phases: focusing on data related to each step of the composition and in-context learning the single-step composition function. They provide experimental and theoretical evidence that demonstrates how CoT significantly reduces the sample complexity of in-context learning (ICL) and facilitates the learning of complex functions that non-CoT methods struggle with.

Strengths

1. This work explores CoT in the learning of MLPs thatis a novel perspective. Such a simplicity can also help to dig the into inner machenism of CoT and ICL. 2. This work experimentally compare three schemes (ICL, CoT-I, and CoT-I/O), providing valuable insights into their differences and the benefits of CoT prompting. 3. The paper provides a formalized theorem that explains how a transformer architecture can realize the CoT process for MLPs.

Weaknesses

1. The Chain-of-Thought (CoT) implementation presented in this study appears to be limited in terms of reasoning steps, which may not be fully consistent with the original motivation behind proposing CoT as a method for complex reasoning tasks. In most existing work, CoT has the potential to assist in various complex reasoning tasks as it provides the model with step-by-step guidance on solving the input. Each reasoning step should include both the intermediate state and the process through which this state is generated from the previous one. This is typically implemented using natural language descriptions or formal language expressions. However, the CoT implementation in this paper focuses on the intermediate states in MLPs and does not provide information on how each intermediate state is produced. By only presenting intermediate states rather than complete reasoning steps, the CoT explored in this study primarily reflects the capability of state transition rather than reasoning. 2. The study appears to have limited analysis on generalization capabilities, which is a fundamental aspect of in-context learning (ICL). This study assumes that there is no distribution shift between the training and test datasets. However, a fundamental aspect of in-context learning (ICL) is the ability to learn new tasks from in-context examples. Consequently, it is crucial to examine the performance in a generalization setting. For example, [1], which shares a similar philosophy with this work, extensively discussed the behavior of ICL on out-of-distribution prompts. [1] Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes.

Questions

1. Do the conclusions drawn in this study remain consistent in an out-of-distribution setting? For example, what might occur if the MLPs in the test and training sets exhibit notable differences, such as having different widths?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

1 poor

Contribution

2 fair

Limitations

This paper does not explicitly address its limitations. The most significant limitation of the study is the unclear contribution towards answering the question of how the Chain-of-Thought (CoT) method can assist in solving complex reasoning tasks.

Thank you!

Thank you for your clarifications and running the additional experiments. Given the commitment for restructuring the paper, I vote to accept the paper but will not increase my score further. As a few of the reviewers also had troubles understanding the paper, I feel that this was indeed a shortcoming of the paper that hopefully will be resolved if accepted. I encourage the other reviewers to rethink their rather, in my opinion, low scores given that the authors will invest time to improve the papers presentation. I find the results very valuable for the audience of NeurIPS and should be discussed at the conference.

Thanks for the feedback.

Thanks very much for the authors' feedback, which solves most of my concerns. As I mentioned in my original review, I believe the contribution of this paper is valuable and inspiring. Considering the author will improve the presentation in the next version, I would increase my score from 4 to 6. I am looking forward to see the final version.

Thank You!

I would like to raise the score from 4 to 5, as the additional experiments addressed my concerns regarding generalization. Nevertheless, I still do not comprehend how this study contributes to elucidating the mechanism of CoT in practical reasoning tasks. Consequently, I can only assign a borderline rating for this work.