PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yunji Wang, Junjie Yao, Linyu Liu, Pinyan Lu, Zhi-Qin John Xu

arXiv:2610.11833v1cs.AI

TL;DR

The paper asks why gradient-based training discovers Fourier structure in modular addition, rather than merely showing that such circuits can represent the task. It uses probability signatures, conditional statistics of the training distribution, to analyze leading dynamics. The analysis accounts for Fourier-mode behavior and faster early loss decrease under random label noise, and suggests prediction of neuron structures for other operators.

  • Problem

    Why gradient-based training selects Fourier-structured parameters from the modular-addition data distribution remains unresolved, although prior work identified such circuits.

    Source: §1 · p. 2

  • Method

    The paper applies probability signatures, sampling-derived conditional statistics, to characterize leading gradient interactions and derive Fourier modes from the data-side operators.

    Source: §3.2 · p. 4, §4.1 · p. 5

  • Results

    For modular addition, conditional-distribution operators recover the Fourier basis and approximately decoupled leading dynamics; random label noise yields faster initial loss decrease, and the operator view suggests predictions for other binary operations.

    Source: §7 · p. 12

  • Takeaways & Limitations

    The operator viewpoint provides a data-side route to explaining Fourier-mode selection and predicting structures for tasks beyond modular addition.

    Source: §7 · p. 12

    The Fourier-basis derivation assumes no noise and training on the entire clean dataset, with a finite alphabet whose prescribed size P is prime.

    Source: §3.1 · p. 3, §4.1 · p. 5

Abstract

from arXiv · show

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.

1 Introduction

Prior work identified Fourier circuits in modular-addition networks, but why gradient training selects them from the data distribution remained unresolved. This paper uses probability signatures to connect training-distribution statistics to leading gradient dynamics, addressing Fourier structure, noisy-label loss behavior, and extensions to other operators.

1 Introduction

  • Grokking is a delayed test-performance improvement after prolonged training, even when training loss has already become very small.In modular addition, its onset accompanies a distinctive structure in the weights.Source: §1 · p. 1
  • The open question is why gradient-based training discovers Fourier-structured parameters, since representability alone does not imply that training will select them.Prior work characterized the circuits in trained networks, but the forward path from the modular-addition distribution remained unexplained.Source: §1 · p. 2
  • Probability signatures explain the emergence of dominant Fourier modes by linking training-distribution statistics to leading gradient interactions.The analysis asks which statistics enter leading parameter dynamics and how they determine mode selection.Source: §1 · p. 2
  • Uniformly random label noise consistently produces a faster initial loss decrease than clean data, a difference measurable through training-distribution statistics.The paper frames this as a puzzle in the early memorization regime.Source: §1 · p. 2
  • The probability-signature analysis extends to other tasks, including XOR, to predict structures network parameters may form before training.The paper studies a two-layer embedding MLP on modular addition with randomly corrupted labels.Source: §1 · p. 2

2 Related Work

Related work describes grokking, Fourier circuits in modular arithmetic, and learning with corrupted labels. This paper positions its contribution as a forward, data-distribution-based account of Fourier structure and one early-time noise phenomenon, rather than a general theory of label noise.

Grokking

  • Prior grokking studies explain delayed generalization through optimization trajectories, weight decay, initialization, and transitions between training regimes [Liu et al., 2023; Kumar et al., 2024].This paper connects weight decay to Fourier-structure formation through a threshold.Source: §2 · p. 3

Modular arithmetic

  • Earlier modular-arithmetic work characterized sparse Fourier circuits and diagnostics such as frequency matching and phase alignment [Nanda et al., 2023; Michaud et al., 2023].The paper instead derives the Fourier basis as common eigenvectors of probability signatures, without assuming the circuit in advance.Source: §2 · p. 3

Learning with corrupted labels

3 Preliminaries

