Dichotomy of Early and Late Phase Implicit Biases Can Provably Induce Grokking

Q1: The paper comes across as dense and can be difficult to read. One reason is that a number of settings for theory and experiments are considered. Another is that the statements of many of the theoretical results are relatively complex, with several quantifications and subtle details, and often without much explanation.

A: Thanks for pointing out this issue! We have polished the writing and changed the organization in the revised version. Now we have included a lot of explanations for our assumptions and theorems.

Q2: The theoretical results about grokking for regression tasks, and in particular for overparameterised matrix completion, are not illustrated experimentally.

A: Thanks for the suggestion. In the revised version, we have added an experiment on matrix completion for the multiplication table; please see the “Empirical validation” paragraph of Sec. 3.3.3 for more details.

Q3: The conclusion of the paper (Section 5) does not seem to indicate any directions for future work.

A: Thanks for pointing this out. We have now added two possible future directions in the last section. First, our work only studies the training dynamics with large initialization and small weight decay, but these may not be the only source of the dichotomy of the implicit biases. Also, our work focuses on understanding the cause of grokking but does not study how to make neural nets generalize without so much delay in time. We leave it to future work to explore the other sources of grokking and practical methods to eliminate grokking.

Q4: In the first sentence in "Our Contributions", it is claimed that homogeneous networks allow the ReLU activation, but Assumption 3.1 assumes smoothness

A: Thanks for pointing this out. We have clarified this issue in the paragraph following Assumption 3.1. Indeed, the specific examples on sparse linear classification and matrix completion that we consider in this paper satisfy this assumption, but for non-smooth functions, we believe that our analysis can be generalized by invoking clarke differential, similar to previous works on homogeneous nets (Lyu and Li, 2020; Ji and Telgarsky, 2020).

Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks.In International Conference on Learning Representations, 2020.

Ziwei Ji and Matus Telgarsky. Directional convergence and alignment in deep learning. Advances in Neural Information Processing Systems, 33:17176–17186, 2020a.

Q5: Code was not submitted with the paper.

A: Thanks for pointing it out. We have included our code in the supplementary material.

Q6: What is the purpose of the section on sharpness reduction with label noise (Section 4.2), since it mostly consists of statements of results from Li et al. (2021)?

A: The main purpose of this section is to point out that sharpness reduction is another form of implicit bias that provably causes grokking. Although this section does not include original theoretical contributions, we still think it valuable to notice the connection between the grokking phenomenon and previous works.

Main Concerns

Q1: The core claim of "sharp transition in test accuracy" is not supported by any of the generalization analyses. Without such an analysis, the best one can conclude is that the models can fit the training data and reach perfect training accuracy eventually.

A: We thank the reviewer for raising the concern.

1. In our paper, we have instantiated our theorems on regression tasks on the matrix completion example in Sec. 3.3.3. We can prove that initially gradient flow can fit the observed entries but make no progress on the remaining ones, but eventually it can generalize to unobserved entries and recover the ground truth.
2. For classification tasks, we added generalization bounds in the updated version of our paper. Concretely, in the “empirical validation: grokking” part of Sec. 3.2.3, where we consider the sparse linear classification task, we added generalization bounds (proved in Section E), showing that the sample complexity required for achieving good generalization in the kernel regime is much higher than that of the rich regime.
3. In the general cases, our main results for both classification (Thm. 3.4 and 3.5) and regression (Thm. 3.9 and 3.10) settings highlight a sharp kernel-to-rich transition from $T=\frac{1-c}{\lambda}\log\alpha$ to $T=\frac{1+c}{\lambda}\log\alpha$. As the reviewer noticed, our results do not directly lead to generalization guarantees. However, for general deep learning, it is known to be very hard to provide any tight generalization bounds. In the literature, existing generalization bounds for DL are usually vacuous in real-world settings, and no generalization measures are known to be fully predictive of the actual generalization error in practice. Hence, the line of works on implicit bias usually uses the following strategy: analyze the training dynamics and show how the training method shapes a neural net’s properties that are known to be correlated with generalization, then exemplify the result in simple settings with generalization analysis. This is indeed the strategy we take.

Q2: Comparison with Moroshko et al. (2020)

