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

We thank you for your thoughtful feedback.

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.

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.

To 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.

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.

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.

Modular Multiplication with $f_{inv\\_sqrt}$

Scalar Product with $f_{inv\\_sqrt}$

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

For 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}$.

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

Task a. is closely related to our original task (taking $K=0$)

Task b.

This 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)

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.

We thank you for your thoughtful feedback.

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”).

We 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”).

Multi hop Question Answering: We synthetically implement the associative recall 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))$.

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.

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.

We will include these additional results and analysis in the revised version.

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.

$X_1$ = at least half of the elements are zeros, $X_2$ = at maximum half of the elements are zeros.
When we ran the CL baselines, we modified three things:

1. data mix:
   i) Using $X _1$ up to $T_1$, then $X_2$ until the end
   ii) Using $X_1$ up to $T_1$, then $X_1$ union $X_2$ until the end
2. Thresholds:
   i) $T_1$ is either 1% or 3% or 10% of the training
   ii) $T_1$ is when train_loss($X_1$) < eps, where we chose eps = {1e-2, 1e-3}
3. lr and weight decay:
   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)

In table 4 of the paper we reported the best choice, which is in bold above.

Re: weaknesses

• 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.

Below, we measure the number of samples needed to get <0.005 loss and 90% test accuracy.

• 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.
• See above for response to generalization experiment limitations

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.

We thank you for your thoughtful feedback.

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.

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.

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.

We 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.

Results for Different Values of $N$ and $q$

Results for Different Values of $N$ (non-powers of 2) and $q$

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

1. What happens if you train 4x longer?

2. What happens if your model is 4x larger (embed dim goes from 256 to 512)?

We 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.

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.