The Slingshot Mechanism: An Empirical Study of Adaptive Optimizers and the Grokking Phenomenon

Introduction

Recently, the grokking phenomenon was proposed by , in the context of studying the optimization and generalization aspects in small, algorithmically generated datasets. Specifically, grokking refers to a sudden transition from chance level validation accuracy to perfect generalization, long past the point of perfect training accuracy, i.e., Terminal Phase of Training (TPT). This curious behavior contradicts the common belief of early stopping in the overfitting regimes, and calls for further understandings of the generalization behavior of deep neural networks.

In the literature, it has been suggested that in some scenarios, marginal improvements in validation accuracy appears in TPT, which seem to directly support grokking. For example, it has been shown in that gradient descent on logistic regression problems converges to the maximum margin solution, a result that has been since extended to cover a wider setting . A key finding in shows that when training on linearly separable data with gradient descent using logistic regression, the classifier’s margin slowly improves at a rate of $\mathcal{O}(\frac{1}{\log t})$, while the weight norm of the predictor layer grows at a rate of $\mathcal{O}(t)$, where $t$ is the number of training steps. While specified for gradient descent, Wang et al. showed that similar results also hold for adaptive optimizers (such as Adam and RMSProp). Taking these results into consideration, one could reasonably hypothesise that deep nonlinear networks could benefit from longer training time, even after achieving zero errors on the training set.

In this paper, we provide in depth empirical analyses to the mechanism behind grokking. We find that the phenomenology of grokking differs from those predicted by in several key aspects. To be concrete, we find that grokking occurs during the onset of another intriguing phenomenon directly related to adaptive gradient methods (see Algorithm 1 for a generic description of adaptive gradient methods). In particular, leveraging the basic setup in , we make the following observations:

1. During the TPT, training exhibits a cyclic behaviour between stable and unstable regimes. A prominent artifact of this behaviour can be seen in the norm of a model’s last layer weights, which exhibits a cyclical behavior with distinct, sharp phase transitions that alternate between rapid growth and plateaus over the course of training.

2. The norm grows rapidly sometime after the model has perfect classification accuracy on training data. A sharp phase transition then occurs when the model missclassifies training samples. This phase change is accompanied by a sudden spike in training loss, and a plateau in the norm growth of the final classification layer.

3. The features (pre-classification layer) show rapid evolution as the weight norm transitions from rapid growth to a growth plateau, and change relatively little at the norm growth phase.

4. Phase transitions between norm growth and norm plateau phases are typically accompanied by a sudden bump in generalization as measured by classification accuracy on a validation set, as observed in a dramatic fashion in .

5. It is empirically observed that grokking as reported in almost exclusively happens at the onset of Slingshots, and is absent without it.

We denote the observations above as the Slingshot Effect, which is defined to be the full cycle starting from the norm growth phase, and ending in the norm plateau phase. And empirically, a single training run typically exhibits multiple Slingshot Effects. Moreover, while grokking as described in might be data dependent, we find that the Slingshot Mechanism is pervasive, and can be easily reproduced in multiple scenarios, encompassing a variety of models (Transformers and MLPs) and datasets (both vision, algorithmic and synthetic datasets). Since we only observe Slingshot Effects when training classification models with adaptive optimizers, our work can be seen as empirically characterizing an implicit bias of such optimizers. Finally, while our observations and conclusions hold for most variants of adaptive gradient methods, we focus on Adam in the main paper, and relegate all experiments with additional optimizers to the appendix.

The findings in this paper have both theoretical and practical implications that go beyond characterizing Grokking. A prominent feature of the Slingshot Mechanism is the repeating phase shifts between stable and unstable training regimes, where the unstable phase is characterized by extremely large gradients, and spiking training loss. Furthermore, we find that learning at late stages of training have a cyclic property, where non trivial feature adaptation only takes place at the onset of a phase shift. From a theoretical perspective, this is contradictory to common assumptions made in the literature of convergence of adaptive optimizers, which typically require $L$ smooth cost functions, and bounded stochastic gradients, either in the $L_{2}$ or $L_{\infty}$ norm, decreasing step sizes and stable convergence . From the apparent generalization benefits of Slingshot Effects, we cast doubt on the ability of current working theories to explain the Slingshot Mechanism. Practically, our work presents additional evidence for the growing body of work indicating the importance of the TPT stage of training for optimal performance .

In an era where the sheer size of models are quickly becoming out of reach for most practitioners, our work suggest focusing on improved methods to prevent excessive norm growth either implicitly through Slingshot Effects or through other forms of explicit regularization or normalization.

Related Work

