Emergence in non-neural models: grokking modular arithmetic via average gradient outer product

We thank the reviewers for their detailed feedback. We address their questions and comments below.

I felt that some details are still unclear and may need further investigation: What is the fundamental mechanism behind the delayed generalization observed here? In other words, why does generalization accuracy behave so differently from the continuous progress measure?

There is a sharp transition in feature alignment. Once the features are aligned beyond a certain threshold, the accuracy (and loss) both improve sharply. Thus the distinction is not between continuous or discrete measures but between measures of loss and measures of feature alignment. Why there is a discrepancy between our progress measures and accuracy / loss and formalizing the mechanism behind the delayed generalization is an important question. A simplified direction to understand this process theoretically would be to analyze how random circulant features enable generalization on modular arithmetic.

It appears that the jump in accuracy occurs when the continuous progress measure saturates. Is this consistent? Have you identified a specific threshold in the progress measures?

We have not observed a direct relationship between the jump in accuracy and saturation of progress measures. While it might look like it happens in Figure 2, it is hard to visually evaluate feature improvements as the alignment is close to 1 and the curve looks flat even as alignment continues to improve.

Generalization does not occur below a certain threshold of the training data fraction. What is the interplay between this threshold and your results?

Both training iterations and number of samples will improve feature quality. Thus, as generalization error is a sharp function of this quality, we see sharp transitions with respect to both.

Finally, since the paper focuses on the specific task of modular arithmetic, the scope of the results is naturally limited.

Like prior works (e.g. [1,2]), we focus on the setting of modular arithmetic because (1) these tasks clearly exhibit the sharp transition from trivial-to-perfect generalization, and (2) there exist analytic solutions to the task that we can compare the learned algorithms to. In particular, there is substantial evidence that the algorithm implemented by neural networks is the Fourier Multiplication Algorithm[2,3]. Moreover, it appears that non-parametric methods without the ability to learn features are unable to learn modular arithmetic, hence these tasks are a strong test-bed for the predictor power of feature learning through AGOP.

[1] Power et al., Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets, Mathematical Reasoning in General Artificial Intelligence Workshop, ICLR 2021.

[2] Nanda et al., Progress measures for grokking via mechanistic interpretability, ICLR, 2023.

[3] Zhong, et al. "The clock and the pizza: Two stories in mechanistic explanation of neural networks." NeurIPS 2023.

I believe that further investigation of the behaviour of the progress measures for different ratios could be beneficial. For example, I could be interesting to see the AGOP alignment (of the current iteration versus the final iteration) and the circulant deviation at several different ratios (below and above the critical threshold). Fig. 10 is also interesting in this context: the AGOP still changes quite a lot from fraction 40% to 45%, while the test accuracy is already 1 in both of them. It could be interesting to see this graph for more fractions (larger than 45) and also for the circulant deviation measure.

Please find a new figure here showing circulant deviation and AGOP alignment evolution over time for modular addition at training fractions 5%, 15% 25%, 35%, 45%, 55%, …, 95%: https://ibb.co/ytKK0p5. As the reviewer points out, we observe that the AGOP alignment can continue to increase after test accuracy is at 100%.

We provide another new plot showing that the circulant deviation at the final iteration rapidly decreases with training size, and is very close to 0 for training fractions larger than 55%: https://ibb.co/9Hp1F3yr.

Typos: Line 70: "and and test accuracy remain at the constant" Around Eq. (6), $\alpha$ probably should be $s$ (by the way, a brief justification for the choice $s=1/2$ would be nice).

Thank you for noticing this, we will fix this typo in our revision. The choice of s=½ is motivated by the case of two-layer linear networks where the NFA is exact with this exponent [1].

[1] Radhakrishnan, Belkin, Drusvyatskiy. “Linear Recursive Feature Machines provably recover low-rank matrices”, PNAS 2025.

We thank the reviewer for their in-depth feedback on our submission. We will address their comments and questions here.

Theorem 5.1 could come earlier in the paper and also be written formally as done in Appendix "well-known results for multiplicative group" may not be well known to ML audience. It would be better to state these group theory references formally in math symbols instead of words (and add refs).

Thank you for this feedback. We will include the Theorem formally in the main text in the updated manuscript, and more clearly reference Appendix E where we discuss the re-ordering.

Please use the standard definition of Jacobian without transpose in Def 2.1 and write G = Jf ^T Jf to avoid confusion.

Thank you for this suggestion, we will use the standard definition to avoid confusion.

Why is AGOP measured by vectorizing the matrices? Why not $\text{alignment}(A, B) = \frac{\text{Tr}(A^T B) }{| A | | B |}$. I believe this is the standard way to measure the alignment of matrices AND it would capture similarities of eigenvectors, unlike the vectorized alignment.