A: While Moroshko et al. (2020) also characterize the implicit bias in the kernel and rich regimes, our results differ from theirs in two crucial aspects:

1. We consider running gradient flow with weight decay, while they do not use weight decay.
2. Our results highlight a sharp transition from kernel to rich regimes, while their paper does not imply such transition bounds.

Other Questions

Q3: Can the authors add more examples in text with datasets commonly used in ``Grokking'' literature?

A:

1. The grokking phenomenon was originally observed on modular arithmetic datasets. In the revised version of the paper, we conducted experiments on modular addition and showed that enlarging the initialization scale or reducing the weight decay indeed delays the sharp transition in test accuracy. This is consistent with our theory, which predicts the grokking time to be scaled with $\frac{1}{\lambda}$ and $\log(\text{init scale})$,
2. Beyond algorithmic datasets, Liu et al. (2023) observed that grokking can also be induced by using a large initialization and a small but nonzero weight decay. This is true for many tasks, including image classification on MNIST and sentiment analysis on IMDB. This can be seen as an important justification of our theory.

Ziming Liu, Eric J Michaud, and Max Tegmark. Omnigrok: Grokking beyond algorithmic data. In The Eleventh International Conference on Learning Representations, 2023.

Q4: One point that the paper makes is about identifying a transition time from memorization to generalization and attributes it to different inductive biases. This is not true when weight decay = 0, which weakens the importance of the first phase as it may not even be necessary to understand ``Grokking''.

A: In Sec 4 (now Sec B.1), we have included a discussion on the training regime without weight decay, where the perfect generalization is delayed. However, the transition in test accuracy is extremely slow: in Fig 2 (now Fig 4), we have shown that the transition lasts from $10^{10^2}$ to $10^{10^6}$ in continuous time. (The experiments are done via some simulation with unnaturally growing LR.) This means this grokking phenomenon is not relevant to the practice. It also does not quite satisfy the most strict definition of grokking, where the transition is required to be “sharp”.

In contrast, in our main setting, we study training with large initialization and small weight decay, where the transition in the implicit bias is very sharp: $\frac{1-c}{\lambda} \log \alpha$ leads to the kernel predictor, but increasing it slightly to $\frac{1+c}{\lambda} \log \alpha$ leads to a KKT solution of min-norm/max-margin problems associated with the neural net.

Main Questions

Q1: What are the limitations of the analysis done in this paper? Ignoring the proof difficulties, do you think the change of implicit biases can be observed and explain grokking in more common settings?

A:

1. We added new experiments on modular addition in Sec. 2 to confirm that enlarging the initialization scale or reducing the weight decay indeed delays the sharp transition in test accuracy. Beyond algorithmic dataset, Liu et al. (2023) observed that large initialization and small weight can induce grokking on many popular tasks, including image classification, sentiment analysis and molecules. Thus, we believe that our theoretical insights on implicit bias can have broad implications in these settings as well.
2. One limitation is that our analysis can only handle the training dynamics with large initialization and small weight decay, but these may not be the only source of the dichotomy of the implicit biases. We further discuss in Sec B (previously Sec 4) that the late phase implicit bias can also be induced by implicit margin maximization and sharpness reduction, though the transition may not be as sharp as the weight decay case.

Q2: What would be the effects of stochasticity in GD (SGD) and large step size (distancing from gradient flow)?

A: As the training pipeline of deep learning can induce many confounding factors (e.g., stochastic noise, finite step size, weight decay, etc.), in this paper we focus on identifying simple yet insightful theoretical setups where grokking with sharp transitions can be rigorously proved and its mechanism can be intuitively understood. This is the reason why we focus on gradient flow in our analysis, and we leave it as future work to explore the effects of noise and large LR.

Q3: Can you further explain Remark 3.11?

A: Remark 3.11 (now Remark 3.13) states that if we use random initialization as stated in Theorem 3.10 (now Theorem 3.12) with small variance $\sigma^2$, then gradient flow on the matrix completion problem exhibits a transition from the kernel regime to rich regime. We need to add noise to initialization here because gradient flow may get stuck at saddle points for some initial points that form a zero-measure set. The result of the rich regime is already stated in Theorem 3.10. For the kernel regime, Corollary 3.9 (now Corollary 3.11) only applies to deterministic initialization, but if we add sufficiently small noise to the initialization, the trajectory of gradient flow would remain close to the original one for arbitrarily long time period, so gradient flow would still learn a solution that is close to the one in Corollary 3.9 (now Corollary 3.11).

