---
title: Modular-Addition Transformers
url: https://www.emergentmind.com/topics/modular-addition-transformers
type: topic
---

# Modular-Addition Transformers

Searching arXiv for recent papers on modular-addition transformers and related mechanistic interpretations.
to=arxiv_search  天天中彩票腾讯json_object  新天天彩票 unfortunately single_search queries not supported?
to=arxiv_search  天天中彩票充值ి{"query":"modular addition transformers mechanistic interpretability", "max_results": 10}
to=arxiv_search  天天中彩票投注_json  彩经彩票{"query":"Teaching Transformers Modular Arithmetic at Scale 2410.03569", "max_results": 5}
to=arxiv_search  天天中彩票被{"query":"Teaching Transformers Modular Arithmetic at Scale", "max_results": 5}
Modular-addition transformers are transformer architectures trained to compute addition in a finite cyclic group, typically either the two-input map $(a,b)\mapsto (a+b)\bmod p$ or the more general $N$-input map $(x_1,\dots,x_N)\mapsto \sum_{i=1}^N x_i \bmod q$. The topic has become a standard testbed for algorithmic generalization, grokking, and mechanistic interpretability because the target function is algebraically simple while the learned internal solutions are nontrivial. Recent work spans stylized one-layer transformers analyzed through Fourier circuits, mechanistic accounts of grokking, and scaled encoder-only transformers that reach regimes such as $N=256$ and $q=3329$, which are directly relevant to lattice-based cryptography [2402.09469] [2407.12332] [2410.03569].

## 1. Formal task definitions and problem regimes

The canonical modular-addition task fixes a prime modulus $p$ and inputs $a_1,\dots,a_k\in\mathbb Z_p$, with target
\[
S=(a_1+\cdots+a_k)\bmod p\in\mathbb Z_p.
\]
A standard dataset formulation is
\[
D_p=\{((a_1,\dots,a_k),y): y=(a_1+\cdots+a_k)\bmod p\},
\]
which is used to study both one-hidden-layer neural networks and one-layer transformers [2402.09469].

The grokking literature also distinguishes regression and classification variants. In the regression version, each example is a one-hot triple $x=(e_a,e_b,e_c)\in\{0,1\}^{3p}$ with target $y=p\cdot \mathbf 1[a+b\equiv c\pmod p]$ under squared loss. In the classification version, inputs are pairs $(e_a,e_b)\in\{0,1\}^{2p}$ and targets are one-hot vectors $e_{(a+b\bmod p)}\in\{0,1\}^p$ under cross-entropy [2407.12332].

A separate line of work reformulates the task at much larger scale as
\[
\mathbb Z_q^N\longrightarrow \mathbb Z_q,\qquad (x_1,\dots,x_N)\mapsto \sum_{i=1}^N x_i \bmod q,
\]
with $N$ up to $256$ and prime moduli $q$ up to $3329$. That regime is motivated by cryptographic settings in which prior work on machine learning had only handled $N\le 6$ and $q\le 1000$, whereas systems such as CRYSTALS-Kyber involve $N\simeq 256$ and $q\sim 3329$ [2410.03569].

## 2. Fourier structure and learned computational circuits

A central empirical and theoretical result is that modular-addition transformers organize their internal computation around Fourier features. In the one-layer theoretical construction, each hidden unit or attention head aligns with exactly one frequency $\zeta\in\{1,\dots,(p-1)/2\}$, and the final linear layer recombines the resulting modes through a learned inverse DFT to recover the one-hot encoding of the sum modulo $p$ [2402.09469].

In learned depth-1 transformer circuits, the embedding matrix is approximately organized into sinusoidal blocks. For a token $x$, the learned embedding can be analyzed as a concatenation of coordinates of the form $\bigl[\cos(\omega_k x),\sin(\omega_k x)\bigr]$, where $\omega_k=2\pi k/p$. This representation makes the addition law compatible with the trigonometric identities
\[
\cos(\alpha+\beta)=\cos\alpha\cos\beta-\sin\alpha\sin\beta,\qquad
\sin(\alpha+\beta)=\sin\alpha\cos\beta+\cos\alpha\sin\beta,
\]
so that attention and MLP components can implement addition by moving between Fourier channels and their pairwise interactions [2402.16726].

For one-layer transformers, the attention matrices themselves admit a Fourier-factorized description. A head’s low-rank key-query product can be written as an outer product of cosines,
\[
[W^K W^{Q\top}]_{i,j}
=
\cos\!\bigl(2\pi\zeta\,i/p+\phi_k\bigr)\,
\cos\!\bigl(2\pi\zeta\,j/p+\phi_q\bigr),
\]
which, after multiplication by one-hot inputs, realizes a frequency mode of the sum. The softmax then converts these scores into a peaked distribution, while the value stage carries the corresponding cosine feature of candidate outputs [2402.09469].

This Fourier mechanism is often described as a “clock” algorithm. A common misconception is that transformers solving modular addition learn a digital carry-ripple circuit. In the one-layer four-head setting studied for full-precision, binary, and ternary networks, none of the models learn such a discrete digital circuit; instead, all three encode continuous sine/cosine channels at learned frequencies and use attention to sum those channels across the two inputs [2405.17703].

