Clustering and Alignment: Understanding the Training Dynamics in Modular Addition
Keywords: grokking, modular addition, training dynamics, particle system, effective theory, transformer, weight decay
TL;DR: I propose an explanation of the training dynamics in a small neural network by modelling the embedding vectors as a particle system
Abstract: Recent studies have revealed that neural networks learn interpretable algorithms for many simple problems. However, little is known about how these algorithms emerge during training. In this article, I study the training dynamics of a small neural network with 2-dimensional embeddings on the problem of modular addition. I observe that embedding vectors tend to organize into two types of structures: grids and circles. I study these structures and explain their emergence as a result of two simple tendencies exhibited by pairs of embeddings: clustering and alignment. I propose explicit formulae for these tendencies as interaction forces between different pairs of embeddings. To show that my formulae can fully account for the emergence of these structures, I construct an equivalent particle simulation where I show that identical structures emerge. I discuss the role of weight decay in my setup and reveal a new mechanism that links regularization and training dynamics. To support my findings, I also release an interactive demo available at https://modular-addition.vercel.app/.
Track: Main track
Submitted Paper: No
Published Paper: No
Submission Number: 14
Loading