Understanding Grokking: Insights from Neural Network Robustness

W1: In Figure 4, how are the hyperparameters (learning rate, weight decay) selected? In the standard training of Figure 4b, the generalization speed may be much faster with proper tuning. It would make the generalization speed-up method more convincing by carefully tuning the hyperparams.

A: In Figure 4, the hyperparameters are all selected based on prior work [1], which we clearly stated in the preliminary section (section 3). And in Figure 4 (b), the standard training and perturb training share the same hyperparameters (like learning rate, weight decay) to make the comparison fair.

[1] OMNIGROK: GROKKING BEYOND ALGORITHMIC DATA

Q1: Where did you define a and b shown in Corollary 4.3? The only place a,b are defined is Sec 3, where (a,b) represents the input pair.

A: The a and b are hyperparameters and are just defined in the context of corollary 4.3, sorry for the abuse of notation.

Q2: Why do you set L=1/2, a=1925, b=500? Does Corollary 4.3's predicted accuracy change with learning rate and weight decay?

A: We set these hyperparameters so that the calculated predicted accuracy can match the empirically observed accuracy well. The predicted accuracy may change with learning rate and weight decay as these may affect the learning dynamics and thus affect the suitable parameters a and b from the defined equation of hyperparameters a and b.

Q3: [1] mentioned that there is almost no grokking phenomenon for the MNIST experiments under standard training, and they induce clear grokking by setting large initialization scale. But Figure 2a and 2c show clear grokking for MNIST. What initialization scale is used in those experiments?

A: Thank you for your question and careful reading. We adopt the same setting as [1] and setting large initialization scale.

We sincerely appreciate your valuable time devoted to reviewing our manuscript. We would like to gently remind you of the approaching deadline for the discussion phase. We have diligently addressed the issues you raised, providing detailed explanations. Given the importance of your feedback in refining and improving the work, we would greatly appreciate it if you could review the rebuttal at your earliest convenience.

W1: Limited technical contributions

A: The theorem is based on a lemma from literature, but the point is we are the first to derive a more suited theorem in the grokking literature that explains the role of weight norm decay. And the assumption we made are very mild and are very common among literature. What's more, the paper is not just about this theorem, we also provide many other contributions like the degrokking strategy and metrics and thus study the grokking phenomenon in depth.

W2: Limited justifications and interpretations for the observations

A: It is hard to provide a theorem of why perturb training will result in faster commutative law. But we may have an intuitive explanation as follows: Because of perturb training, when input pair (a, b)'s representation is perturbed a bit the prediction will still remain true. As the representation of (a, b) and (b, a) is much similar than the similarity of (a, b) and other (a', b'), then the prediction of pair (a, b) may become same for the prediction of (b, a), making the commutative law becomes valid. For the better correlation of metrics, mainly conclude this from an inspection of Figures. For example, in Figure 8, the metric changes sharply when test accuracy groks which we also marked in the picture and discussed in the respective section.

W3: Poor paper writing

A: Thank you for your reading, we may will the symbols more clearly to further improve readability. And the answers to your questions are as follows:

(a) The ReLU transformer is a transformer with a ReLU activation function. One layer is the standard meaning of one attention block and MLP block.

(b): Onehot(y) notation is a standard notation that represents a vector with only the y-th coordinate 1 and others 0.

(c): a and b are hyperparameters just defined in the context of corollary 4.3 and L is defined in theorem 4.2 as the lip constant. The use of abbreviations like train-acc is because that the name in full will make the picture ugly and not clear.

We sincerely appreciate your valuable time devoted to reviewing our manuscript. We would like to gently remind you of the approaching deadline for the discussion phase. We have diligently addressed the issues you raised, providing detailed explanations. Given the importance of your feedback in refining and improving the work, we would greatly appreciate it if you could review the rebuttal at your earliest convenience.

W1 & Q1: The study and experiments are based primarily on small-scale tasks and datasets, making it difficult to assess the generality of the theory in larger, more complex datasets or deeper networks. While the authors demonstrate a relationship between l2 weight decay and grokking, the applicability of this theory to more complex datasets or large-scale models (e.g., modern deep neural networks) is not discussed. For instance, does this phenomenon persist when the dataset and task complexity increase? Could similar l2 weight decay behavior be observed on more complex tasks?

A: The use of rather small-scale tasks and datasets in our study was influenced by the prevalence of such practices in previous literature, which all focus on small-scale settings. And we have already stated this limitation in our conclusion part.

W2 & Q2: Although the authors propose a perturbation strategy for degrokking, they do not clearly differentiate this approach from traditional data augmentation methods or discuss its unique advantages in accelerating generalization. The proposed degrokking strategy introduces perturbations to enhance robustness, thereby accelerating generalization. How does this method fundamentally differ from traditional data augmentation techniques? If the goal is to accelerate generalization, could conventional data augmentation methods (e.g., random noise) achieve similar effects? Additionally, the authors’ theory appears to support the idea that data augmentation enhances generalization—would it be helpful to further clarify the relationship between the two?

