Clustering and Alignment: Understanding the Training Dynamics in Modular Addition

TL;DR

This paper investigates how neural networks trained on modular addition develop interpretable structures in embeddings, revealing clustering and alignment tendencies through explicit formulas and particle simulations, and explores the impact of regularization on training dynamics.

Contribution

It introduces explicit interaction force formulas for embedding structures and demonstrates their emergence via particle simulations, advancing understanding of training dynamics in neural networks.

Findings

Embedding vectors form grid and circle structures.

Clustering and alignment tendencies explain structure emergence.

Weight decay influences training dynamics and regularization.

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/.