## 3. Optimization, grokking, and sample-complexity structure

The modular-addition benchmark is also central to the theory of grokking. Early in training, when the kernel regime approximately holds, permutation-equivariant models behave like their NTK linearization. In that regime, no permutation-equivariant model can achieve small population error on modular addition unless it sees at least a constant fraction of all possible data points. For regression, if a kernel method is equivariant under permutations of $(a,b,c)$ and the number of training triples satisfies $n<Cp^3$, then the expected population loss obeys
\[
\mathbb E[L_2]\ge p/2.
\]
For classification, any permutation-equivariant kernel on $(a,b)$ versus one-hot $c$ needs $n=\Omega(p^2)$ examples to beat chance [2407.12332].

Late in training, the picture changes. The same work shows that models eventually escape the kernel regime, and once they do, bounded $\ell_\infty$ solutions can generalize with substantially fewer examples. For regression, a width-$h$ interpolator with $\|\theta\|_\infty\le r$ satisfies
\[
L_2\le O\!\Bigl(\frac{p^2}{n}\,(r^6h^2)(\log^3 n)\Bigr),
\]
so $n=\tilde \Omega(p^2)$ suffices. For classification, a width-$h$ zero-training-error solution with $\ell_\infty$-normalized margin within a constant factor of the maximum satisfies an error bound leading to $n=\tilde\Omega(p^{5/3})$ for near-perfect generalization. Empirically, train loss reaches zero while test performance remains poor until the empirical NTK changes substantially; only after leaving the lazy regime does feature learning emerge and test accuracy jump [2407.12332].

A complementary max-margin analysis gives an explicit width requirement for Fourier solutions. In the one-layer setting, if $m$ denotes width, then achieving the maximum normalized $L_{2,k+1}$ margin on $D_p$ requires
\[
m\ge 2^{2k-2}(p-1),
\]
with each hidden neuron living on exactly one nonzero frequency. This links the learned Fourier decomposition to the implicit bias of weight-decay training at small $\lambda$, for which global minimizers converge to a solution achieving the maximum normalized margin $\gamma^*$ [2402.09469].

## 4. Scaling to large $N$ and cryptographic moduli

The most substantial departure from the classic small-$p$ toy setting is the encoder-only transformer training pipeline introduced for large-scale modular arithmetic. That work proposes three changes: more diverse training data, an angular embedding, and a custom loss function. The goal is to learn modular sums for values of $N$ up to $256$ and prime moduli $q$ up to $3329$ [2410.03569].

The training distribution is no longer uniform over $\mathbb Z_q^N$ alone. Instead, it interleaves sparsity sampling and tail sampling. Sparsity sampling chooses the number of nonzero positions using
\[
f(z)\propto \frac{1}{\sqrt{z+1}}
\quad\text{or}\quad
\frac{1}{z+1+\sqrt N},
\]
creating a curriculum of easier, sparser sums. Tail sampling targets rare wrap-around cases by drawing examples whose pre-modulo average lies in the tails of the uniform sum-distribution; only a small fraction of training examples, $\lesssim 0.05\%$, are drawn from this distribution [2410.03569].

The embedding is explicitly modular-aware. Each integer $a\in\mathbb Z_q$ is mapped to an angle
\[
\phi(a)=\frac{2\pi}{q}a,
\qquad
\mathrm{Embed}(a)=\bigl(\cos\phi(a),\sin\phi(a)\bigr)\in\mathbb R^2,
\]
so that adjacency between $0$ and $q-1$ is preserved geometrically. The loss combines a mean-squared term with a regularizer that both blows up at the origin and encourages predictions to remain near the unit circle:
\[
\ell_\alpha(x',y';\phi)
=
\alpha\Bigl(x'^2+y'^2+\tfrac1{x'^2+y'^2}\Bigr)
+(1-\alpha)\bigl[(\cos\phi-x')^2+(\sin\phi-y')^2\bigr],
\quad \alpha=0.01.
\]
Once $\|(x',y')\|\approx 1$, the loss is minimized at both $\phi'=\phi$ and $\phi'=\phi+2\pi$, so wrapping becomes a loss minimum rather than a discontinuity [2410.03569].

The model itself is an encoder-only transformer with $L=12$ layers, hidden dimension $256$, and $H=4$ attention heads, with positional encodings added to the input angular embeddings. The optimizer is Adam with $\mathrm{lr}=10^{-4}$, a $1$ k-step linear warmup, and cosine learning-rate decay. Training uses batch size $256$ per GPU on $8\times$ V100 GPUs, with $2.56\times 10^6$ examples per epoch per GPU for $30$ epochs, and no regularization beyond the custom loss [2410.03569].

