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

2024/11/18 by Musat, Tiberiu · 1 citation
#Computation and Language (cs.CL) #FOS: Computer and information sciences #Machine Learning (cs.LG)

paper · doi:10.48550/arxiv.2411.12118

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.

Cited by

• On the Emergence of Induction Heads for In-Context Learning

Related

• Is Random Attention Sufficient for Sequence Modeling? Disentangling Trainable Components in the Transformer
• Systematic Generalization and Emergent Structures in Transformers Trained on Structured Tasks
• Memorization Capacity of Multi-Head Attention in Transformers
• The End of Transformers? On Challenging Attention and the Rise of Sub-Quadratic Architectures
• Talking Heads: Understanding Inter-layer Communication in Transformer Language Models