cs.LGFeb 18, 2026
SaveOn the Mechanism and Dynamics of Modular Addition: Fourier Features, Lottery Ticket, and Grokking
Organizations: Department of Statistics and Data Science, Yale University
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.
Figures & tables
Figure 1: An illustration of the primary analytical technique and results. Discrete Fourier Transform (DFT) is utilized to quantitatively interpret the mechanism of learned models within the feature space, revealing the training dynamics that result in consistent feature learning. Figure (a) shows the neural network architecture — we adopt a two-layer fully connected neural network to learn the modular addition task. The inputs and are represented as one-hot vectors in , denotes the activation function, and the width of the neural network is denoted by . Figure (b) illustrates the technique of DFT. We apply DFT to the weights at the input and output layers, respectively. Each neuron involves two weight vectors, which lead to two magnitudes and phases. (See Observation 1 in § 3 .) Figure (c) illustrates some of our key empirical observations — phase alignment (Observation 2), phase symmetry (Observation 3), and lottery ticket mechanism (Observation 6).
Figure 2: Learned parameters under the full random initialization with and ReLU activation using AdamW. Figure (a) plots a heatmap of the learned parameters for the first 10 neurons after Discrete Fourier Transform (DFT, see § 5.1 ) grouped with frequency. Each row in the heatmap corresponds to the Fourier components of a single neuron’s parameters. The plot clearly reveals a single-frequency pattern: each neuron exhibits a large, non-zero value focused on only one specific frequency component, confirming a highly sparse and specialized frequency encoding. We remark that since only 10 out of 512 neurons are shown, not all frequencies appear in this sample. The same single-frequency pattern holds across all 512 neurons, which collectively cover all 11 frequencies (see Observation 3). Figure (b) further examines the periodicity by plotting line plots of the learned parameters for three neurons, each overlaid with a trigonometric curve fitted via DFT. The fitted curve aligns almost perfectly with the actual one.
Figure 3: Visualizations of learned phases with neurons. Figure (a) plots the relationship among the normalized and , with all points lying around the line . Figure (b) shows the uniformity of the learned phases within a specific group . The left panel displays for on unit circles, and the points are nearly uniformly distributed. The right panel quantifies this symmetry by computing the averages of and , all of which are close to zero. Figure (c) presents violin plots of the magnitudes and . The tight distribution of these values around their mean suggests that the neurons learn nearly identical magnitudes.
| Loss | 0.000 | 0.000 | 0.051 | 4.246 | 3.891 | ||||
| Accuracy | 1.000 | 1.000 | 1.000 | 1.000 | 1.000 | 1.000 | 1.000 | 0.041 | 0.036 |
Table 1: Adaptivity of the learned parameterization. We evaluate the robustness of the trained model by replacing the ReLU with various alternative functions at test time. As shown in the table, the model maintains perfect prediction accuracy when using the absolute value function, even-order polynomials, or the exponential function. This demonstrates that the learned features are not strictly dependent on the original activation but rather on its underlying even-order components.
Figure 4: Illustration of the lottery-ticket mechanism under the fully random initialization. Figure (a) plots the dynamics of every frequency for a specific neuron, with the red curve tracing the trajectory of the frequency that eventually dominates. In the left-hand plot, misalignment levels are rescaled to for clarity. Typically, the winning frequency is the one that starts with a comparatively larger initial magnitude and a smaller misalignment. Figure (b) plots the contour of the magnitude with various after steps. The contours are symmetric about , and reproduce the trend seen in Figure (a): initial states with larger magnitude and lower phase misalignment yield higher final magnitudes after the same training duration.
Figure 5: Progress measure of grokking behavior. The shaded regions mark three distinct phases: an initial memorization phase, followed by two generalization phases. Figures (a) and (b) plot the train-test loss and accuracy curve, where the network first overfits the training data to achieve a near-zero training loss while the test loss remains high. Figure (c) visualizes the dynamics of average phase alignment level, measured by . Figure (d) tracks the evolution of the average neuron-wise frequency sparsity level, as measured by the inverse participation ratio (IPR) of the Fourier coefficients, alongside the -norm of the parameter.
Figure 6: An illustration of the three stages of grokking dynamics and their main driving force.
Figure 7: Heatmap of the output logits with quadratic activation, with boxes indicating predicted higher values.
Figure 8: Visualizations of the alignment behavior and neuron evolution dynamics with . Figure (a) illustrates the dynamics of the normalized phase difference given by ( 5.9 ). Initialized randomly on the unit circle, the gradient flow will always drive to , regardless of the initial half-space. Figure (b) plots the dynamics of magnitudes and phases of the feature frequency for a specific neuron during the initial stage of training. and evolves to align, and magnitudes and starts growing rapidly once the phases are well-aligned.
Figure 9: Learned feature and dynamics of parameters initialized at Assumption 5.1 with and ReLU activation. Figure (a) shows heatmaps of the parameters after DFT at initialization and at the end of the initial stage. Similar to the quadratic activation (see Figure 14 ), the single-frequency pattern is approximately maintained, with small values emerging at frequencies “ ”, “ ” for , and “ ”, “ ” for . Figure (b) plots the dynamics of a specific neuron . Here, the phase quickly aligns, i.e., , and the magnitudes and grow rapidly and synchronously.
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 10: Heatmaps of trained model from Figure 5 at the end of the memorization stage . The left panel displays the data distribution: dark blue entries represent training data, light blue entries are test data whose symmetric counterparts are in the training set, and white entries (outlined in red) are the remaining held-out test data. The middle panel shows the model’s accuracy, demonstrating that it has perfectly memorized all training data and their symmetric variants but completely fails to generalize to the held-out data. Finally, the right panel visualizes the model’s post-softmax output on the correct answer for each data point, further confirming the accuracy results.
Figure 11: Data distribution during the memorization stage. The first panel illustrates the data partitioning, which, unlike in Figure 10 , uses the following scheme: white entries denote test data, dark brown entries represent common (symmetric) training data, and light brown entries (outlined in red) denote rare (asymmetric) training data. The remaining three plots track the model’s accuracy, demonstrating a two-stage memorization scheme. At initialization, the model performs at a low, chance-level accuracy. However, after approximately 1000 steps, it masters the common symmetric training data, but its performance on rare asymmetric data drops to zero, overwriting any initially correct random predictions. By the end of the memorization stage, the model finally memorizes these rare data points, achieving 100% training accuracy
Figure 12: Heatmaps of parameters after applying discrete Fourier transform along training epoches for the first neurons with under train-test split setup. At the end of the memorization stage (step 2200), a single-frequency pattern has started to emerge, accompanied by noisy perturbations in other frequencies. This initial "perturbed Fourier solution" is subsequently refined, as weight decay prunes the noisy, non-feature frequencies to reveal the final, clean pattern.
| Part I: Frequency Diversity Ablation. | |||||
| Loss | Freqs | Freqs | Freqs | Freqs | Full Freqs |
| Average | |||||
| Standard Deviation | |||||
| Part II: Phase Diversity Ablation. | |||||
| Loss | |||||
Table 2: Performance of the predictor under different ablation configurations. For the frequency ablation study, the average and standard deviation of the loss are reported across all possible combinations of frequencies of the specified size . The results show that the fully diversified parametrization achieves the lowest CE loss, confirming its maximum efficiency under the fixed constraints of model scale and neuron budget .
Figure 13: Output logits for the predictor under different ablation configurations, evaluated across four distinct query points . The true prediction label is indicated by the dashed vertical line in each panel. The fully diversified parametrization yields the largest logit gap between the ground truth and incorrect labels, signifying maximal prediction confidence.
Figure 14: Heatmaps of parameters after applying discrete Fourier transform along training epoches for the first neurons initialized under Assumption 5.1 with and quadratic activation. At the initial stage, these neurons preserve the single-frequency pattern by evolving only the Fourier coefficients corresponding to the initial frequency , while keeping the others throughout.
| Part I: Dynamics of Original Parameters. | |
| Part II: Dynamics of Dicrete Fourier Coefficients. | |
Table 3: Summarization of the approximate dynamics during the initial stage. Please refer to § B.3.2 , § B.4 and § B.5 for formalized arguments and detailed derivations.
Figure 15: Training dynamics of a specific decoupled neuron characterized by ( B.47a ) and ( B.47b ) with identical initial scales and different phase difference . Figure (a) plots the dynamics of phases, phase difference, and the magnitudes with , whose behavior is detailedly characterized in Theorem B.7 . The difference decreases monotonically to , while the magnitudes first decay slightly when and then increase rapidly when falls below . Figure (b) plots the dynamics under where is initialized closer to the convergence point, resulting in a shorter convergence time compared to the case in Figure (a). Moreover, the simplified dynamics shown in Figure (b) align well with the full dynamics in Figure 8 with the same initialization, indicating the effectiveness of the approximation.