{"id":"169c5206-731e-4f47-a7f8-c8090a1bdf39","arxiv_id":"2502.02407","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":1,"one_line_summary":"SAM's failure in language modeling is traced to a dominant 'logit path' that minimizes sharpness spuriously, and the proposed Functional-SAM, which removes that path, improves validation loss over AdamW and SAM.","lead":"The authors show that Sharpness-Aware Minimization fails in language modeling because it mostly reduces sharpness by making the model overconfident, not by flattening the learned function. They propose Functional-SAM, which perturbs only the network function, and it beats AdamW and SAM on validation loss across model scales.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"The spurious-logit mechanism is not tested: preconditioned SAM, which keeps the logit path, already recovers most of the gain, and τ is never measured on the proposed algorithms' own trajectories.","rationale":"The reader's verdict is CONDITIONAL and I agree with that. My stress-test focuses on the single most load-bearing step: the causal attribution. The paper is transparent that the functional-path explanation is a hypothesis, and the empirical improvement of Functional-SAM is real as a reported result. However, the internal comparison between preconditioned SAM and preconditioned Functional-SAM in Tables 2–3 weakens the attribution: preconditioning alone, without any removal of the logit path, converts SAM from harmful to helpful and captures roughly half or more of the total gain over AdamW. This is an internal-data tension, not just a disagreement with an outside consensus. The missing measurement of τ along the proposed algorithms' trajectories is the concrete gap that separates a supported mechanism from an untested hypothesis. If the proposed check shows that preconditioned SAM has τ_logit≈1 and still improves over AdamW, the paper's central explanatory claim would need to be revised to a geometry-mismatch story, even though the algorithmic contribution might remain. This does not change the verdict from CONDITIONAL: the algorithmic claims are plausible and the authors flag the hypothesis status, but the headline causal claim should not be accepted at face value until the mechanism is directly measured.","tokens_in":18963,"tokens_out":9910,"duration_ms":95279,"concrete_test":"Train the 23.9M Nanodo model on C4 for 10K steps with preconditioned SAM, preconditioned Functional-SAM, and preconditioned Angle-SAM at φ=π/2 (logit-only), using matched ρ values. Every 500 steps, compute τ_logit, τ_func, and τ_cross (Eq. 9) along each method's own trajectory. If preconditioned SAM improves over AdamW while τ_logit stays near 1, the spurious-logit diagnosis is not the driver of the gains. Additionally, if the φ=π/2 logit-only variant (which suppresses the functional path and keeps the logit path) does not worsen validation loss relative to AdamW, then the logit path is not harmful once the perturbation geometry is matched; this would directly contradict the paper's causal story.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The central claim is that SAM fails in language modeling because the logit path dominates the sharpness gradient (τ_logit≈1), and that Functional-SAM succeeds by suppressing that spurious path. This causal attribution is not actually supported by the paper's own tables. In Table 2, preconditioned SAM—which uses the unmodified SAM gradient, including the δ_logit term, with only the perturbation geometry matched to Adam—already improves over AdamW (3.55 vs 3.57 at 23.9M; 3.44 vs 3.45 and 3.27 vs 3.28 at larger sizes), and Functional-SAM adds only 0.01–0.03. Table 3 shows the same pattern. So the data are at least as consistent with the hypothesis that the key defect of SAM in this setting is the mismatch between the SAM perturbation and the Adam geometry, not the dominance of the logit path per se. The manuscript labels the functional-path explanation a hypothesis, but it never tests it on the trajectories of the proposed methods: Figures 2 and 4 measure τ along AdamW training only, and Table 4 reports final H_L/H_G traces without isolating the functional Hessian contribution. Without τ measurements during Functional-SAM/preconditioned-SAM training, the mechanism that the title and abstract assert ('avoiding spurious sharpness minimization') remains unverified.","agreement_with_reader":"partial"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper studies why Sharpness Aware Minimization (SAM) underperforms in language modeling. It decomposes the SAM sharpness gradient into a logit path (through the loss derivative with respect to the logits) and a functional path (through the function Jacobian), defines normalized contributions tau_logit, tau_func, and tau_cross, and measures them on Transformer language models and vision Transformers. The measurements show that tau_logit is close to 1 in language modeling, whereas the two paths are more balanced in vision. Motivated by the hypothesis that logit-path sharpness reduction is spurious, the authors propose Functional-SAM, which perturbs only the Jacobian, and preconditioned SAM, which applies the AdamW second-moment preconditioner to the SAM perturbation. Experiments on C4 with Nanodo models from 2M to 1.2B parameters report consistent validation-loss improvements over AdamW and SAM in fixed-length and Chinchilla-style training regimes, together with Hessian-trace and one-shot pruning analyses.","tokens_in":19203,"tokens_out":7536,"duration_ms":78570,"significance":"If the causal story holds, the paper offers a useful explanation of why naive SAM fails in NLP and a practical modification that makes sharpness regularization viable for language modeling at scale. The mathematical decomposition in Section 3 is clean, the provided JAX snippets match the described update rules, and the empirical comparisons include a wide range of model sizes, including a 1.2B-parameter run. The equal-step gains over AdamW are modest but consistent across scales and training regimes, and the paper goes beyond validation loss with Hessian and pruning evidence. The main weakness is that the central causal attribution is not directly tested, which is why the paper needs revision rather than acceptance in its current form.","major_comments":[{"comment":"The causal claim that SAM fails because the logit path dominates sharpness reduction is not tested on the trajectories of the compared methods. The tau measurements are taken while training with AdamW, not while training with SAM, preconditioned SAM, or Functional-SAM. At the same time, preconditioned SAM, which does not remove the logit-path term and only changes the perturbation geometry, already captures most of the gain over AdamW (e.g., 3.55 versus 3.57 at 23.9M and 3.27 versus 3.28 at 117.9M in Table 2). The authors should either measure tau_logit and tau_func along the trajectories of the proposed methods, or provide a controlled ablation that isolates the logit path while holding the perturbation geometry fixed, for example the Angle-SAM variant with phi=pi/2 discussed in Appendix D. Without one of these, the improved results are also consistent with an Adam-geometry mismatch and do not establish the title's claim of avoiding spurious sharpness minimization.","section":"Section 3.3, Figures 2 and 4; Section 5.3, Tables 2 and 3"},{"comment":"The proposed algorithms are meant to reduce functional sharpness, but the reported Hessian metrics do not isolate the functional component. Table 4 reports lambda_max(H_L), tr(H_L), and tr(H_G), but not tr(H_F) or lambda_max(H_F). Moreover, the reported lambda_max(H_L) is larger for preconditioned Functional-SAM (6.20) than for SAM (2.71), so the maximum eigenvalue of the full Hessian does not support the claim that the proposed method finds a flatter solution under the metric that the paper itself emphasizes. The authors should report the functional-Hessian counterparts and state explicitly which flatness measure Functional-SAM is designed to reduce.","section":"Section 5.5, Table 4"},{"comment":"The main empirical results are presented as single numbers without error bars or repeated seeds. The improvements over AdamW are 0.01 to 0.04 in validation loss, which is within the range where run-to-run variability matters. The statement that standard deviations are at the fourth decimal is made for the 2M prototype in Section 5.2 and cannot automatically be assumed for the larger models in Tables 2 and 3. The paper should report multiple seeds, error bars, or an equivalent statistical analysis for the final tables, or explicitly state the computational constraints that prevent this and hedge the claims accordingly.","section":"Section 5.2 and Section 5.3, Tables 1-3"},{"comment":"The random-matrix argument for preconditioning assumes free independence of A and B and an isotropic random vector v, whereas in the actual network setting the relevant vector is the training gradient and the matrices H_G and H_F are related to that gradient. The appendix already says the conclusion is a suspicion, but the main text in Section 4.2 presents preconditioned SAM as an indirect way to promote the functional path. Either soften this motivation or add a direct empirical check, such as measuring the norms of H_G M^{-1} g and H_F M^{-1} g during training, to show that the proposed preconditioner indeed upweights the functional component relative to the logit component.","section":"Appendix B.1, Eqs. (13)-(17)"}],"minor_comments":[{"comment":"The code snippet contains a typo: 'hax_aux' should be 'has_aux'; the text in Listing 1 also has 'dF_dtheta_fn' defined with a function that takes only theta, which is correct but should be aligned with the surrounding notation.","section":"Appendix C, Listing 1"},{"comment":"The PDF text contains visible spacing artifacts such as 'ofSA M' and 'SA M' that should be fixed in the final version.","section":"Abstract and Section 1"},{"comment":"The definitions of tau_logit, tau_func, and tau_cross have a denominator of ||nabla_theta SP||^2, which is undefined if the sharpness gradient is exactly zero; a short remark on this degenerate case would be helpful.","section":"Section 3.2, Eq. (9)"},{"comment":"At the 1.2B scale the tables show NaN for SAM and preconditioned SAM; the warm-up mitigation numbers (3.07 and 3.08) are mentioned only in the text. These values should be reported in the main tables, or the tables should clearly state that the proposed method comparison at this scale is only preconditioned Functional-SAM.","section":"Section 5.3, Tables 2 and 3"},{"comment":"The relationship to FisherSAM and to Zhang et al. (2025) is described only briefly; a sentence stating the precise difference between the proposed preconditioning and those alternatives would help the reader judge novelty.","section":"Section 4.2"}],"recommendation":"major_revision","confidential_remarks":"The paper is within scope and the empirical contribution is potentially useful. My main concern is the gap between the diagnosis and the intervention: the central claim is causal, but the manuscript currently measures tau only on AdamW trajectories and does not directly test whether removing the logit path, rather than changing the perturbation geometry, is what drives the gains. A targeted experiment measuring tau on the proposed methods' own trajectories, or a Logit-SAM ablation, would resolve this. I do not see concerns about citation practice or novelty disclosure, though the relation to concurrent preconditioning work should be made precise."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"The genuinely new thing here is the logit/functional decomposition of the SAM sharpness gradient (Eqs. 6 and 10) and the Functional-SAM update that follows. The derivation is clean, the JAX snippet makes the implementation obvious, and the empirical observation that SAM underperforms AdamW in language modeling while costing 2x compute is reproduced across scales. The paper also deserves credit for labeling the functional-path explanation a hypothesis in Section 3.3 and for acknowledging the 2x compute cost in the Discussion.\n\nThe soft spots are real, and they sit right on the central claim. The tables do not support the causal mechanism as cleanly as the abstract suggests. Preconditioned SAM, which keeps the logit term in the gradient untouched and only changes the perturbation geometry, already captures most of the gain over AdamW—in Table 2, 3.57 to 3.55 at 23.9M, with Functional-SAM adding only 0.02. The Chinchilla table shows the same pattern. Meanwhile, tau is measured only along AdamW trajectories (Figures 2 and 4), never on preconditioned SAM or Functional-SAM. So we have no direct evidence that suppressing the logit path is what drives the improvement. The random-matrix argument in Appendix B.1 is a heuristic, not evidence for these actual preconditioners on these models.\n\nThe gains themselves are modest, on the order of 0.02 to 0.04 validation loss, and the multi-scale tables carry no error bars; the only reported standard deviation is for the 2M prototype. At billion scale the comparison rests on post hoc NaN handling for SAM and preconditioned SAM. And the experiments are equal-steps, not equal-compute—there is no 2x-compute AdamW baseline, which matters when the selling point is that sharpness regularization can work for LLMs.\n\nNone of this is fatal. The paper is transparent about most of it, and the decomposition plus the algorithmic variants are worth having in the literature. I would cite it for the logit/functional split. But the title and abstract overstate the conclusion. A serious referee should ask for tau measurements on the proposed methods' own trajectories, error bars in the main tables, an equal-compute AdamW baseline, and a more careful framing of the mechanism as hypothesis rather than demonstrated fact. That is a revision, not a rejection.\n\nYes, send this to peer review. The diagnostic framework and the algorithm are solid contributions even if the headline mechanism is not yet proven.","headline":"Genuinely useful decomposition of SAM's sharpness gradient, but the causal story about logit-path dominance is not actually tested in the paper—preconditioned SAM recovers most of the gain and tau is never measured on the proposed algorithms.","tokens_in":19749,"tokens_out":2703,"would_cite":true,"duration_ms":28501,"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":"The paper claims that SAM's failure in language modeling comes from its descent being dominated by logit sharpness, and that Functional-SAM, which perturbs only the function Jacobian, improves validation loss over AdamW and SAM across…","keywords":["Sharpness-Aware Minimization","language modeling","generalization","Hessian decomposition","Gauss-Newton decomposition","functional Hessian","preconditioning","large language models"],"falsifier":"Train a language model with a deliberately logit-path variant (Angle-SAM with $\\phi=\\pi/2$, i.e., Logit-SAM) at the same compute; if that variant matches or beats AdamW's validation loss, the claim that logit sharpness reduction is spurious is refuted. A complementary check is to measure $\\tau_{\\mathrm{logit}}$ and $\\tau_{\\mathrm{func}}$ during actual SAM and Functional-SAM training rather than only along AdamW trajectories; if the functional path does not become dominant during Functional-SAM training, the proposed mechanism is not what drives its gains.","tokens_in":18758,"feed_emoji":"📉","tokens_out":13690,"duration_ms":107734,"temperature":0.7,"pith_summary":"Sharpness-Aware Minimization (SAM), a curvature regularizer that helps vision models, consistently hurts language modeling. This paper argues the reason is that in language models SAM mostly reduces sharpness through the logits—making the network overconfident—rather than through the function the network computes, and calls that reduction spurious. To test the diagnosis, the authors derive a decomposition of SAM's sharpness gradient into a logit path and a functional path, measure that the logit path dominates in language while the two are balanced in vision, and then remove the logit path in a new update rule, Functional-SAM. Combined with a preconditioned perturbation that uses AdamW's second-moment statistics, the method improves validation loss over both AdamW and SAM in fixed-length and compute-optimal training at scales up to about 1.2 billion parameters. If correct, this makes sharpness regularization usable for large language models, not just vision.","feed_headline":"SAM's logit shortcut, fixed: Functional-SAM beats AdamW","feed_subtitle":"Decomposing sharpness into logit and functional paths turns a failing regularizer into a working one for LLMs.","key_machinery":"The central object is the decomposition of the sharpness gradient into two additive routes, $\\delta_{\\mathrm{logit}}$ and $\\delta_{\\mathrm{func}}$, obtained by applying the product rule to $\\nabla_\\theta\\|\\nabla_\\theta\\mathcal{L}\\|$ and identifying the two terms with the GGN and functional Hessian components of the loss Hessian. The paper measures their normalized squared norms ($\\tau_{\\mathrm{logit}}$ and $\\tau_{\\mathrm{func}}$) to diagnose which route dominates, and builds Functional-SAM by keeping only the $\\delta_{\\mathrm{func}}$ route in the finite-difference perturbation update. Preconditioned SAM modifies the perturbation to $\\rho M^{-1}\\epsilon^*$, using Adam's second-moment estimate to reduce the logit route.","core_discovery":"Starting from the penalty form of SAM, $\\min_\\theta \\mathcal{L}(\\theta)+\\rho\\|\\nabla_\\theta\\mathcal{L}(\\theta)\\|$, the paper uses the product rule and the Gauss-Newton decomposition of the Hessian, $\\nabla_\\theta \\mathrm{SP}=\\delta_{\\mathrm{logit}}+\\delta_{\\mathrm{func}}=H_G\\,\\epsilon^*+H_F\\,\\epsilon^*=H_L\\,\\epsilon^*$, to split the sharpness-reduction signal into a logit path (through the generalized Gauss-Newton (GGN) term $H_G$, reducible by making softmax outputs more one-hot) and a functional path (through the functional Hessian $H_F$, tied to the geometry of the network function). Measuring the normalized contributions $\\tau_{\\mathrm{logit}}$, $\\tau_{\\mathrm{func}}$, and $\\tau_{\\mathrm{cross}}$ on Transformers, the paper finds that in vision the two paths are comparable, while in language modeling $\\tau_{\\mathrm{logit}}\\approx 1$ throughout training, so SAM's updates mostly shrink logit sharpness rather than flattening the function. The proposed remedy, Functional-SAM, computes $g_{\\mathrm{Func\\text{-}SAM}}=-\\nabla_\\theta F(\\theta+\\rho\\epsilon^*)\\cdot \\nabla_F \\mathcal{L}(\\theta)$, perturbing only the Jacobian of the function and discarding the logit contribution, and the authors argue and show that preconditioning the perturbation with the AdamW preconditioner $M^{-1}$ further favors the functional path. The paper's central claim is that this distinction between logit and functional sharpness explains SAM's failure in NLP and that emphasizing the functional path makes sharpness minimization work for language modeling.","pith_inferences":["Inference: The causal direction may be weaker than the experiments suggest: the paper measures $\\tau$ values along AdamW trajectories and hypothesizes that functional sharpness is the beneficial object, so the gains of Functional-SAM support the hypothesis without independently proving the mechanism.","Inference: The logit-path diagnosis may also explain why SAM helps in fine-tuning settings where parameters move less and logits are constrained, and would predict that Logit-SAM should hurt pre-training while possibly helping fine-tuning; this is testable.","Inference: The same decomposition could be applied to other sharpness variants and to other output-rich tasks such as speech or multi-label classification to decide whether those settings are logit-dominated before running expensive SAM experiments.","Inference: Because preconditioned Functional-SAM also avoids the numerical instabilities that plain SAM shows at billion-parameter scale, the method may be of practical value even if the sharpness-decomposition narrative were later revised."],"forward_implications":["In language modeling, SAM's poor validation loss is attributable to the logit path: at equal compute the baseline SAM underperforms AdamW, while Functional-SAM outperforms both in the same fixed step budget.","Preconditioning the SAM perturbation with the optimizer's own second-moment statistics is enough to shift the balance toward the functional path and improve validation loss, so the inner perturbation and outer optimizer geometries should match.","Functional-SAM and its preconditioned version reach flatter solutions (lower maximum Hessian eigenvalue and trace) than AdamW, and those lower-curvature solutions are the ones with better validation loss.","The combination works across model sizes from 23.9M to 1.2B parameters and in both fixed-length and Chinchilla-style training, so curvature regularization can scale to large language models.","The decomposition into logit and functional paths suggests an interpolation (Angle-SAM) between Functional-SAM and Logit-SAM, which could be tuned per domain instead of always using the SAM default."],"supporting_citations":[{"why":"Defines SAM and the min-max sharpness objective that the paper diagnoses and modifies.","marker":"Foret et al., 2020"},{"why":"Introduces the penalty-SAM form and the neglected functional Hessian component that motivates the functional path.","marker":"Dauphin et al., 2024"},{"why":"Provides the Gauss-Newton decomposition of the Hessian used to split sharpness into logit and functional routes.","marker":"Schraudolph, 2002"},{"why":"Defines the functional Hessian component of the loss Hessian used in the decomposition.","marker":"Singh et al., 2021"},{"why":"Supplies the decoder-only Transformer implementation and model family used in the language modeling experiments.","marker":"Liu et al., 2024"},{"why":"Supplies the C4 dataset used for all language modeling training.","marker":"Raffel et al., 2020"},{"why":"Defines the Chinchilla compute-optimal training regime used in the scaling experiments.","marker":"Hoffmann et al., 2022"},{"why":"Defines AdamW, the baseline optimizer and the source of the second-moment preconditioner used in preconditioned SAM.","marker":"Loshchilov and Hutter, 2019"},{"why":"Supplies the observation that gradients align with principal Hessian eigenspaces, motivating inverse preconditioning to downweight the logit route.","marker":"Gur-Ari et al., 2018"}],"fun_headline_variants":["Functional-SAM sidesteps logit shortcut, beats AdamW for LLMs","SAM's NLP failure traced to logit sharpness; Functional-SAM fixes it","Decomposing sharpness: Functional-SAM beats AdamW in language models","Not all curvature is useful: Functional-SAM broadens SAM to LLMs","Avoiding spurious sharpness: Functional-SAM broadens SAM to LLMs"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The load-bearing premise is that functional sharpness, not logit sharpness, is what drives generalization in language modeling—and that the sharpness split measured along AdamW trajectories is representative of SAM's own updates; if either gives way, the diagnosis of spurious minimization collapses.","fun_headline_variants_meta":{"raw":{"variants":["Functional-SAM sidesteps logit shortcut, beats AdamW for LLMs","SAM's NLP failure traced to logit sharpness; Functional-SAM fixes it","Decomposing sharpness: Functional-SAM beats AdamW in language models","Not all curvature is useful: Functional-SAM broadens SAM to LLMs","Avoiding spurious sharpness: Functional-SAM broadens SAM to LLMs"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000253,"raw_usage":{"total_tokens":1649,"prompt_tokens":1111,"completion_tokens":538,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":727,"completion_tokens_details":{"reasoning_tokens":434}},"tokens_in":727,"tokens_out":538,"duration_ms":5726,"temperature":1.0,"reasoning_tokens":434,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-09T12:15:15.468291+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Train a language model with a deliberately logit-path variant (Angle-SAM with $\\phi=\\pi/2$, i.e., Logit-SAM) at the same compute; if that variant matches or beats AdamW's validation loss, the claim that logit sharpness reduction is spurious is refuted. A complementary check is to measure $\\tau_{\\mathrm{logit}}$ and $\\tau_{\\mathrm{func}}$ during actual SAM and Functional-SAM training rather than only along AdamW trajectories; if the functional path does not become dominant during Functional-SAM training, the proposed mechanism is not what drives its gains.","supporting_citations":[],"review_version":1}