Learning Modular Exponentiation with Transformers

TL;DR

This paper explores how Transformer models learn modular exponentiation, revealing emergent numerical reasoning, specialized circuits, and generalization behaviors that enhance interpretability and efficiency in cryptographic computations.

Contribution

It demonstrates that Transformers can internalize modular arithmetic through specialized circuits and exhibits emergent reasoning capabilities, advancing interpretability in neural models.

Findings

Reciprocal operand training improves performance

Sudden generalization across related moduli observed

A subgraph of attention heads suffices for full task performance

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.