A simple and interpretable model of grokking modular arithmetic tasks

I thank the reviewer for positive comments and questions!

Strengths: The proposed setup is simple and interpretable. The analytic solutions could be valuable in analyzing the grokking phenomenon for the tasks of modular arithmetic.

Weaknesses:

The constructed analytic solutions for the tasks of modular arithmetic has potential but specifically, this manuscript does not produce too much new insights for understanding grokking. For instance, one can easily construct analytic solutions for learning k-sparse parity with two-layer ReLU networks, where we can reproduce the grokking phenomenon.

I am not familiar with the result. I suspect that either solution is not found by optimization (of which there are examples) or only makes sense for ReLU networks.

To further elaborate, the solution I presented works for arbitrary number of variables and over $\mathbb{Z}_p$ with any $p$. $k$-sparse parity is a particular case of this setting with $p=2$ and linear modular function with $k$ non-zero coefficients. I would be very interested to compare the two solutions if the reviewer can provide a reference where such solution for ReLU network is constructed.

The authors have empirically shows a peak around 0 for in the found solution, satisfying equality they propose. However, the presentation falls short of providing adequate evidence that the found weights have the periodic structure of the analytic solution. It is imperative that the authors supplement their work with further empirical evidence or a comprehensive theoretical analysis to elucidate how the weights progressively evolve toward the analytic solution during the training process.

There must have been a poor explanation on my part: The very notion of a phase is not meaningful unless the weights are periodic. Consequently, the very presence of the phases implies that the rows of the weight matrix are periodic. The phases are extracted by fitting the rows of the weight matrices with $\cos(\omega + \phi)$ and then extracting $\phi$.

I further demonstrate the periodicity in the Supplementary Material, Section C on dynamics (see general comments above and Fig.6 of Supplementary material), where inverse participation ratio (IPR) of the Fourier-transformed weights is defined and measured as a function of time. Growth of IPR indicates that weights become increasingly more periodic (i.e. more localized in Fourier space) over time.

Further investigation into the minimal data amount of grokking occurrences is warranted. Does the order of the minimal amount is $O(p^2)$ . If not, it necessitates a more suitable definition of the fraction as presented in equation (3).

The equation (3) defines an $\alpha$ which is simply a "dimensionless" (meaning independent of $p$) measure of how much data was used for training. It is a hyperparameter that is independent of $p$ by definition. The minimal amount of training data $\alpha_c$ does weakly depend on $p$ and slowly decreases as $p$ gets larger. This effect has been observed in [1] where linear functions of many variables over $Z_p$ are learnt and larger $p$ lead to "simpler" learning task.

References:

[1] Emily Wenger, Mingjie Chen, François Charton, Kristin Lauter, ``SALSA: Attacking Lattice Cryptography with Transformers''

Some mathematical oversight.

The definition in (6) might lead to the misconception that the weights form an matrix $Nxp^2$.

Fixed!

The factor $\frac{1}{N}$ is missing at the beginning of (11).

I thank the reviewer for their diligence! This choice was deliberate to lighten the equations. I state: ``(we drop the normalization factors)'' above Eq.(8) and restore the $\frac{1}{N}$ factor in Eq.(16).

The $\frac{1}{N}$ factors are not critical, unless we wish to have a finite large-$N$ limit, which is the normalization I opted for in the end.

I thank the reviewer for the positive comments and questions!

Strengths:

This paper is a timely and important contribution to the growing literature on grokking. It offers a class of problems with explicit solutions, so that grokking can be studied in great depth.

Weaknesses:

While simplicity and explicit solution are a strength, it also limits the scope of the paper in terms of covering the grokking phenomenon in general. Moreover, it is desirable to study the dynamics of the optimizers in reaching the exact solutions, but the paper did not make such an attempt.

I have addressed both of these questions in the general comments part of my response.

I added experiments and some discussion with other architectures: Transformers and pure attention layers (please see general comments for detailed description of those).

Discussion of the dynamics appears in the Supplementary material. In the case of modular addition, the network starts moving towards a periodic solution right away (as I show by measuring inverse participation ratio of Fourier-transformed weights over time) and this motion accelerates after a certain point, at this point generalization picks up.

More details about the dynamics are the subject of work in progress.

Questions:

The transition from memorization to generalization appears to be a continuous process of Occam's razor, i.e., gradually reducing the complexity of the model while maintaining the training error. Converging to periodic features is also of this nature. Is this the correct understanding of grokking?

