Structure Development in List-Sorting Transformers

TL;DR

This paper investigates how a one-layer attention transformer learns to organize its attention mechanisms for list sorting, revealing natural tendencies towards simpler structures and the influence of training data on internal organization.

Contribution

It demonstrates the emergence of two main attention modes in transformers during sorting tasks and links these to data features and training dynamics, advancing understanding of model development.

Findings

Attention heads organize into vocabulary-splitting and copy-suppression modes.

Vocabulary-splitting occurs regardless of regularization, indicating inherent simplicity bias.

Training data features influence the model's internal structure development.

Abstract

We study how a one-layer attention-only transformer develops relevant structures while learning to sort lists of numbers. At the end of training, the model organizes its attention heads in two main modes that we refer to as vocabulary-splitting and copy-suppression. Both represent simpler modes than having multiple heads handle overlapping ranges of numbers. Interestingly, vocabulary-splitting is present regardless of whether we use weight decay, a common regularization technique thought to drive simplification, supporting the thesis that neural networks naturally prefer simpler solutions. We relate copy-suppression to a mechanism in GPT-2 and investigate its functional role in our model. Guided by insights from a developmental analysis of the model, we identify features in the training data that drive the model's final acquired solution. This provides a concrete example of how the training data shape the internal organization of transformers, paving the way for future studies that could help us better understand how LLMs develop their internal structures.