Mechanistic Insights into Grokking from the Embedding Layer

TL;DR

This work analyzes grokking and delayed generalization, arguing that trainable embedding layers induce the phenomenon in MLPs solving modular arithmetic tasks. It identifies two mechanisms—sparse embedding updates for rare tokens and bilinear coupling between embeddings and first-layer weights that creates saddle points—and proposes frequency-aware sampling and embedding-specific learning-rate adjustments, notably an adaptive scheme with $c = \frac{\eta_E}{\eta_W}$ proportional to $\frac{\sigma_{\max}(\mathbf{E})}{\sigma_{\max}(\mathbf{W})} \cdot \frac{f_W}{f_E}$, empirically set to $10$, to accelerate convergence. The study also demonstrates that balancing update dynamics via Adam-LR improves stability and generalization, with Hessian analyses showing more balanced curvature between embedding and downstream weights. While focused on MLPs, the findings have broader implications for Transformer optimization where bilinear interactions also play a role, offering practical strategies to mitigate grokking and enhance generalization in bilinear systems.

Abstract

Grokking, a delayed generalization in neural networks after perfect training performance, has been observed in Transformers and MLPs, but the components driving it remain underexplored. We show that embeddings are central to grokking: introducing them into MLPs induces delayed generalization in modular arithmetic tasks, whereas MLPs without embeddings can generalize immediately. Our analysis identifies two key mechanisms: (1) Embedding update dynamics, where rare tokens stagnate due to sparse gradient updates and weight decay, and (2) Bilinear coupling, where the interaction between embeddings and downstream weights introduces saddle points and increases sensitivity to initialization. To confirm these mechanisms, we investigate frequency-aware sampling, which balances token updates by minimizing gradient variance, and embedding-specific learning rates, derived from the asymmetric curvature of the bilinear loss landscape. We prove that an adaptive learning rate ratio, \(\frac{η_E}{η_W} \propto \frac{σ_{\max}(E)}{σ_{\max}(W)} \cdot \frac{f_W}{f_E}\), mitigates bilinear coupling effects, accelerating convergence. Our methods not only improve grokking dynamics but also extend to broader challenges in Transformer optimization, where bilinear interactions hinder efficient training.

Figures (13)

• Figure 1: Heatmaps for (a) additive group $(\text{mod } 6)$ and (b) multiplicative group $(\text{mod } 7)$. The two groups are isomorphic despite differing appearances.
• Figure 2: Training and validation accuracies of the MLP model on modular arithmetic tasks, trained with Adam. Left two: Addition task, without (first) and with (second) embeddings. Right two: Multiplication task, without (third) and with (fourth) embeddings. In the embedding-free cases, training and validation accuracies increase together only for addition; multiplication fails to generalize. In contrast, models with embeddings reach 100% training accuracy in both tasks, but only begin generalizing after a delay exhibiting the grokking phenomenon.
• Figure 3: Gradient heat maps of the MLP model at random optimization steps. Sparse columns in the embedding gradients reflect the absence of certain tokens in sampled batches, leading to uneven learning dynamics and contributing to delayed generalization.
• Figure 4: Training and validation accuracies for the modular datasets with a learning rate of 0.001, comparing different sampling strategies (random, uniform, skewed) across batch sizes (64, 128, 512). Each row corresponds to a batch size, and each column represents a dataset. The x-axis is logarithmic to emphasize convergence trends. Uniform sampling promotes faster generalization and convergence compared to random sampling, but the performance gap diminishes for batch sizes exceeding 512. Skewed sampling, while fitting the training data well, fails to generalize, highlighting the detrimental impact of imbalanced token updates.
• Figure 5: Sensitivity of test accuracy to the learning rate ratio $c = \eta_E / \eta_W$ across four tasks. Small $c$ leads to under-updating, large $c$ causes instability, and $c = 10$ consistently balances convergence and stability.

Theorems & Definitions (1)

• Proposition 4.1