Pith. sign in

REVIEW 5 major objections 6 minor 1 cited by

Mechanistic Insights into Grokking from the Embedding Layer

T0 review · 5 major / 6 minor · reviewed 2026-08-07 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗ Review JSON

Pith's one-line read Trainable embeddings are the main source of grokking, and a learning-rate ratio derived from the bilinear coupling between embeddings and weights shortens the delay.

desk verdict A genuinely new empirical observation about embeddings and grokking, wrapped in a theory that overreaches—deserves review but needs major revision. read the letter →

arxiv 2505.15624 v1 pith:LQ7N56XB submitted 2025-05-21 cs.LG cs.CL

H.V.AlquBoj , Hilal AlQuabeh , Velibor Bojkovic , Munachiso Nwadike , Kentaro Inui This is my paper · ORCID
classification cs.LGcs.CL
keywords grokkingembeddinglayerbilinearcouplingmodulararithmeticlearningrateratioAdamoptimizertokensamplinggeneralizationdelay
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

This paper argues that the embedding layer is the main trigger of grokking, the long gap between perfect training accuracy and sudden test generalization, in MLPs trained on modular arithmetic. It shows that an MLP without embeddings generalizes immediately on modular addition, while adding an embedding layer produces the memorization-then-generalization pattern known as grokking. The delay is traced to two mechanisms: rare tokens receive almost no gradient updates and are slowly eroded by weight decay, and the bilinear product of embeddings and first-layer weights creates an asymmetric loss landscape with saddle points that slow optimization. To counteract this, the paper proves a learning-rate ratio, $\eta_E/\eta_W \propto \sigma_{\max}(E)/\sigma_{\max}(W) \cdot f_W/f_E$, that equalizes effective update sizes, and shows empirically that a fixed ratio of about ten, combined with balanced token sampling, substantially shortens the grokking gap across modular addition, multiplication, division, and sum-of-squares. A reader should care because the same bilinear embedding-weight interaction also appears in Transformer attention, so the diagnosed imbalance may explain broader optimization difficulties in Transformers.

What carries the argument

The load-bearing object is the bilinear coupling between the embedding matrix $E$ and the first-layer weight matrix $W$, through the term $z(EW)$. Because the gradient with respect to $E$ is proportional to $W^\top \nabla L$ and the gradient with respect to $W$ is proportional to $E^\top \nabla L$, the update sizes are governed by the other matrix's largest singular value: $\|\nabla E\| \propto \sigma_{\max}(W)$ and $\|\nabla W\| \propto \sigma_{\max}(E)$. The paper's key identity is the learning-rate ratio $c = \eta_E/\eta_W \propto (\sigma_{\max}(E)/\sigma_{\max}(W)) \cdot (f_W/f_E)$, where $f_E$ and $f_W$ are update frequencies, together with the practical recipe of approximately $c=10$ and frequency-aware sampling that minimizes gradient variance. This identity is what turns the observed Hessian asymmetry between $E$ and $W$ into an actionable optimizer (Adam-LR) that balances the two parameter blocks.

What would settle it

Run the same modular-arithmetic MLP with uniform token frequencies so $f_E = f_W$, and initialize $E$ so $\sigma_{\max}(E) \approx \sigma_{\max}(W)$; if grokking still shows the same delay and $c=10$ still accelerates, the singular-value and frequency ratio is not the operative mechanism. Alternatively, measure $\|\nabla E\|/\|\nabla W\|$ under Adam across training and check whether it tracks $\sigma_{\max}(W)/\sigma_{\max}(E)$; if it drifts by an order of magnitude, the core scaling assumption fails.

Watch

Extended reading notes

Core claim

