精读笔记

Problem Setting

《Gradient descent at the Edge of Stability: free energy model and kinetic description of the two-layer network》（arXiv preprint / 2026）关心的不是“GD 为什么能在 ηλmax≈2 附近训练”这个经验现象本身，而是如何给这个不稳定但有结构的离散动力学写出一个有效宏观描述。关键矛盾是：EoS 中的 fast oscillation 不是可平均掉的高频噪声，它的协方差会通过 Hessian 曲率反作用于平均轨迹；但如果完全保留离散动力学，又没有可分析的低维对象。

以前的 gradient flow 路线在这里本质失效，因为 η 不小，离散步的稳定边界本身就是动力学机制。Central Flow 抓住了 curvature-induced drift，但它假设 Σ 已经处在稳定边界上的平衡/约束态，因此更像 relaxed macroscopic limit。本文要解决的是更细一层的问题：当 Σ 没有瞬时松弛、loss/sharpness 仍在产生 spike 时，能否仍然追踪振荡包络，并找到一个比 raw loss 更自然的下降量。

Motivation

已有路线缺的是一个“振荡能量”的状态变量。EoS 现象中，loss 和 sharpness 的瞬时值高度振荡，直接看 E(θ̃k) 容易误判优化方向；只看平均路径 E(θ) 又漏掉了振荡在高曲率方向携带的能量。作者的核心动机是把这种能量显式纳入宏观状态，而不是把它当作扰动项。

这也解释了为什么会想到 free energy：Taylor 展开 E(θ+√ηδθ) 后，平均能量自然出现 E(θ)+(η/2)H:Σ。这个项不是装饰性的正则化，而是离散不稳定动力学中真实 loss 的二阶平均近似。换句话说，论文的缺口判断是：EoS 下的优化监控量应从 loss 改成 free energy，动力学变量应从 θ 改成 (θ,Σ)。

Core Idea

核心思想是将 GD 的 EoS 轨迹视为 slow-fast system：真实参数 θ̃k 被分解为平均轨迹 θk 加上 √η 量级的快速振荡 δθk，而振荡不被消去，只通过其协方差 Σk 进入宏观方程。这样得到的平均轨迹不仅有普通梯度项 −η∇E，还多了一个由 Hessian 空间变化率驱动的漂移 −(η²/2)∇∇²E:Σ；这正是离散振荡将系统从尖锐区域推走的机制。

和 prior 的本质区别在于：Central Flow 相当于只看 Σ relaxed 后的中心路径，而本文把 Σ 作为动态变量。这个改变很重要，因为 EoS 中 sharpness 的下降、自稳定、再上升，本质上都发生在 Σ 与 H 的耦合中。作者引入的 inductive bias 是“二阶振荡统计足以决定短时间 envelope”，而不是“轨迹可由平衡稳定边界约束决定”。这使模型更适合解释 spike/envelope，但也把有效性压在二阶闭合假设上。

Method

1. θ/Σ 有效 ODE：从 θ̃=θ+√ηδθ 出发，对 ∇E 做 Taylor 展开并对 fast variable 平均，得到 θ 的 drift；再对 δθ 的线性化更新取二阶矩，得到 Σ 的更新。它解决的是 raw GD 过细、gradient flow 过粗的问题；核心变化是把离散不稳定性的影响变成一个显式 curvature-covariance coupling。

2. Σ 动力学：dΣ/dt=−η(ΣH+HΣ)+η²HΣH。这个式子是本文比 Central Flow 更有信息量的地方。若在 H 的特征方向上看，Σ 的增长/衰减由离散稳定因子控制，ηλ 接近 2 时会出现非消失振荡；低于阈值的方向会被收缩，高曲率方向则表现出边界效应。它提供了 sharpness increase/decrease 的动态解释，而不是只给一个平衡约束。

3. Free energy：F=E+(η/2)H:Σ。它解决的是在振荡 regime 中 loss 本身不是 Lyapunov quantity 的问题。文中证明在 ||H||∞≤2/η 的条件下，F 的导数包含 −η||∇θF||² 以及一个非正曲率耗散项，因此 F 比 E 更接近真实可下降量。这个机制是论文最清楚的理论支点。

4. Two-layer mean-field kinetic equation：在宽两层网络中，把 Σ 写成低秩 V⊗V，并考虑经验分布 fN(θ,v) 的极限，得到扩展相空间上的 kinetic PDE。它解决的是有限维 (θ,Σ) 不能直接描述无限宽粒子系统的问题；核心变化是把 EoS 振荡方向 v 纳入 mean-field 粒子状态，使标准 Wasserstein gradient flow 被扩展为带 oscillation coordinate 的 free-energy flow。

Key Insight / Why It Works

最关键的 insight 是：EoS 的宏观 drift 来自“振荡协方差 × 曲率变化”，而不是来自噪声、显式正则或普通 gradient flow 修正。若 Hessian 在某方向太大，离散映射在该方向产生近二周期振荡；这些振荡通过 ∇H 的空间变化推动平均轨迹离开尖锐区域。这一机制解释了 self-stabilization，也解释了为什么只看 θ 的 loss 会错：一部分能量存在于 oscillation covariance 中。

真正有效的部分大概率是二阶矩闭合和 free energy 视角。Σ 方程虽然形式简单，但抓住了离散稳定因子 (I−ηH) 对协方差的作用；F 则把 Taylor 平均 loss 写成可监控量，因此实验中它能比 raw loss 更好地贴合 envelope 并不意外。相比之下，高维 CIFAR 实验中只跟踪 top Hessian eigenspace、少量 eigenvectors 不频繁重算，更像 engineering approximation；它说明局部 Hessian 顶部子空间可能稳定，但不是理论核心。

