A Hessian View of Grokking in Mathematical Reasoning

Authors: zhenshuo  zhang
jerry w liu
christopher  ré
hongyang r zhang
Venue: NeurIPS_cc_2024_Workshop_MATH-AI
Published: 2024-09-26 06:49:06

Abstract: Mathematical reasoning is a central problem in developing more intelligent language models. An intriguing phenomenon observed in mathematical arithmetics is grokking, where the training loss of a transformer model stays near zero for an extended period until the validation loss finally reduces to near zero. In this work, we approach this phenomenon through a view of the Hessian of the loss surface. The Hessian relates to the generalization properties of neural networks as it can capture geometric properties of the loss surface, such as the sharpness of local minima. We begin by noting in our experiments that high weight decay is essential for grokking to occur in several arithmetic tasks (trained with a GPT-2 style transformer model). However, we also find that the training loss is highly unstable and exhibits strong oscillations. To address this issue, we consider adding regularization to the Hessian by injecting isotropic Gaussian noise to the weights of the transformer network, and find that this combination of high weight decay and Hessian regularization can smooth out the training loss during grokking. We also find that this approach can accelerate the grokking stage compared to existing methods by at least $50\%$ measured on seven arithmetic tasks. Finally, to understand the precise cause of grokking, we consider a Hessian-based measurement for multi-layer networks and find that this measure yields non-vacuous estimates of the generalization errors observed in practice. We hope these empirical findings can facilitate future research towards understanding grokking (and generalization) in mathematical reasoning.

Main paper: 
 introduction: 
 Mathematical reasoning in large neural networks [1] is a central issue in the design of more intelligent, interactive language models, especially in scenarios that require precise, step-by-step logical operations. An intriguing phenomenon that has been observed in arithmetic tasks is grokking [19] , where a transformer model exhibits a delayed yet sudden generalization of training data even as the training curve has converged. In this paper, we analyze grokking behavior in arithmetic tasks by examining the Hessian of the transformer model's loss surface.Power et al. [19] demonstrate that for a range of modular arithmetic tasks, grokking can occur on two-layer transformers, whereby the training loss remains near zero for a long period until the validation loss also converges to zero. However, the training loss can exhibit dramatic oscillations. To motivate this work, we begin by applying regularization methods to SGD and examine their effect on grokking. First, we find that high weight decay is crucial for grokking in arithmetic tasks. In other words, we find that the validation accuracy does not increase to near perfect when weight decay is moderate (Fig. 1a and 1b ). Moreover, even after adding high decay, the training curve still exhibits notable variations, leading to unstable training (with training accuracy reduced to zero; see Fig. 1c ). To address these issues, we consider regularization methods that regularize the Hessian. These methods are known to improve generalization by reducing the sharpness of solutions in the loss surface [6, 22] . In particular, we consider regularizing the Hessian by first adding noise to the weights of the transformer before computing its gradient. We observe that this noise injection provides an approximately unbiased estimate of the trace of the Hessian. Surprisingly, we find that, along with high weight decay, this noise injection algorithm can now smooth out the oscillations in the training curve. Moreover, we also find that this Hessian regularization can reduce the number of grokking steps during training by at least 28%, measured on three arithmetic tasks.Finally, to understand these results, we develop a preliminary theoretical analysis. We examine a Hessian-based generalization measure motivated by PAC-Bayes analysis [17, 11, 2] . We find that by measuring a Hessian-vector product on the weight space, we can provide a non-vacuous estimate of the generalization errors. Note that the phenomenon of delayed generalization is known since classical works on boosting [3] . Our contribution is to provide a Hessian view of this phenomenon since the Hessian can be measured from data.In summary, we find that by using high weight decay and noise injection to regularize the Hessian, we can effectively reduce the instability that has commonly been observed for training transformers on arithmetic tasks. Second, this combined regularization can further reduce the number of grokking steps. Third, we develop a Hessian-based measurement that can give a non-vacuous estimate of the generalization error. We hope these findings can facilitate future research on understanding grokking in the mathematical reasoning of large models. 