One important motivation of the work was to flesh out the final solution in as much detail as possible to address this very question. I am not aware of any quantitative (meaning computable) measure of complexity that monotonically decreases over the course of training for all learnable modular functions and all hyperparameter and initialization choices. In particular, sparsity in Fourier space (which is only relevant for addition) or weight norms (which increase according to Fig.1 of my work) are not relevant measures. Consequently, I am skeptical of explanations based on any notion of complexity.

The best I can presently offer is: for sufficient amount of data ($\alpha > \alpha_c$) most low training loss minima are generalizing, while for insufficient amount of data $(\alpha < \alpha_c)$ most low training loss minima are memorizing. The critical amount of data $\alpha_c$ depends strongly on the modular operation we are trying to learn and weakly on the optimizer. How this transition happens is not clear and is an exciting research direction.

I thank the reviewer for the positive comments and questions!

Strengths: The analysis of grokking help to understand and dynamics of model training and how to achieve good generalization The theoretical results are applicable for general modular functions. Follow-up work could leverage on these results.

Weaknesses: Simple architecture and tasks (two layer MLP, modular arithmetic) could limit the applications and extensions of this work The given analytical solution does not help much in understanding how grokking happens as the latter occurs earlier than achieving 100$\%$ test accuracy.

I have addressed some of these concerns in the general comments above. I have added experiments with transformers and with pure attention layers, and described the results above.

Part of this question is addressed in the Supplementary Material on dynamics (also, see general comments above). In the case of modular addition the network starts moving towards a periodic solution right away (as I show by measuring inverse participation ratio of Fourier-transformed weights over time) and this motion accelerates after a certain point, at this point generalization picks up. The progress measure -- inverse participation ratio -- is motivated by the analytic solution: it measures periodicity of the weights.

More details about the dynamics are the subject of the work in progress.

Questions: How does the analytical solution help understanding grokking ? Neural networks are known to converge to local minima. I wonder if there are potentially other analytical solutions and why it seems that model training leads to the same solution.

In the present work I could only answer where the dynamics converges to. The questions why and how are extremely complicated even in the case of the simplest possible model. Detailed understanding of the dynamics is still work in progress and will appear in a separate paper.

There are memorizing solutions, however I am not aware of any other generalizing solutions.

I thank the reviewer for the positive comments and questions!

Strengths: Grokking is an interesting and exciting phenomenon that is worth careful study. The paper is technically sound. The presentation and organization is clear.

Weaknesses: To provide an analytical solution and interpretability of the model, this paper focuses on a very simple model (definitely not used in practice) and arithmetic function to be learned, which limits its impact on practical models currently in use, such as CNN, Transformer.

Some of the concerns are addressed in the general comments above.

