PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Jianliang He, Leda Wang, Siyu Chen, Zhuoran Yang

arXiv:2602.16849v1cs.LGmath.OCstat.ML

TL;DR

The paper asks how two-layer networks learn and combine Fourier features to solve modular addition and how those features emerge under standard training. It uses mechanistic analysis and gradient-flow theory to show that diversification, frequency competition, and phase dynamics explain the learned solution and grokking. The resulting account describes majority voting over neuron signals and a three-stage transition from memorization to generalization.

  • Problem

    Prior work had not provided a conclusive end-to-end explanation of how finite-neuron networks under standard training form Fourier features, combine them into modular-addition solutions, and grok.

  • Method

    The paper combines systematic experiments with rigorous analysis of two-layer networks, using Fourier-domain interpretation, phase dynamics, gradient flow, and the ODE comparison lemma.

  • Results

    The network uses frequency diversification and phase symmetry to aggregate noisy neuron signals, while within-neuron frequency competition selects winners based on initial magnitude and phase alignment.

  • Takeaways & Limitations

    The paper characterizes grokking as three stages in which weight decay prunes residual noise and refines a perturbed Fourier representation into a clean representation for generalization.

  • Takeaways & Limitations

    The quadratic activation analysis can fail when a few well-aligned neurons dominate, although normalized or spherical gradient descent may mitigate this issue.

Abstract

from arXiv · show

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.

1 Introduction

This paper gives an end-to-end account of how two-layer networks learn Fourier features for modular addition, combine them into a global solution, and transition from memorization to generalization.

  • Motivation: Prior work identified Fourier features and grokking, but finite-neuron dynamics under standard training remained incompletely explained.The paper addresses this gap through mechanistic and theoretical analysis.
  • Mechanistic Interpretation: The network combines frequency diversification with phase symmetry to aggregate biased neuron signals into a global approximation of the modular-addition logic.Frequency diversification covers the required Fourier components, while phase symmetry supports noise cancellation across neurons.
  • Training Dynamics: Within each neuron, competing frequencies produce a single winner determined by initial spectral magnitude and phase misalignment.The winning component grows exponentially faster than its competitors under the analyzed gradient flow.
  • Grokking: Grokking unfolds in three stages: memorization, followed by two generalization phases in which weight decay removes residual noise and refines Fourier features.The process is driven by competition between loss minimization and weight decay.

2 Preliminaries

The paper studies modular addition as a classification task using a two-layer neural network, softmax outputs, and cross-entropy loss.

  • Modular Addition: The modular-addition dataset contains every pair in Z_p^2 with label (x + y) mod p, then separates examples into disjoint training and test sets.Test accuracy measures prediction of unseen input pairs.
  • Two-Layer Neural Network: The model uses M hidden neurons without biases, with either canonical or trainable input embeddings.Each input x is assigned an embedding h_x in R^d, and each neuron has hidden and output-layer parameters.
  • Model and Objective: Experiments primarily use ReLU activation, while theoretical interpretations use the quadratic activation σ(x) = x^2.The paper applies softmax to network outputs and trains with cross-entropy loss.
  • Model and Objective: Each input pair is transformed into neuron-wise hidden representations, linearly combined into logits, and converted into a categorical distribution.The correct label is represented by a one-hot vector.

3 Empirical Findings

The trained network organizes modular-addition solutions into sparse Fourier features whose phases and magnitudes coordinate across neurons. Training dynamics explain both feature emergence from random initialization and grokking’s progression from memorization to generalization.

  • Fourier Features: Each neuron learns a trigonometric parameterization with one active Fourier frequency, producing a sparse and specialized frequency representation.The input and output parameters are cosine waves with learned magnitudes and phases.
  • Phase Coupling: The output phase equals twice the input phase modulo 2π, indicating layer-wise coupling in Fourier space.The learned phase relation is (2ϕm − ψm) mod 2π = 0.
  • Model Symmetry: For large networks, each frequency is represented, phases are approximately uniform within frequency groups, and neuron magnitudes remain similar.Within-group phase symmetry supports collective coordination without a single neuron dominating.
  • Activation Robustness: Replacing ReLU at inference preserves perfect accuracy for activations with nonzero even-order components but fails for x and x^3.The findings identify even-order components as the key activation property for the learned parameterization.
  • Feature Emergence: Within each neuron, frequencies compete, and the component with larger initial magnitude and smaller phase misalignment becomes dominant.The winning component accelerates rapidly while competitors grow more slowly.
  • Grokking: Grokking proceeds through memorization and two generalization stages as loss minimization and weight decay successively shape the frequency representation.The memorization phase fits training data while unseen-test loss remains high; later weight decay promotes frequency-domain sparsification and delayed generalization.

