Modular Addition

Latest papers

12 papers2 papers in the last four weeks, against 2 the four weeks before. 0.0% of all new papers.

Oct 30, 2026

Nonconformal Distortion and Local Orientation Reversal in the Fourier Circuit of a Modular-Addition Transformer

We study the local differential geometry of a Fourier-based modular-addition circuit in a one-layer transformer whose algorithm has previously been reverse-engineered. For each of the five key Fourier frequencies, we restrict the model's multilayer perceptron (MLP) contribution to a corresponding Fourier input plane and readout, and measure its directional Beltrami coefficient under continuous input perturbations. Whereas nondegenerate restrictions of ideal complex multiplication are conformal, the trained maps exhibit nonconformal distortion (median

|μ|=0.304

0.446

), spatial organization in the Beltrami argument, and regions of local orientation reversal. The median across frequencies decreases from

0.939

at initialization to

0.390

after training. The key-frequency planes also show lower median distortion than most random-plane controls, although their distributions overlap and the coordinate frames differ in conditioning. Across 500 checkpoints in one training run, most of the aggregate reduction in Beltrami magnitude occurs before the principal increase in the measured Fourier-circuit formation. At the individual-input level, Beltrami magnitude shows negligible association with the correct-logit margin (

|ρ|\leq0.026

). Late-checkpoint measurements across six runs show comparable distortion for most selected frequencies. These results provide a quantitative characterization of local deformation in a frequency-projected component of an identified neural circuit, complementary to its algorithmic reconstruction.

Eduardo Fuentes, Gonzalo Riera, Héctor Allende-Cid

AI Mechanisms Lab, Chile · WriteWise, Chile · Facultad de Matemáticas, Pontificia Universidad Católica de Chile, Santiago, Chile +2

Neural Representation Geometry Mechanistic Interpretability Modular Addition

Ask PDF

Oct 8, 2026

Probability-Signature Dynamics: Unpacking Modular Addition Learning Within Two-Layer Networks

Neural networks trained on modular addition tasks often develop Fourier-structured representations that support exact generalization. While prior work has identified these Fourier circuits, the mechanism by which gradient-based training selects them from the data distribution remains unclear. We address this question using probability signatures, which express leading gradient interactions through conditional statistics of the training distribution. For modular addition, these signatures are cyclic shift operators and are diagonalized by the discrete Fourier transform, yielding approximately decoupled Fourier-mode dynamics. This explains the emergence of Fourier sparsity, frequency matching, and phase alignment. The same framework resolves a puzzle under label noise: corrupted examples can show faster early loss decrease than clean examples, despite lacking a coherent generalization rule. We show that noise increases conditional label collisions, strengthening early shared-coordinate reinforcement. Finally, this method can be applied to other operators. Taking XOR as an example, we observed the predicted frequency in experiments.

Yunji Wang, Junjie Yao, Linyu Liu +2

School of Mathematical Sciences, Shanghai Jiao Tong University · Huawei Technologies Ltd · Shanghai University of Finance and Economics +1

Neural Network Training Dynamics Frequency-Domain Feature Learning Modular Addition

Ask PDF

Sep 22, 2026

A Spectral Theory of Grokking: Weight Decay induces Feature Learning

In grokking an early fit to the training data separates from a much later improvement in generalization. During this delay, training can move from a fixed neural tangent kernel (NTK) regime to one in which task-relevant kernel eigendirections continue to evolve. We provide a quantitative theory for how this transition from lazy to rich learning can produce delayed generalization. For homogeneous networks trained with squared loss and

L_2

weight decay, we show that a finite residual remains after memorization, with larger residual fractions in target components associated with smaller NTK eigenvalues. These residuals feed back into the dynamics of the NTK itself, and projecting the resulting dynamics onto task-relevant spectral directions yields a reduced system in which residual-driven kernel growth competes with weight decay. This system predicts that the grokking timescale is controlled by the product of learning rate and weight decay, that feature learning slows logarithmically near a critical decay above which task-aligned NTK structure can no longer support generalization, and that stronger decay can prevent fitting altogether. We test these predictions in modular addition. In a homogeneous MLP, task-aligned Fourier structure continues to emerge in the NTK after training accuracy has saturated, and an 84

\times

90-grid of trained networks across varying learning rate and weight decay recovers the predicted phase geometry and inverse-product scaling of the generalization time with learning rate and weight decay. A one-block Transformer shows similar macroscopic phase structure in a 42

\times

