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

Jianliang He; Leda Wang; Siyu Chen; Zhuoran Yang

cs.LGFeb 18, 2026

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

Authors: Jianliang He, Leda Wang, Siyu Chen, Zhuoran Yang

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

x

and

y

are represented as one-hot vectors in

\mathbb{R}^{p}

\sigma(\cdot)

denotes the activation function, and the width of the neural network is denoted by

M

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

Panel of Figure 2 — Figure 2: Learned parameters under the full random initialization with $p=23$ 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 $(p-1)/2=11$ 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.

Panel of Figure 3 — Figure 3: Visualizations of learned phases with $M=512$ neurons. Figure (a) plots the relationship among the normalized $2\phi_{m}$ and $\psi_{m}$ , with all points lying around the line $y=x$ . Figure (b) shows the uniformity of the learned phases within a specific group $\mathcal{N}_{k}$ . The left panel displays $\iota\phi_{m}$ for $\iota\in\{1,2,3,4\}$ on unit circles, and the points are nearly uniformly distributed. The right panel quantifies this symmetry by computing the averages of $\cos(\iota\phi_{m})$ and $\sin(\iota\phi_{m})$ , all of which are close to zero. Figure (c) presents violin plots of the magnitudes $\alpha_{m}$ and $\beta_{m}$ . The tight distribution of these values around their mean suggests that the neurons learn nearly identical magnitudes.

$\sigma(x)$	$\max\{x,0\}$	$\|x\|$	$x^{2}$	$x^{4}$	$x^{8}$	$\log(1+e^{2x})$	$e^{x}$	$x$	$x^{3}$
Loss	$1.194\times 10^{-8}$	0.000	0.000	$3.1\times 10^{-5}$	0.051	$1.2\times 10^{-3}$	$6.5\times 10^{-4}$	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.

Panel of Figure 4 — Figure 4: Illustration of the lottery-ticket mechanism under the fully random initialization. Figure (a) plots the dynamics of every frequency $k$ for a specific neuron, with the red curve tracing the trajectory of the frequency that eventually dominates. In the left-hand plot, misalignment levels ${\mathcal{D}}_{m}^{k}$ are rescaled to $[-\pi,\pi)$ 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 $\beta_{m}^{k}$ with various $(\beta_{m}^{k}(0),{\mathcal{D}}_{m}^{k}(0))$ after $10,000$ steps. The contours are symmetric about $\pi$ , 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.

Panel of Figure 5 — 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 $m^{-1}\sum_{m=1}^{M}|\sin({\mathcal{D}}_{m}^{\star})|$ . 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 $\ell_{2}$ -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.

Panel of Figure 8 — Figure 8: Visualizations of the alignment behavior and neuron evolution dynamics with $\kappa_{\sf init}=0.02$ . Figure (a) illustrates the dynamics of the normalized phase difference ${\mathcal{D}}_{m}^{\star}(t)$ given by ( 5.9 ). Initialized randomly on the unit circle, the gradient flow will always drive ${\mathcal{D}}_{m}^{\star}(t)$ to $0$ , regardless of the initial half-space. Figure (b) plots the dynamics of magnitudes and phases of the feature frequency for a specific neuron $m$ during the initial stage of training. $2\phi_{m}^{\star}$ and $\psi_{m}^{\star}$ evolves to align, and magnitudes $\alpha_{m}^{\star}$ and $\beta_{m}^{\star}$ starts growing rapidly once the phases are well-aligned.

Panel of Figure 9 — Figure 9: Learned feature and dynamics of parameters initialized at Assumption 5.1 with $p=23$ 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 “ $3k^{\star}$ ”, “ $5k^{\star}$ ” for $\theta_{m}$ , and “ $2k^{\star}$ ”, “ $3k^{\star}$ ” for $\xi_{m}$ . Figure (b) plots the dynamics of a specific neuron $m$ . Here, the phase quickly aligns, i.e., $\psi_{m}^{\star}\approx 2\phi_{m}^{\star}$ , and the magnitudes $\alpha_{m}^{\star}$ and $\beta_{m}^{\star}$ 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 $20$ neurons with $p=23$ 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	$1$ Freqs	$2$ Freqs	$4$ Freqs	$8$ Freqs	Full Freqs
Average	$1.64$	$6.02\times 10^{-1}$	$2.88\times 10^{-2}$	$2.99\times 10^{-8}$	$7.41\times 10^{-15}$
Standard Deviation	$2.01\times 10^{-2}$	$8.79\times 10^{-2}$	$1.55\times 10^{-2}$	$1.07\times 10^{-7}$	$-$
Part II: Phase Diversity Ablation.
	$[0,0.4\pi)$	$[0,0.8\pi)$	$[0,1.2\pi)$	$[0,1.6\pi)$	$[0,2\pi)$
