{"id":"1a18848f-c0f0-4fc3-87b0-593ccac660ca","arxiv_id":"2505.15624","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":5.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":2,"one_line_summary":"Trainable embeddings in a simple MLP cause delayed generalization (grokking) on modular arithmetic, and a higher embedding learning rate plus balanced sampling accelerates it.","lead":"This paper claims that trainable embedding layers are what make small neural networks grok, suddenly generalize after overfitting, on modular arithmetic: adding embeddings to an MLP creates the delayed generalization, and removing them lets the network generalize right away. It proposes two fixes, more even token sampling and a higher learning rate for embeddings, and shows these fixes speed up the delayed generalization.","discovery_kind":"extension","skeptic_critique":{"model":"deepseek-v4-flash","headline":"The theory is derived for SGD but all experiments use Adam; Adam's per-coordinate normalization breaks both the rare-token update-frequency premise and the singular-value ratio proof, so the claimed proof does not cover the reported results.","rationale":"The reader's weakest assumption identified the same load-bearing point: the proof in Appendix B assumes raw gradient norms scale with the coupled matrix's largest singular value and that these relations survive Adam and batch sampling. I agree, and I would sharpen it: the mismatch is not merely an unproven proportionality. Adam's update is approximately scale-invariant in the gradient, so the entire update-ratio argument in Eq. 23 cannot hold for the optimizer used in the experiments. This matters because the paper's headline claim is a proof of the learning-rate ratio, and Adam-LR is presented as derived from that ratio. If the derivation does not apply to Adam, c=10 is a tuned constant, and the theoretical contribution collapses to a heuristic. The empirical observation that a 10x embedding learning rate accelerates grokking may still be true, so the conditional verdict remains appropriate: the paper should either prove the ratio for Adam, relabel Prop. 4.1 as a heuristic, rerun the experiments with SGD to test the raw-gradient theory, or provide code and multi-seed statistics to support the empirical claims. I do not see a reason to move the verdict to reject, because the central empirical directions (balanced sampling and larger embedding LR) are plausible and the deficiencies are addressable; but they are real and currently unresolved.","tokens_in":16098,"tokens_out":13277,"duration_ms":123267,"concrete_test":"Re-run the Fig. 6 comparison (Adam-LR vs Adam) with plain SGD under identical seeds, dataset splits, weight decay, and a learning-rate grid, measuring grokking step for c in {1, 10}. If c=10 does not accelerate SGD, the reported benefit is attributable to Adam's per-coordinate normalization rather than to the singular-value ratio in Prop. 4.1; as a supporting analytical check, verify whether Eq. 23 still holds when the raw update eta*grad is replaced by Adam's normalized step.","verdict_should_be":"UNCHANGED","load_bearing_attack":"Appendix B and Eqs. 3-4 model updates as eta times the raw gradient: E[update] = -eta*lambda*e_i - eta*p_i*grad, and then derive an update ratio ||Delta E||/||Delta W|| proportional to eta_E*sigma_max(W)*f_E/(eta_W*sigma_max(E)*f_W). But Section 5 states that every experiment uses the Adam optimizer. Adam does not implement this update rule. For a rare token absent from many batches, the second moment v decays; when the token reappears, the normalized step is roughly eta*sign(g), not eta*p_i*grad, so sparse tokens do not simply stagnate. Likewise, if the raw gradient is rescaled by a constant s, both m and sqrt(v) rescale by s, making Adam's update magnitude approximately invariant to the gradient-norm proportionality in Eq. 22. Consequently, the ratio c proportional to sigma_max(E)/sigma_max(W)*f_W/f_E in Prop. 4.1 is not the effective update-balance condition for Adam, and the choice c=10 in Fig. 5 is not derived from the stated theory; it is a fitted hyperparameter. The appendix also claims standard initialization gives sigma_max(E) much greater than sigma_max(W), but for the stated dimensions (E roughly V by 128, W roughly 512 by 512) random-matrix spectral norms give the opposite ordering, further undermining the numerical justification of c=10.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","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.","tokens_in":16349,"tokens_out":8439,"duration_ms":74029,"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":[{"comment":"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.","section":"Section 5 and Appendix B (Prop. 4.1)"},{"comment":"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).","section":"Appendix B, Eq. (22) and Prop. 4.1"},{"comment":"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.","section":"Section 4.3, Eqs. (8)-(9)"},{"comment":"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.","section":"Section 4.2 and Appendix A"},{"comment":"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.","section":"Figures 4-6 and Sections 5.1-5.2"}],"minor_comments":[{"comment":"Equation (5) has a sign error: the first term should be ∇θ_t L^T (θ_{t+1} - θ_t), not ∇θ_t L^T (θ_{t+1} + θ_t).","section":"Section 4.1, Eq. (5)"},{"comment":"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.","section":"Section 3.3 and Appendix B"},{"comment":"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.","section":"Figure 6 caption"},{"comment":"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.","section":"Abstract and Section 1"},{"comment":"The reference list contains incomplete entries (e.g., [9] 'PhD thesis, none') and missing publication details; please verify and complete all references.","section":"References"},{"comment":"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.","section":"Figure 4 caption"}],"recommendation":"major_revision","confidential_remarks":"The paper's use of 'prove' in the abstract is considerably stronger than the supporting derivations, and the selection of c=10 from a sensitivity sweep is presented as a derived constant. This is a presentation-integrity issue that should be addressed. In addition, no code is released, which limits reproducibility of the empirical results. The empirical core—embeddings induce grokking and balanced sampling plus a higher embedding LR helps—could be publishable after a major revision that either fixes the theory or reframes the paper as an empirical study."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Colleague,\n\nThe thing to know: this paper has a real empirical finding. Explicit trainable embeddings turn a scalar-input MLP that generalizes immediately on modular addition into a grokking network, and bumping the embedding learning rate (c=10) plus balanced token sampling shortens the grokking delay. That is new and worth taking seriously. The theory attached to it is not a proof, and the paper would be much better if it said so.\n\nWhat's good: the experimental comparison across addition, multiplication, division, and sum-of-squares is reasonably broad; the Hessian analysis in Fig 7 gives some support to the claim that E and W have different curvature; the balanced sampling result (uniform beats skew/random at small batch sizes) is sensible and consistent with the rare-token story.\n\nWhere it falls down: the abstract and Proposition 4.1 say 'we prove' that ηE/ηW ∝ σmax(E)/σmax(W) · fW/fE. Appendix B derives that under SGD-style raw gradient updates, with a proportionality that is asserted, not proven. All experiments use Adam. Adam's per-coordinate normalization breaks the rare-token stagnation premise (a token absent for many steps doesn't just have its update scaled by p_i; the second moment decays and the sign-dominated update fires) and also makes the update invariant to constant rescaling of gradient norms, which guts the σmax ratio argument. So c=10 is a tuned constant, not a derived quantity. The appendix claim that σmax(E) >> σmax(W) at initialization also looks wrong for the stated dimensions (E is V x 128, W is 512 x 128), so the numerical justification for c=10 is shaky.\n\nAlso Eq 4 uses p_i but the correct per-batch inclusion probability depends on batch size and sampling scheme; Eq 9 mixes η inconsistently. There are no seeds, no error bars, no code—for a phenomenon as seed-sensitive as grokking, that matters.\n\nBottom line: the empirical core is plausible and could be a useful contribution after revision. The 'proof' should be relabeled as a heuristic or repaired for Adam; the fitted constant should be reported as fitted; multi-seed statistics and code should be released. This deserves referee time—it is a solid draft with an overclaim, not a waste of time. I would send it to review with the expectation of major revision. I would not cite it as it stands.","headline":"A genuinely new empirical observation about embeddings and grokking, wrapped in a theory that overreaches—deserves review but needs major revision.","tokens_in":16939,"tokens_out":2920,"would_cite":false,"duration_ms":24832,"reading_group":"maybe","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"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.","keywords":["grokking","embedding layer","bilinear coupling","modular arithmetic","learning rate ratio","Adam optimizer","token sampling","generalization delay"],"falsifier":"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.","tokens_in":15845,"feed_emoji":"⚡","tokens_out":8270,"duration_ms":66563,"temperature":0.7,"pith_summary":"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.","feed_headline":"A 10x embedding learning rate cuts grokking delay","feed_subtitle":"Rare tokens stagnate and bilinear coupling stalls training; a 10x embedding learning rate accelerates grokking","key_machinery":"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.","core_discovery":"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.","pith_inferences":["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."],"forward_implications":["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."],"supporting_citations":[{"why":"Introduces grokking and the modular arithmetic datasets that the paper uses as its central testbed.","marker":"[19]"},{"why":"Provides the closest prior theoretical analysis of grokking modular addition with one-hot inputs; the paper extends this by treating embeddings as trainable and bilinearly coupled.","marker":"[15]"},{"why":"Supplies the effective-theory view of representation learning and the Fourier-analysis methodology reused to show that embeddings learn frequency-localized structure.","marker":"[12]"},{"why":"Shows grokking beyond algorithmic data and supports the claim that grokking is architecture-agnostic, which the paper builds on before isolating the embedding layer.","marker":"[13]"},{"why":"Gives the Hessian-based 'block heterogeneity' explanation of why Adam outperforms SGD on Transformers, which the paper attributes to bilinear embedding-weight coupling.","marker":"[26]"},{"why":"Attributes Adam's advantage to heavy-tailed gradient noise; the paper situates its own curvature-based account against this explanation.","marker":"[25]"},{"why":"Argues that gradient noise is not the main factor behind the SGD-Adam gap and points to sign descent, motivating the paper's focus on curvature and update scales.","marker":"[10]"},{"why":"Prescribes width-based per-layer scaling under independent layer evolution, which the paper contrasts with its own singular-value and frequency-based coupling scaling.","marker":"[24]"},{"why":"Defines the Transformer attention mechanism whose multiplicative query-key-value interactions the paper claims share the same bilinear coupling.","marker":"[22]"}],"fun_headline_variants":["Embeddings drive grokking; adaptive LR cuts delay","Adaptive embedding LR ratio accelerates grokking","Rare-token stagnation causes grokking; LR fix helps","Frequency-aware LR reduces grokking delay","Bilinear coupling stalls grokking; adaptive LR speeds"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"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.","fun_headline_variants_meta":{"raw":{"variants":["Embeddings drive grokking; adaptive LR cuts delay","Adaptive embedding LR ratio accelerates grokking","Rare-token stagnation causes grokking; LR fix helps","Frequency-aware LR reduces grokking delay","Bilinear coupling stalls grokking; adaptive LR speeds"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000949,"raw_usage":{"total_tokens":4077,"prompt_tokens":996,"completion_tokens":3081,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":612,"completion_tokens_details":{"reasoning_tokens":3017}},"tokens_in":612,"tokens_out":3081,"duration_ms":18906,"temperature":1.0,"reasoning_tokens":3017,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-07T15:14:13.418691+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"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.","supporting_citations":[{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"Supplies the effective-theory view of representation learning and the Fourier-analysis methodology reused to show that embeddings learn frequency-localized structure."},{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"Shows grokking beyond algorithmic data and supports the claim that grokking is architecture-agnostic, which the paper builds on before isolating the embedding layer."},{"cited_title":"Zhang, S","cited_arxiv_id":null,"evidence_quote":"Attributes Adam's advantage to heavy-tailed gradient noise; the paper situates its own curvature-based account against this explanation."},{"cited_title":"Yang and E","cited_arxiv_id":null,"evidence_quote":"Prescribes width-based per-layer scaling under independent layer evolution, which the paper contrasts with its own singular-value and frequency-based coupling scaling."}],"review_version":1}