Grokking Tickets: Lottery Tickets Accelerate Grokking

Minegishi Gouki, Yusuke Iwasawa, Yutaka Matsuo

Primary Area: visualization or interpretation of learned representations

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Keywords: grokking, lottery ticket hypothesis, weight norm, pruning, generalization

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.

Abstract:

Grokking is one of the most surprising puzzles in neural network generalization: a network first reaches a memorization solution with perfect training accuracy but poor generalization, but with further training, it reaches a perfectly generalizable solution.
We aim to analyze the mechanism of grokking from the lottery ticket hypothesis, identifying the process to find the lottery tickets (good sparse subnetworks) as the key to describing the transitional phase between memorization and generalization.
Firstly, with the lottery tickets identified via the magnitude pruning after perfect generalization, we show that the lottery tickets drastically accelerate grokking compared to the dense networks on various configurations (MLP and Transformer, and an arithmetic and image classification task).
We also show that the speedup is significant even when compared with the dense networks with the same weight norm as the lottery tickets.
Besides, the speedup only happens when training ``good'' subnetworks are identified at the generalization solution. Specifically, speedup does not happen when using tickets identified at the memorization solution or transition between memorization and generalization or when pruning networks at the initialization (Random pruning, Grasp, SNIP, and Synflow).
The results indicate that the weights norm of network parameters is not enough to explain the process of grokking, but the importance of finding good subnetworks to describe the transition from memorization to generalization.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 6981