2025-11-24T09:25:18.470449

Rigorous dynamical mean field theory for stochastic gradient descent methods

Gerbelot, Troiani, Mignacco et al.

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.

academic

Rigorous dynamical mean field theory for stochastic gradient descent methods

基本信息

论文ID: 2210.06591
标题: Rigorous dynamical mean field theory for stochastic gradient descent methods
作者: Cédric Gerbelot, Emanuele Troiani, Francesca Mignacco, Florent Krzakala, Lenka Zdeborová
分类: math-ph, cs.IT, cs.LG, math.IT, math.MP, stat.ML
发表时间: 2023年11月29日（arXiv v3版本）
论文链接: https://arxiv.org/abs/2210.06591

摘要

本文为一阶梯度优化方法（如SGD、Nesterov加速等）在高维渐近行为下建立了严格的闭式方程。这些方程与统计物理中的动态平均场理论（DMFT）的离散化形式完全一致。证明方法基于迭代高斯条件化技术，明确描述了有效动力学中记忆核的形成机制，并支持非可分离更新函数，从而可处理具有非单位协方差矩阵的数据集。论文还提供了针对具有广泛批量大小和恒定学习率的SGD的数值实现。

研究背景与动机

要解决的问题

本文旨在为随机梯度下降（SGD）及其变体在高维数据上的精确动力学行为提供严格的数学证明。具体而言，需要刻画这些算法在学习M估计器、浅层神经网络等模型时的渐近性质。

问题的重要性

理论基础缺失：尽管SGD是现代机器学习的核心优化工具，但对其高维动力学的精确理解长期停留在启发式物理方法层面
实践指导需求：精确的理论描述可以指导学习率、批量大小等超参数的选择
物理与数学的桥梁：将统计物理中的DMFT方法严格化，为跨学科研究提供坚实基础

现有方法的局限性

物理方法非严格：早期DMFT推导40,41,14,15基于启发式论证，缺乏数学严格性
连续时间限制：现有严格工作11主要关注梯度流的连续时间极限，而实际算法运行在离散时间
数据矩阵限制：先前严格结果11要求数据矩阵具有i.i.d.次高斯元素和单位协方差，限制了应用范围
确定性算法：未能处理SGD的随机性（如mini-batch采样、热噪声等）

研究动机

本文旨在克服上述局限，为离散时间随机优化算法建立严格的DMFT方程，并扩展到更广泛的数据分布和算法类别。

核心贡献

严格的离散时间DMFT方程：首次为离散时间一阶梯度方法（包括SGD、动量方法、Langevin算法等）建立了精确的高维渐近方程
迭代高斯条件化证明技术：提出了比现有AMP（近似消息传递）方法更直接简洁的证明框架，明确展示记忆核的形成机制
非可分离更新函数支持：允许处理具有任意良态协方差矩阵的数据，通过非可分离更新函数实现
广泛的算法覆盖：统一框架涵盖：
- 具有广泛批量大小的多轮SGD
- Polyak重球法和Nesterov加速梯度
- Langevin动力学（包含热噪声）
- 时变学习率和正则化
数值实现：提供了自洽方程的数值求解器，在teacher-student感知机模型上验证了理论预测

方法详解

任务定义

考虑以下经验风险最小化问题：

$\hat{w} \in \inf_{w \in \mathbb{R}^{d \times q}} L(Xw, y) + F(w)$

其中：

$X \in \mathbb{R}^{n \times d}$ ：设计矩阵（数据）
$y = \Phi_0(Xw^*) \in \mathbb{R}^n$ ：标签（由真实参数 $w^* \in \mathbb{R}^{d \times q}$ 生成）
$L, F$ ：可微的损失和正则化函数
$q$ ：有限的输出维度（如隐藏单元数）
$n, d \to \infty$ 且 $n/d = \alpha$ （高维极限）

通过一阶梯度方法求解：

$w^{t+1} = w^t - \gamma_t \left( X^\top \nabla L_t(Xw^t, y) + \nabla F(w^t) \right)$

理论框架架构

通用迭代形式

将算法重写为增量形式：

$v^{t+1} = h_t(\{v^k\}_{k=0}^t) + X^\top g_t(r^t)$ $r^t = X \sum_{k=0}^t v^k$

其中：

$v^t = w^t - w^{t-1}$ ：权重增量
$h_t, g_t$ ：伪Lipschitz连续的更新函数
$r^t$ ：预激活值

有效动力学（主定理3.2）

