Learning Modular Exponentiation with Transformers

• URL: http://arxiv.org/abs/2506.23679v1
• Date: Mon, 30 Jun 2025 10:00:44 GMT
• Title: Learning Modular Exponentiation with Transformers
• Authors: David Demitri Africa, Sara M. Kapoor, Theo Simon Sorg,
• Abstract summary: We train a 4-layer encoder-decoder Transformer model to perform modular exponentiation.<n>We find that reciprocal training leads to strong performance gains, with sudden generalization across related moduli.<n>These results suggest that transformer models learn modular arithmetic through specialized computational circuits.
• Score: 0.0
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Neural ODE Transformers: Analyzing Internal Dynamics and Adaptive Fine-tuning [30.781578037476347]
  We introduce a novel approach to modeling transformer architectures using highly flexible non-autonomous neural ordinary differential equations (ODEs)<n>Our proposed model parameterizes all weights of attention and feed-forward blocks through neural networks, expressing these weights as functions of a continuous layer index.<n>Our neural ODE transformer demonstrates performance comparable to or better than vanilla transformers across various configurations and datasets.
  arXiv  Detail & Related papers  (2025-03-03T09:12:14Z)
• Interpreting Affine Recurrence Learning in GPT-style Transformers [54.01174470722201]
  In-context learning allows GPT-style transformers to generalize during inference without modifying their weights. This paper focuses specifically on their ability to learn and predict affine recurrences as an ICL task. We analyze the model's internal operations using both empirical and theoretical approaches.
  arXiv  Detail & Related papers  (2024-10-22T21:30:01Z)
• Algorithmic Capabilities of Random Transformers [49.73113518329544]
  We investigate what functions can be learned by randomly transformers in which only the embedding layers are optimized. We find that these random transformers can perform a wide range of meaningful algorithmic tasks. Our results indicate that some algorithmic capabilities are present in transformers even before these models are trained.
  arXiv  Detail & Related papers  (2024-10-06T06:04:23Z)
• Principled Understanding of Generalization for Generative Transformer Models in Arithmetic Reasoning Tasks [5.522116934552708]
  Transformer-based models excel in various tasks but their generalization capabilities, especially in arithmetic reasoning, remain incompletely understood.<n>This paper develops a unified theoretical framework for understanding the generalization behaviors of transformers in arithmetic tasks.
  arXiv  Detail & Related papers  (2024-07-25T11:35:22Z)
• Learning on Transformers is Provable Low-Rank and Sparse: A One-layer Analysis [63.66763657191476]
  We show that efficient numerical training and inference algorithms as low-rank computation have impressive performance for learning Transformer-based adaption. We analyze how magnitude-based models affect generalization while improving adaption. We conclude that proper magnitude-based has a slight on the testing performance.
  arXiv  Detail & Related papers  (2024-06-24T23:00:58Z)
• Towards Empirical Interpretation of Internal Circuits and Properties in Grokked Transformers on Modular Polynomials [29.09237503747052]
  Grokking on modular addition has been known to implement Fourier representation and its calculation circuits with trigonometric identities in Transformers.<n>We show that the transferability among the models grokked with each operation can be only limited to specific combinations.<n>Some multi-task mixtures may lead to co-grokking, where grokking simultaneously happens for all the tasks.
  arXiv  Detail & Related papers  (2024-02-26T16:48:12Z)
• Understanding Addition in Transformers [2.07180164747172]
  This paper provides a comprehensive analysis of a one-layer Transformer model trained to perform n-digit integer addition. Our findings suggest that the model dissects the task into parallel streams dedicated to individual digits, employing varied algorithms tailored to different positions within the digits.
  arXiv  Detail & Related papers  (2023-10-19T19:34:42Z)
• Uncovering mesa-optimization algorithms in Transformers [61.06055590704677]
  Some autoregressive models can learn as an input sequence is processed, without undergoing any parameter changes, and without being explicitly trained to do so. We show that standard next-token prediction error minimization gives rise to a subsidiary learning algorithm that adjusts the model as new inputs are revealed. Our findings explain in-context learning as a product of autoregressive loss minimization and inform the design of new optimization-based Transformer layers.
  arXiv  Detail & Related papers  (2023-09-11T22:42:50Z)
• End-to-End Meta-Bayesian Optimisation with Transformer Neural Processes [52.818579746354665]
  This paper proposes the first end-to-end differentiable meta-BO framework that generalises neural processes to learn acquisition functions via transformer architectures. We enable this end-to-end framework with reinforcement learning (RL) to tackle the lack of labelled acquisition data.
  arXiv  Detail & Related papers  (2023-05-25T10:58:46Z)
• On the Generalization and Adaption Performance of Causal Models [99.64022680811281]
  Differentiable causal discovery has proposed to factorize the data generating process into a set of modules. We study the generalization and adaption performance of such modular neural causal models. Our analysis shows that the modular neural causal models outperform other models on both zero and few-shot adaptation in low data regimes.
  arXiv  Detail & Related papers  (2022-06-09T17:12:32Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.