Mechanistic Insights into Grokking from the Embedding Layer

TL;DR

This paper reveals that embeddings are crucial for grokking in neural networks, identifying mechanisms like embedding dynamics and bilinear coupling, and proposes adaptive learning rates to improve training efficiency.

Contribution

The study uncovers the role of embeddings in grokking, analyzes underlying mechanisms, and introduces adaptive learning rate strategies to enhance convergence in Transformers.

Findings

Embedding dynamics cause delayed generalization in MLPs.

Bilinear coupling introduces saddle points affecting training.

Adaptive learning rates mitigate bilinear effects, speeding up convergence.

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{\eta_E}{\eta_W} \propto \frac{\sigma_{\max}(E)}{\sigma_{\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.