Prime by Prime: A Reference for Emergent Modularity

Neural networks can break a task into subtasks and learn to solve each independently, in a different part of the model. But what determines whether this modular organization emerges during learning remains poorly understood. One obstacle is that the relevant subtasks and their internal representations are often unknown. We address this by studying transformers trained to compute $(a+b)\mathrm{mod}n$ and $(a\times b)\mathrm{mod}n$, where $n=p_1\cdots p_k$ is a product of distinct primes. These tasks admit a simple decomposition: represent the inputs through their remainders $a\mathrm{mod}{p}_i$ and $b\mathrm{mod}{p}_i$, perform arithmetic separately for each prime, and reconstruct the unique answer modulo $n$ consistent with the results, as guaranteed by the Chinese remainder theorem. We ask whether transformers discover this decomposition through training. We find that models trained jointly on addition and multiplication (i.e., learning the ring ${\mathbf{Z}}_n$) consistently do so, with evidence at three levels: (i) errors occur approximately independently across primes, (ii) each prime’s computation occupies its own representational subspace, and (iii) these computations largely rely on distinct groups of neurons. This modular organization critically depends on the training task: addition alone produces entangled or partially modular solutions that vary across seeds, including fused subtasks and asymmetric dependencies between primes. Why training on the ring induces this modular solution, while addition alone does not, remains an open theoretical question and an interesting direction for future work.