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

We would like to thank the reviewer for their time and insightful comments about our work and address them below. We have included the reviewers' questions in bold for reference.

“Their method is essentially a different way of describing the already-known features in modular arithmetic [1,2]...Since the RFM method is also not new to this work [3], I am unable to identify the significant contribution of this paper...My understanding is that the circulant features are identical to the Fourier-features found in literature (up to a global transformation learned by kernel regression). Is this assertion correct?”

We note that the connection to block circulant features is new and has not been observed in prior work. Circulant embeddings are not directly equivalent to Fourier features in neural networks – for example, RFM does not have a notion of a neuron. Further, the feature transformations in neural networks are inherently non-linear, as they consist of a weight matrix composed with an activation function, while circulant matrices perform a strictly linear transformation.

There are many feature maps that can lead to generalization for modular arithmetic and it is, a priori, unclear which feature map set AGOP would recover, if any. We provide an example of another (low rank) solution for modular arithmetic that is recovered by neither RFM nor neural networks in Appendix L.

To the best of our knowledge, there is no other work that has shown a non-parametric model class can learn modular arithmetic tasks. Personally we have found the performance of RFM surprising as a simple kernel machine (the first iteration of RFM) does not perform better than a random guess, yet features improve through RFM iteration.

“The feature-matrix of a network that solves this task must also respect this symmetry. This symmetry constraint readily gives the random circular matrix as the feature-matrix, making it trivial in some sense”

It is not a priori clear why RFM would learn a symmetry constraint from data, especially given that a simple kernel machine does not perform better than a random guess. Further, even if the generalizing solution found by the model is invariant under this symmetry, it is not clear that the feature matrix learned by RFM satisfies this constraint.

“Continuous progress measures for this task also already exist in literature, and not new to this work [1,2]”

Thank you for your comment. The progress measures in those works rely on the specific structure of neural networks (neuron specialization) and do not apply to models other than neural networks. In contrast, our measures apply to both neural networks and kernel machines. We will clarify this in the revision.

"Do the authors understand why there is a difference in the progress measures between add/sub and mul/div tasks in Figure 5 (B)? Is it related to the fact that mul/sub operations have to deal with 0 separately?”

Regarding the difference in progress measures between add/sub and mul/div in Figure 5(B) we do not observe this difference in the trend for RFM. As both RFM and neural networks need to handle the 0 case separately, and RFM does not show this difference, we believe this could be an artifact of the neural network training.

“Can the analysis presented in the work be used to predict the sample complexity of modular arithmetic tasks?”

Our work provides a relatively simple model, RFM, that exhibits grokking with modular arithmetic. We hope that this model will enable an analysis of sample complexity for these tasks, though it is beyond the scope of our paper. We note that an analysis of sample complexity for neural networks trained by gradient descent has not yet been provided.

We once again thank the reviewer for their feedback, if this response addressed your concern please consider raising the score. We look forward to reading the reviewers comments on our rebuttal.

We would like to thank the reviewer for their time and insightful comments about our work and address them below. We have included the reviewers' questions in bold for reference.

“This work focuses on a particular task, providing evidence that there may be tasks in which skills can emerge purely as a result of feature learning and independent of loss, but the scope of the results extending to other task and the use case of this particular kernel machine setting is less clear.”

Like others, we study grokking with modular arithmetic as the task requires feature learning to achieve non-trivial accuracy for non-parametric learning procedures. Further, neural networks and RFM exhibit sharp transitions between trivial and non-trivial error throughout training iterations. In particular, our results are a counterexample to (widely argued) theory that emergence is a “mirage” due to the disconnect between continuous loss function and discrete accuracy. Hence, any explanation for emergence must account for our results, which demonstrate that emergence occurs both in accuracy and loss.

“There is still a gap between theory and experiment where there lacks a theoretical proof that grokking will occur, or identify what the mechanism is behind learning modular arithmetic from a finite set of samples.”

We note that no prior work establishes a theoretical explanation for how gradient descent induces grokking in real neural networks. Prior works argue that the solutions to modular arithmetic tasks with neural networks satisfy certain margin maximization conditions, while not demonstrating that gradient descent actually performs this maximization. In our work, we provide a proof that kernel machines trained with circulant features can learn the Fourier Multiplication Algorithm (Appendix G), the same generalizing solution found by neural networks.

“I’m not very familiar with the prior literature on RFMs and the use of AGOPs to explain feature learning for neural networks; I know there is context sprinkled throughout in section 2, but it might be useful to reframe Appendix A as a more consolidated review of its introduction and previous use cases.”

Thank you for the suggestion. We will expand in Appendix A on the introduction of RFM and prior demonstrations of feature learning including deep neural collapse, edge detection in convolutional neural networks, and low-rank matrix completion (to name a few).

“How does the grokking behavior trend with the training fraction?”

