Grokking › Key Parameters and Model Configuration

Understanding which hyperparameters matter most for grokking, and how they interact, is needed for both recreating the phenomenon and applying its lessons to practical training problems. The parameters fall into two categories: those that affect the memorization side of the tradeoff and those that affect the algorithmic side. Parameters that make memorization harder (more training examples, stronger regularization) tend to accelerate grokking. Parameters that make the algorithmic solution easier to stand for (appropriate model size, good architecture) also accelerate grokking. The interplay between these two categories determines whether and when grokking occurs.

Understanding how each hyperparameter affects grokking dynamics helps you design experiments and interpret results. When grokking does not occur in an experiment, diagnosing which parameter is responsible, too little weight decay, too few training steps, wrong model size, requires understanding the role each parameter plays. When grokking occurs faster than expected, understanding why helps you apply those insights to other settings.

The core hyperparameters are:

• vocab_size: Size of the token vocabulary ($p + 1$ for numbers 0 to $p-1$ plus the equals token). Larger vocabularies increase the dataset size quadratically, which can either help or hurt grokking depending on whether the additional data primarily helps memorization or generalization.
• d_model: Dimension of embeddings and hidden states (128 in our experiments). This controls the capacity available for both memorization and the algorithmic solution. Too small, and the model cannot stand for either solution well; too large, and memorization becomes too easy.
• nhead: Number of attention heads (4 heads in our experiments). Multiple heads allow the model to attend to multiple patterns simultaneously, which may be important for learning the compositional operations needed for modular arithmetic.
• num_layers: Number of transformer encoder layers (2 layers in our experiments). Deeper models have more representational capacity and may be able to learn more complex algorithms, but also have more capacity for memorization.
• dim_feedforward: Dimension of the feedforward network in each layer (512). This controls the capacity for internal transformations within each layer.
• dropout: Set to 0.0 because dropout introduces its own regularization that can obscure or interfere with clean grokking dynamics driven by weight decay alone.
• weight_decay: Regularization strength that drives the transition from memorization to generalization (1.0 in our main experiments). This is the most necessary hyperparameter for grokking behavior.
• learning_rate: Step size for optimization (1e-3). Balances training speed with stability. Both too-high and too-low values can disrupt grokking.
• n_epochs: Number of complete passes through the training data. Thousands to millions may be needed to observe full grokking in larger experiments.
• train_fraction: Proportion of data used for training (0.5 creates a challenging generalization task). Lower fractions make memorization cheaper and delay grokking; higher fractions give more signal for generalization.
• batch_size: Number of examples per training batch (512). Larger batches give more stable gradient estimates, which generally helps grokking by giving cleaner signals for the weight decay to work against.