Lucas Prieto @lucas-prieto.bsky.social · 10/01/2025Thank you to my co-authors Tolga Birdal, Pedro Mediano and Melih Barsbey for all their help! 000
Lucas Prieto @lucas-prieto.bsky.social · 10/01/2025This simple intervention removes the delay in generalization characteristic of grokking, leading to train and test performance increasing in tandem! 100
Lucas Prieto @lucas-prieto.bsky.social · 10/01/2025Since the gradient does not fully align with the NLM direction, we tried ignoring the component of the gradient that aligns with the direction of the weights, preserving only the part of the gradent that is orthogonal to the NLM direction. 100
Lucas Prieto @lucas-prieto.bsky.social · 10/01/2025For homogeneous models, naïvely scaling the logits corresponds to scaling the weights along their current direction, and we observe that after fitting the trainign data, MLP and transformer gradients do align with the direction of the weights. 100
Lucas Prieto @lucas-prieto.bsky.social · 10/01/2025Why do tasks like modular addition lead to SC? We find ease of overfitting to be a crucial aspect. Beyond the point of 100% training accuracy, models can decrease the CE loss by scaling the logits without changing the decision boundary. We call this Naïve Loss Minimization (NLM). 100
Lucas Prieto @lucas-prieto.bsky.social · 10/01/2025This simple intervention prevents SC and leads to grokking without weight decay, with weight norms increasing during generalization. 100
Lucas Prieto @lucas-prieto.bsky.social · 10/01/2025To validate this hypothesis we replace the exponential in the Softmax with a softer ramp function s(x). 100
Lucas Prieto @lucas-prieto.bsky.social · 10/01/2025Looking at modular addition, we notice that models can start grokking without weight decay but generalization stops suddenly when a large fraction of the samples face SC. As expected, this point is delayed when we increase floating point precision. 100
Lucas Prieto @lucas-prieto.bsky.social · 10/01/2025These "absorption" errors, lead to zeros in the loss and the gradient, putting an end to learning! 100
Lucas Prieto @lucas-prieto.bsky.social · 10/01/2025Why do we not see grokking by default when using cross-entropy loss? The surprising answer is a specific kind of floating point error in the Softmax we call Softmax Collapse (SC)! 100
Lucas Prieto @lucas-prieto.bsky.social · 10/01/2025paper: arxiv.org/pdf/2501.04697 code: github.com/LucasPrietoA...github.comGitHub - LucasPrietoAl/grokking-at-the-edge-of-numerical-stabilityContribute to LucasPrietoAl/grokking-at-the-edge-of-numerical-stability development by creating an account on GitHub. 100
Lucas Prieto @lucas-prieto.bsky.social · 10/01/2025I am excited to share our new paper: “Grokking at the Edge of Numerical stability”! We show that floating point errors in the Softmax play a surprising role in grokking, explaining among other things, why weight decay seems necessary for grokking in most cases! 🧵 181
