Paper: Dichotomy of Early and Late Phase Implicit Biases Can Provably Induce Grokking (ICLR-2024)

Reviewer comment (Theory):
The biggest concern is that the core claim of a "sharp transition in test accuracy" is not supported by any of the analysis made in the paper. While the paper states theorems showing that two different inductive biases are at work and provides corollaries for diagonal linear networks, there are no statements, analyses, or proofs in the main paper related to generalization risks or guarantees.

Author response:
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.

Reviewer comment (Novelty):
Without addressing the above concern, the best conclusion one can make is that the models can fit the training data and eventually reach perfect training accuracy, but this line of work has already been done by Moroshko et al. (Moroshko).

Author response:
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.

Reviewer comment (Theory):
In Section 3.1.3 on page 5 of the draft, the paper claims that implicit biases imply a transition in test accuracy. As noted in the weaknesses, this is not supported by any analysis (e.g., bounds).

Author response:
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.

Reviewer comment (Experiments):
No empirical results are provided that help the reader crystallize their understanding of the work in terms of when the transition might occur in practice. More examples with datasets commonly used in the "Grokking" literature would be helpful.

Author response:
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.

Reviewer comment (Theory):
There is very limited discussion on the possible extensions and limitations of this model. Moreover, results are in the case that the initialization scale goes to infinity which doesn't show (at least immediately) if these observations can be observed in the more common settings.

Author response:
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.

Reviewer comment (Experiments):
Additional experiments can clarify other potential limitations of the analysis (e.g., gradient flow).

Author response:
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.

Reviewer comment (Theory):
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?

Author response:
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.

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

Author response:
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.

Reviewer comment (Theory):
Can you further explain Remark 3.11?

Author response:
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).

Reviewer comment (Reproducibility):
Can you further explain the learning rate and time scaling for Figure 2? (E.g., the $\log(t)$ has often very large values.)

Author response:
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.

Reviewer comment (Theory):
Is there any proof provided for Theorem 4.1?

Author response:
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.

Reviewer comment (Writing):
What is $q(\theta)$ in Appendix A.2? Also $r_i$'s?

Author response:
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.

Reviewer comment (Presentation):
In Figure 1, what is the norm presented? (Whether it's L1/L2 and if it's for $w$ or $u,v$?)

Author response:
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.

Reviewer comment (Theory):
I think 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 in the main text.

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

Reviewer comment (Theory):
In the paragraph before Corollary 3.5, isn't the kernel feature $(2x, -2x)$?

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

Reviewer comment (Theory):
In the paragraph of Empirical Validation: Grokking, it would be nice if some references for $O(d)$ and $O(k\log d)$ 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.

Author response:
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.