Teaching Transformers Modular Arithmetic at Scale

Abstract

Modular addition is, on its face, a simple operation: given N elements in Z_q, compute their sum modulo q. Yet, scalable machine learning solutions to this problem remain elusive: prior work trains ML models that sum N 6 elements mod q 1000. Promising applications of ML models for cryptanalysis-which often involve modular arithmetic with large N and q-motivate reconsideration of this problem. This work proposes three changes to the modular addition model training pipeline: more diverse training data, an angular embedding, and a custom loss function. With these changes, we demonstrate success with our approach for N = 256, q = 3329, a case which is interesting for cryptographic applications, and a significant increase in N and q over prior work. These techniques also generalize to other modular arithmetic problems, motivating future work.

Tasks

Reproductions