arXiv · 2506.23679
Learning Modular Exponentiation with Transformers
Abstract
Modular exponentiation is crucial to number theory and cryptography, yet remains largely unexplored from a mechanistic interpretability standpoint. We train a 4-layer encoder-decoder Transformer model to perform this operation and investigate the emergence of numerical reasoning during training. Utilizing principled sampling strategies, PCA-based embedding analysis, and activation patching, we examine how number-theoretic properties are encoded within the model. We find that reciprocal operand training leads to strong performance gains, with sudden generalization across related moduli. These synchronized accuracy surges reflect grokking-like dynamics, suggesting the model internalizes shared arithmetic structure. We also find a subgraph consisting entirely of attention heads in the final layer sufficient to achieve full performance on the task of regular exponentiation. These results suggest that transformer models learn modular arithmetic through specialized computational circuits, paving the way for more interpretable and efficient neural approaches to modular exponentiation.
Explore related subjects
Keep this discovery
David Demitri Africa, Sara M. Kapoor, Theo Simon Sorg, Challenger Mishra. 2025-06-30. Learning Modular Exponentiation with Transformers. https://arxiv.org/abs/2506.23679
Cite the original work for its findings. Save a collection to share your selection of sources.