The central claim is that trainable embeddings, not network depth or activation choice, are what make grokking appear in modular-arithmetic MLPs. When the first layer multiplies a learned embedding matrix $E$ by a weight matrix $W$, the loss landscape becomes bilinear: gradients with respect to $E$ scale with the spectrum of $W$ and vice versa, rare tokens are updated only when they happen to appear in a batch, and weight decay pushes all embeddings toward zero in the meantime. These effects delay stabilization of $E$, and generalization arrives only once embeddings settle into task-aligned, frequency-localized representations, so the delay shows up as grokking. The paper's constructive contribution is the proof that setting the embedding learning rate relative to the downstream rate as $c \propto \sigma_{\max}(E)/\sigma_{\max}(W) \cdot f_W/f_E$ cancels the update imbalance; with this Adam-LR recipe ($c \approx 10$), the Hessian eigenvalues of $E$ and $W$ become more balanced, validation accuracy rises earlier, and grokking is accelerated on all four tested modular tasks.

Load-bearing premise

The derivation assumes that embedding and weight gradient norms scale with the other matrix's largest singular value, $\|\nabla E\| \propto \sigma_{\max}(W)$ and $\|\nabla W\| \propto \sigma_{\max}(E)$, and that this proportionality survives Adam's per-parameter normalization and random batch sampling; if that scaling fails, the proposed learning-rate ratio is not justified and $c=10$ becomes a fitted constant.

Editorial extensions

If this is right

  • MLPs with trainable embeddings on modular arithmetic will grok, whereas embedding-free MLPs generalize immediately on addition but fail on multiplication; the embedding layer is the trigger for the delayed generalization.
  • Rare-token stagnation is driven by sampling probability and weight decay, so dataset splits that skew token frequencies reliably produce suboptimal generalization plateaus.
  • Uniform token sampling accelerates grokking at moderate batch sizes, while at large batch sizes random and uniform sampling converge because per-batch token coverage is sufficient.
  • Setting $\eta_E/\eta_W \approx \sigma_{\max}(E)/\sigma_{\max}(W) \cdot f_W/f_E$, practically around 10, balances embedding and weight update scales, reduces the Hessian eigenvalue gap, and speeds up grokking on modular addition, multiplication, division, and sum-of-squares.
  • The same bilinear interaction appears in Transformer attention through query-key-value projections, so the imbalance diagnosed here should also hinder Transformer optimization and is addressable by similar block-wise learning-rate corrections.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • If the coupling account is right, grokking should also be controllable by initializing $\sigma_{\max}(E)$ close to $\sigma_{\max}(W)$ or by normalizing the two blocks, independent of learning rate; this is testable but not tested in the paper.
  • The paper derives an entropy-regularized softmax sampling rule over gradient norms in Appendix A but does not run it as a standalone optimizer; a direct comparison of that rule against uniform sampling would isolate how much of the speedup comes from sampling versus learning-rate scaling.
  • The Hessian split the paper observes, where $W$ drives early progress and $E$ tunes late generalization, predicts that freezing $E$ after partial training should preserve or accelerate grokking, an experiment not reported here.
  • Because the ratio formula depends on singular values that drift during training, a fully adaptive version that recomputes $c$ periodically could outperform the fixed $c=10$; the paper only uses a constant multiplier.
Share X Bluesky LinkedIn Reddit HN

Signed reviews

No signed human review yet.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

5 major / 6 minor

Summary. The paper studies grokking in MLPs on modular arithmetic and claims that the embedding layer is the key component causing delayed generalization. It proposes two mechanisms: sparse gradient updates and weight decay for rare tokens, and bilinear coupling between embeddings and first-layer weights. It then proposes frequency-aware sampling and a higher embedding learning rate, claiming a proof that the optimal learning-rate ratio scales as σmax(E)/σmax(W) · fW/fE and setting it to c=10. Experiments show that adding embeddings induces grokking in MLPs and that both balanced sampling and c=10 accelerate convergence.

Significance. If the empirical claims are upheld, the paper offers a simple and instructive demonstration that trainable embeddings can create grokking-like dynamics in MLPs, and that practical interventions—balancing token frequencies and raising the embedding learning rate—can markedly reduce the delay. The Hessian and rank-evolution analyses are interesting probes. However, the theoretical contribution as presented is not reliable: the SGD-based derivation does not match the Adam experiments, the key proportionality assumptions are unverified, and the claimed proof of the learning-rate ratio is effectively circular with the fitted constant c=10. The paper also lacks statistical rigor and does not compute the proposed formula from measured quantities. The central qualitative finding is plausible, but the mechanistic explanation and the 'proof' need substantial revision.