The Slingshot Mechanism we uncover here is reminiscent of the catapult mechanism described in al. . Lewkowycz et al. show that loss of a model trained via gradient descent with an appropriately large learning rate shows a non-monotonic behavior —the loss initially increases and starts decreasing once the model "catapults" to a region of lower curvature —early in training. However, the catapult phenomenon differs from Slingshot Effects in several key aspects. The catapult mechanism is observed with vanilla or stochastic gradient descent unlike the Slingshot Mechanism that is seen with adaptive optimizers including Adam and RMSProp . Furthermore, the catapult phenomenon relates to a large initial learning rate, and does not exhibit a repeating cyclic behavior. More intriguingly, Slingshot Effects only emerge late in training, typically long after the model reaches perfect accuracy on the training data.

Cohen et al. describe a "progressive sharpening" phenomenon in which the maximum eigenvalue of the loss Hessian increases and reaches a value that is at equal to or slightly larger than $2/\eta$ where $\eta$ is the learning rate. This "progressive sharpening" phenomenon leads to model to enter a regime Cohen et al. call Edge of Stability where-in the model shows non-monotonic training loss behavior over short time spans. Edge of Stability is similar to the Slingshot Mechanism in that it is shown to occur later on in training. However, Edge of Stability is shown for full-batch gradient descent while we observe Slingshot Mechanism with adaptive optimizers, primarily Adam or AdamW .

As noted above, the Slingshot Mechanism emerges late in training, typically longer after the model reaches perfect accuracy and has low loss on training data. The benefits of continuing to training a model in this regime has been theoretically studied in several works including . Soudry et al. show that training a linear model on separable data with gradient using the logistic loss function leads to a max-margin solution. Furthermore Soudry et al. prove that the loss decreases at a rate of $O(\frac{1}{t})$ while the margin increases much slower $O(\frac{1}{\log t})$, where $t$ is the number of training steps. Soudry et al. also note that the weight norm of the predictor layer increases at a logarithmic rate, i.e., $O({\log(t)})$. Lyu and Li generalize the above results to homogeneous neural networks trained with exponential-type loss function and show that loss decreases at a rate of $O(1/t(\log(t))^{2-2/L})$. This is, where $L$ is defined as the order of the homogenous neural network. Although these results indeed prove the benefits of training models, their analyses are limited to gradient descent. Moreover, the analyses developed by Soudry et al do not predict any phenomenon that resembles the Slingshot Mechanism. Wang et al. show that homogenous neural networks trained RMSProp or Adam without momentum do converge in direction to the max-margin solution. However, none of these papers can explain the Slingshot Mechanism and specifically the cyclical behavior of the norm of the last layer weights.

The Slingshot Mechanism

We use the training setup studied by Power et al. in the main paper as a working example to illustrate the Slingshot Mechanism. In this setup, we train decoder-only Transformers on a modular division dataset of the form $a\div b=c$, where $a$, $b$ and $c$ are discrete symbols and $\div$ refers to division modulo $p$ for some prime number $p$, split into training and validation sets. The task consists of calculating $c$ given $a$ and $b$. The algorithmic operations and details of the datasets considered in our experiments are described in Appendix B. The Transformer consists of 2 layers, of width 128 and 4 attention heads with approximately 450K trainable parameters and is optimized by Adam . For these experiments we set learning rate to 0.001, weight decay to 0, $\beta_{1}=0.9$, $\beta_{2}=0.98$, $\epsilon=10^{-08}$, linear learning rate warmup for the first 10 steps and minibatch size to 512 which are in line with the hyperparameters considered in .

Figure 2 shows the metrics of interest that we record on training and validation samples for modular division dataset. Specifically, we measure 1) train loss; 2) train accuracy; 3) validation loss; 4) validation accuracy; 5) last layer norm: denoting the norm of the classification layer’s weights and 6) feature change: the relative change of features of the l-th layer ($h^{l}$) after the t-th gradient update step $\frac{\|h^{l}_{t+1}-h^{l}_{t}\|}{\|h^{l}_{t}\|}$. We observe from Figure 2b that the model is able to reach high training accuracy around step 300 while validation accuracy starts improving after $10^{5}$ steps as seen in Figure 2d. Power et al. originally showed this phenomenon and refer to it as grokking. We observe that while the validation accuracy does not exhibit any change until much later in training, the validation loss shown in Figure 2c exhibits a double descent behavior with an initial decrease, then a growth before rapidly decreasing to zero. Seemingly, some of these observations can be explained by the arguments in and their extensions to adaptive optimizers . Namely, at the point of reaching perfect classification of the training set, the cross-entropy (CE) loss by design pressures the classification layer to grow in norm at relatively fast rate. Simultaneously, the implicit bias of the optimizer coupled with the CE loss, pushes the direction of the classification layer to coincide with that of the maximum margin classifier, albeit at a much slower rate.