The preliminaries specify a finite-alphabet binary-operation learning task, a two-layer MLP, and its training setup. They define probability signatures as sampling-derived conditional statistics used to characterize leading gradient terms.

3.1 Settings

  • The task predicts a binary operation on a finite alphabet Ω = {0, 1, …, P−1}, with prescribed prime P.The input is a pair (a, b), and the target is F(a, b).Source: §3.1 · p. 3
  • Clean and test labels use modular addition, while noise labels are uniform random values; training combines clean and noisy examples.The training-to-test ratio is 8:2, and the noise fraction is α.Source: §3.1 · pp. 3–4
  • The model is a two-layer MLP with one-hot inputs, softmax outputs, cross-entropy loss, L2 weight decay, and a continuous gradient-flow approximation.The hidden dimension D is much larger than P; weights are independently initialized with scale controlled by γ.Source: §3.1 · p. 4

3.2 Probability signatures

  • Probability signatures are sampling-derived statistics, independent of training, that characterize main gradient terms across training stages.The framework is extended to the model with shared embedding structure.Source: §3.2 · p. 4
  • The gradient-flow proposition derives leading parameter dynamics under full-batch cross-entropy training, small initialization, and a smooth activation near zero.Its conditional expectations are taken under the training distribution with one input position fixed; higher-order terms are O(ϵ^3).Source: §3.2 · p. 4
  • The signatures comprise conditional probabilities and joint conditional-probability tensors for the two input positions and output.The paper denotes these statistics by φ, Φ, and ψ.Source: §3.2 · pp. 4–5

4 Theoretical Analysis

The analysis derives leading training dynamics from probability signatures of the data distribution. For modular addition, these signatures yield independent Fourier-mode dynamics and explain frequency selection, while weight decay filters which modes grow.

4 Theoretical Analysis

  • The gradient-flow analysis separates terms by their associated probability signatures and characterizes the dominant dynamics under stated conditions.For sampled data, the analysis also estimates the gap from the full-addition signature.Source: §4 · p. 5, §4.2 · p. 7

4.1 The Fourier basis derived from the signatures

  • The modular-addition signatures are shift operators whose common eigenbasis is the DFT, which approximately decouples the dynamics at each nonzero frequency.Each frequency's three-variable system evolves independently.Source: §4.1 · pp. 5–6
  • Growing neurons develop a dominant frequency and phase alignment; the first frequency to blow up becomes dominant under the proposition's stated conditions.Weight decay provides a unified survival threshold, while initial weights and random asymmetry shape which frequencies cross it.Source: §4.1 · p. 6

4.2 Estimation of noise

  • The noise-estimation theorem approximates the effect of sampling by comparing the sampled dataset with the full-addition signature.The passage describes this as bridging the gap caused by sampling.Source: §4.2 · p. 7

5 Validation on Fourier frequency

The experiments connect neuron growth with Fourier structure and check the framework's predictions against training behavior. They report that growth concentrates among high-sparsity neurons, while frequency matching and phase alignment emerge during training.

5 Validation on Fourier frequency

  • With GELU, noise ratio α = 0.1, and initialization scale γ = 0.5, λ = 0.0005 regularizes training so the model generalizes while ignoring noise.The passage says weight decay strongly affects training direction.Source: §5 · p. 8
  • Norm growth is concentrated among high-sparsity neurons, and its timing coincides with improving accuracy.Most neuron norms shrink under weight decay, while a few grow and the newly grown neurons have relatively high sparsity.Source: §5 · p. 8
  • The estimated gradient achieves cosine similarity ≥0.9 with the actual gradient in the reported numerical experiments.In a separate full-batch check, predicted dominant frequencies agree with the converged frequencies of the actual GELU network for 93% of valid neurons.Source: §5 · pp. 8–9
  • For neurons reaching high sparsity, the dominant frequencies of U_d, V_d, and W_d tend to agree, with their phases tending to align.Figure 3 shows frequency matching and phase alignment emerging together during training.Source: §5 · pp. 8–9, Figure 3 · p. 9

