跳转至

Does Weight Decay Enhance Training Stability?

会议: NeurIPS 2026
arXiv: 2605.16622
领域: 学习理论
关键词: 权重衰减, 稳定性边缘, 神经正切核, 相变, 动力学阻尼

一句话总结

本文通过将带权重衰减的梯度下降建模为欠阻尼谐振子,揭示了权重衰减在稳定性边缘(EoS)下通过全局几何对齐项引发的动力学阻尼与相变机制,证明了参数空间曲率下降能直接约束神经正切核谱半径并增强特征学习稳定性。

研究背景与动机

在现代深度学习优化实践中,权重衰减(Weight Decay)早已超越了经典的吉洪诺夫正则化(Tikhonov regularization)或岭回归范畴。传统理论将其视为主观假设空间的静态惩罚项,用于限制模型容量或改善适定性;而在实际的大模型预训练(如基于 AdamW 的优化)中,研究人员和从业者普遍依赖权重衰减来“稳定训练动力学”。然而,关于权重衰减究竟通过何种具体动力学机制稳定训练过程,在优化理论界长期缺乏严格的因果性解释与定量动力学刻画。

当前对神经网络训练稳定性的理解通常分散在三个相互交织但视角各异的维度:控制前向激活尺度的信号传播稳定性(如 \(\mu\text{P}\) 框架)、衡量函数空间泛化与特征演化的神经正切核(NTK)动力学,以及刻画参数空间局部曲率演化的稳定性边缘(Edge of Stability, EoS)。在全量梯度下降中,渐进变锐(progressive sharpening)通常使最大 Hessian 特征值 \(\lambda_{\max}(\nabla^2 L)\) 持续上升并最终悬停在临界不稳定性阈值 \(2/\eta\) 附近。经典启发式观点认为,加入衰减系数 \(\gamma\) 仅会将有效损失修改为 \(\tilde{L} = L + \frac{\gamma}{2}\|\theta\|^2\),从而将 EoS 阈值平移至 \(2/\eta - \gamma\)。但该静态假设无法解释为何从业者观察到的动力学振荡被大幅削弱,更无法预测曲率何时会发生非线性的结构性骤降。

本文的研究目标正是解构权重衰减在 EoS 制度下的动力学本质。研究切入点在于打破局部曲率分析的局限,将参数矢量的全局位置与损失锐度梯度的相互作用引入离散更新方程中。核心 idea:将带权重衰减的优化轨迹形式化为受全局对齐项调控的欠阻尼谐振子,发现全局对齐量 \(c_y = \langle \nabla S, \theta\rangle\) 跨越临界分岔阈值时会诱发极限环崩溃相变,从而使锐度在远低于理论边界处稳定并从函数空间抑制 NTK 谱爆炸。

方法详解

整体框架

本文构建了一套将局部高阶曲率与全局参数坐标完全解耦并重构的离散动力学分析框架。以全量梯度下降(GD)在最大锐度方向与渐进变锐方向张成的二维局部坐标系为研究对象,系统推导了权重衰减所引入的“局部摩擦力”与“全局几何对齐项”。方法流程涵盖从局部流形泰勒展开、欠阻尼振子化简,到全局相变分岔判定,最终桥接至经验 NTK 的谱分析与目标对齐。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["梯度下降参数更新<br/>含权重衰减项 θ_t+1 = (1-ηγ)θ_t - η∇L"] --> B["局部动力学投影与阻尼建模<br/>分解至最大锐度 u 与锐度梯度 ∇S"]
    B --> C["全局对齐与极限环崩溃判定<br/>计算 c_y = ⟨∇S, θ*⟩ 及其分岔临界点"]
    C -->|c_y < c_y^crit| D["CNN 模式:阻尼振荡<br/>锐度震荡呈指数衰减收敛至 2/η-γ"]
    C -->|c_y ≥ c_y^crit| E["MLP 模式:相变崩溃<br/>极限环坍缩,锐度稳定于远低于 2/η-γ 阈值"]
    D --> F["函数空间 NTK 谱约束与目标对齐<br/>通过外尔不等式抑制 λ_max(Θ) 并提升 KTA"]
    E --> F

关键设计

1. 局部动力学投影与阻尼建模:将权重衰减解构为负反馈回路中的线性阻尼

