CoT-ICL Lab: A Petri Dish for Studying Chain-of-Thought Learning from In-Context Demonstrations

2025-02-21

AI Summary

The mechanistic role of chain-of-thought (CoT) prompting in in-context learning (ICL) remains poorly understood. Method: We propose the first fine-grained, controllable CoT-ICL modeling framework. By generating synthetic tokenized datasets, we decouple causal structure from token mapping functions, enabling precise control over ICL complexity. Using decoder-only Transformers (≤700M parameters), we conduct embedding-space and attention-map analyses to isolate architectural and functional factors. Contribution/Results: We establish that model depth is critical for effective CoT utilization in few-shot settings; moderate restriction of function diversity improves causal structure learning. Empirically, CoT accelerates accuracy transitions across all model scales; shallow models match deep-model performance given sufficient in-context examples. Our work provides the first interpretable, theory-grounded empirical foundation for CoT-enhanced ICL, bridging mechanistic insight with controllable experimental design.

Technology Category

• Machine Learning: Causal Learning
• Reasoning under Uncertainty: Causality
• Cognitive Modeling & Cognitive Systems: Conceptual Inference and Reasoning

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

We introduce CoT-ICL Lab, a framework and methodology to generate synthetic tokenized datasets and systematically study chain-of-thought (CoT) in-context learning (ICL) in language models. CoT-ICL Lab allows fine grained control over the complexity of in-context examples by decoupling (1) the causal structure involved in chain token generation from (2) the underlying token processing functions. We train decoder-only transformers (up to 700M parameters) on these datasets and show that CoT accelerates the accuracy transition to higher values across model sizes. In particular, we find that model depth is crucial for leveraging CoT with limited in-context examples, while more examples help shallow models match deeper model performance. Additionally, limiting the diversity of token processing functions throughout training improves causal structure learning via ICL. We also interpret these transitions by analyzing transformer embeddings and attention maps. Overall, CoT-ICL Lab serves as a simple yet powerful testbed for theoretical and empirical insights into ICL and CoT in language models.

Problem

• Study chain-of-thought learning in language models.
• Generate synthetic datasets for in-context learning.
• Analyze transformer embeddings and attention maps.

Innovation

• Generates synthetic tokenized datasets
• Decouples causal structure from processing
• Trains decoder-only transformers effectively