major comments (5)
  1. [Section 5 and Appendix B (Prop. 4.1)] The theoretical derivation is for SGD, but all experiments use Adam. Under Adam, the per-coordinate normalization changes the update rule fundamentally: if the raw gradient is scaled by a constant s, both the first and second moment estimates scale, making the Adam update approximately invariant to that scaling. Consequently, the proportionality ∥ΔE∥/∥ΔW∥ ∝ (η_E σmax(W) f_E)/(η_W σmax(E) f_W) in Eq. (23) does not hold for Adam, and the choice c=10 is not derived from the stated theory; it is a fitted hyperparameter. The authors must either provide an Adam-specific analysis or restrict the theoretical claims to SGD and add matching experiments.
  2. [Appendix B, Eq. (22) and Prop. 4.1] The assumptions ∥∇E∥ ∝ σmax(W) and ∥∇W∥ ∝ σmax(E) are not proven and are not verified empirically. Moreover, the derivation is circular: substituting these proportionality assumptions into Eq. (23) yields a ratio in which the σ factors cancel exactly under the proposed choice c ∝ σmax(E)/σmax(W) · fW/fE, so Eq. (24) is an identity under the ansatz rather than a substantive constraint. The paper also states that standard initialization ensures σmax(E) ≫ σmax(W), but for the stated dimensions (E roughly V×128, W roughly 512×128 under default 1/√fan-in scaling) random-matrix spectral norms do not support this ordering. No measured singular values or update frequencies are reported to validate Eq. (24).
  3. [Section 4.3, Eqs. (8)-(9)] There is a dimensional inconsistency between the ODE and its solution. The ODE in Eq. (8), de_i/dt = -λ e_i - p_i ∇e_i L, contains no learning rate η, but the solution in Eq. (9) includes the term η p g/λ. For constant g, the correct solution is C e^{-λt} - p_i g/λ, not C e^{-λt} - η p_i g/λ. This error propagates to the convergence-time bound T ≥ (1/λ) ln(C/ε) and should be corrected or the scaling must be explained.
  4. [Section 4.2 and Appendix A] The abstract claims the paper investigates 'frequency-aware sampling, which balances token updates by minimizing gradient variance,' but the experiments use three data-split strategies (random, uniform, skewed) chosen heuristically; the entropy-regularized optimal probabilities derived in Appendix A are not implemented or tested. In addition, the optimization in Eq. (7) has a degenerate vertex solution (all probability mass on the largest gradient norm), which is not the uniform distribution used in the paper. The connection between the proposed sampling theory and the experimental splits needs to be made explicit, or the abstract and contribution bullets should be revised to match what is actually done.
  5. [Figures 4-6 and Sections 5.1-5.2] The experimental claims lack statistical support: no number of seeds, no error bars or confidence intervals, and no ablations that separate the effect of sampling from the learning-rate ratio. In particular, Fig. 5 shows the sensitivity sweep for c but no variance across runs, and Section 5.2 reports only a single comparison per dataset. Given that the central theoretical claim concerns the optimality of a specific c, the paper should report the distribution of steps-to-accuracy across multiple seeds and, ideally, the measured quantities entering Eq. (24) (σmax(E), σmax(W), fE, fW) to demonstrate that the predicted ratio matches the empirically optimal c.
minor comments (6)
  1. [Section 4.1, Eq. (5)] Equation (5) has a sign error: the first term should be ∇θ_t L^T (θ_{t+1} - θ_t), not ∇θ_t L^T (θ_{t+1} + θ_t).
  2. [Section 3.3 and Appendix B] The architecture dimensions are inconsistent: with input dimension 4d and hidden dimension 4d, the first-layer weight is (4d)×(4d), but Appendix B writes W∈R^{4d×d}. Please reconcile the dimensions.
  3. [Figure 6 caption] The caption says 'Adam-LR scales the embedding learning rate based on the singular values of the embedding matrix,' but the actual method multiplies the embedding learning rate by a fixed c=10; clarify that the scaling is fixed, not computed from observed singular values.
  4. [Abstract and Section 1] The word 'prove' is used in the abstract and in the contribution list. Given that Appendix B is a heuristic proportionality argument with unverified assumptions, please rephrase to 'argue' or 'propose' unless a rigorous proof is supplied.
  5. [References] The reference list contains incomplete entries (e.g., [9] 'PhD thesis, none') and missing publication details; please verify and complete all references.
  6. [Figure 4 caption] The caption states 'Each row corresponds to a batch size,' but the layout within each subfigure shows three batch sizes per dataset; please adjust the wording to match the actual figure arrangement.

