Prime Fourier Embeddings: A Principled Basis for Modular Arithmetic

ICML 2026 Workshop - 3rd AI for Math Workshop: Website

Abstract

Numbers have algebraic structure that standard neural embeddings often fail to expose. We introduce Prime Fourier Embeddings (PFE), which encode integers as prime-indexed (cos, sin) pairs derived from the harmonic analysis of Q, providing a pre-structured representation in which modular arithmetic reduces to selecting the relevant prime channel rather than discovering algebraic structure from scratch. We prove that any linear map equivariant with respect to the product group action on PFE must be block-diagonal with one independent block per prime — a consequence of Schur’s lemma applied to the resulting character decomposition. For square-free composite moduli, the Chinese Remainder Theorem predicts which prime channels are task-relevant. Both predictions are confirmed empirically: ablation studies show specialization ratios exceeding 500× between task-relevant and task-irrelevant channels, with perfect in-distribution test accuracy across all square-free composite moduli tested.

Hyunsang Hwang
Hyunsang Hwang
Ph.D. Student

Nice to meet you. And you?

Suhyun Bae
Suhyun Bae
M.S. Student

The cake is not a lie!

Donghun Lee
Donghun Lee
Associate Professor

Connecting artificial intelligence and mathematics, in both directions.