Paper: Grokking Tickets: Lottery Tickets Accelerate Grokking (ICLR-2024)

Reviewer comment (Writing):
The writing needs improving, and at times, I find it hard to follow the text. For instance, the term "grokking tickets" is used 12 times before its formal definition presented, which presents a certain impediment to readers.

Author response:
As you pointed out, we have defined Grokking tickets in the Abstract and at the beginning of the Introduction. Additionally, for readability, we have made changes to the writing, such as defining comparison models and correcting typographical errors.

Reviewer comment (Novelty):
I am uncertain about the practical significance of the findings. Specifically, I understand that one of the conclusions is that tickets found during the $C_{mem}$ phase exhibit better performance. However, this seems to align with the common practice in LTH research~\cite{chen2021unified, gan2022playing}, where networks are thoroughly trained and then pruned based on optimal validation scores during each pruning iteration. Therefore, the question arises as to what contribution the findings in the paper make to the discovery of improved lottery tickets.

Author response:
The novelty of this research in relation to the Lottery Tickets hypothesis is as follows: (1) Lottery Tickets as an explanation for the Grokking phenomenon. (2) Lottery Tickets that acquire good structures for generalization, not just a reduction in parameters. 
- Regarding point (1), in Appendices G and H, we conducted analyses on how Grokking tickets are discovered. 
In Appendix G, building upon prior research [Nanda 2022] [Zhong 2023], we investigated the structures that emerge during the transition phase in the modular addition task. In the task of modular addition, it is known that periodic representations are necessary for generalization. Our paper also demonstrates similar results, and it was found that in the case of Grokking tickets, good representations are acquired more quickly.
In Appendix H, following the [Paganini 2020], we examined the structures acquired mid-way through grokking using the Jaccard distance [P. Jaccard 1901] . As can be seen from Section 4, the results show that structures which gradually generalize (lottery tickets) are being acquired. Especially during the transition phase, the distance of the masks decreases rapidly, indicating that the discovery of structures is important for generalization.

- Regarding point (2),  in Appendices I and J, We conducted an analysis demonstrating that generalization is accelerated not merely due to the reduction of parameters through Lottery Tickets, but through the discovery of good structures.
In Appendix I, to respond to the reviewer's concern that the lottery tickets learn faster than the original dense model has already been demonstrated in the original LTH paper and many other previous work. We conducted experiments using Grokking tickets, where we
applied the mask only to the initial values and not during the course of learning. We refer to this model as the ’Re-Dense model’. By doing this, the number of parameters optimized by the optimizer becomes equal in both the Dense and Re-Dense models, differing only in their initial values.  Base Model and Re-Dense Model have the same number of parameters, with the only difference being in their initial values (further, the difference is just the application of a mask to the same initial values). It is evident that they generalize faster than the Base Model. In Appendix J, we conducted additional experiments. One experiment involved masking only the completely 'dead' units (18%) of Grokking tickets. If Grokking tickets were accelerating generalization simply by reducing the number of learnable parameters, as suggested by existing studies, then this method should also accelerate grokking. However, as shown in Appendix J, this method not only did not accelerate generalization but actually slowed it down. Conversely, it was also confirmed that reviving the completely dead units in Grokking tickets does not change the speed of generalization compared to the original Grokking tickets.

We believe these results provide important analysis for research in the Lottery Ticket (LT) field.

Reviewer comment (Experiments):
While the existing results do indeed support the author's assertions, the use of overly simplistic tasks (modular addition and MNIST classification), basic baselines (MLP and single-head Transformer), and a lack of additional experimental configurations (such as performance under multi-layer MLPs or multi-head attention) leave me with reservations regarding the validity of the conclusions.

Author response:
Following the reviewer's comments, we have added experiments with wide MLPs and multi-head Transformers in Appendix F.
Consistent with the claims in the main text, it becomes clear in the case of Grokking tickets that the speed of generalization is accelerated.

Reviewer comment (Presentation):
The paper could be better organized. To name a few, in 2.1 ‘abuse’ -> ‘abuse of’; in 5.1 ‘leaders’->’readers’. Excessive space is devoted to some trivial aspects, such as formulas for weight initialization.

Author response:
As you pointed out, we have defined Grokking tickets in the Abstract and at the beginning of the Introduction. Additionally, for readability, we have made changes to the writing, such as defining comparison models and correcting typographical errors.

Reviewer comment (Experiments):
It is advisable to include additional backbones and experimental setups to comprehensively validate the claims.

Author response:
Following the reviewer's comments, we have added experiments with wide MLPs and multi-head Transformers in Appendix F.
Consistent with the claims in the main text, it becomes clear in the case of Grokking tickets that the speed of generalization is accelerated.

