# NeuralGrok: Accelerate Grokking by Neural Gradient Transformation

arXiv: https://arxiv.org/abs/2504.17243 · HTML: https://mlsanity.com/paper/2504.17243/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 14 of its 20 references that are in the map.
- First posted on arXiv 24 Apr 2025.

## Builds on (most cited first, 14 of 14)

- [Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets](https://mlsanity.com/paper/2201.02177/) (2022)
- [Train faster, generalize better: Stability of stochastic gradient descent](https://mlsanity.com/paper/1509.01240/) (2015)
- [Towards Understanding Grokking: An Effective Theory of Representation Learning](https://mlsanity.com/paper/2205.10343/) (2022)
- [Grokking as the Transition from Lazy to Rich Training Dynamics](https://mlsanity.com/paper/2310.06110/) (2023)
- [The Slingshot Mechanism: An Empirical Study of Adaptive Optimizers and the Grokking Phenomenon](https://mlsanity.com/paper/2206.04817/) (2022)
- [On Generalization Error Bounds of Noisy Gradient Methods for Non-Convex Learning](https://mlsanity.com/paper/1902.00621/) (2019)
- [Deep Networks Always Grok and Here is Why](https://mlsanity.com/paper/2402.15555/) (2024)
- [Unifying Grokking and Double Descent](https://mlsanity.com/paper/2303.06173/) (2023)
- [Grokking at the Edge of Numerical Stability](https://mlsanity.com/paper/2501.04697/) (2025)
- [Grokfast: Accelerated Grokking by Amplifying Slow Gradients](https://mlsanity.com/paper/2405.20233/) (2024)
- [On the Overlooked Pitfalls of Weight Decay and How to Mitigate Them: A Gradient-Norm Perspective](https://mlsanity.com/paper/2011.11152/) (2020)
- [Deep Grokking: Would Deep Neural Networks Generalize Better?](https://mlsanity.com/paper/2405.19454/) (2024)
- [Progress Measures for Grokking on Real-world Tasks](https://mlsanity.com/paper/2405.12755/) (2024)
- [Inconsistency, Instability, and Generalization Gap of Deep Neural Network Training](https://mlsanity.com/paper/2306.00169/) (2023)

## Built on by (most cited first, 2 of 2)

- [Structure-Specific Representational Priors Causally Control the Grokking Delay](https://mlsanity.com/paper/2607.04333/) (2026)
- [Is Grokking Worthwhile? Functional Analysis and Transferability of Generalization Circuits in Transformers](https://mlsanity.com/paper/2601.09049/) (2026)

## Abstract (arXiv)

Grokking is proposed and widely studied as an intricate phenomenon in which generalization is achieved after a long-lasting period of overfitting. In this work, we propose NeuralGrok, a novel gradient-based approach that learns an optimal gradient transformation to accelerate the generalization of transformers in arithmetic tasks. Specifically, NeuralGrok trains an auxiliary module (e.g., an MLP block) in conjunction with the base model. This module dynamically modulates the influence of individual gradient components based on their contribution to generalization, guided by a bilevel optimization algorithm. Our extensive experiments demonstrate that NeuralGrok significantly accelerates generalization, particularly in challenging arithmetic tasks. We also show that NeuralGrok promotes a more stable training paradigm, constantly reducing the model's complexity, while traditional regularization methods, such as weight decay, can introduce substantial instability and impede generalization. We further investigate the intrinsic model complexity leveraging a novel Absolute Gradient Entropy (AGE) metric, which explains that NeuralGrok effectively facilitates generalization by reducing the model complexity. We offer valuable insights on the grokking phenomenon of Transformer models, which encourages a deeper understanding of the fundamental principles governing generalization ability.