Reproduce “Mechanistic Insights into Grokking from the Embedding Layer” on OpenResearch

About this paper

H.V.AlquBoj, Hilal AlQuabeh, Velibor Bojkovic, Munachiso Nwadike, Kentaro Inui

A mechanistic analysis reveals that embedding layers are the primary drivers of the grokking phenomenon in neural networks, showing how their bilinear coupling and token frequency-dependent updates lead to delayed generalization. The study proposes Adam-LR, a modified optimizer that applies different learning rates to embedding parameters, which accelerates grokking by balancing update magnitudes and mitigating initialization sensitivity.

Key insights

• Embedding layers are not passive input transformers but active drivers of the grokking phenomenon due to their bilinear coupling with downstream weights.
• Token frequency during training creates an imbalance in embedding updates, causing rare tokens to lag in optimization progress.
• The bilinear coupling between embedding weights and the first layer weights introduces a complex optimization problem with saddle points and increased sensitivity to initialization.

Summary by alphaXiv.