针对传统 EoS 理论将优化动力学仅视为无摩擦弹簧系统的局限,本文拓展了 Damian 等人的自稳定化(Self-Stabilization)框架。在参考鞍点 \(\theta^*\) 处(此时锐度 \(S(\theta^*) = \lambda_{\max}(\nabla^2 L(\theta^*)) = 2/\eta\)),将参数位移 \(\delta_t = \theta_t - \theta^*\) 投影至最大锐度主特征向量 \(u\) 与锐度梯度方向 \(\nabla S\),建立局部标量坐标 \(x_t = \langle u, \delta_t \rangle\) 与 \(y_t = \langle \nabla S, \delta_t \rangle\)。结合三阶自稳定恒等式 \(\nabla^3 L(u, u, \cdot) = \nabla S\),在剔除高阶微扰后推导出离散更新系统: $\(x_{t+1} \approx -x_t (1 + \eta y_t + \eta \gamma) - \eta \gamma c_x\)$ $\(y_{t+1} \approx y_t (1 - \eta \gamma) + \eta \left(\alpha - \beta \frac{x_t^2}{2}\right) - \eta \gamma c_y\)$ 其中 \(\alpha = -\langle \nabla L, \nabla S\rangle > 0\) 代表渐进变锐驱动力,\(\beta = \|\nabla S\|^2\)。连续包络展开显示,锐度偏差 \(V(\tau)\) 严格遵循二阶线性常微分方程 \(r^2 + \gamma r + 2(\alpha + \gamma^2) = 0\)。特征根实部恒为 \(-\gamma/2\),直接证明了权重衰减在局部充当线性阻尼器,使锐度振荡幅度以 \(e^{-\frac{\gamma}{2}\tau}\) 的速率呈指数衰减。

2. 全局对齐与极限环崩溃判定:全局相互作用引发的非平凡相变

针对局部理论无法解释“锐度为何会跌落至远低于经典极限”的核心痛点,本文发现了全局坐标投影项 \(c_y = \langle \nabla S, \theta^* \rangle\) 对系统能量势阱的本质重构。对稳态周期-2 轨道 \(x_t = (-1)^t \Delta + M\) 进行两步周期时均化积分,解得弹跳平衡中心漂移为 \(M = -\eta \gamma c_x / 2\)(微观非对称偏移),而极限环静止振幅 \(\Delta\) 则满足: $\(\Delta = \sqrt{\frac{2(\alpha + \gamma^2 - \gamma c_y)}{\beta}}\)$ 当优化轨迹的全局演化使得几何对齐项 \(c_y\) 满足临界分岔判据时: $\(c_y \ge c_y^{\text{crit}} = \frac{\alpha}{\gamma} + \gamma\)$ 根号下实数解条件被破坏,使得物理极限环振幅 \(\Delta \to 0\) 发生结构性崩溃。一旦自稳定项 \(-\frac{\beta}{2} x_t^2\) 消失,系统锐度将直接脱离 \(2/\eta - \gamma\) 边界,跌落至极低的新稳态定点 \(y^* = \frac{\alpha}{\gamma} - c_y \ll -\gamma\)。该机制解释了为何 MLP 中较大的 \(\gamma\) 会因 \(1/\gamma\) 缩放快速击穿临界阈值并呈现曲率断崖式下降,而 CNN 的几何结构使得 \(c_y\) 全程低于 \(c_y^{\text{crit}}\),因而仅表现为振荡平抑。

3. 函数空间 NTK 谱约束与目标对齐:将曲率平抑转化为表征泛化增益

