# Making Hard Problems Easier with Custom Data Distributions and Loss Regularization: A Case Study in Modular Arithmetic

arXiv: https://arxiv.org/abs/2410.03569 · HTML: https://mlsanity.com/paper/2410.03569/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 32 references that are in the map.
- First posted on arXiv 4 Oct 2024.

## Builds on (most cited first, 19 of 19)

- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Progress measures for grokking via mechanistic interpretability](https://mlsanity.com/paper/2301.05217/) (2023)
- [GPT-NeoX-20B: An Open-Source Autoregressive Language Model](https://mlsanity.com/paper/2204.06745/) (2022)
- [Neural Turing Machines](https://mlsanity.com/paper/1410.5401/) (2014)
- [Grokking modular arithmetic](https://mlsanity.com/paper/2301.02679/) (2023)
- [Failures of Gradient-Based Deep Learning](https://mlsanity.com/paper/1703.07950/) (2017)
- [Length Generalization in Arithmetic Transformers](https://mlsanity.com/paper/2306.15400/) (2023)
- [Linear algebra with transformers](https://mlsanity.com/paper/2112.01898/) (2021)
- [Why Do You Grok? A Theoretical Analysis of Grokking Modular Addition](https://mlsanity.com/paper/2407.12332/) (2024)
- [Learning the greatest common divisor: explaining transformer predictions](https://mlsanity.com/paper/2308.15594/) (2023)
- [SALSA: Attacking Lattice Cryptography with Transformers](https://mlsanity.com/paper/2207.04785/) (2022)
- [Provable Advantage of Curriculum Learning on Parity Targets with Mixed Inputs](https://mlsanity.com/paper/2306.16921/) (2023)
- [Grokking Modular Polynomials](https://mlsanity.com/paper/2406.03495/) (2024)
- [SALSA PICANTE: a machine learning attack on LWE with binary secrets](https://mlsanity.com/paper/2303.04178/) (2023)
- [Analyzing the Nuances of Transformers' Polynomial Simplification Abilities](https://mlsanity.com/paper/2104.14095/) (2021)
- [Learning to Compute Gr"obner Bases](https://mlsanity.com/paper/2311.12904/) (2023)
- [Salsa Fresca: Angular Embeddings and Pre-Training for ML Attacks on Learning With Errors](https://mlsanity.com/paper/2402.01082/) (2024)
- [Benchmarking Attacks on Learning with Errors](https://mlsanity.com/paper/2408.00882/) (2024)
- [Machine learning for modular multiplication](https://mlsanity.com/paper/2402.19254/) (2024)

## Built on by (most cited first, 4 of 4)

- [TAPAS: Datasets for Learning the Learning with Errors Problem](https://mlsanity.com/paper/2510.08797/) (2025)
- [Learning Pseudorandom Numbers with Transformers: Permuted Congruential Generators, Curricula, and Interpretability](https://mlsanity.com/paper/2510.26792/) (2025)
- [Improving ML Attacks on LWE with Data Repetition and Stepwise Regression](https://mlsanity.com/paper/2604.03903/) (2026)
- [Learning Large-Scale Modular Addition with an Auxiliary Modulus](https://mlsanity.com/paper/2605.07648/) (2026)

## Abstract (arXiv)

Recent work showed that ML-based attacks on Learning with Errors (LWE), a hard problem used in post-quantum cryptography, outperform classical algebraic attacks in certain settings. Although promising, ML attacks struggle to scale to more complex LWE settings. Prior work connected this issue to the difficulty of training ML models to do modular arithmetic, a core feature of the LWE problem. To address this, we develop techniques that significantly boost the performance of ML models on modular arithmetic tasks, enabling the models to sum up to N=128 elements modulo q ≤ 974269. Our core innovation is the use of custom training data distributions and a carefully designed loss function that better represents the problem structure. We apply an initial proof of concept of our techniques to LWE specifically and find that they allow recovery of 2x harder secrets than prior work. Our techniques also help ML models learn other well-studied problems better, including copy, associative recall, and parity, motivating further study.