4 Mechanistic Interpretation of Learned Model

The learned model combines single-neuron trigonometric signals through diversified frequency-phase pairs. Phase symmetry cancels residual noise across neurons, allowing the aggregate logits to approximate the correct modular-addition indicator.

  • Single-Neuron Contribution and Majority Voting: Each neuron contributes a primary signal peaking at j = (x+y) mod p alongside frequency- and phase-dependent residual noise.The primary cosine term identifies the correct sum, while the remaining terms create spurious, input-dependent peaks.
  • Single-Neuron Contribution and Majority Voting: Although individual neurons are biased by their frequency-phase views, aggregating diverse votes enables perfect-accuracy prediction of the correct sum.The network distills diverse but noisy local signals through majority voting.
  • Majority-Voting Approximates Indicator via Overparameterization: Full diversification requires balanced frequency coverage, homogeneous neuron magnitudes, and uniform higher-order phase symmetry.The definition formalizes the frequency-phase structure needed for collective signal aggregation.
  • Majority-Voting Approximates Indicator via Overparameterization: For sufficiently large neuron groups, the diversified model approximates the modular-addition indicator within an arbitrarily small error under the stated magnitude condition.The bound uses a ≳(Np)^−1 · log(p/ϵ), and the true-signal peak exceeds noise peaks by aNp/8.
  • Majority-Voting Approximates Indicator via Overparameterization: Under full diversification, the true-signal logit exceeds spurious peaks, so softmax concentrates on x + y mod p when magnitudes grow sufficiently.The residual noise arises from the flawed indicator structure but is dominated after aggregation.
  • Majority-Voting Approximates Indicator via Overparameterization: Fully diversified parameterization is the most parameter-efficient ablation and produces the largest ground-truth-versus-incorrect-label logit gap.The comparison uses the same neuron-budget constraint across limited-frequency and non-uniform-phase alternatives.

5 Training Dynamics for Feature Emergence

The paper analyzes feature emergence by tracking two-layer-network parameters in Fourier space under gradient flow. Early dynamics decouple across neurons and frequencies, preserve single-frequency structure, and drive phase alignment through coupled magnitude-phase evolution.

  • 5.1 Background: Discrete Fourier Transform: The analysis applies a DFT-based change of basis to decompose input and output parameters into frequency components represented by magnitudes and phases.A prime modulus p > 2 is chosen for analytical convenience, and sine-cosine pairs are reparameterized in polar form.
  • 5.2 Gradient Flow: Under quadratic activation, complete-data cross-entropy training is studied through continuous-time gradient flow with a small, structured initialization.The theoretical setup uses one-hot embeddings and an initialization containing one frequency per neuron.
  • 5.3 Initial-Stage Dynamics: During the initial stage, small parameter magnitudes make the softmax nearly uniform, while loss minimization becomes approximately average-margin optimization.Despite near-uniform predictions, gradients still learn feature directions.
  • 5.4 Preservation of Single-Frequency Pattern: Initial-stage neuron dynamics decouple, so each neuron’s parameter evolution depends only on its own input and output parameters.This decoupling permits neuron-wise analysis of Fourier-feature emergence.
  • 5.4 Preservation of Single-Frequency Pattern: Non-feature frequencies remain negligible during the initial stage, preserving the approximately single-frequency pattern as feature magnitudes grow.The theorem assumes log M/M ≲ c−1/2 · (1 + o(1)), κinit = o(M−1/3), and Cend ≍ κinit.
  • 5.5 Neuron-Wise Phase Alignment: Phase alignment is coupled to magnitude growth through an exact conservation law, with opposite phase rotations converging toward alignment.Theorem 5.3 further characterizes alignment time through initial misalignment, tolerance, initialization scale, and modulus.

6 Theoretical Extensions