Circularity Check

2 steps flagged · score 6.0 of 10

The adaptive ratio 'proof' is a by-construction balance condition, and c=10 is chosen by the same sensitivity sweep later reported as validation.

  1. fitted input called prediction [Section 4.3 (Prop. 4.1) and Fig. 5 / Section 5.2]
    "In practice, we set c=10, guided by empirical singular value trends and supported by sensitivity analysis (see Fig. 5, §5.2). This adjustment improves convergence and stability, especially under sparse embedding updates common in skewed token distributions."

    The advertised ratio is never evaluated from measured singular values or update frequencies; instead c=10 is chosen because the sensitivity sweep in Fig. 5 shows that c=10 consistently balances convergence and stability. Section 5.2 then reports that Adam-LR, which scales the embedding learning rate by a factor of 10, significantly accelerates grokking. The confirmation is therefore the selection criterion, not an independent prediction: the successful value was fitted on the same tasks and then presented as the validated outcome.

  2. self definitional [Appendix B, Eqs. (23)-(24); Proposition 4.1]
    "To equalize update scales under cross-entropy loss, the learning rate ratio c=ηE/ηW should satisfy: c∝ σmax(E)/σmax(W) · fW/fE, where σmax(·) denotes the largest singular value and fE, fW are the respective update frequencies (see appendix B for details). For proportional updates (∥∆E∥ ∼ ∥∆W∥), the ratio c=ηE/ηW must satisfy: c∝ σmax(E)/σmax(W) · fW/fE. (24)"

    Equation 23 defines the effective update ratio as proportional to ηE σmax(W) fE / (ηW σmax(E) fW). Solving for c that makes this ratio equal to 1 is exactly what produces the formula in Proposition 4.1. Hence the proposition restates the definition of equal update scales rather than deriving why that balance mitigates bilinear coupling or accelerates grokking; the acceleration claim depends on the empirical sweep, not on the proof. The proof is thus self-definitional for the central ratio.

full rationale

The paper's main empirical content—embeddings induce grokking in MLPs on modular arithmetic, uniform token sampling reduces the delay, and a larger embedding learning rate helps—is supported by direct experiments and is not circular. There is no load-bearing self-citation or imported uniqueness theorem. The circularity is concentrated in the advertised proof of the adaptive learning-rate ratio. Prop. 4.1 and Appendix B only solve for the c that makes the SGD effective update ratio equal to one; this is a balance condition by construction, not a proof of accelerated convergence. Moreover, the concrete value c=10 is not obtained by measuring σmax(E), σmax(W), fW, or fE; it is picked from a sensitivity sweep that already shows c=10 works, and then the same value is reported as the success of the proposed method. The SGD-based derivation is also applied to Adam experiments without accounting for Adam's normalization, but that is a correctness gap, not circularity. Because the central ratio and its numerical value reduce to a fitted choice, the score is 6.

Assumptions & free parameters 2 free parameters · 6 assumptions · 0 invented entities

The analysis rests on a series of stated but unproven premises about gradient magnitude scaling, batch sampling, gradient stabilization, and initialization spectra. The count below shows the paper's theory section has more assumptions than theorems.

free parameters (2)
  • c = ηE/ηW = 10
    Set to 10 in all Adam-LR experiments after a sensitivity sweep (Fig 5); the formula in Prop 4.1 is not used to compute it.
  • weight decay λ = 0.001
    Hand-chosen, standard in grokking experiments; the embedding dynamics analysis in Eqs 3-9 depends on it, and convergence time scales as 1/λ.