These insights motivate us to measure the classifier’s last layer norm during training. We observe in Figure 2a that once classification reaches perfect accuracy on the training set, the classification layers norm exhibits a distinct cyclic behavior, alternating between rapid growth and plateau, with a sharp phase transition between phases. Simultaneously, the training loss retains a low value in periods of rapid norm growth, and then wildly fluctuating in periods of norm plateau. Figure 2e and Figure 2f shows the evolution of the relative change in features output by each layer in the Transformer. We observe that the feature maps are not updated much during the norm growth phase. However, at the phase transition, we observe that the feature maps receive a rapid update, which suggests that the internal representation of the model is updating.

In an attempt to ascertain the generality of Slingshot Effects as an optimization artifact, we run similar experiments with additional architectures, datasets, optimizers, and hyperparameters. We use all algorithmic datasets as proposed in , as well as frequently used vision benchmarks such as CIFAR-10 , and even synthetic Gaussian dataset. For architectures, we use Transformers, MLPs and deep linear models (see figure 1). We find abundant evidence of Slingshot Effects in all of our experiments with Adam, AdamW and RMSProp. We are unable to observe Slingshot Effects with Adagrad and also with stochastic gradient descent (SGD) or SGD with momentum, pointing to the generality of the mechanism across architectures and datasets. We refer the reader to Appendix A for the full, detailed description of the experiments.

We hypothesize that the norm growth continues until the curvature of the loss surface becomes large, effectively “flinging" the weights to a different region in parameter space as small gradient directions get amplified, reminiscent of the mechanics of a slingshot flinging a projectile. We attempt to quantify how far a model is flung by measuring the cosine distance between a checkpoint during optimization and initial parameters. Specifically, we divide the model parameters into representation (pre-classifier) parameters and classifier (last layer) parameters and calculate how far these parameters have moved from initialization. We show that checkpoints collected after a model experiences Slingshot has a larger representation cosine distance. We defer the reader to the appendix for further details.

By design, adaptive optimizers adapt the learning rate on a per parameter basis. In toy, convex scenarios, the $\epsilon$ parameter provably determine whether the algorithm will converge stably. To illustrate this, we take inspiration from , and consider a quadratic cost function $\mathcal{L}(A,B,C)=\frac{1}{2}x^{\top}Ax+B^{\top}x+C,A\in\mathcal{R}^{d\times d},x,B\in\mathcal{R}^{d},C\in\mathcal{R}$, where we assume $A$ is symmetric and positive definite. Note that the global minimum of this cost is given by $x^{\star}=-A^{-1}B$. The gradient of this cost with respect to $x$ is given by $g=Ax+B$. Consider optimizing the cost with adaptive optimization steps of the simple form $x_{t+1}=x_{t}-\mu\frac{g}{|g|+\epsilon}=x_{t}-\mu\frac{Ax_{t}+B}{|Ax_{t}+B|+\epsilon}$ where $\mu$ is a learning rate, and the division and absolute operations are taken element wise. Starting from some $x_{0}$, the error $e_{t}=x_{t}-x^{\star}$ evolves according to:

Note that the condition $\|A\|_{s}<\frac{2\epsilon}{\mu}$ where $\|\cdot\|_{s}$ denotes the spectral norm, implies that the mapping $\mathcal{M}_{t}$ is a contraction for all values of $t$, and hence convergence to the global optimum is guaranteed (This is in contrast to gradient descent, where the requirement is $\|A\|_{s}<\frac{2}{\mu}$). Note that the choice of $\epsilon$ crucially controls the requirement on the curvature of the cost, represented by the the spectrum of $A$ in this case. In other words, the smaller $\epsilon$, the more restrictive the requirements on the top eigenvalue of $A$. In , it was observed that full batch gradient descent increases the spectral norm of the hessian to its maximum allowed value. We therefore hypothesize that for deep networks, a small value for $\epsilon$ requires convergence to a low curvature local minimum, causing a Slingshot Effect when this does not occur. Moreover, we may reasonably predict that increasing the value of $\epsilon$ would lift the restriction on the curvature, and with it evidence of Slingshot Effects.

Figure 3 shows evidence consistent with the hypothesis that Slingshot Effects occur in the vicinity of high loss curvature, by measuring the local loss surface curvature along the optimization trajectory. Let $\mathcal{H}_{t}$ denote the local hessian matrix of the loss, and $u_{t}$ the parameter update at time $t$ given the optimization algorithm of choice. We use the local curvature along the trajectory of the optimizer, given by $\frac{1}{\|u_{t}\|^{2}}u_{t}^{\top}\mathcal{H}_{t}u_{t}$, as a curvature measure. Across the arithmetic datasets from , whenever the last layer weight norm plateaus, the curvature measure momentarily peaks and settles back down.

We next observe from Figure 2a that the training loss value also spikes up around the time step when the weight norm transitions from growth to plateau. A low training loss value suggests that the gradients (and their moments) used as inputs to the optimizer are small, which in turn can cause the $\epsilon$ hyperparameter value to play a role in calculating updates. Our hypothesis here is that the Slingshot Effect should eventually disappear with a sufficiently large $\epsilon$. To confirm this hypothesis, we run an experiment where we vary $\epsilon$ while retaining the rest of the setup described in the previous section.