A: There may be a misunderstanding. The goal of presenting the perturbation-based strategy is to show that improving robustness may contribute to fast generalization, thus further validating our main theorem. This strategy is mainly used to validate the theory and we do not aim to show advantage to data augmentation methods. And we discussed in our paper that in some cases, data augmentation can be seen as a sort of perturbation and may explain why grokking does not happen in standard cases when augmentations are employed. The main relationship is that the perturb strategy can be seen as a special case of augmentation, but it is just one of the ways we consider that can easily enhance robustness and not the only way. And also we are the first to employ this strategy in the grokking literature.

Q3: Although the decay in the l2 weight norm is related to grokking, how does this phenomenon relate to the model’s learning dynamics and generalization ability? This relationship could be more complex in larger models. Can the authors’ theoretical explanation be extended to more diverse model architectures and tasks?

A: The decay of $l_2$ weight norm may correlate with the model's generalization in a somewhat complex manner [1]. Our paper's theoretical explanation does not rely on specific model architectures as long as it satisfies the basic settings the theorem requires. When tasks and datasets are more complex, it may violate the assumption of the theorem, for e.g. the training sets may not be fitted well in some complex cases.

[1] A PAC-BAYESIAN APPROACH TO SPECTRALLY-NORMALIZED MARGIN BOUNDS FOR NEURAL NETWORKS

We sincerely appreciate your valuable time devoted to reviewing our manuscript. We would like to gently remind you of the approaching deadline for the discussion phase. We have diligently addressed the issues you raised, providing detailed explanations. Given the importance of your feedback in refining and improving the work, we would greatly appreciate it if you could review the rebuttal at your earliest convenience.

W1: Figures' x-axis switch between logarithmic and linear, make it very hard to do a comparison among figures. Especially in Figure 3, curves are compressed by logarithmic x-axis on larger step numbers make it very difficult to understand the difference.

A: We sometimes switch between logarithmic and linear to make the plotted pattern more visible because in cases log axis may make the changing pattern not very clear. And in Figure 3, there may be a misunderstanding, the difference between the two curves is small is just what we needed as we want the predicted accuracy to match the real one.

W2: It is vague that how authors choose the transition point in all figures, for example, in Fig 1(a), the transition point for weight norm is selected at the beginning of weight norm start decrease, but in Fig 9(b), where PE has a similar curve, the transition point is selected much later to convey the point that PE is more aligned with Test Acc. The vague selection also happens in 8(b) 9(a) 10(b) 14(a) and etc. I suggest authors provide a consistent principle in choosing those points to reduce bias in result presenting.

A: There may be a misunderstanding. Figure 1 and other Figures are different. As we pointed in the paper, Figure 1's weight norm decays before the test accuracy changes, while the other proposed metrics correlate more closely with test accuracy change. Thus the transition point in Figure 1 is different as it does not align with the transition of test accuracy.

W3: I didn't really get the definition of PMI, in ln 394, $W_1$ is the first layer so the $W_2$ is the second layer? Is the PMI defined as the layer-wise mutual information on feature embedding gram matrix? also, in ln 399, It is not clear to my that $z^{(2)}_i$ is defined on whether entire $W$ or $W_2$?

A: $W_1$ is the first layer parameter and $W_2$ are the parameters of network except the first layer. PMI can be understood as mutual information between the output of first layer and final output layer. $z^{(2)}_i$ is defined on entire $W$.

W4: Results and experiments are mostly based on shallow network and no deep network results.

A: The adoption of small-scale tasks and datasets in our study follows the research paradigm in existing literature, as related works focus on small-scale settings. This limitation has been noted in our conclusion section.

W5: Combine 3 and 4, if I understand correctly PMI and PE will be extremely costly when network being larger and deeper, what will be the complexity of PMI and PE regarding the number of layers and representation dimensions?

A: The complexity will increase as the network goes deeper, as the computation relies on the output of the network, but the cost is not that big as the inference speed on GPU is quite fast. The cost reliance on dimension is basically on the calculation of Gram matrix, which is also quite efficient on GPU.

W6: Ln 429 states that in Fig. 9, the PE changes sharply but 9(a) doesn't, and 9(b) also changes earlier than test acc (see weakness 2).

A: We think the changing pattern can be seen as sort of evident when task accuracy drastically changes, thus making PE valuable. But of course, the criteria of whether it is timely or not is somewhat subjective and we offer this metric as it is theoretically nuanced and may have some value in our own opinion.

Q1: What is meaning of X-axis in figure 1 and 2? There are several more figures don't have x-axis labels.

A: It means the number of optimization steps.

Q2: What is Algo refer to in Fig 2?

A: As we stated in the caption of this Figure. we use Algo to indicate Modulo Addition dataset to save the space for image caption.

Q3: In Fig. 7, the logit distance has a significant scale difference(0-15000 vs. 0-150), does standard training also has the "chaotic behavior" if we zoom-in?

A: No, the case of standard training will have the distance nearly 0.

We sincerely appreciate your valuable time devoted to reviewing our manuscript. We would like to gently remind you of the approaching deadline for the discussion phase. We have diligently addressed the issues you raised, providing detailed explanations. Given the importance of your feedback in refining and improving the work, we would greatly appreciate it if you could review the rebuttal at your earliest convenience.