[论文解读] The high-dimensional asymptotics of first order methods with random data

Michael Celentano, Cheng Chen+1 more|arXiv (Cornell University)|Dec 14, 2021

Theoretical and Computational Physics被引用 15

一句话总结

本文建立了对具有 i.i.d. 子高斯条目的随机数据矩阵上的一阶优化方法（如梯度流）的高维渐近表征。通过一种新颖的近似消息传递（AMP）算法约化方法，证明了其动力学在 k 维中普遍收敛到由动力学平均场理论（DMFT）所确定的非线性随机过程的确定性系统，该系统与子高斯分布的具体形式无关。

ABSTRACT

We study a class of deterministic flows in ${\mathbb R}^{d imes k}$, parametrized by a random matrix ${\boldsymbol X}\in {\mathbb R}^{n imes d}$ with i.i.d. centered subgaussian entries. We characterize the asymptotic behavior of these flows over bounded time horizons, in the high-dimensional limit in which $n,d o\infty$ with $k$ fixed and converging aspect ratios $n/d oδ$. The asymptotic characterization we prove is in terms of a system of a nonlinear stochastic process in $k$ dimensions, whose parameters are determined by a fixed point condition. This type of characterization is known in physics as dynamical mean field theory. Rigorous results of this type have been obtained in the past for a few spin glass models. Our proof is based on time discretization and a reduction to certain iterative schemes known as approximate message passing (AMP) algorithms, as opposed to earlier work that was based on large deviations theory and stochastic processes theory. The new approach provides a unified view of a general class of algorithms and implies that the high-dimensional behavior of the flow is universal with respect to the distribution of the entries of ${\boldsymbol X}$. As specific applications, we obtain high-dimensional characterizations of gradient flow in some classical models from statistics and machine learning, under a random design assumption.

研究动机与目标

• 表征在 n, d → ∞ 且 n/d → δ、k 固定时，随机设计设置下一阶优化流的高维极限。
• 建立这些流的普遍渐近描述，且独立于数据矩阵条目的子高斯分布形式。
• 为高维统计与机器学习中梯度流及相关算法的分析提供严谨且统一的框架。
• 推导显式的固定点方程（即 DMFT 方程），以 k 维随机过程的形式描述优化轨迹的极限行为。

提出的方法

• 将优化流形式化为由具有 i.i.d. 中心化子高斯条目的随机数据矩阵 X ∈ R^{n×d} 驱动的 R^{d×k} 上的确定性动力系统。
• 通过时间离散化近似连续流，并将其简化为类似于近似消息传递（AMP）算法的迭代格式。
• 建立流的 empirical distribution 收敛到由非线性 SDE 系统控制的确定性 k 维随机过程。
• 证明极限动力学满足一个固定点方程，该方程涉及信号与噪声的协方差结构，其推导基于函数空间上的压缩映射论证。
• 将该框架应用于广义线性模型和带随机特征的神经网络训练等具体模型。
• 利用旋转不变性与 Stein 引理简化并求解固定点方程，将其与近端算子及 Hessian 结构联系起来。

实验结果

研究问题

• RQ1 在一阶优化方法于具有 i.i.d. 随机数据矩阵的高维设置下，其行为如何？
• RQ2 梯度流的高维动力学是否可普遍描述，而与数据的子高斯分布形式无关？
• RQ3 优化轨迹的极限行为在协方差与均值演化方面如何体现？
• RQ4 动力学平均场理论（DMFT）的固定点方程如何从高维连续时间流中导出？
• RQ5 在高维极限下，流收敛到其极限分布的收敛速率是多少？

主要发现

• 一阶流的高维极限收敛到由非线性 SDE 系统控制的 k 维随机过程，且独立于数据矩阵条目的子高斯分布形式。
• 极限动力学满足一个涉及 Hessian 逆与噪声协方差的固定点方程，可通过函数空间上的压缩映射论证求解。
• 在正则化风险强凸的条件下，流收敛到极限分布的速率呈指数快速，满足 ∥θt − θ∞∥L2 ≤ Ce−c0t（c0 > 0）。
• 参数的极限协方差结构由一个固定点方程表征，该方程将信号先验、噪声与损失函数的 Hessian 联系起来。
• 该框架普遍适用于广义线性模型和带随机特征的两层神经网络，可提供优化路径的精确渐近描述。
• 所推导的 DMFT 方程等价于近端算子与基于 Hessian 的更新，将优化动力学与高维统计估计联系起来。