cs.LGMay 8, 2026

Learning Large-Scale Modular Addition with an Auxiliary Modulus

Authors: Hanato KikuchiRyosuke MasuyaKazuhiko KawamotoHiroshi Kera

Organizations: Chiba University · National Institute of Informatics

Abstract

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 KqKq 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 NN, large modulus qq, 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=64N=64 and q=974269q=974269, our method trained on 100K samples achieves 97.0%97.0\% ττ-accuracy at τ=0.05τ=0.05, while the sparse method achieves only 9.5%9.5\% with the same data size and 93.9%93.9\% even when extended to 1M samples.

Explore similar work

Jun 16, 2026cs.LG

The Discrete-Log Clock: How a Transformer Learns Modular Multiplication

When small transformers grok modular multiplication, prior work reports that the learned embedding has a "dense" Fourier spectrum requiring all frequencies. This contrasts with modular addition, where only a sparse set of key frequencies suffices. We show this density is an artifact of analyzing in the wrong basis. The natural Fourier transform for multiplication is not the standard additive DFT but the multiplicative character transform, which decomposes functions on the multiplicative group (Z/pZ)(\mathbb{Z}/p\mathbb{Z})^* into its irreducible representations. Applying this transform to a grokked transformer trained on abmod113a \cdot b \bmod 113, we find the embedding spectrum becomes highly sparse (Gini coefficient 0.58 vs. 0.07 in the additive basis) with only 4 key frequencies carrying significant energy. Furthermore, 96.9% of MLP neurons are cleanly tuned to a single multiplicative frequency, and neuron activation heatmaps reveal 2D-periodic structure when reordered by the discrete logarithm. These results demonstrate the transformer reduces multiplication to addition in discrete-log space, implementing a "Discrete-Log Clock" algorithm analogous to Nanda et al.'s Clock algorithm for addition. The methodology generalizes: matching the analysis basis to the algebraic structure of the task reveals interpretable structure where standard tools see noise.
Huu Danh Nguyen
May 27, 2026cs.LG

Learning High-Dimensional Parity Functions with Product Networks using Gradient Descent

Parity functions are fundamental Boolean operations with critical applications across machine learning, cryptography, and error correction. Yet, learning high-dimensional parity functions poses significant challenges: in a general setting, standard neural network architectures typically require exponential sample complexity, making gradient-based optimization intractable for large number of inputs NN. We demonstrate that compact product-based neural architectures combined with stochastic data sparsity (Bernoulli inputs with pe1/Np_e \leq 1/N) and appropriate hyperparameter choice enable efficient parity learning, with theoretical guarantees of convergence. Experiments validate our theory across dimensions up to N=100,000N = 100{,}000, with empirical evidence showing optimal hyperparameter choices for pep_e and learning rate αα, as well as polynomial complexity scaling laws. This work establishes fundamental connections between architectural inductive bias and data sparsity, opening new possibilities for neural arithmetic, structured reasoning, binary neural networks, and machine learning applied to automated protocol discovery.
Guillaume Larue, Louis-Adrien Dufrène, Quentin Lampin +2
May 31, 2026cs.LG

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 \href{https://github.com/DL-World-Models/mnist-math}{accessed here}.
Divyansh Jha, Yuanfang Xie, Varan Mehra +1