assumptions (6)
  • standard math The loss L(θ,{e_i}) is β-smooth.
    Used in Eq 5 to bound the expected loss decrease.
  • ad hoc to paper The embedding gradient stabilizes to a constant g during training.
    Needed to solve the ODE Eq 8 and obtain Eq 9; no evidence given that gradients are constant.
  • ad hoc to paper Gradient norms of E and W are proportional to σmax of the other matrix.
    Appendix B Eq 22; true only under alignment or orthogonality, not generally, and ignores Adam normalization.
  • domain assumption The probability that token i appears in a batch is p_i, so the expected gradient contribution is η p_i ∇ei L.
    Eq 4; ignores the batch size factor (1-(1-p_i)^B) and treats token sampling as per-example.
  • domain assumption PyTorch default initialization yields σmax(E) ≫ σmax(W).
    Appendix B; asserted without derivation and likely false for the stated dimensions (V≈100, d=128 vs W of size 512×d).
  • standard math Weight decay drives nuclear norm regularization of the product EW.
    Cited to [8]; used in Appendix C.2 to interpret rank evolution.

how reviews work

0 comments
Cite this review

Pith. "Pith review of Mechanistic Insights into Grokking from the Embedding Layer." pith.science (2026). https://pith.science/paper/LQ7N56XB

@misc{pith2026250515624,
  author       = {Pith},
  title        = {Pith review of: Mechanistic Insights into Grokking from the Embedding Layer},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/LQ7N56XB}},
  note         = {Machine review of arXiv:2505.15624}
}
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.

Figures

Figures reproduced from arXiv: 2505.15624 by the authors.