To address the way in which grokking relates to the training fraction we note that Figure 4 and Appendix Figure 6 suggest a scaling law for some values of p by plotting test accuracy vs. training fraction for RFM as well as kernels using random circulant features. To expand on this, we have generated similar plots for additional values of p and a larger set of training fraction percentages to more clearly show the (empirical) relationship between test accuracy and training fraction. In these new figures we have also provided expanded comparisons of RFM to neural networks trained with AdamW and weight decay as well as to kernels with random circulant features. The new figures can be seen here: https://imgur.com/a/1TxDY8J.

We will update the final version of our paper with additional discussion about this and include the new plots. We observe that neural networks have somewhat better sample complexity, nevertheless the nature of the transition appears very similar.

“The RFM and NN learn a similar mechanism for the tasks (i.e. implement the Fourier multiplication algorithm)-- in general, to what extent are these learned solutions for tasks exhibiting emergence ‘universal’? Could there be tasks where learned mechanism differs across algorithm or architecture, highlighting a difference in inductive bias?”

There is some evidence that the solutions and features for a broader class of arithmetic tasks are similar. For instance he authors of [1] study pretrained LLMs and find that they utilize Fourier features in order to solve arithmetic problems. As we have shown the Fourier Multiplication Algorithm learned by transformers for grokking modular addition is also implemented by the quadratic kernel when trained with 100% of the data (which is given in our theoretical result and we have empirically validated this theorem). This is certainly suggestive of a broader “universal” solution/feature set in the case of arithmetic problems. We also note that these architectures are capable of learning multi-index models (see the discussion in Appendix A) , which is a different and very common setting. As discussed in the conclusion, unifying these settings remains an open challenge.

We once again thank the reviewer for the feedback.

[1] Tianyi Zhou and Deqing Fu and Vatsal Sharan and Robin Jia, Pre-trained Large Language Models Use Fourier Features to Compute Addition, 2024. arXiv preprint 2406.03445.

“I'm curious if the authors tried training RFMs on other synthetic tasks that have been popular for studying grokking”

We have not run experiments on other group structures but we expect that similar grokking will occur. For example it has also been observed that neural networks exhibit hidden progress when learning sparse parities (Barak et al., 2022).

We would like to thank the reviewer for their time and insightful comments about our work and address them below. We have included the reviewers' questions in bold for reference.

“More discussion on the importance of data is needed. Is the amount of data required for grokking comparable between this algorithm and other optimization methods like GD/Adam?”

We have run further experiments comparing test accuracy vs. training data fraction for RFM, kernels with random circulant features, and neural networks trained with AdamW. Figures for these can be seen here: https://imgur.com/a/1TxDY8J.

We will update the final version of our paper with additional discussion about this and include the new plots. We observe that neural networks have somewhat better sample complexity, nevertheless the nature of the transition appears very similar.

“The insight that feature learning is central to grokking is not completely new, for instance, the work in https://arxiv.org/abs/2310.06110 demonstrates the train loss of a neural network decreases much earlier than its test loss can arise due to a neural network transitioning from lazy training dynamics to a rich, feature learning regime.

Regarding the transition from lazy-to-rich dynamics, one of the arguments in our work is that the phenomenon we observe for grokking modular arithmetic in RFM can only be due to feature learning. In our setting the training loss is identically zero and we further observe in the first few iterations that feature learning takes place immediately. For instance, in Figure 2B we observe across addition, subtraction, multiplication, and division that both measures of AGOP alignment and circulant deviation steadily improve at a near-linear rate in the first iterations of RFM in spite of train loss being zero and test loss / accuracy equivalent to random guessing. Thus we do not see evidence for lazy-to-rich transition in RFM training.

“The mechanism behind the emergence of circular features during algorithm execution remains unclear, which seems crucial for understanding grokking, particularly from a dynamics perspective. From the curves I saw from the paper, I would expect some kind of lazy-to-rich transition happens here as well.

Could the authors provide intuition about how RFM leads to this grokking behavior from the perspective of their iterative algorithm?”

RFM succeeds in feature learning by recursively estimating the target function, then using this estimate to learn a better representation of the data. As the estimated features improve, the estimate of the target function improves, and vice versa. At a high level, we expect that grokking occurs because the error of the learned predictor is a sharp function of the feature quality. While feature quality increases gradually through training (indicated by our progress measures), it is only once the features are estimated to an error below a certain threshold that these features can be utilized to obtain better-than-trivial predictions.

“What are the implications of using a non-prime modulus?”

Whether or not $p$ is prime does not significantly affect our conclusions. The behaviors we observe for non-prime values are identical for addition and subtraction. For multiplication and division, however, we observe slightly different accuracy curves even though both tasks are able to generalize with RFM. The reason is that multiplication with non-prime modulus contains divisors of zero, e.g., 7, 11, 14…, for modulus p=77. For division, we must also restrict the training data to those whose denominators are invertible (do not share common divisors with the modulus).

We have run experiments on all four operations for p = 77 using a training fraction of 50% and include a figure of this setting here: https://imgur.com/a/80aPEJB.

We would be happy to include these additional plots in the Appendix of our final paper draft and add an expanded discussion in the main text pointing out this observation.

We once again thank the reviewer for their feedback, if this response addressed your concern please consider raising the score. We look forward to reading the reviewers comments on our rebuttal.