45-grid, as well as the same transition-time scaling despite violating exact homogeneity. Together, these results provide a mechanistic derivation connecting post-fit feature learning to both the onset of generalization and its phase structure in the learning rate and weight decay plane.

Lenz Pracher, Pascal de Jong, Oskar Lieshaus +2

Arnold-Sommerfeld-Center for Theoretical Physics, Ludwig-Maximilians-Universität München, Munich, Germany · Department of Applied Physics, Standford University, Standford, CA, USA · Department of Mathematics, University of Cambridge, Cambridge, UK

Neural Network Generalization Weight Decay Modular Addition Neural Tangent Kernel Grokking

Ask PDF

Sep 15, 2026

Symmetry without a manifold: intrinsic dimension on orbits

The standard geometric derivation of neural scaling exponents takes the intrinsic dimension of a data manifold as its input. On modular addition in

\mathbb{Z}_p

that derivation has no input. The exact algebraic solution is an orbit of

\mathbb{Z}_p

acting by isometries. Transitivity alone makes the ratio statistic underlying the standard dimension estimator a point mass, so the estimator is undefined, and here the two nearest neighbour distances coincide exactly. Breaking the symmetry at scale

ε

returns a number, but one that tracks

1/ε

with no scale free plateau. We show that the failure is general, since on any finite orbit of a group acting by isometries the estimator reports the resolution at which the set is probed rather than a dimension. What replaces the power law is exponential in hidden width,

L(h)=L_\infty+A\exp(-c\,h^α)

, with

R^2

between 0.982 and 0.995 against 0.857 to 0.906 for a power law admitting the same floor and fitted under the same protocol. Where the data supply is sufficient the rate belongs to the regulariser rather than to the group, since weight decay moves

c

by a factor of 47 while group order moves it by 1.10, a residual below seed to seed resolution, for every fixed

α

between 0.75 and 2. The critical width falls with group order rather than rising, against capacity counting that assigns a fixed number of neurons to each irreducible representation.

Chon-Fai Kam, Miloud Bessafi, Frédéric Cadet

