How Transformers Get Rich: Approximation and Dynamics Analysis

TL;DR

This paper provides a detailed theoretical and dynamic analysis of how transformers implement induction heads, revealing an abrupt transition from lazy to rich in-context learning mechanisms during training.

Contribution

It formalizes induction head mechanisms and analyzes their implementation and training dynamics, offering new insights into transformer in-context learning.

Findings

Transformers can efficiently implement induction heads.

Training exhibits an abrupt transition from lazy to rich mechanisms.

Dynamics analysis reveals the training process of induction heads.

Abstract

Transformers have demonstrated exceptional in-context learning capabilities, yet the theoretical understanding of the underlying mechanisms remains limited. A recent work (Elhage et al., 2021) identified a ``rich'' in-context mechanism known as induction head, contrasting with ``lazy'' $n$-gram models that overlook long-range dependencies. In this work, we provide both approximation and dynamics analyses of how transformers implement induction heads. In the {\em approximation} analysis, we formalize both standard and generalized induction head mechanisms, and examine how transformers can efficiently implement them, with an emphasis on the distinct role of each transformer submodule. For the {\em dynamics} analysis, we study the training dynamics on a synthetic mixed target, composed of a 4-gram and an in-context 2-gram component. This controlled setting allows us to precisely characterize the entire training process and uncover an {\em abrupt transition} from lazy (4-gram) to rich (induction head) mechanisms as training progresses.