NeuralGrok: Accelerate Grokking by Neural Gradient Transformation

• URL: http://arxiv.org/abs/2504.17243v2
• Date: Fri, 25 Apr 2025 03:44:09 GMT
• Title: NeuralGrok: Accelerate Grokking by Neural Gradient Transformation
• Authors: Xinyu Zhou, Simin Fan, Martin Jaggi, Jie Fu,
• Abstract summary: We propose NeuralGrok, a gradient-based approach that learns an optimal gradient transformation to accelerate generalization of transformers in arithmetic tasks.<n>Our experiments demonstrate that NeuralGrok significantly accelerates generalization, particularly in challenging arithmetic tasks.<n>We also show that NeuralGrok promotes a more stable training paradigm, constantly reducing the model's complexity.
• Score: 54.65707216563953
• License: http://creativecommons.org/licenses/by-sa/4.0/
• Abstract: 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.

Related papers

• Exact, Tractable Gauss-Newton Optimization in Deep Reversible Architectures Reveal Poor Generalization [52.16435732772263]
  Second-order optimization has been shown to accelerate the training of deep neural networks in many applications. However, generalization properties of second-order methods are still being debated. We show for the first time that exact Gauss-Newton (GN) updates take on a tractable form in a class of deep architectures.
  arXiv  Detail & Related papers  (2024-11-12T17:58:40Z)
• What Improves the Generalization of Graph Transformers? A Theoretical Dive into the Self-attention and Positional Encoding [67.59552859593985]
  Graph Transformers, which incorporate self-attention and positional encoding, have emerged as a powerful architecture for various graph learning tasks. This paper introduces first theoretical investigation of a shallow Graph Transformer for semi-supervised classification.
  arXiv  Detail & Related papers  (2024-06-04T05:30:16Z)
• Generalization of Scaled Deep ResNets in the Mean-Field Regime [55.77054255101667]
  We investigate emphscaled ResNet in the limit of infinitely deep and wide neural networks. Our results offer new insights into the generalization ability of deep ResNet beyond the lazy training regime.
  arXiv  Detail & Related papers  (2024-03-14T21:48:00Z)
• Linear Time GPs for Inferring Latent Trajectories from Neural Spike Trains [7.936841911281107]
  We propose cvHM, a general inference framework for latent GP models leveraging Hida-Mat'ern kernels and conjugate variational inference (CVI) We are able to perform variational inference of latent neural trajectories with linear time complexity for arbitrary likelihoods.
  arXiv  Detail & Related papers  (2023-06-01T16:31:36Z)
• Understanding the Generalization Ability of Deep Learning Algorithms: A Kernelized Renyi's Entropy Perspective [11.255943520955764]
  We propose a novel information theoretical measure: kernelized Renyi's entropy. We establish the generalization error bounds for gradient/Langevin descent (SGD/SGLD) learning algorithms under kernelized Renyi's entropy. We show that our information-theoretical bounds depend on the statistics of the gradients, and are rigorously tighter than the current state-of-the-art (SOTA) results.
  arXiv  Detail & Related papers  (2023-05-02T01:17:15Z)
• Learning Non-Vacuous Generalization Bounds from Optimization [8.294831479902658]
  We present a simple yet non-vacuous generalization bound from the optimization perspective. We achieve this goal by leveraging that the hypothesis set accessed by gradient algorithms is essentially fractal-like. Numerical studies demonstrate that our approach is able to yield plausible generalization guarantees for modern neural networks.
  arXiv  Detail & Related papers  (2022-06-09T08:59:46Z)
• Thermodynamics-informed graph neural networks [0.09332987715848712]
  We propose using both geometric and thermodynamic inductive biases to improve accuracy and generalization of the resulting integration scheme. The first is achieved with Graph Neural Networks, which induces a non-Euclidean geometrical prior and permutation invariant node and edge update functions. The second bias is forced by learning the GENERIC structure of the problem, an extension of the Hamiltonian formalism, to model more general non-conservative dynamics.
  arXiv  Detail & Related papers  (2022-03-03T17:30:44Z)
• Orthogonal Graph Neural Networks [53.466187667936026]
  Graph neural networks (GNNs) have received tremendous attention due to their superiority in learning node representations. stacking more convolutional layers significantly decreases the performance of GNNs. We propose a novel Ortho-GConv, which could generally augment the existing GNN backbones to stabilize the model training and improve the model's generalization performance.
  arXiv  Detail & Related papers  (2021-09-23T12:39:01Z)
• Fractal Structure and Generalization Properties of Stochastic Optimization Algorithms [71.62575565990502]
  We prove that the generalization error of an optimization algorithm can be bounded on the complexity' of the fractal structure that underlies its generalization measure. We further specialize our results to specific problems (e.g., linear/logistic regression, one hidden/layered neural networks) and algorithms.
  arXiv  Detail & Related papers  (2021-06-09T08:05:36Z)
• The Neural Tangent Kernel in High Dimensions: Triple Descent and a Multi-Scale Theory of Generalization [34.235007566913396]
  Modern deep learning models employ considerably more parameters than required to fit the training data. Whereas conventional statistical wisdom suggests such models should drastically overfit, in practice these models generalize remarkably well. An emerging paradigm for describing this unexpected behavior is in terms of a emphdouble descent curve. We provide a precise high-dimensional analysis of generalization with the Neural Tangent Kernel, which characterizes the behavior of wide neural networks with gradient descent.
  arXiv  Detail & Related papers  (2020-08-15T20:55:40Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.