Causal In-Context Learning in Transformers: Training Dynamics Across Heterogeneous Interventional Data

Shanyun Gao, Murat Kocaoglu, Qifan Song

Keywords: Transformer; In-Context Learning; Causal Discovery; Interventional Data; Heterogeneous Data

Abstract: Large language models (LLMs) have shown strong abilities in generalization, adaptation, and reasoning. Among these capabilities, in-context learning (ICL) has emerged as a defining property of LLMs. ICL enables models to separate and exploit heterogeneous data arising from different environments or tasks while producing context-dependent solutions. In this work, we show that such in-context learning capabilities can arise even in a simple two-layer Transformer with a single attention head per layer when trained on heterogeneous data with causal structure. Through an analysis of training dynamics during next-token prediction on diverse interventional datasets, we observe two phases: an initial phase that learns the parents of a target variable, followed by a later phase that captures the full Markov boundary. By leveraging the crucial role of positional embeddings, we show that in-context learning can be used for local causal discovery under certain interventions without knowing from which environment or interventional distribution each prompt is drawn.

Submission Number: 26