In fact $\text{Tr}(A^T B)$ equals the inner product of the vectorized A and B, and $||A||,||B||$ are equivalent to the norms of their vectorizations, thus our formulation is indeed a standard measure of matrix similarity.

We thank the reviewer for their positive comments and thorough feedback. We will answer their questions below.

While the paper is well-grounded in the literature on grokking and feature learning, it could benefit from more discussion of works on implicit bias in optimization and how these biases shape feature learning.

We would appreciate specific references that the reviewer feels are relevant. We would be happy to extend the discussion to include them.

On the downside, the experiments are limited to modular arithmetic tasks, so it is uncertain how well these insights will carry over to more complex, real world dynamics. Also, the paper focuses heavily on RFMs, which may limit how easily it can be generalized to other architectures.

Applying RFM is crucial for establishing that (1) grokking is not tied to neural networks, as it is a non-neural model, (2) it is not tied to gradient based optimization or training loss, and (3) can be induced solely through feature learning. Further, RFM allows us to isolate the ability of feature learning through AGOP to induce generalization in surprising settings. As mentioned in another response, there is also substantial evidence that the algorithm implemented by neural networks is the Fourier Multiplication Algorithm. RFM being able to recover the FMA (in our theoretical setting) and the features learned by neural networks empirically suggest that feature learning through AGOP is a useful lens to study neural networks. This is especially true as non-parametric methods without the AGOP appear unable to learn modular arithmetic tasks.

Have you tested your approach on any tasks beyond modular arithmetic to see if the delayed feature evolution and grokking phenomenon persist?

We have ongoing results where grokking occurs with RFM for other algorithmic tasks including sparse parities, the Chinese remainder theorem number representations, and certain group structures.

How sensitive are your progress measures to different kernel choices or initializations in RFMs?

The choice of kernel can affect the rate at which grokking occurs in RFM or whether it happens at all. Quadratic and Gaussian kernels are suited for modular arithmetic due to the target function being quadratic in Fourier space, while RFM with the Laplace kernel does not appear to generalize. In all of our experiments, we use the default initialization for RFM with the identity matrix.

Can you provide more insights into how the reordering using the discrete logarithm affects the observed feature structure?

Re-ordering by the discrete logarithm only affects the visualization of the features, as RFM is invariant to re-ordering of the input coordinates. We re-order coordinates for multiplication and division tasks in order to visualize the circulant structure, which is otherwise obscured. The features are circulant after re-ordering because modular multiplication/division (excluding zero element) are equivalent to addition/subtraction after transformation with the discrete logarithm. Note that it is necessary to know the structure of that re-ordering in order to invert it and recover the “underlying algorithm”.

We thank the reviewer for their detailed feedback. We address their questions and comments below.

First, the paper focuses exclusively on modular arithmetic -- but grokking is a more general phenomenon and it has been observed also in language models among others. This has to be clarified somehow: even though the paper helps us better understand the process of learning structured features in a very specific task that relates to modulo arithmetic, its claims and evidence cannot be generalized beyond such tasks.

Like prior works (e.g. [1,2]), we focus on the setting of modular arithmetic because (1) these tasks clearly exhibit the sharp transition from trivial-to-perfect generalization, and (2) there exist analytic solutions to the task that we can compare the learned algorithms to. In particular, there is substantial evidence that the algorithm implemented by neural networks is the Fourier Multiplication Algorithm [2,3]. Moreover, it appears that non-parametric methods without the ability to learn features are unable to learn modular arithmetic, hence these tasks are a strong test-bed for the predictor power of feature learning through AGOP.

[1] Power et al., Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets, Mathematical Reasoning in General Artificial Intelligence Workshop, ICLR 2021.

[2] Nanda et al., Progress measures for grokking via mechanistic interpretability, ICLR, 2023.

[3] Zhong, et al. "The clock and the pizza: Two stories in mechanistic explanation of neural networks." NeurIPS 2023.

The paper certainly contributes in the fundamental understanding of grokking -- but it does not explain some key characteristics of grokking (such as, what determines how many long it will take, say in epochs, until the model learns those structured features?)

While this question is beyond the scope of this paper, we do observe the rate at which features evolve through the AGOP alignment and circulant deviation progress measures. We find that while the features make progress linearly through iteration, generalization exhibits a sudden phase transition, suggesting error is a sharp function of feature quality specifically. We are exploring theoretical analyses to potentially identify the threshold when better features lead to improved generalization.

Is there an info theoretic way to show why circulant matrices emerge?

One possibility to understand why circulant matrices emerge on these tasks is through their effect on the Mahalanobis kernel. As circulant matrices are diagonalized by the DFT matrix, we can view inner products through the block-circulant matrix as applying the DFT to each of the one-hot encoded integer arguments, then re-weighting the frequency components by the (complex) eigenvalues of the circulant sub-blocks. We suspect the frequency reweighting induces linear separability of the DFT vectors, enabling generalization with kernel regression. We are currently exploring this direction.