The Factorization Was the Circuit: Chinese Remainder Theorem in Transformers

Emergent algorithms can arise in neural networks trained on simple tasks. Twenty small transformers (1-layer, no LayerNorm) trained to solve $a^b\mathrm{mod}p$ over a range of primes $p$ exhibit perfect generalization after grokking. We reverse engineer the emergent mechanism and find that the model first rewrites the answer as $g^{\gamma }$, and $\gamma ={\log }_g(a)\cdot b\mathrm{mod}q$, where $q=p-1$ and $g$ is a primitive root of $p$. It then splits that product using the $\text{Chinese Remainder Theorem}$ (CRT): computing $\gamma \mathrm{mod}{m}_i$ on independent streams, where $m_i$ are the prime-power factors of $q$; then recombining the residues in the unembedding. Similar to prior work, we learn that the model uses a sparse set of "key" frequencies, which index the irreducible representations. Unlike prior work, this set is completely deterministic, as it is the union of disjoint bands of the Fourier spectrum at the multiples of $q/m_i$. We call these bands the CRT bands. Through causal tests, we show that each band controls its own corresponding stream independently. Furthermore, for prime-power streams $m_i={\ell }^e$ (e.g., $\mathbb{Z}/16\mathbb{Z}$), the irreducible representations of these groups are filtered by order, dual to the quotient tower: $\mathbb{Z}/{\ell }^e\mathbb{Z}\twoheadrightarrow \cdots \twoheadrightarrow \mathbb{Z}/\ell \mathbb{Z}$. With the a priori knowledge of these irreducible representations and their frequencies, we are able to show the streams' independence causally, as well as isolate and observe the mechanism's circuitry formation during grokking.