Search papers, labs, and topics across Lattice.
This paper provides a mechanistic interpretation of how two-layer neural networks learn modular addition, explaining how individual neurons combine Fourier features into a global solution. It formalizes a diversification condition consisting of phase symmetry and frequency diversification, proving that these properties allow the network to approximate an indicator function for modular addition. The paper further explains the emergence of these features via a lottery ticket mechanism and characterizes grokking as a three-stage process involving memorization and two generalization phases.
Two-layer neural nets solve modular addition via a "flawed indicator function" built from noisy single-frequency Fourier features, where phase symmetry enables a majority-voting noise cancellation scheme.
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.