regularization of loss surface hessian: 
 Previous works have indicated that grokking requires an appropriate choice of weight decay [18, 14] . However, using weight decay alone can still lead to oscillations of the training curve. To address this issue, we explore an alternative approach, where we regularize the loss Hessian matrix, which can provide more fine-grained control on the loss surface, such as sharpness. To instantiate the Hessian regularization, we add a random noise variable to the weight matrices of a transformer network. In particular, let ℓ(f W (x), y) denote the loss of a neural net f W (parameterized by W ), given an input pair (x, y). Let U be a random sample from an isotropic Gaussian (with the same dimension as W ), whose variance has been scaled by σ 2 . We consider the following noise injection update:EQUATIONfor some learning rate η. In particular, σ 2 determines the level of regularization in this procedure.To see that this update regularizes the Hessian, we notice that equation ( 1) is equivalent to applying SGD to the stochastic optimization objective ofE U [ℓ(f W +U (x), y)] ≈ ℓ(f W (x), y) + σ 2 2 ∇ 2 ℓ(f W (x), y) + O(σ 3 ).In practice, we add the noise injection along both the positive and negative directions of U . This helps eliminate the variance that appears from the first-order Taylor's expansion term above [22] .3 Experimental Results 
results on arithmetic tasks: 
 Our experimental setup follows the work of Power et al. [19] . We focus on evaluating the grokking phenomena with arithmetic tasks, which correspond to equations of the form (a • b) mod p = c. "a," "•," "b," "=," and "c" are separated tokens, where "c" is the prediction goal. For each task, we generate a ∈ [p] and b ∈ [p], resulting in a total of p 2 (in our settings, p = 97) unique data. Specifically, we select a 2 + ab + b 2 = c to illustrate our findings, and more experiments of different equations are shown in Appendix A. We compare our regularization method to naive SGD and SAM [6] . SAM is based on a constrained minimax optimization formulation that penalizes the worst-case perturbations.#1: Stabilizing the training curves. We observed that all the approaches can induce grokking in our experiments, as shown in Figure 2 . After the training loss converges, we observe a sudden increase in validation accuracy to over 99%. Although all these methods can exhibit stable training before convergence, we find that during the grokking phase, SGD experiences sharp fluctuations in the training curve, with training accuracy dropping close to zero, which also caused the validation accuracy to drop to nearly zero. By contrast, both SAM and our noise-injection method can maintain stable training loss values during the grokking phase, which avoids dramatic fluctuations. #2: Reducing the steps of grokking steps. We also report the comparison of the number of training steps from the point where the training accuracy has converged to near 100% to the point where validation accuracy convergences to near 100%. We report the results for different approaches in Table 1 . We observe that our approach requires fewer steps than SGD. In some tasks, our approach doesn't even need grokking steps to generalize. We also note that SAM, which penalizes the largest eigenvalue of the Hessian, requires more steps than noise injection. We also evaluate our findings on the Needle-in-a-Haystack task, following the setting of Zhong et al. [23] . Specifically, we have an input sequence[m 1 , c 1 , m 2 , c 2 , ..., m k , c k , m u ],where m i are different markers and c i are corresponding values. The last element is a marker m u , u ∈ [1, k] which indicates the goal marker. The model is trained to learn to search for the marker in the previous sequence and give the corresponding value c u .More surprisingly, we observe that the grokking phenomenon does not occur when using SGD, although the training curve also experiences slight fluctuations, and the validation accuracy remains low. The number of grokking steps of SAM and noise injection are 12288 and 5632, respectively: see Figure 3 .  
