Selective induction heads: how transformers select causal structures in context

This research explores how transformers adapt to changing causal structures in data, which is crucial for understanding their success in language processing. They introduce a new test using interleaved Markov chains with varying "lags" or dependencies. The paper shows that a three-layer transformer can learn to identify the correct lag and predict the next token, a process termed selective induction heads. A detailed construction for how attention weights achieve this is provided, demonstrating that this mechanism converges to the maximum likelihood solution and is implemented by both specially designed and standard trained transformers. Experimental results confirm this theoretical understanding and show that transformers learn this lag-selection ability.