Figure 4 shows the results for various values of $\epsilon$ considered in this experiment. We first observe that the number of Slingshot Effect cycles is higher for smaller values of $\epsilon$. Secondly, smaller values of $\epsilon$ cause grokking to appear at an earlier time step when compared to larger values. More intriguingly, models that show signs of grokking also experience Slingshot Effects while models that do not experience Slingshot Effects do not show any signs of grokking. Lastly, the model trained with the largest $\epsilon=10^{-5}$ shows no sign of generalization even after receiving 500K updates.

2 Effects on Generalization

In order to understand the relationship between Slingshot Effects and neural networks generalization, we experiment with various models and datasets. We observe that models that exhibit Slingshot tend to generalize better, which suggests the benefit of training models for a long time with Adam and AdamW . More surprisingly, we observe that Slingshots and grokking tend to come in tandem.

We follow the setting in Power et al. and generate several datasets that represent algorithmic operations and consider several training and validation splits. This dataset creation approach is consistent with the methodology used to demonstrate grokking . The Transformer is trained with AdamW with a learning rate of 0.001, weight decay set to 0, and with learning rate warmup for 500K steps. We consider $\epsilon$ of AdamW as a hyperparameter in this experiment. Figure 5 summarizes the results for this experiment where the x-axis indicates the algorithmic operation followed by the training data split size. As can be seen in Figure 5, Slingshot Effects are seen with lower values of $\epsilon$ and disappear with higher values of $\epsilon$ which confirms the observations made in Section 3 with modular division dataset. In addition, models that exhibit Slingshot Effects and grokking (shown in green) tend to generalize better than models that do not experience Slingshot Effects and grokking (shown in red).

For further validation of Slingshot Effects and generalization, we train a Vision Transformer (ViT) on CIFAR-10 . The ViT consists of 12 layers, width 384 and 12 attention heads trained on fixed subsets of CIFAR-10 dataset . The ViT model described above is trained with 10K, 20K, 30K, 40K and 50K (full dataset) training samples. We train the models with the following learning rates: $10^{-04}$, $3.10e^{-04}4$ and $10^{-03}$ and with a linear learning rate warmup for the 1 epoch of optimization. We consider multiple learning rates to study the impact of this hyperparameter on Slingshot taking inspiration from where the authors report observing grokking over a narrow range of learning rates . Figure 6 shows a plot of the highest test accuracy for a set of hyperparameters (learning rate, number of training samples) as a function of the number of training samples from which we make the following observations. The best test accuracy for a given set of hyperparameters is typically achieved after Slingshot phase begins during optimization. The checkpoints that achieve the highest test accuracy are labeled as "post-slingshot" and shown in green in Figure 6. While post-Slingshot checkpoints seem to enjoy higher test accuracy, there are certain combinations of hyperparameters that lead to models that show better test accuracy prior to the start of the first Slingshot phase. We label these points as "pre-slingshot" (shown in blue) in Figure 6. The above observations appear to be consistent with our finding that training long periods of time may lead to better generalization seen with grokking datasets .

We conduct experiments with MLPs on synthetic data where the synthetic data is a low dimensional embedding projected to higher dimensions via random projections. This design choices is critical with showing the existence of the Slingshot Effect with synthetically generated data. We find that using low dimensional data does not lead to any Slingshots. With this dataset, we show that generalization occurs late in training with Adam. Specifically, we tune $\epsilon$ in Adam and show that the optimizer is highly sensitive to this hyperparameter. These observations are consistent with the behavior reported above with Transformers and on algorithmic datasets as well as standard vision benchmark such as CIFAR-10. We refer the reader to Appendix A.1 for complete description and details of these experiments.

3 Drawbacks and Limitations

While the Slingshot Mechanism exposes an interesting implicit bias of Adam that often promotes generalization, due to it’s arresting of the norm growth and ensuing feature learning, it also leads to some training instability and prolonged training time. In the Appendix we show that it is possible to achieve similar levels of generalization with Adam on the modular division dataset using the same Transformer setup as above, while maintaining stable learning, in regimes that do not show a clear Slingshot Effect. First we employ weight decay, which causes the training loss values to converge to a higher value than the unregularized model. In this regime the model does not become unstable, but instead regularization leads to comparable generalization, and much more quickly. However, it is important to tune the regularization strength appropriately. Similarly, we find that it is possible to normalize the features and weights using the following scheme to explicitly control norm growth: $w=\frac{w}{\lVert w\rVert},f(x)=\frac{f(x)}{\lVert f(x)\rVert},$ where $w$ and $f(x)$ are the weights and inputs to the classification layer respectively, the norm used above is the $L_{2}$ norm, and $x$ is the input to the neural network. This scheme also results in stable training and similar levels of generalization. In all cases the effects rely on keeping the weight norms from growing uncontrollably, which may be the most important factor for improving generalization. These results suggest that while the Slingshot Mechanism may be an interesting self-correcting scheme for controlling norm growth, there are likely more efficient ways to leverage adaptive optimizers to similar levels of generalization without requiring the instability that is a hallmark of the Slingshot effect. Finally, we lack a satisfactory theoretical explanation for the Slingshot Mechanism, and hence removed all attempts at a more rigorous mathematical definition, which we feel would only serve as a distraction.

