Selective Induction Heads: How Transformers Select Causal Structures In Context

2025-09-09

AI Summary

Existing studies model induction heads using fixed-structure Markov chains, failing to capture the context-dependent dynamism inherent in natural language. Method: We propose the “Interleaved Delayed Markov Chain” framework, enabling Transformers to adaptively select optimal causal structures across varying contexts. We identify and construct a selective induction head circuit wherein self-attention and in-context learning jointly orchestrate dynamic structural switching; we theoretically prove its convergence to the maximum-likelihood solution. Contribution/Results: Experiments on a three-layer Transformer demonstrate that the model accurately infers context-optimal time delays and precisely replicates target tokens, simultaneously selecting causal structure during sequence prediction. This work overcomes the limitations of static causal modeling and establishes a novel paradigm for understanding the dynamic inductive mechanisms underlying Transformer architectures.

Technology Category

• Machine Learning: Causal Learning
• Knowledge Representation and Reasoning: Action, Change, and Causality
• Reasoning under Uncertainty: Causality

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• 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

Abstract

Transformers have exhibited exceptional capabilities in sequence modeling tasks, leveraging self-attention and in-context learning. Critical to this success are induction heads, attention circuits that enable copying tokens based on their previous occurrences. In this work, we introduce a novel framework that showcases transformers' ability to dynamically handle causal structures. Existing works rely on Markov Chains to study the formation of induction heads, revealing how transformers capture causal dependencies and learn transition probabilities in-context. However, they rely on a fixed causal structure that fails to capture the complexity of natural languages, where the relationship between tokens dynamically changes with context. To this end, our framework varies the causal structure through interleaved Markov chains with different lags while keeping the transition probabilities fixed. This setting unveils the formation of Selective Induction Heads, a new circuit that endows transformers with the ability to select the correct causal structure in-context. We empirically demonstrate that transformers learn this mechanism to predict the next token by identifying the correct lag and copying the corresponding token from the past. We provide a detailed construction of a 3-layer transformer to implement the selective induction head, and a theoretical analysis proving that this mechanism asymptotically converges to the maximum likelihood solution. Our findings advance the understanding of how transformers select causal structures, providing new insights into their functioning and interpretability.

Problem

• How transformers dynamically handle varying causal structures in context
• Understanding how transformers select correct causal dependencies between tokens
• Mechanism for identifying appropriate token relationships in changing contexts

Innovation

• Interleaved Markov chains with varying lags
• Selective Induction Heads circuit implementation
• Dynamic causal structure selection mechanism