Loss	$4.82$	$2.00\times 10^{-3}$	$1.19\times 10^{-9}$	$3.54\times 10^{-7}$	$7.41\times 10^{-15}$

Appendix

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

|\mathcal{K}|

. The results show that the fully diversified parametrization achieves the lowest CE loss, confirming its maximum efficiency under the fixed constraints of model scale

\alpha_{m}\beta_{m}^{2}=1

and neuron budget

M=128

Panel of Figure 13 — Figure 13: Output logits for the predictor under different ablation configurations, evaluated across four distinct query points $(x,y)$ . 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 $20$ neurons initialized under Assumption 5.1 with $p=23$ 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 $k^{\star}$ , while keeping the others $0$ throughout.

Part I: Dynamics of Original Parameters.
$\theta_{m}[j](t)$	$\partial_{t}\theta_{m}[j](t)\approx 2p\cdot\alpha_{m}^{\star}(t)\cdot\beta_{m}^{\star}(t)\cdot\cos(\omega_{k}j+\psi_{m}^{\star}(t)-\phi_{m}^{\star}(t)),\quad\forall j\in[p]$
$\xi_{m}[j](t)$	$\partial_{t}\xi_{m}[j](t)\approx p\cdot\alpha_{m}^{\star}(t)^{2}\cdot\cos(\omega_{k^{\star}}j+2\phi_{m}^{\star}(t)),\quad\forall j\in[p]$
Part II: Dynamics of Dicrete Fourier Coefficients.
$g_{m}[2k^{\star}](t)$	$\partial_{t}g_{m}[2k^{\star}](t)\approx\sqrt{2}\cdot p^{3/2}\cdot\alpha_{m}^{\star}(t)\cdot\beta_{m}^{\star}(t)\cdot\cos\big(\psi_{m}^{\star}(t)-\phi_{m}^{\star}(t)\big)$
$g_{m}[2k^{\star}+1](t)$	$\partial_{t}g_{m}[2k^{\star}+1](t)\approx-\sqrt{2}\cdot p^{3/2}\cdot\alpha_{m}^{\star}(t)\cdot\beta_{m}^{\star}(t)\cdot\sin\big(\psi_{m}^{\star}(t)-\phi_{m}^{\star}(t)\big)$
$r_{m}[2k^{\star}](t)$	$\partial_{t}r_{m}[2k^{\star}](t)\approx p^{3/2}/\sqrt{2}\cdot\alpha_{m}^{\star}(t)^{2}\cdot\cos\big(2\phi_{m}^{\star}(t)\big)$

Appendix

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.

Panel of Figure 15 — Figure 15: Training dynamics of a specific decoupled neuron characterized by ( B.47a ) and ( B.47b ) with identical initial scales $\alpha_{m}^{\star}(0)=\beta_{m}^{\star}(0)$ and different phase difference ${\mathcal{D}}_{m}^{\star}(0)$ . Figure (a) plots the dynamics of phases, phase difference, and the magnitudes with ${\mathcal{D}}_{m}^{\star}(0)\in(\pi/2,\pi)$ , whose behavior is detailedly characterized in Theorem B.7 . The difference decreases monotonically to $0$ , while the magnitudes first decay slightly when ${\mathcal{D}}_{m}^{\star}(t)\in(\pi/2,\pi)$ and then increase rapidly when ${\mathcal{D}}_{m}^{\star}(t)$ falls below $\pi/2$ . Figure (b) plots the dynamics under ${\mathcal{D}}_{m}^{\star}(0)\in(0,\pi/2)$ where ${\mathcal{D}}_{m}^{\star}$ 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.

Lottery Ticket Hypothesis

See all 10 papers

Oct 8, 2026

Probability-Signature Dynamics: Unpacking Modular Addition Learning Within Two-Layer Networks

