Dissecting Chain-of-Thought: Compositionality through In-Context Filtering and Learning (2305.18869v2)

Abstract: Chain-of-thought (CoT) is a method that enables LLMs to handle complex reasoning tasks by decomposing them into simpler steps. Despite its success, the underlying mechanics of CoT are not yet fully understood. In an attempt to shed light on this, our study investigates the impact of CoT on the ability of transformers to in-context learn a simple to study, yet general family of compositional functions: multi-layer perceptrons (MLPs). In this setting, we find that the success of CoT can be attributed to breaking down in-context learning of a compositional function into two distinct phases: focusing on and filtering data related to each step of the composition and in-context learning the single-step composition function. Through both experimental and theoretical evidence, we demonstrate how CoT significantly reduces the sample complexity of in-context learning (ICL) and facilitates the learning of complex functions that non-CoT methods struggle with. Furthermore, we illustrate how transformers can transition from vanilla in-context learning to mastering a compositional function with CoT by simply incorporating additional layers that perform the necessary data-filtering for CoT via the attention mechanism. In addition to these test-time benefits, we show CoT helps accelerate pretraining by learning shortcuts to represent complex functions and filtering plays an important role in this process. These findings collectively provide insights into the mechanics of CoT, inviting further investigation of its role in complex reasoning tasks.

The paper "Dissecting Chain-of-Thought: Compositionality through In-Context Filtering and Learning" explores the mechanics of the Chain-of-Thought (CoT) method in LLMs, particularly transformers. The authors aim to understand how CoT enables complex reasoning by decomposing tasks into simpler steps, which significantly enhances in-context learning (ICL) capabilities.

Key Contributions

1. Compositional Function Decomposition:

The paper investigates the application of CoT to in-context learning of compositional functions, specifically multi-layer perceptrons (MLPs). The analysis reveals that CoT operates by decomposing the task into two phases:

• Filtering Phase: Isolating data pertinent to each compositional step.
• Learning Phase: In-context learning the function for each individual step.

2. Experimental and Theoretical Evidence:

Through experiments and theoretical analysis, the paper shows that CoT effectively reduces the sample complexity necessary for ICL. This reduction makes learning more efficient compared to non-CoT methods, particularly in handling complex functions.

3. Transition in Transformers:

The authors illustrate how transformers can evolve from standard in-context learning to mastering compositional functions using CoT. This transition is achieved by:

• Incorporating additional layers.
• Utilizing the attention mechanism for data filtering required by CoT.

4. Training Efficiency:

Beyond test-time benefits, the paper demonstrates that CoT facilitates accelerated pretraining. This acceleration is attributed to the ability of CoT to learn shortcuts in representing complex functions. The data-filtering process intrinsic to CoT aids in this pretraining phase, highlighting its integral role.

Implications

The findings have significant implications for the development of more efficient and capable LLMs. By understanding the mechanics of CoT, researchers can enhance the ability of transformers to handle complex reasoning tasks, offering a pathway to more robust and interpretable AI systems. The insights garnered from this paper pave the way for further investigations into the underlying principles of CoT and its potential applications in various domains requiring advanced problem-solving capabilities.