nonvacuous generalization error estimates with hessian: 
 Toward rigorously understanding the above empirical results, we consider the PAC-Bayes analysis framework. In particular, we consider a linear PAC-Bayes bound [16, 2] , which holds with probability 1δ for any δ > 0:EQUATIONAbove, Q is the posterior hypothesis distribution of the learning algorithm. P is the prior distribution of the learning algorithm. C > 0 is an upper bound on the loss value. L and L refer to the expected and empirical risks. For example, in the context of fine-tuning foundation models, one may view P as the weight of the pretrained model (plus some small perturbations), and Q is the fine-tuned model weight [11] . For a complete statement, see Lemma B.2 in Appendix B.Derivation of a Hessian-based measure: Let U ∼ Q be a random variable drawn from a posterior distribution Q. We are interested in the perturbed loss, ℓ Q (f U (x), y), which is the expectation of ℓ(f U (x), y) over U . Using Taylor's expansion, we get thatEQUATIONwhere Σ i is the population covariance matrix of the perturbation added to layer i, and ∇ 2 i is the Hessian matrix with respect to the weights at layer i of f W . See Lemma B.1 in Appendix B for the complete statement of this result.Based on equation (3), next, we apply the PAC-Bayes bound from equation ( 2) to an L-layer transformer neural network f W parameterized by W . We note that the KL divergence between the prior and posterior distributions, which are both Gaussian, is equal toL i=1 Σ -1 i , v i v ⊤ i, where v i is the distance between the initialized weight and the trained weight at layer i.Next, it remains to minimize the sum of the Hessian estimate, and the above KL divergence in the PAC-Bayes bound will lead to a different covariance matrix for every layer. Let ∇ 2 i + denote the truncated Hessian matrix where we set the negative eigenvalues of ∇ 2 i to zero. We have thatEQUATIONBy applying equations ( 4) and (3) back to equation (2), and minimizing over β, we will derive an upper bound on the generalization error (between L(f W ) and L(f W )) that is equal to:EQUATIONwhere n is the size of the sample set and D is the unknown distribution where the samples are drawn.Having introduced the Hessian measure, we now report the results from measuring the above α in the grokking experiments and compare α with the empirically observed generalization errors. The results are shown in Figure 4 below. We can see that α now gives a nonvacuous upper bound on the generalization error. Importantly, while this can also be achieved with standard methods such as k-fold cross-validation, the Hessian can reveal more structures (e.g., sharpness) of loss surfaces. 5 Related WorkGrokking. The grokking phenomenon, first proposed by Power et al. [19] , illustrates that with continued training over several epochs, the validation loss eventually decreases and converges after training loss does not decrease further after converging. Extending the study of grokking, Liu et al. [15] conducted experiments across diverse datasets, including images, language, and graphs, expanding the area of grokking. Previous research predominantly focused on how training configurations influence grokking. Davies et al. [5] explored the relationship between grokking and double descent concerning pattern learning. Huang et al. [9] examined the impact of model and dataset sizes on grokking. Nanda et al. [18] highlighted the critical role of weight decay in grokking, noting that insufficient weight decay prolongs the process of grokking. Thilak et al. [20] linked grokking to the slingshot mechanism, interpreting it as a form of implicit regularization. More recently, Lee et al. [14] introduced Grokfast, a method designed to accelerate grokking by amplifying the gradients' low-frequency components. Theoretical investigations of why grokking occurs have recently been studied [21] . In particular, Xu et al. [21] provably demonstrate grokking in two-layer ReLU networks trained by gradient descent on XOR cluster data where a constant fraction of the training labels are flipped.Hessian and optimization algorithms. Historical studies on second-order methods for training multi-layer networks primarily focus on optimization methods like Newton or quasi-Newton and employ the Hessian matrix to adjust learning rates [13, 12, 4] . Although they estimate the spectral information of the Hessian by computing Hessian-vector products, they do not explore the dynamics of the Hessian throughout training. In the Neural Tangent Kernel (NTK) analysis [10] , the Hessian matrix is treated as a random features matrix, which remains fixed during training. The estimation of spectral density through Stochastic Lanczos Quadrature is discussed by Ghorbani et al. [7] . Additionally, Grosse et al. [8] have investigated scaling up influence functions in large neural networks, which includes innovative techniques for computing Hessian-inverse vector products. 
