Is Stochastic Gradient Descent Effective? A PDE Perspective on Machine Learning processes

TL;DR

This paper models stochastic gradient descent (SGD) in machine learning using PDEs, revealing how it concentrates around local minima and escapes suboptimal points, especially in non-convex and degenerate cases.

Contribution

It provides a PDE-based analysis of SGD's dynamics in non-convex settings, including new bounds on escape times and convergence behavior under degenerate diffusion.

Findings

SGD concentrates near local minima in the drift regime.

Stochastic fluctuations enable escape from suboptimal minima.

New bounds on Mean Exit Time for non-convex, degenerate cases.

Abstract

In this paper we analyze the behaviour of the stochastic gradient descent (SGD), a widely used method in supervised learning for optimizing neural network weights via a minimization of non-convex loss functions. Since the pioneering work of E, Li and Tai (2017), the underlying structure of such processes can be understood via parabolic PDEs of Fokker-Planck type, which are at the core of our analysis. Even if Fokker-Planck equations have a long history and a extensive literature, almost nothing is known when the potential is non-convex or when the diffusion matrix is degenerate, and this is the main difficulty that we face in our analysis. We identify two different regimes: in the initial phase of SGD, the loss function drives the weights to concentrate around the nearest local minimum. We refer to this phase as the drift regime and we provide quantitative estimates on this concentration phenomenon. Next, we introduce the diffusion regime, where stochastic fluctuations help the learning process to escape suboptimal local minima. We analyze the Mean Exit Time (MET) and prove upper and lower bounds of the MET. Finally, we address the asymptotic convergence of SGD, for a non-convex cost function and a degenerate diffusion matrix, that do not allow to use the standard approaches, and require new techniques. For this purpose, we exploit two different methods: duality and entropy methods. We provide new results about the dynamics and effectiveness of SGD, offering a deep connection between stochastic optimization and PDE theory, and some answers and insights to basic questions in the Machine Learning processes: How long does SGD take to escape from a bad minimum? Do neural network parameters converge using SGD? How do parameters evolve in the first stage of training with SGD?