Prime Fourier Embeddings: A Principled Basis for Modular Arithmetic

Standard neural embeddings often struggle with modular arithmetic, frequently requiring long training periods (grokking) to discover the underlying algebraic structure. This paper investigates whether providing a pre-structured representation—specifically one that aligns with the Chinese Remainder Theorem (CRT)—can eliminate the need for this discovery process and improve generalization.

The authors introduce Prime Fourier Embeddings (PFE), which represent integers as a collection of (cos, sin) pairs indexed by prime numbers. This construction is derived from the harmonic analysis of the rational numbers, specifically the adelic character factorization, which ensures that each prime channel acts independently. The authors prove that any linear map equivariant to the product group action on these embeddings must be block-diagonal, meaning there is no cross-prime or cross-depth coupling in the learned weights.

Empirical ablation studies confirm that PFE successfully isolates task-relevant information. When training on modular arithmetic tasks, the model exhibits high specialization, with task-relevant prime channels showing significantly higher importance than irrelevant ones (specialization ratios exceeding 500x). The model achieves perfect in-distribution test accuracy across all tested square-free composite moduli, validating that the PFE structure provides a complete basis for the task.

By transforming modular arithmetic from a problem of representational discovery into one of channel selection, PFE provides a principled alternative to heuristic embeddings. This approach demonstrates how incorporating mathematical symmetries—in this case, the prime-local structure of modular arithmetic—directly into the architecture can lead to more robust and interpretable models that avoid the fragility associated with learning algebraic structure via gradient descent.