在高维极限下， $(v^t, r^t)$ 的分布由以下低维随机过程刻画：

$\nu^{t+1} = \theta^t \Gamma_t + h_t(\{\nu^k\}_{k=0}^t) + \sum_{k=0}^{t-1} \theta^k R_g(t,k) + u^t$

$\eta^t = \sum_{k=0}^{t-1} g^k(\eta^k) R_\theta(t,k) + \omega^t$

其中：

$\theta^t = \sum_{k=0}^t \nu^k$ ：有效权重
$\eta^t$ ：有效预激活
$u^t, \omega^t$ ：协方差为 $C_g(s,t), C_\theta(s,t)$ 的高斯过程

关键量定义：

响应核（记忆效应）： $R_\theta(t,s) = \lim_{d \to \infty} \frac{1}{d} \sum_{i=1}^d \mathbb{E}\left[\frac{\partial \theta^t_i}{\partial u^s_i}\right]$
$R_g(t,s) = \lim_{d \to \infty} \frac{1}{d} \sum_{i=1}^n \mathbb{E}\left[\frac{\partial \bar{g}^t_i}{\partial \omega^s_i}(\eta^t)\right]$
瞬时响应： $\Gamma_t = \lim_{d \to \infty} \frac{1}{d} \sum_{i=1}^n \mathbb{E}\left[\frac{\partial g^t_i}{\partial \eta^t_i}(\eta^t)\right]$
协方差： $C_\theta(t,s) = \lim_{d \to \infty} \frac{1}{d} \mathbb{E}[(\theta^t)^\top \theta^s]$
$C_g(t,s) = \lim_{d \to \infty} \frac{1}{d} \mathbb{E}[g^s(\eta^s)^\top g^t(\eta^t)]$

技术创新点

1. 迭代高斯条件化技术

核心思想：在每个时间步，将数据矩阵 $X$ 条件化到已观测的历史信息 $\mathcal{S}_t = \sigma(v^0, \ldots, v^t, r^0, \ldots, r^{t-1})$ 上。

正交分解（引理A.1）：

$X | \mathcal{S}_t \stackrel{d}{=} P_{M_{t-1}} X + X P_{W_t} - P_{M_{t-1}} X P_{W_t} + P^\perp_{M_{t-1}} \tilde{X} P^\perp_{W_t}$

其中：

$M_{t-1} = [m^0 | \cdots | m^{t-1}]$ ， $m^t = g_t(r^t)$
$W_t = [w^0 | \cdots | w^t]$
$\tilde{X}$ ： $X$ 的独立副本

关键洞察：

投影到历史子空间的部分产生记忆核
正交部分产生新的高斯噪声
通过归纳法精确控制各项的渐近行为

2. 记忆核的显式构造

通过Stein引理（引理A.3），将投影系数与偏导数联系：

$\frac{1}{d} \mathbb{E}[(\omega^s)^\top \omega^t] = \sum_{k=0}^{t-1} C_\theta(s,k) \alpha^{t,*}_k + C_\theta(s,t-1)$

其中 $\alpha^{t,*}$ 是投影系数的极限，满足：

$\alpha^{t,*} = \lim_{n,d \to \infty} \mathbb{E}\left[\left(\frac{1}{d} \Theta^\top_{t-1} \Theta_{t-1}\right)^{-1} \frac{1}{d} \Theta^\top_{t-1} (\theta^t - \theta^{t-1})\right]$

这明确展示了记忆如何通过历史迭代的投影累积。

3. 非可分离函数处理

对于协方差为 $\Sigma$ 的数据，通过变换 $\tilde{w} = \Sigma^{1/2} w$ 重写优化问题：

$\tilde{w}^{t+1} = \tilde{w}^t - \gamma \left( X^\top \nabla L(X\tilde{w}^t) + \Sigma^{-1/2} \nabla F(\Sigma^{-1/2} \tilde{w}^t) \right)$

正则化项变为非可分离函数 $\Sigma^{-1/2} \nabla F(\Sigma^{-1/2} \cdot)$ ，但仍可纳入框架。

4. 随机效应的统一处理

Mini-batch采样：通过独立Bernoulli变量 $s^t \in \{0,1\}^n$ 建模， $s^t_i \sim \text{Bern}(b)$
热噪声（Langevin）：在 $h_t$ 中添加 $\sqrt{T} z^t$ ， $z^t \sim \mathcal{N}(0, I_d)$
动量：在 $h_t$ 中包含历史增量项（如Polyak的 $\beta v^t$ ）