Conclusion

We have empirically shown that optimizing deep networks with cross entropy loss and adaptive optimizers produces the Slingshot Mechanism, a curious optimization anomaly unlike anything described in the literature. We have provided ample evidence that Slingshot Effects can be observed with different neural architectures and datasets. Furthermore, we find that Grokking almost always occurs in the presence of Slingshot Effects and associated regions of instability in the Terminal Phase of Training (TPT). These results in their pure form absent explicit regularization, reveal an intriguing inductive bias of adaptive gradient optimizers that becomes salient in the TPT, characterized by cyclic stepwise effects on the optimization trajectory. These effects often promote generalization in ways that differ from non-adaptive optimizers like SGD, and warrant further study to be able to harness efficiently. There are open question remaining to be answered, for instance 1) What’s the causal factor of the plateau of weight norm growth? 2) Are there better ways of promoting generalization without relying on this accidental training instability? Answering these questions w ill allow us to decouple optimization and regularization, and ultimately to control and improve them independently.

References

Appendix A Slingshot Effects across Architectures, Optimizers and Datasets

This section provides further evidence of the prevalence of Slingshot across architectures and optimizers on subsets of CIFAR-10, testing setups beyond the specific setup consider by Power et al. . In these experiments, we focus solely on characterizing the optimization properties of various setups described below. The small sample sizes are used in order to more easily find regimes where different architectures can converge to fit the training data fairly quickly.

We use cross-entropy loss to optimize the models with AdamW in the following experiments. The following experiments are implemented in PyTorch .

For further validation, we train a Vision Transformer (ViT) with 12 layers that has 10 million parameters on a small sample of the CIFAR-10 dataset . In this setup, we use a learning rate to 0.001, no weight decay, $\beta_{1}=0.9$, $\beta_{2}=0.95$, $\epsilon=1e-08$ and minibatch size of 128. We choose a sample size of 1000 training samples for computational reasons, as we wish to observe multiple cycles of the Slingshot Mechanism extremely late in training. The input images are standardized to be in the range $$. No data augmentation is used in our training pipeline. Due to the extremely small sample size, we focus our attention on the training metrics since no generalization is expected. Figure 7a (respectively Figure 7b) shows a plot of training loss (respectively training accuracy) and last layer norm evolution during the latter stages of training. Multiple Slingshot stages are observed in these plots (5 clear cycles), which can be seen by the sharp transition of the weight norm from high growth to plateau.

A.0.2 CNN on 200 samples from CIFAR-10

We consider a VGG-like architecture that has been adapted for CIFAR-10 dataset.We use the VGG11 architecture without batch normalization from https://github.com/kuangliu/pytorch-cifar in this experiment. The model is trained with 200 randomly chosen samples from CIFAR-10 training split and with full-batch AdamW . The hyperparameters used for the optimizer include a learning rate of $0.001$, weight decay$=0$, $\beta_{1}=0.9$, $\beta_{2}=0.95$, and $\epsilon=1e-08$. As with ViT, no data augmentation is used in these experiments other than standardizing the input to be in the range $$. We observe the prescenece of multiple Slingshot stages with CNN from Figure 8a and Figure 8b. These experiments suggest that Slingshot Effects are not restricted to Transformers architecture alone.

We repeat the CNN-based described above but with a VGG-like model that includes batch normalization .We use the VGG11 architecture with batch normalization from https://github.com/kuangliu/pytorch-cifar in this experiment. The training setup is identical to the one described for CNN wihtout batch normalization. We observe the prescenece of multiple Slingshot stages with CNN from Figure 9a and Figure 9b. The weight norm does not decrease during training as opposed to the weight norm dynamics for CNN wihtout batch normalization seen in Figure 8. These experiments suggest that Slingshot Effects can be seen with standard neural network training components including batch normalization.

A.0.3 MLPs on 200 samples from CIFAR-10

The next architecture we consider is a deep (6 layers) fully connected network trained on a small sample of 200 samples belonging to the CIFAR-10 dataset with full-batch AdamW optimizer. The optimizer’s hyperparameters are set as following: learning rate $=0.001$, weight decay $=0$, $\beta_{1}=0.9$, $\beta_{2}=0.95$, and $\epsilon=1e-08$. As with the ViT setup above we do no use data augmentation for training this model. Figure 10a (respectively Figure 10ab) shows a plot of training loss (respectively training accuracy) and last layer norm evolution during the latter stages of training. Multiple Slingshot stages are observed in this setup as well. These experiments further suggest that the Slingshot mechanism is prevalent in simple models as well.

