A Counterexample to Fourier Alignment in Single-Neuron Modular Addition

Technical Summary: A Counterexample to Fourier Alignment in Single-Neuron Modular Addition

Problem Statement
The paper addresses a specific question regarding the representation-theoretic properties of neural networks trained on modular arithmetic, a task central to understanding "grokking" and algorithmic generalization. Specifically, it investigates the problem raised in [Cla26d] concerning a single ReLU neuron trained via exact full-table cross-entropy on addition in the cyclic group $\mathbb{Z}/p\mathbb{Z}$ (where $p \ge 5$ is prime).

The core question is whether, conditional on the neuron being initially active, the normalized parameter vector $(u, v, w)$ must converge to a direction supported on a single nonzero real Fourier class (i.e., a specific frequency $\zeta$ and its conjugate $-\zeta$). Previous theoretical work on two-layer networks suggested that individual neurons tend to select single-frequency Fourier features, with frequencies competing based on initial spectral magnitudes. This paper seeks to determine if this "single-frequency alignment" is a necessary consequence of training dynamics for the simplest possible model: a single neuron.

Methodology
The author employs a rigorous analytical approach using Clarke differential inclusions to model continuous-time training dynamics, acknowledging the non-differentiability of the ReLU activation at zero. The methodology involves:

1. Construction of Specific Initial Conditions: The author defines a set of initial parameters where exactly one input gate (corresponding to inputs $(0,0)$) is active, while all other gates are strictly inactive (negative pre-activation).
2. Dynamical System Analysis: They analyze the gradient flow of the cross-entropy loss. By exploiting symmetry, they reduce the high-dimensional system to a low-dimensional ordinary differential equation (ODE) governing the active gate's preactivation and the output weights.
3. Invariant Analysis: The author identifies a conserved quantity (or a quantity with a specific monotonic behavior) relating the output weight differences and the input preactivation. This allows them to prove that the active gate's preactivation decreases linearly to zero in finite time.
4. Clarke Trajectory Selection: Upon the preactivation reaching zero, the author utilizes the Clarke generalized derivative framework. They demonstrate that by selecting the derivative of ReLU at zero to be 0 (a standard convention in automatic differentiation), the trajectory enters a "dead" state where all gradients vanish, and the parameters freeze.
5. Fourier Spectral Analysis: The author computes the Fourier energy distribution of the frozen parameter vectors. They show that the limiting direction is a scalar multiple of a centered point mass vector, which distributes its energy equally across all nonzero real Fourier classes.

Key Contributions and Results

• Refutation of Single-Frequency Alignment: The paper provides a negative solution to the problem posed in [Cla26d]. It constructs a counterexample where the normalized parameters converge to a limit that is not supported on a single Fourier class. Instead, the Fourier energy is distributed equally among all $p-1$ nonzero frequency classes. For $p \ge 5$, the fraction of energy in any single class is $2/(p-1)$, which is strictly less than 1.
• Finite-Time Death and Freezing: The counterexample demonstrates that a single neuron can reach a completely inactive state ("dying ReLU") in finite time. Once dead, the parameters freeze at a limit determined by the initial conditions and the dynamics leading up to the death.
• Robustness and Probability: The counterexample is not an isolated artifact.
  • It holds on an open set of initial conditions, implying it occurs with positive probability under standard Gaussian initialization.
  • Appendix A strengthens these results, showing that:
    • Failure occurs for every Clarke trajectory from an open set of initial conditions (removing the need for specific trajectory selection).
    • The phenomenon persists under the convention $\text{ReLU}'(0)=0$.
    • It holds for smooth dead-zone approximations of ReLU ($C^\infty$ functions that are zero for $z \le 0$).
    • It occurs in fixed-step full-batch gradient descent, where the neuron dies in a single step.
• Persistent Active Counterexample: The appendix also presents a "persistent sliding" trajectory where the neuron remains active (one gate stays positive) but the normalized weights still converge to a multi-frequency limit, failing single-frequency alignment even without the neuron dying.

Significance and Claims
The paper claims that single-frequency, or monosemantic, Fourier structure is not a general consequence of training a single neuron on modular addition.

• Mechanism of Failure: The failure arises from two familiar mechanisms: the neuron dying (entering a dead region) or the neuron remaining active while effectively "memorizing" a single table entry (the $(0,0)$ case) rather than learning the general algorithm. The persistent trajectory realizes the "memorization" side of the memorization-to-generalization dichotomy associated with grokking in the smallest possible model.
• Limitations: The author explicitly notes that their results rely on the existence of an exact dead region in the activation function. The question remains open for everywhere-positive activations (such as softplus), where the "dying" mechanism cannot occur.
• Context: The work clarifies that while Fourier structure is present in the solution space, the training dynamics do not force the selection of a single frequency. The results distinguish the behavior of single-neuron models from the multi-frequency selection mechanisms observed in wider networks (e.g., width $m \ge 4(p-1)$ as in [MEO+24]).

In summary, the paper establishes that the hypothesis of inevitable single-frequency alignment in single-neuron modular addition is false, providing a concrete counterexample where the limiting direction is maximally non-pure in terms of Fourier classes.