Grokking behavior in networks linked to structure of Latin squares
Mapping the Solution-Set Structure of Grokking on Latin Squares
Machine Learning
Summary
Grokking happens when a neural network suddenly shifts from just memorizing to truly understanding a task after some delay. The authors study this effect using mathematical objects called Latin squares, which represent structured problems, to see when networks can generalize well from partial data. They find that for certain structured tasks tied to groups, networks almost always generalize well, but for random Latin squares, good generalization becomes rare. The work helps explain when and why grokking happens by showing how problem structure influences the learning outcomes.
What this means in practice
- •For machine learning engineers: Design network training schemes knowing when structured tasks will reliably generalize from partial data to improve model robustness.
- •For cryptography developers: Use insights on Latin squares and group-based structures to assess cryptographic algorithms that rely on combinatorial designs for security.
Authors
Zihan Gu, Lulu Zhao, Tianyi Ma, Hua Zhang, Wenbo Su, Yanan Cao, Yue Hu
Abstract
Grokking is a delayed transition from memorisation to generalisation. A network first fits observed examples while remaining near chance on held-out inputs, then generalises abruptly. Prior work has focused mainly on this delay, finding that the network's main loss becomes weight decay after the fitting phase. We name the solution to this optimization problem as the minimum-norm endpoint, and study which tasks make all minimum-norm endpoints generalise from partial observations. Specifically, we model pairwise input tasks as edge colourings of $K_{q,q}$, with Latin colourings as our structured family, and classify each observed instance as all-good, mixed, or no-good by its endpoint set. Furthermore, we prove two contrasting results: (1) Finite-abelian-group families are asymptotically all-good from partial observations; (2) Under a uniform random Latin square, all-good tasks become asymptotically rare for observation-based learners, leaving mixed or no-good behaviour. Moreover, experiments show that relabellings of addition generalise despite their dense polynomial representations. Finally, on tasks of composite order, finite-training results suggest that predictors learn a coarse part of the output before the full label.