针对曲率稳定是否等同于模型泛化提升的争论,本文建立了参数曲率到函数空间动力学的严格数学映射。在均方误差(MSE)损失下,网络经验 Hessian 矩阵精确分解为 Gauss-Newton 矩阵与残差阵之和 \(H = J^T J + R\)。由外尔不等式(Weyl's inequality),经验 NTK 矩阵 \(\Theta = J J^T\) 的最大特征值受控于: $\(|\lambda_{\max}(H) - \lambda_{\max}(\Theta)| \le \|R\|_2\)$ 随着经验风险收敛使得残差项 \(\|R\|_2 \to 0\),参数空间被下压的锐度稳态 \(S^*\) 直接对函数空间 NTK 施加了紧确的谱半径上界 \(\lambda_{\max}(\Theta) \le S^*\)。进一步基于秩-1 NTK 近似 \(K_t \approx c_t^2 X^T X + \|v_t\|^2 \hat{y} \hat{y}^T\),稳态 \(c_t^{*2} \approx n S^* / \lambda_1\) 的受压迫下降使得模态比值 \(\alpha_t = \|v_t\|^2 / c_t^2\) 显著放大。依据 Bunch-Nielsen-Sorensen 长期方程,这一比例放大迫使 NTK 主特征向量强力旋转并朝向真实目标向量 \(Y\) 深度对齐,阐明了权重衰减提升表征学习鲁棒性的深层机制。

损失函数 / 训练策略

优化使用全量梯度下降(Full-batch GD),单步更新严格遵循 \(\theta_{t+1} = (1 - \eta \gamma)\theta_t - \eta \nabla L(\theta_t)\)。学习率设置考虑架构初始化锐度差异:MLP 设为 \(\eta = 0.02\)(初始锐度约 10,远离 \(2/\eta = 100\));CNN 设为 \(\eta = 0.01\)(初始锐度约 90,低于对应阈值 200)。权重衰减系数 \(\gamma\) 扫掠范围涵盖 \([10^{-5}, 3 \times 10^{-2}]\)。锐度与三阶张量信息通过二阶自动微分与 Lanczos 迭代算法联合求解,避免了直接存储 \(\mathcal{O}(P^2)\) 完整 Hessian 的开销。

实验关键数据

主实验

实验基于 CIFAR-10 数据集规范化前 5,000 样本子集,系统评估了不同架构在全量梯度下降及多样权重衰减下的 EoS 稳定行为与收敛指标。

架构配置 学习率 \(\eta\) 权重衰减 \(\gamma\) EoS 进入步数 理论基准 \(2/\eta - \gamma\) 实际稳定锐度 \(\lambda_{\max}\) 稳态 NTK 最大特征值 \(\lambda_{\max}(\hat{\Theta})/N\)
MLP (ReLU) 0.02 0.000 (无正则) 3,200 100.00 99.85 1.82
MLP (ReLU) 0.02 0.005 4,600 99.995 99.60 1.68
MLP (ReLU) 0.02 0.015 (相变临界) 7,100 99.985 86.40 1.41
MLP (ReLU) 0.02 0.020 (相变区) 8,900 99.980 79.80 1.25
MLP (ReLU) 0.02 0.025 (深度相变) 14,200 99.975 68.20 1.08
CNN (2-Conv) 0.01 0.000 (无正则) 850 200.00 199.50 (混沌发散振荡) 2.45
CNN (2-Conv) 0.01 0.010 1,400 199.99 199.20 (弱阻尼收敛) 2.15
CNN (2-Conv) 0.01 0.020 2,150 199.98 199.10 (强阻尼快速稳定) 1.89

消融实验

针对 MLP 架构探究全局几何对齐项 \(c_y\)、临界分岔阈值 \(c_y^{\text{crit}}\) 与极限环存在性的对应关系。

权重衰减 \(\gamma\) 达到稳态前 \(c_y\) 峰值 分岔临界值 \(c_y^{\text{crit}}\) 是否击穿临界点 (\(c_y \ge c_y^{\text{crit}}\)) 极限环振幅 \(\Delta\) 稳态行为描述
\(1 \times 10^{-4}\) 12.4 1,450.2 否 0.048 标准 EoS 极限环,曲率紧贴 \(2/\eta\) 振荡
\(5 \times 10^{-3}\) 18.6 128.5 否 0.035 轻度阻尼,极限环稳定于 \(2/\eta - \gamma\)
\(1.5 \times 10^{-2}\) 24.3 21.8 是 (步数约 5,800) \(\to 0\) (坍缩) 极限环瓦解,锐度提前自发下落锁定于 ~86
\(2.0 \times 10^{-2}\) 26.8 14.2 是 (步数约 4,200) \(\to 0\) (坍缩) 提前相变,锐度稳定于 ~80 且伴随非单调损失
\(3.0 \times 10^{-2}\) 31.5 8.6 是 (步数约 2,100) \(\to 0\) (坍缩) 极早相变,渐进变锐极大延缓,稳态曲率极低

关键发现

  • 渐进变锐速率全面迟滞:无论是在 MLP 还是 CNN 中,增加 \(\gamma\) 均显著平抑了训练早期的曲率增长斜率,使得首次触发单步损失非单调增加(即 EoS 启动时间)成倍推迟。
  • 架构依赖的双重稳定分支:在 CNN 中,由于多尺度卷积参数流形的几何约束,全局对齐项 \(c_y\) 始终被压制在临界点之下,表现为振荡振幅随 \(\gamma\) 呈指数收敛;而在全连接 MLP 中,\(c_y^{\text{crit}} \propto 1/\gamma\) 的快速衰减促使系统跨越分岔点,触发相变使曲率稳定于 \(2/\eta\) 下方多达 20%~30%。
  • 曲率不等于稳定性判据:在 MLP 发生相变后(如 \(\gamma = 0.02, \lambda_{\max} \approx 80 \ll 99.98\)),单步损失变化量 \(\Delta L_t\) 依旧在正负间剧烈震荡,完全破坏了梯度下降下降引理(Descent Lemma),证明低曲率并不意味着系统离开 EoS。

亮点与洞察

  • 全局对齐项的动力学引入:突破了传统 EoS 分析仅关注局部二阶展开的认知局限,首次数学化地证明了全局参数向量 \(\theta\) 与锐度梯度 \(\nabla S\) 的内积 \(c_y\) 是调控局部流形极限环存亡的决定性分岔参数。
  • 重新审视曲率度量:揭示了在存在正则化扰动的复杂动力系统中,简单的凸/二次启发式曲率阈值(如 \(2/\eta\))不再是判断优化轨迹是否稳定的可靠指标。
  • 跨空间稳定性桥梁:将参数空间的局部相变理论与函数空间的经验 NTK 特征值界和核目标对齐(KTA)建立了内生因果链条,给出了权重衰减能够实质提升神经网络表征质量的坚实数理支撑。

局限与展望

  • 算力与规模受限:受限于步进计算 Lanczos 顶特征值及三阶张量外积的高昂复杂度 \(\mathcal{O}(d \cdot n)\),实验主要局限于小型两层感知机及浅层卷积网络;在超大参数规模 Transformer 中的泛化表现有待进一步验证。
  • 依赖全量梯度下降假设:核心动力学模型依赖于确定性全量梯度下降构建的负反馈回路,尚未完全推广至小批量随机梯度下降(Mini-batch SGD)与带有自适应预条件矩阵的优化器(如 AdamW)。
  • 架构几何依赖机理未完全闭环:CNN 为何能在轨迹全周期中强力限制 \(c_y\) 增长而避开相变,其深层对称性与权重共享机理仍需更深入的微分几何工具给予形式化证明。

相关工作与启发

  • vs Damian et al. (Self-stabilization, 2022):Damian 等人阐明了三阶导数在无正则化优化中提供自稳定负反馈的基本原理;本文在此基础上引入了非零权重衰减项,导出了显式的阻尼谐振子包络方程及极限环崩溃分岔准则。
  • vs Cohen et al. (Edge of Stability, 2021):Cohen 等人首次定义了 EoS 现象与 \(2/\eta\) 稳定悬停假说;本文打破了经典 \(2/\eta - \gamma\) 的平移认知,发现了由全局对齐引发的非线性亚阈值稳定相变。
  • vs Jiang et al. (NTK Evolution at EoS, 2025):Jiang 等人建立了 EoS 期间锐度下降与核目标对齐增强的联系;本文则通过形式化证明展示了权重衰减正是通过相变压低稳态常数 \(c_t^{*2}\),进而系统性驱动 KTA 增强的内生引擎。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次揭示权重衰减在 EoS 下诱发极限环相变的全局机制,视角开创且极具洞见。
  • 实验充分度: ⭐⭐⭐⭐ 理论推导完备,在标准基准小集与合成模型上进行了严格的 Lanczos 谱跟踪,受限于计算复杂度未拓展至大模型。
  • 写作质量: ⭐⭐⭐⭐⭐ 逻辑推导层次清晰,物理图景与优化动力学模型结合严丝合缝。
  • 价值: ⭐⭐⭐⭐⭐ 从根本上解答了权重衰减提升训练稳定性的学术直觉谜题,对优化器与调度策略设计具有关键理论指导意义。