Clustering and Alignment: Understanding the Training Dynamics in Modular Addition
This work provides insights into the fundamental mechanisms of neural network training dynamics, which is incremental but offers a new perspective on regularization effects.
The paper investigates how neural networks develop interpretable algorithms during training by analyzing the training dynamics of a small network on modular addition, revealing that embedding vectors organize into grid and circle structures due to clustering and alignment tendencies, with explicit formulae proposed to model these interactions and validated through particle simulations.
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/.