Bridging Lottery Ticket and Grokking: Understanding Grokking from Inner Structure of Networks

We thank the reviewer for the constructive feedback. Please let us know if our responses in the following address your concerns.

We revised the paper based on the reviewers’ comments, and the major edit was highlighted with coloring (purple). Please also check the updated manuscript.

W1

The hyperparameters (e.g. learning rate and weight decay) are not tuned.

We acknowledge the variability of grokking speed ($\tau_{\mathrm{grok}} = t_{\mathrm{train}} / t_{\mathrm{test}}$) with hyperparameter settings such as learning rate and weight decay. To clarify this, Appendix M includes experiments that illustrate how $\tau_{\mathrm{grok}}$ changes under various configurations. As you noted, hyperparameter tuning can influence the observed speed of grokking. However, our primary focus is not on the absolute speed of grokking but on the structural insights gained through the use of lottery tickets. Specifically, our study aims to explain how lottery tickets reduce delayed generalization from a structural perspective, as detailed in Sections 4 and 5. The findings in Section 3 serve as preparatory groundwork.

Thus, while speed changes are observed, they are not the central message of our work. To clarify this emphasis, we revised our terminology throughout the paper, including in the Abstract, Introduction, and Section 3, replacing statements like “65 times faster” with “reduce delayed generalization.”

This better reflects the essence of our findings and avoids overemphasizing specific speed improvements. We hope these clarifications address your concerns and provide a clearer context for our results.

W2 & Q1

Why the network structure can be measured by Jaccard distance? It seems intuitive that when the model undergoes a phase transition (from memorization to generalization), its weight norm will change rapidly, causing the Jaccard distance to also increase rapidly.

Why the Jaccard distance can be used as a progress measure? The dynamics of Jaccard distance and that of test accuracy change show similar trend.

While it is true that the weight norm decreases during the phase transition (as shown in Figure 6), we explain in Section 4.1 that weight norm alone cannot fully account for the phenomenon. This motivates the need for a new progress measure that better captures the underlying dynamics of structural changes during grokking. To address this, we propose the Jaccard distance as a novel progress measure.

The Jaccard distance directly quantifies the distance between the structures of two networks, providing insights into the evolution of the inner structure as it transitions from memorization to generalization. Unlike the weight norm, which only reflects a global property of the weights, the Jaccard distance captures structural changes more explicitly, as evidenced by its alignment with test accuracy trends.

Thank you again for your valuable feedback on our paper. As we have not yet received a response, we would like to kindly remind Reviewer 2Wju to review these revisions.

We would greatly appreciate it if you could consider whether our responses adequately address your comments.

We thank the reviewer for the constructive feedback. Please let us know if our responses in the following address your concerns.

We revised the paper based on the reviewers’ comments, and the major edit was highlighted with coloring (purple). Please also check the updated manuscript.

W1

It seems that, instead of using the LTH as a tool to understand the grokking behavior of neural network training, the paper focus more on the grokking behavior of the pruned subnetworks. From this perspective, though they consist of interesting findings, the results of the paper lacks practical implications: we are not sure how the grokking behavior of the subnetworks is going to shed light on the grokking behavior of neural network in general from an "inner structure" perspective.

We would like to clarify that our paper does not solely focus on pruned subnetworks. Instead, we build upon the concept of pruned subnetworks to derive insights about grokking behavior and extend the analysis to explore the characteristics of these subnetworks in depth. Specifically:

1. Dynamics of Structure Acquisition (Section 5.1, Figure 7):
   Using the Jaccard distance as a metric for structural change, we analyze how the subnetwork evolves during grokking. This goes beyond simply observing pruned subnetworks and provides a dynamic perspective on how subnetworks adapt and contribute to grokking over time.

2. Regularization Beyond Weight Decay (Section 5.2, Table 1):
   From the observation that good structures are crucial for generalization, we introduce the edge-popup algorithm to directly explore structural optimization during training. By combining this structural exploration with traditional weight decay regularization, we demonstrate further improvements in generalization speed. These findings offer practical implications for designing new regularization strategies for grokking, extending beyond the conventional L2-based weight decay.

In summary, our paper leverages pruned subnetworks not as the sole focus but as a foundation to uncover broader principles of grokking behavior. These findings provide practical implications, including a novel perspective on regularization and structural adaptation, which extend beyond the context of pruned subnetworks alone.

We hope this addresses your concern. Please let us know if further clarification is needed.

W2

The paper does not decouple the behavior of "shorter grokking period" from the "easier training" property of the winning tickets.

We agree that winning tickets often exhibit "easy-to-train" properties, such as achieving generalization in fewer epochs, and we acknowledge that this observation itself is not novel. Our paper builds upon this known behavior by specifically connecting it to the grokking phenomenon, which introduces a unique perspective and deeper implications.

For instance, as discussed in **Section 5.1 **, we observe abrupt changes in Jaccard distance during grokking, indicating structural evolution in the subnetworks. Furthermore, **Figure 8 ** reveal that these grokking tickets acquire task-specific structures, such as periodicity in Modular Addition, which are critical for solving the task and are directly tied to the grokking process.