6 Application of Probability Signature

Probability signatures explain why uniformly random label noise can reduce loss faster early in training and suggest how to predict representations for other operations. For XOR, the analysis predicts a Hadamard basis, consistent with observed weight patterns.

6.1 The puzzle of early noise

  • Uniform random labels can produce a faster early loss decrease than modular-addition labels because their expected shared-coordinate counts are higher.For clean/noise mixtures, the predicted noise-loss decrease is faster, consistent with experiments.Source: §6.1 · p. 10
  • The paper attributes faster early noise-loss decrease to greater conditional label collisions among examples sharing an input coordinate.The loss-slope comparison is derived by counting sample pairs that share coordinates.Source: §6.1 · p. 10

6.2 Extension of the Method: XOR as an example

  • When joint probability-signature operators share an eigenbasis, transforming parameters into it approximately decouples leading-order dynamics, allowing the basis to be predicted from the data.For XOR, the Hadamard basis simultaneously diagonalizes the operators.Source: §6.2 · p. 11
  • For XOR, the analysis predicts Hadamard-like weight structure, and experiments report growth of high-sparsity neurons with frequency matching.The reported observations include high correlation between frequency sparsity and norm.Source: §6.2 · pp. 11–12

7 Conclusion and Discussion

The paper uses probability signatures to account for Fourier structure in modular addition, faster early loss decrease under label noise, and Hadamard structure in XOR. Its dynamical analysis is limited to dominant gradient terms within a particular training regime.

7 Conclusion and Discussion

  • For modular addition, conditional-distribution operators recover the Fourier basis and approximately decouple leading-order dynamics, accounting for dominant frequencies, frequency matching, and phase alignment.Source: §7 · p. 12
  • Under uniformly random label noise, conditional label collisions among examples sharing an input coordinate explain faster initial loss decrease for corrupted examples than clean examples.Source: §7 · p. 12
  • For XOR, the joint signature identifies the Hadamard basis as a candidate coordinate system, consistent with observed trained-weight patterns.Source: §7 · p. 12
  • The analysis approximates empirical gradient flow with idealized full-distribution signatures and corrections, and does not claim validity throughout the entire training process.Source: §7 · p. 12

A Model and training details

The appendix specifies the network and training setup, then compares activation functions and explains their differing generalization behavior through the second-order term in the gradient approximation.

A Model and training details

  • The model uses input embedding matrices U and V and output matrix W, with hidden width D = 2000 and modulus P = 113.Source: A Model and training details · p. 14
  • Adam experiments use a 10^-3 initial learning rate and 10^-5 minimum; AdamW generally needs about 1000x more weight decay for similar effects.Too little decay can prevent generalization, while too much can prevent convergence; the authors find Adam provides better regularization in this task.Source: A Model and training details · pp. 15–16

B Discussion on Activation Functions

  • Quadratic and ReLU activations can generalize well under appropriate weight decay, while linear activation learns poorly and tanh has low test generalization.Source: B Discussion on Activation Functions · p. 16
  • The analysis attributes this contrast to the second-order term: quadratic activation contains it, whereas tanh lacks it and linear activation does not learn well.The authors note that ReLU's success is compatible with the analysis because ReLU is nondifferentiable near zero.Source: B Discussion on Activation Functions · p. 17

C.1 Gradient Flow Approximation

The appendix tests gradient approximations, frequency selection, and ReLU behavior. The approximated gradients closely match the actual gradients, while observed frequency concentration and a sufficient Fourier construction support the proposed dynamics.

C.1 Gradient Flow Approximation

  • The approximated gradients have 90% similarity to actual gradients, and a quadratic-activation ablation shows that second-order main terms suffice for test generalization.Source: §C.1 · p. 18
  • The gradient approximation focuses on the dominant terms; above scale h0, higher-order terms dominate, and weight decay can then limit further parameter growth.An appropriate decay places its effect between the second-order and higher-order regimes.Source: §C.1 · p. 18

