{"id":"242c69f6-9970-4d82-aa73-56561a7665cb","arxiv_id":"2506.02285","paper_version":2,"verdict":"CONDITIONAL","confidence":"HIGH","novelty_score":4.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":0,"one_line_summary":"Gradient norm blow-up at the end of LLM training is caused by an interaction between weight decay, normalization layers, and the learning rate schedule, and a corrected weight decay term removes it.","lead":"During long LLM training runs, the gradient norm spikes near the end because weight decay and the shrinking learning rate push the gradient-to-weight ratio upward. The paper proposes a simple corrected weight decay term that removes the spike and slightly lowers loss, offering a practical fix for large-scale training.","discovery_kind":"extension","skeptic_critique":{"model":"deepseek-v4-flash","headline":"Transformer residual structure violates the ⟨g, x⟩ = 0 premise behind Eq. (2), so the LLM tail-blowup mechanism is asserted rather than derived; a direct measurement of the cross term is needed to decide whether the paper's main claim holds.","rationale":"The reader's weakest assumption identifies the same load-bearing concern I would raise: the orthogonality condition ⟨g, x⟩ = 0 is exact only for layers directly followed by a normalization operation, and Transformer layers do not have that structure. The paper itself flags this in Section 6, making it an explicit limitation rather than an overlooked one. I agree with the CONDITIONAL verdict because the ImageNet experiments provide credible support for the equilibrium mechanism in directly normalized layers, while the LLM case, which motivates the paper, remains an extrapolation. A direct measurement of the normalized cross term is the decisive check: if it is large, the derivation leading to Eq. (2) does not apply to the layers the correction modifies, and the observed improvement could be a regularization-schedule effect rather than the proposed mechanism. I do not base the verdict on the 'lower loss' claim because that is secondary to the main mechanism; even if the loss comparison is confounded by total weight decay, the gradient-norm explanation could still be correct for the cases where the theory's assumptions hold.","tokens_in":7563,"tokens_out":8226,"duration_ms":84396,"concrete_test":"Instrument a 120M Llama training run under AdamW with the same setup as Figure 4, and log the smoothed per-layer normalized cross term C_t = ⟨g_t, x_t⟩ / (‖g_t‖ · ‖x_t‖) for every linear layer. If the median |C_t| in the second half of training is materially above zero (e.g., > 0.1), then Eq. (1)'s premise is violated for Transformers and the paper's LLM mechanism is not derived; if |C_t| remains close to zero, the residual-structure concern is resolved.","verdict_should_be":"UNCHANGED","load_bearing_attack":"Section 3's steady-state ratio (Eq. 2) is derived from Eq. (1), which drops the cross term ⟨g_t, x_t⟩ under the assumption that the layer is immediately followed by a normalization operation. In the Llama-style Transformers that motivate the paper (Figure 1), no linear layer has this exact structure: QKV projections are preceded by RMSNorm, and attention/MLP output projections feed a residual stream before any normalization. The loss is therefore not scale-invariant in any individual layer's weights, so ⟨g_t, x_t⟩ is not forced to zero. Section 6 explicitly concedes that the Transformer structure 'no longer corresponds to the direct linear-followed-by-norm structure that our theory covers,' yet the correction is applied to every linear layer except the output. If the cross term is material, Eq. (2) does not govern the LLM tail dynamics, and the empirical success of AdamC could instead be explained by the correction simply reducing late-training weight decay, rather than by the proposed equilibrium mechanism. This is the load-bearing gap for the paper's headline LLM claim; the ImageNet ResNet results, where conv layers are directly BatchNormed, support the theory but do not validate the Transformer extrapolation.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"This short note addresses the rapid increase of gradient norms near the end of long LLM training runs. The paper models a normalized linear layer under weight decay and learning-rate scheduling, deriving the steady-state ratio ||g_t||/||x_t|| = sqrt(2λ/γ_t) (Section 3, Eq. 2). As γ_t decays, this target ratio grows, which the paper identifies as the cause of the tail blow-up. The proposed fix is a corrected weight decay λ̂_t = λ γ_t/γ_max applied only to normalized layers, yielding the AdamC and SGDC optimizers. Experiments include a 200B-token 120M-parameter Llama-3 run on FineWeb-Edu and ImageNet ResNet-50 training; the paper reports that the correction eliminates the late-training gradient increase and gives lower loss.","tokens_in":7745,"tokens_out":9707,"duration_ms":80555,"significance":"If valid, the paper offers a simple, parameter-free explanation and a practical optimizer correction for a phenomenon observed in real LLM training. The theoretical steady-state prediction in Eq. (2) is falsifiable and is convincingly validated on ImageNet ResNet-50 with a flat schedule (Fig. 2) and a cosine schedule (Fig. 3), and the five-seed ImageNet accuracy comparison (Fig. 5) is a genuine strength. The main weaknesses are the extrapolation of the theory to Transformers—where the orthogonality premise does not hold—and the thin single-run LLM evidence for the lower-loss claim. The paper is clearly written and the algorithm is straightforward to implement, but these gaps prevent the headline LLM claim from being fully established.","major_comments":[{"comment":"The steady-state ratio (Eq. 2) is derived from Eq. (1), whose second equality drops the cross term ⟨g_t, x_t⟩ under the premise that the layer is immediately followed by a normalization operation. In the Llama-style architecture used in Figure 4, no linear layer has this structure: the QKV projections are preceded by RMSNorm, and the attention/MLP output projections feed the residual stream before any normalization. The text itself concedes this ('The application of normalization operations in Transformers no longer corresponds to the direct linear-followed-by-norm structure that our theory covers'), yet the correction is applied to essentially every linear layer. Because the cross term is not guaranteed to vanish, the observed LLM gradient-norm behavior and the success of AdamC cannot be attributed to the proposed mechanism without a direct measurement of ⟨g_t, x_t⟩ (or of the residual corrections) on at least one representative Transformer layer. This is the central load-bearing gap for the LLM claim.","section":"Section 6, Figure 4"},{"comment":"The uncoupled update x_{t+1} = x_t - γ_t g_t - λ_t x_t gives the steady-state ratio ∥g_t∥/∥x_t∥ = √(2λ_t)/γ_t. Substituting λ̂_t = λ γ_t/γ_max yields √(2λ γ_t/γ_max)/γ_t = √(2λ/(γ_max γ_t)), not √(2λ/γ_max) as claimed. The proposed correction makes the ratio schedule-independent only when the weight-decay term is multiplied by the learning rate, as in Algorithm 1's line 12, where the decay is γ_t^2 λ/γ_max. The manuscript should correct this derivation, define SGDC's update precisely (coupled or uncoupled), and, if the uncoupled form is intended, use λ γ_t^2/γ_max as the corrected decay.","section":"Section 5"},{"comment":"The 'lower loss throughout training' claim for LLMs rests on a single 200B-token run with no error bars or seed variation. The final loss difference (2.461 vs 2.457) is within plausible run-to-run variation. Given that the Transformer extension is already heuristic, the loss claim needs at least three independent runs or a checkpoint-level uncertainty estimate.","section":"Figure 4"}],"minor_comments":[{"comment":"The first display equation in Section 5 uses an undefined symbol η; this is likely a typo for γ_t and should be fixed.","section":"Section 5"},{"comment":"Line 12's decay term γ_t^2 λ/γ_max should be explicitly linked to λ̂_t = λ γ_t/γ_max in the text, since the appearance of γ_t^2 is otherwise surprising.","section":"Algorithm 1"},{"comment":"The phrase 'it's relation' should be 'its relation'.","section":"Section 2"},{"comment":"The y-axis labels appear as '10 2' and '10 1' due to missing superscripts; please format them as 10^{-2} and 10^{-1}.","section":"Figures 2 and 3"},{"comment":"The Hölder-inequality approximation for ∥x_t∥_{A_t} is explicitly non-rigorous; a sentence noting that the infinity-norm interpretation is heuristic would help prevent over-reading.","section":"Section 4.2"}],"recommendation":"major_revision","confidential_remarks":"The paper's positioning relative to Xie et al. (2023) is fair, and the self-citations are not excessive. The main editorial concern is that the Transformer extrapolation is presented more strongly than the evidence supports; if the journal values empirical rigor, the direct cross-term measurement should be requested before acceptance."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Short version: the paper gives a clean steady-state explanation for why gradient norms rise at the end of training under cosine schedules, and proposes a one-line fix: scale weight decay by γt/γmax. The ImageNet ResNet experiments genuinely support the mechanism. The LLM part, which is the headline, is where it gets soft: the theory assumes ⟨g, x⟩ = 0, which holds for a layer directly followed by normalization, but Llama's residual structure doesn't satisfy that, and the paper admits it. So the paper doesn't actually derive the Transformer tail blow-up; it asserts it by analogy and then shows the fix still works empirically.\n\nWhat's new: the ratio interpretation and the specific corrected decay. The underlying equilibrium is Van Laarhoven's, and Xie et al. already suggested scheduling weight decay, but I don't think they gave this particular constructive form or connected it to the observed gradient-norm shape. The math in Section 3 is correct under the stated assumptions, and the figures showing layer ratios tracking sqrt(2λ/γ) are convincing for ResNet.\n\nSoft spots, in order. First, the Transformer extrapolation is the load-bearing claim. The cross term ⟨g, x⟩ is not zero for attention and MLP projections; measuring it directly would decide whether Eq. (2) governs the LLM tail or whether AdamC just reduces late-training weight decay. The paper's own Section 6 concession doesn't help; it applies the correction to all linear layers anyway. Second, the \"lower loss\" evidence is thin: the ImageNet gain is ~0.1% and the LLM gain is 0.004 loss, with no error bars on the LLM run. A matched baseline controlling total weight decay would be needed to show the correction, not just less decay, is responsible. Third, novelty is moderate, but the paper is honest about prior work.\n\nWho is this for? People training long LLM runs and optimizer theorists. It's a useful note, not a breakthrough. I'd send it out for review; the mechanism is plausible, the fix is cheap, and a referee can ask for the direct measurement that would settle the Transformer question.","headline":"A clean equilibrium explanation for end-of-training gradient blow-up on normalized layers, with a cheap fix that works on ImageNet—but the Transformer headline rests on an unverified analogy.","tokens_in":8334,"tokens_out":2408,"would_cite":true,"duration_ms":22736,"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":"Weight decay interacting with the learning-rate schedule causes the end-of-training gradient spike in LLMs, and a corrected decay term removes it.","keywords":["gradient norm","weight decay","learning rate schedule","normalization layers","large language model training","AdamW","layer balancing","AdamC"],"falsifier":"Train a small transformer with a cosine learning-rate schedule, log the per-layer inner product $\\langle g_t, x_t\\rangle$ alongside the gradient and weight norms, and compare the observed ratio $\\|g_t\\|/\\|x_t\\|$ to the predicted $\\sqrt{2\\lambda/\\gamma_t}$. If the ratio follows the prediction even when the cross term is non-negligible, the orthogonality assumption is harmless; if the trajectories diverge once the cross term is included, the orthogonality assumption is what breaks.","tokens_in":7263,"feed_emoji":"📉","tokens_out":10703,"duration_ms":80167,"temperature":0.7,"pith_summary":"The paper argues that the rapid rise in gradient norm seen near the end of long LLM training runs is not a mysterious instability but a predictable side effect of weight decay pulling normalized layers toward a moving equilibrium. For a layer directly followed by a normalization operation, the ratio of gradient norm to weight norm is driven to $\\|g_t\\|/\\|x_t\\| = \\sqrt{2\\lambda/\\gamma_t}$, where $\\gamma_t$ is the current learning rate, so as the schedule decays the target gradient magnitude grows. The same steady-state argument is carried over to AdamW, and uncoupled weight decay by itself does not remove the problem. The paper proposes a correction, $\\hat{\\lambda}_t = \\lambda \\gamma_t/\\gamma_{\\max}$, that keeps the target ratio constant, and shows that the corrected optimizers AdamC and SGDC eliminate the end-of-training spike while giving lower loss on a 120M-parameter LLM and on ImageNet ResNet-50.","feed_headline":"Scale weight decay with the LR schedule to stop gradient spikes","feed_subtitle":"A time-varying weight-decay coefficient keeps the gradient-to-weight target fixed, lowering loss in LLM and ImageNet training.","key_machinery":"The central object is the ratio $\\|g_t\\|/\\|x_t\\|$ for a layer immediately followed by a normalization operation. The load-bearing identity is the orthogonality condition $\\langle g_t, x_t\\rangle = 0$, which simplifies the norm dynamics of the weight update to a form whose steady state is $\\|g_t\\|/\\|x_t\\| = \\sqrt{2\\lambda/\\gamma_t}$. The paper uses this moving equilibrium both to explain the end-of-training gradient spike and to design the corrected weight-decay coefficient $\\hat{\\lambda}_t = \\lambda \\gamma_t/\\gamma_{\\max}$, which freezes the target ratio at a constant value.","core_discovery":"The central discovery is that weight decay and the learning-rate schedule together set a time-varying target for the gradient-to-weight norm ratio of normalized layers. Starting from the SGD update with weight decay, the orthogonality identity $\\langle g_t, x_t\\rangle = 0$ for a layer followed by normalization reduces the squared-norm update to $\\|x_{t+1}\\|^2 = (1 - \\gamma \\lambda)^2 \\|x_t\\|^2 + \\gamma^2 \\|g_t\\|^2$. Setting $\\|x_{t+1}\\| = \\|x_t\\|$ at the steady state and dropping the small $\\lambda^2$ term yields $\\|g_t\\|/\\|x_t\\| = \\sqrt{2\\lambda/\\gamma_t}$. When $\\gamma_t$ decays toward zero under a cosine or similar schedule, this target ratio blows up, which is the mechanism behind the observed gradient-norm increase. For AdamW the same argument holds in a weighted norm, and the paper shows the original Adam's coupled weight decay does not admit a uniform ratio, which it presents as an explanation for AdamW's empirical advantage. The proposed corrected decay, $\\hat{\\lambda}_t = \\lambda \\gamma_t/\\gamma_{\\max}$, fixes the time dependence and, according to the experiments, removes the gradient blow-up while improving loss.","pith_inferences":["If the ratio $\\|g_t\\|/\\|x_t\\|$ is the right control parameter, monitoring it during training could give an early signal that the schedule is pushing the model into the tail-blow-up regime before the loss degrades.","The same ratio-balancing logic suggests the corrected weight decay might also make the optimal weight-decay coefficient less sensitive to schedule length and warmup, since it removes the schedule's time dependence from the steady-state target.","The paper applies the correction to all linear layers except the output; a more precise rule based on whether a layer's activations are normalized, or on the measured cross term $\\langle g_t, x_t\\rangle$, could be tested to see whether it further improves transformer training.","A natural stress test is to apply AdamC to much longer runs or to schedules with more aggressive decay, where the uncorrected spike is largest; the correction should keep the gradient norm flat if the mechanism is complete."],"forward_implications":["On the 120M-parameter LLM trained for 200B tokens, AdamC removes the second-half gradient-norm increase and yields lower loss than AdamW throughout training.","On ImageNet ResNet-50, SGDC eliminates the rapid end-of-training gradient spike that SGDM exhibits, while slightly improving test accuracy (77.07% vs 76.95%).","Because the corrected weight decay keeps the gradient-to-weight ratio constant, it also stabilizes weight norms, which otherwise decay rapidly under a cosine schedule.","The analysis implies that AdamW's decoupled weight decay, not just its adaptive steps, is what allows all normalized layers to settle at the same ratio; Adam's coupled decay leaves layers unbalanced."],"supporting_citations":[{"why":"Supplies the steady-state analysis of weight decay and the gradient-to-weight ratio that the paper extends.","marker":"Van Laarhoven (2017)"},{"why":"Defines LayerNorm, the normalization operation whose orthogonality property anchors the derivation.","marker":"(Ba et al., 2016)"},{"why":"Defines BatchNorm, the other normalization operation covered by the orthogonality assumption.","marker":"(Ioffe and Szegedy, 2015)"},{"why":"Introduced decoupled weight decay (AdamW), the starting point for the paper's correction.","marker":"(Loshchilov and Hutter, 2019)"},{"why":"Previously observed gradient-norm growth with weight decay and proposed scheduling weight decay; the paper contrasts its precise mechanism and exact correction with this earlier bound.","marker":"(Xie et al., 2023)"},{"why":"Provides the rotational-equilibrium view of weight decay that independently supports the layer-balancing conclusion.","marker":"(Kosson et al., 2024)"},{"why":"Gives the effective learning rate for momentum, which the paper uses to adjust the steady-state ratio for SGD with momentum.","marker":"(Defazio and Gower, 2021)"}],"fun_headline_variants":["Fix gradient spikes by scaling weight decay with LR","Time-varying weight decay removes end-of-training gradient spikes","Coupling weight decay to LR schedule fixes final gradient spikes","Fix gradient blowup by making weight decay track LR"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The argument rests on assuming the gradient and the weights of each normalized layer are exactly orthogonal, an identity that holds for a layer directly followed by a normalization layer but is only approximate for transformer layers, where the paper nevertheless applies the correction to essentially all linear layers.","fun_headline_variants_meta":{"raw":{"variants":["Fix gradient spikes by scaling weight decay with LR","Time-varying weight decay removes end-of-training gradient spikes","Coupling weight decay to LR schedule fixes final gradient spikes","Fix gradient blowup by making weight decay track LR"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000779,"raw_usage":{"total_tokens":3402,"prompt_tokens":862,"completion_tokens":2540,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":478,"completion_tokens_details":{"reasoning_tokens":2477}},"tokens_in":478,"tokens_out":2540,"duration_ms":18040,"temperature":1.0,"reasoning_tokens":2477,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-07T11:26:49.631974+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Train a small transformer with a cosine learning-rate schedule, log the per-layer inner product $\\langle g_t, x_t\\rangle$ alongside the gradient and weight norms, and compare the observed ratio $\\|g_t\\|/\\|x_t\\|$ to the predicted $\\sqrt{2\\lambda/\\gamma_t}$. If the ratio follows the prediction even when the cross term is non-negligible, the orthogonality assumption is harmless; if the trajectories diverge once the cross term is included, the orthogonality assumption is what breaks.","supporting_citations":[{"cited_title":"and Szegedy, C","cited_arxiv_id":null,"evidence_quote":"Defines BatchNorm, the other normalization operation covered by the orthogonality assumption."},{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"Previously observed gradient-norm growth with weight decay and proposed scheduling weight decay; the paper contrasts its precise mechanism and exact correction with this earlier bound."}],"review_version":1}