Reviewer comment (Novelty):
The paper attempts to make a connection to the Lottery Ticket Hypothesi (LTH). Sparsity is a reasonable hypothesis that has been explored by Merrill et. al. previously in literature. However, I am not convinced that connecting LTH to this work is necessary. Merrill 2023 make observations about sparse networks without invoking LTH.

Author response:
To clarify the differences from prior research, we added a section on related works.
As mentioned in section 6, while Merrill et. al  prunes neurons with smaller weight norms, we use the lottery ticket (LT) method for pruning weights. Besides, we tested arithmetic tasks where the importance of sparsity is not obvious and the standard benchmark in the grokking, while they have only conducted experiments in tasks like (n-k)-parity where the importance of sparsity is obvious. 
To be applicable to more general tasks and architectures, we need to use the Lottery Tickets for analysis of grokking using the Lottery Ticket (LT) method.

Reviewer comment (Novelty):
Nanda et. al. (Nanda 2023) show that the network after generalization consists of a few sinusoids, i.e., finds a sparse solution. So sparsity being an explanation has been shown in prior literature. Also suggest that the authors consult Gromov 2023 for the solution found a.

Author response:
The novelty of our study compared to previous research can be summarized in four points: 1) Methodology (Lottery Tickets), 2) Tasks, 3) Architecture, and 4) Claims.
1) We conducted analysis using the Lottery Tickets (LT) method. The necessity of analysis through LT is as described above.

2) Our tasks are not (n-k) parity where the optimal structure being sparse  is self-evident, but we have also validated in more general tasks such as MNIST.

3) Unlike Merrill 2023 or Nanda 2022, we have not confined our validation to just one architecture; instead, we have tested across multiple architectures.

4) Furthermore, we assert that it's not just sparsity that's important, but the acquisition of good structures is crucial (refer to section 5.2, Appendix G and H).

Reviewer comment (Novelty):
The paper notes in Section 5.1 (ablation of weight norms) that a subnetwork has increased weight norm after a generalizing solution is found by gradient descent while other weights end up with smaller norms. This is the same observation made in Merrill 2023 where the authors study subset parity learning problem.

Author response:
Indeed, similar observations have been confirmed, but as mentioned above, our results indicate that not just sparsity, but also the acquisition of good structures is important for generalization.  In section 5.2 in the original version of our paper, we compared various methods fixed at the same level of sparsity, and these results showed that Grokking tickets have effects beyond just reducing parameters. This shows that in explaining grokking, it's not just sparsity that's key, but also the acquisition of good structures is crucial for generalization.

Reviewer comment (Novelty):
The observation that the sparse generalized solution optimizes faster than the dense network while interesting is not critical to understanding Grokking. It is known that regularization or lowering capacity of models via regularization does help shorten the time between fitting and generalization in algorithmic datasets. This observation falls in-line with the above.

Author response:
As suggested in existing research, one reason why generalization is faster with Grokking tickets is due to the reduction in the number of learnable parameters. However, our study suggests that grokking is not just about reducing the effective number of parameters being learned, but also about acquiring structures appropriate to the task. For instance, in section 5.2 in the original version of our paper, we compared various methods fixed at the same level of sparsity, and these results showed that Grokking tickets have effects beyond just reducing parameters.
To analyze this in detail, we conducted additional experiments. One experiment involved masking only the completely 'dead' units (18%) of Grokking tickets (see Appendix J). If Grokking tickets were accelerating generalization simply by reducing the number of learnable parameters, as suggested by existing studies, then this method should also accelerate grokking. However, as shown in Appendix J, this method not only did not accelerate generalization but actually slowed it down. Conversely, it was also confirmed that reviving the completely dead units in Grokking tickets does not change the speed of generalization compared to the original Grokking tickets.
We have added these discussions in the revised section 6.4.

Reviewer comment (Experiments):
Given the empirical nature of the paper, the number of datasets considered in the analysis appear to be inadequate. Power 2022 construct various algorithmic datasets. Gromov 2023 use a MSE solution with MLP for addition. Have the observations been confirmed in more settings than the ones considered in the paper?

Author response:
Following the reviewers’ comments, we conducted experiments across various task configurations and with different architectures. 
In Appendix D, similar to MLP experiments, we conduct experiments on single-head transformer architecture. Similar to the experiments conducted in the paper, similar results were obtained on single-head transformer architecture as well.
In Appendix E,  we conduct experiments on multi-head transformer architecture and multi-layer MLP. Similar to the experiments conducted in the paper, generalization is faster due to Grokking tickets. on multi-head transformer architecture and multi-layer MLP as well.
In Appendix F, referencing [Power 2022], we conduct experiments on diverse tasks. Similar to the experiments conducted in the paper, generalization is faster due to Grokking tickets　on diverse tasks.