Thus, while faster generalization is a known property of winning tickets, our paper goes beyond this by examining how the discovery of task-specific good structures underpins the grokking phenomenon.

To clarify these points, we have expanded the discussion in Section 6.

We hope this clarifies our contribution, and we are happy to elaborate further if needed.

Thank you again for your valuable feedback on our paper. As we have not yet received a response, we would like to kindly remind Reviewer GMh9 to review these revisions.

We would greatly appreciate it if you could consider whether our responses adequately address your comments.

We thank the reviewer for the constructive feedback. Please let us know if our responses in the following address your concerns.

We revised the paper based on the reviewers’ comments, and the major edit was highlighted with coloring (purple). Please also check the updated manuscript.

W1 & Q5

it is still unclear how specific/different "periodicity" can be regarded as good representations.

Why does periodicity in representations correlate with good performance?

We have addressed this in Appendix L, referencing (Nanda et al.)[1], demonstrating periodicity's significance in modular addition tasks. Below is a brief summary:

The task, defined as predicting $c \equiv (a + b) \mod p$, inherently involves periodicity due to its modular arithmetic structure. Periodic representations align with this structure by encoding inputs $a$ and $b$ into a Fourier basis as sine and cosine components of key frequencies $w_k = \frac{2k\pi}{p}$.

Using trigonometric identities, the model computes:

$$\cos(w_k(a+b)) = \cos(w_k a) \cos(w_k b) - \sin(w_k a) \sin(w_k b),$$

and logits are derived to ensure constructive interference at $c \equiv (a + b) \mod p$ :

$$\cos(w_k(a+b-c)) = \cos(w_k(a+b))\cos(w_k c) + \sin(w_k(a+b))\sin(w_k c).$$

This mechanism enables the model to generalize effectively by leveraging the modular arithmetic structure. We will also refer to this in Section 5.3 of the main text for further clarity.

If anything remains unclear, we are happy to provide further clarification.

[1] Nanda et al., Progress measures for grokking via mechanistic interpretability, https://arxiv.org/abs/2301.05217

W2 & Q3

the paper mainly explores grokking within fully connected networks and lacks a broad examination across various network architectures. For example, adding experiments across a wider range of architectures (e.g., Transformers) would increase the generalizability and appeal of the findings.

How does the finding generalize to other neural network structures (e.g., Transformer) and language-based tasks?

We would like to clarify that our paper includes experiments with Transformers as well as MLPs. Specifically:

1. Figure 2-(b) in the main text presents results from Transformer-based experiments. These results demonstrate that, similar to MLPs, grokking tickets result in less delay of generalization compared to the base model in Transformers.

2. Additionally, we provide results from an NLP task (text sentiment analysis) using an LSTM architecture in Appendix J. In this task, we observed test accuracy improving almost simultaneously with training accuracy from the very beginning of the optimization process.

Across a variety of tasks and architectures, including Modular Arithmetic, MNIST, regression, sparse parity, and an NLP task, our findings consistently show that having a good structure explains grokking. These tasks span architectures such as MLPs, Transformers, and LSTMs, supporting the generalizability of our claims.

We hope this clarifies the breadth of our experiments and how they extend across diverse tasks and architectures. Please let us know if further clarification is needed.

W2 & Q2

The experimental setup is comprehensive but could be further extended to strengthen the generality of the results.

Although the paper does very nicely on modular arithmetic and MNIST, how would grokking tickets generalize to a few domains, perhaps such as NLP, like using a model such as BERT?

To address the concern about the generality of grokking tickets across different domains, we have included an analysis in Appendix J, where we evaluated grokking tickets on a sentiment analysis task using the IMDb dataset (Maas et al., 2011)[2]. This dataset contains 50,000 movie reviews classified as positive or negative, processed with the 1,000 most frequent words and tokenized into arrays of indices. For classification, we employed a two-layer LSTM model, and the experimental details are provided in Appendix J.

The results show the base model rapidly achieves 100% training accuracy; however, its test accuracy remains at 50%(chance rate) until approximately 10k optimization steps, after which it begins to improve. On the other hand, the grokking ticket demonstrates a different behavior, with test accuracy improving almost simultaneously with training accuracy from the very beginning of the optimization process. This suggests that the grokking ticket's structural properties facilitate faster and more efficient generalization.

These findings extend the applicability of grokking tickets beyond modular arithmetic and MNIST to NLP tasks such as sentiment analysis. They also underscore the potential of grokking tickets to uncover meaningful patterns in different domains, further demonstrating their versatility and generality.

[2] Maas et al., Learning Word Vectors for Sentiment Analysis, https://aclanthology.org/P11-1015/

W3 & Q4

the authors do account for the role of structural similarity in terms of Jaccard similarity in their grokking explanation, one might further explore what the implications of such a metric would be

Could the authors go into additional detail about which variations in structural similarity correlate with stages of grokking?

Our hypothesis is that during grokking, the network acquires a good structure that is beneficial for generalization. For the modular addition task, this corresponds to a periodic structure, as demonstrated in Figure 8 (bottom).