Neural networks trained on modular addition tasks often develop Fourier-structured representations that support exact generalization. While prior work has identified these Fourier circuits, the mechanism by which gradient-based training selects them from the data distribution remains unclear. We address this question using probability signatures, which express leading gradient interactions through conditional statistics of the training distribution. For modular addition, these signatures are cyclic shift operators and are diagonalized by the discrete Fourier transform, yielding approximately decoupled Fourier-mode dynamics. This explains the emergence of Fourier sparsity, frequency matching, and phase alignment. The same framework resolves a puzzle under label noise: corrupted examples can show faster early loss decrease than clean examples, despite lacking a coherent generalization rule. We show that noise increases conditional label collisions, strengthening early shared-coordinate reinforcement. Finally, this method can be applied to other operators. Taking XOR as an example, we observed the predicted frequency in experiments.

Yunji Wang, Junjie Yao, Linyu Liu +2

School of Mathematical Sciences, Shanghai Jiao Tong University · Huawei Technologies Ltd · Shanghai University of Finance and Economics +1

Neural Network Training Dynamics Frequency-Domain Feature Learning Modular Addition

Ask PDF

Mar 24, 2026

Latent Algorithmic Structure Precedes Grokking: A Mechanistic Study of ReLU MLPs on Modular Arithmetic

Grokking-the phenomenon where validation accuracy of neural networks on modular addition of two integers rises long after training data has been memorized-has been characterized in previous works as producing sinusoidal input weight distributions in transformers and multi-layer perceptrons (MLPs). We find empirically that ReLU MLPs in our experimental setting instead learn near-binary square wave input weights, where intermediate-valued weights appear exclusively near sign-change boundaries, alongside output weight distributions whose dominant Fourier phases satisfy a phase-sum relation

φ_{\mathrm{out}} = φ_a + φ_b

; this relation holds even when the model is trained on noisy data and fails to grok. We extract the frequency and phase of each neuron's weights via DFT and construct an idealized MLP: Input weights are replaced by perfect binary square waves and output weights by cosines, both parametrized by the frequencies, phases, and amplitudes extracted from the dominant Fourier components of the real model weights. This idealized model achieves 95.5% accuracy when the frequencies and phases are extracted from the weights of a model trained on noisy data that itself achieves only 0.23% accuracy. This suggests that grokking does not discover the correct algorithm, but rather sharpens an algorithm substantially encoded during memorization, progressively binarizing the input weights into cleaner square waves and aligning the output weights, until generalization becomes possible.

Anand Swaroop

Grokking Neural Network Interpretability ReLU Neural Networks Neural Network Memorization

Ask PDF

Apr 28, 2026

Feature Repulsion and Spectral Lock-in: An Empirical Study of Two-Layer Network Grokking

Tian (2025) proves a repulsion theorem (Theorem 6) for the matrix

B = (\widetilde{F}^\top \widetilde{F} + ηI)^{-1}

during the interactive feature-learning stage of grokking: similar features have negative off-diagonal entries

B_{j\ell}

, producing an effective repulsive force that drives them apart. However, the theorem does not specify when this mechanism becomes empirically observable, nor whether it leaves a measurable spectral signature in the parameter updates. We test this directly on Tian's modular addition setup (

M = 71

K = 2048

, MSE loss) and observe a clear structure-mechanism dissociation. The predicted sign rule holds robustly on the top-200 most-similar feature pairs across activations (empirical sign-match rising from 0.865 to 0.985 on

σ= x^2

across 5 seeds, and saturating at 1.000 on

σ= \operatorname{ReLU}

). However, the spectral signature in the parameter updates is strongly activation-dependent. With

σ= x^2

, a simple slope detector on the rolling eigengap

σ_2 / σ_3

ΔW

fires in 15/15 grokking seeds at epoch 174 (IQR [173,174]) and in 0/15 non-grokking controls, with 229

\times

late-stage magnitude separation; the spectrum is rank-2. In contrast, with

σ= \operatorname{ReLU}

, the detector never fires and the spectrum remains effectively rank-1. This dissociation aligns with Tian's Theorem 5 distinction between focused (power-law) and spreading (ReLU) memorization: while the sign structure of

B

depends only on

\widetilde{F}^\top \widetilde{F}

, how feature repulsion translates into weight updates critically depends on the activation derivative

σ'

Yongzhong Xu

Neural Network Training Dynamics Modular Arithmetic Grokking

Ask PDF

Abstract

Lottery Ticket Hypothesis

Explore similar work

Probability-Signature Dynamics: Unpacking Modular Addition Learning Within Two-Layer Networks

Latent Algorithmic Structure Precedes Grokking: A Mechanistic Study of ReLU MLPs on Modular Arithmetic

Feature Repulsion and Spectral Lock-in: An Empirical Study of Two-Layer Network Grokking