The extensions analyze frequency competition from multi-frequency initialization, ReLU leakage, and quantitative gradient reinforcement. They show that initial magnitude and phase alignment select a dominant frequency, while ReLU approximately preserves the learned single-frequency structure with decaying harmonic leakage.

  • 6.1 Lottery Mechanism under Multi-Frequency Initialization: Under multi-frequency initialization, frequency dynamics are decoupled and each component evolves within its own orthogonal subspace.This self-contained evolution creates the competitive setting required for the lottery ticket mechanism.
  • 6.1 Lottery Mechanism under Multi-Frequency Initialization: Better phase alignment accelerates magnitude growth, while larger magnitudes accelerate alignment, forming a self-reinforcing competition.The growth and alignment rates are governed by cosine and sine functions of phase misalignment.
  • 6.1 Lottery Mechanism under Multi-Frequency Initialization: The winning frequency is selected by initial spectral magnitude and phase alignment, with larger magnitude and smaller misalignment producing faster growth.The ODE comparison lemma formalizes this lottery-ticket selection rule.
  • 6.1 Lottery Mechanism under Multi-Frequency Initialization: The winning component dominates rapidly under the corollary’s stated dominance-time bound and high-probability conditions.The bound includes a probability of at least 1 − e^{Θ(p−c)} under mild conditions.
  • 6.2 Dynamics Beyond Quadratic Activation: Quadratic activation can produce unstable, highly imbalanced neurons because initially better-aligned neurons grow faster and dominate the output.The paper notes that normalized or spherical gradient descent can mitigate this imbalance.
  • 6.2 Dynamics Beyond Quadratic Activation: With ReLU, the single-frequency pattern is approximately preserved, while leakage appears at higher harmonics and decays with increasing multiplicative factor.Observed leakage occurs at odd multiples for θm and at frequencies such as 2k⋆ and 3k⋆ for ξm.
  • 6.2 Dynamics Beyond Quadratic Activation: ReLU gradient leakage decays as 1/r^2 relative to the dominant frequency, while the dominant-frequency gradient reinforces the existing feature direction without phase rotation.For the output layer, even harmonics receive zero leakage under the stated parity constraint.

7 Conclusion

The paper reverse-engineers modular addition in two-layer networks, linking Fourier-domain majority voting, frequency competition, and grokking’s three-stage dynamics.

  • The network implements modular addition through a Fourier-domain majority-voting algorithm based on phase alignment and model symmetry.
  • Frequencies compete within each neuron, with the winning component determined by initial magnitude and phase misalignment.
  • Grokking comprises memorization followed by two generalization phases, as weight decay transforms a perturbed Fourier representation into a clean, generalizable solution.
  • These findings provide insights into feature-learning dynamics that may extend to more general tasks.

A.1 Detailed Interpretation of Grokking Dynamics in Section 3.3

Grokking unfolds as memorization followed by two generalization phases, as loss minimization and weight decay progressively change the model’s dynamics. The model first exploits symmetry, then cleans Fourier structure while test accuracy approaches 100%.

  • Progress measures: The progress measure for grokking is defined over model outputs and the parameters θm and ξm.The paper introduces this measure to track the dynamics discussed in the grokking analysis.
  • Overall dynamics: The three-phase process is governed by a shifting balance between loss minimization and weight decay.The initial stage is loss-gradient dominated, while the two generalization stages increasingly reflect weight-decay effects.
  • Phase I: Memorization: Phase I reaches 100% training accuracy and roughly 70% test accuracy, but fails on truly unseen held-out test data.The model succeeds on training examples and test examples whose symmetric counterparts were included in training.
  • Phase II: Loss-Driven Norm Growth with Rapid Feature Cleanup: During Phase II, loss minimization continues reducing training loss while weight decay prunes unnecessary frequencies and slows parameter-norm growth.The loss gradient pushes parameters toward larger scale, whereas weight decay steers their direction and removes noisy components.
  • Phase III: Slow Cleanup Driven Solely by Weight Decay: Phase III begins when training loss is near-zero and test accuracy approaches 100%, allowing weight decay to dominate and slowly reduce the parameter norm.Test accuracy gradually converges to 100% during this slow final cleanup.

A.2 Ablations Studies for Fully-Diversified Parametrization

The ablations test frequency diversification and phase uniformity under fixed computational constraints. Results show that the fully diversified parametrization is most efficient, while restricted frequency or phase diversity degrades performance.

  • Frequency diversification: Frequency ablations restrict the learned frequencies to subsets with |K| = 1, 2, 4, or 8 while keeping selected-frequency phases uniformly distributed.The study averages loss across all frequency combinations of each specified size.
  • Phase uniformity: Phase-uniformity ablations retain all frequencies but restrict phases to intervals [0, ιπ) for ι ∈ {0.4, 0.8, 1.2, 1.6}, degrading performance.The results identify phase-distribution range as a second important dimension of diversification.
  • Ablation setup: Full frequency and phase diversification achieves the lowest CE loss under fixed model-scale and neuron-budget constraints.The comparison uses equivalent computational budgets with M = 128 and αmβm^2 = 1.
  • Frequency diversification: CE loss drops from 1.64 with |K| = 1 to 7.41 × 10^-15 with the full frequency set.This frequency ablation demonstrates the importance of spectral richness.
  • Prediction confidence: The fully diversified parametrization produces the largest logit gap between the true label and incorrect alternatives.This indicates the greatest prediction confidence among the ablation configurations shown.
  • Single-frequency stability: Under sufficiently small initialization, neurons preserve their initial single-frequency pattern throughout the early training process.The empirical result agrees with Theorem 5.2 and reflects approximately decoupled neuron dynamics.