C.2 Frequency Selectivity

  • About 50% of neurons developed a dominant frequency, and 0.93 of that subset converged to the same frequency under shared initialization.Source: §C.2 · p. 19

D Discussion on the ReLU activation

D.1 Fourier-mode is able to do modular addition under RELU

  • Under the stated Fourier construction, the ReLU network predicts modular addition correctly with probability tending to 1 as width D grows.Source: §D.1 · p. 20

D.2 Experiment Details of ReLU activation

The ReLU experiments show behavior similar to GELU, while the appendix analyzes early example-wise memorization and develops gradient-flow and noise-estimation calculations.

D.2 Experiment Details of ReLU activation

  • ReLU exhibits weight-decay selectivity and frequency-phase scatter evolution similar to GELU, suggesting these phenomena are not specific to smooth activations.Source: §D.2 · p. 22

D.3 Early Memorization

  • The memory score accumulates rapidly early in training, then gradually decreases under the norm-reducing effect of weight decay.Source: §D.3 · p. 25

E.1 Calculation of the gradient flow

  • The gradient-flow calculation retains constant and first-order terms in the hidden representation, with the remaining higher-order terms grouped into a remainder of order ϵ^3.Source: §E.1 · pp. 26–27

E.2 Proof of Dominance of the gradient

  • The proof concludes its gradient-flow expression after adding the weight-decay term.Source: §E.2 · p. 29

E.3 Proof of Estimation of noise

  • The noise-estimation error depends on sampling ratio ρ and noise ratio α; too few clean training samples can make the actual and predicted gradients differ substantially.Source: §E.3 · p. 30

E.4 Proof of dominant frequency and phase locking

The proof shows how Fourier-like mode competition selects a dominant frequency and aligns its phase, then establishes an early-loss inequality and the common Walsh–Hadamard eigenbasis for XOR operators.

E.4 Proof of dominant frequency and phase locking

  • A growing mode’s energy fraction tends to one as its energy diverges while competing modes remain finite.For phase locking, δ = 0 is stable while amplitudes remain nonzero; continuous random initialization avoids the unstable δ = π equilibrium with probability one, yielding δ(t) → 0 modulo 2π.Source: §E.4 · p. 31

E.5 Proof of the early-loss theorem

  • Sharing an input gives E[f(a + b)f(a + c)] ≥ E[f(a + b)f(c + d)], strictly for activations such as GELU and ReLU.Source: §E.5 · pp. 31–32

E.6 Propositions of Hadamard basis

  • Every Walsh vector is a common eigenvector of the XOR operators, which are simultaneously diagonalized by the normalized Walsh–Hadamard basis.Source: §E.6 · pp. 33–34

F.1 Experiment on Modular Multiplication Task

The modular multiplication experiments compare early loss across moduli, while a bijective-noise construction tests how label collisions affect which data lose first. The observed loss patterns match the stated predictions.

F.1 Experiment on Modular Multiplication Task

  • For modular multiplication, M = 60 showed a significantly faster early loss decline than M = 59 and M = 61, consistent with more concentrated labels and larger K0.For the prime moduli 59 and 61, the signatures are more uniformly distributed, leading to smaller K0 and K1.Source: §F.1 · pp. 34–35

F.2 Experiment on the Bijective Noise

  • The quadratic-residue Legendre-symbol construction gives bijective labels in each coordinate and is not isomorphic to addition.It is used to replace uniform random noise in the experiment.Source: §F.2 · pp. 35–36
  • With bijective noise, the noisy data lose first when α < 0.5, but the ordering reverses when α > 0.5.The observed pattern is reported as consistent with the theoretical prediction in Figure 18.Source: Figure 18 · p. 35, §F.2 · p. 36
Loading 2610.11833v1…