Unified View of Grokking, Double Descent and Emergent Abilities: A Comprehensive Study on Algorithm Task

Thank you for the valuable comments that help us improve our work. Here is our response:

Questions about framework.

1. How do we arrive at this framework?

   We developed this framework based on two key experimental findings: the decreasing curve of critical dataset size for generalization and the increasing curve of the model's memorization capacity, as shown in our Preliminary Study.

2. why is it good?

   We believe the framework has two advantages:

   Firstly, it is consistent with existing work and links various phenomena. For example, a vertical line in the right matches Varma's results[1]. This framework can demonstrate the relationship between double descent and grokking. We can imagine that if the critical dataset curve also increases with model size, then double descent would not occur, which violate double descent theory.

   Secondly, a good framework helps predict model's performance. In Section 4.3, we make predictions based on it that increasing difficulty of generalization would result in a more pronounced double descent, which was also confirmed by our experiments.

3. How would it connect to a more theoretical account?

   In the paper, we analyzed the framework through the lens of competition between memorization and generalization. Here, we provide an optimization-based analysis.

   With the AdamW optimizer, both the gradient from cross-entropy loss and weight decay influence optimization. The memorization capacity curve relates to the loss, aiming for near-zero training loss. Beyond this curve, non-zero loss pushes the model toward generalization. Below the curve, near-zero loss highlights the impact of weight decay. The critical dataset size curve represents where weight decay doesn't favor memorization or generalization. Thus, different training dynamics also reflect the varying effects of cross-entropy loss and weight decay.

make the writing sharper

We also noticed similar issues in the current content due to the numerous findings and experimental results. To address this, we will add key takeaways at the end of the introduction, directly linking to the corresponding sections. Here are some examples of takeaways:

1. Critical dataset size for generalization decreases with model size. -> Sec 2.2
2. Memorization capacity increases with model size. -> Sec 2.3
3. The training dynamics can be categorized into four types. -> Sec 3

[1] Explaining grokking through circuit efficiency. Varma et al.

Thank you for the valuable comments that help us improve our work. Here is our response:

Tag in Figure 4

We acknowledge that final performance does not fully support the labels for different stages, especially grokking. Visualizing both final performance and training dynamics in a single figure is challenging. Thus, we propose another metric for measuring delayed generalization during grokking. We define $S_t(M)$ as the training steps to achieve 99% of maximum training accuracy and $S_v(M)$ for validation accuracy. The difference, $S_{delay}(M) = S_v(M) - S_t(M)$, quantifies the delayed generalization. A large $S_{delay}(M)$ indicates grokking. Below are the $S_{delay}$ values for several grokking experiments in Figure 4:

All models exhibit delayed generalization. As the model size and data size increase, the number of delayed steps for generalization decreases significantly.

The performance on the memorization examples in the multitask experiment

In all our multitask experiments, the performance on memorization examples quickly achieved 100%.

lower bound of the memorization dataset size for observing the effect in Section 5

We conducted additional experiments with fewer memorization data points. The results are shown below, where the columns represent model size, and the rows represent number of memorization data:

The impact of memorization data decreases with fewer data points, which is expected. A small amount of memorization data acts as noise, and the model can typically ignore small noise and identify correct patterns. These results also suggest a lower bound for memorization data required to exhibit emergence lies between 1000 and 2000 points when using 3000 generalization data.

Experimental setting in Figure 1

We will add a brief explanation in the main article to provide better clarity. Currently, the specific experiment settings are described in the caption of Figure 1.

(1c) doesn't seem to show "multiple plateaus"

We will rephrase it as "multiple stages of increase", as shown in the curve around 20,000 steps in Figure 1(c).

Thank you for the valuable comments that help us improve our work. Here is our response:

Could engage with the relevant literature much more thoroughly.

We appreciate you pointing out the additional related work that we missed. We will incorporate these references into the appropriate section in our revised version.

Writing of Section 5

In Section 5, we did not intend to imply that implicit multi-task learning during the pretraining stage is the sole reason for the emergence of abilities. We agree that the emergence of abilities in current LLMs is driven by a multitude of complex factors and requires further investigation. Our goal with the simulation experiment was to offer an alternative perspective on understanding emergent abilities in LLMs, potentially spurring further research. We acknowledge that some of our wording may have been too strong, leading to misunderstandings, and will clarify this in the revised version.

Shape of memorization capacity

Figure 3 shows a non-linear relationship between the model's hidden size and memorization capacity. Initially, memorization capacity increases rapidly with small hidden sizes, then transitions to a linear increase. The curve's flattening towards the end might be misleading due to the limited training data (12,769), which underestimates the model's capacity. Therefore, a definitive conclusion requires further verification with different training settings, such as various learning rates, data quantities, and task types. Further research on this topic would be interesting.

Regarding the impact of layer number on memorization capacity, we conducted experiments with a hidden size of 32. The results are shown below. We observed that increasing the number of layers does increase the model's memorization capacity, but this increase is less significant compared to increasing the hidden size. For instance, comparing row 2 and row 5, despite row 2 having more parameters, it exhibits lower memorization capacity.

Legend in Figure 4; Use fewer evaluative words/phrases; Typos

Thank you for highlighting these issues. We will make the necessary modifications to address all these problems.

Thank you for the valuable comments that help us improve our work. Here is our response:

We acknowledge that a limitation of our study is that all experiments and analyses are conducted on algorithmic tasks, as noted in our Limitations Section. The primary reason for choosing these tasks is their simplicity, absence of noisy data, and lack of overlap between the training and validation sets. This allows us to clearly distinguish the model’s behaviors in terms of memorization and generalization.

Moreover, we believe our framework can also make sense for more realistic models and tasks, as the concepts of memorization and generalization remain pertinent in these contexts. Notably, there is existing research that has observed grokking in real-world tasks such as MNIST and IMDB, though the effects are less pronounced compared to algorithmic tasks [1]. We are currently attempting to conduct experiments on real-world tasks to further validate our framework. However, these experiments are highly time-consuming due to the extremely long training periods required for observing grokking. Should we obtain additional results, we will be pleased to include them in our revised paper.

[1] Ziming Liu, Eric J. Michaud, Max Tegmark. Omnigrok: Grokking Beyond Algorithmic Data. ICLR 2023