Figure 1
Figure 1. Heatmaps for (a) additive group (mod 6) and (b) multiplicative group (mod 7). The two groups are iso￾morphic despite differing appearances. The Transformer model [22] utilizes a self-attention mechanism to capture dependencies between tokens. In this framework, embeddings map input tokens to high￾dimensional vectors, which are processed through atten￾tion layers. These embeddings help the model capture contextualize… view at source ↗
Figure 2
Figure 2. Training and validation accuracies of the MLP model on modular arithmetic tasks, trained [PITH_FULL_IMAGE:figures/full_fig_p005_2.png] view at source ↗
Figure 3
Figure 3. Gradient heat maps of the MLP model at random optimization steps. Sparse columns in [PITH_FULL_IMAGE:figures/full_fig_p006_3.png] view at source ↗
Figures from the paper (10 more)
Figure 4
Figure 4. Figure 4: Training and validation accuracies for the modular datasets with a learning rate of 0.001, [PITH_FULL_IMAGE:figures/full_fig_p007_4.png]
Figure 5
Figure 5. Figure 5: Sensitivity of test accuracy to the learn [PITH_FULL_IMAGE:figures/full_fig_p008_5.png]
Figure 6
Figure 6. Figure 6: Performance comparison of Adam-LR and Adam optimizers on four algorithmic datasets. [PITH_FULL_IMAGE:figures/full_fig_p009_6.png]
Figure 7
Figure 7. Figure 7: Maximum eigenvalues of the Hessian with respect to embedding weights ( [PITH_FULL_IMAGE:figures/full_fig_p010_7.png]
Figure 8
Figure 8. Figure 8: Token probabilities in the training and test sets under different sampling strategies. Im [PITH_FULL_IMAGE:figures/full_fig_p013_8.png]
Figure 9
Figure 9. Figure 9: Maximum eigenvalues of the Hessian with respect to embedding weights ( [PITH_FULL_IMAGE:figures/full_fig_p014_9.png]
Figure 10
Figure 10. Figure 10: Rank evolution during training for three optimization setups: Adam (wd=0.001), Adam [PITH_FULL_IMAGE:figures/full_fig_p015_10.png]
Figure 11
Figure 11. Figure 11: Discrete Fourier analysis of learned embedding representations across tasks. For each [PITH_FULL_IMAGE:figures/full_fig_p016_11.png]
Figure 12
Figure 12. Figure 12: Training and validation accuracies for the modular multiplication dataset for learning rate [PITH_FULL_IMAGE:figures/full_fig_p017_12.png]
Figure 13
Figure 13. Figure 13: Training and validation accuracies for the modular multiplication dataset for learning rate [PITH_FULL_IMAGE:figures/full_fig_p017_13.png]

Discussion (0). Continue with ORCID to comment.

Forward citations

Cited by 1 Pith paper

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score. Full citation record

  1. At-Grok Is Not Converged:A Measurement-Validity Audit for Grokking Representation Metrics

    cs.LG 2026-07 accept novelty 6.5 of 10

    Embedding effective rank at grokking is a transient that overstates the converged floor by 3–5× (MLP) / 1.3–1.5× (transformer), and compression lags generalization by order T_grok, modulated by LayerNorm.

Reference graph

Works this paper leans on

29 extracted references · 12 canonical work pages · cited by 1 Pith paper

  1. [1]

    Davies, L

    X. Davies, L. Langosco, and D. Krueger. Unifying grokking and double descent.arXiv preprint arXiv:2303.06173, 2023

    arXiv 2023

  2. [2]

    Doshi, T

    D. Doshi, T. He, A. Das, and A. Gromov. Grokking modular polynomials.arXiv preprint arXiv:2406.03495, 2024

    arXiv 2024

  3. [3]

    S. Fan, R. Pascanu, and M. Jaggi. Deep grokking: Would deep neural networks generalize better?arXiv preprint arXiv:2405.19454, 2024

    arXiv 2024

  4. [4]

    A. Gromov. Grokking modular arithmetic.arXiv preprint arXiv:2301.02679, 2023

    arXiv 2023

  5. [5]

    X. S. Huang, F. Perez, J. Ba, and M. V olkovs. Improving transformer optimization through better initialization. InInternational Conference on Machine Learning, pages 4475–4483. PMLR, 2020

    work page 2020

  6. [6]

    A. I. Humayun, R. Balestriero, and R. Baraniuk. Deep networks always grok and here is why. arXiv preprint arXiv:2402.15555, 2024

    arXiv 2024

  7. [7]

    Jeffares, A

    A. Jeffares, A. Curth, and M. van der Schaar. Deep learning through a telescoping lens: A simple model provides empirical insights on grokking, gradient boosting & beyond.Advances in Neural Information Processing Systems, 37:123498–123533, 2024

    work page 2024

  8. [8]

    Kobayashi, Y

    S. Kobayashi, Y . Akram, and J. V on Oswald. Weight decay induces low-rank attention layers. Advances in Neural Information Processing Systems, 37:4481–4510, 2024. 10

    work page 2024

Show all 29 references
  1. [9]

    Kumar.Grokking as the transition from lazy to rich training dynamics

    T. Kumar.Grokking as the transition from lazy to rich training dynamics. PhD thesis, none, 2024

    2024

  2. [10]

    Kunstner, J

    F. Kunstner, J. Chen, J. W. Lavington, and M. Schmidt. Noise is not the main factor behind the gap between sgd and adam on transformers, but sign descent might be.arXiv preprint arXiv:2304.13960, 2023

    2023 arXiv

  3. [11]

    J. Lee, B. G. Kang, K. Kim, and K. M. Lee. Grokfast: Accelerated grokking by amplifying slow gradients.arXiv preprint arXiv:2405.20233, 2024

    2024 arXiv

  4. [12]

    Z. Liu, O. Kitouni, N. S. Nolte, E. Michaud, M. Tegmark, and M. Williams. Towards un- derstanding grokking: An effective theory of representation learning.Advances in Neural Information Processing Systems, 35:34651–34663, 2022

    2022

  5. [13]

    Z. Liu, E. J. Michaud, and M. Tegmark. Omnigrok: Grokking beyond algorithmic data. InThe Eleventh International Conference on Learning Representations, 2022

    2022

  6. [14]

    K. Lyu, J. Jin, Z. Li, S. S. Du, J. D. Lee, and W. Hu. Dichotomy of early and late phase implicit biases can provably induce grokking.arXiv preprint arXiv:2311.18817, 2023

    2023 arXiv

  7. [15]

    M. A. Mohamadi, Z. Li, L. Wu, and D. J. Sutherland. Why do you grok? a theoretical analysis of grokking modular addition.arXiv preprint arXiv:2407.12332, 2024

    2024 arXiv

  8. [16]

    Nakkiran, G

    P. Nakkiran, G. Kaplun, Y . Bansal, T. Yang, B. Barak, and I. Sutskever. Deep double descent: Where bigger models and more data hurt.Journal of Statistical Mechanics: Theory and Experiment, 2021(12):124003, 2021

    2021

  9. [17]

    Nanda, L

    N. Nanda, L. Chan, T. Lieberum, J. Smith, and J. Steinhardt. Progress measures for grokking via mechanistic interpretability.arXiv preprint arXiv:2301.05217, 2023

    2023 arXiv

  10. [18]

    Pan and Y

    Y . Pan and Y . Li. Toward understanding why adam converges faster than sgd for transformers. arXiv preprint arXiv:2306.00204, 2023

    2023 arXiv

  11. [19]

    Power, Y

    A. Power, Y . Burda, H. Edwards, I. Babuschkin, and V . Misra. Grokking: Generalization beyond overfitting on small algorithmic datasets.arXiv preprint arXiv:2201.02177, 2022

    2022 arXiv

  12. [20]

    Prieto, M

    L. Prieto, M. Barsbey, P. A. Mediano, and T. Birdal. Grokking at the edge of numerical stability. arXiv preprint arXiv:2501.04697, 2025

    2025 arXiv

  13. [21]

    Varma, R

    V . Varma, R. Shah, Z. Kenton, J. Kramár, and R. Kumar. Explaining grokking through circuit efficiency.arXiv preprint arXiv:2309.02390, 2023

    2023 arXiv

  14. [22]

    A. Vaswani. Attention is all you need.Advances in Neural Information Processing Systems, 2017

    2017

  15. [23]

    Z. Xu, Z. Ni, Y . Wang, and W. Hu. Let me grok for you: Accelerating grokking via embedding transfer from a weaker model.arXiv preprint arXiv:2504.13292, 2025

    2025 arXiv

  16. [24]

    Yang and E

    G. Yang and E. J. Hu. Tensor programs iv: Feature learning in infinite-width neural networks. In M. Meila and T. Zhang, editors,Proceedings of the 38th International Conference on Machine Learning, volume 139 ofProceedings of Machine Learning Research, pages 11727–11737. PMLR,...

    2021

  17. [25]

    Zhang, S

    J. Zhang, S. P. Karimireddy, A. Veit, S. Kim, S. Reddi, S. Kumar, and S. Sra. Why are adaptive methods good for attention models?Advances in Neural Information Processing Systems, 33:15383–15393, 2020

    2020

  18. [26]

    block heterogeneity

    Y . Zhang, C. Chen, T. Ding, Z. Li, R. Sun, and Z.-Q. Luo. Why transformers need adam: A hessian perspective.arXiv preprint arXiv:2402.16788, 2024. 11 Appendix A Optimizing for Sampling Porbability Uniform Importance Assumption If we assume that all gradients are equally impor...

    2024 arXiv

  19. [27]

    Initialize a random vectorv 0 with the same dimensionality as the parameters[E,W]. 14 Figure 10: Rank evolution during training for three optimization setups: Adam (wd=0.001), Adam- LR (wd=0.001 with learning rate ratio), and Adam with stronger weight decay (wd=0.005). While a...

  20. [28]

    Compute the Hessian-vector productHv k using automatic differentiation: Hvk =∇ θ (∇θL ·vk), whereθ= [E,W]

  21. [29]

    Figure 9 shows the maximum eigenvalues of the Hessian with respect to E and W during training

    Normalize the vector and update the eigenvalue estimate: vk+1 = Hvk ∥Hvk∥ , σ max ≈v ⊤ k Hvk. Figure 9 shows the maximum eigenvalues of the Hessian with respect to E and W during training. The results highlight distinct curvature properties for E and W, reflecting their roles ...

Pith tools

Reviewed August 7, 2026 · model on record in the stance chip above.