On a held-out uniform test set of $100$ k examples, the reported performance is: for $(N,q)=(20,257)$, MSE $0.04\times 10^{-4}$ and $99.9\%$ exact accuracy; for $(20,769)$, MSE $0.03\times 10^{-4}$ and $98.2\%$ exact; for $(20,3329)$, MSE $0.04\times 10^{-4}$ and $57.0\%$ exact; for $(100,3329)$, MSE $0.42\times 10^{-4}$ and $20.7\%$ exact; and for $(256,3329)$, MSE $0.46\times 10^{-4}$ and $16.4\%$ exact, while $\tau=0.5\%$ accuracy remains $99.6\%$. Baseline MLPs or transformers trained on uniform data alone fail, with exact accuracy approximately $0\%$ even on the $(20,257)$ problem. The same approach retains $\tau=0.5\%$ accuracy of approximately $98\%$ at $(N,q)=(384,3329)$ [2410.03569].

## 5. Mechanistic interpretability, compression, and representational geometry

Mechanistic work on modular-addition transformers has progressively moved from neuron-by-neuron inspection toward analytic and geometric descriptions of full circuits. In one-layer “pizza” transformers, the MLP can be rigorously interpreted as a quadrature scheme. Under an infinite-width lens, the ReLU MLP
\[
f(x)=W_2\cdot \phi(W_1x)
\]
is approximated by an integral kernel, and each neuron corresponds to a rectangle in a numerical-integration rule for a trigonometric identity. The resulting Riemann-sum error satisfies
\[
E\le L\cdot \sum_{i=1}^N w_i\max(\Delta_i,\Delta_{i-1})
\le L\cdot W\cdot \max_i \Delta_i,
\]
and, for roughly uniform spacing, $E\le L(2\pi)^2/N$. In the trained pizza MLP analyzed in that work, $L\approx 2$ and $N=512$ yield $E\lesssim 0.05$, while the relative error is approximately $0.02$, providing a non-vacuous proof in $O(\text{parameter-count})$ time [2412.03773].

A second development replaces local circuit sketches with manifold-level descriptions. Earlier interpretations had associated trainable-attention models with “Clock” circuits and uniform-attention models with “Pizza” circuits. A later geometric-topological analysis argues that this is not the case: both architectures implement the same algorithm via topologically and geometrically equivalent representations. The first-layer frequency-specific manifolds are modeled either as a filled $2$-D disc or as a discrete torus $S^1\times S^1$, and deeper layers collapse these structures to the circle of correct-sum logits. Across $703$ independent seeds of each architecture, principal-component geometry, phase-alignment distributions, MMD tests, and persistent homology are used to argue that the learned circuits are equivalent at the representation level rather than fundamentally distinct [2512.25060].

A further unification is the “approximate Chinese Remainder Theorem” interpretation. In that account, neurons activate exclusively on approximate cosets, frequency clusters behave like modular subsystems, and the full network aggregates a small set of frequencies to isolate the correct output. The theory predicts that universally learned solutions in deep networks with trainable embeddings or more than one hidden layer require only $O(\log n)$ features, and the paper reports empirical confirmation across MLPs and transformers [2505.18266].

## 6. Architectural variants, robustness, and extensions beyond plain addition

The modular-addition transformer literature also serves as a comparative laboratory for architecture and representation. In binary, ternary, and full-precision one-layer four-head transformers trained on modulus $P=113$, all models learn similar algorithms rather than qualitatively different ones. Binary BitNet and ternary BitNet $b_{1.58}$ approximate the same Fourier “clock” mechanism, and the study concludes that quantization does not make the learned solution fundamentally more interpretable; binary networks in fact exhibit a higher Fourier noise floor than full-precision models [2405.17703].

A distinct extension studies text-based modular arithmetic at the character level. For $p=97$, a two-layer transformer with learned absolute positional embeddings can achieve high in-distribution accuracy yet fail catastrophically under position shift or out-of-distribution natural-language templates. The identified failure mode is that a baseline model is near-perfect at one absolute position but collapses when the same arithmetic expression is shifted. A training recipe combining explicit expression boundary markers, a steps-based position curriculum, diverse template mixtures, and multi-variant consistency regularization substantially improves robustness while maintaining high in-distribution accuracy; an ALiBi-style ablation fails to learn the task under that setup [2601.04283].

Modular addition also functions as a source task for transfer and for comparison with other modular operators. In modular polynomials and related arithmetic, Fourier Frequency Sparsity and Fourier Coefficient Ratio are proposed as progress measures. Addition is characterized by a sharp drop in Fourier Frequency Sparsity while Fourier Coefficient Ratio stays approximately $1$, whereas other operators exhibit different signatures. Transfer from grokked addition is limited rather than universal: pre-grokked addition models help on linear expressions such as $2a+b$, but transferability is reported only for specific combinations, and challenging non-factorizable polynomials do not exhibit equally clear patterns [2402.16726].

At larger scale, the same large-$N$ training pipeline used for modular addition extends to other modular functions $h:\mathbb Z_q^N\to\mathbb Z_q$, including
\[
\Bigl(\sum_i x_i^2\Bigr)^2 + x_1 \pmod q,
\]
for which the reported exact accuracy exceeds $90\%$ on $N=20,q=257$. This suggests a broader role for modular-aware embeddings, losses, and training distributions in neural computation over finite rings, including tasks motivated by Learning With Errors and related cryptosystems [2410.03569].

Source: https://www.emergentmind.com/topics/modular-addition-transformers