{"id":"48e720ec-69be-429c-b92a-fcd34d6f98b4","arxiv_id":"2606.05326","paper_version":1,"verdict":"UNVERDICTED","confidence":"LOW","novelty_score":7.0,"correctness_risk":"unknown","formal_verification":"none","parameter_count":0,"one_line_summary":"Derives an effective free energy model and mean-field kinetic equation for gradient descent dynamics at the edge of stability in two-layer networks, interpreted as a Wasserstein-2 gradient flow.","lead":"The paper proposes a continuous-time effective model for gradient descent at the edge of stability that tracks average trajectories coupled to oscillation covariances through an effective free energy combining risk and curvature terms. A smart generalist might read it to understand how to model and predict the spiky loss behavior that appears during training of some neural network architectures.","discovery_kind":"new_method","skeptic_critique":{"model":"grok-4.3","headline":"Mean-field limit derivation assumes clean timescale separation for oscillations that the paper claims to relax, risking inconsistency in the kinetic equation.","rationale":"The reader's weakest_assumption correctly isolates the separation-plus-mean-field step as the least secure link; the numerical evidence cited in the abstract does not directly probe the comparable-timescale regime for the two-layer mean-field object, so the UNVERDICTED status is unchanged.","tokens_in":1762,"tokens_out":327,"duration_ms":12320,"concrete_test":"Take the two-layer network and effective-model ODEs from the paper; set the learning-rate parameter so that the observed oscillation period is within a factor of 2–3 of the mean-weight evolution timescale; integrate both the microscopic dynamics and the proposed kinetic equation for 10^4 steps and check whether the predicted loss-envelope matches the simulated spikes to within 10% relative L2 error on the upper envelope.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The central construction derives a kinetic equation for the joint law of weights and fluctuations as the W2 gradient flow of a macroscopic free energy, under the regime of 'stable non-vanishing oscillations' for wide two-layer nets. This requires that time-averaging of the fast covariance produces a closed, autonomous equation for the slow marginal even when the oscillation period is not asymptotically small compared with the drift of the mean weights. If the averaging step implicitly uses a separation that is only approximate, the resulting PDE may not be the correct effective description and the free-energy interpretation would rest on an unclosed moment hierarchy.","agreement_with_reader":"agree"},"referee_report":{"model":"grok-4.3","summary":"The manuscript studies gradient descent dynamics in the Edge of Stability regime, where large learning rates induce persistent oscillations in loss and sharpness. It proposes a continuous-time effective model coupling the slow average trajectory with the time-averaged covariance of fast oscillations, identifies an effective free energy (risk plus curvature-related entropic term) as the natural monitor, and derives a mean-field limit for wide two-layer networks yielding a novel kinetic equation for the joint law of weights and fluctuations. This PDE is shown to be the Wasserstein-2 gradient flow of a macroscopic free energy. Numerical evidence on matrix factorization and CIFAR-10 is provided to support accuracy in tracking oscillation envelopes.","tokens_in":1887,"tokens_out":466,"duration_ms":18490,"significance":"If the mean-field closure and free-energy interpretation hold without requiring strict timescale separation, the work supplies a new analytic tool for oscillatory training regimes that standard averaged analyses cannot capture, with direct relevance to understanding loss spikes and effective dynamics in overparameterized networks.","major_comments":[{"comment":"Mean-field limit section (derivation of the kinetic equation as W2 gradient flow): the closure of the time-averaged covariance into an autonomous equation for the slow marginal is asserted to hold under 'stable non-vanishing oscillations' even when periods are comparable to the mean-weight drift. The derivation steps that justify this closure (without implicit separation assumptions) must be made fully explicit, as any residual dependence on fast-slow separation would render the resulting PDE non-autonomous and undermine the free-energy interpretation.","section":"Mean-field limit derivation"}],"minor_comments":[{"comment":"Notation for the effective free energy functional should be introduced with an explicit equation number at first appearance to aid cross-referencing with the kinetic equation.","section":null},{"comment":"The numerical section would benefit from an additional panel or table quantifying the error between the predicted envelope and observed spikes across multiple random seeds.","section":null}],"recommendation":"major_revision","confidential_remarks":"The manuscript fits the math.OC scope well; however, the citation list appears light on prior Edge-of-Stability analyses that already relax strict separation, which the editor may wish to flag for completeness."},"author_rebuttal":{"model":"grok-4.3","summary":"We thank the referee for their careful reading and constructive feedback. We address the single major comment below.","responses":[{"response":"We agree that the steps justifying the autonomous closure must be stated explicitly. The manuscript already asserts that the closure holds under the sole assumption of stable non-vanishing oscillations (without requiring strict timescale separation), because the time average is taken with respect to the fast periodic orbit whose shape is determined instantaneously by the current slow state. In the revision we will expand the mean-field section with a self-contained derivation that (i) defines the averaging operator along the fast orbit, (ii) shows that the resulting covariance functional depends only on the slow marginal, and (iii) verifies that the resulting kinetic equation remains autonomous and is the Wasserstein-2 gradient flow of the macroscopic free energy. No additional separation hypothesis is introduced.","revision_made":"yes","referee_comment":"[Mean-field limit derivation] Mean-field limit section (derivation of the kinetic equation as W2 gradient flow): the closure of the time-averaged covariance into an autonomous equation for the slow marginal is asserted to hold under 'stable non-vanishing oscillations' even when periods are comparable to the mean-weight drift. The derivation steps that justify this closure (without implicit separation assumptions) must be made fully explicit, as any residual dependence on fast-slow separation would render the resulting PDE non-autonomous and undermine the free-energy interpretation."}],"tokens_in":1347,"tokens_out":310,"duration_ms":20997,"standing_objections":[]},"desk_editor":{"model":"grok-4.3","letter":"The main takeaway is that the authors build a continuous-time model coupling the slow average weights to the time-averaged covariance of fast oscillations, then package the whole thing as an effective free energy that adds a curvature-linked entropic term to the usual risk. For wide two-layer networks they pass to a mean-field limit and obtain a kinetic equation for the joint law of weights and fluctuations, which they cast as the Wasserstein-2 gradient flow of a macroscopic free energy. They also claim the construction works even when oscillation and drift timescales are comparable.\n\nWhat the paper does reasonably well is supply numerical evidence on matrix factorization and CIFAR-10 that the model tracks the envelope of the loss and sharpness spikes. That is concrete and addresses a regime people actually see in training.\n\nThe soft spot is the mean-field derivation itself. The stress-test note is on point: if the time-averaging step that closes the equation for the slow marginal still leans on some form of separation (even while the text says it relaxes that separation), the resulting PDE may not be fully autonomous and the free-energy interpretation rests on an incomplete moment hierarchy. I would want to see the precise closure argument and any error estimates before accepting the kinetic equation as exact rather than approximate. Scope is also limited to two-layer nets, so the work does not yet speak to deeper architectures.\n\nThis is for people already working on mean-field limits, Wasserstein geometry, or edge-of-stability dynamics. A reader comfortable with those tools will get the most from the kinetic construction and the numerics.\n\nIt is worth sending to peer review. The topic is practical, the proposed objects are new on their face, and the numerics give something to evaluate, even if the derivations will need tightening.","headline":"The paper gives a usable effective free energy plus Wasserstein kinetic equation for tracking oscillations in two-layer nets at the edge of stability, with decent numerical checks, though the mean-field closure under comparable timescales needs a close look.","tokens_in":2365,"tokens_out":444,"would_cite":false,"duration_ms":19822,"reading_group":"maybe","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"grok-4.3","headline":"Gradient descent with large learning rates follows an effective free energy that includes an entropic curvature term.","keywords":["edge of stability","gradient descent","effective free energy","mean-field limit","kinetic equation","Wasserstein gradient flow","two-layer networks","oscillations"],"falsifier":"Running gradient descent on a two-layer network with large learning rate and checking if the effective free energy decreases monotonically while the original loss oscillates.","tokens_in":2659,"feed_emoji":"","tokens_out":656,"duration_ms":21831,"temperature":0.7,"pith_summary":"The paper studies gradient descent when the learning rate is large enough to cause persistent oscillations in the loss and sharpness. It proposes a continuous-time model that couples the average weight trajectory with the covariance of its fast oscillations. This model shows that the relevant quantity is an effective free energy combining the risk with a curvature-related term. For wide two-layer networks, it derives a kinetic equation for the joint distribution of weights and fluctuations, interpreted as a gradient flow in Wasserstein space. Numerical tests on matrix factorization and CIFAR-10 confirm it captures oscillation envelopes.","feed_headline":"Effective free energy controls gradient descent at edge of stability","feed_subtitle":"Model tracks average trajectory and fast oscillations via curvature-augmented energy for two-layer nets","key_machinery":"The effective free energy, which combines the risk functional with a curvature-related entropic term, and the associated Wasserstein-2 gradient flow of the macroscopic free energy in the mean-field limit.","core_discovery":"In the Edge of Stability regime, the dynamics of gradient descent are captured by an effective model tracking the average trajectory and time-averaged covariance of oscillations. The natural quantity to monitor is an effective free energy that augments the original risk functional with a curvature-related entropic term. For wide two-layer networks under stable oscillations, a mean-field limit yields a kinetic equation for the joint distribution of weights and fluctuations, which is the Wasserstein-2 gradient flow of a macroscopic free energy.","pith_inferences":["If the free energy model holds, it suggests that optimization at the edge of stability can be understood as minimizing a modified landscape that penalizes high curvature.","This kinetic description might extend to other network architectures beyond two-layer networks if similar separation of timescales applies.","Testing the model on deeper networks could reveal whether the mean-field limit generalizes."],"forward_implications":["The envelope of oscillations can be tracked even when their dynamics evolve on similar timescales as the averaged weights.","The model predicts spikes that occur during training of some neural network architectures.","For wide two-layer networks, the joint distribution of weights and fluctuations follows a novel kinetic equation.","This allows accurate capture of oscillation behavior in tasks like matrix factorization and deep learning on CIFAR-10."],"fun_headline_variants":["Free energy model tracks GD at edge of stability with oscillations","Effective free energy governs oscillations in two-layer network GD","Mean field kinetic equation for free energy in edge of stability","Wasserstein gradient flow describes two-layer net free energy dynamics"],"cache_read_input_tokens":2112,"weakest_assumption_plain":"That persistent oscillations can be separated into a slow average trajectory plus fast time-averaged covariance, and that the mean-field limit applies to wide two-layer networks with stable oscillations.","fun_headline_variants_meta":{"raw":{"variants":["Free energy model tracks GD at edge of stability with oscillations","Effective free energy governs oscillations in two-layer network GD","Mean field kinetic equation for free energy in edge of stability","Wasserstein gradient flow describes two-layer net free energy dynamics"]},"model":"grok-4.3","cost_usd":0.00489,"raw_usage":{"total_tokens":2405,"prompt_tokens":683,"num_sources_used":0,"completion_tokens":65,"cost_in_usd_ticks":48899500,"prompt_tokens_details":{"text_tokens":683,"audio_tokens":0,"image_tokens":0,"cached_tokens":256},"completion_tokens_details":{"audio_tokens":0,"reasoning_tokens":1657,"accepted_prediction_tokens":0,"rejected_prediction_tokens":0}},"tokens_in":683,"tokens_out":65,"duration_ms":15230,"temperature":1.0,"reasoning_tokens":1657,"cache_read_input_tokens":256,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-06-28T04:58:47.409977+00:00","model_set":{"reader":"grok-4.3"},"falsifier":"Running gradient descent on a two-layer network with large learning rate and checking if the effective free energy decreases monotonically while the original loss oscillates.","supporting_citations":[],"review_version":1}