The Jaccard similarity metric directly reflects how the network structure (as captured by the magnitude pruning mask) evolves during training. Figure 7 shows that the structure changes abruptly, corresponding with the sharp increase in test accuracy, indicating a significant structural transformation of the network. When we delve deeper and combine these findings with the results on periodic structures (Figure 8), we observe that during grokking, the network rapidly transitions to a periodic structure—one that is highly suitable for the task—at the point when test accuracy improves dramatically.

W4 & Q3

it may be more helpful to know how this might impact the practices in the real world in terms of training or pruning neural networks.

the results show generalization with the appropriate subnetworks instead of depending on weight decay, do the authors have any insights on practical takeaways?

Our findings that highlight the importance of good subnetworks suggest a new perspective on regularization. While weight decay traditionally regularizes the L2 norm of weights, our results imply that effective regularization should focus on discovering good structures within the network. In this sense, weight decay can be seen as a proxy for encouraging such structure discovery indirectly.

For example, as shown in Table 1, adding the edge-popup algorithm, which explicitly explores structure (optimizes masks), to weight decay resulted in faster generalization. This demonstrates that integrating structural exploration into training can enhance performance beyond what weight decay alone achieves.

These insights suggest that practitioners may improve generalization by incorporating methods that directly optimize beneficial structures rather than solely relying on traditional regularization techniques like weight decay. Our results pave the way for developing new, structure-oriented regularization techniques to better leverage the benefits of grokking tickets in practical applications.

To make this point clearer, we have revised Section 5.2 to elaborate on the implications of the edge-popup algorithm's results.

We thank the reviewer for the constructive feedback. Please let us know if our responses in the following address your concerns.

We revised the paper based on the reviewers’ comments, and the major edit was highlighted with coloring (purple). Please also check the updated manuscript.

W1 & Q1

the theoretical motivation for why subnetworks (grokking tickets) are enough to explain grokking.

It would be great if you provide more theoretical insight into why it is enough for generalization to have the presence of certain subnetworks (grokking tickets)?

While not fully grounded in theory, but it is evident that task-adaptive structures contribute significantly to generalization. For instance, as demonstrated in Neyshabur [1], which is based on the theory of Minimum Description Length (MDL), incorporating $\beta$-Lasso regularization into fully connected MLPs facilitates the emergence of locality—resembling the structures found in CNNs—leading to improved performance in image-related tasks. This insight aligns closely with the motivation of our work, as mentioned in the introduction of our paper: to analyze the delayed generalization in grokking from the perspective of network structure.

While the relationship between good structures and generalization has been extensively studied in deep learning, the connection between grokking and network structural properties remains underexplored. Our research seeks to bridge this gap by investigating how specific subnetworks (grokking tickets) contribute to generalization.

We hypothesize that the reason grokking tickets generalize well is their ability to acquire a task-adaptive structure, similar to how CNNs adapt to image data. This hypothesis is substantiated by the findings in Section 5.3, where we show that the subnetworks acquire periodicity—a critical characteristic for Modular Addition tasks. This periodic structure in Modular Addition tasks can be interpreted analogously to the local structures in image tasks captured by CNNs. These results provide insight into the importance of a network's inner structure for achieving generalization.

To make this point clearer, we have revised the following sections of our paper:

• Section 5.3: We have updated this section to make it clearer that grokking tickets acquire task-specific, beneficial structures. The revised text highlights how these structures contribute to generalization.
• Appendix K: Related to Q7 and Q8, we have added the visualizations of weight matrices and masks of the grokking tickets after generalization. This addition provides a clearer illustration of how grokking tickets exhibit task-relevant structures (periodicity).
• Abstract and Introduction: We have updated these sections to reflect the changes and emphasize our motivation and the key insights regarding the role of task-adaptive structures in grokking tickets.

We believe these revisions address the reviewer's concerns and offer a clear explanation of why subnetworks (grokking tickets) are sufficient to account for the phenomenon of grokking.

[1] Neyshabur, Towards Learning Convolutions from Scratch, https://arxiv.org/abs/2007.13657

Q8

Can the authors provide some illustrations-for example, weight heatmaps or connectivity graphs of subnetworks-probing the evolution through multiple training phases?

We have included visualizations of the weight matrices after generalization and the masks of the grokking ticket in Figure 22 (Appendix K). In these results, periodic patterns are observed in the weight matrices, and the masks of the grokking ticket reflect these characteristics. This indicates that the grokking ticket has acquired structures that are beneficial for the task (Modular Addition).

These visualizations additionally demonstrate the role of grokking tickets in delayed generalization through their ability to acquire periodic structures.

Q7

Are the authors willing to provide more analysis into why such classical PaI methods do not elicit generalization?

We have provided additional analysis in Appendix K, specifically in Figure 23, where we compare the masks (structures) obtained by pruning-at-initialization (PaI) methods such as Random, GraSP, SNIP, and SynFlow. Unlike the results of the grokking ticket shown in Figure 22, these PaI methods do not exhibit periodic structures.

This comparison highlights the superiority of the grokking ticket in acquiring structures that are more conducive to the Modular Addition task, further emphasizing its advantage over traditional PaI methods.