{"id":"deea42f4-44dc-43ea-9974-99c4bc1a40a3","arxiv_id":"2501.04697","paper_version":2,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":7.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":0,"one_line_summary":"Grokking fails without regularization because Softmax floating point errors (Softmax Collapse) stop gradients; removing them or removing the logit-scaling gradient direction restores and accelerates grokking.","lead":"Without weight decay, grokking tasks can push networks into floating point errors in the Softmax that freeze learning; the paper calls this Softmax Collapse and shows it can be fixed with a new activation. The authors also identify a gradient component that scales logits without changing predictions, and an optimizer that removes it, which makes networks generalize immediately instead of after a long overfitting phase.","discovery_kind":"unification","skeptic_critique":{"model":"deepseek-v4-flash","headline":"The direction removed by ⊥Grad is formally NLM only for homogeneous networks; for the biased MLPs/transformers actually trained, scaling all parameters does not exactly scale logits, so the causal role of NLM is not established for those models.","rationale":"The paper has genuine independent support: the SC mechanism is concrete, the artificial-SC intervention in App. B.1 is a clean causal check, precision changes shift behavior in the predicted direction, and StableMax is a simple intervention that reproducibly enables grokking without regularization. The reader's conditional verdict reflects these strengths. My stress-test focuses on the second leg of the central claim, namely that NLM explains the delay and causes SC. That leg rests on ⊥Grad removing the NLM component. The formal definition of NLM is only proved for positively homogeneous networks; for the biased models that appear in the main experiments, the claim that the removed radial component is exactly the NLM direction is an extrapolation. The last-layer-homogeneity comment in Sec. 4.2 is insufficient because ⊥Grad projects onto the full parameter vector, and scaling all parameters in a biased network does not yield exact logit scaling. Fig. 5 only shows gradient-weight alignment, not that the aligned direction satisfies the defining equation of Def. 5. This is a real soft spot, but it is testable and the existing evidence is suggestive. If the proposed test shows approximate logit scaling, the concern is resolved; if not, the causal story for biased architectures would need reworking. Since the reader already flagged this as the weakest assumption, my read does not move the verdict; it sharpens the exact experiment needed.","tokens_in":19410,"tokens_out":9380,"duration_ms":94957,"concrete_test":"Take a biased MLP checkpoint trained on addition mod 113 at the point of 100% train accuracy (as in Fig. 5b). Compute θ, the gradient g=∇L(θ), and the radial component g_rad=(θ·g/||θ||^2)θ. Simulate one small update θ'=θ−ηg_rad with ||θ'−θ|| ≈ 1e−3||θ||. For all training inputs, measure (i) the relative L2 error between f(θ';x) and c·f(θ;x) for the best scalar c, and (ii) whether the cross-entropy loss decreases. Repeat across several checkpoints and for the one-layer transformer with biases. If the radial update does not approximately scale logits (relative error above, say, 10%) while still reducing loss, then the direction removed by ⊥Grad is not an NLM direction and the interpretation of Figs. 6–7 as 'preventing NLM' is unsupported; if the logits scale within tolerance, the reader's concern is resolved.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The formal argument for NLM (Def. 5, Sec. 4.2) shows that d_NLM(θ)=αθ satisfies f(θ+d_NLM;x)=c f(θ;x) only under positive homogeneity (Def. 6). The paper extends this to biased MLPs and transformers by appealing to the last layer being homogeneous and to empirical gradient-weight alignment (Fig. 5). But the actual intervention ⊥Grad projects the gradient onto the hyperplane orthogonal to the full parameter vector θ (Def. 7, Eq. 12), not onto the last-layer parameters. For a biased network, the map θ→cθ is not a logit-scaling operation: with two hidden layers and an output bias, f(cθ;x)=c^2 W_out ReLU(W_2 ReLU(W_1 x+b_1)+b_2)+c b_out, which is not c^L f(θ;x) when b_out≠0. Thus the radial component removed by ⊥Grad need not be the NLM direction of Def. 5; removing it could accelerate generalization by preventing generic weight-norm growth or lazy training rather than by specifically preventing logit scaling. Fig. 5 measures cosine similarity between weights and gradients, which shows alignment but does not verify the defining property f(θ+d;x)=c f(θ;x). The paper's own Limitations paragraph acknowledges the missing quasi-homogeneous formal characterization, making this the weakest load-bearing step for the causal story on the architectures actually used.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper proposes that grokking failures without regularization are caused by floating-point absorption errors in the Softmax, termed Softmax Collapse (SC): once the correct-class logit is much larger than the others, the softmax denominator collapses to that logit, the cross-entropy loss becomes exactly zero, and gradients from correctly classified samples vanish, halting learning. The paper further introduces Naive Loss Minimization (NLM), a gradient component aligned with scaling the weights/logits that reduces CE loss after 100% training accuracy without changing predictions, and argues that NLM both delays generalization and eventually causes SC. Two interventions are proposed: StableMax, a numerically stable replacement for Softmax that enables grokking without regularization, and ⊥Grad, an optimizer that projects out the radial gradient component and removes the overfitting delay. The claims are supported by experiments on modular arithmetic, sparse parity, and MNIST, with additional validation on GPT-2 and ResNet. The paper also offers post-hoc explanations for why weight decay and MSE loss induce grokking.","tokens_in":19735,"tokens_out":8696,"duration_ms":83294,"significance":"If the central claims hold, the paper reframes a widely studied phenomenon as a numerical-stability problem rather than purely an implicit-bias or regularization effect, and it provides two simple, actionable interventions. The strengths are the explicit formal definitions (SC and NLM), the direct causal interventions (artificial SC, StableMax, ⊥Grad), the breadth of experimental settings, and the released code. The SC portion is well supported: the onset of SC coincides with stalled generalization, artificial SC stops generalization, and StableMax restores grokking. However, the NLM causal story has a load-bearing formal gap for the biased architectures actually used in the main experiments, as the paper itself acknowledges in its Limitations paragraph. The single-seed main figures also leave the quantitative robustness of the claims under-addressed.","major_comments":[{"comment":"The NLM direction is defined by Eq. (9), which requires f(θ+d_NLM(θ);x)=c f(θ;x). For positively homogeneous networks this is satisfied by d_NLM(θ)=αθ, as shown in Sec. 4.2. However, the MLPs and transformers actually trained include bias terms (Fig. 5b, 5c), and for biased networks f(cθ;x) is not generally equal to c f(θ;x); scaling all parameters scales hidden pre-activations by different powers of c and leaves bias terms with their own scaling, so the defining logit-scaling property is not satisfied. The paper acknowledges this in the Limitations paragraph but still uses ⊥Grad, which removes the projection of the gradient onto the full parameter vector θ (Eq. 12), and interprets the resulting speedup as specifically preventing NLM. Fig. 5 measures only cosine similarity between weights and gradients, which does not verify the logit-scaling property. A direct test is needed: for trained biased models, measure how f((1+α)θ;x) compares to c f(θ;x), or measure the change in logits when perturbing along θ; alternatively, restrict the NLM claims to homogeneous models and modify ⊥Grad to project onto the last-layer parameters only. Without this, the claim that ⊥Grad removes NLM on the architectures studied is not established.","section":"Sec. 4.2, Def. 5; Sec. 5.1, Def. 7; Fig. 5; Limitations"},{"comment":"The main empirical claims—SC onset coinciding with stalled generalization, StableMax inducing grokking, and ⊥Grad removing the delay—are each illustrated with single training runs and no error bars or multiple seeds. Grokking dynamics are known to be seed-sensitive, and the quantitative timing of SC onset and generalization can vary across runs. Reporting mean curves with standard deviations or at least a few seeds would make the load-bearing qualitative claims robust to stochasticity. Table 1 reports seeds only for the realistic settings, not for the central grokking experiments.","section":"Figs. 2, 4, 6, 7"},{"comment":"The description of the artificial SC intervention is internally inconsistent. The text says the goal is to 'set the gradients from the correct classes to zero', but the implementation is described as 'multiplying the logits for the right classes by 0'. Multiplying the correct logit by zero changes the softmax probability and produces a large negative gradient for the correct class, not a zero gradient. If the intervention actually zeroes the gradients, the sentence should say so; if it zeroes the logits, the experiment does not implement SC as defined in Def. 3 and its negative result would not support the SC explanation.","section":"App. B.1, Fig. 8"}],"minor_comments":[{"comment":"The caption sentence 'MLPs with (a) and without (b) bias terms' is reversed relative to the subcaptions, which correctly label (a) as without bias and (b) as with bias; the text in Sec. 4.2 refers to Fig. 5b for biased models, so the caption should be corrected.","section":"Fig. 5 caption"},{"comment":"The proof of Prop. 2 is more convoluted than necessary and the normalization step is not clearly justified. The result follows directly from Cauchy-Schwarz: ⟨−∇⊥L,∇L⟩ = −(‖∇L‖² − (θ·∇L)²/‖θ‖²) ≤ 0, with equality iff ∇L is parallel to θ. Consider replacing the current argument with this two-line derivation.","section":"App. A, proof of Prop. 2"},{"comment":"The right panel of Fig. 4 is labeled '2-hot input' and 'random binary input', while Sec. 4.1 describes a random binary vector of dimension 14. Please clarify whether the inputs are 2-hot in 14 dimensions or dense random binary vectors; the two descriptions are not equivalent.","section":"Fig. 4, right panel"},{"comment":"The definition of ⊥Grad assumes θ_t ≠ 0 in the projection formula (Eq. 12). Since training is initialized with small random weights and the projection is well defined after the first step, this is not a practical issue, but the zero-parameter case should be noted for completeness.","section":"Sec. 5.1, Def. 7"}],"recommendation":"major_revision","confidential_remarks":"The paper is honest about its limitations and the experimental evidence for the SC mechanism is strong. The main gap is the formal mismatch between the NLM definition and the biased architectures used in the headline experiments; this is fixable by adding a direct empirical test of logit scaling or by narrowing the claims. No concerns about citation practice or novelty disclosure."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"The paper's core observation is that unregularized grokking often stalls because floating-point absorption errors in the softmax (Softmax Collapse) zero out gradients, and that this is driven by a post-overfitting gradient alignment that scales logits without changing predictions (NLM). That is a new way to tie together the delayed onset, the role of weight decay, and why MSE behaves differently. The two interventions, StableMax and ⊥Grad, are clean and directly test the claimed causes; the artificial SC experiment in App. B.1 is a good control.\n\nCredit where it is earned: the paper builds honestly on Lyu & Li and Ji & Telgarsky for gradient-direction alignment, and the precision and stable-sum experiments in Fig. 2 support the SC mechanism. The Fourier sparsity result without weight decay is a nice bonus. The authors also state their limitations in plain language.\n\nSoft spots: the main grokking figures are single runs without error bars, which matters for a paper making temporal claims about when generalization does or does not happen. The transformer results with StableMax are promised in the text but not shown in the main figures. More importantly, the stress-test concern is real: the NLM direction is formalized only for positively homogeneous networks. For the biased MLPs and transformers actually trained, scaling all parameters does not exactly scale logits, and ⊥Grad projects out the full-parameter radial component, not a last-layer logit-scaling direction. The authors acknowledge this in the limitations paragraph, and Fig. 5 shows per-layer alignment, but that does not close the gap. So the causal role of NLM in biased models is plausible, not proven. The empirical speedup from ⊥Grad stands regardless, but its attribution specifically to NLM is softer.\n\nWho this is for: researchers working on grokking, implicit bias, or numerical stability in deep learning. The interventions are worth knowing even if the theory needs tightening.\n\nRecommendation: send to a serious referee. The empirical core is strong enough to merit referee time, but the authors should be pushed to add seeds and error bars, and either extend the quasi-homogeneous analysis or temper the NLM attribution for biased networks.","headline":"A genuinely new mechanistic account of grokking, with clean interventions and an honest limitation that keeps the causal loop from fully closing for biased networks.","tokens_in":20247,"tokens_out":3894,"would_cite":true,"duration_ms":38196,"reading_group":"yes","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"Grokking without regularization fails because floating-point errors in Softmax zero out gradients, and two targeted fixes restore it.","keywords":["grokking","softmax collapse","floating-point absorption errors","naive loss minimization","logit scaling","StableMax","⊥Grad","cross-entropy loss"],"falsifier":"Measure, after 100% training accuracy on modular addition, whether the radial gradient component that ⊥Grad removes ever changes predicted labels on the test set; if it does, that component is not purely logit scaling and the NLM account of the delay needs revision. A second check: train the same biased MLP with ⊥Grad on a loss, such as mean-squared error with bounded targets, where scaling logits does not reduce the loss; if it still generalizes faster than vanilla SGD, the speedup is not explained by blocking NLM.","tokens_in":19244,"feed_emoji":"🧮","tokens_out":12590,"duration_ms":110705,"temperature":0.7,"pith_summary":"Grokking—the sudden jump from memorization to generalization after prolonged overfitting—typically needs weight decay or another regularizer; this paper claims the real blocker in cross-entropy training is numerical rather than statistical. After a model reaches 100% training accuracy on a grokking task, the gradient aligns with a direction that merely scales up the logits, lowering the loss without changing any prediction. That scaling drives logits so far apart that floating-point addition inside Softmax absorbs the smaller exponentials, making the loss and the gradients for correctly classified samples exactly zero; the paper calls this Softmax Collapse and shows that it ends learning, sometimes before the test accuracy has moved. The paper validates the claim twice: a numerically stable replacement for Softmax (StableMax) lets models grok with no regularization, and an optimizer that removes the logit-scaling gradient component (⊥Grad) makes generalization happen without the long overfitting delay. If this account is right, the central mystery of grokking—why generalization is delayed and why regularization seems necessary—is largely a story about numerical stability.","feed_headline":"Grokking is blocked by Softmax round-off, not missing regularization","feed_subtitle":"A Softmax fix and a gradient projection make delayed generalization disappear without weight decay.","key_machinery":"The central object is the decomposition of the gradient after overfitting into two components: the NLM component, which for a positively homogeneous network (a network whose outputs scale by a constant when all weights are scaled by a constant) points along the current weight vector $\\theta$ and only rescales the logits (so $f(\\theta+d_{\\mathrm{NLM}}(\\theta);x)=c f(\\theta;x)$ for some $c>1$), and the orthogonal component that actually changes predictions. The second piece of machinery is the absorption-error condition that defines SC: when the exponent gap between the true-class exponential and every other term in the Softmax sum exceeds the significand precision, the sum collapses to the true-class term, the cross-entropy loss is exactly zero, and the gradients from correctly classified samples vanish. StableMax replaces the exponential in Softmax with the piecewise function $s(x)=x+1$ for $x\\ge 0$ and $s(x)=1/(1-x)$ for $x<0$, which grows linearly rather than exponentially and avoids the extreme summands that cause absorption. $\\perp$Grad projects the gradient onto the hyperplane orthogonal to the weight vector, $\\nabla_\\perp L(\\theta_t)=\\nabla L(\\theta_t)-\\frac{\\theta_t^\\top\\nabla L(\\theta_t)}{\\theta_t^\\top\\theta_t}\\theta_t$, removing the NLM component so only the prediction-changing part drives updates. Together these two mechanisms carry the argument: SC explains why training stops, NLM explains the delay before it stops, and each intervention isolates one link of the chain.","core_discovery":"The paper's central claim is that grokking fails without regularization because training runs into Softmax Collapse (SC), a floating-point absorption error in the Softmax sum. For a correctly classified sample, once the true-class logit is so large that $\\sum_k e^{z_k} \\doteq e^{z_y}$ under floating-point arithmetic, the cross-entropy loss evaluates to exactly zero and the gradient contribution from that sample vanishes, so learning stops even though the model has not generalized. SC is the last step of a longer chain: after reaching 100% training accuracy, cross-entropy gradients align with the naive loss minimization (NLM) direction, $d_{\\mathrm{NLM}}(\\theta)=\\alpha\\theta$ for positively homogeneous networks, which scales all logits by a constant and lowers the loss without altering predictions; the same alignment is observed empirically in the biased MLPs and transformers actually trained. Two interventions support the chain: StableMax, a modified Softmax whose unbounded 'exponential' grows only linearly, prevents SC and produces grokking without regularization, while ⊥Grad, which updates only the gradient component orthogonal to the weight vector, removes the delay in generalization. In this account, weight decay works, MSE loss works on shallow networks, and label smoothing behaves differently for a common reason: whether the loss can be reduced indefinitely by scaling logits, and whether floating-point collapse can be triggered.","pith_inferences":["The paper's diagnosis would predict that grokking curves shift with floating-point precision in a quantifiable way: identical runs in float16, float32, and float64 should stall at precisely the epochs where the SC fraction crosses a threshold; running the same sweep in bfloat16 would be a direct check.","Because NLM is a property of losses that are monotonically decreased by logit scaling, any training scheme that bounds or penalizes logit norms should eliminate the delay; the paper's appendix logit-regularization experiment supports this, and the same logic suggests fixed-norm or normalized-gradient variants would grok without either StableMax or ⊥Grad.","The formal gap for biased networks matters: if a quasi-homogeneous analysis shows that the radial gradient component is not purely logit-scaling in the presence of biases, the causal story would need to be revised, so a theoretical characterization of quasi-homogeneous models is the natural next step.","In mixed-precision training beyond grokking benchmarks, SC may be an unexamined regularizer that silently zeroes gradients for well-classified examples; replacing Softmax with a sub-exponential variant could change training dynamics in settings where losses are driven near zero."],"forward_implications":["Weight decay induces grokking by counteracting NLM: it pulls weights back along the same radial direction, so logit scaling stops once the loss reduction no longer outweighs the regularization penalty.","MSE loss groks on shallow networks without regularization because scaling logits cannot reduce an MSE loss indefinitely, so the NLM mechanism is absent.","Temperature scaling and float64 only delay SC rather than prevent it; only sub-exponential replacements like StableMax avoid the collapse entirely.","⊥Grad reaches 100% test accuracy with no preceding overfitting phase and, in the tested settings, outperforms the best-tuned weight decay with no extra hyperparameters.","Slingshot spikes may be a mechanism that prevents full SC, explaining why adaptive optimizers can occasionally grok without weight decay."],"supporting_citations":[{"why":"Introduced the grokking phenomenon and the modular-arithmetic benchmark, and provided the observation that grokking can occur without regularization on larger datasets.","marker":"Power et al. (2022)"},{"why":"Established the standard transformer setup and weight-decay-induced progress measures for modular addition, and documented slingshot spikes whose disappearance under float64 the paper connects to SC.","marker":"Nanda et al. (2023)"},{"why":"Supplied the sparse parity task and the finding that SGD learns parities with hidden progress, a key grokking dataset used for validation.","marker":"Barak et al. (2022)"},{"why":"Proved that homogeneous networks converge in direction and maximize margin after interpolation, which is the theoretical basis for identifying scaling along the weights as an NLM direction.","marker":"Lyu & Li (2020)"},{"why":"Provided the MLP architecture and the 'Goldilocks zone' view of weight norms, and the setting of scaling logits to delay or induce grokking.","marker":"Liu et al. (2023a)"},{"why":"Showed shallow networks with MSE loss grok without regularization, which the paper explains as MSE preventing NLM.","marker":"Kumar et al. (2024)"},{"why":"Observed slingshot dynamics with adaptive optimizers that can lead to grokking without weight decay, whose mechanism the paper hypothesizes is avoiding full SC.","marker":"Thilak et al. (2022)"},{"why":"Argued weight decay prevents floating-point issues and increases effective learning rate in mixed-precision training, supporting the paper's account of why weight decay matters.","marker":"D'Angelo et al. (2023)"}],"fun_headline_variants":["Grokking's real blocker: Softmax round-off, not missing decay","Fix Softmax collapse, and grokking appears without weight decay","Why grokking stops: floating-point Softmax kills gradients","StableMax: grokking without regularization by stopping collapse","Softmax collapse, not regularization, is what gates grokking"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The load-bearing premise is that, for the biased MLPs and transformers actually trained, the gradient component pointing along the current weights is exactly the part that only scales the logits without changing predictions; the formal proof covers only bias-free networks in which scaling all weights scales all outputs, and for the biased models the paper relies on empirical alignment rather than a theorem.","fun_headline_variants_meta":{"raw":{"variants":["Grokking's real blocker: Softmax round-off, not missing decay","Fix Softmax collapse, and grokking appears without weight decay","Why grokking stops: floating-point Softmax kills gradients","StableMax: grokking without regularization by stopping collapse","Softmax collapse, not regularization, is what gates grokking"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.00019,"raw_usage":{"total_tokens":1430,"prompt_tokens":1130,"completion_tokens":300,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":746,"completion_tokens_details":{"reasoning_tokens":209}},"tokens_in":746,"tokens_out":300,"duration_ms":3293,"temperature":1.0,"reasoning_tokens":209,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-10T21:25:57.348665+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Measure, after 100% training accuracy on modular addition, whether the radial gradient component that ⊥Grad removes ever changes predicted labels on the test set; if it does, that component is not purely logit scaling and the NLM account of the delay needs revision. A second check: train the same biased MLP with ⊥Grad on a loss, such as mean-squared error with bounded targets, where scaling logits does not reduce the loss; if it still generalizes faster than vanilla SGD, the speedup is not explained by blocking NLM.","supporting_citations":[],"review_version":1}