conclusion: 
 This paper explored the issue of generalization in arithmetic tasks. We analyze the grokking phenomenon through a view of the Hessian matrix of the loss surface. We find that using a high weight decay and noise injection can smooth out the oscillations commonly observed in SGD training of arithmetic tasks. Another benefit of this regularization is that we could accelerate the grokking stage, reducing the number of training steps required for model generalization. Finally, we find that a Hessian-based measurement can give a nonvacuous estimate of the generalization errors in various modular arithmetic tasks.  
 
 Appendix: First, we state the result of Taylor's expansion of the perturbed loss. Lemma B.1. For any i = 1, 2, • • • , L, let U i ∈ R didi-1 be a random vector sampled from a Gaussian distribution with mean zero and variance Σ i . Let the posterior distribution Q be centered at W i and perturbed with an appropriately reshaped U i at every layer. Then, there exists a fixed value C 1 > 0 that does not grow with n, such that the following holds for any x ∈ X and y ∈ {1, . . . , k}:EQUATIONNext, we state the PAC-Bayes bound, which can be found in the PAC-Bayes literature (e.g., [16, 2] ). Lemma B.2. Suppose the loss function ℓ(x, y) lies in a bounded range [0, C] given any x ∈ X with label y. For any β ∈ (0, 1) and δ ∈ (0, 1], with probability at least 1δ, the following holdsEQUATIONThe 4th Workshop on Mathematical Reasoning and AI at NeurIPS 2024 (MATH-AI 24). 


Decision: The ideas presented in the work are interesting. However, as the reviewers point out, the empirical evidence is a little bit lacking and would benefit from more extensive experiments. We suggest that the authors consider including additional experiments to further validate their claims in the camera-ready version.

Reviews: review: title: title: The novelty and contribution of this work is limited
main_review: main_review: This paper investigates the grokking behavior that occurs in arithmetic tasks. The authors propose to address this issue by injecting Gaussian noise into the transformer network, which helps accelerate the grokking stage.

Strengths:
The paper attempts to address the grokking phenomenon which is a common behavior in training, so the technique in this work might be applicable in other settings as well.

Weaknesses:
My concerns regarding this paper are its novelty and its significance. 1) The noise injection technique during the training phase of Deep Learning models is well-known, so the novelty of this work is lacking and its contribution is limited. 2) The technique in this work addresses the training phase but does not really improve the performance of the model. And so, its results are not really interesting. It would be more exciting if the work proposes a method to improve the arithmetic capabilities of LLM, which is still a problem for popular LLMs such as GPT-4o or o1.
title: title: Insightful study but weak relation to reasoning
main_review: main_review: This paper presents a Hessian regularization, which is shown to smooth out the oscillations and accelerate the grokking stage. The authors also find a Hessian-based measurement to estimate the generalization errors. 

The idea of making use of Hessian’s ability to capture geometric properties of the loss surface is insightful. Starting from this, the proposals of the new regularization and the new measurement are both novel. 
The experimental setup, following some existing work, is also confirmatory. With the hyperparameters clearly listed in the appendix, the discussion around the experimental results is convincing. 

One concern is that the paper focuses on modulo arithmetic tasks only, and the generalizability to other mathematical reasoning, the focus of the workshop, is unclear ("analyzing complex information, identifying patterns and relationships, and drawing logical conclusions from evidence"). 

Overall: insightful study but weak relation to reasoning
title: title: The paper is well-aligned with the workshop's theme and clearly presents its contributions, but could benefit from broader analysis across arithmetic tasks and more clarity on computational efficiency and methodology details.
main_review: main_review: **Strength**

_Alignment with the workshop_

The paper aligns really well with the theme of the workshop. The paper proposes regularization of loss surface hessian (equation 1, shows the weight update) to speed up grokking in arithmetic tasks and training stability (reduce variation in training accuracy in the later part of training)

_Writing Style_
* The paper is written well and conveys the core idea without confusion
* The experimental results are presented in a straightforward manner, highlighting the effectiveness of noise injection in stabilizing training.
* The conclusions effectively tie together theoretical insights with empirical observations, offering valuable directions for future research on grokking.


---

**Weakness**


_Limited analysis_

* Though the proposed approach is very effective in increasing grokking speed, the paper only analyzes and shows results on 3 different arithmetic equations. Therefore it is not clear if the approach can generalize for all arithmetic equations.
* Also the paper does not analyze the computation efficiency, the proposed approach requires computing gradients twice with noice added and subtracted from weights (so two forward pass and two backward passes are needed) 

_Clarity_
* Equation 1 shows the weight update with loss surface hessian regularization, but the equation doesnt show any weight decay, therefor it is not very clear whether weight decay is used for results shown in Table1.