A.0.4 Deep linear models

We train a 6 layer linear model with 200 samples belonging to CIFAR-10 with full-batch AdamW . The optimizer’s hyperparameters are set as following: learning rate $=0.001$, weight decay $=0$, $\beta_{1}=0.9$, $\beta_{2}=0.95$, and $\epsilon=1e-08$. Figure 11a and Figure 11b show the training loss and accuracy behavior observed during optimization. Multiple Slingshot stages are observed with this architecture as well.

A.0.5 Different Optimizers

In this set of experiments, we study the training loss behavior of deep linear models optimized full-batch with AdamW , RMSProp and full-batch gradient descent (GD). The six layer model is trained with 200 samples. The hyperparameters used for optimizing the model with various optimizers are described in Table 1. Figure 11 shows the training loss and accuracy behavior of the three optimizers considered in this experiment. We observe Slingshot behavior with AdamW and RMSProp from Figure 11 while Slingshot behavior is absent with standard gradient descent. This observation suggests that the normalization used in adaptive optimizers to calculate the update from gradients may lead to Slingshot behavior.

A.1 Slingshot with MLP and Synthetic Dataset

In this section, we provide empirical evidence that Slingshot Effects are observed with a synthetic dataset in a fully-connected architecture. The small dimensional dataset, like the Grokking dataset of Power et al. , allows us to easily measure of sharpness, given by $\frac{1}{\|u_{t}\|^{2}}u_{t}^{\top}\mathcal{H}_{t}u_{t}$ where $u_{t}$ is the optimizer’s update vector and $\mathcal{H}_{t}$ is the Hessian at step $t$, to examine the interplay between Slingshot Effects and generalization.

In Appendix A, we have empirically shown that the existence of the Slingshot phenomenon on a small subset of CIFAR-10 dataset with Vision Transformers (ViTs). We now study the impact that Slingshot has on the generalization ability of ViTs by training a model on all 50000 samples in CIFAR-10 training dataset. The ViT used here is a larger model than the one considered in A to account for larger dataset size. The ViT model consists of 12 layers, width 384 and 12 attention heads and is optimized by AdamW . For this experiment, we set the learning rate to 0.0001, weight decay to 0, $\beta_{1}=0.9$, $\beta_{1}=0.95$ and $\epsilon=10^{-08}$, minibatch size of 512 and linear learning rate warmup for $1$ epoch of optimization. Figure 13 shows the results of experiment with full CIFAR-10 dataset. Multiple Slingshots can be observed in these plots similar to the plots described in Appendix A. We observe from Figure 13d that the test accuracy peaks in epochs following a Slingshot with the maximum recorded test accuracy occurring very late in optimization. This observation suggests that the Slingshot can have a favorable effect on generalization consistent with the behavior observed in the main paper with division dataset.

A.1.1 Abalation Study

In this section, we train a toy model on a synthetically generated dataset with the aim of analysing the effect of different hyper parameters on the Slingshot Mechanism. We construct a 128-dimensional dataset with Scikit-learn that has 3 informative dimensions that represents a 8-class classification problem. The class centers are the edges of a 3-dimensional hypercube around which clusters are data are sampled from a standard normal distribution. The other 125-dimensions are also filled at random to create a high-dimensional dataset used in our experiments. We generate 256 training and validation samples for this dataset and use a minibatch size of 128 in all the experiments described in the following.

Figure 12 shows the training and validation metrics when we optimize a 4-layer fully-connected network (FCN) with Adam using a learning rate of 0.001, $\beta_{1}=0.9$, $\beta_{1}=0.95$, no weight decay and $\epsilon=1e^{-08}$. Note that we use this value of $\epsilon$ in our first experiment as this is the default value proposed in Kingma and Ba . These experiments are implemented in JAX .

In the next set of experiments with synthetic data, we tune $\epsilon$ value for Adam to understand its impact on test accuracy. Figure 14 shows a plot of the maximum validation accuracy achieved by models trained with Adam as a function of time (epoch). We observe that Adam reaches its best test accuracy late in optimization with $\epsilon=10^{-5}$ yielding the highest validation accuracy. Furthermore, the best accuracy is achieved with a model that experiences Slingshot during optimization. This observation is consistent with our findings for ViT training with CIFAR-10 dataset described in the main paper and Appendix A.1.