所有这些独立于 $X$ 的随机性可直接融入条件化框架。

证明核心步骤（以 $r^t$ 为例）

归纳假设：假设定理对 $r^0, \ldots, r^{t-1}, v^0, \ldots, v^t$ 成立。

目标：证明 $r^t$ 的渐近分布。

步骤1：条件化 $r^t | \mathcal{S}_t = r^{t-1} + (X P_{W_{t-1}} + P_{M_{t-1}} X P^\perp_{W_{t-1}} + P^\perp_{M_{t-1}} \tilde{X} P^\perp_{W_{t-1}}) v^t$

步骤2：逐项分析

第一项： $r^{t-1}$ 由归纳假设控制
第二项： $X P_{W_{t-1}} v^t = \sum_{k=0}^{t-1} r^k \alpha^{t,*}_k$ （投影系数）
第三项：产生记忆核 $\sum_{k=0}^{t-1} g^k(\eta^k) R_\theta(t,k)$
第四项：新高斯噪声 $\tilde{\omega}^t \sim \mathcal{N}(0, C^\perp_{v,t} \otimes I_n)$

步骤3：协方差匹配通过Stein引理验证组合噪声 $\omega^t = \sum_{k=0}^{t-1} \omega^k \alpha^{t,*}_k + \omega^{t-1} + \tilde{\omega}^t$ 具有正确的协方差结构 $C_\theta(s,t)$ 。

步骤4：提升条件使用伪Lipschitz函数的浓度性质（引理A.2），从条件分布提升到边缘分布。

实验设置

数据集

Teacher-Student二分类感知机：

输入： $x_\mu \sim \mathcal{N}(0, I_d)$ ， $\mu = 1, \ldots, n$
标签： $y_\mu = \text{sign}(x^\top_\mu w^*)$ ，其中 $w^* \sim \mathcal{N}(0, \frac{1}{d} I_d)$
参数： $d = 1000$ ， $\alpha = n/d \in \{0.9, 3\}$

损失函数

Logistic损失： $l(r, y) = \log(1 + e^{-yr})$
岭正则化： $F(w) = \frac{\lambda}{2} \|w\|^2_2$ ， $\lambda \in \{0.5, 1\}$

算法配置

学习率： $\gamma \in \{0.02, 0.04, 0.06\}$
批量大小： $b \in \{0.2, 0.5, 1.0\}$ （占数据集比例）
初始化： $w^0_i \sim \mathcal{N}(0, \frac{1}{d})$ i.i.d.

评价指标

余弦相似度（与教师向量）： $\frac{m^t}{\sqrt{C_\theta(t,t)}}$ 其中 $m^t = \lim_{d \to \infty} \mathbb{E}[(w^*)^\top w^t]$ 是磁化强度。

数值求解方法

自洽迭代（算法5.1）：

初始化响应核 $R_g, R_\theta$ 和辅助函数 $\Gamma_t, \nu_t$ 的猜测
在固定核下数值积分DMFT方程，生成随机过程 $\{\eta^t, \theta^t\}$
通过对生成过程平均更新核和辅助函数
重复直至收敛（图3显示收敛非常快）

实验结果

主要结果

学习率和批量大小的影响（图2）

观察：

完美匹配：理论曲线（连续线）与 $d=1000$ 的有限维模拟（点）几乎完全重合
学习率效应：
- $\gamma = 0.02$ ：收敛慢但稳定
- $\gamma = 0.04$ ：收敛速度适中
- $\gamma = 0.06$ ：初期振荡，但最终达到相似性能
批量大小效应：
- $b = 0.2$ ：噪声大，收敛慢但可能逃离局部最优
- $b = 1.0$ ：噪声小，收敛快且平滑

数值精度：即使在中等维度（ $d=1000$ ）下，理论预测的准确性也非常高，无需额外平均。

收敛速度（图3）

自洽迭代性能：

在2500次随机过程采样下，5-10次迭代即可收敛
使用70%新核+30%旧核的混合策略稳定收敛
磁化强度 $m^t$ 的理论值与模拟完全一致

样本分裂情况（定理4.1）

简化场景验证：

每步使用新数据矩阵 $A^t$ （样本分裂）
得到马尔可夫动力学（无记忆核）： $\omega^{t+1} = (1 - \gamma_t \alpha \mathbb{E}[f''(z^t)]) \omega^t + \gamma_t u^t$
图1显示即使在 $n=50, d=100$ 的极低维度下也能完美匹配