Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Clustering and Alignment: Understanding the Training Dynamics in Modular Addition

Tiberiu Musat

TL;DR

This article studies the training dynamics of a small neural network with 2-dimensional embeddings on the problem of modular addition and proposes explicit formulae for these tendencies as interaction forces between different pairs of embeddings.

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

Clustering and Alignment: Understanding the Training Dynamics in Modular Addition

TL;DR

This article studies the training dynamics of a small neural network with 2-dimensional embeddings on the problem of modular addition and proposes explicit formulae for these tendencies as interaction forces between different pairs of embeddings.

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/.
Paper Structure (27 sections, 5 equations, 6 figures, 2 tables, 2 algorithms)

This paper contains 27 sections, 5 equations, 6 figures, 2 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. My contribution
  3. Architecture
  4. Grids and circles: the key to successful generalization
  5. Co-evolution of embeddings and linear layers
  6. Clustering
  7. Alignment
  8. Particle simulation
  9. Conclusion
  10. Algorithms
  11. Circle detection
  12. Grid imperfections
  13. Weight decay and the magnitude of weights and biases
  14. The role of weight decay
  15. Embedding vectors
  16. ...and 12 more sections

Figures (6)

  • Figure 1: Embedding vectors self-organize into grids (left) and circles (right).
  • Figure 2: Sometimes embedding vectors self-organize into imperfect grids.
  • Figure 3: Embedding vectors (left); classifier formed by the combined linear and output layers (background, right) and sums of embedding pairs in the training set (markers, right).
  • Figure 4: Embedding vectors (left); classifier formed by the combined linear and output layers (background, right) and sums of embedding pairs in the training set (markers, right).
  • Figure 5: Particles form circles (left), grids (center), and imperfect grids (right).
  • ...and 1 more figures