In these experiments, we aim to study the impact of Adam/AdamW optimizer’s $\beta_{1}$ and $\beta_{2}$ hyperparameters on Slingshot. We use the synthetic data described above and set the learning rate of $0.001$ and $\epsilon=10^{-08}$ for this analysis. Figure 15 and Figure 16 shows the results of this study. We observe from Figure 15 that the Slingshot Mechanism is fairly robust to the values of $\beta_{1}$ and $\beta_{2}$. Figure 15a-Figure 15c show that Slingshot is even observed with $\beta_{1}$ and $\beta_{2}$ set to which effectively disables exponential moving averaging of gradient moments in Adam . Figure 15g-Figure 15i provide an example of hyperparameters that fail to induce Slingshot. We observe from Figure 15 that models that experience Slingshot tend to reach their best test accuracy during the later stages of training. Specifically, we observe from Figure 15b, Figure 15e and Figure 15k that the best validation accuracy occurs after $60000$ epochs. These examples provide further evidence about an interesting implicit bias of Adam. Figure 16 shows more examples of hyperparameters that do not induce Slingshot Effects. Finally, we observe from Figure 16 that hyperparameters that provide higher validation accuracy are from models that experience Slingshot Effects.

Appendix B Slingshot and Grokking

We use the empirical setup described by Power et al. to describe the Slingshot Mechanism. The following section describes relevant details including datasets, architecture and optimizer used in our experiments.

The model used a decoder-only Transformer with causal attention masking. The architecture used in all our experiments consists of 2 decoder layers with each layer of width 128 and 4 attention heads.

We train the architecture described above with Adam optimizer in most of our experiments unless noted otherwise. The learning rate is set to $0.001$ and with linear learning rate warmup for the first $10$ steps. We use $\beta_{1}=0.9$, $\beta_{2}=0.98$ for Adam’s hyperparameters. The Transformers are optimized with cross-entropy (CE) loss that is calculated on the output tokens for a given binary operation.

The Transformer is trained on small algorithmic datasets that consists of sequences that represent a mathematical operation. The following operations are used in our experiments:

$c=a^{2}+b^{2}\pmod{p}$ for $0\leq a,b<p$

$c=a^{2}+b^{2}+ab\pmod{p}$ for $0\leq a,b<p$

$c=a^{2}+b^{2}+ab+b\pmod{p}$ for $0\leq a,b<p$

$c=a^{3}+ab^{2}+b\pmod{p}$ for $0\leq a,b<p$

$c=[a\div b\pmod{p}$ if $b$ is odd, otherwise $a-b\pmod{p}$] for $0\leq a,b<p$

$c=a\cdot b\cdot a^{-1}$ for $a,b\in S_{5}$

$c=[a+b\pmod{p}$ if $a$ is even, otherwise $a*b\pmod{p}$] for $0\leq a,b<p$

$c=[a+b\pmod{p}$ if $a$ is even, otherwise $a-b\pmod{p}$] for $0\leq a,b<p$

where $p=97$ and with the dataset split in training and validation data. Each equation in the dataset is of the form $(a)(op)(b)(=)c$ where (x) represents the token used to represent x. We refer to Power et al. for a detailed description of the datasets

B.1 Analysis of Parameter Dynamics

A common observation is that intermediate representations tend to evolve beyond simple scale increase during phase transitions from norm growth to plateau. In order to empirically quantify this effect, we train the Transformer described in Appendix B with modular addition, multiplication and division datasets using Adam with learning rate set to $0.001$ and $\beta_{1}=0.9$ and $\beta_{2}=0.98$. We calculate the cosine distance between the representation and classification parameters from their initial values where the cosine distance is given by

where $d^{repr}$ ($d^{clf}$) denotes cosine distance for representation (respectively classification) parameters, $w^{repr}_{t}$ (resp. $w^{clf}_{t}$) denotes representation (resp. classification) parameters at time $t$ with $w^{repr}_{0}$ ($w^{clf}_{t}$) indicating the initial representation (resp. classification) parameters where the norm used above is the Euclidean norm.

Figure 17 shows the dynamics of the loss, accuracy and cosine distance recorded during training. We observe that the classification parameters move farther away from initialization faster than the representation parameters. More interestingly, we observe from Figure 17c and Figure 17f that the representation parameters travel farther from initialization for training runs that experience Slingshot. These trials use $\epsilon=10^{-08}$ and $\epsilon=10^{-07}$ and experience Slingshot Effects. In contrast, we see from Figure 17i and Figure 17l that the representation distance remains low for models trained with $\epsilon=10^{-05}$ and $\epsilon=10^{-04}$. The models trained with higher $\epsilon$ values do not experience Slingshot Effects. These results suggest that Slingshot may have a beneficial effect in moving the representation parameters away from initialization which eventually helps with model generalization. Figure 18 and Figure 19 show a similar trend for multiplication and division datasets respectively.

B.2 SGD Optimization

