Learning Large-Scale Modular Addition with an Auxiliary Modulus

TL;DR

This paper introduces a covariate-shift-free training method for large-scale modular addition, enabling scalable and sample-efficient learning even with large inputs and small datasets.

Contribution

It proposes an auxiliary modulus technique that maintains input distribution consistency, improving learning performance over existing methods.

Findings

The method achieves high accuracy on large input and modulus sizes.

It outperforms sparse methods in low-data regimes.

Experiments demonstrate scalability and efficiency.

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