Technical Summary: Prime Fourier Embeddings

Problem Statement

Standard neural embeddings often fail to expose the inherent algebraic structure of integers, particularly regarding modular arithmetic. While the Chinese Remainder Theorem (CRT) dictates that addition modulo a composite number $N$ decomposes into independent additions modulo its prime factors, standard learned embeddings do not naturally reflect this decomposition. Consequently, models trained on modular arithmetic tasks often suffer from "grokking"—a phenomenon where models endure prolonged memorization phases before suddenly generalizing. This delay is attributed to a representational failure: when the input lacks explicit prime-local structure, the network must reconstruct these algebraic relationships from scratch via gradient descent, a process that is data-hungry and fragile.

Existing approaches have limitations: continuous encodings (e.g., xVal) preserve magnitude but discard periodicity, while periodic approaches like Fourier Number Embeddings (FoNE) often use composite bases (e.g., base-10) that entangle distinct prime structures (e.g., 2-adic and 5-adic) within shared channels, failing to isolate prime-local information.

Methodology: Prime Fourier Embeddings (PFE)

The authors introduce Prime Fourier Embeddings (PFE), a principled embedding scheme derived from the harmonic analysis of the rational numbers $\mathbb{Q}$.

• Construction: PFE encodes an integer $a$ as a collection of prime-indexed $(\cos, \sin)$ pairs. For a prime $p$ and depth $d$ (where $0 \le d < D$), the embedding is defined as:
  $$\text{PFE}_{p,d}(a) = \left[ \cos\left(\frac{2\pi a}{p^{d+1}}\right), \sin\left(\frac{2\pi a}{p^{d+1}}\right) \right] \in \mathbb{R}^2$$
  This corresponds to the real decomposition of the character $\chi_1$ of the group $\mathbb{Z}/p^{d+1}\mathbb{Z}$.
• Theoretical Basis: The design is motivated by Pontryagin duality and the factorization of adelic characters. The embedding implements local $p$-adic characters directly, ensuring that the residue structure of any integer is readable from the corresponding prime channel without cross-prime coupling.
• Implementation: In experiments, the authors use a fixed set of 16 primes ($P_0 = \{3, \dots, 59\}$) and a depth cap of $D=3$, resulting in an embedding dimension of 96. The embedding features are deterministic and non-trainable.

Key Theoretical Contributions

The paper establishes a rigorous structural guarantee for any linear map operating on PFE:

1. Block-Diagonal Decomposition Theorem (Theorem 3.1): Any linear map $W$ that is equivariant with respect to the product group action on PFE must be block-diagonal.
   • There is one independent block per prime $p$ and depth $d$.
   • There is zero coupling between blocks of distinct primes.
   • There is zero coupling between blocks of distinct depths within the same prime.
2. Proof Mechanism: This result follows from Schur's Lemma. Each PFE block represents a non-isomorphic irreducible representation (a character) of the additive group $\mathbb{Z}$. Since the characters for distinct $(p, d)$ pairs are non-isomorphic, any equivariant intertwiner between them must be zero.
3. CRT Prediction: For a task with modulus $N = p_1 \cdots p_k$, the theory predicts that only the $k$ prime channels corresponding to the factors of $N$ will carry task-relevant information, while all other channels remain irrelevant.

Experimental Results

The authors validate these predictions through systematic ablation studies on modular addition tasks $(a+b) \mod N$.

• Experiment 1 (Single-Prime Tasks): Models were trained on tasks with varying prime moduli. When the specific prime row corresponding to the task modulus was ablated (zeroed), test accuracy dropped significantly. In contrast, ablating irrelevant prime rows caused near-zero accuracy drops.
  • Result: Specialization ratios (diagonal drop vs. off-diagonal drop) exceeded 500$\times$ across most configurations, confirming that the network utilizes only the relevant prime channel.
• Experiment 2 (Composite Moduli): Models were trained on square-free composite moduli (e.g., $N=15, 21, 105$).
  • Result: Ablating prime factors of $N$ caused large accuracy drops, while ablating non-factors caused negligible drops. The factor/nonfactor ratio again exceeded 500$\times$.
  • Accuracy: The models achieved 100% in-distribution test accuracy across all tested square-free composite moduli and input ranges (provided the input range was sufficiently large relative to $N$).

Significance and Claims

The paper claims that PFE transforms the learning of modular arithmetic from a representational discovery problem into a selection problem.

• Principled vs. Arbitrary: Unlike heuristic frequency choices, the prime basis is mathematically grounded in the adelic Fourier decomposition of $\mathbb{Q}$. This ensures the necessary algebraic structure is present from initialization.
• Mechanistic Interpretability: The work demonstrates that embedding design can be guided by the mathematics of the objects being represented. The internal structure of the network is characterized before training (via representation theory) and verified after training (via targeted ablations).
• Falsifiable Theory: The Block-Diagonal Decomposition Theorem provides a provable structural constraint on network behavior. The empirical confirmation of this constraint offers a path toward more reliable and transparent numerical reasoning in neural networks, moving beyond post-hoc observation to mathematically grounded theory.

The authors note that while the block-diagonal structure is proven, the specific mechanism by which gradient descent converges to an equivariant solution (rather than breaking symmetry) remains an open question for future work.