Mechanism and Emergence of Stacked Attention Heads in Multi-Layer Transformers

TL;DR

This paper investigates how multi-layer transformers develop attention mechanisms to solve retrieval tasks, revealing that attention heads emerge sequentially guided by an implicit curriculum, which enhances understanding of their internal workings.

Contribution

It introduces the retrieval problem as a reasoning task, demonstrates transformers can solve it without fine-tuning, and uncovers the sequential emergence of attention heads during training.

Findings

Transformers solve retrieval tasks with minimal layers.

Attention heads emerge in a specific sequence during training.

Implicit curriculum guides the emergence of attention mechanisms.

Abstract

In this paper, I introduce the retrieval problem, a simple yet common reasoning task that can be solved only by transformers with a minimum number of layers, which grows logarithmically with the input size. I empirically show that large language models can solve the task under different prompting formulations without any fine-tuning. To understand how transformers solve the retrieval problem, I train several transformers on a minimal formulation. Successful learning occurs only under the presence of an implicit curriculum. I uncover the learned mechanisms by studying the attention maps in the trained transformers. I also study the training process, uncovering that attention heads always emerge in a specific sequence guided by the implicit curriculum.