In this appendix, we show that Slingshot Effects are not seen during Transformer training with stochastic gradient descent (SGD) with momentum to support our claim in the main paper. To this end, we use train the Transformer described in in Appendix B on modular division dataset with a 50/50 train/validation split using SGD with momentum. We use a mini-batch size of $512$ which requires the optimizer to take $10$ steps per epoch for dataset split described above. We set momentum to $0.9$ and use the following learning rates: $0.001$, $0.01$ and $0.1$ and run the optimizer for 1500000 steps. The number of steps used here is 3 times larger than the steps used to run Adam/AdamW in this work which is chosen to give SGD additional time to reach convergence. Figure 20 shows the usual loss and accuracy metrics calculated on training and validation data as well as the weight norm of the classifier layer. We observe that there is no evidence of Slingshot with SGD. Lastly, we do not see any evidence of Grokking or generalization with this setup as well.

B.3 Slingshots with Additional Datasets

In this appendix, we provide evidence of Slingshot Effects on additional datasets from Power et al Grokking work. The datasets are created by a subset of mathematical operations defined in Appendix B. Each operation can have multiple datasets that depends on the train/validation split ratio. We use the training setup described in B on $18$ separate datasets. Figure 21 - Figure 38 shows the results the datasets described in this appendix. We observe Slingshot Effects and generalization with all $18$ datasets. These results suggest the prevalence of Slingshot Effects when large models are trained with adaptive optimizers, specifically Adam .

Appendix C Controlling Instability Through Normalization and Norm Constraints

Training instability is the hallmark of the Slingshot Mechanism, yet as seen in previous sections, the Slingshot Effect typically results in improved performance, and Grokking. In this section, we explore whether it is possible to maintain stable training, without sacrificing performance. To this end, we explore how constraining and regularizing the weights of the network affect the Slingshot behaviour, and overall performance.

Weight decay is a commonly used regularization approach to improve the generalization performance of neural networks. Power et al. show that weight decay has the largest positive effect on alleviating Grokking. Weight decay naturally controls the size of the parameters and consequently their norm growth. We study the effect of weight decay on stability of training Transformers with Grokking datasets in this section. We use weight decay values from $0,0.1,0.2,0.4,0.6,0.8and1.0$ with AdamW optimizer. Figure 39 shows the results for division dataset. We observe from Figure 39 that as weight decay strength increases, both Slingshot Effects and Grokking phenomenon disappear with the model reaching high validation accuracy quickly as seen in Figure 39e and Figure 39f. However, we observe that the model experiences instability as can been seen with the loss plots in Figure 39b and Figure 39c or the accuracy plots in Figure 39e and Figure 39f. A similar trend is observed for addition and multiplication datasets in Figure 40 and Figure 41 respectively.

The results shown above indicate that Slingshot may not be the only way to achieve good generalization. Both Slingshot and weight decay prevent the norms from growing unbounded and achieve high validation accuracy as seen in plots described above. While weight decay shows different weight norm dynamics, this regularization does not decrease training instability. These results suggest the need for alternative approaches to improve training stability.

C.2 Features and parameter normalization

A second approach that we use to explicitly control weights and feature norm is by normalizing the features and weights via the following scheme: $w=\frac{w}{\lVert w\rVert},f(x)=\frac{f(x)}{\lVert f(x)\rVert},$ where $w$ and $f(x)$ are the weights and inputs to the classification layer respectively, the norm used above is the $L_{2}$ norm, and $x$ is the input to the neural network. We take the cosine similarity of the normalized weights and features and divide this value by a temperature value that we treat as a hyperparameter in these experiments. The operation is given by: $y=\frac{w\cdot f(x)}{\tau}$ where $\tau$ represents the temperature hyperparameter. We use temperature values from $0.1,0.25,0.5,0.75,1.0$ for these experiments.

Figure 42 shows the results of Transformer training on division dataset described in Appendix B that is split evenly into train and validation sets. We observe that the model displays training instability evidenced by norm behavior and also loss behavior in Figure 42a at lower temperature values. We observe that $\tau=0.25$ provides a good compromise between fitting training data while showing no training instability as seen in Figure 42b. This hyperparameter value also results in Grokking as validation accuracy improves late in training as can be seen from Figure 42e. These together suggest that bounding weights and features norm helps stabilize training without sacrificing training performance.

We validate the normalization scheme with two additional datasets namely multiplication and division from Appendix B. Figure 43 shows the results for training Transformers with multiplication dataset that is split evenly into train and validation sets. We observe from Figure 43 that a proper temperature value can stabilize training and with some tuning can provide a compromise between training stability and generalization. Specifically, $\tau=0.25$ allows the model to fit the training data and reach almost perfect validation accuracy as seen from Figure 43b and Figure 43e.

Finally, we repeat the above experiments with subtraction dataset and show the results in Figure 44. This dataset shows that while a properly tuned temperature can help the model achieve almost perfect generalization, training instability shows up very late in optimization. This observation can be seen from Figure 44b and Figure 44d. This result suggests that more work remains to be done with understanding and stabilizing the training behavior of large neural networks.