这不是 scaling、retrieval、data coverage 或 benchmark trick 类贡献，而是一个更好的 inductive bias：把离散优化器的非稳定模式当作状态变量。它本质上是 coarse-graining / kinetic modeling，而不是优化算法改进。文中没有证明二阶统计足够，也没有证明真实深网的 δθ 满足该 scaling；因此其有效性目前更像“短时间局部渐近模型”而非完整训练动力学定理。

Relation To Prior Work

最接近的是 Central Flow、Rod Flow、EoS self-stabilization 理论和两层网络 mean-field Wasserstein gradient flow。

相对 Central Flow，本文的实质新增信息是让 Σ 动起来。Central Flow 已经有 −(1/2)∇H:Σ 这类中心漂移，但它把 Σ 限制在 H≈2/η 的稳定边界核空间里，相当于假设快变量已经平衡。本文则试图描述 relaxation 过程本身，因此能覆盖 short-time envelope 和 spike。这个区别不是符号层面的，而是时间尺度建模的区别。

相对 Rod Flow，本文没有显式利用 near-two-periodicity 的几何结构，而是使用协方差/二阶矩描述。因此它更粗、更可解析，也可能更 general；但若真实 EoS 主要由二周期 rod-like mode 支配，Rod Flow 可能在相位结构上更精确。作者自己也暗示本文可能位于 Central Flow 与 Rod Flow 之间。

相对 mean-field neural network 文献，标准结果是 ρ(θ) 在 Wasserstein space 中做 risk 的 gradient flow；本文把粒子状态扩展到 (θ,v)，使振荡方向参与宏观 free energy。这个 kinetic PDE 是实质创新，但推导依赖非 maximal update regime 和低秩 Σ 假设，机器学习实践相关性仍有限。

看似新的“free energy”其实来自 Taylor 平均能量，是经典 coarse-graining 思想在 EoS GD 上的重组；真正新的是把它和 Σ 动力学、EoS stability boundary、mean-field kinetic limit 连在一起。

Dataset / Evaluation

实验覆盖低维矩阵分解与 CIFAR-10 上的小型 CNN/MLP。它们验证的核心 claim 是：模型 (3.4) 能在短时间尺度上追踪 loss/sharpness 振荡 envelope，free energy 比 raw loss 更适合作为下降监控量。这个 evaluation 对论文主张是基本对齐的，因为论文不是声称训练更好，而是声称建模更准。

但证据范围很有限。矩阵分解是可控 toy setting，适合验证机制但不能代表深网；CIFAR-10 使用小数据子集和小模型，且高维中无法维护完整 Σ，只跟踪 Hessian 顶部子空间。这意味着实验更多验证了“在 top eigenspace 近似下模型可用”，而不是完整协方差模型在真实规模深网中成立。

此外，长时间尺度上模型精度下降，作者也承认可能是高阶项被忽略。benchmark 没有验证 generalization、optimizer bias 是否真的导向 flatter minima，也没有检验 SGD、momentum、Adam 或大规模训练。因此 evaluation 支持 short-time effective dynamics，不支持更广的训练理论声明。

Limitation

最大限制是形式推导多于严格有效性理论。θ̃=θ+√ηδθ 的 scaling 是关键，但文中未充分说明何时成立；fast oscillation 的 ensemble/time average 需要某种 ergodicity 或周期稳定性，文中也没有证明。Σ 的二阶闭合忽略了三阶以上 fluctuation statistics，而长时间误差很可能正来自这些项。

模型还隐含 Hessian 局部结构相对稳定，尤其在高维实验里直接假设 top eigenvectors 不怎么旋转。这个假设在短窗口可能成立，但在真实训练跨阶段变化时不可靠。若 Hessian eigenspace 快速旋转、多个 sharp modes 交替主导，Σ 动力学可能需要更高阶或非局部状态。

mean-field kinetic 方程的实践上限也明显：推导 regime 不对应常见 maximal update；N→∞ 时 quadratic term 被丢弃，可能削弱对大步长有限宽网络的解释力；低秩 Σ 假设虽然数值上看似合理，但不是一般事实。它把原问题从“理解 GD 轨迹”转移为“理解扩展相空间 free-energy flow”，数学上优雅，但离真实深网优化仍有距离。

最后，实验增益来源不清：free energy 贴合 loss envelope 可能主要来自 Taylor 二阶校正，而不一定证明完整 ODE 的所有项都必要；高维近似中 eigenvector 固定策略可能贡献了稳定性，也可能掩盖模型缺陷。

Takeaway

• 1. EoS 下不应只建模平均权重；振荡协方差是决定宏观 drift 的状态变量。
• 这个视角可迁移到 momentum、Adam、SGD noise 等其他离散/随机优化动力学。
• 2. Free energy 比 loss 更适合作为 unstable discrete dynamics 的监控量。
• 以后研究 EoS 的“下降性”时，直接看 E(θ) 可能是错的对象。

一句话总结

这篇论文把 Edge-of-Stability 中的 GD 振荡从“需要平均掉的副作用”提升为“携带曲率能量并驱动平均轨迹的状态变量”，用 θ–Σ free-energy 动力学连接了 Central Flow、振荡包络建模与两层网络 mean-field kinetic theory。