[论文解读] Rigorous dynamical mean field theory for stochastic gradient descent methods

Cédric Gerbelot, Emanuele Troiani+3 more|arXiv (Cornell University)|Oct 12, 2022

Stochastic Gradient Optimization Techniques被引用 6

一句话总结

本文在高维设置下，为离散时间随机梯度下降（SGD）及相关一阶优化方法建立了严格的动力学平均场理论（DMFT）方程。证明了在具有通用协方差矩阵的高斯数据上，SGD 的渐近动力学可被精确描述为离散 DMFT 方程，这些方程通过自洽核捕捉记忆效应，并能通过磁化和范数动力学精确预测泛化性能。

ABSTRACT

We prove closed-form equations for the exact high-dimensional asymptotics of a family of first order gradient-based methods, learning an estimator (e.g. M-estimator, shallow neural network, ...) from observations on Gaussian data with empirical risk minimization. This includes widely used algorithms such as stochastic gradient descent (SGD) or Nesterov acceleration. The obtained equations match those resulting from the discretization of dynamical mean-field theory (DMFT) equations from statistical physics when applied to gradient flow. Our proof method allows us to give an explicit description of how memory kernels build up in the effective dynamics, and to include non-separable update functions, allowing datasets with non-identity covariance matrices. Finally, we provide numerical implementations of the equations for SGD with generic extensive batch-size and with constant learning rates.

研究动机与目标

• 为离散时间随机梯度下降（SGD）及相关优化算法提供动力学平均场理论（DMFT）方程的严格推导。
• 将先前的 DMFT 结果从 i.i.d. 子高斯数据扩展到具有任意协方差矩阵的一般高斯数据。
• 通过推导有效随机积分微分方程中的显式记忆核，对动力学中的记忆效应进行建模。
• 通过高维极限下的磁化和范数动力学，实现对泛化性能（例如估计器与真实值之间的余弦相似度）的精确预测。
• 通过大规模小批量和恒定学习率下的 SGD 数值实现，验证离散 DMFT 框架。

提出的方法

• 通过迭代高斯条件化和微扰理论分析 SGD 中权重更新的无限维极限，推导出离散时间 DMFT 方程。
• 引入一组自洽的随机积分微分方程，描述估计器权重向量和残差误差的有效动力学。
• 通过损失函数对历史梯度的泛函导数计算记忆核，捕捉动力学中的非马尔可夫效应。
• 围绕参考轨迹进行一阶微扰展开以推导有效动力学，其各项在高维极限下收敛于高斯过程。
• 推导出磁化（与真实值的余弦相似度）和权重范数的闭式方程，从而实现对泛化性能的预测。
• 对具有恒定学习率和大规模小批量的 SGD 实现推导出的 DMFT 方程，验证渐近预测结果。

实验结果

研究问题

• RQ1 能否为具有非 i.i.d. 数据（包括非单位协方差矩阵）的离散时间 SGD 推导出严格的 DMFT 方程？
• RQ2 在高维情形下，由梯度历史引起的记忆效应如何在 SGD 的渐近动力学中体现？
• RQ3 离散 DMFT 方程在多大程度上能准确描述 SGD 在高维极限下的泛化性能？
• RQ4 DMFT 框架能否扩展至包含非可分更新函数和非 i.i.d. 数据，而不仅限于 i.i.d. 子高斯设计？
• RQ5 在 DMFT 框架下，离散时间 SGD 动力学与连续时间梯度流之间的确切关系是什么？

主要发现

• 本文证明，即使在非 i.i.d. 高斯数据下，离散时间 DMFT 方程也能在高维极限下对 SGD 动力学提供精确的渐近描述。
• 记忆核自然地从损失函数的泛函导数中产生，并被显式计算，从而捕捉过去梯度对当前更新的影响。
• 泛化性能完全由磁化 $ m^t = \text{corr}(\bm{w}^*, \bm{w}^t) $ 描述，其演化遵循包含损失导数和噪声协方差的自洽方程。
• 有效动力学由一组涉及记忆核 $ R_g(t,t') $、具有协方差 $ C_g(t,t') $ 的噪声过程 $ u^t $ 和自正则化项 $ \tilde{\nu}^t $ 的随机积分微分方程描述。
• DMFT 方程的数值实现能准确预测高维设置下估计器范数和与真实值余弦相似度的演化。
• 该框架适用于一大类一阶方法，包括 SGD、Nesterov 加速法以及基于动量的算法，在经验风险最小化框架下均适用。