Reviewer comment (Novelty):
My major concern is that while the combination of LTH with Grokking is new to the community, the empirical findings shown in this paper is somehow mediocre. For instance, i. it is not surprising to see that the subnetworks obtained during the generalization phase is crucial for grokking speedup.

Author response:
>it is not surprising to see that the subnetworks obtained during the generalization phase is crucial for grokking speedup.

As you mentioned, the acceleration of grokking in a subnetwork is not surprising.
To address your concern, we have added the following two analysis.
1) As other reviewers have pointed out, the simplest explanation for faster generalization is that the number of learnable parameters decreases. However, our analysis suggests that performance improvement is not just due to a reduction in parameters. Firstly, not any structure works effectively (refer to section 5.2), and secondly, even if only completely 'dead' neurons in grokking tickets (18%) are removed, learning does not accelerate but rather decelerates (refer to Appendix J). These results imply that it's not just about reducing the number of parameters, but the construction of some good structures contributes to the acceleration of generalization.
2) 
As stated in general response (2), to analyze how Grokking tickets are acquired during the course of learning, we examined the periodicity of weights (refer to Appendix G) and the similarity of masks (refer to Appendix H).
Additionally, the results of the newly added analysis of representations in grokking tickets indicate that better representations are acquired earlier with grokking tickets. While the relationship between grokking and representation learning has been widely discussed, our study connects the perspectives of representation learning and structure acquisition in grokking.

Reviewer comment (Novelty):
The definition of the Grokking Tickets has no essential difference than the original Lottery Tickets and can be covered by the original ones since the original LTH does train the dense model to the end.

Author response:
>The definition of the Grokking Tickets has no essential difference from the original Lottery Tickets and can be covered by the original ones since the original LTH does train the dense model to the end.

As stated in General Response (2), we are not proposing a new LT method; rather, we are employing the Lottery Ticket (LT) method as a means to analyze grokking. Therefore, Grokking tickets are encompassed within the definition of LT.
However, as mentioned above, experiments in Appendix I and J, demonstrate that Grokking tickets are not merely sparse (i.e., having fewer learning parameters), but also acquire good structures for generalization.

Reviewer comment (Novelty):
Perhaps, the authors can emphasize the contribution of the papers from the perspective of Grokking. Why the grokking tickets are important for Grokking?

Author response:
>Perhaps, the authors can emphasize the contribution of the papers from the perspective of Grokking. Why the grokking tickets are important for Grokking?

As mentioned in the General Response as well, in section 6, we have detailed the comparison between previous studies on grokking and our research. Especially in existing studies, explanations often focus on weight norms and weight decay [Liu 2022, Power 2022]. In contrast, our research not only examines weight norms but also focuses on the network structure. We determined that grokking occurs due to the discovery of lottery tickets.
Although Merrill et al. emphasize the importance of sparsity, they do not extend their claim to the acquisition of good structures.
Therefore, in response to the question 'Why are Grokking tickets important for Grokking?', our answer lies in the acquisition of good structures.

Reviewer comment (Evaluation):
Besides the speedup, does Grokking tickets bring any performance benefits over the dense one?

Author response:
>Besides the speedup, does Grokking tickets bring any performance benefits over the dense one?

Of course, as we are finding extremely sparse solutions, I believe it is highly relevant from the perspective of learning efficiency. For instance, the model we used in this study maintained its generalization performance even after pruning about 81% of its edges. (Conversely, this led to an earlier occurrence of grokking.)
As also mentioned in the General Response, the relationship with representation learning is conceivable, and our findings are likely to provide important insights from the perspective of model interpretability.

Reviewer comment (Experiments):
The experiments are confined to a narrow set of tasks and architectures, which might not fully demonstrate the generalizability of the proposed method.

Author response:
## Expand Experimental Scope
>Limited Experimental Scope
>Would the authors consider expanding their experimental evaluation to include a wider variety of tasks and network architectures to confirm the robustness of their findings?

Following the reviewers’ comments, we conducted experiments across various task configurations and with different architectures. 
In Appendix D, similar to MLP experiments, we conduct experiments on single-head transformer architecture. Similar to the experiments conducted in the paper, similar results were obtained on single-head transformer architecture as well.
In Appendix E,  we conduct experiments on multi-head transformer architecture and multi-layer MLP. Similar to the experiments conducted in the paper, generalization is faster due to Grokking tickets. on multi-head transformer architecture and multi-layer MLP as well.
In Appendix F, referencing [Power 2022], we conduct experiments on diverse tasks. Similar to the experiments conducted in the paper, generalization is faster due to Grokking tickets　on diverse tasks.