B.1 Proof of Proposition 4.2

The proof combines phase-aligned Fourier expansions, frequency-group summation, and a softmax margin bound. It shows that diversified neurons produce a predictor concentrated on the correct modular sum, with an explicit approximation guarantee.

  • Neuron contributions: Phase alignment and homogeneous scaling transform each neuron’s contribution into cosine terms involving x + y − j and related frequencies.The proof uses the condition 2ϕm − ψm = 0 mod 2π and then groups neurons by frequency.
  • Fourier summation: Balanced frequencies and high-order phase symmetry yield a summed Fourier expression that isolates the modular-addition signal.The derivation expands the grouped contribution and sums it across frequencies.
  • Logit margin: For x ≠ y, the correct-label logit exceeds every incorrect label by aNp/8, with an even larger margin when x = y.This margin is the key input to the softmax concentration bound.
  • Softmax concentration: The softmax prediction satisfies ||smax ◦ f(x, y; ξ, θ) − e_(x+y) mod p||_1 ≤ p · exp(−aNp/8).The bound follows from the softmax lemma applied to the correct-label margin.
  • Error requirement: To achieve error ϵ, it suffices to choose a ≳ (Np)^−1 · log(p/ϵ).This gives the explicit scale required for the approximation guarantee.
  • Initial-stage approximation: During the initial stage, small-scale parameters make softmax gradient components approximately uniform with controllable error.This approximation supports the simplified gradient-flow analysis.
  • Gradient-flow dynamics: The approximate gradient dynamics depend primarily on each neuron’s own parameters, so neurons are almost fully decoupled up to small approximation errors.The analysis tracks original parameters, Fourier coefficients, magnitudes, and phases to characterize the central flow.

B.4 Proof of Theorem 5.2: Single-Frequency Preservation

The proof shows that, under small random initialization, the target Fourier frequency grows while non-feature frequencies remain controlled by approximation errors. High-probability phase alignment and bounded errors preserve a single-frequency structure during the initial stage.

  • Initial-stage control: The initial stage ends when the parameter scale first reaches Cend, with tinit bounded on the order of 1/(pκinit) times a logarithmic factor.The proof defines tinit as the first hit time at which the parameters reach the prescribed threshold.
  • Initial-stage control: With probability greater than 1 − M^-c, the feature frequency dominates non-feature frequencies throughout the initial interval.The bound holds uniformly for t ∈ (0, tinit].
  • Initialization and phase alignment: Random initialization supplies a well-aligned neuron with high probability, enabling the feature coefficient’s growth and the single-frequency preservation argument.The phase event is obtained from concentration under the initialization assumptions.
  • Initial-stage control: The feature coefficient can grow rapidly even when perturbed by bounded approximation noise, while non-feature frequencies remain governed by controlled error terms.The proof establishes this behavior during the initial interval defined by the threshold Cend.
  • Initialization and phase alignment: Zero initialization eliminates non-feature frequency terms initially, while the proof bounds their later evolution through approximation-error estimates.The argument uses the stated initialization and inequalities controlling the error terms.

B.5 Proof of Theorem 5.3: Phase Alignment

The phase-alignment proof reduces the dynamics to a tractable Fourier-domain system and shows that phase differences converge to zero while limiting phases retain a uniform distribution. The convergence time depends on the initial phase difference through a symmetry-based two-case analysis.

  • Convergence time: For initial phase differences in (0, π/2), convergence takes Δtδ, whereas differences in (π/2, π) require 2Δtδ.The second case is decomposed using the symmetry around the π/2 hitting time.
  • Approximation: The reduced dynamics align well with the full dynamics under the same initialization, supporting the effectiveness of the approximation.Figure 15 compares the simplified and full trajectories for a decoupled neuron.
  • Phase convergence: The phase difference converges to zero, implying the limiting phase-alignment condition 2ϕ⋆ = ψ⋆.The limiting phases are determined by initialization, while the phase difference vanishes during training.
  • Phase distribution: The phase dynamics preserve a uniform marginal distribution over phases for every time t and neuron m.The result states ϕ⋆m(t) and ψ⋆m(t) are independently uniformly distributed in their respective phase coordinates.
  • Phase distribution: The limiting joint phase measure is concentrated on the periodic line 2ϕ = ψ while retaining uniform marginals.The limiting measure is represented as the pushforward T#λunif with T(φ) = (φ, 2φ) mod 2π.

