On the Mechanism and Dynamics of Modular Addition: Fourier Features, Lottery Ticket, and Grokking

TL;DR

This paper analyzes how two-layer neural networks learn modular addition, revealing the mechanisms of feature combination, phase symmetry, and frequency diversification, and explaining phenomena like grokking through a comprehensive theoretical framework.

Contribution

It provides a full mechanistic and theoretical explanation of feature learning, training dynamics, and generalization in neural networks solving modular addition, including the lottery ticket hypothesis and grokking.

Findings

Neurons learn Fourier features with phase alignment and frequency diversification.

Phase symmetry enables majority voting to cancel noise and identify correct sums.

Grokking is characterized as a three-stage process involving memorization and two generalization phases.

Abstract

We present a comprehensive analysis of how two-layer neural networks learn features to solve the modular addition task. Our work provides a full mechanistic interpretation of the learned model and a theoretical explanation of its training dynamics. While prior work has identified that individual neurons learn single-frequency Fourier features and phase alignment, it does not fully explain how these features combine into a global solution. We bridge this gap by formalizing a diversification condition that emerges during training when overparametrized, consisting of two parts: phase symmetry and frequency diversification. We prove that these properties allow the network to collectively approximate a flawed indicator function on the correct logic for the modular addition task. While individual neurons produce noisy signals, the phase symmetry enables a majority-voting scheme that cancels out noise, allowing the network to robustly identify the correct sum. Furthermore, we explain the emergence of these features under random initialization via a lottery ticket mechanism. Our gradient flow analysis proves that frequencies compete within each neuron, with the "winner" determined by its initial spectral magnitude and phase alignment. From a technical standpoint, we provide a rigorous characterization of the layer-wise phase coupling dynamics and formalize the competitive landscape using the ODE comparison lemma. Finally, we use these insights to demystify grokking, characterizing it as a three-stage process involving memorization followed by two generalization phases, driven by the competition between loss minimization and weight decay.