Reviewer comment (Experiments):
The paper does not provide an in-depth analysis of the scalability of the approach or a comparison with other state-of-the-art methods across varied settings.

Author response:
## The scalability of our  method
>How do the authors envision the scalability of the proposed method, and how does it compare with other pruning or training acceleration techniques in terms of computational efficiency and performance?

As mentioned in the General Response, in the Appendix G and H, unlike other pruning methods, we investigated whether the network structure changes similarly during grokking and generalizes, from the perspective of Lottery Tickets (LT). This analysis is beneficial not only for research on efficient learning methods but also for the interpretability of models.

Reviewer comment (Theory):
Can the authors elaborate on the theoretical foundations that might explain the observed acceleration in learning due to grokking tickets?

Author response:
>Theoretical Underpinning
>Can the authors elaborate on the theoretical foundations that might explain the observed acceleration in learning due to grokking tickets? 

To address your concern, we have added the following two analisis.
As other reviewers have pointed out, the simplest explanation for faster generalization is that the number of learnable parameters decreases. However, our analysis suggests that performance improvement is not just due to a reduction in parameters. Firstly, not any structure works effectively (refer to section 5.2), and secondly, even if only completely 'dead' neurons in grokking tickets (18%) are removed, learning does not accelerate but rather decelerates (refer to Appendix J). These results imply that it's not just about reducing the number of parameters, but the construction of some good structures contributes to the acceleration of generalization.

Reviewer comment (Experiments):
Would the authors consider expanding their experimental evaluation to include a wider variety of tasks and network architectures to confirm the robustness of their findings?

Author response:
## Expand Experimental Scope
>Limited Experimental Scope
>Would the authors consider expanding their experimental evaluation to include a wider variety of tasks and network architectures to confirm the robustness of their findings?

Following the reviewers’ comments, we conducted experiments across various task configurations and with different architectures. 
In Appendix D, similar to MLP experiments, we conduct experiments on single-head transformer architecture. Similar to the experiments conducted in the paper, similar results were obtained on single-head transformer architecture as well.
In Appendix E,  we conduct experiments on multi-head transformer architecture and multi-layer MLP. Similar to the experiments conducted in the paper, generalization is faster due to Grokking tickets. on multi-head transformer architecture and multi-layer MLP as well.
In Appendix F, referencing [Power 2022], we conduct experiments on diverse tasks. Similar to the experiments conducted in the paper, generalization is faster due to Grokking tickets　on diverse tasks.

Reviewer comment (Experiments):
How do the authors envision the scalability of the proposed method, and how does it compare with other pruning or training acceleration techniques in terms of computational efficiency and performance?

Author response:
## The scalability of our  method
>How do the authors envision the scalability of the proposed method, and how does it compare with other pruning or training acceleration techniques in terms of computational efficiency and performance?

As mentioned in the General Response, in the Appendix G and H, unlike other pruning methods, we investigated whether the network structure changes similarly during grokking and generalizes, from the perspective of Lottery Tickets (LT). This analysis is beneficial not only for research on efficient learning methods but also for the interpretability of models.

Reviewer comment (Evaluation):
While this work made a lot of interesting observations, profound analysis and investigation behind empirical results are lacking.

Author response:
>While this work made a lot of interesting observations, profound analysis and investigation behind empirical results are lacking.
Following the reviewer comments, we conducted three additional analyses to investigate the implication behind empirical results. 

(a) Analyzing similarity of structure (lottery tickets) among memorization, transition, and generalization phase of training. 
Appendix H, following the [Paganini 2020], we examined the structures acquired mid-way through grokking using the Jaccard distance [P. Jaccard 1901] . As can be seen from Section 4, the results show that structures which gradually generalize (lottery tickets) are being acquired. Especially during the transition phase, the distance of the masks decreases rapidly, indicating that the discovery of structures is important for generalization.

(b) Connection between grokking tickets to existing analysis. 
Recently,  [Nanda 2022] investigated the weights of networks during grokking, and found that the weights exhibit clear periodic patterns after the grokking occurs. We follow the similar experiments to connect the grokking tickets and the representation learning perspective of the grokking in Appendix G. The results show that the grokking tickets help to learn task-appropriate representations, i.e. unlike the dense model, grokking tickets help to learn the periodicity in weights at the early stages of training.

Reviewer comment (Writing):
The writing quality of this work is not satisfactory.

Author response:
>The writing quality of this work is not satisfactory.

Thank you for carefully reading our paper. To improve the clarity, we proof-read the paper again and made necessary revisions. A significant change is the addition of a section on related work in Section 6. We also have clearly defined grokking tickets and other comparison models in the abstract and introduction.