Dipartimento di Fisica e Chimica, Universit`a degli Studi di Palermo, via Archirafi 36, I-90123 Palermo, Italy · EnergyLab, Universit´e de La R´eunion, F-97715 Saint-Denis, France · Universit´e Paris Cit´e and Universit´e de La R´eunion, BIGR, INSERM UMR S1134, F-75014 Paris, France +1

Intrinsic Dimensionality Modular Addition Neural Scaling Laws

Ask PDF

Aug 5, 2026

A Counterexample to Fourier Alignment in Single-Neuron Modular Addition

We give a negative solution to MAIS-O60. We first construct an example in which an initially active ReLU neuron becomes completely inactive in finite time and thereafter remains frozen at a limit whose Fourier energy is equally distributed among all nonzero real frequency classes. The counterexample holds on an open set of initial conditions and therefore occurs with positive probability under Gaussian initialization. An appendix prepared by GPT-5.6 Sol strengthens the counterexample by showing that the same failure can occur for every Clarke trajectory from an open set of initial conditions, under the convention

\mathrm{ReLU}'(0)=0

, for smooth dead-zone approximations of ReLU, and for fixed-step full-batch gradient descent. Thus, single-frequency alignment is not a general consequence of training a single neuron on modular addition.

Gautam Neelakantan Memana

Shallow Neural Networks ReLU Neural Networks Frequency-Domain Feature Learning Modular Addition Neural Network Training Dynamics

Ask PDF

Jun 8, 2026

Beyond Neural Collapse: Task-Intrinsic Geometry Governs Neural Representations in Modular Arithmetic

While neural collapse (NC) predicts that a

K

-class-balanced classifier should organize terminal representations as a

(K-1)

-dimensional simplex equiangular tight frame (ETF), modular addition consistently enters a different regime: networks compress to a two-dimensional cyclic geometry in which both classifier weights and token embeddings lie on circles. We refine the explanation of this phenomenon in three directions. First, we formalize a layerwise non-uniform training mechanism: downstream classifier weights are driven by dense cross-entropy gradients into a rank-2 equiangular configuration before upstream embeddings fully reorganize, and once this classifier plane forms, backpropagated feature gradients constrain embedding motion to the same plane while weight decay suppresses orthogonal components. Second, after this subspace locking, the induced in-plane dynamics admit an entropy-regularized transport interpretation on

S^1

; combined with modular-addition labels, this reduces embedding formation to phase alignment, whose minimizers are single-frequency characters of

\mathbb{Z}/P\mathbb{Z}

and hence equal-angle points on a circle. Third, we quantify why this solution prevails over NC: a simplex ETF gains only an

O(1)

advantage in cross-entropy, whereas the cyclic rank-2 solution enjoys a

Θ(K)

advantage under Schatten or weight-decay surrogates, yielding a critical threshold

λ_{\mathrm{crit}} = Θ(1/K)

. Our results explain both why classifier weights move first and why embeddings subsequently align with them, showing that grokking on modular arithmetic is governed not by maximal separation alone but by a task-structured trade-off between separation, symmetry, and complexity.

Hu Tan, Kuo Gai, Shihua Zhang

Academy of Mathematics and Systems Science, Chinese Academy of Sciences, Beijing 100190, China; and School of Mathematical Sciences, University of Chinese Academy of Sciences, Beijing 100049, China · Shanghai Institute for Mathematics and Interdisciplinary Sciences (SIMIS), Shanghai, China · Academy of Mathematics and Systems Science, Chinese Academy of Sciences, Beijing 100190, China; School of Mathematical Sciences, University of Chinese Academy of Sciences, Beijing 100049, China; and Key Laboratory of Systems Health Science of Zhejiang Province, School of Life Science, Hangzhou Institute for Advanced Study, University of Chinese Academy of Sciences, Chinese Academy of Sciences, Hangzhou 310024, China

Neural Representation Geometry Classification Modular Addition Neural Collapse Neural Network Training Dynamics Modular Arithmetic+1

Ask PDF

May 31, 2026

BRo-JEPA: Learning Modular Arithmetic in Latent Space

Can neural networks learn abstract algebraic rules, or do they merely memorize training patterns? We investigate this using MNIST digits as states and modular arithmetic operations as actions in a JEPA-style latent world model. Standard supervised baselines and JEPA models with additive operation embeddings fit seen operations but fail to extrapolate reliably to unseen ones. To bridge this gap, we introduce a block-rotation predictor that imposes the circular structure of modulo-10 arithmetic in latent space. This enables strong zero-shot generalization, with the best ResNet-based JEPA block-rotation model achieving 99.46% zero-shot and 99.46% rollout accuracy. Our results suggest that latent world models can learn symbolic transformation rules when architecture matches the structure of the problem. Our code can be accessed here.

Divyansh Jha, Yuanfang Xie, Varan Mehra +1

1Georgia Institute of Technology · 2NYU Langone Health

Zero-Shot Learning Joint-Embedding Predictive Architecture Latent World Models Modular Arithmetic Algorithmic Generalization Modular Addition

Ask PDF

May 29, 2026

Assign and Add: A Mechanistic Study of Compositional Arithmetic

Large language models are able to compose skills in order to perform complex tasks, many of which might not have been seen during training. The details of how exactly this composition occurs remain elusive. In this paper, we study a mechanism for compositional generalization in transformers by considering a simple controlled setting involving variable assignment and modular addition. By partitioning our training data into disjoint sets, we observe that small transformers are able to generalize to previously unseen combinations of variables and numbers. Our mechanistic analysis shows that the same ``modular addition'' MLP module is used whether the inputs are given directly or indirectly through a separate variable assignment mechanism. We also analyze the training dynamics from an empirical lens, which reveals three phases of learning: first, modular addition is learned, then the structure required for variable assignment, and finally a refinement phase where the model generalizes to some hard sequences not seen in training. Finally, we provide a theoretical framework to explain how compositionality emerges from training dynamics. These results suggest that compositional generalization can be a natural consequence of the compositionality of internal mechanisms in~transformers.

Brady Exoo, Alberto Bietti, John Sous

Yale University · Flatiron Institute

Transformer Interpretability Transformer Modular Addition Compositional Generalization

Ask PDF

May 8, 2026

Learning Large-Scale Modular Addition with an Auxiliary Modulus

Learning parity functions, more general modular addition, is a challenging machine learning task due to its input sensitivity. A recent study substantially scaled modular addition learning in both the number of summands and the modulus. Its key idea is to increase zeros in training sequences, reducing the effective number of summands and thus controlling training difficulty; however, this induces covariate shift between training and test input distributions. This study theoretically and empirically analyzes this side effect and proposes a covariate-shift-free method for modular addition. Specifically, we introduce an auxiliary modulus

Kq

during training, which reduces wrap-around frequency and problem difficulty while preserving the same input distribution across training and testing. Experiments show strong scalability and sample efficiency: even for large input length

N

, large modulus

q

, and small datasets -- where the sparse method fails to learn -- our method achieves equal or better match accuracy and relaxed

τ

-accuracy. For example, at

N=64

and

q=974269

, our method trained on 100K samples achieves

97.0\%

τ

-accuracy at

τ=0.05

, while the sparse method achieves only

9.5\%

with the same data size and

93.9\%

even when extended to 1M samples.

Hanato Kikuchi, Ryosuke Masuya, Kazuhiko Kawamoto +1

Chiba University · National Institute of Informatics

Boolean Function Learning Distribution Shift Covariate Shift Modular Arithmetic Data-Efficient Learning Modular Addition

Ask PDF

Apr 20, 2026

A Mechanistic Study of Algorithmic Generalization in Diffusion Models on Modular Arithmetic

Despite their empirical success, how diffusion models generalize remains poorly understood from a mechanistic perspective. In this work, we provide a controlled mechanistic account of algorithmic learning in diffusion models. By utilizing the grokking phenomenon on an image-based modular addition task as a tractable testbed, we dissect the model's internal dynamics across two distinct visual regimes. In a single-image regime, mechanistic dissection reveals that the model implements modular addition by encoding operands as periodic representations and composing them at the self-attention--feedforward interface in a way that recovers the trigonometric addition identities. In a diverse-image regime with high intraclass variability, we find that the iterative sampling process partitions the task into an arithmetic computation phase and a subsequent visual refinement phase, with their transition precisely captured by an internal frequency-entropy signal. Together, our analysis clarifies the mechanics of algorithmic generalization of diffusion models in a controlled setting, while leaving open how broadly these mechanisms transfer beyond modular addition.

Joon Hyeok Kim, Yong-Hyun Park, Mattis Dalsætra Østby +1

Department of Computer and Information Science, University of Pennsylvania, Philadelphia, PA 19104, USA

Pixel-Space Diffusion Models Mechanistic Interpretability Modular Addition Grokking Diffusion Models

Ask PDF

Feb 18, 2026

On the Mechanism and Dynamics of Modular Addition: Fourier Features, Lottery Ticket, and Grokking

We present a comprehensive analysis of how two-layer neural networks learn features to solve the modular addition task. Our work provides a full mechanistic interpretation of the learned model and a theoretical explanation of its training dynamics. While prior work has identified that individual neurons learn single-frequency Fourier features and phase alignment, it does not fully explain how these features combine into a global solution. We bridge this gap by formalizing a diversification condition that emerges during training when overparametrized, consisting of two parts: phase symmetry and frequency diversification. We prove that these properties allow the network to collectively approximate a flawed indicator function on the correct logic for the modular addition task. While individual neurons produce noisy signals, the phase symmetry enables a majority-voting scheme that cancels out noise, allowing the network to robustly identify the correct sum. Furthermore, we explain the emergence of these features under random initialization via a lottery ticket mechanism. Our gradient flow analysis proves that frequencies compete within each neuron, with the "winner" determined by its initial spectral magnitude and phase alignment. From a technical standpoint, we provide a rigorous characterization of the layer-wise phase coupling dynamics and formalize the competitive landscape using the ODE comparison lemma. Finally, we use these insights to demystify grokking, characterizing it as a three-stage process involving memorization followed by two generalization phases, driven by the competition between loss minimization and weight decay.

Jianliang He, Leda Wang, Siyu Chen +1

Department of Statistics and Data Science, Yale University

Lottery Ticket Hypothesis Modular Addition Neural Network Training Dynamics Frequency-Domain Feature Learning Neural Network Generalization Grokking

Ask PDF

Oct 31, 2024

A Mechanistic Study of Transformers Training Dynamics

Large-scale pretraining of transformers has been central to the success of foundation models. However, the scale of those models limits our understanding of the mechanisms at play during optimization. In this work, we study the training dynamics of transformers in a controlled and interpretable setting. On the sparse modular addition task, we demonstrate that specialized attention circuits, called clustering heads, can be implemented during gradient descent to solve the problem. Our experiments show that such pathways naturally emerge during training. By monitoring the evolution of tokens via a visual sandbox, we uncover a two-stage learning and the occurrences of loss spikes due to the high curvature of normalization layers. Our findings provide several insights into patterns observed in more practical settings, such as the pretraining of large language models.

Ambroise Odonnat, Wassim Bouaziz, Vivien Cabannes

Inria, Univ. Rennes 2 · Mistral AI; work done while at Meta. · FAIR at Meta.

Transformer Interpretability Attention Mechanisms Modular Addition Neural Network Training Dynamics Attention Head Specialization Transformer Training Dynamics

Ask PDF