Making Hard Problems Easier with Custom Data Distributions and Loss Regularization: A Case Study in Modular Arithmetic
Authors: Eshika Saxena
Alberto Alfarano
Emily Wenger
Kristin E. Lauter
Venue: ICML 2025 poster (2025)
Keywords: transformers
modular arithmetic
math
cryptography
Primary area: applications
TL;DR: We introduce two techniques, varying the diversity of training data and introducing a regularized loss function, to improve transformer learning on modular arithmetic, cryptanalysis, and other synthetic tasks

Abstract: 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 \le 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.

Making Hard Problems Easier with Custom Data Distributions and Loss Regularization: A Case Study in Modular Arithmetic
Eshika Saxena 1
Alberto Alfarano" 1
Emily Wenger t 122 Kristin Lauter
Abstract
Recent work has shown that machine learning (ML) models can, under certain conditions, be used to break LWE by recovering the secret (Wenger et al., 2022: Li et al., 2023a;b; Stevens et al., 2024). The attack trains ML models to predict b given A, which requires learning to compute the (noisy) modular arithmetic operation A (mod q)- If models predict with even low aceuracy,
Recent work showed that ML-based attacks on Learning with Errors (LWE), a hard problem used in post-quantum cryptography, outperform classi- cal 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, 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 mod- ulo q S 974269. Our core innovation is the use of custom training data distributions and a care- fully designed loss function that better represents the problem structure. We apply an initial proof of concept of our techniques LWE specifically and find that they allow recovery of 2x harder se- crets than prior work. Our techniques also help ML models learn other well-studied problems bet- ter, including copy, associative recall, and parity, motivating further study.
be recovered.
However, ML attacks struggle recover LWE secrets with many nonzero entries. limiting their efficacy. Prior work links this struggle to ML models' inability to learn modu- lar arithmetic at scale, since recovering LWE secrets with more nonzero entries requires summing more elements mod q (Wenger al.. 2024). Given the widespread adoption of LWE-based cryptosystems, exploring the potential of these nascent ML attacks which can exploit statistical depen- dencies algebraic attacks may miss-is critical to improve understanding of LWE's practical security.
ML models' struggle to perform modular arithmetic extends beyond eryptanalysis. Numerous works have documented the poor performance of ML models on modular arithmetic problems with many terms and/or large prime moduli (Pala- mas, 2017; Gromov, 2023; Abbe et al., 2023; Mohamadi et al., 2024; Doshi et al., 2024). This is surprising because ML models can learn other complex math tasks such as symbolic regression, linear algebra, discovering Lyapunov functions, computing Gröbner bases, polynomial simplifica- tion, and computing the greatest common divisor (Charton et al., 2021; Charton, 2022; Alfarano et al., 2024; Kera et al., 2024; Agarwal et al., 2021: Charton, 2024). Modular arith- metic, on its face, seems easier, but scalable ML. solutions remain elusive.
1. Introduction
Modular arithmetic is a key component of many crypto- graphic hard problems, including Learning with Errors (LWE), which is the basis for several newly standardized post-quantum cryptosystems (Chen et al., 2022). The LWE problem involves determining secret vector given two pieces of information: A E Z Zmxx, matrix sampled uni- formly at random mod and b = A E 2,m. the noisy dot product of A with the secret vector s mod g. When s binary, which often happens in practical applications of LWE, computing b from A requires computing a subset sum mod g, which is modular addition problem.
Thus, our goal is simple: train ML models that can sum many elements mod 4. In considering this challenge, we make two key observations. First, modular arithmetic problems lie along gradient of difficulty based on how many times the sum "wraps" around the modulus. Prior work observed that models learn better from examples that wrap fewer times around the modulus q. e.g. Σns S 2q (Wenger et al., 2024). Second, the geometry of the mod- ular field disrupts traditional notions of gradient descent, as and q are close (Shalev-Shwarl Shwartz et al., 2017). Prior work proposed angular embedding for model input/outputs that
Equual contribution Equal contribution 'Meta AI Duke Uni- versity. Correspondence to: Eshika Saxena <eshika @ mema.com>.
Proceedings of the nnnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Making Hard Problems Easier with Custom Data Distributions and Loss Regularization
better represented this geometry (Stevens et al., 2024), but this structure was not not explicitly used during optimization.
2024). Table summarizes the best prior results on modular addition. The best existing methods train models that sum N I 6 elements for moduli up g 1000.
Our contribution. Building on these observations, we hy- pothesize that two changes to the training process-ustee custom training data distributions and problem-specific loss regularization-mmy help models better learn modular arith- metic. The use of custom data distributions is inspired by recent prior work that demonstrated the efficacy of custom data in related domains (Charton & Kempe; Abbe et al., 2023). Crafting training data that simultancously exposes models to easy and harder versions of modular arithmetic (e.g. elements that wrap fewer or more times around q) could help the model unlock the structure of the problem. Loss regularization is widely used in ML training, but could be adapted to better represent modular arithmetic specifi- cally, as noted by prior work (Gromov, 2023; Jelassi et al., 2023). Building on this, we:
Prior work has laid groundwork for analytically understand- ing how models learn modular arithmetic (Gromov, 2023; Doshi et al., 2024; Nanda et al., 2023). Modular arithmetic is a common task used study "grokking" as the model suddenly transitions from memorization to generalization this task. Some works go to show that ML models are able to inherently learn the problem by encoding the modular arithmetic structure into the model weights (Nanda et al., 2023; Doshi et al., 2024). However, prior work consis- tently shows (Table 1) that models struggle to learn modular arithmetic as the number of summed terms N and modulus q increase, motivating our search for new learning methods.
Table I. Summary of prior work on ML-enabled modular ad- dition. Best N and g are bold.
Propose a novel training data distribution that in- cludes both easy and hard variants of the modular arith- metic problem. This distribution innriiennneeee as a function over training elements rather than a curricu- lum-smoothly adjusts the hardness of elements seen by the model, enabling efficient and effective learning.
Paper
# Terms
Mod (g)
% Accuracy
(N)
Nanda al. (2023)
2
53, 109,
100
Transformer
113, 401,401
Mohamadi et al. (2024)
2
433
100
2-layer MLP
Design a custom loss function with a penalty term specific to modular arithmetic. This function discour- ages model convergence at local, unhelpful minima, enabling discovery of global optima appropriate to our problem setting.
Doshi al. (2024)
11, 23
97.1
2-layer MLP
Gromov (2023)
97
100
2-layer MI.P
Jelassi et (2023)
100, 1000
73
Encoder-only transformer
Abbe et al. (2023)
2
2
100
4-layer ML.
Key results. Our methods enable ML models to perform modular addition for a variety of N and g, up to = 128 and q = 974269. This significantly outperforms prior work, which summed N < 6 elements mod q S 1000. When we apply our methods to the LWE problem specifically, we find we can recover 2x harder secrets than prior work.
Custom training data distributions. Some prior work has explored how changing the training data distribution affects performance on arithmetic tasks. For example, Charton & Kempe show that repeating training examples improves transformer performance on arithmetic tasks like modular multiplication, greatest common divisor, and eigenvalue calculation. Mohamadi et al. (2024) briefly investigated the effect of training data on modular arithmetic and found that models need to be trained constant fraction of all possible modular arithmetic behaviors for given N and q to generalize. These works motivate further investigation into how the distribution of training data can affect model performance on arithmetic tasks.
Although our motivation is LWE-specific, our methods could aid learning on other problems as well. In $4.2, we highlight several interesting findings about how models learn modular arithmetic, which yield insights that could aid learning in other settings. We also apply our methods to other
problems (copy, associative recall, and parity) and find they enable models to learn better, demon- strating their potential to generalize. Our code is available at: https://itthub..om httpaa/l..ooeeebbooooeeceeerrree
Related work has investigated the impact of changing the training data distribution over time in a predetermined pat- tern, technique known as curriculum learning (CL) (Ben- gio et al., 2009). Abbe et al. (2023) show that the parity function can be learned via curriculum learning strategies. We note that our approach is different than CL because we do not change the data distribution over time, so our approach requires less tuning than CL. Our experiments
2. Related Work
ML for modular arithmetic. Previous work has inves- tigated whether ML models can learn modular arithmetic operations (Palamas, 2017; Lauter et al., 2024; Gromov, 2023; Abbe et al., 2023; Mohamadi et al., 2024; Doshi et al.,
2

---

Making Hard Problems Easier with Custom Data Distributions and Loss Regularization
show that CL is highly sensitive to the curriculum parame- ters, whereas our approach is simpler and ensures consistent convergence (see Table 4).
uniform density) and fnw [nw_eart(2) VN-=#1 (similar to the distribution chosen by Allen-Zhu (2025)) where x means the functions are rescaled by a constant such that the sum of f over all z in its domain equals 1. We compare these to a baseline of fdefault, which is the PDF of the number of zeros in a when a is drawn uniformly from 2 Most prior work on ML for modular arithmetic used faetautt for training data generation, as noted previously.
Effect of loss functions. Several works attribute models' failure to learn harder modular addition problems to the complexity of the loss space (Gromov, 2023; Jelassi et al., 2023). Shalev-Shwartz et al. (2017) highlight the limi- tations of gradient-based based methods on different problems, including learning random parities. These motivate adding loss regularization based on the problem to overcome these challenges.
Effect of f on training data. Figure 1 shows the sparsity of training data elements created using these sampling strate- gies with NN = 16 and 257. Note that fine is biased towards creating "harder" elements with few Os but still al- lows some easier examples to exist, whereas fdefault rarely creates examples with many zeros, meaning that the model must learn only from hard examples. During testing, we evaluate models examples drawn uniformly random from 2 since general modular addition is the target task.
3. Methodology
Following prior work (Jelassi et al., 2023), we train encoder- only transformer models to add N elements mod (fixed N and q for each model). We also leverage angular em- beddings proposed by Stevens et al. (2024), which map input and output elements mod 0 to points on the unit circle. These better represent the structure of modular arithmetic since 0 and 2x, which corresponds to q, are close under this geometry. This section describes our key innovations to the training pipeline, custom training data and loss regular- ization, then gives an overview of our end-to-end training procedure and evaluation metrics.
fuceept
funi
0.15
Probabiltty
0.05
3.1. Innovation 1: Augmenting Training Data with Sparse Vectors.
10
12
14
16
Number of non-zero elements
Most prior work on ML for modular arithmetic trains mod- els using randomly generated (a, b) pairs, where is drawn uniformly random from E i.e. a consists of elements [e1, .2...., avl, ag E Zg and b = Σ a; (mod q) (Je- lassi et al., 2023; Doshi et al., 2024). Based on observations about the importance of training data diversity from Wenger et al. (2024) and Mohamadi et al. (2024), models learn bet- ter when they see more examples of "simpler" versions of the target operation. Seeing these simplified problems may help models understand the modular arithmetic structure and learn better. Thus, we propose adding additional sparse vectors to the training data, in which more coordinates of a are 0, alongside more general vectors.
Figure 1. Probability of number of non-zero elements in each training data element when N = 16 and q = 257 for fuvvearos funi and faetault sampling distributions.
3.2. Innovation 2: Loss Regularization to Avoid Model Collapse.
Following Stevens et al. (2024), we use an angular embed- ding to represent transformer inputs and outputs. Practi- cally, the embedding encodes an integer E Zag as an angle é = 2n 1 and then as a point (cos(o), ()), sin( E RR² on the unit circle. Based on this, we initially chose an MSE loss for training. However, we observed that in harder settings the model would often converge to bad local minima like the origin. To address this issue, we use a custom loss function during training that combines mean squared error (MSE) loss with extra regularization term. Given a prediction of the form (a', 2²) and ground truth (x cos o, y sin Φ), this loss takes the form:
&
Data generation approach. We selectively add sparse vectors to the training data by employing a sampling distri- bution f in our data generation procedure, which controls the number of nonzero elements in a training data element a. To generate a, we first sample from f to get a number n E [1, N1. the number of nonzero elements in a. We then sample the nonzero elements of a from Z7, pad this with zero string of length N n, and shuffle the resulting vector to ensure Os are randomly distributed.
1
i = 222 22² +
x22 + y2 + ((x - x')² + (y - y))2),
We experiment with two f distributions: [mmi(2) = N (i.e.
a 10-
3

---

Making Hard Problems Easier with Custom Data Distributions and Loss Regularization
The first term penalizes the model for predicting the origin by driving the loss to infinity if x' 0.3/ 0. It also encourages the model to predict y/) ) on the unit circle (the first term is minimized with z²² +2² = 1). The second term is the standard MSE loss. After some training x' and y' are close to the unit circle, so we can approximate x' and y' as cos o' and sin ø'. Under this condition, the MSE loss function component becomes:
against the ground truth of (cos 2 8, sin s).
To get complete picture of model performance, com- pute the following metrics: Mean Squared Error (MSE) of angle predictions and % accuracy with a margin of er- ror (r) relative to q. MSE help us to evaluate the model's performance in terms of closeness between the predicted and ground truth angles (MSE). r-accuracy enables us to measure whether the model learns the approximate function behavior, even if exact accuracy is low. The formulae for these metrics are:
t (cos の - cos d' + (sin Ø - sin sin 2²)²
=2- 2 cos@ & cos ø' sin @ sin ø'
=2 2 co d')
This loss component will be minimized when cos(ó- 00) P 1, which occurs Φ 3 Ø 0 and & o' 2T. In the modular arithmetic setting, we want 0 and 2x to be understood as "close" in the loss space, so this loss term
MSE
=
[D]
Σ
((cos Φ - cos cos Φ²)² + (sin o - -snn )²)²)
1
T r-accuracy = [D]
ZED
correctly describes the desired behavior.
4. Results and Further Exploration
3.3. Model Training and Evaluation
4.1. Key Results
We implement the proposed changes and train encoder-only transformers to sum N elements mod q, using the following settings and evaluation metrics.
Best results on modular addition. Our methods enable models to learn modular addition of up to = 128 ele- ments mod q. We present best results across a range of N and q values in Table Overall, the MSE near across N and q, showing that the model converges and learns well. Notably, 7 0.5% accuracy is almost perfect for all models. This means that in almost all cases, an "incorrect" model prediction is still within 0.5% of q.
Parameter selection. We experiment with N (16, 32, 64, 128} to identify trends as N increases. Be- cause we are interested in the LWE application, we use prime moduli, which are commonly used in that setting. We use q = {257, 3329, 42899, 974269], including one (q- = 3329) used in standardized I. L.E--based eryptosystem, CRYSTALS-KYBER (Avanzi al., 2021). We also tested with non- non-prime modulus q 2 and obtained similar re- sults, as shown in Appendix A. We select N = 16, 257 as our base case because the sample space is large enough to ensure models generalize.
Comparison to baseline. We compare our results a baseline encoder-only transformer with standard MSE loss and the Sdetault distribution, as is used in some of the prior work (Jelassi et al., 2023). With the baseline approach, we observe 1.3% T = 0.57-accu on N = 16, q 257 data (our base case) with the same number of training data samples as we used. In other words, the model does not learn the task at all. In comparison, our methods achieve 99.8% 8% on the same problem. Unlike Gromov (2023) and Doshi et al. (2024), we do not observe grokking in models because we use a very small fraction of data from the possible sample space (2.76- 10-31 when N = 16 and = 257). As such, our models gradually learn with a standard training loss behavior and do not overfit.
Training procedure. All experiments were imple- mented in Python with Pytorch. We train the transformer models with a hidden dimension of 256, 4 attention heads, and 4 encoding layers on batches of 250 examples, using the Adam optimizer (Kingma & Ba, 2015) with a learning rate of 3- 10-3, an initial linear warm-up phase of 1,000 optimization steps, and cosine scheduling. All experiments run on 1 V100 GPU with 32 GB of memory. The models were trained with 10M distinct samples for total 100M samples. Training time is around 3 hours.
4.2. Deeper Exploration of Our Methods
Evaluation metrics. We generate a held-out test set Drest of size 100,000 that distinet from the training set and contains examples drawn uniformly from Z, To eval- uate model performance on Dusst we take the final hid- den state of the transformer and pass it through a linear layer to produce an output of the form 1a),))). We project this point onto the unit cirele, producing (cos ø', sin d') = (cos 2T s sin 2T 3²). The model prediction then compared q
Beyond simply improving ML performance on modular arithmetic, our methods offer interesting lessons on the learning process. Here, we present results that validate our approach and give more general insight about how models learn modular arithmetic.
Models learn easy examples before hard ones. A key motivation for our approach is that models should initially
4

---

Making Hard Problems Easier with Custom Data Distributions and Loss Regularization
Table Our methods perform consistently well adding e [16, 32, 64, 128) elements mod y E [257, 3329, 42899, 974269). All metrics are computed on held out test set. MSE is mean squared error,
80
0.5% Accuracy is percentage of predictions within 0. 0.005q of right answer and 7 = 1% Accuracy is percentage of predictions within 0.01g of right answer (sce $3). The models perform with consistently low MSE and very high r-accuracy.
60
Accuracy
40
nonzero elements securacy
nonzeno acouracy
nonzero elements accuracy
20
nunzero elements accunccy
50 nonsero elements secuncy
# Terms (N)
Mod (g)
MSE
T 0.5% Accuracy
T 1% Accuracy
nonzero elements sccuracy
20
40
100
16
257
0.00
99.8%
100.0%
Training completion (%)
3329
0.00
99.7%
100.0%
42899
0.00
99.7%
100.0%
Figure 2. The model learns to sum fewer nonzero elements carlier than more complex examples. Model accuracy (N 64, q = 237) after cach epoch on unseen test set stratified by number of nonzero elements, As the number of nonzero elements increases, it takes longer to reach perfect accuracy.
974269
0.00
99.7%
100.0%
32
257
0.00
99.5%
100.0%
3329
0.00
99.4%
100.0%
42899
0.00
99.4%
100.0%
974269
0.00
99.5%
100.0%
64
257
0.01
98.9%
99.4%
This strongly indicates that these models need to see sparse training examples to generalize.
3329
0.01
97.4%
99.4%
42899
0.01
97.4%
99.4%
974269
0.01
98.2%
99.4%
Table 3. Sampling the training data from fine produces the best accuracy results across N and q. 7 = 0.5% Accuracy is percentage of predictions within 0.005g of right answer (see $3), KI. divergence is the level of similarity hetween the training and testing datasets. With default sampling fdefault, the model does not learn at all. Distributions with KL divergence that is not too high or too low enable the model to perform best.
128
257
0.04
96.1%
98.2%
3329
0.04
92.9%
98.0%
42899
0.05
94.1%
97.9%
974269
0.04
93.3%
97.4%
learn from the sparse training examples before learning the full task. To validate that this occurs, we train model on dataset
N 64, q = 257 and monitor its performance Dval drawn from the same distribution as Derain. Figure 2 shows model accuracy on samples with 1 to 64 nonzero elements over the training epochs. Here, we see that the model initially performs better on sparse examples (e.g. I non-zero element) and then becomes accurate on more complex examples in later epochs. This suggests that these models first learn simpler sums and build that knowledge to learn more complex sums, supporting our use of sparsity sampling in creating training data.
Terms (N)
Mod (Q)
Training Data f
T 0.5% Accuracy
KL divergence
257
Sdetault fivv fani
1.2%
0.0
99.8%
25.2
99.7%
35,4
32
257
fdefault fun jarr fani
1.3%
0.0
99.5%
49,8
98.9%
71.5
64
257
fdefault S fine sqrt funi
1.3%
0.0
98.9%
98.1
95.3%
144.0
Sparse data elements are critical for learning. As de- scribed in $3, we construct more diverse training datasets by sampling elements with sparsity defined by PDF f. Here, we explore how different PDFs (factault finv haw_sqrt. and funi. see $3) affect model performance. We report two metrics: T = 0.5% accuracy of models and the Kullback-Lein (KL) divergence between the training and testing datasets. KL divergence quantifies the similarity between training dataset Drrain. constructed using function f, and Drest, sam- pled from the set Z uniformly at random, i.e. faefault. As Table 3 shows, the accuracy difference between models trained with the default sampling (fdefault) and any other distribution f is stark. The exact same architecture has neg- ligible accuracy if we do not modify the training dataset sparsity distribution and achieves over 90% when we do.
128
257
Sdatialt
1.3%
0.0
five sqrt fani
96.1
193.7
92.7%
289.5
KL divergence Durain/ impacts accuracy. We ob- serve that models trained using f that produce very low (Rs 0) or very high Drrsin/Drost KL divergence generalize worse than ff with mid-range KL divergence. Models trained with the default fdefauts distribution have 0 Pursin/DPess KL divergence, since the train/test distributions tical, and model accuracy is 0%. On the other hand, the uniform sparsity function funi diverges too far from the test distribution (there are too many "simple" samples in the funi distribution), resulting in lower accuracy. A "Goldilocks" KI. divergence needed, and we find that distributions with
almost iden-
5

---

Making Hard Problems Easier with Custom Data Distributions and Loss Regularization
fewer sparse training elements, like /iww_aget. work best.
The data distribution approach is agnostic of the element chosen for sparsity. Regarding choosing include 0 as the element for sparsity, we actually found that including random but fixed number instead of O yields similar results. It is not the number itself but instead the shift in distribution that leads to performance improvements. The key is chang- ing the KL divergence between the train and test sets, which is irrespective of the element chosen for sparsity.
Our data distribution produces more consistent results than curriculum learning. We explore how our approach compares to standard CL. since both involve improving model learning by showing the model easier versions of the task. However, CL requires (a) defining phases of the curriculum by ranking the difficulty of each samples and de- ciding when to switch phases, (b) determining how to adjust the weight decay and learning rate parameters, (c) defin- ing the data mixture across different phases. Our approach instead fixes these parameters during training by implic- itly encoding curriculum "phases" via our data distribution, resulting in more static but easier-to-train method.
To prove this, run experiment where we substitute the O with an arbitrary integer K and shift the distribution by sampling more Kss compared to all remaining elements in z We used a random K for each different to ensure K is not a factor. We report the results in Table 5.
For comparison, we implement CL as follows. First, train the model for fixed horizon on easier samples only, then train on the entire dataset, as in Abbe et al. (2023). Table 4 indicates that CL results in noisy outcomes, especially with large N, when evaluating averuge accuracy and variance over 8 trials. Conversely, our method, which fixes the data distribution during training, ensures convergence. We also explored alternative curriculum designs (e.g. training easier samples until training loss reaches a certain threshold) and conducted extensive hyperparameter tuning on weight decay, learning rate, alternative difficulty ranking, and data mixture. These alternatives yield results similar to those in Table 4. We provide the specific parameter sets tested in Appendix F and report the best results in Table 4. Ultimately, while it may be possible to design an optimal curriculum, our approach is simpler and task agnostic (see $6).
Table S. Our approach works well consistently across different integers used for sparsity. We train the models with the best data parameters from $4 and include random number K instead of in the data to make the examples sparse. T = 1% Accuracy is percentage of predictions within 0.01q of right answer (see $3)
# Terms
Mod (4)
K
T = 1% Accuracy
(N)
32
257
160
100%
3329
3176
100%
42899
24606
100%
974269
79062
100%
64
257
160
99.2%
3329
3176
99.3%
42899
24606
99.4%
974269
79062
99.3%
128
257
160
97.9%
3329
3176
98.3%
42899
24606
97.9%
Table 4. Our approach is more consistent and stable than CL, especially for larger N. We chose the best performing curriculum to compare to our approach. We report the average 7 = 0.5% accuracy (see $3) and the variance across 8 trials.
974269
79062
97.9%
Repeating examples helps. In applications like LWE crypt- analysis, the amount of training data may be limited, so we consider whether models can learn from fewer samples. Interestingly, recent work shows that it may be beneficial for models to encounter samples multiple times during train- ing, suggesting that limited training data may not be as problematic as anticipated (Charton & Kempe). To test whether repeated examples help in our setting, we train models with d = (..LM, 1M, 10M, 100M) distinct sam- ples from the fivv finy_aurt _qqrt sampling distribution over a fixed training budget of b = 100M samples. This means that the model sees each distinct sample b/d times during training, so 100M /0. 1M = 1000 times in the d = 0.1M case. We find that as long as the model is exposed to at least 1M distinct
0.5% Accuracy
# Terms (N)
Mod (q)
Our approach
CL
16
257
99.7% + 0.0%
99.7% + 0.0%
3329
99.6% + 0.0%
99.5% + 0.0%%
42899
99.6% + 0.1%
99.4% + 0.0%
974269
99.6% + 0.0%
99.4% + 0.0%
32
257
99.3% + 0.1%
98.8% + 0.4%
3329
99.2% + 0.1%
98.0% + 0.5%
42899
99.2% + 0.1%
97.8% + 1.2%
974269
99.4% + 0. 1%
98.5% + 0.4%
64
257
98.5% + 0.3%
95.7% +: 2.8%
3329
97.3% + 0.1%
91.2% + 2.5%
42899
97.3% + 0.2%
91.3% + 0.9%
to 100 repeats in our ex- periments-it it can readily learn the underlying algorithm (see Table 6).
974269
98.0% + 0.2%
93.9% + 3.2%
128
257
95.8% + 0.5%
85.1% + 10.3%
3329
92.2% + 0.7%
82.2% + 8.8%
42899
92.9% + 0.8%
80.0%
12.2%
Loss regularization prevents model collapse when the task is hard. Finally, we evaluate how the regularization
974269
92.1% + 0.4%
77.3% + 7.5%
6

---

Making Hard Problems Easier with Custom Data Distributions and Loss Regularization
is negligible, around 0.04% relative difference.
Table Seeing some-but but not too meny-repeated examples aids performance. We train the models with 257 and with different numbers of repeated examples (all with the finv distribution, angular embedding, and custom loss) and evaluate the same test set for all. data repeats is computed as 100M dd, where d is the number of distinct training examples out of the total 1 training budget. т = 0.5% Accuracy is percent of predictions within 0.005g of right answer (see $3).
5. Application: Cryptanalysis of LWE
Having established that our methods enable ML models to better learn modular addition, we now apply them to our motivating example: improved ML-enabled attacks on the Learning with Errors (LWE) problem. The goal of these attacks is find given A and b, where b A (mod q). A E Zmxx is a uniformly random m X n matrix, the secret 8 E Z is a vector with length n.
T 0.5% Accuracy
# data repeats
- 16
N - 32
N - 64
N 128
1000
99.3%
97.7%
95.7%
86.1%
For this proof-of-concept vector e so we can better study if models learn the harder part of the problem: computing A 8 (mod q). If models recover 8 without noise, recovery with noise should be pos- sible based on experiments from Wenger et al. (2022). We assume is binary and sparse, based on proposed implemen- tations for LWE in the popular homomorphic encryption setting (Bossuat et al., 2024).
evaluation, we ignore the error
100
99.8%
99.2%
98.4%
92.7%
10
99.8%
99.5%
98.9%
96.1%
1
99.7%
99.25
97.2%
91.5%
term in our custom loss function affects model performance. We train several models with varying N and q on two ver- sions of the loss function from $3.2: one with a 10-4, activating our additional term, and one with is standard MSE loss. We find that when the task is still easy, regularization is not necessary (Table 7). However, as task difficulty increases, the regularization term increases the probability of success. We explore explanations for this in $5.
0.0, which
Following Stevens et al. (2024), we train an encoder-only model on (A, b) pairs to predict b from A. We use 30 mil- lion distinet data examples and a total training budget of billion examples, adding our improvements to the loss function to the training pipeline. We evaluate attack perfor- mance by measuring number of successful secret recoveries out of 20 random model initializations. We use the binary distinguisher of Wenger et al. (2022) to recover secrets.
Table 7. Model performance with custom and MSE loss are comparable when trained on the modular addition. We train the models with q 257 and with the best training data parameters identified in $4 with angular embeddings and evaluate on the same test set for all. 0.5% Accuracy is percentage of predictions within 0.005g of right answer (see $3).
We report the results in Table 8: for N 64 and q (257, 3329, 42899, 974269), the model successfully recov- ers binary secrets with Hamming weight (number of nonzero elements in the secret) up to 6. This result is significant be- cause Wenger et al. (2024) show that ML attacks on LWE traditionally only work when secret recovery requires effec- tively summing < 3 elements mod q. Thus, our methods appear to double the performance of ML. attacks on LWE, paving the way for future study.
t 0.5% Accuracy
Terms (N)
Mod (g)
Custom Loss
MSE Loss
a = 10-4
a = 0
16
257
99.8%
99.9%
3329
99.7%
99.7%
42899
99.7%
99.6%
Effect of regularization term in LWE setting, In the prior section ($4.2), we noted that the regularization term in our loss function aided learning in more difficult problem set- tings. To explain this, we consider the (ithout task set- ting. Without regularization, for given batch B with ground truth (æ, y) and model predictions (x', y²), the MSE loss to optimize is l = ||:: x=113 + ||y y |12 At the beginning of the training, we inspect model predictions, and we observe they are approximately equal, as shown in Appendix E. Using this and given b = a (mod q) is uniformly distributed, we conclude that, on average, the model loss i R EIE (x' cos 2n (2 + (y/ sin 2T (²)) Since ΣD cos 27 7 E. sin 2n = 0, the minimum loss is achieved when (a',8') = (0.0). This encourages the model to produce predictions with decreasing magnitude at
974269
99.7%
99.7%
32
257
99.5%
99.6%
3329
99.4%
99.2%
42899
99.4%
99.3%
974269
99.5%
98.9%
Our approach has minimal computational overhead. The computational overhead of our approach is minimal compared to the baseline, as we still generate the same num- ber of data samples and simply modify the distribution used for generation. Similarly, the loss regularization term does not introduce any additional cost (besides the negligible cost of calculating the term itself) as we have a standard training loop. We time the difference between the custom loss and standard loss experiments and report that the time difference
7

---

Making Hard Problems Easier with Custom Data Distributions and Loss Regularization
each training epoch, ultimately degenerating predictions centered around the origin. Adding the regularization term encourages the model to avoid the origin, allowing faster and better convergence (see Table 8 for comparison).
positions. Table 9 shows that for N 16 and y 257, achieve an accuracy exceeding 95%+ for these functions.
Table 9. With our methods, models can learn other modular arithmetic functions with good accuracy (N 16, 2 257). % Accuracy percentage of predictions exactly correct.
We experimented with other alternatives to loss regulariza- tion, such as estimating the angle distance between predic- tions, but found that the regularization term most effectively ensured consistent convergence.
Function
% Accuracy
1==1,k=11 (a) + ... + aN aa mod
95.1%
Table 8. Model performs better when trained with our custom loss on the (.ithouut),, problem. We train the models with the best data parameters from $4 and evaluate if the model successfully recovers the secret. Recovery % is the percent of secrets recovered out of 20 different model initialization.
Rj=1,k=3 = (a) + ... aN)
+ +aj mod
96.2%
hy=2,k=1 = (of +... a3) + mod
95.5%
Modular multiplication and scalar product. We test both modular multiplication and the scalar product of two vectors mod g. respectively in Table 10 and Table 11. The angular embedding is designed for addition, so we use standard to- ken embedding in multiplication experiments and compare Jinv_sqr to fdefault- For both tasks, the model with finv_sar performs well for smaller q, but declines for larger q. In all scenarios and in both tasks, we see that the t = 1% accuracy for fdeaault is around 2%. We also note that the scalar product more difficult, because the model needs to learn both both multiplications and additions.
Recovery %
# Terms
Hamming weight
Mod (g)
Custom Loss
MSE Loss a 0
(N)
a
10-2
64
257
%
0%
3329
20 %
0%
42899
15 %
0%
974269
15 %
0%
128
257
10%
3329
%
0%
42899
15 %
0%
Y74269
10%
0%
256
257
10%
Table 10. Our methods also lead to better performance on mod- ular multiplication. Wc train the models with the best data param- eters from $4. The model performance when trained using fdetsalts is 2% all different settings. T 1% Accuracy is percentage of predictions within 0.01g of right answer (sec $3).
3329
%
42899
15 %
0%
974269
15 %
0%
6. Beyond Modular Addition
Terms (N)
Mod (q)
1% Accuracy
Finally, consider whether our methods can aid learning in settings beyond modular addition, including other modular arithmetic tasks and broader functions.
97
100%
257
98%
3329
3%
32
97
100%
6.1. Other Modular Arithmetic Tasks
257
75%
3329
3%
We investigate whether methods enable ML models to learn other modular arithmetic functions beyond addition. We train models to predict outputs from these functions, using the same setup as before: encoder-only transformer model with modified data distribution, angular embedding, and custom loss.
64
97
100%
257
65%
3329
3%
6.2. Synthetic Tasks
Asymmetric Functions. Doshi et al. (2024) conjectured that two-layer MIPs can only learn functions that can be represented as h(g (a1), ...). gN (ax)) and cannot extend beyond this class. We introduce a class of functions h Z Z.g outside the aforementioned class, where hi.k ΣΣ - at) + +ff to show that our approach helps models learn other modular arithmetic functions.
Finally, we explore whether the data distribution techniques alone aid transformer performance on well-studied synthetic tasks, including:
Copy task (Graves (2014)): given vector of size N where each element is sampled from vocabulary V. output an exact copy of the vector. A random guess is correct with probability V-N
For these problems, we add a positional embedding in the transformer since these functions depend on input sequence
Associative recall task (Graves (2014)): given N keys and N values sampled from two distinct vocabularies
8

---

Making Hard Problems Easier with Custom Data Distributions and Loss Regularization
Table Our methods also lead to better performance on the scalar product task. We train the models with the best data parameters from $4. The model performance when trained using Jdefauls is 2% on all different settings. = 1% Accuracy is percentage of predictions within 0.01g right answer (see $3).
Table 12. Training data with more diverse problem lengths yields better accuracy results across different tasks. Partic- ularly on the Associative recall and Parity, the model fails to learn with fdefault % Accuracy is percentage of predictions exactly correct.
# Terms
Mod
1% Accuracy
% Accuracy
(N)
(g)
2
97
100%
Task
max length
fdctault
Jnnr
funi
257
100%
Copy
32
100.0% 0%
100,0% %
100.0%
3329
3%
64
100.0%
100,0 %
100.0%
128
94.3%
100,0 %
100.0%
4
100%
256
81.4%
98.1 %
97.4%
257
30%
3329
3%
Associative recall
8
32.5%
100.0 %
100.0
16
6.6%
100.0 %
100.0 %
8
97
78%
32
3.4%
100.0 %
100.0 %
257
2%
64
1.8%
100.0 %
1.8%
3329
3%
Parity
50.3%
100.0%
100.0%
50.6%
99.8%
100.0%
50.0%
%
50.2%
Vx and Vv, retrieve the correct value of one key. A random guess is correct with probability N-1
256
50.2%
99.4% %
50,2%
Selective copy
32
100.0%
100.0%
100.0%
64
100.0%
100.0%
100.0%
Parity task (Abbe et al. (2023)): given binary vector of size N, output the parity of the vector. A random guess is correct with probability 50%.
128
83.4%
100.0
100.0%
256
57.2%
100.0
99.3%
Selective copy (Gu & Dao): given a vector of size N where T 16 non-zero elements are sampled from vocabulary V and N T elements are equal to zero, output copy of the vector, discarding all elements equal to O.
significant improvement over prior work. We also demon- strate that these techniques can be applied to the Learning with Errors (LWE) problem in cryptography to train ML models that recover 2x harder secrets than prior work.
For these different generative tasks, we trained a decoder- only model with rotary positional embedding ((Black al., 2022)), 4 layers, and hidden dimension of 256 on batches of 250 examples for 10M total samples. Similar Sec- tion 3.1, sample each problem length from a distribution to vary between and max_ mex_le,tth instead of fixing it to be always max_length When we use fdefault we sample problems with length max_length. During evaluation, we calculate the accuracy on problems with length max Jength only. We do not use the custom loss during training, as it was designed for modular arithmetic.
Furthermore, findings provide valuable insights into how models learn modular arithmetic, which can inform the development of more effective ML algorithms for wide range of tasks. By applying our methods to other well-studied problems, we demonstrated their potential to generalize and improve learning outcomes.
Several interesting directions remain for future work. As the number of terms N increases, our models have slightly higher MSE and lower accuracy. This motivates future work to understand how to improve performance as N scales. An- other avenue is more exploration into transferring our tech- niques to other settings, such as ML-enabled cryptanalysis. While our method achieves success on q used in real cryp- tosystems and N close to real-world use cases (N 512 used in practice (Avanzi et al., 2021)) on Learning without Errors, more experimentation is needed to transfer modular addition knowledge to Learning with Errors with large N and Possible approaches include pretraining on general modular addition and fine-tuning on specific settings, but future research should consider creative approaches.
As demonstrated in Table 12, incorporating mix of easier and harder data can significantly enhance model perfor- mance. This is particularly evident in the Associative recall and Parity, where models trained with Sdefault are essentially making predictions at random.
7. Conclusion
This work introduces two key changes to the training pro- cess to help ML models learn modular addition. These techniques- caahiiuus--vrrying the diversity of training data and intro- ducing a regularized loss function centttionennbll ML models to add hundreds of elements mod a large q with high accuracy,

---

Making Hard Problems Easier with Custom Data Distributions and Loss Regularization
Acknowledgments
Charton, F. Learning the greatest common divisor: explain- ing transformer predictions. In ICLR, 2024.
The authors wish to thank Zeyuan Allen-Zhu for initially proposing the sampling distribution strategy. We also thank François Charton for discussions on the data repetition effect and Mohamed Malhou and Andrey Gromov for helpful conversations arithmetic tasks.
Charton, and Kempe, J. Emergent properties with re- peated examples. In NeurIPS 2024 Workshop on Scien- tific Methods for Understanding Deep Learning.
Charton, F., Hayat, A., and Lample, G. Learning advanced mathematical computations from examples. Proc. of ICLR, 2021.
Impact Statement
This paper presents work whose goal advance the field of Machine Learning. There are many potential societal consequences of our work, none which feel must be specifically highlighted here.
Chen, L., Moody, D., Liu, Y.K.., et al. PQC Stan- dardization Process: Announcing Four Candidates be Standardized, Plus Fourth Round Candi- dates. US Department of Commerce, NIST 2022. https://care. .ist wow.llts..o///ee//222222 pag-ccaddddtess-to-ee--ttaae rondd-
References
Abbe, E., Cornacchia, E.,, and Lotfi, A. Provable advantage of curriculum learning on parity targets with mixed inputs. Proc. of NeurlPS, 2023.
Doshi, D., He, T., Das, A., and Gromov, A. Grokking modular polynomials. arXiv preprini arXiv:2406.03495. 2024.
Agarwal, V., Aditya, S., and Goyal, N. Analyzing the nuances of transformers' polynomial simplification abili- ties, 2021. URL 14095.
Graves, A. Neural turing machines. arXiv preprint arXiv:1410.5401, 2014.
2204.
Gromov, A. Grokking modular arithmetic. arXiv preprint arXiv:2301.02677 2023.
Alfarano, A., Charton, F., and Hayat, A. Global іyаpunov functions: a long-standing open problem mathematics, with symbolic transformers. Proc. of NeurIPS, 2024.
Gu, A. and Dao, T. Mamba: Linear-time sequence model- ing with selective state spaces. First Conference on Language Modeling.
Zhu, Z. Physics of Language Models: Part 4.1, Ar- chitecture Design and the Magic of Canon Layers. SSRN Electronic Journal, May 2025, https:/ /ts//sssn.////////// com/ abst rast=5533222
Jelassi, S., d' Ascoli, S., Domingo-Enrich, C., Wu, Y., Li, Y., and Charton, Length generalization in arithmetic transformers. arXiv preprint arXiv:2306. 15400, 2023.
Avanzi, R., Bos, J., Ducas, L., Kiltz, E., Lepoint, T., Lyuba- shevsky, V., Schanck, J. M., Schwabe, P., Seiler, G., and Stchlé, D. CRYSTALS Krrke-r (version 3.02) Submis- sion to round 3 of the NIST post-quantum project. 2021. Available at ht ttttps:: httptlp:///qq-exe// /ptpcrrystall..............
Kera, H., Ishihara, Y., Kambe, Y., Vaccon, T. and Yokoyama, K. Learning to compute gröbner bases. Ad- vances Neural Information Processing Systems, 37: 33141-33187, 2024.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. In Proc. of ICLR, 2015.
Bengio, Y., Louradour, J., Collobert, R., and Weston, J. Curriculum learning. Proc. of ICML, 2009.
Lauter, K., Li, C. Y., Maughan, K., Newton, R., and Srivas- tava, M. Machine learning for modular multiplication. In Research Directions Number Theory- Women in Num- bers Springer, 2024.
Black, S., Biderman, S., Hallahan, E., Anthony, Q. Gao, L., Golding, L., He, H., Leahy, C., McDonell, K., Phang J., et al. -0t-neox-20b: An open-source autoregressive lan- guage model. In Workshop on Challenges Perspectives in Creating Large Language Models, 2022.
Li, C. Y., Sotáková, J., Wenger, E., Malhou, M., Garcelon, E., Charton, F,, and Lauter, K. Salsa Picante: A Machine Learning Attack on LWE with Binary Secrets. In Proc. of ACM CCS, 2023a.
Bossuat, J.PP., Cammarota, R., Cheon, J. H., Chillotti, I., ct al. Security guidelines for implementing homomorphic encryption. Cryptology ePrint Archive, 2024.
Li, C. Y., Wenger, E., Allen-Zhu, Z., Charton, F., and Lauter, K. E. SALSA VERDE: a machine learning attack on LWE with sparse small secrets. In Proc. of NeurIPS, 2023b.
Charton, Linear algebra with transformers. Transactions in Machine Learning Research, 2022.
10

---

Making Hard Problems Easier with Custom Data Distributions and Loss Regularization
Mohamadi, M. A., Z., Wu, L., and Sutherland, D. Why do you grok? theoretical analysis grokking modu- lar addition. In International Conference Machine Learning, Pp. 35934-35961 PMLR, 2024.
Nanda, N., Chan, L., Lieberum, T., Smith, J., and Stein- hardt, J. Progress measures for grokking via mechanistic interpretabillity. In Proc. of ICLR, 2023.
Palamas, Investigating the ability of neural net- works to learn simple modular arithmetic. 2017. https: tttts::/pproecc-areee archive. inff..f... ed. msc/20172390/mms hatpa/900mmcpppro__ pdf.
Shalev-Shwartz, S., Shamir, O., and Shammah, S. Failures of gradient-based deep learning. In International Con- ference on Machine Learning, pp. 3067-3075. PMLR, 2017.
Stevens, S., Wenger, E., Li, C. Y., Nolte, N., Saxena, E., Charton, F., and Lauter, K. Salsa fresca: Angular em- beddings and pre-training for ml attacks on learning with errors. Cryptology ePrint Archive, 2024.
Wenger, E., Chen, M., Charton, F., and Lauter, K. Salsa: Attacking lattice cryptography with transformers. In Proc. of NeurlPS, 2022.
Wenger, E.,, Saxena, E., Malhou, M., Thieu, E., and Lauter, K. Benchmarking attacks on learning with errors. 2025 IEEE Symposium on Security and Privacy (SP), pp. 58-58, IEEE Computer Society, 2024.

---

Making Hard Problems Easier with Custom Data Distributions and Loss Regularization
A. Additional results
We report additional results using our approach Table and Table 14 for different non-prime and in Table 15 for different N that are not powers of 2. We see similar trends as Table 2.
Table 13. Our approach is also successful for non-prime = 2 and N {16, 32, 64, 128) elements. All metrics are computed on held out set. MSE mean squared error, r 0.5% Accuracy is percentage predictions within 0.005q right answer, and = 1% Accuracy is percentage of predictions within 0.01g of right answer (see $3 for details).
Terms (N)
Mod (g)
MSE
0.5% Accuracy
7 1% Accuracy
16
2
0.00
99.8%
100.0%
32
2
0.00
99.6%
100.0%
64
2
98.1%
99.5%
128
2
0.02
95.7%
98.4%
Table 14, Our approach also works for different non-prime q and N E {16,32 32, 64, 128}. We train the models with the best data parameters from $4, All metrics are computed on a held out test set. т 1% Accuracy is percentage of predictions within 0.01q of right answer (see $3 for details).
Table 15. Our approach also works with powers of and prime q. We train the models with the best data parameters from $4. All metrics test set. 7 = 1% Accuracy is percentage of predictions within 0.01g of right answer (see $3 for details).
computed held out
# Terms
Mod (4)
1% Accuracy
(N)
16
1728
100%
Terms (N)
Mod (g)
1% Accuracy
100000
100%
1048576
100%
20
257
100%
10000001
100%
3329
100%
32
1728
100%
42899
100%
100000
100%
974269
100%
1048576
100%
257
99.7%
10000001
100%
3329
99.6%
64
1728
99.5%
42899
99.7%
100000
99.3%
974269
99.6%
1048576
99.4%
101
257
98.6%
10000001
99.8%
3329
98.8%
128
1728
98.0%
42899
98.9%
100000
98.2%
974269
98.5%
1048576
98.1%
10000001
98.8%
We also see that our method not fully robust to high N, but perhaps a longer training time or larger model is needed for higher N. We experiment both with longer training time and a larger model and see that both help with performance on high N. We present these results in Tables 16, 17, and 18.
B. Sample Efficiency Comparison
We investigate why our method succeeds and find that our sampling technique allows for linear sample complexity, while fdefault needs an exponential sample complexity to tackle the problem. This helps to explain why our proposed sampling strategy so effective.
In Table 19, we measure the number samples needed to get < 0.005 loss and 90% test accuracy.
12

---

Making Hard Problems Easier with Custom Data Distributions and Loss Regularization
Table 16. With our approach, performance declines
Table Training 4x longer improves performance on high N. All metrics are computed on a held out test set. MSE is mean squared error and 1%
with higher N. All metrics are computed on held out test set. MSE is mean squared enror Accuracy is percentage of predictions within 0.01g of right answer (see $3 for details).
1%
Accuracy is percentage of predictions within 0.01g of right answer (see $3 details).
Terms (N)
Mod
MSE Loss
1% Accuracy
# Terms
Mod (e)
MSE
1% Accuracy
(g)
(N)
Loss
256
257
0.15
90.4%
256
257
0.08
94.8%
3329
0.14
92.7%
3329
0.08
95.1%
256
42899
0.18
91.2%
256
42899
0.09
95.0%
256
974269
0.17
90.6%
256
974269
0.10
94.5%
Table 18. A 4x larger model also improves perfor- mance on high N. All metrics are computed on a held test set. MSE is mean squared error and 1% Accuracy is percentage of predictions within 0.01g of right answer (see $3 details).
Terms (N)
Mod
MSE
1% Accuracy
(g)
Loss
256
257
0.07
96.2%
256
3329
0.07
96.7%
256
42899
0.08
96.1%
256
974269
0.09
95.8%
Table Number samples needed get < 0.005 and 90% test accuracy with g = 3329.
N
Sdeiault
Jinv sqrt. (with best Saetauls setting)
finv sqrt (with our best setting)
4.5M
4.1M
0.6M
9
7.1M
1.9M
0,45M
12
12.85M
2.6M
0.95M
15
51.1M
8.15M
1.3M
18
Never
9.35M
1.75M
C. Polynomial Sum Tasks
We additionally run two symbolic polynomial tasks mod qF (a) sum N polynomials with degree < K and (b) sum 2 polynomials with degree < K.
For both tasks, we fix = 3329 and we encode the polynomial as am, Om-1,.... , an separate polynomials with <SEP> token, and ask the model to predict the coefficients of the polynomial sum. We add RoPE positional embedding. We report the results comparing finv vs faefault.
For all experiments, model can quickly predict the degree of the polynomial sum, i.e. the number of tokens to output before the <EOS> token. Therefore the real difference between the two strategies is predicting the right coefficients. We say that the model found the solution if the maximum error across all coefficients is less than 0.01g.
Task (a) is closely related to our original modular addition task (with K 0), and we see the success of the distribution fuuv_eurt Results are reported in Table 20.
Task (b) can be decomposed into finding the right index to attend to and summing two numbers mod q (which can be completely memorized with O(q²) even without angular embedding). We show the results on Table 21.
D. XE as MSE-like loss
When selecting an appropriate loss function for our model, we had to choose between a MSE-like loss (regression) or a
13

---

X Loss Making Hard Problems Easier with Custom Data Distributions and Loss Regularization
Table 20. On summing polynomials with degree
S K task, the Juv_aqrt distribution shows good per- formance. The standard distribution faetautt fails to learn the task.
Table 21. On summing two polynomial with degree
I K task, the fuv distribution shows spectac- ular performance, while fdetault fails to learn the task.
N K
Correct %
Correct % (fdefault)
(fav_sqrt)
16
1
99.5%
0%
Correct % Unnv_arr
Correct (fdetauls)
4
99.2%
0%
16
99.3%
0%
99.5%
0%
64
98.8%
0%
128
99.2%
0%
98.7%
0%
16
98.5%
0%
cross-entropy-like (XE) loss (classification). The choice particularly delicate for arithmetic problems, especially if the model decodes token from vocabulary (with XE) or angular representation (with MSE). Our claim is that, in particular case, MSE is actually superior compared to XE. In particular let x E Z, be point, and let the vector E Z where all components are zeros except for the i-th component, which is 1. Let g Z needs to learn (in particular g(x) 1 (mod q) for the modular arithmetic task and g(z) secret E Z is binary unknown vector, for the LWE task). The finite difference in the i-th direction can be defined as
Zg be the function the model (mod q). where the
(g)()) = g( + e,) 9(x).
It's easy to show that for these tasks, max [r)( < 1. This equivalently means that the difference between the function evaluated on closely related inputs remains relatively small. In this case, observed that well-trained model using XE loss can correctly approximate the aforementioned g(x).
Table 22. Models perform better when trained with MSE-like loss. Models trained on the best settings identified in $4 and evaluated on the same test set all. т = 0.5% Accuracy is percentage of predictions within 0.005g of right answer (see $3)
0.5% Accuracy
# Terms
Mod
Custom Loss
MSE Loss
(N)
(g)
16
257
99.8%
99.9%
99.5%
3329
99.7%
99.7%
99.2%
42899
99.7%
99.6%
98.1% 18.
974269
99.7%
99.7%
96.4%
32
257
99.5%
99.6%
98.5%
3329
99.4%
99.2%
96.8%
42899
99.4%
99.3%
94.3%
974269
99.5%
98.9%
91.0%
After inspecting the cross- cross-entropy output, we noticed that XE loss learned bell-shape curve centered at the correct solution with a tiny variance (in the order of 0.01g) which effectively forces the XE loss into a MSE-like loss. This suggests us to prefer a MSE loss and decode the output as regression and as Table 22 shows, we achieve best results using MSE-like loss.
E. Visualizing model predictions in LWE setting
We analyze the model's output predictions after 100 steps of the training to validate that, when the task is particularly hard, MSE loss cannot be used without the regularization term. To do this, we pass input sequences to the model and extract their outputs. We plot the outputs, coloring them based on y = x (mod g) of the input sequence æ. Figure presents this analysis for three models trained with N = 64, q = 257 and different settings for the hamming weight (number of nonzero elements in the secret) and loss: (a) hamming weight 3 with MSE Loss; (b) hamming weight with MSE Loss; (c) and hamming weight 6 with Custom Loss.
14

---

Making Hard Problems Easier with Custom Data Distributions and Loss Regularization
12
12-
0005
-11
-15
-12
-200 -035 -850
-125
023
135
850
1.0
-30097--00022- EIK0S 80000 8.00Z oveca 680055 830000
-2.0 -0.5
-050
-023
025 0.00
0.55
1.0
(a) hamming weight 3, MSE Loss
(b) hamming weight 6, MSE Loss
(c) hamming weight 6, Custom Loss
Figure 3. Final model output for different hamming weights and losses. Plots show the model's final outpur. Points with the same color have the same output y = x mod (i.c. they should be close together in representation).
As Figure shows, model trained with MSE loss successfully Iearn the problem when hamming weight is 3. However, for the setting without the regularization term and hamming weight 6, the model fails to learn, and the final output is visually meaningless, Additionally, looking at the magnitude, we empirically verify that predictions for a single batch are tied together and very close to the origin. Once we add the regularization term, the model can learn the task with hamming weight 6. This implies that for small hamming weight, the custom loss not as important, likely because the problem is simpler, but for larger hamming weight, the custom loss enables learning.
F. Curriculum Learning Hyperparameter Analysis
We provide the different configurations we tried to get a fair comparison between curriculum learning and our sampling strategy. In bold, we indicate the best setting.
When we ran the CL baselines, we modified three things:
1. Thresholds:
(a) Ti either 1% or 3% or 10% of the training
(b) Ti when train_loss eps, where we chose eps 1e-2, le-3
2. Data mix: we divide the training dataset X into two disjoint sets X1 and X2 where X, is the subset of such that each instance contains at least half of the elements equal to zeros.
(a) Train the model using Xı up Ti, then X2 until the end
(b) Train the model using X; up to Ti, then the entire X until the end
3. Learning rate and weight decay:
(a) We experimented with choices of learning rate (le-5, 3e-5, le-4) and choices of weight decay (0.03, 0.1, 0.3)
15


Reviews:
[{"review_id": "2JyhXxkB3e", "reviewer": "Reviewer_ZgiU", "summary": "The paper addresses the challenge machine learning models face in learning modular arithmetic, specifically in the context of the Learning with Errors (LWE) problem. It proposes two techniques: (i) using a designed data distribution that mixes sparse and dense modular arithmetic instances, and (ii) introducing a custom loss function with angular embedding and regularization to discourage convergence to trivial local minima. Experimental evaluation demonstrates an increase in performance of ML models on modular arithmetic task.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your thoughtful feedback. \n\n**Re: limited generalization experiments:** Per your suggestion, we conducted additional experiments with more $N$ and $q$ values. These results are presented in the response to reviewer Qesb (“additional experiments on challenging settings”).\n\nWe also conducted experiments on additional arithmetic tasks (product of n numbers, scalar product of two vectors, polynomial sum) and synthetic tasks (multi hop question answering and selective copy). We present the results for the synthetic tasks below and present the other results in our response to reviewer Eby5 (“other tasks beyond addition”).\n\n*Multi hop Question Answering:* We synthetically implement the [associative recall](https://arxiv.org/pdf/2212.14052) with multi hop. Given $N$ pairs $(a_i, b_i)$ which represent a random permutation from $\\\\{1, 2, …, N\\\\}$ to $\\\\{1, 2, …, N\\\\}$, we want to find the 2-th successor, i.e. $\\\\sigma(\\\\sigma(x))$. \n\n| # max_length | layers | $f_{default}$ | $f_{inv\\\\_sqrt}$ | $f_{uni}$ |\n| --- | --- | --- | --- | --- |\n| 16 | 4 | 7% | 100% | 100% |\n| 32 | 8 | 3% | 96% | 99% |\n| 64 | 12 | 2% | 93% | 94% |\n| 128 | 12 | 1% | 91% | 90% |\n\n*Selective copy:* Given a vector of size $N$ where each element is sampled from vocabulary $V$, output a selective copy of the vector (all tokens different from the <JUNK> token). This task was introduced by [Mamba paper](https://arxiv.org/pdf/2312.00752). \n\n| # max_length | $f_{default}$ | $f_{inv\\\\_sqrt}$ | $f_{uni}$ |\n| --- | --- | --- | --- |\n| 32 | 100% | 100% | 100% |\n| 64 | 100% | 100% | 100% |\n| 128 | 83% | 100% | 100% |\n| 256 | 57% | 100% | 99% |\n\nWe find that our method is robust to these additional tasks. Specifically, we vary the distribution of the input length with our custom distributions, and we find that this still leads to performance improvements on these additional tasks. We are also happy to experiment with any other tasks the reviewers suggest. \n\nWe will include these additional results and analysis in the revised version. \n\n**Re: hyperparameter sensitivity:** We will include more details on the hyperparameter sensitivity analysis in the revised version. While modifying certain parameters in the curriculum can slightly improve performance, the CL approach is much more involved and requires more tuning to the specific task. Our approach is simpler and provides consistent improvement over sampling from the default distribution. We provide the specific parameters here for your reference. \n\n$X_1$ = at least half of the elements are zeros, $X_2$ = at maximum half of the elements are zeros. \\\nWhen we ran the CL baselines, we modified three things: \n1. data mix: \\\n            i) Using $X _1$ up to $T_1$, then $X_2$ until the end \\\n            ii) **Using $X_1$ up to $T_1$, then $X_1$ union $X_2$ until the end** \n2. Thresholds: \\\n            i) $T_1$ is either 1% or 3% or 10% of the training \\\n            ii) $T_1$ is when train_loss($X_1$) < eps, where we chose eps = {**1e-2**, 1e-3} \n3. lr and weight decay: \\\n            i) We experimented with 3 choices of lr (1e-5, **3e-5**, 1e-4) and 3 choices of weight decay (from 0.03, **0.1**, 0.3) \n\nIn table 4 of the paper we reported the best choice, which is in bold above. \n\n**Re: weaknesses**\n* We are happy to include more explanation on the observed improvements. We did investigate why our method succeeds and found that our sampling technique allows for a linear sample complexity, while $f_{default}$ needs an exponential sample complexity to tackle the problem. This helps to explain why our proposed sampling strategy is so effective. \n\nBelow, we measure the number of samples needed to get <0.005 loss and 90% test accuracy.\n\n| $N$ | $f_{default}$ | $f_{inv-sqrt}$ (with best $f_{default}$ setting) | $f_{inv\\\\_sqrt}$ (with our best setting) |\n| --- | --- | --- | --- |\n| 6   | 4.5M | 4.1M   | 0.6M  |\n| 9   | 7.1M | 1.9M | 0.45M |\n| 12  | 12.85M | 2.6M | 0.95M |\n| 15  | 51.1M | 8.15M  | 1.3M |\n| 18  | Never | 9.35M | 1.75M |\n\n* Experimenting on the full practical cryptographic setting is an important area for future work, and this work provides the foundation for improving performance on practical settings. \n* See above for response to generalization experiment limitations\n\n**Re: computational overhead:** The computational overhead of our approach is minimal compared to the baseline, as we still generate the same number of data samples and simply modify the distribution used for generation. Similarly, the loss regularization term does not introduce any additional cost (besides the negligible cost of calculating the term itself) as we have a standard training loop. We timed the difference between the custom loss and standard loss experiments and saw that there was a 0.04% difference. We will include this explanation in the revised version of the paper.", "claims_and_evidence": "Overall, the paper presents convincing empirical evidence supporting the claimed improvements. Experiments clearly demonstrate that custom data distributions and the novel loss function substantially enhance accuracy across various problem complexities. However, the claim regarding generalization to a other set of arithmetic and synthetic tasks, while promising, could benefit from further detailed experiments and additional baseline comparisons (for example check other N values).", "methods_and_evaluation_criteria": "The proposed methods (custom data distribution sampling and the custom regularization of the loss function) are well-motivated by observations from prior literature and empirical insights. The evaluation criteria, including mean squared error (MSE) and accuracy metrics, are standard for the problem context.", "theoretical_claims": "The paper does not propose new theoretical claims requiring formal proofs and builds on theoretical insights from previous works.", "experimental_designs_or_analyses": "The experimental evaluation is sound, and the experiments are comprehensive (different q, number of terms N). However, there is limited hyperparameter sensitivity analysis. Specifically, it needs a more detailed and explicit sensitivity analysis for the improvements over CL.", "relation_to_broader_scientific_literature": "I think the paper has sufficient novelty, the authors make two contributions which 1) Augmenting Training Data with Sparse Vectors. 2)  Loss Regularization to Avoid Model Collapse. Both of which in relation to prior work are new contributions.", "essential_references_not_discussed": "Mostly discussed, however this area is not my primary area, I may have missed some references.", "other_strengths_and_weaknesses": "Strengths:\n* The paper is well written and structured\n* The experimental evaluation is sound.\n* Demonstrated significant improvements over SOTA works.\n\nWeaknesses:\n* Lack of detailed theoretical explanation for observed improvements limits understanding of the underlying mechanisms.\n* Experiments do not fully address the practical cryptographic setting (i.e., noisy LWE instances).\n* Generalization experiments, while promising, are limited and do not fully substantiate claims of broader applicability.", "other_comments_or_suggestions": "* It would be beneficial to discuss the computational overhead introduced by these modifications compared to the baseline approach.", "questions_for_authors": "I do not have any specific questions for the authors.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MVyfbebrjZ", "reviewer": "Reviewer_Qesb", "summary": "The paper improves machine learning attack baselines on LWE by training models to do modular arithmetic better. It uses custom training data and a special loss function, allowing the model to sum up to 128 elements modulo q ≤ 974269. It also shows improvements on other tasks like copy, associative recall, and parity.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your thoughtful feedback. \n\n**Re: comparisons:** In the paper, we compare our methods to the standard training approach with regular loss and default data distribution for the arithmetic, synthetic and the cryptography tasks (see Tables 3, 6, 7, and 9). We also provide a comparison to curriculum learning for modular addition (Table 4), where we show that our method is more consistent at learning the task compared to the curriculum learning approach. The lattice estimator provides theoretical cost estimates (in bit-operations) for attacking certain instantiations of LWE-based cryptosystems. Since the proof-of-concept we present in this work does not attack actual LWE-based cryptosystems, comparing against the lattice estimator would not be a fair comparison. Future work expanding on this should implement a true LWE setup and can then compare against other LWE attack methods. We welcome any additional suggestions of comparisons we could make in this work. \n\n**Re: cost/efficiency comparison:** We will include an analysis of the cost/efficiency of our method compared to the baselines in the revised version. The computational overhead of our approach is minimal compared to the baseline, as we still generate the same number of data samples and simply modify the distribution used for generation. Similarly, the loss regularization term does not introduce any additional cost (besides the negligible cost of calculating the term itself) as we have a standard training loop. We timed the difference between the custom loss and standard loss experiments and saw that there was a 0.04% difference.\n\n**Re: additional experiments on challenging settings:** As mentioned in our response to reviewers ZgiU and Eby5, we conducted additional experiments on more challenging settings and other tasks, which we will include in the revised version. \n\nWe conducted additional experiments on more values of $N$ and $q$, including non powers of 2 and non primes, for modular addition. We find that our method is robust to increased $q$s and non-powers of 2 $N$s. \n\n**Results for Different Values of $N$ and $q$**\n| $N$ | $q$        | tau=1% acc |\n| --- | ---      | ---       |\n| 16 | 1728     | 100.0%    |\n| 16 | 100000   | 100.0%    |\n| 16 | 1048576  | 100.0%    |\n| 16 | 10000001 | 100.0%    |\n| 32 | 1728     | 100.0%    |\n| 32 | 100000   | 100.0%    |\n| 32 | 1048576  | 100.0%    |\n| 32 | 10000001 | 100.0%    |\n| 64 | 1728     | 99.5%     |\n| 64 | 100000   | 99.3%     |\n| 64 | 1048576  | 99.4%     |\n| 64 | 10000001 | 99.8%     |\n| 128 | 1728    | 98.0%     |\n| 128 | 100000  | 98.2%     |\n| 128 | 1048576 | 98.1%     |\n| 128 | 10000001| 98.8%     |\n\n**Results for Different Values of $N$ (non-powers of 2) and $q$**\n| $N$ | $q$        | tau=1% acc |\n| --- | ---      | ---       |\n| 20 | 257     | 100.0% |\n| 20 | 3329    | 100.0% |\n| 20 | 42899   | 100.0% |\n| 20 | 974269  | 100.0% |\n| 49 | 257     | 99.7%     |\n| 49 | 3329    | 99.6%     |\n| 49 | 42899   | 99.7%     |\n| 49 | 974269  | 99.6%     |\n| 101 | 257    | 98.6%     |\n| 101 | 3329   | 98.8%     |\n| 101 | 42899  | 98.9%     |\n| 101 | 974269 | 98.5%     |\n\nWe also see that our method is not fully robust to high $N$, but perhaps a longer training time or larger model is needed for higher $N$. \n\n| $N$ | $q$        | MSE loss | tau=1% acc |\n| --- | ---      | ---       | ---       |\n| 256 | 257 | 0.15  | 90.4% |\n| 256 | 3329 | 0.14 | 92.7%  |\n| 256 | 42899 | 0.18 | 91.2% |\n| 256 | 974269 | 0.17 | 90.6% |\n\n**1. What happens if you train 4x longer?**\n| $N$ | $q$        | MSE loss | tau=1% acc |\n| --- | ---      | ---       | ---       |\n| 256 | 257 | 0.08 | 94.8% |\n| 256 | 3329 | 0.08 | 95.1% |\n| 256 | 42899 | 0.09 | 95.0% |\n| 256 | 974269 | 0.10 | 94.5% |\n\n**2. What happens if your model is 4x larger (embed dim goes from 256 to 512)?**\n| $N$ | $q$        | MSE loss | tau=1% acc |\n| --- | ---      | ---       | ---       |\n| 256 | 257 | 0.07 | 96.2%  |\n| 256 | 3329 | 0.07 | 96.7%  |\n| 256 | 42899 | 0.08 | 96.1% |\n| 256 | 974269 | 0.09 | 95.8% |\n\nWe also conducted experiments on additional arithmetic tasks (product of n numbers and scalar product of two vectors) and synthetic tasks (multi hop question answering and selective copy). These results are presented in our response to reviewer Eby5 (“other tasks beyond addition”). We find that our method is robust to these additional tasks. Specifically, we vary the distribution of the input length with our custom distributions, and we find that this still leads to performance improvements on these additional tasks. We are also happy to experiment with any other tasks the reviewers suggest. \n\n**Re: embedding question:** Thanks for the question. We use the angular embedding from Stevens et al. (2024) in our work. Could you please clarify your question on the embedding? We want to address your question appropriately but didn’t quite understand it well enough to answer.", "claims_and_evidence": "The paper claims that its methods enable ML models to perform modular addition for up to 128 elements modulo q ≤ 974269—far exceeding prior limits of N ≤ 6 and q ≤ 1000. While the evidence presented supports these claims, it lacks comparisons with other released methods such as the lattice-estimator baseline (https://github.com/malb/lattice-estimator). Without such baselines, it is hard to fully assess the real gains and stability improvements.", "methods_and_evaluation_criteria": "The methods generally make sense by using distribution aligned with difficulty and better domain knowledge.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "The experimental design is generally sound, showcasing improved performance on LWE and other tasks. However, the training process consumes a large amount of data and computation, and a relative comparison with existing approaches is missing. It would strengthen the paper if the authors compared the amortized cost and practical efficiency against established methods. Moreover, while some settings are shown to be easier, it remains unclear whether the method has been evaluated under truly hard conditions. More experiments on challenging settings are needed to validate the robustness of the approach.", "relation_to_broader_scientific_literature": "The technique could help contribute a way for attacking LWE", "essential_references_not_discussed": "Unsure", "other_strengths_and_weaknesses": "The work is overall promising and sheds light on how modular arithmetic tasks can probe the stability of ML models. Addressing the missing baselines and providing deeper comparisons will further solidify the paper’s contributions.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "I am confused on how the encoded embedding can map to multiple coefficient? I think the space of embedding would be constrained to decode a large number of terms due to expressiveness?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Xw7Rafu5Kz", "reviewer": "Reviewer_Eby5", "summary": "This paper proposes a new training strategy and loss function for successful modular additions and other operations. The critical observation is the utility of sparse samples. Models learn modular addition better on sparse samples, and if sparse and dense samples are mixed, it learns from sparse samples and then generalizes to dense samples. Previous studies sample training samples uniformly, and this leads the sampling to concentrate on non-sparse elements. The authors proposed to sample numbers from a distribution in which large numbers are sampled less (and thus, more chances of sampling 0). Further, the new loss function includes an additional term that encourages the model's output to be on the unit circle. The idea is based on the observation that in hard settings, the model's output collapses to the origin. In the experiments, the proposed sparsity-aware sampling and the custom loss both are demonstrated to improve the accuracy significantly.\n\n**update after rebuttal.**\n\nI appreciate the author's full elaborations and answers to my concerns, which greatly deepened my understanding of their work. The explanation and additional results address my concerns so I suppose that this work is worth being presented in the main conference.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your thoughtful feedback. \n\n**Re: claims:** We will include revisions in the final paper to clarify the claims more generally to not assume addition. Thanks for pointing out the clarification on $0$ and $q-1$ in the modular field, we will update this in the final version. \n\n**Re: sparsity:** Regarding choosing to include $0$ as the element for sparsity, we actually found that including a random but fixed number instead of $0$ yields similar results. It is not the number itself but instead the shift in distribution that leads to performance improvements. The key is changing the KL divergence between the train and test sets, which is irrespective of the element chosen for sparsity.\n\nTo show this, we ran an additional experiment where we substitute the $0$ with an arbitrary integer $K$ and shift the distribution by sampling more $K$s compared to all remaining elements in $Z_p^N$. We used a random $K$ for each different $q$ to ensure $K$ is not a factor. We can also run with more $K$ to show that this trend persists. We found similar results in the multiplication case.\n\n| $N$ | $q$        | $K$       | tau=1% acc |\n| --- | ---      | ---     | ---       |\n| 32 | 257     | 160    | 100.0% |\n| 32 | 3329    | 3176   | 100.0% |\n| 32 | 42899   | 24606  | 100.0% |\n| 32 | 974269  | 79062  | 100.0% |\n| 64 | 257 | 160    | 99.2% |\n| 64 | 3329 | 3176   | 99.3% |\n| 64 | 42899 | 24606  | 99.4% |\n| 64 | 974269  | 79062  | 99.3% |\n| 128 | 257  | 160 | 97.9% |\n| 128 | 3329 | 3176 | 98.3% |\n| 128 | 42899 | 24606 | 97.9% |\n| 128 | 974269 | 79062 | 97.9% |\n\n**Re: other tasks beyond addition:** We also conducted additional experiments on many other tasks based on your feedback. We present results on the synthetic tasks in our response to reviewer ZgiU and the rest below.\n\n*Modular multiplication and scalar product:* The angular embedding is designed for addition, so we use standard token embedding in multiplication experiments and compare $f_{inv-sqrt}$ to $f_{default}$. We test both modular multiplication and the scalar product of two vectors mod $q$. For both tasks, the model with $f_{inv-sqrt}$ performs well for smaller $q$, but declines for larger $q$. The scalar product is more difficult due to it requiring both multiplications and additions. Still, $f_{default}$ performs around 0% (on acc tau=1%) acc for all settings in both tasks, so $f_{inv\\\\_sqrt}$ is still an improvement. \n\n**Modular Multiplication with $f_{inv\\\\_sqrt}$**\n| $N$  | $q$      | tau=1% acc |\n|----|--------|------------|\n| 4  | 97     | 100% |\n| 4  | 257    | 100% |\n| 4  | 3329   | 51% |\n| 8  | 97     | 100% |\n| 8  | 257    | 100% |\n| 8  | 3329   | 32% |\n| 16 | 97     | 100% |\n| 16 | 257    | 98% |\n| 16 | 3329   | 25% |\n| 32 | 97     | 100% |\n| 32 | 257    | 75% |\n| 32 | 3329   | 13% |\n| 64 | 97     | 100% |\n| 64 | 257    | 65% |\n| 64 | 3329   | 3% |\n\n**Scalar Product with $f_{inv\\\\_sqrt}$**\n| $N$  | $q$ | tau=1% acc |\n|----|--------|------------|\n| 2  | 97  | 100% |\n| 2  | 257 | 100% |\n| 2  | 3329 | 98%  |\n| 4  | 97  | 100% |\n| 4  | 257 | 92% |\n| 4  | 3329 | 83% |\n| 8  | 97 | 78% |\n| 8  | 257 | 38% |\n| 8  | 3329 | 15% |\n\n*Polynomial sum:* We also run on two symbolic polynomial tasks mod $q$: sum $N$ polynomials with degree $\\\\leq K$ and sum 2 polynomials with degree $\\\\leq K$.\n\nFor both tasks, we encode the polynomial as $a_m, a_{m-1}, …, a_0$, separate polynomials with <SEP> token, and ask the model to predict the coefficients of the polynomial sum. We add a RoPE positional embedding to help the model learn how to count. We fix $q=3329$. We report the results comparing $f_{inv\\\\_sqrt}$ vs $f_{default}$.\n\nFor all experiments, the model is quickly able to predict the degree of the polynomial sum (the number of tokens to output before the <EOS> token), so the real difference between the two strategies is predicting the right coefficients. We say that the model found the solution if the maximum error across all coefficients is less than $0.01q$.\n\nTask a. is closely related to our original task (taking $K=0$)\n| $N$ | $K$ | Correct % for $f_{inv-sqrt}$ | Correct % for $f_{default}$ |\n|----|----|---------------------------|-------------------------------|\n| 16 | 1  | 99.5% | 0% |\n| 16 | 4  | 99.2% | 0% |\n| 16 | 16 | 99.3% | 0% |\n| 64 | 1  | 98.8% | 0%|\n| 64 | 4  | 98.7% | 0% |\n| 64 | 16 | 98.5% | 0%|\n\nTask b.\n| $K$ | Correct % for $f_{inv\\\\_sqrt}$ | Correct % for $f_{default}$ |\n|----|---------------------------|-------------------------------|\n| 64 | 99.0% | 94.3% |\n| 128| 99.1% | 92.1% |\n\nThis task is a bit easier because it can be decomposed into find the right index to attend to (and RoPE is expert on this) + sum two numbers mod $q$ (which can be completely memorized with O($q^2$) even without angular embedding)\n\n**Re: references:** Thanks for sharing the references, we will update the related work section to include these references. We agree that investigating finite-field polynomial tasks is an interesting avenue for future work.", "claims_and_evidence": "The claims and evidence are generally reasonable and convincing, but I have a concern that the explanation and the method sometimes assume modular addition, and the authors should include more general discussions.\n\nFor example, [l.113] says\n> [l.113, left]  Because 0 and q-1 are \"close\" in a modular field, ...\nbut this needs several remarks. First, modular field does not equip any distance as it violates the triangle inequality. Thus, being \"close\" is not precise even with quotations. The two numbers 0 and q-1 might appear \"close\" because of the unit 1 in the addition, but the field also equips multiplication. Similarly, including many 0 makes problems sparser/easier for addition task, but not necessarily for others. For the multiplication task 1 should be the number that makes the task \"sparser.\"", "methods_and_evaluation_criteria": "The proposed method and evaluation criteria are reasonable. One of the concerns is that the explanation of the proposed method is biased toward modular addition and not general enough. The experiments also put a great focus on the modular addition, and the extended task (Table 8) is also taking additive formats. It would be better for authors to make the explanation more general, and verify them in the experiments.", "theoretical_claims": "No theoretical claims are made.", "experimental_designs_or_analyses": "The experiments carefully examine the claims with variations in the number of terms and modulus. The recent repetition technique is also taken into account, and beyond the modular addition task, LWE cryptanalysis and several standard benchmark tasks are included in the experiments. However, as mentioned in #Methods And Evaluation Criteria, the experiments are strongly based on modular addition. The experiments also include other tasks in the end, but Parity is also an additive task, and the other tasks are not very arithmetic. Including more arithmetic tasks and generalizing the proposed method will make this study even more impactful.", "relation_to_broader_scientific_literature": "This study has a strong contribution to the modular arithmetic learning. It is interesting to see if the proposed method works for symbolic computation that involves modular arithmetic. For example, symbolic regression over modular field, or polynomial sum, multiplication, reduction, with finite-field coefficients.", "essential_references_not_discussed": "The following paper can be added as the literature of ML applications to hard math problems in [l.37, right]．\n1. Learning to compute Gröbner bases, Hiroshi Kera, Yuki Ishihara, Yuta Kambe, Tristan Vaccon, Kazuhiro Yokoyama, NeurIPS'24\n\nThis also shows that the polynomial computations over finite-field is unsuccessful due to errors in coefficient predictions. There are also several studies handling easier polynomial tasks (not necessarily working on finite field), and as mentioned above, it should be an interesting future work to see if the proposed method also works successfully on these tasks on finite field.\n\n2. Do Transformers Understand Polynomial Simplification? Vishesh Agarwal, Somak Aditya, Navin Goyal, ICLR'21", "other_strengths_and_weaknesses": "Important srengths and weaknesses have been mentioned in other cells. This work is solid and has sufficient contributions to this topic. The proposed methods are simple and easy to adopted.", "other_comments_or_suggestions": "The paper is written well in general. As I commented above, the explanation should be carefully revised and make it clear whether it writes about addition or more general cases.", "questions_for_authors": "Please refer to the other cells. Particularly, I want a clear explanation of the proposed method beyond addition.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This work develops machine learning techniques that achieve enhanced performance for modular arithmetic. A key application is developing stronger ML attacks against the Learning with Errors (LWE) problem. The methods are evaluated experimentally and exhibit 
 increased performance over prior work. Overall, the reviewers agreed that this is a solid contribution.
Decision: Accept (poster)