B.5.1 Proof of Auxiliary Lemma B.8, B.9 and B.10

The auxiliary lemmas reduce the Fourier-domain dynamics from a high-dimensional system to a low-dimensional, effectively one-dimensional flow. They establish invariant quantities, phase symmetry, and monotone phase-difference behavior needed for the alignment theorem.

  • Reduced dynamics: Under single-frequency initialization, the 2p-dimensional dynamics can be tracked through a reduced four-particle system.The reduction follows from Fourier-domain structure and preservation of the single-frequency pattern.
  • Phase dynamics: The phase itself cannot be followed directly by a smooth ODE because atan2 introduces jumps from −π to π.The proof instead tracks exponentiated phase variables to avoid these discontinuities.
  • Invariant quantities: The reduced flow has initialization-dependent conserved quantities that constrain the evolution of magnitudes and phase differences.Lemma B.9 identifies constant relationships throughout the gradient flow.
  • Phase dynamics: The phase difference remains in the half-space determined by its initialization and never reaches the boundary where a modulo-2π jump would occur.This invariant-region property makes the phase difference directly trackable.

C.1 Proof of Corollary 6.1: Phase Lottery Ticket

The phase lottery-ticket proof shows that decoupled frequencies preserve their initial ranking and that the frequency with the smallest initial phase gap becomes dominant. Its advantage grows from a near-tie into rapid separation under the gradient flow.

  • Proof strategy: The multi-frequency analysis inherits the single-frequency lemmas because neuron decoupling and Fourier orthogonality make frequencies evolve independently.The proof treats the single-frequency dynamics as the template for each frequency component.
  • Winner selection: The winning frequency is the one with the smallest initial phase difference, equivalently the largest initial phase alignment, under equal initial magnitudes.Order preservation ensures that this identity remains unchanged throughout training.
  • Winner selection: Frequency rankings by magnitude remain fixed throughout the gradient flow, with larger magnitudes corresponding to smaller initial phase differences.This follows from neuron decoupling, orthogonality, and the cooperative-system rank-preservation argument.
  • Dominance growth: The winner-to-runner-up magnitude ratio first exceeds one and then grows super-exponentially toward infinity.The proof separates the initial escape from the stationary ratio 1 and the subsequent accelerated growth.
  • Dominance growth: At the end of the initial growth interval, the ratio satisfies ρm(t1) ≳ 1 + π^4p^-4(c+1) for sufficiently large p.This establishes a positive advantage before the super-exponential phase.

D Comparison with Existing Results

The comparison distinguishes prior analyses based on modified losses or mean-field limits from this paper’s standard cross-entropy treatment of finite-neuron frequency dynamics.

  • Prior settings: Prior work studies two-layer networks for group multiplication, a generalization of modular addition, using a modified ℓ2-loss to reduce constant-frequency noise.This paper instead analyzes standard cross-entropy loss with small initialization.
  • This paper’s dynamics: The paper’s standard CE loss naturally decouples frequency dynamics, while the constant frequency has zero gradient and remains zero under zero-initialization.The network output is normalized by 1/2p within the loss calculation.
  • Fourier representation: The paper represents weights with complex Fourier coefficients and establishes their one-to-one correspondence with real cosine-sine DFT components.Conjugate symmetry determines the coefficients at paired frequencies k and p − k from a single real component.
  • Loss and dynamics: Earlier loss analysis decomposes the modified loss into per-frequency terms whose variables are third-order monomials of Fourier coefficients.The cited mean-field treatment derives heuristic gradient dynamics under symmetric Gaussian initialization as M →∞.
  • Structural consequences: The appendix ODE solution supports phase alignment and single-frequency sparsity, while mean-field work models parameters as a continuous distribution and Wasserstein gradient flow.For matched frequencies, the relevant imaginary component converges to zero; for distinct frequencies, the corresponding phase-coupling term vanishes.
Loading 2602.16849v1…