I have addressed this comment by making further experiments with transformers and pure-attention networks. Initially, I chose not to include these architectures because grokking was first empirically observed in two-layer transformers, while the blogpost [https://www.neelnanda.io/blog/interlude-a-mechanistic-interpretability-analysis-of-grokking] empirically considered one-layer transformer.

My objective was to distill the simplest possible model of the phenomenon rather than focus on models that are more realistic, but cannot be understood in as much detail. Since the identical effects (jump in generalization as a function of time and amount of data) are present in all cases -- simple and complex --, it is reasonable to make as much progress as possible in the simplest case. (Such research philosophy has been very successful in theoretical physics so far).

If the model really learns the arithmetic function, it will be interesting to see whether the model generates accurate results for OOD data, e.g., training with the data from [0, 10], testing with the data from [1000, 1100].

One advantage of the analytic solution is that the reviewer does not need to take my word on it! The learnt feature maps are described by equations and are made from trigonometric functions. In the main text I prove analytically that these features perform modular operations including addition and show empirically that similar features are learnt during optimization.

The model learns modular arithmetic. That is, arithmetic over a finite field $\mathbb{Z}_p$ rather than over $\mathbb Z$. The modulus $p$ appears explicitly in the analytic solution and in the dimension of the input and output of the architecture. In the transformer case, $p$ will appear as vocab_size, unless a representation of integers with base $B<p$ is used (in which case ``tokenization scheme'' is determined by the pair $(B,p)$). If numbers larger than $p$ are presented to the model, it will round them down modulo $p$.

Meta learning many different operations in one network is beyond the scope of this paper.

Questions: "Instead, large width leads to redundant representations: Each frequency appears several times with different random phases ultimately leading to a better wave interference" Since they are identical frequencies, will combining them provide a more concise representation?

Unfortunately, combining the neurons with same frequencies is not possible because in addition to frequency the trigonometric functions contain a random phase, which is different for different neurons. To be concrete, neurons with $\cos(\omega + \phi_1)$ and $\cos(\omega + \phi_2)$ cannot be easily replaced by a single neuron. What is possible, however, is to simply keep erasing the ``redundant'' neurons. This will keep performance as measured by accuracy constant (up to some point), but will decrease the train and test loss.

Besides the amount of data, how is grokking affected by the training data distribution?

This questioned was mentioned as open in the paper: ``The value of $\alpha_c$ further depends on how training set is sampled from the entire dataset; the appropriate choice of the sampling method may thus improve the data efficiency.''

Unfortunately, I had very little progress in this direction. Empirically it appears that sampling the finite dataset with the uniform probability gives the best performance. We did some experiments along the lines of work [1], but no benefit was seen.

References:

[1] Mansheej Paul, Surya Ganguli, Gintare Karolina Dziugaite, ``Deep Learning on a Data Diet: Finding Important Examples Early in Training''

I thank the reviewer for the positive comments on my work and questions!

Strengths: The paper's presentation is clear and to the point, and the construction of the weights is succinctly explained. The experimental evidence is convincing. Mechanistic interpretability is also a highly interesting direction overall.

Weaknesses: Literature review is missing some recent work: There is a growing body of work on learning single-index and multi-index functions (see e.g., "Online stochastic gradient descent on non-convex losses from high-dimensional inference" by Ben Arous et al., and "SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics" by Abbe et al.) which shows similar grokking effects. It could be interesting to understand how these relate to the arithmetic grokking effect.

I was not aware of that line of work. Presently, it is not clear to me what the relation is. The point I make in my paper is that grokking is tightly connected to feature learning. The exact shape of the learning curve is not sufficient. I now cite this work in the new version.

More crucially, there was a paper called "Progress measures for grokking via mechanistic interpretability" which appeared online in Jan., 2023 and was published in ICLR 2023. This paper also seems to derive the Fourier-based solution to the grokking task. This seems unfortunate, because it seems that at the time that this paper was written either the authors are unaware of this other paper, or that this paper was written concurrently and has been to a good extent subsumed by that other paper. Could the authors comment on this? This is the main weakness in my mind.

My work appeared online before "Progress measures for grokking via mechanistic interpretability" paper and, consequently, I opted not to cite it in the present version. I have now added the citation with a comment on the timeline.

I would like to further emphasize that the (very interesting) paper "Progress measures for grokking via mechanistic interpretability" (2301.05217v1), which appeared after mine, does not contain the analytic solution I have described in my work. The tables with trigonometric identities are obtained empirically and serve as interpretation of the activations measured in experiment. The authors of 2301.05217 are aware of my work, but choose to not cite it.

The analysis only gives an Ansatz for the final solution of the weights, but does not explain why more/less data leads to finding it, and why there is a sharp jump in the algorithm's loss from not finding the Ansatz to finding the Ansatz. In other words, the paper only predicts the final weights but does not give an interpretation of what is driving the dynamics of the grokking process.

The current version of the paper concludes: ``Given the simplicity of our setup, the basic explanation for grokking must be quite banal: At some point in training (and assuming there is enough data), the only way to decrease the training loss is to start learning the right features.''. This is my current opinion on the dynamics: all low training loss minima generalize if there is enough data.

I do, however, completely agree with the reviewer: there is no clear picture for the fraction of the data and for what exactly changes when this fraction crosses its critical value. It is still work in progress.

Questions: What is meant by "Functions of the form f(n, m) + g(n, m) mod p are more difficult to grok: they require more epochs and larger $\alpha$"? Why do you need both f(n,m) and g(n,m) here?

Indeed, more difficult means more epochs and more data. Two examples are shown in Fig. 10 of the Supplementary Material. The form $f(n, m) + g(n, m) \quad \textrm{mod} \quad p$ is a slight abuse of notation meant to indicate ``functions of $n$ and $m$ that cannot be written as $F(f_1(n) + f_2(m))$''. To be concrete, $n^2 + m^2 + 2nm$ is not in this class, while $n^2 + m^2 + nm$ is.

Typos: "a lightening" -> "an enlightening"

This is indeed a typo, however I meant a more humble "a lightning"! Fixed now!

"we do not observer"

Fixed!