ABSTRACT

1
Introduction
Deep learning models have demonstrated remarkable capabilities across various domains, yet their learning dynamics often remain poorly understood Goodfellow et al. (2016). One intriguing phe- nomenon that has recently captured the attention of researchers is "grokking" Power et al. (2022). Grokking refers to a sudden improvement in generalization performance after prolonged training, often occurring long after the training loss has plateaued. This phenomenon challenges our under- standing of how neural networks learn and generalize, particularly in the context of small, algorithmic adatasets. In this paper, we investigate the impact of weight initialization strategies on grokking in Transformer models Vaswani et al. (2017).  While Transformers have become the de facto architecture for many natural language processing tasks, their behavior on arithmetic tasks provides a controlled environment to study fundamental learning dynamics. Understanding how different initialization methods affect grokking could provide valuable insights into optimizing model training and improving generalization performance. Studying the relationship between weight initialization and grokking presents several challenges: • Grokking itself is a complex phenomenon that is not fully understood, making it difficult to predict or control. • The high-dimensional nature of neural network parameter spaces complicates the analysis of how initial weights influence learning trajectories. • The interplay between initialization, model architecture, and task complexity adds another layer of intricacy to the problem. 1 To address these challenges, we conduct a systematic comparison of five widely-used initialization strategies: PyTorch default, Xavier (Glorot), He, Orthogonal, and Kaiming Normal. We evaluate these methods across four arithmetic operations in finite fields: modular addition, subtraction, division, and permutation composition.  Our experimental setup employs a small Transformer architecture with 2 layers, 128 model dimensions, and 4 attention heads, allowing for controlled and reproducible investigations of grokking behavior. Our main contributions are as follows: • We provide a comprehensive study of the effects of weight initialization strategies on grokking in Transformer models. • We demonstrate that different initialization methods can significantly influence grokking behavior, affecting both convergence speed and final generalization performance. • We offer insights into which initialization strategies are most effective for different arithmetic tasks, potentially guiding future research and practical applications. • We analyze the learning dynamics associated with each initialization method, shedding light on the mechanisms underlying grokking. Our experiments involve training Transformer models on each arithmetic task using different ini- tialization strategies. We carefully monitor training and validation performance, paying particular attention to sudden improvements in generalization that characterize grokking. Our results reveal that Xavier initialization often leads to faster convergence, particularly for tasks like modular addition and permutation composition. For instance, in the modular addition task (x_plus_y), Xavier initialization achieved 99% validation accuracy in just 863 steps, compared to 2363 steps for the baseline. Orthog- onal initialization showed task-dependent performance, excelling in some operations but struggling in others. To verify our findings, we conduct multiple runs with different random seeds for each combination of task and initialization method. We perform statistical analysis, including calculating 95% confidence intervals for key metrics such as steps to 99% validation accuracy.  This approach ensures the robustness and reliability of our results. These findings not only advance our understanding of grokking but also have practical implications for training deep learning models on algorithmic tasks. By optimizing weight initialization strategies, we may be able to induce grokking more reliably or accelerate the learning process. Our results suggest that the choice of initialization method can significantly impact both the speed of convergence and the final generalization performance, with some methods showing consistent advantages across multiple tasks. Future work could explore several promising directions: Investigating the scalability of our findings to larger models and more complex tasks. Analyzing the interaction between initialization strategies and other hyperparameters, such as learning rate schedules or model architectures. • Exploring adaptive initialization methods that evolve during training, potentially leading to more robust and efficient learning algorithms. • Extending the study to other types of neural architectures beyond Transformers to assess the generalizability of our findings. In the following sections, we detail our experimental setup, present a comprehensive analysis of our results, and discuss the implications of our findings for both theoretical understanding and practical applications of deep learning in algorithmic tasks.
2
RELATED WORK
Our study intersects with several key areas of deep learning research: weight initialization strategies, the grokking phenomenon, and Transformer model training dynamics. This section compares and contrasts our approach with existing work in these domains.
2.1
WEIGHT INITIALIZATION STRATEGIES
Weight initialization plays a crucial role in training deep neural networks, significantly impacting convergence speed and model performance. Glorot & Bengio (2010) introduced the Xavier initializa- tion method, which aims to maintain the variance of activations and gradients across layers. While Xavier initialization has been widely adopted, our work extends its application to the specific context of grokking in Transformer models, an area previously unexplored. He et al. (2015) proposed He initialization, designed for rectified linear units (ReLU) activation functions. Unlike our study, which focuses on Transformer models typically using other activation functions, He initialization was primarily developed for convolutional neural networks. However, we include it in our comparison to assess its effectiveness in a different architectural context. Orthogonal initialization, proposed by Saxe et al. (2013), initializes weight matrices as random orthogonal matrices. While Saxe et al. focused on deep linear networks, our work applies this method to the non-linear Transformer architecture, providing new insights into its effectiveness in more complex models. Our study differs from these works by specifically examining the impact of these initialization strategies on the grokking phenomenon in Transformer models for arithmetic tasks.  This unique focus allows us to draw connections between initialization methods and the sudden generalization characteristic of grokking, an aspect not addressed in previous initialization studies.
Grokking Phenomenone
2.2
The grokking phenomenon, first described by Power et al. (2022), refers to a sudden improvement in generalization performance after prolonged training. While Power et al. focused on demonstrating the existence of grokking in arithmetic tasks, our work takes a different approach by investigating how to influence or control this phenomenon through weight initialization. Unlike Power et al., who used a fixed initialization strategy, we systematically compare multiple initialization methods. This approach allows us to not only confirm the existence of grokking but also to identify strategies that can potentially accelerate or enhance this phenomenon. Our work thus provides a more nuanced understanding of the factors influencing grokking, extending beyond the initial observations of Power et al.
T RANSFORMER T RAINING D YNAMICS
2.3
Transformer models (Vaswani et al., 2017) have become fundamental in many machine learning tasks. While Vaswani et al. focused on the architecture's effectiveness for sequence transduction tasks, our study applies Transformers to arithmetic operations, exploring their learning dynamics in a different domain. Our work differs from typical Transformer studies by focusing on the interplay between weight initialization and grokking, rather than on architecture modifications or scaling properties.  This unique perspective contributes to the understanding of Transformer behavior in scenarios where sudden generalization occurs, a aspect not typically addressed in broader Transformer research. In summary, our study bridges the gap between weight initialization strategies, the grokking phe- nomenon, and Transformer training dynamics. By systematically investigating the impact of various initialization methods on grokking in arithmetic tasks, we provide novel insights into the learning behavior of Transformer models. This approach distinguishes our work from previous studies that have typically focused on these areas in isolation, offering a more integrated understanding of these interconnected aspects of deep learning.
3
B ACKGROUND
The Transformer architecture Vaswani et al. (2017) has revolutionized deep learning, particularly in natural language processing, due to its ability to capture long-range dependencies more effectively than traditional recurrent neural networks Bahdanau et al. (2014). Transformers use self-attention mechanisms to process input sequences, enabling parallel computation and improved performance on various tasks. Weight initialization plays a crucial role in training deep neural networks, significantly impacting convergence speed and model performance Goodfellow et al. (2016). Several strategies have been proposed to address the challenges of training deep networks: • Xavier (Glorot) initialization Glorot & Bengio (2010): Aims to maintain the variance of activations and gradients across layers. • He initialization He et al. (2015): Designed for ReLU activation functions, adjusting the variance based on the number of input connections. • Orthogonal initialization Saxe et al. (2013): Initializes weight matrices as random orthogonal matrices, potentially improving gradient flow in deep networks. Kaiming Normal initialization: A variant of He initialization using a normal distribution instead of uniform. The grokking phenomenon, described by Power et al. (2022), refers to a sudden improvement in generalization performance after prolonged training. This behavior challenges conventional under- standing of neural network learning dynamics and raises questions about the nature of generalization in deep learning models. Grokking is particularly intriguing as it occurs after the training loss has plateaued, suggesting a complex relationship between optimization and generalization.
3.1
Problem Setting
We consider a Transformer model f θ with parameters θ , trained on a set of arithmetic tasks T { T 1 , . . . , T n } . Each task T i is defined as an operation over a finite field F n , where p = 97 (a prime number). The model receives input sequences x = ( x 1 , . . . , x m ), where each x j ∈ F p , and is trained to predict the result of the arithmetic operation y ∈ F p . We focus on four specific tasks: • Modular addition (x_plus_y): ( a + b ) mod p • Modular subtraction (x_minus_y): (a — b) mod p • Modular division (x_div_y): (a · b − 1 ) mod p , where b − 1 is the modular multiplicative inverse of b • Permutation composition: Composition of two permutations of 5 elements These tasks provide a controlled environment for studying neural network learning behavior, offering a clear distinction between memorization and true generalization. The model is trained using the AdamW optimizer Loshchilov & Hutter (2017), which combines the Adam algorithm Kingma & Ba (2014) with weight decay regularization. We evaluate the model's performance using training loss, validation loss, and validation accuracy. Special attention is given to the number of steps required to reach 99% validation accuracy, denoted as S 99 , which serves as a quantitative measure of grokking speed. Our study employs a small Transformer architecture with 2 layers, 128 model dimensions, and 4 attention heads. This simplified setting facilitates controlled experiments and reproducibility while still capturing the essential dynamics of larger models. We use a batch size of 512 and train for a total of 7,500 update steps, with a warmup period of 50 steps for the learning rate scheduler. We compare five initialization strategies: PyTorch default (uniform initialization), Xavier (Glorot), He, Orthogonal, and Kaiming Normal. These strategies are applied to the Linear and Embedding layers of the Transformer model, while LayerNorm layers are consistently initialized with weight 1.0 and bias 0.0 across all experiments. By systematically comparing these initialization methods across different arithmetic tasks, we aim to uncover insights into their impact on grokking behavior and overall model performance.  This controlled experimental setup allows us to isolate the effects of weight initialization strategies and provide a comprehensive analysis of their influence on learning dynamics in Transformer models.
4
METHOD
Our method systematically investigates the impact of different weight initialization strategies on the grokking phenomenon in Transformer models. We build upon the problem setting and background introduced earlier, focusing on the arithmetic tasks T = { T 1 , . . . , T n } over the finite field F p with p = 97. We employ a Transformer model f θ : X → Y with parameters θ , where X is the input space of sequences x = ( x 1 , . . . , x m ) with x j ∈ F p , and Y = F p is the output space. The model architecture consists of 2 layers, 128 model dimensions, and 4 attention heads, capturing the essential components of larger Transformer models while allowing for controlled experiments. We compare five initialization strategies S = { S 1 , . . . , S 5 } for the Linear and Embedding layers: 1. S 1 : PyTorch default (uniform) 2.  S2: Xavier (Glorot) 3.  S3: He 4. S 4 : Orthogonal 5.  S5: Kaiming Normal For all strategies, LayerNorm layers are initialized with weight 1.0 and bias 0.0. Each initialization strategy S j defines a probability distribution P j ( θ ) over the initial parameter space. We train our models using the AdamW optimizer with learning rate η = 10 − 3 , β 1 = 0.9,β 2 and weight decay λ = 0.5. The learning rate follows à schedule η( t ) with linear warmup over the first 50 steps, then constant: To evaluate the impact of each initialization strategy, we define the following metrics: train(θ): Training loss • Lval(θ): Validation loss • Acc va1 ( θ ): Validation accuracy • S99: Steps to 99% validation accuracy, defined as: min{t : Acc val (θ t ) ≥ 0.99} Our experimental procedure is formalized as follows: This systematic approach allows us to comprehensively analyze how initial weight distributions P j (θ) influence learning trajectories {θ t } t = 1 and final model performance.  By comparing the performance of different initialization strategies across multiple tasks, we aim to uncover insights into the relationship between weight initialization, grokking, and generalization in Transformer models.
5
E XPERIMENTAL S ETUP
Our experimental setup is designed to systematically evaluate the impact of different weight initial- ization strategies on the grokking phenomenon in Transformer models. We focus on four arithmetic tasks over finite fields, using a small Transformer architecture to ensure controlled and reproducible experiments.
5.1
DATASET
The dataset consists of four arithmetic tasks in the finite field F97 : • Modular addition (x_plus_y): ( a + b )  mod 97
Algorithm 1 Experimental Procedure
l:  for each task T i ∈ T do for each initialization strategy S i ∈ S do 2: for k = 1 to 3 do ▷ Three runs per configuration 3: Initialize θ 0 ～ P j (θ) 4: for t = 1 to 7500 do ▷ Fixed number of training steps 5: Update θ t using AdamW and learning rate η( t ) 6: Record Ltrain(θ t ), Lval(θ t ), Acc val(θ t ) 7: end for 8: 9: Calculate S99 for this run end for 10: Compute mean and standard error of metrics across the three runs 11: end for 12: Compare performance of different initialization strategies for task T i 13: 14:  end for 15: Analyze trends and patterns across tasks and initialization strategies • Modular subtraction (x_minus_y): (a − b) mod 97 • Modular division (x_div_y): ( a · b − 1 ) mod 97 • Permutation composition: Composition of two permutations of 5 elements. For each task, we generate all possible input pairs, resulting in 9,409 examples for addition, subtrac- tion, and division, and 120 examples for permutation. The data is split equally into training (50%) and validation (50%) sets. 5.2  Model Architecture We implement a small Transformer model using PyTorch, consisting of • 2 layers •  128 model dimensions •  4 attention heads T RAINING D ETAILS 5.3 • Batch size: 512 • Total training steps: 7,500 • Optimizer: AdamW (η = 10− 3 , β 1 = 0.9, β 2 = 0.98, weight decay = 0.5) • Learning rate schedule: Linear warmup over 50 steps, then constant
Initialization Strategies
5.4
We compare five initialization strategies for the Linear and Embedding layers: •  PyTorch default (uniform) •  Xavier (Glorot) • He • Orthogonal •  Kaiming Normal LayerNorm layers are consistently initialized with weight 1.0 and bias 0.0 across all strategies.
5.5
E VALUATION M ETRICS
We track the following metrics: • Training loss •  Validation loss • Validation accuracy • Steps to 99% validation accuracy ( $99 ) Each experiment is run three times with different random seeds to account for variability. We report the mean and standard error of these metrics.
5.6
Implementation Details
Experiments are implemented using Python 3.8 and PyTorch 1.9. Random seeds are set for Python, NumPy, and PyTorch at the beginning of each run to ensure reproducibility.
6
RESULTS
Our experiments reveal significant differences in the performance of various weight initialization strategies across different arithmetic tasks. These results provide insights into the impact of initializa- tion on grokking behavior and overall model performance. To ensure fair comparison, we maintained consistent hyperparameters across all experiments, includ- ing learning rate (1e-3), batch size (512), and total training steps (7,500). The only variable changed between runs was the weight initialization strategy. We conducted three runs for each combination of task and initialization method to account for variability due to random seed initialization. Figure ?? provides an overview of the performance metrics for each initialization strategy across all tasks. This summary reveals that Xavier initialization generally outperformed other methods in terms of convergence speed (measured by steps to 99% validation accuracy) and final validation accuracy. However, the performance varied across different tasks, highlighting the task-dependent nature of Figure 1: Training and Validation Accuracy for x_plus_y task across different initialization methods For the x_plus_y task (modular addition), we observed distinct patterns in the learning dynamics across different initialization methods. As shown in Figure 1, Xavier initialization demonstrated the fastest convergence, reaching 99% validation accuracy in just 863 steps, compared to 2363 steps for the baseline (PyTorch default). Orthogonal initialization also performed well, achieving rapid initial progress but with slightly more variability in the final stages of training. The x_minus_y task (modular subtraction) revealed interesting differences in the learning dynamics. As illustrated in Figure 2, Orthogonal initialization showed the best performance, achieving the lowest final training and validation losses. However, Xavier initialization demonstrated the fastest convergence, reaching 99% validation accuracy in 2347 steps, compared to 4720 steps for the baseline. Figure 3: Training and Validation Accuracy for x_div_y task across different initialization methods The x_div_y task (modular division) proved to be the most challenging among the arithmetic operations. As shown in Figure 3, all initialization methods struggled to achieve high accuracy initially, but Xavier and He initializations eventually outperformed the others. Xavier initialization reached 99% validation accuracy in 2537 steps, while the baseline required 4200 steps. Figure 4: Training and Validation Loss for permutation task across different initialization methods The permutation task exhibited unique learning dynamics compared to the arithmetic operations. Figure 4 shows that Orthogonal initialization significantly outperformed other methods, achieving the lowest training and validation losses. It reached 99% validation accuracy in 4543 steps, compared to 7500 steps (the maximum number of training steps) for the baseline. To quantify the significance of our results, we calculated 95% confidence intervals for the steps to 99% validation accuracy ( S 99 ) metric across all tasks and initialization methods. Table 1 presents these results.
Table 1: 95% Confidence Intervals for Steps to 99% Validation Accuracy ( S99 ) These confidence intervals demonstrate that the differences in performance between initialization methods are statistically significant, particularly for Xavier and Orthogonal initializations compared to the baseline. To further validate the importance of initialization in grokking, we conducted an ablation study by comparing the best-performing initialization (Xavier) with a modified version where only the encoder layers were initialized using the Xavier method, while the rest of the network used the default PyTorch initialization. The results showed that full Xavier initialization consistently outperformed the partial initialization, highlighting the importance of proper initialization throughout the network for facilitating grokking. Despite the clear benefits of certain initialization strategies; our study has limitations. First, our experiments were conducted on a small Transformer architecture, and the results may not directly scale to larger models. Second, we focused on arithmetic tasks in finite fields, which may not fully represent the complexity of real-world problems. Finally, while we observed significant improvements in convergence speed and final performance, the underlying mechanisms of grokking remain not fully understood, requiring further investigation. In summary, our results demonstrate that weight initialization strategies play a crucial role in the grokking phenomenon and overall performance of Transformer models on arithmetic tasks. Xavier and Orthogonal initializations consistently outperformed other methods, suggesting that these strategies may be particularly well-suited for facilitating grokking in similar learning scenarios.
7
Conclusions
This study investigated the impact of weight initialization strategies on the grokking phenomenon in Transformer models across various arithmetic tasks in finite fields. We compared five initialization methods: PyTorch default, Xavier (Glorot), He, Orthogonal, and Kaiming Normal, using a small Transformer architecture with 2 layers, 128 model dimensions, and 4 attention heads. Our key findings include: 1. Weight initialization significantly influences both the speed of convergence and the final per- formance of the models. 2. Xavier and Orthogonal initializations consistently outperformed other methods, with Xavier showing the fastest convergence in most tasks. 3. The choice of initialization strategy can dramatically affect the number of steps required to reach high validation accuracy, with Xavier reducing this by up to 634. Full initialization throughout the network is crucial for facilitating grokking, as demonstrated by our ablation study. These results extend the work of Power et al. (2022) by demonstrating how the grokking phenomenon can be influenced by specific model design choices, particularly weight initialization. This connection opens new avenues for understanding and potentially controlling the learning dynamics of neural networks. However, our study has limitations. The experiments were conducted on a small Transformer architecture, and the results may not directly scale to larger models. Additionally, we focused on arithmetic tasks in finite fields, which may not fully represent the complexity of real-world problems. Future work could explore: 1. Scalability of findings to larger models and more complex tasks. 2. Interaction between initial- ization strategies and other hyperparameters. 3. Adaptive initialization methods that evolve during training. 4. Extension to other neural architectures beyond Transformers. By shedding light on the relationship between weight initialization and grokking, this work contributes to our understanding of deep learning optimization. These insights could lead to more efficient training strategies, faster convergence, better generalization, and potentially reduced computational requirements for training large models. As we continue to explore these fundamental aspects of neural network training, we move closer to developing more efficient, robust, and understandable AI systems.  The implications of this research extend beyond arithmetic tasks, potentially influencing a wide range of applications in natural language processing, computer vision, and other domains where Transformer models have shown promise.
R EFERENCES
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Beñgio.  Neural machine translation by jointly learning to align and translate. arXiv preprint arXiv: I409.0473 , 2014. Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. pp. 249–256, 2010. Ian Goodfellow, Yoshua Bengio, Aaron Courville, and Yoshua Bengio. Deep learning, volume 1. MIT Press, 2016. Kaiming He, X. Zhang, Shaoqing Ren, and Jian Sun.  Delving deep into rectifiers:  Surpassing human-level performance on imagenet classification. 2015 IEEE International Conference on Computer Vision (ICCV) , pp. 1026–1034, 2015. Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014. Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization. arXiv  preprint arXiv:1711.05101, 2017. Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gen- eralization beyond overfitting on small algorithmic datasets. arXiv preprint arXiv:2201.02177, 2022. Andrew M. Saxe, James L. McClelland, and S. Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. CoRR , abs/1312.6120, 2013. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.