Attention with Markov: A Framework for Principled Analysis of Transformers via Markov Chains

TL;DR

This paper introduces a Markov chain-based framework to analyze transformer models, explaining why single-layer transformers often fail to learn bigram distributions and how deeper models succeed, supported by theoretical and empirical evidence.

Contribution

It provides a novel Markov chain framework for analyzing transformers, characterizes the loss landscape of single-layer models, and explains their empirical failure to learn bigram distributions.

Findings

Single-layer transformers often get trapped in local minima representing unigram distributions.

Deeper transformers reliably learn the in-context bigram distribution.

Theoretical analysis matches empirical observations of model behavior.

Abstract

Attention-based transformers have achieved tremendous success across a variety of disciplines including natural languages. To deepen our understanding of their sequential modeling capabilities, there is a growing interest in using Markov input processes to study them. A key finding is that when trained on first-order Markov chains, transformers with two or more layers consistently develop an induction head mechanism to estimate the in-context bigram conditional distribution. In contrast, single-layer transformers, unable to form an induction head, directly learn the Markov kernel but often face a surprising challenge: they become trapped in local minima representing the unigram distribution, whereas deeper models reliably converge to the ground-truth bigram. While single-layer transformers can theoretically model first-order Markov chains, their empirical failure to learn this simple kernel in practice remains a curious phenomenon. To explain this contrasting behavior of single-layer models, in this paper we introduce a new framework for a principled analysis of transformers via Markov chains. Leveraging our framework, we theoretically characterize the loss landscape of single-layer transformers and show the existence of global minima (bigram) and bad local minima (unigram) contingent on data properties and model architecture. We precisely delineate the regimes under which these local optima occur. Backed by experiments, we demonstrate that our theoretical findings are in congruence with the empirical results. Finally, we outline several open problems in this arena. Code is available at https://github.com/Bond1995/Markov .