Clustering Head: A Visual Case Study of the Training Dynamics in Transformers
This provides insights into transformer training mechanisms, but it is incremental as it focuses on a specific synthetic task.
The paper tackles the sparse modular addition task by examining how transformers learn it, revealing 'clustering heads' as circuits that learn invariants and analyzing training dynamics, including two-stage learning and loss spikes.
This paper introduces the sparse modular addition task and examines how transformers learn it. We focus on transformers with embeddings in $\R^2$ and introduce a visual sandbox that provides comprehensive visualizations of each layer throughout the training process. We reveal a type of circuit, called "clustering heads," which learns the problem's invariants. We analyze the training dynamics of these circuits, highlighting two-stage learning, loss spikes due to high curvature or normalization layers, and the effects of initialization and curriculum learning.