Research map / Learning theory and algorithms / Prediction and decision losses
Making Hard Problems Easier with Custom Data Distributions and Loss Regularization: A Case Study in Modular Arithmetic
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes a sparsity-based training data distribution and a regularized angular MSE loss that let encoder-only transformers learn modular addition of up to 128 terms and recover binary LWE secrets with Hamming weight up to 6.
Problem
ML models struggle to learn modular arithmetic (summing many elements mod q), which limits ML-based attacks on Learning with Errors (LWE) to few nonzero secret entries.
Approach
Augment training data with sparse vectors sampled from a number-of-zeros distribution f (inv_sqrt) instead of uniform sampling, and add a regularization term penalizing predictions near the origin to the angular MSE loss, so the model avoids collapse and learns easier-to-harder structure without a curriculum.
Key results
- Modular addition with N=128 terms and q up to 974269 reaches 92-96% tau=0.5% accuracy, versus 1.3% for the default distribution baseline at N=16, q=257.
- On LWE with N=64, the method recovers binary secrets with Hamming weight up to 6, versus about 3 for prior ML attacks (claimed doubling).
- Curriculum learning gives 77-86% accuracy at N=128 with higher variance, while the proposed distribution gives 92-96% with low variance.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add sparse training examples using f_inv_sqrt instead of default sampling (ablation) | Encoder-only transformer modular addition · Modular addition N=16, q=257 | τ=0.5% accuracy: 1.2% -> 99.8% | Helped |
| Use f_inv_sqrt rather than uniform sparsity sampling (ablation) | Encoder-only transformer modular addition · Modular addition N=16, 32, 64, 128; q=257 | τ=0.5% accuracy: 99.7% vs 99.8%; 98.9% vs 99.5%; 95.3% vs 98.9%; 92.7% vs 96.1% | Helped |
| Add sparse examples to the training data | Encoder-only transformer modular addition · Modular addition N=16, q=257 | τ=0.5% accuracy: 1.3% -> 99.8% | Helped |
| Add custom loss regularization (ablation) | Encoder-only LWE model · LWE without errors, N=64, secret weight=6, q=257 | Secret recovery: 0% -> 15% | Helped |
| Add custom loss regularization (ablation) | Encoder-only transformer modular addition · Modular addition N=16, q=257 | τ=0.5% accuracy: 99.9% vs 99.8% | No clear effect |
| Increase training repeats from 10 to 1000 (ablation) | Encoder-only transformer modular addition · Modular addition N=128, q=257 | τ=0.5% accuracy: 96.1% -> 86.1% | Hurt |
| Increase training duration by 4× (ablation) | Encoder-only transformer modular addition · Modular addition N=256, q=257 | τ=1% accuracy: 90.4% -> 94.8% | Helped |
| Increase model size by 4× (ablation) | Encoder-only transformer modular addition · Modular addition N=256, q=257 | τ=1% accuracy: 90.4% -> 96.2% | Helped |
| Replace fixed-length training with diverse problem lengths using f_inv_sqrt | Decoder-only transformer associative recall · Associative recall, max length 64 | Accuracy: 1.8% -> 100.0% | Helped |
| Replace fixed-length training with diverse problem lengths using f_inv_sqrt | Decoder-only transformer parity · Parity, max length 256 | Accuracy: 50.2% -> 99.4% | Helped |
| Replace f_default with f_inv_sqrt sampling | Transformer modular multiplication · Modular multiplication N=16, q=3329 | τ=1% accuracy: 2% -> 3% | Helped |
| Replace f_default with f_inv_sqrt sampling | Transformer scalar product · Scalar product N=8, q=257 | τ=1% accuracy: 2% vs 2% | No clear effect |
Surprising findings
- Including a random fixed integer K instead of 0 as the sparse element yields similar gains, suggesting the distribution shift rather than the element itself matters.
- Repeating a moderate number of distinct samples (about 10 repeats) gave the best accuracy, and very few distinct samples (0.1M) still failed to generalize as well.
Limitations
- Performance declines as N grows (e.g., about 90% at N=256 with tau=1%).
- LWE experiments ignore the error term and assume binary sparse secrets, so noisy LWE is not demonstrated.
- Modular multiplication and scalar product results degrade for larger q (e.g., 3% at q=3329).
Experimental setup
Datasets: synthetic modular addition data with sparse vectors (f_inv_sqrt) synthetic LWE (A, b) pairs synthetic copy, associative recall, parity, selective copy data
Benchmarks: held-out modular addition test set (100,000 examples) LWE binary secret recovery copy associative recall parity selective copy
Metrics: MSE tau-accuracy (0.5%, 1%) exact accuracy secret recovery %
Baselines: default uniform training distribution curriculum learning standard MSE loss cross-entropy loss
Builds on: angular embedding for modular arithmetic ML attacks on LWE (Wenger et al.) curriculum learning grokking studies
Closest papers by meaning
- TAPAS: Datasets for Learning the Learning with Errors Problem (NeurIPS 2025 · similarity 0.90)
- Mind the Gap: A Practical Attack on GGUF Quantization (ICML 2025 · similarity 0.84)
- A Reliable Cryptographic Framework for Empirical Machine Unlearning Evaluation (NeurIPS 2025 · similarity 0.84)
- MathAttack: Attacking Large Language Models towards Math Solving Ability (AAAI 2024 · similarity 0.83)
- Reconstruction Attacks on Machine Unlearning: Simple Models are Vulnerable (NeurIPS 2024 · similarity 0.82)
- Exploiting LLM Quantization (NeurIPS 2024 · similarity 0.82)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.