[论文解读] Optimal learning rate schedules in high-dimensional non-convex optimization problems

Stéphane d’Ascoli, Maria Refinetti+1 more|arXiv (Cornell University)|Feb 9, 2022

Stochastic Gradient Optimization Techniques被引用 4

一句话总结

本文首次对高维非凸优化中的最优学习率调度进行了理论分析，采用随时间衰减的学习率 $\eta(t) = \eta_0 t^{-\beta}$ 的朗之万动力学方法。研究发现，在纯非凸景观中，$\beta < 1$ 可通过避免鞍点来加速优化；而在信号恢复问题中，采用两阶段策略——探索阶段保持恒定 $\eta$，收敛阶段采用 $\beta = 1$ 的衰减——可实现最优性能。

ABSTRACT

Learning rate schedules are ubiquitously used to speed up and improve optimisation. Many different policies have been introduced on an empirical basis, and theoretical analyses have been developed for convex settings. However, in many realistic problems the loss-landscape is high-dimensional and non convex -- a case for which results are scarce. In this paper we present a first analytical study of the role of learning rate scheduling in this setting, focusing on Langevin optimization with a learning rate decaying as $η(t)=t^{-β}$. We begin by considering models where the loss is a Gaussian random function on the $N$-dimensional sphere ($N ightarrow \infty$), featuring an extensive number of critical points. We find that to speed up optimization without getting stuck in saddles, one must choose a decay rate $β&lt;1$, contrary to convex setups where $β=1$ is generally optimal. We then add to the problem a signal to be recovered. In this setting, the dynamics decompose into two phases: an \emph{exploration} phase where the dynamics navigates through rough parts of the landscape, followed by a \emph{convergence} phase where the signal is detected and the dynamics enter a convex basin. In this case, it is optimal to keep a large learning rate during the exploration phase to escape the non-convex region as quickly as possible, then use the convex criterion $β=1$ to converge rapidly to the solution. Finally, we demonstrate that our conclusions hold in a common regression task involving neural networks.

研究动机与目标

• 为解决现代机器学习中常见设置——高维非凸优化下学习率调度缺乏理论理解的问题。
• 分析学习率衰减策略 $\eta(t) = \eta_0 t^{-\beta}$ 对非凸景观中优化动力学的影响。
• 确定在纯非凸和信号恢复设置下，最小化损失的最优衰减指数 $\beta$。
• 在使用随机梯度下降（SGD）的神经网络回归任务中，通过实际深度学习设置验证理论发现。

提出的方法

• 采用带球形约束和时变学习率 $\eta(t) = \eta_0 t^{-\beta}$ 的优化问题建模，通过温度 $T$ 对 SGD 的噪声进行建模。
• 分析两种模型：球形的夏里宁-基尔帕特里克（SK）模型（$p=2$）和带信号的矩阵张量（SMT）模型（$p>2$），均在 $N \to \infty$ 极限下定义于 $N$-维球面上的高斯随机损失函数。
• 应用统计物理方法，包括凯斯-赖斯方法和动态复制理论，计算关联函数与响应函数，并推导有效动力学。
• 推导出用于强制球形约束的拉格朗日乘子 $z_\infty$ 的渐近值，该值决定了非凸相中的阈值损失。
• 识别出两阶段动力学行为：初始探索阶段在粗糙的非凸区域进行，随后在存在信号时进入凸盆地的收敛阶段。
• 在教师-学生神经网络回归任务中，使用 $B=1$ 的小批量 SGD 进行验证，实证确认了最优调度策略。

实验结果

研究问题

• RQ1 在纯非凸、高维景观中，加速优化的最优学习率衰减指数 $\beta$ 是什么？
• RQ2 当存在可恢复信号时，其如何改变非凸优化中的最优调度策略？
• RQ3 在具有神经网络的真实深度学习设置中，能否验证最优 $\beta$ 的理论预测？
• RQ4 在信号恢复问题中，标志着从探索到收敛转变的动力学相变是什么？
• RQ5 噪声（温度 $T$）与学习率衰减之间的相互作用如何影响收敛速度和最终损失？

主要发现

• 在纯非凸景观中（如 $p=2$ 的 SK 模型），最优衰减速率为 $\beta = 1/2$，远慢于凸设置中使用的 $\beta = 1$。
• 对于更高阶的非凸性（$p > 2$），最优 $\beta$ 小于 1，表明为高效逃离鞍点，需要更慢的衰减。
• 当存在信号时，最优策略为两阶段调度：探索阶段保持恒定学习率以快速逃离非凸区域，随后在凸收敛阶段采用 $\beta = 1$ 的衰减。
• 从探索到收敛的转换时间 $t_{\text{cross}}$ 与“专业化”相变的结束时间一致，此时在恒定学习率下损失趋于平台。
• 在 $B=1$ 的教师-学生神经网络回归任务中，过早衰减学习率会导致高损失，而将衰减延迟至 $t_{\text{cross}}$ 之后可实现零损失收敛。
• 在低温度下，非凸相中的理论阈值损失与温度 $T$ 呈线性关系，证实了推导出的 $z_\infty$ 和 $\ell_{\text{th}}$ 表达式的鲁棒性。