Q4: Can you further explain the learning rate and time scaling for Figure 2?

A: The time scaling in Figure 2 (now Figure 4) is in log scale or double log scale. The main point of this figure is to show that implicit margin maximization can take an extremely long time to improve generalization. To simulate gradient flow for a very long time, we set the learning rate as $\frac{\eta}{\mathcal{L}(\theta(t))}$, where $\eta$ is a constant. The main intuition behind this is that the Hessian of $\mathcal{L}(\theta)$ can be shown to be bounded by $O(\mathcal{L}(\theta) \cdot \mathrm{poly}(||\theta||))$, and gradient descent stays close to gradient flow when the learning rate is much lower than the reciprocal of the smoothness. Since we view gradient descent as a simulation of gradient flow, the variable “t” in the figure stands for the continuous time, calculated by summing the learning rates up to this point. We have included these details in the revision.

Q5: Is there any proof provided for Theorem 4.1?

A: We do not view Theorem 4.1 (now Theorem B.1) as a part of our novel theoretical contributions because it directly follows from the previous paper (Li et al., 2021). Still, we present this theorem since it is worth pointing out that (Li et al., 2021) implicitly demonstrated one possible cause of grokking, which does not seem to be known in the literature.

Q6: Can you clarify … [about a proof step in Lemma A.5 and B.2]

A: At the beginning of Sec. A.2 (now Sec. C.2), we define $T_{\max} := \inf\\{ t \ge 0 :\|\frac{e^{\lambda t}}{\alpha}\theta(t) - \bar{\theta}\| > \epsilon_{\max}\\}$. In our proof of Lemma A.5, for all $t \le \min\\{T_{\max}, \frac{1}{\lambda}(\log \alpha - \frac{1}{L} \log \log A + \Delta T)\\}$, we can derive the bound $\|\frac{1}{\alpha}e^{\lambda t}\theta(t) - \bar{\theta}\| \le \epsilon_{\max}$, so by definition we must have $T_{\max}\ge \frac{1}{\lambda}(\log \alpha - \frac{1}{L} \log \log A + \Delta T)$.

Q7: What are q and r?

A: Sorry for the confusion of notations here. The function $q_i(\theta)$ is defined as $y_i f_i(\theta)$ (following Lyu & Li, 2020), and $r_i(\theta) = \ell(f_i(\theta); y_i) = \exp(-q_i(\theta))$ (following Moroshko et al., 2020). We have updated the paper.

Other Questions

Q8: What is the norm in Figure 1?

A: The norm in figure 1 is the parameter norm, i.e., $\\|\theta\\|_2$, where $\theta$ is rigorously defined in Sec. 3.2.3. We have made this explicit in the revised paper.

Q9: The analyses require the data to be linearly separable for the classification setting. Also, the model should be able to interpolate in the regression setting. Although these assumptions are reasonable for overparametrized model, it would be beneficial if they were stated more clearly.

A: Thanks for pointing this out. We have made these two assumptions explicit in Assumptions 3.3 and 3.8 in the revised version.

Q10. In the paragraph before Corollary 3.5, isn't the kernel feature (2x, -2x)?

A: Thanks for pointing out this typo. We have updated the paper.

R1. In the paragraph of Empirical Validation: Grokking, it would be nice if some references for and sample complexities were provided. Further, if some references/theoretical results for poor generalization of L2 margin are provided, it would complete the theoretical part of grokking for classification.

A: Thanks for the suggestions. In the revised version of the paper, we provide generalization guarantees for both L2 and L1 margin maximization in Sec. E. Our guarantees indicate that L2-max margin classifier requires $O(kd)$ samples to achieve good generalization, while L1-max margin classifier only requires $\tilde{O}(k^2)$ samples. Here since we are considering sparse vectors, $k$ is much smaller than $d$. As a result, L1 max-margin has a much smaller sample complexity.