{"id":"a5fc0f53-da5d-4e9a-9429-93faa7c642b0","arxiv_id":"2504.17243","paper_version":2,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":6,"one_line_summary":"NeuralGrok trains a softmax MLP to transform gradients via bilevel optimization, accelerating grokking on arithmetic tasks and correlating with reduced gradient entropy.","lead":"A small auxiliary neural network is trained to reshape the gradients of a transformer while it learns modular arithmetic, and this speeds up the moment when the model starts to generalize. The paper also introduces a gradient-entropy measure that tracks the switch from memorization to generalization.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"The central speedup rests on a 94–168-example Douter used for the amplifier's one-step lookahead, with no split-ratio or seed ablations; the authors themselves flag the 'narrow validation objectives' (Section 5).","rationale":"The reader's weakest assumption is exactly the Douter proxy: that minimizing loss on a tiny, same-distribution subset after a single transformed gradient step is a reliable signal for true generalization. My stress-test reading converges on the same point. The arithmetic is stark: for the two-variable mod-97 tasks, Douter is about 94 examples; for the five-variable mod-7 task, about 168 examples. The amplifier is updated using a one-step lookahead on that set, and no split-ratio ablation or seed variation is reported. The paper itself concedes in Section 5 and Appendix E that the narrow Douter objective may cause task-specific overfitting, which is a self-asserted limitation that should weigh on the verdict. I also note the missing ablation isolating the learned softmax weighting from plain gradient normalization; Section 3.2 shows gradient rescaling alone changes training dynamics, so the unique contribution of the neural-amplifier is not cleanly identified. These concerns are exactly why the existing CONDITIONAL verdict is appropriate, so I do not recommend changing it. The proposed concrete test would provide the missing evidence without requiring a change in the paper's core framing.","tokens_in":11411,"tokens_out":7453,"duration_ms":79257,"concrete_test":"Run NeuralGrok on (a+b) mod 97 and (ac+bd-e) mod 7 with at least 10 random seeds and three Dinner:Douter splits (e.g., 99:1, 49:1, 4:1), reporting median and IQR of steps to 95% test accuracy. Include one condition with Douter labels shuffled and one with the amplifier's softmax weights fixed to uniform (plain gradient normalization) under the same inner-loop batch schedule and T. If the speedup over fixed-uniform normalization disappears when Douter labels are shuffled, or if step counts vary by more than about 2x across splits or seeds, the central claim is not established.","verdict_should_be":"UNCHANGED","load_bearing_attack":"NeuralGrok's outer loop (Algorithm 2) optimizes the amplifier against Douter, a 2% slice of Dtrain: roughly 94 examples for the mod-97 tasks and roughly 168 for (ac+bd-e) mod 7. Every reported step-to-95% figure comes from a single run with no error bars, so the headline speedups (e.g., 1896 vs 8853) could reflect run-to-run grokking variance rather than a reliable method effect. More importantly, the claimed mechanism is that minimizing one-step loss on this tiny subset teaches a gradient transformation that produces held-out generalization. The paper never varies the Dinner:Douter ratio, never ablates the outer-loop signal (e.g., shuffled labels or a fixed random transformation trained with the same pipeline), and Section 3.2 only ablates the rescaling coefficient c. The authors explicitly acknowledge in Section 5 and Appendix E that the 'narrow validation objectives (Douter)' may force the amplifier into local task geometries. If the outer-loop signal is just fitting 94–168 examples, the observed acceleration may be memorization of Douter or a side effect of gradient normalization, not evidence for a learned generalizable gradient transformation.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper proposes NeuralGrok, a bilevel optimization method that trains an auxiliary MLP (the neural-amplifier) to transform the gradients of a transformer base model during grokking. The transformation is a softmax-weighted elementwise scaling of each gradient entry followed by renormalization to a constant norm c. The method is evaluated on five modular arithmetic tasks: (a+b) mod 97, (a-b) mod 97, (a*b) mod 97, (a*a-b) mod 97, and (a*c+b*d-e) mod 7. The central claim is that NeuralGrok accelerates grokking, reaching 95% test accuracy in fewer optimization steps than standard training and Grokfast-MA, with the hardest task showing 1896 steps versus 8853 for Grokfast-MA. The paper also introduces an Absolute Gradient Entropy (AGE) metric and argues that NeuralGrok stabilizes training and reduces model complexity, in contrast to weight decay, which can cause instability.","tokens_in":11703,"tokens_out":5452,"duration_ms":51345,"significance":"If the central claim holds, the paper makes a useful empirical contribution: it extends the Grokfast line of work by showing that a learned, gradient-modulating module can shorten the memorization phase on arithmetic tasks, and it documents that simple gradient normalization enables generalization on a task where standard training fails. The authors provide code and compare against external baselines on a disjoint test set, which are strengths. However, the current evidence is preliminary: all headline numbers come from single runs, the learned softmax component is not ablated against fixed transformations under the same protocol, and the outer-loop validation signal uses a tiny 2% subset of the training set without robustness checks. The AGE analysis is also qualitative and limited to one task. With added seeds, ablations, and split-ratio analysis, the empirical claims would be substantially strengthened.","major_comments":[{"comment":"All results in Table 1 and Figure 1 come from a single run per setting. Grokking step-to-generalization is known to vary substantially across seeds, including in the original Grokfast ablations, so the reported speedups (e.g., 1896 vs 8853 on (ac+bd-e) mod 7) may reflect run-to-run variance rather than a reliable method effect. The authors should report results over at least 5-10 seeds with medians and spreads (e.g., IQR or min-max) for each method and task.","section":"Section 3.1, Table 1"},{"comment":"The experiments do not isolate the learned softmax component p from the fixed rescaling/normalization. The only ablation in Section 3.2 varies the rescaling coefficient c, and Figure 3 shows that standard training with plain gradient normalization already generalizes on the hard task, reaching 95% at roughly 29770 steps. To support the claim that the neural-amplifier 'learns an optimal gradient transformation,' the authors should compare NeuralGrok, under the same Dinner/Douter split and outer-loop budget, against (i) a fixed uniform p with the same norm rescaling, (ii) a fixed random p, and (iii) an amplifier trained with the same bilevel pipeline but with shuffled labels in Douter. Without such controls, the acceleration cannot be attributed to the learned modulation rather than to the normalization itself.","section":"Section 3.2, Algorithm 2"},{"comment":"The outer-loop signal uses Douter, only 2% of Dtrain: roughly 94 examples for the mod-97 tasks and about 168 examples for the five-argument task. The paper neither varies the Dinner:Douter ratio nor tests the reliability of the one-step lookahead proxy on Douter, and it does not check seed sensitivity of the split. The authors themselves note in Section 5 and Appendix E that the 'narrow validation objectives (Douter)' may force the amplifier into local task geometries, so the central mechanism—that minimizing loss on this tiny subset teaches a generalizable gradient transformation—remains unvalidated. Please add split-ratio ablations (e.g., 98:2, 95:5, 90:10) and a control in which Douter labels are shuffled, to test whether the outer-loop signal is necessary for the speedup.","section":"Section 3 (data split), Section 5, Appendix E"},{"comment":"The reported 'minimal optimization steps' count only base-model inner-loop steps on Dinner; the additional amplifier updates in the outer loop and the extra hyperparameter tuning are not accounted for. Also, NeuralGrok trains the base model on 98% of Dtrain while baselines use 100% (with the 2% set aside as Douter), so the comparison mixes data-seen and compute-cost differences. The authors should report total wall-clock time or total optimizer steps including outer-loop updates, and state explicitly that base-model steps are the reported metric. This will clarify how much of the apparent acceleration is due to the transformation versus the comparison protocol.","section":"Section 3.1, Table 1"},{"comment":"The AGE analysis is qualitative: transition windows in Figure 5 are marked by hand, no quantitative definition or automated detector is given, and all entropy curves come from a single seed on one task (a+b mod 97). The statement that AGE 'consistently correlates' with phase transitions is therefore not supported by measurement. At minimum, define the window boundaries, show multiple seeds and tasks, and compare AGE against the existing AWE and weight-norm baselines to justify the claimed explanatory role.","section":"Section 4, Figures 5-6"}],"minor_comments":[{"comment":"Calling the elementwise multiplication p·g a 'rotation' is inaccurate, since p has nonnegative entries; it is a direction-changing positive rescaling, not a rotation. Please rephrase.","section":"Section 2, Eq. (2)"},{"comment":"The Appendix D heading says '(axc+bxd-e) mod 97' while the text and Figure 11 are about (ac+bd-e) mod 7; please fix the inconsistency.","section":"Appendix D, title"},{"comment":"Please define the meaning of '-' in Table 1; presumably it indicates failure to reach 95% within the step budget, but this should be stated explicitly.","section":"Table 1, caption"},{"comment":"There are several typos and wording issues: 'arguements' should be 'arguments', 'nuance reasoning mechanisms' should be 'nuanced reasoning mechanisms', and 'Godilocks zone' should be 'Goldilocks zone'. Please proofread.","section":"Section 3.1, Section 5"},{"comment":"In Algorithm 2, the copy M'(θ) is updated to θ' = θ - ηθ g'_θ, and the meta-gradient is written as ∇φ L(θ', Douter) = ∇φ L(θ - ηθ G(φ, gθ), Douter). This equality is only correct if the dependence of θ' on φ is made explicit; please write θ'(φ) to avoid confusion.","section":"Algorithm 2"}],"recommendation":"major_revision","confidential_remarks":"The manuscript's central claim may well be correct, but the evidence is currently single-run and lacks the controls needed to credit the learned-transformation mechanism. The missing seed variance analysis, fixed-transformation ablations, and outer-loop proxy validation are all fixable within the paper's scope, so major revision seems appropriate rather than rejection. I also note that the authors' own limitation statements in Section 5 and Appendix E corroborate the need for a deeper analysis of the tiny Douter validation set."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"First thing to know: this paper does something new—it trains a small MLP to transform gradients via a softmax weighting plus norm rescaling in a bilevel loop, explicitly to shorten grokking—and it reports speedups on five modular arithmetic tasks. The central claim is plausible, but the evidence lacks error bars over seeds, and the learned softmax is not isolated from plain gradient normalization as cleanly as the title implies.\n\nWhat is actually new: the closest prior work, Grokfast, uses a fixed low-pass filter on gradients. Here the transformation is learned per task, which is a meaningful difference. The outer-loop optimization against a small validation subset is a reasonable application of learned-optimizer ideas to grokking. The paper also documents that gradient normalization alone stabilizes training on these tasks better than weight decay, and that the combination is best on the hard task—useful empirical information.\n\nThe paper does some things well: the ablation on the rescaling coefficient c is thorough, and the robustness of NeuralGrok to c is shown. The limitations section is candid: the authors explicitly note the narrow Douter validation objectives and the poor cross-task transfer, and they include transfer experiments rather than hiding them. The AGE metric is a natural extension of AWE to gradients; it tracks phase transitions and could be a handy monitoring tool, though it is descriptive, not causal.\n\nSoft spots, in order of severity: first, no seeds or error bars anywhere. Grokking step counts are high-variance; a single run claiming 1896 vs 8853 could be within run-to-run spread. Second, the outer loop uses Douter of only 94–168 examples, and the paper never varies the Dinner:Douter split or replaces the outer-loop signal with a fixed or random transformation to show the learned component alone is responsible. The comparison to standard training with gradient normalization (Figure 3) does suggest the softmax adds speed beyond normalization, since NeuralGrok reaches 900 steps where plain norm takes 3450, but it's not the same training pipeline, so the isolation is imperfect. Third, the interpretability story rests on qualitative entropy dynamics, not a mechanism.\n\nWho this is for: researchers working on grokking or small-transformer generalization will get genuine value from the method and the honest failure modes. Outside that subfield, the significance is modest.\n\nRecommendation: this deserves peer review, not desk rejection. The revision should add multiple seeds, a fixed-transformation ablation, and a split-ratio sensitivity check. If those hold up, I'd cite it.","headline":"A genuinely new learned-gradient-transformation approach to accelerating grokking, honestly reported, but the speedup needs seed-level replication and a cleaner isolation of the learned component.","tokens_in":12192,"tokens_out":3729,"would_cite":false,"duration_ms":35161,"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":"NeuralGrok learns to reshape gradients and reaches 95% test accuracy on five modular arithmetic tasks faster than standard training or Grokfast.","keywords":["grokking","gradient transformation","bilevel optimization","modular arithmetic","neural-amplifier","generalization","absolute gradient entropy","transformer training"],"falsifier":"Train the same bilevel loop but replace $D_{\\text{outer}}$ with a corrupted or shuffled version (e.g., random labels or only even inputs). If NeuralGrok still reaches 95% test accuracy at the same speed, the outer-loop signal is not driving generalization-oriented gradient shaping; alternatively, if the speedup vanishes under a different 2% slice of the same training set, the result hinges on the specific slice rather than the method.","tokens_in":11181,"feed_emoji":"⚡","tokens_out":6035,"duration_ms":48872,"temperature":0.7,"pith_summary":"The paper proposes NeuralGrok, a method that trains a small auxiliary MLP — the neural-amplifier — to transform each gradient before it updates a transformer trained on modular arithmetic. The amplifier weights gradient entries with a softmax and rescales the gradient to constant norm, and it is tuned by a bilevel loop on a tiny validation slice of the training set. The paper shows that this accelerates grokking on five arithmetic tasks, with up to 4.67× fewer steps than Grokfast-MA on the hardest task, and that test accuracy stays stable after generalization, unlike standard training or weight-decay baselines. It also introduces the Absolute Gradient Entropy (AGE) metric, reporting that transformed gradients stay low-entropy while original gradients spike, linking the speedup to reduced model complexity.","feed_headline":"Learned gradient filter makes transformers grok 1.4–4.7× faster","feed_subtitle":"On five modular arithmetic tasks, it hits 95% test accuracy in up to 4.7× fewer steps, and stays stable.","key_machinery":"The neural-amplifier $G(\\varphi)$ is an MLP that maps each gradient entry to a softmax weight and then rescales the whole gradient to constant norm $c$: $g' = c \\cdot p \\cdot g / \\|p \\cdot g\\|_2$. It is trained by bilevel optimization: the inner loop updates the base transformer with transformed gradients, and the outer loop (every $T$ steps) updates the amplifier to minimize loss on $D_{\\text{outer}}$, using a one-step lookahead through a copy of the base model. The transformation is a learned rotation plus fixed rescaling, which the paper interprets as suppressing task-specific noise in cyclic modular arithmetic. The diagnostic is Absolute Gradient Entropy, $H(G) = -\\sum_i |g_i| \\ln |g_i|$, which rises during the memorization phase and falls during generalization; the amplifier keeps transformed gradients at lower AGE than original gradients.","core_discovery":"The central claim is that grokking's long memorization phase can be shortened by learning a task-specific gradient transformation rather than relying on hand-designed filters or regularization. In the inner loop, gradients $g$ from a training subset pass through the neural-amplifier, which produces a softmax distribution $p$ over gradient entries and forms $g' = c \\cdot p \\cdot g / \\|p \\cdot g\\|_2$; the base transformer is updated with $g'$. In the outer loop, the amplifier is updated to minimize loss on $D_{\\text{outer}}$, a 2% slice of the training data, evaluated through a one-step copy of the base model. On five modular arithmetic datasets — $(a+b) \\bmod 97$, $(a-b) \\bmod 97$, $(a\\cdot b) \\bmod 97$, $(a^2-b) \\bmod 97$, and $(ac+bd-e) \\bmod 7$ — this reaches 95% test accuracy in fewer steps than standard training and Grokfast-MA (e.g., 1896 vs 8853 steps on the hardest task), and test accuracy does not collapse after generalization. The paper also reports that standard gradient normalization alone, without the amplifier, stabilizes training better than weight decay on these tasks.","pith_inferences":["A stress test of the outer-loop proxy would be to permute, shrink, or corrupt $D_{\\text{outer}}$: if the speedup survives a corrupted validation signal, the amplifier is not actually tracking generalization.","The rotation-plus-rescaling view suggests a testable hypothesis: the amplifier learns to de-emphasize gradient entries tied to answer tokens or high-frequency components; inspecting learned $p$ values per parameter group would reveal the pattern.","If the AGE drop is causal rather than correlational, directly penalizing gradient entropy during training should reproduce part of the speedup; the paper does not run that experiment.","The low cross-task transfer implies practical value depends on per-task meta-tuning, so a natural extension is a shared amplifier trained across tasks with a meta-objective."],"forward_implications":["If correct, the neural-amplifier gives a drop-in gradient preprocessing step that shortens grokking on modular arithmetic transformers without changing the base architecture.","The bilevel formulation implies that a small validation slice can drive fast generalization, suggesting gradient shaping can substitute for careful weight-decay tuning.","Gradient normalization alone is reported to stabilize training and enable generalization on a task where standard training fails, so gradient magnitude control may be the operative regularizer in these settings.","The AGE metric provides a per-step complexity signal that rises in memorization and falls in generalization, offering a candidate progress measure for grokking.","Because learned transformations transfer poorly across even similar arithmetic tasks, the speedup is task-specific, which bounds direct reuse but points toward meta-learning shared transformations."],"supporting_citations":[{"why":"Supplies the grokking phenomenon, the modular arithmetic dataset construction, and the 50/50 train/test split used by all methods.","marker":"Power et al. (2022)"},{"why":"Provides the Grokfast-MA and Grokfast-EMA baselines that NeuralGrok must beat, and the precedent of manipulating gradients (low-pass filtering) to accelerate grokking.","marker":"Lee et al. (2024)"},{"why":"Introduces Absolute Weight Entropy, which the paper adapts into the Absolute Gradient Entropy diagnostic for tracking phase transitions.","marker":"Golechha (2024)"},{"why":"Offers the weight-norm/Godilocks-zone explanation of grokking that the paper contrasts with its entropy-based complexity account.","marker":"Liu et al. (2023)"},{"why":"Frames the generalization gap that the bilevel outer-loop objective is conceptually designed to minimize.","marker":"Johnson & Zhang (2023)"}],"fun_headline_variants":["Neural gradient filter speeds grokking up to 4.7×","Learn to transform gradients to grok faster","Grokking acceleration via learned gradient modulation","NeuralGrok: hit 95% accuracy in far fewer steps","Gradient transformation cuts grokking time on arithmetic"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The amplifier is trained on a 2% slice of the training set, and the method assumes that minimizing loss on that tiny, same-distribution slice after a single transformed gradient step is a reliable proxy for held-out test generalization.","fun_headline_variants_meta":{"raw":{"variants":["Neural gradient filter speeds grokking up to 4.7×","Learn to transform gradients to grok faster","Grokking acceleration via learned gradient modulation","NeuralGrok: hit 95% accuracy in far fewer steps","Gradient transformation cuts grokking time on arithmetic"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000593,"raw_usage":{"total_tokens":2819,"prompt_tokens":1025,"completion_tokens":1794,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":641,"completion_tokens_details":{"reasoning_tokens":1713}},"tokens_in":641,"tokens_out":1794,"duration_ms":12003,"temperature":1.0,"reasoning_tokens":1713,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-16T10:45:51.368711+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Train the same bilevel loop but replace $D_{\\text{outer}}$ with a corrupted or shuffled version (e.g., random labels or only even inputs). If NeuralGrok still reaches 95% test accuracy at the same speed, the outer-loop signal is not driving generalization-oriented gradient shaping; alternatively, if the speedup vanishes under a different 2% slice of the same training set, the result hinges on the specific slice rather than the method.","supporting_citations":[{"cited_title":"Inconsistency, Instability, and Generalization Gap of Deep Neural Network Training","cited_arxiv_id":"2306.00169","evidence_quote":"Frames the generalization gap that the bilevel outer-loop objective is conceptually designed to minimize."}],"review_version":1}