Learning Large-Scale Modular Addition with an Auxiliary Modulus

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 $τ$-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.

• The paper introduces an auxiliary modulus approach for modular addition to eliminate covariate shift, resulting in improved model accuracy.
• The method allows stable learning of modular addition tasks across varied architectures, ensuring robustness against hyperparameter changes.
• Empirical analysis demonstrates significant performance gains compared to sparse input methods, particularly with large-scale modular operations.

Covariate-Shift-Free Modular Addition Learning via Auxiliary Modulus

Introduction

Learning modular addition functions, especially at scale, is a non-trivial task for neural architectures due to intense input sensitivity and high misgeneralization rates. Previous approaches to address large-scale modular addition—critical in algorithmic reasoning and cryptographic settings—bias training distributions by enforcing sparsity in the input, thereby reducing problem difficulty and accelerating convergence. However, this sparsification leads to a covariate shift between training and test distributions, which manifests as generalization gaps and structural brittleness.

This work, “Learning Large-Scale Modular Addition with an Auxiliary Modulus” (2605.07648), proposes an alternative approach: instead of altering the input distribution or decomposing the task through explicit curriculum, an auxiliary modulus task is introduced during training. This enables sample-efficient learning of addition modulo large $q$ without covariate shift. The central contribution is the theoretical motivation, empirical demonstration, and architectural robustness of this auxiliary-modulus paradigm.

Covariate Shift in Sparse Modular Addition Training

The sparsification technique [24] increases the prevalence of zeros in the training input, essentially emulating smaller-sum modular addition tasks by reducing the effective number of summands. While effective from an optimization perspective—mitigating the vanishing gradient issue for large moduli—it induces substantial distribution shift, as the test data is composed of fully dense uniform integer vectors. The authors formalize and empirically validate this covariate shift, providing a lower bound on the resultant generalization gap. The gap scales strongly with both modulus $q$ and summand count $N$, reaching irreducible magnitudes for realistic parameters, corresponding to adverse test-time performance unless the sparse method is very finely tuned.

Notably, the sparse method is architecturally sensitive, suffering training collapse under modest shifts in layer normalization placement, bias inclusion, weight initialization, or dropout rate. Such sensitivity renders the method unsuitable as a standardized solution for modular arithmetic learning in large and heterogeneous systems.

Auxiliary Modulus Training Scheme

The paper introduces a training methodology where, instead of modulating the input difficulty via sparsity, the model is tasked, with probability $r$, to perform modular addition under modulus $Kq$ (for integer $K > 1$), and with probability $1-r$, under target modulus $q$. Inputs remain unchanged—uniformly sampled from $[0, q-1]^N$—thereby eliminating covariate shift. The expanded modulus reduces the effective wrap-around rate, thus attenuating the optimization difficulty associated with high-modular wrap frequency. This “auxiliary modulus” task is implemented either through extension of output vocabulary (for token-categorical models) or through angular embedding (for continuous output models).

Theoretical analysis shows that the expected wrap count, the principal measure of problem difficulty, is comparably reduced as with input sparsity, without altering the data distribution. Empirically, the method achieves high accuracy even with small $r$ and moderate $q$0, and exhibits performance stability across wide hyperparameter ranges.

Experimental Analysis

Benchmarked across a wide spectrum of configurations ($q$1, $q$2), the proposed method achieves consistently high match accuracy and T-accuracy. The improvements are especially pronounced under sample-constrained regimes. For example, at $q$3 and $q$4, 100K train samples suffice for 97% T-accuracy (T=0.05); the sparse method attains only 9.5%, even when scaling data size by $q$5 achieves just 93.9%.

Critically, the auxiliary modulus method outperforms the sparse baseline even when allowed a tenth of the data, and maintains high accuracy as $q$6 and $q$7 scale. The method’s performance is robust to model configuration and hyperparameter settings (e.g., different $q$8 choices), with minimal performance degradation under architectural variation (such as with or without dropout, bias, pre-/post-norm ordering).

Angular embedding variants further improve sample efficiency. With $q$9, the method yields 93.0% match accuracy with 100K samples, compared to the baseline’s 63.8% even at 1M samples.

Discussion and Implications

The auxiliary modulus strategy eliminates the primary failure mode of previous approaches—covariate shift—while preserving or improving their sample efficiency. This demonstrates that the challenge in modular addition learning is primarily in controlling optimization landscape difficulty without data distribution manipulation, and that auxiliary-task frameworks provide a robust mechanism for this.

The theoretical analysis reveals that, by leveraging a highly related but smoother auxiliary task (modulus $N$0), the model is encouraged towards learning group-theoretic periodic structure (as seen in mechanistic interpretability studies) while remaining exposed to the full support of the input space. This resonates with recent trends in using auxiliary objectives and multitask learning for enhancing model robustness and generalization. The approach is agnostic to network architecture, readily porting between categorical-token and continuous-angle parameterizations.

Practically, this method is highly relevant to end-to-end learning scenarios in cryptographic protocol analysis, algorithmic reasoning, and symbolic AI, where distribution mismatch between training corpus and deployment setting is a chronic concern. The ability to train models on large modular addition tasks with high reliability and minimal hyperparameter tuning enables scalable neural algorithmic search, post-quantum cryptography attacks (such as on LWE-based systems), and the future integration of deep modular arithmetic into larger neuro-symbolic frameworks.

Future Directions

The method’s main constraint is its current specificity to the modular addition operation. Extending auxiliary modulus training to other modular functions (e.g., modular multiplication, polynomial evaluation) is both a compelling theoretical and practical open problem, and may require more elaborate auxiliary objective design. Additionally, further work could establish principled frameworks for hyperparameter ($N$1) selection based on task difficulty metrics and establish tighter connections between expected wrap statistics and empirical performance. Integration with broader multitask or meta-learning objectives is also a natural next step.

Conclusion

The introduction of auxiliary modulus training constitutes a robust and highly sample-efficient technique for learning large-scale modular addition without covariate shift. The approach substantially outperforms the prevailing sparse input method in both match and T-accuracy, is robust to model and training hyperparameters, and operates effectively across a range of neural architectures. Its contributions clarify the source of sample complexity in modular addition, and establish auxiliary objective design—without data distribution manipulation—as a preferred paradigm for algorithmic reasoning tasks in deep learning (2605.07648).