stochastic gradient descent, supervised learning, mass concentration, mean exit time, asymptotic behaviour, entropy method.

. Pri: 35Q68, 68T07, 35K65, 35B40.Sec: 60J60, 35J70, 35K15

• Departamento de Matemáticas, Universidad Autónoma de Madrid,
  ICMAT - Instituto de Ciencias Matemáticas, CSIC-UAM-UC3M-UCM,
  Campus de Cantoblanco, 28049 Madrid, Spain.

• E-mail: davide.barbieri@uam.es Web-page: https://sites.google.com/view/davidebarbieri/

• E-mail: matteo.bonforte@uam.es Web-page: https://verso.mat.uam.es/~matteo.bonforte/

• E-mail: ibarrondo.peio@gmail.com;

Acknowledgments↩︎

D.B., P.I. were partially supported by the research Projects PID2022-142202NB-I00 and PID2019-105599GB-I0 all funded by the Spanish Ministry of Science and Innovation MCIN/AEI/10.13039/501100011033. M.B., P.I. were partially supported by the research Projects MTM2017-85757-P, PID2020-113596GB-I00 and PID2023-150166NB-I00 all funded by the Spanish Ministry of Science and Innovation MCIN/AEI/10.13039/501100011033. D.B., M.B., P.I. acknowledge the financial support from the Spanish Ministry of Science and Innovation, through the “Severo Ochoa Programme for Centres of Excellence in R&D” CEX2019-000904-S and CEX2023-001347-S, and by the European Union’s Horizon 2020 research and innovation programme under the Marie Skłodowska-Curie grant agreement no. 777822. P.I. was partially funded by the FPU-grant FPU19/04791 from the Spanish Ministry of Universities.

1 Introduction↩︎

In recent decades, Machine Learning has gained significant attention across various scientific communities, thanks to its practical utility and wide range of applications. One of the main challenges in Machine Learning is choosing the weights or parameters of a neural network to minimize a non-convex loss function based on a given data set. This is a non-convex optimization problem formulated as follows: \[\arg\min\limits_{\theta\in\mathbb{R}^d}L(\theta):=\arg\min\limits_{\theta\in\mathbb{R}^d}\frac{1}{N}\sum_{i=1}^{N}L_i(\theta)\] where \(L,L_i:\mathbb{R}^d\rightarrow\mathbb{R}_+\) for \(i=1,\dots,N\) and \(\theta\in\mathbb{R}^d\) are the parameters of the model. The function \(L = \frac{1}{N}\sum_{i=1}^{N}L_i\) represents the total loss function of the given data set with \(N\) samples, while \(L_i\) represents the loss associated to the \(i^{\rm th}\) training sample. A classical approach to this problem is to perform a gradient descent, where given some initial weights \(\theta_0\in\mathbb{R}^d\) and a constant learning rate \(\eta>0\) the weights at step \(n\ge 0\) are updated by \[\theta_{n+1}=\theta_n-\eta\nabla L(\theta_n)\quad\qquad\forall n\ge 0\,.\] This learning algorithm is computationally too expensive to implement since at each step we have to compute \(N\) gradients of loss functions \(L_i\). Moreover, this method may lead the parameters to a local minimum of \(L\) instead of reaching the global minimum. In order to avoid these problems, the use of the stochastic gradient descent (SGD) has proven to be effective. Defining \(\lbrace \gamma_n\rbrace_{n\ge 1}\) as i.i.d. uniform random variables with values on \(\lbrace 1,\dots,N\rbrace\), the basic SGD algorithm reads \[\theta_{n+1}=\theta_n-\eta\nabla L_{\gamma_n}(\theta_n)\quad\qquad\forall n\ge 0\,.\] Implementing this method yields very good results in practice, but the mathematical theory behind it remains poorly understood, see [1] for an overview. The goal of this manuscript is to shed light on this learning process and to present new research directions in this field.

In 2017, Li, Tai and E opened a new framework for analyzing the SGD learning process in [2], which they examined in more detail in [3]. They proved that the SGD iteration is a first order approximation, in the Euler-Maruyama discretization sense, of an associated stochastic differential equation (SDE). See Section 1.2 for a discussion of this model and related work. This manuscript is devoted to the study of the Fokker-Planck equation for such a continuous SDE model, that is the drift-diffusion PDE for the transition probability.

For this purpose, we will examine the interplay between the drift term governed by the loss function and the degenerate diffusion provided by the randomness in SGD. In the initial stages of the learning process, the dynamics of the parameters primarily follows the drift, tending to concentrate around the local minima of the loss function \(L\). Nevertheless, under mild condition on the degeneracy of the diffusion coefficients, we estimate the training duration required for the neural network to escape from non-optimal local minima. In this manuscript, we present a quantitative analysis of the two regimes of motion in the Fokker-Planck evolution: the drift regime in Section 2 and the diffusion regime in Section 3.

Additionally, in Section 4, we address the problem of the existence of steady states for the parameters distributions and the asymptotic convergence towards them. As a consequence of the high degeneracy of this problem, the differential operator describing the evolution of the transition probability resembles an intermediate case between a transport equation and a diffusion equation, see Section 4.4. By means of simple yet representative examples we analyze the extreme cases to get an idea of the possible scenarios. We provide two types of results. On one hand, we introduce a new variant of the SGD learning process, and exploit the recent results by Porretta [4] about convergence to steady states in non-degenerate Fokker-Planck equations. This will also allow us to prove existence of steady measures for the general case, which to the best of our knowledge was not known. On the other hand, we review some results from Arnold and Erb [5] based on the entropy method (á la Bakry-Émery), and see when they can apply to our situation. We also provide a new entropy computation, which differs from the classical one for non-degenerate Fokker-Planck equations, and represents a (small) novelty in this direction. These results hold only for constant diffusion matrices, hence representing a local scenario around a minimum of the cost function. Finally, in Section 5, we provide a number of open questions and directions for future research.

Beside the interest in the Machine Learning applications, which has motivated this work, our analysis could be also of interest in the study of stochastic differential equations and their long time behaviour.

1.1 Main results↩︎

For more generality, in this paper we will consider the mini-batch Stochastic Gradient Descent. Instead of selecting randomly one sample in each step, we choose randomly a batch of samples \(B_n\) of size \(\mathcalboondox{b}\ge 1\). Namely, the SGD with constant learning rate \(\eta>0\) and batch size \(|B_n|=\mathcalboondox{b}\ge1\) reads as follows \[\label{SGD32Main}\theta_{n+1}=\theta_n-\frac{\eta}{\mathcalboondox{b}}\sum_{b_i\in B_n}\nabla L_{b_i}(\theta_n)\qquad\forall n\ge0\,,\tag{1}\] with \(B_n=\lbrace b_i\rbrace_{i=1}^{\mathcalboondox{b}}\in\lbrace 1,2,\dots,N\rbrace^\mathcalboondox{b}\) being a batch of indexes chosen with uniform distribution in each step \(n\). In [3] and [6], it is deduced that we can approximate 1 with the continuous stochastic process defined by \[\begin{align} \label{SDE32intro} \text{d}X_t=-\nabla L(X_t)\,{\rm d}t+\sqrt{\frac{\eta}{\mathcalboondox{b}}Q(X_t)}\,\text{d}W_t\,, \end{align}\tag{2}\] with \(Q(x) = \tfrac{1}{N}\sum_{i=1}^{N}\nabla L_i(x)\otimes\nabla L_i(x)-\nabla L(x)\otimes\nabla L(x)\) being a nonnegative matrix for any \(x\in\mathbb{R}^d\). Observe that \(Q\) is generically not invertible. Indeed, since we have \[Q(x)=T(x)^*\, T(x)\qquadwith\qquad T(x) = \begin{pmatrix} \frac{1}{N}\nabla L_1(x)-\nabla L(x) \\ \vdots \\ \frac{1}{N}\nabla L_N(x)-\nabla L(x) \end{pmatrix} \in \mathbb{R}^{N \times d}\,,\] then the rows of \(T(x)\) sum zero at every \(x\in\mathbb{R}^d\). Here we have used \(T(x)^*\) to denote the transpose of \(T(x)\). Specifically, in the common case of overparametrization [7], we have that \(Q(x)\in\mathbb{R}^{d\times d}\) satisfies the rank condition rank\((Q(x))\le N-1< d\). As a consequence, the noise in 2 is degenerate.

In addition, the ratio between the learning rate and the mini-batch is usually considered very small, which diminishes the noise. We call this ratio the effective learning rate and we denote it by \[\varepsilon^2:=\frac{\eta}{2\mathcalboondox{b}}\,.\]

One of the advantages of the continuous approximation 2 is that we can use the theory of PDEs to analyze the behaviour of the transition probability and its evolution in time. More specifically, the transition probability \(\rho(t,x)\) associated to the process in 2 with an initial distribution \(\rho_0(x)\) satisfies the following Fokker-Planck type equation: \[\label{FP32intro} \begin{cases} \;\;\partial_t \rho&=\nabla\cdot\bigg(\varepsilon^2\nabla\cdot\left(Q(x)\rho\right)+\rho\nabla L(x)\bigg)\qquadin\qquad(0,\infty)\times\mathbb{R}^d\,,\\ \rho(0,x)&=\rho_0(x)in \qquad\mathbb{R}^d\,, \end{cases}\tag{3}\] where \(\nabla\cdot A\), with \(A\) being a matrix, denotes the divergence taken columnwise. For the deduction of this equation from Ito’s formula see [8]. The class of solutions that we consider lie on the probability space \(\mathcal{P}_k(\mathbb{R}^d)\) of nonnegative normalized Radon measures with \(k\ge 0\) finite moments.

The existence of weak solutions is guaranteed whenever \(Q\) and \(\nabla L\) belong to \(C(\mathbb{R}^d)\), see [9]. For our purposes, the most natural assumption will be slightly stronger, namely we will assume \(L \in C^{1,1}(\mathbb{R}^d)\).

The main challenge of equation 3 is the degenerate diffusion provided by the nonnegative matrix \(Q\). This equation is somehow an intermediate case between a diffusive Fokker-Planck equation and a pure transport equation. In the first case, when \(Q\) is uniformly elliptic, the diffusion regularizes and asymptotic convergence towards the unique steady state holds whenever \(L\) grows at infinity. However, in the latter case, where the diffusion matrix \(Q\) is zero, the mass tends to concentrate around critical points of \(L\). Specifically, the solution will tend to a sum of Dirac’s deltas weighted according to the mass distribution of the initial datum.

We address this problem by focusing on three main questions outlined below. The first question concerns the initial stage of the learning process described in 1 :

Due to the smallness of the effective learning rate \(\varepsilon>0\), we expect that the drift term, given by \(\nabla L\), will lead the dynamics of the parameters for small times. As detailed in Section 2, if the loss function is \(\lambda\)-convex with \(\lambda>0\) in a ball centered at a local minimum \(x_0\in\mathbb{R}^d\), then we aim to obtain that \[\int_{B_{R(t)}(x_0)}\rho(t,x)\,{\rm d}x\ge\int_{B_{R_0}(x_0)}\rho_0(x)\,{\rm d}x-c(\varepsilon,t) \qquad\forall 0<t<T_\varepsilon\,,\] where \(R(t)\) is a decreasing radius and \(0\le c(\varepsilon,t)\xrightarrow[]{\varepsilon\to 0}0\). In order to give a precise statement, we need to introduce the correct speed of decreasing of \(R(t)\) and a suitable approximation of the characteristic function of a ball.

The result above highlights the strong dependence on the initial parameter distribution during the first stage of the 1 . Actually, it indicates that the mass of \(\rho_0\) within the basins of attraction of \(L\) tends to stay concentrated around the local minima for a while.

Once we know that if the effective learning rate is small enough there is a concentration phenomenon around local minima, the next natural question is the following:

For this purpose, we approach the issue from the perspective of the Mean Exit Time (MET) problem associated to 2 . Given a domain \(\Omega\subset\mathbb{R}^d\), the first exit time for the stochastic process \(X_t\) starting at \(x\in\Omega\) is given by \[\tau_\Omega^x=\inf\lbrace t>0:X_t\not\in\Omega,X_0=x\rbrace\,.\] The MET is the function \[u(x)=\mathbb{E}[\tau_\Omega^x]\,,\] which satisfies the following elliptic equation [11] \[\label{Intro32Mean32Exit32Time} \begin{cases} -\mathcal{A}u(x) = 1, & in\Omega \\ u(x)=0, & on \partial \Omega \end{cases}\tag{5}\] where \[\mathcal{A}u(x) = \varepsilon^2tr\left(Q(x)D^2 u(x)\right)-\nabla L(x)\cdot\nabla u(x).\]

We devote Appendices 7 and 8 to recall the connection between MET and PDEs and other classical results in the case of isotropic and non-degenerate diffusion, leading to the well-known Kramers’ Law [12]–[14].

In order to cope with the diffusion matrix \(Q\) in 2 , which is far from being uniformly elliptic, we provide new estimates for the MET associated to our problem. In this case we will consider viscosity solutions to the problem of MET, since this notion allows us to obtain the desired estimates from comparison principles. This notion applies to the problem of the MET as it is the dual problem, or Ornstein-Uhlenbeck problem, of the Fokker-Planck equation, whose solutions are probability distributions and hence weak solutions in the usual distributional sense 4 , or obtained as limits in the sense of measures, as we do for Theorem 6.

The next two Theorems will be proved in Section 3 along with more general results.

Proving an upper bound for a general degenerate matrix \(Q\) is more challenging, since the MET will not be bounded if there is no diffusion. Nevertheless, assuming some non-degeneracy of the diffusion in just one direction is sufficient to provide an upper estimate.

For the validity of the last result in the case of unbounded domains, see Remark 4.

The last question we consider in this manuscript concerns the asymptotic behaviour of the learning process. We aim to characterize and show the convergence of the parameter distribution for large times. Hence, we can state the last question as follows:

Q3: Does the distribution of neural network parameters converge using SGD?

To address this question we implement two different methods relying on the known literature: a duality method and an entropy method.

Duality Method. In the non-degenerate case we exploit the dual equation of the Fokker-Planck equation 3 , which is of Ornstein-Uhlenbeck type. We found that the recent results of Porretta [4] are especially useful in our setting, when the matrix is non-degenerate, and help to construct the stationary solutions by means of non-degenerate approximations. One of the key points in Porretta’s approach is that the asymptotic stabilization of the solution \(\rho\) of 3 for large times, is equivalent to the oscillation decay of the solution of the dual equation. As far as we know, this method has a drawback: it requires the diffusion matrix \(Q\) to be uniformly elliptic. This is the reason why we introduce in Section 4.1 a variant of the 1 that we call Noisy Stochastic Gradient Descent 6 by adding some gaussian noise in each iteration, something that kills the non-degeneracy. Namely,

\[\label{NSGD}\theta_{n+1}=\theta_n-\frac{\eta}{\mathcalboondox{b}}\sum_{b_i\in B_n}\nabla L_{b_i}(\theta_n)+\eta Z_n\qquad\forall n\ge 0\,,\tag{6}\] with \(\lbrace Z_n\rbrace_{n\ge 0}\) being a family of i.i.d. gaussian processes satisfying for some \(\delta>0\) \[Z_n\sim\mathcal{N}(0,\delta I_{d\times d})\qquad\forall n\ge 0\,.\] We will show that, with the same approximation technique from [2] that obtains 2 from 1 , the diffusion matrix of the Fokker-Planck equation associated to 6 is \(Q_\delta(x) = \varepsilon^2 Q(x) + \eta \delta I_{d\times d}\), so it is uniformly elliptic. The results of [4] prove the convergence of the NSGD to a unique stationary measure, see Corollary 18. This convergence is qualitative in several aspects: on the one hand, it only works in non-degenerate cases, which are not realistic in ML applications. On the other hand the convergence rate towards equilibrium cannot be quantified, since it is obtained by a non-constructive proof. However, the above convergence for non-degenerate Fokker-Planck equations, gives us an approximation that allows to prove the existence of stationary measures for the “original” degenerate problem in the limit \(\delta \rightarrow 0\). Indeed, in Section 4.2 we will prove the following theorem.

Entropy Method. In Section 4.4 we discuss in detail the problem of finding the right entropy for the class of Fokker-Planck equations studied in this paper. This is in general a delicate task, and in particular for our setting it requires to modify in a substantial way the common approach. Indeed, for the standard Fokker-Plank equation: \[\partial_t \rho = \nabla\cdot(\nabla \rho + \rho \nabla L)\] it is immediate to see that \(\rho_\infty=e^{-L}\) is a stationary state, since it annihilates the vector field inside the divergence, that is \(\nabla \rho_\infty= - \rho_\infty \nabla L\). Then, a natural choice for the relative entropy would be the standard weighted \(L^2\)-norm: \(\mathcal{E}[\rho|\rho_\infty]=\int_{\mathbb{R}^d} |\rho/\rho_\infty\, -1|^2 \rho_\infty\,{\rm d}x\). Indeed, in this case, standard arguments allows one to deduce that \(\mathcal{E}\) decreases along the flow. However, these arguments exploit in an essential way the fact that \(\rho_\infty\) annihilates the vector field inside the divergence. In our case it is not true in general that the inner vector field vanishes on the stationary state: we can only rely on the full equation \[\nabla\cdot\bigg(\varepsilon^2\nabla\cdot\left(Q(x)\rho_\infty\right)+\rho_\infty\nabla L(x)\bigg) = 0\] in a suitable weak sense. As a consequence, the computation of the entropy production becomes nontrivial. We were able to overcome this difficulty, and we provide in Lemma 21 a proof of the monotonicity of the entropy \(\mathcal{E}\) along the “degenerate flow”, which we consider a significant novelty, to the best of our knowledge. This is the result that, when combined with suitable weighted Poincaré inequalities, provides exponential convergence towards equilibrium. The validity of such a Poincaré inequality for degenerate quantities is a delicate matter, and in really elementary cases² it may simply fail. Hence, this problem presents its own interest, and it requires a thorough study which goes beyond the scope of this paper.

In the particular situations of a diffusion matrix that is degenerate but constant, and of a quadratic loss function \(L(x)=\frac{1}{2}x^TCx\), equation 3 becomes \[\partial_t\rho=\nabla\cdot\left(Q_0\nabla\rho+\rho Cx\right).\] In this case, the Bakry-Émery [17] approach applies, as proven by Arnold and Erb [5], when the following two conditions are satisfied

1. Confinement potential: \(C\) is positive stable, i.e. all eigenvalues have positive real part.

2. Hörmander’s condition: There are no eigenvectors of \(C\) in ker\(\,Q_0\).

The first assumption ensures the \(\lambda\)-convexity of \(L\), while the second ensures that the diffusion is spread in the whole space by the drift.

When condition [II32intro] fails, the stationary state may not have an \(L^1\) density, showing that our existence Theorem 6 is somehow sharp. We shall consider next an example in which exponential convergence in a suitable topology holds even when condition [II32intro] is not met. More precisely, we will prove convergence for the degenerate non-Hörmander case which reads \[\label{Generalized32simple32FP32intro} \partial_t u=\nabla_{x,y}\cdot\left( \begin{pmatrix} Q_0 & 0 \\ 0 & 0 \end{pmatrix} \begin{pmatrix} \nabla_x u \\ \nabla_y u \end{pmatrix} +u \begin{pmatrix} C_0 & 0 \\ 0 & C_3 \end{pmatrix} \begin{pmatrix} x \\ y \end{pmatrix} \right)\tag{7}\] with \(u:(0,\infty)\times\mathbb{R}^n\times\mathbb{R}^{d-n}\rightarrow\mathbb{R}^d\), assuming only that condition [II32intro] holds for \(Q_0\) and \(C_0\), i.e. for the variables that we have called \(x\).

In this framework, we prove the asymptotic convergence of 7 in the 2-Wasserstein distance and the exponential convergence of the second moments to a steady state of the form

where \(g_\infty(x) = c\,e^{-\frac{x^TK^{-1}x}{2}}\) and \(K \in \mathbb{R}^{d \times d}\) is the unique positive definite matrix satisfying \(2Q_0=C_0K+KC_0\).

We will prove an expanded version of this statement, with a quantitative rate \(\lambda\), in Theorem 28.

This may not seem a realistic model for our ML problem. However, this may be considered as a description of the local behaviour of the Fokker-Planck equation 3 , through a linearization of \(Q\) and \(L\) around the local minima of \(L\). Indeed, in Section 5, we exploit this approximation argument by defining, for \(h \in (-1,1)\) small, \[u_h(t,z)=\rho(t,x_0+ hz)\] where \(x_0\in\mathbb{R}^d\) is a local minimum of \(L\). Thus, the equation for the new function \(u_h = u_h(t,z)\) reads \[\label{linearized32eq} \partial_t u_h=\nabla_z\cdot\left(\varepsilon^2Q(x_0)\nabla_z u_h+u_hD^2L(x_0)z\right)\tag{8}\] up to an error of order \(|h|\).

We conclude this manuscript by presenting some open questions concerning the approximation of the steady state of 3 by \[\label{linearized32approximation} \rho_\infty(x)\approx\sum_{i=1}^{M}m_i(\infty)\;u_{i,\infty}\left(\frac{x-x_i}{h}\right)\,,\tag{9}\] where \(u_{i,\infty}\) denotes the steady state of 8 linearized around the local minimum \(x_i\in\mathbb{R}^d\), and \(m_i(\infty)\) represent the mass partition such that \(\sum_{i=1}^{M}m_i(\infty)=1\).

Notation: Let \(I_{d\times d}\) be the identity matrix of dimension \(d\times d\). We denote by \(\mathcal{M}(\mathbb{R}^d)\) the space of probability measures in \(\mathbb{R}^d\). For these measures, we define the norm \(\|\mu\|_{\mathcal{M}_k}=\|\mu\|_{\mathcal{P}_k}:=\int_{\mathbb{R}^d}(1+|x|^2)^{k/2}\text{d}|\mu|(x)\) and the space \(\mathcal{M}_k(\mathbb{R}^d):=\lbrace \mu\in\mathcal{M}(\mathbb{R}^d):\|\mu\|_{\mathcal{M}_k}<+\infty\rbrace\). Analogously, we write \(\mathcal{P}_k(\mathbb{R}^d)\) for the probability measures with \(k\) finite moments.

1.2 Related work on the SGD and the SDE approximation↩︎

Stochastic gradient descent (SGD) has a long history as an optimization technique in machine learning in general, and for neural networks in particular (see e.g. [18], [19]). In recent years, SGD has been the object of several theoretical and experimental studies, considering different aspects of the algorithm, due to its unique mixture of desirable properties of quality of the results, robustness, and generalization capabilities. Special attention has been devoted to the implicit bias, or implicit regularization, in the optimization (for recent reviews see [20], [21]), motivated by seminal works like [22] which show that neural networks are effectively trained by SGD for a number of parameters that exceeds the number of data point (overparametrized), and that results are qualitatively unaffected by explicit regularization. On the other hand, it has been observed that more efficient adaptive optimization methods may generalize worse than SGD [23], [24]. While implicit regularization in gradient descent algorithms can be attributed to multiple factors, involving interactions between the optimization algorithm and the properties of both the parametrization and the dataset [25]–[29], the stochasticity of SGD is believed to play a relevant role, in particular due to its state-dependent noise [30], [31]. Another crucial aspect of SGD design is the dependence of the results on the ratio between learning rate and batch size, also called linear scaling rule, which appears to be a key factor to obtain so-called flat minima, ensuring better generalization [32], [33] (but see also [34]).

The continuous stochastic approximation of SGD with an SDE that we consider in this paper was introduced and proved in [2]. However, the idea of considering a stochastic iterative algorithm as a discretization of a continuous time stochastic process, or, equivalently, that a continuous-time SDE could be obtained as a limiting procedure, is classical (see e.g. [35]). More recently, an SDE model for SGD has been proposed in [36], [37], where the authors assume a gaussian state-independent covariance matrix as a model for a large number of samples, and in [33], where the authors emphasize that different learning rates with the same scaling to batch size give rise to the same continuous approximation and discuss the common practice of assuming the covariance to be approximately equal to the Hessian of the loss (sometimes called label noise). In [38] a different, but still state-independent, covariance matrix is considered, together with a time-dependent learning rate, which allow the authors to prove convergence to a stationary state with similar arguments to the ones used in the present paper for state-dependent diffusion matrices. Here, for the sake of comparison, we prove existence of a stationary state with a degenerate and nonconstant diffusion, and we provide new arguments in terms of entropy methods to attack the problem to convergence to a stationary state for the original SDE model. A learning rate that decreases with time has also been considered in [39], to prove convergence of SGD for convex optimization problems, again in terms of the SDE continuous model. On the other hand, nonconvex scenarios for stochastic learning were considered already in [40], although using a spatially independent nondegenerate diffusion: we recall that, as mentioned after equation 2 , diffusion in the SDE approximation needs to be degenerate and, in the common overparameterized situations, is highly degenerate for dimensional issues. Also, spatial dependence is highly nontrivial in real world scenario, because it contains the different behaviour of each data with respect to the system parameterization. A general approach to the stochastic approximation with SDE of online learning in terms of the theory of semigroups was given in [6]. A different proof of the stochastic approximation of SGD was provided in [41], where the authors show experimentally the limiting behaviors for general state-dependent noise. Key properties of the diffusion that approximates SGD are given in [15], where the authors discuss the limiting behavior and the escape time for non-degenerate diffusion. We can compare in particular our treatment of Question 2 with [15], where the authors obtain a exponential time for the escape from local minima from non spatially homogeneous Fokker-Planck equations by means of large deviations theory. Due to the technique used, they need nondegenerate diffusion, and can deduce the exponential decay only asymptotically for small learning rates. On the other hand, with a different technique, with Theorem 5 we provide a quantitative exit time result that is not asymptotic, and applies to degenerate diffusion, that is the generic situation.. Two regimes in the SDE model of SGD, that dominated by drift, and the diffusive one, were considered in [42] and applied to elementary architectures: these results should be compared directly with our theoretical study in Section 2, where we can obtain general quantitative results of mass concentration around minima in the drift regime.

Two observed features of the SGD that are reproduced by the SDE continuous model are the scaling law, which is intrinsic in the limiting procedure, and the implicit regularization. This second one has been studied for SDE in a long series of works. In [43], under the assumption of the existence and uniqueness of a steady state obtained by a divergence-free force, the authors show that the Fokker-Planck trajectories are monotone minimizer of the KL-divergence with respect to the steady state, hence justifying a Bayesian inference of the SGD, and show that such a state is not a minimum of the loss function itself but rather of a regularization in terms of the Shannon entropy. In [44] the authors assume the covariance of the SDE to be close to the Hessian of the loss function and deduce different escape time estimates from minima depending on the local spectral properties of the Hessian, hence defining an implicit bias towards flat minima. Still modeling the noise covariance with the Hessian of the loss, [45] discuss yet another form of implicit bias in SGD through a continuous SDE model: they show that in the presence of a manifold of minima, SGD does not randomly walks on that manifolds once it reaches it, but rather continues to minimize the trace of the Hessian, hence stabilizing around the flattest region of the minima. They obtain this result by first observing that the dynamics that is normal to the manifold approximately behaves like a constant-covariance Ornstein-Uhlenbeck process, and deduce asymptotically reaching the manifold of minima, and then study the tangent dynamics. The techiques that we introduce that allow us to obtain a rigorous proof of existence of stationary states for the SDE with the true covariance matrix, and a proof of the decay in the non degenerate case, may provide insight on how to consider the tangent dynamics for a general state dependent noise.

Limitations of the SDE model of SGD have been pointed out in several occasions. In [46] the whole approximation procedure is questioned in terms of heuristic arguments concerning the forcing of the scaling law in the limiting procedure, and the difference between the covariance matrix and the more commonly considered Hessian. More recently, [47] have considered the possible deviations of the SDE approximation of SGD for large values of the quotient between learning rate and batch size, by directly comparing trajectories of SGD with a fine discretization of the SDE. Large learning rates are indeed often used in practice to provide an implicit regularization bias [32], [48], [49].

Extensions of SDE modeling of SGD currently include proof of stochastic approximation for adaptive schemes such as ADAM [50], while [51], [52] have considered non-Gaussian noise.

Another notable mathematical approach to modeling the dynamics of neural network parameters trained via SGD relies on mean field limit techniques. When the number of neurons in a layer grows to infinity, the empirical distribution of the network parameters evolves according to a deterministic partial differential equation, expressed as a Wasserstein gradient flow, see for instance the survey [53]. This perspective allows one to bypass the complexities of individual parameter updates and instead analyze the global behavior of the parameter distribution. However, this approach is limited to shallow networks or to deep linear networks, where a well-posed infinite-width limit can be established. In particular, [54] rigorously proves a Law of Large Numbers and a Central Limit Theorem quantifying fluctuations around the mean field trajectory, while [27] shows that in a deep linear network, the infinite-width limit exhibits exponential convergence to a continuous-time limit.

2 Analysis in the drift regime↩︎

The dynamics of the parameters evolving according to the SGD present two different regimes of behaviour. On the first regime, which we call the drift regime, the diffusion is weaker than the drift term \(\nabla L\) and the latter is the one who determines the behaviour of the learning process. This fact leads to a concentration phenomena around the local minima of the loss function \(L\) whenever the effective learning rate \(\varepsilon>0\) is small enough. As we will discuss in Section 3, once the concentration at small time has occurred, the diffusion tends to spread out the parameters. In this section, we are interested in quantifying this concentration behaviour at small times, and addressing the question Q1 formulated in the introduction: How do parameters evolve in the first stage of training with SGD?

To this end, we study the probability density function \(\rho\) of the stochastic process associated to the SGD. Let us recall that \(\rho\) satisfies the equation \[\label{Fokker-Planck322} \begin{cases} \quad\;\,\partial_t\rho=\nabla\cdot\bigg(\varepsilon^2\nabla\cdot (Q(x)\rho)+\rho\nabla L(x)\bigg)\qquadon\quad(0,\infty)\times\mathbb{R}^d\,,\\[4mm] \rho(0,x)=\rho_0(x)in\quad\mathbb{R}^d\,, \end{cases}\tag{10}\] with \(\nabla\cdot A=(\nabla\cdot A_{1\cdot},\dots,\nabla\cdot A_{d\cdot})\) and \(A\in\mathbb{R}^{d\times d}\). In this framework, we will address the above question by obtaining lower bounds for the following quantity: \[\int_{B_R(x_0)}\rho(t,x)\,{\rm d}x\,,\] with \(x_0\in\mathbb{R}^d\) being a local minimum of \(L\) and \(R>0\) being certain radius. Nevertheless, before performing a local analysis, let us show a global concentration estimates when \(L\) is strictly convex in the whole space \(\mathbb{R}^d\). This can be easily deduced from the following upper bound of the second moments of \(\rho\).

Once we have estimated the second moment of \(\rho\), it is possible to obtain a concentration estimate by Chebyshev’s inequality.

. We remark that Lemma 8 and Corollary 9 does not imply the convergence to \(\delta_{x_0}(x)\) in any sense due to the error term \(\varepsilon^2\frac{2\sigma}{\lambda}\). However, this error seems to be unavoidable whenever there is a diffusion term. Let us consider the case \(Q(x)=\frac{\sigma}{d} I\) and \(L(x)=\frac{\lambda}{2}|x-x_0|^2\) for every \(x\in\mathbb{R}^d\) with \(\sigma,\lambda>0\). In this case, the fundamental solution is smooth and reads as \[\rho(t,x)=\frac{1}{\left(2\pi\frac{\varepsilon^2\sigma}{\lambda d}(1-e^{-2\lambda t})\right)^{d/2}}\exp\left(-\frac{\lambda d|x-x_0|^2}{2\varepsilon^2\sigma(1-e^{-2\lambda t})}\right)\,.\] Moreover, we know that \(\rho(t)\rightarrow\rho_\infty\) uniformly in compacts, in 2-Wasserstein distance and in \(L^2(\mathbb{R}^d,\rho_\infty^{-1})\) with \[\rho_\infty(x)= \frac{1}{\left(2\pi\frac{\varepsilon^2\sigma}{\lambda d}\right)^{d/2}}\exp\left(-\frac{\lambda d|x-x_0|^2}{2\varepsilon^2\sigma}\right)\,.\] However, \(\rho_\infty\) does not concentrate in \(x_0\) due to the exponential tails. Indeed, the second moment is \[\int_{\mathbb{R}^d}|x-x_0|^2\rho_\infty(x)\,{\rm d}x=tr\left(\frac{\varepsilon^2\sigma}{\lambda d}I\right)=\frac{\varepsilon^2\sigma}{\lambda}\,.\]

One may think that this is not a representative example since the diffusion is non-degenerate. Nevertheless, let us consider the degenerate case where the operator is not even hypocoercive.

. Assume \(L(x)=\frac{\lambda}{2}|x|^2\) with global minimum at \(x_0=0\) and the diffusion matrix \[Q(x)=\begin{pmatrix} \sigma & 0\dots &0 \\ 0 & 0\dots &0 \\ \vdots &\;\;\;\; \ddots&\vdots\\ 0&\;\;\;\dots&0 \end{pmatrix}\,,\] with \(\sigma,\lambda>0\). Note that this operator is not hypocoercive since \(\lambda I e_i=\lambda e_i\) and \(e_i\in\ker(Q)\) for any \(i=2,\dots,d\). Given an initial datum \(\rho_0\in\mathcal{P}_2(\mathbb{R}^d)\), the solution to \[\partial_t \rho=\nabla\cdot\left(\varepsilon^2Q\nabla \rho+\lambda x\rho\right)\,,\] is the following: \[\rho(t,x)=\frac{e^{\lambda(d-1)t}}{\sqrt{2\pi\frac{\varepsilon^2\sigma}{\lambda}(1-e^{-2\lambda t})}} \; \int_{\mathbb{R}}\exp\left(-\frac{\lambda|x_1-\tilde{x}_1e^{-\lambda t}|^2}{2\varepsilon^2\sigma(1-e^{-2\lambda t})}\right)\;\rho_0\left(\tilde{x}_1,x'e^{\lambda t}\right)\text{d}\tilde{x}_1\,,\] with \(x'=(x_2,\dots,x_d)\in \mathbb{R}^{d-1}\). Due to the tail decay of \(\rho_0\in\mathcal{P}_2(\mathbb{R}^d)\), the stationary solution is \[\rho_\infty(x)=\frac{1}{\sqrt{2\pi\frac{\varepsilon^2\sigma}{\lambda}}}\; \exp\left(-\frac{\lambda|x_1|^2}{2\varepsilon^2\sigma}\right)\;\delta_{0'}(x')\] See Theorem 28 for a 2-Wasserstein convergence of \(\rho(t)\rightarrow\rho_\infty\) when \(t\rightarrow\infty\). Note that even in this very degenerated case the second moment of the invariant measure is given by \[\int_{\mathbb{R}^d}|x|^2\rho_\infty(x)\,{\rm d}x=\frac{1}{\sqrt{2\pi\frac{\varepsilon^2\sigma}{\lambda}}} \int_{\mathbb{R}}|x_1|^2\exp\left(-\frac{\lambda|x_1|^2}{2\varepsilon^2\sigma}\right)\,{\rm d}x_1=\frac{\varepsilon^2\sigma}{\lambda}\]

Based on this simple example, we aim to estimate this error that the diffusion generates in the concentration phenomena around local minima. As in the example, the main influent parameters are the \(\lambda\)-convexity of the loss function in the basin of attraction, the effective learning rate \(\varepsilon>0\) and the \(L^\infty\)-norm of the matrix \(Q\), which quantify the diffusion. From an application perspective, these parameters establish the likelihood that the weights of the neural network will fall into the nearest local minimum.

2.1 Local mass concentration↩︎

In this section, we show that in the first steps of the learning process, the parameters tends to concentrate with high probability around the local minima of the loss function. Assuming without loss of generality that \(0\) is a local minimum of \(L\), we would like to obtain an estimate on the following quantity, for small values of \(r\) \[\label{local32mass} \mathbb{P}\lbrace X_t\in B_r(0)\rbrace=\int_{B_r(0)}\rho(t,x)\,{\rm d}x\,.\tag{11}\] The main problem is that the variation of this quantity depends on the flux on the boundary of \(B_r(0)\) which is a priori unknown. Namely, if we try to differentiate this quantity using the equation of the probability density function, we obtain \[\begin{align} \frac{\text{d}}{\,{\rm d}t}\int_{B_r(0)}\rho(t,x)\,{\rm d}x& =\int_{B_r(0)}\nabla \cdot\bigg(\varepsilon^2\nabla\cdot(\rho Q)+\rho\nabla L\bigg)\,{\rm d}x\\ &=\int_{|x|=r}\bigg(\varepsilon^2\nabla\cdot(\rho Q)+\rho\nabla L\bigg)\cdot \frac{x}{|x|}\,{\rm d}x \end{align}\] Since this variation is complicated to estimate with 11 to close the differential inequality, we will consider a different quantity with a controllable variation. To motivate this new approach, let us consider the following continuity equation (\(\varepsilon=0\)) \[\label{transport} \begin{cases} \partial_t \mu=\nabla\cdot\left(\mu\nabla L\right)\\ \mu(0,x)=\mu_0(x) \end{cases}\tag{12}\] with \(L\) being \(\lambda\)-convex with \(\lambda>0\) and a global minimum at \(0\). It is well known that if \(\mu_0\in\mathcal{P}_2(\mathbb{R}^d)\), then the solution \(\mu(t)\) converge to \(\delta_0\) in the Wasserstein topology [55], but we are interested in the speed of concentration. Moreover, the associated flow dynamic is \[\begin{cases} \frac{\text{d}}{\,{\rm d}t} \Phi_t(x)&=-\nabla L(\Phi_t(x)) \\ \Phi_0(x)&=x. \end{cases}\] and the solution of 12 is the pushforward \(\mu(t,x)=(\Phi_t)_\#\mu_0(x)\). In order to show the concentration velocity of \(\Phi_t(x)\), let us compute the variation in time of the Euclidean distance from \(\Phi_t(x)\) to the global minimum at 0: \[\begin{align} \frac{\text{d}}{\,{\rm d}t}|\Phi_t(x)-0|^2 & =2 \frac{\text{d}}{\,{\rm d}t}\Phi_t(x)\cdot\left(\Phi_t(x)-0\right)=-2\nabla L(\Phi_t(x))\cdot \left(\Phi_t(x)-0\right) \\ &\le- \lambda |\Phi_t(x)-0|^2\,, \end{align}\] which implies that \[\label{flow32map32convergence} |\Phi_t(x)|\le e^{-\frac{\lambda}{2}t}|x|\qquad\forall t>0.\tag{13}\] Hence, the flow tends to concentrate particles starting from a point \(x\in\mathbb{R}^d\) around \(0\) with exponential velocity. This phenomena can be understood in the sense of the following lemma.

Considering the diffusion equation 10 for \(\rho\). Compared with the result in the lemma above, we expect that the probability of a suitable shrinking ball should be non decreasing up to an error depending on \(\varepsilon\) and before certain time. Namely, our goal is to obtain a lower bound in terms of the initial datum \(\rho_0\) of the form \[\int_{B_{R(t)}(0)}\rho(t,x)\,{\rm d}x\ge \int_{B_{R_0}(0)}\rho_0(x)\,{\rm d}x-c(\varepsilon,t)\] with an appropiate shrinking radius \(R(t)\) and \(c(\varepsilon,t)\) tending to 0 as \(\varepsilon\) goes to 0. Since we want to use the equation of \(\rho\) and test it with the characteristic function \(\chi_{B_{\mathbb{R}(t)}(0)}(x)\), we use the following \(C^{1,2}\) approximation cut-off function \[\label{smooth32characteristic} \varphi(t,r)=\begin{cases} 1, & if\qquad\qquad\qquad\; r\le R(t) \\ 1-\frac{2}{\delta^2R(t)^2}(r-R(t))^2, & if \quad\qquad R(t)<r\le(1+\frac{\delta}{2})R(t) \\ \frac{2}{\delta^2 R(t)^2}((1+\delta)R(t)-r)^2, & if(1+\frac{\delta}{2})R(t)<r\le(1+\delta)R(t)\\ 0, & if(1+\delta)R(t)<r\,. \end{cases}\tag{14}\] for some radius \(R(t)\) to be chosen later and some small constant \(\delta>0\). In terms of this test function, the result reads as follows.

Note that the shrinking condition ?? of the radius \(R(t)\) coincide with the concentration velocity of the flow map 13 associated to the continuity equation 12 . The main idea is that if the radius decreases slower than the flow map concentrates, the mass in the shrinking balls will be almost preserved, up to an error depending on \(\varepsilon\).

Despite the error term not being sharp, it is impossible to prove the estimate ?? with \(C=0\) due to the diffusive nature of the equation of \(\rho\). Indeed, if we consider the case \(Q(x)=I\) in \(\mathbb{R}^2\) with a double well potential \(L(x,y)=(x^2-1)^2+y^2\), it is easy to see that even if the mass of the initial datum is concentrated in one of the local minima, half of the mass will reach the neighbourhood of the other local minimum. A simple proof can be done studying the local mass of the unique steady state \(\rho_\infty(x,y)=c \exp(-\frac{L(x,y)}{\varepsilon^2})\).

Now, we use Theorem 11 with a decreasing radius satisfying \(R'(t)\ge-\frac{\lambda}{2}R(t)\) in order to obtain an explicit control of the error term. Thus, a natural choice \(R(t)\) is the exponentially decreasing radius, which we have considered previously in the non diffusive case in Lemma 10 and satisfies condition ?? with equality.

The above result gives us a lower bound for the mass concentrated around a ball centered in the local minimum of \(L\). We are now in the position to prove Theorem 3.

Letting \(T=T(\varepsilon, C)\) in ?? , with \(T(\varepsilon, C) = \frac{1}{\lambda}\log\left(1 +\frac{R_0^2 \lambda}{C \varepsilon^{2 \alpha}}\right)\), gives that for every \(0<t<T_\varepsilon\) we have \[\begin{align} \int_{\mathbb{R}^d}\varphi(t,|x|)\rho(t,x)\,{\rm d}x &\ge \int_{\mathbb{R}^d}\varphi(0,|x|)\rho_0(x)\,{\rm d}x-\varepsilon^2\frac{C}{R_0^2\lambda}\left(e^{\lambda T_\varepsilon}-1\right)\\ &=\int_{\mathbb{R}^d}\varphi(0,|x|)\rho_0(x)\,{\rm d}x-\varepsilon^{2-2\alpha}. \end{align}\] Finally, we notice that we can replace \(t=0\) by \(t=t_0\) by the time shift invariance of the equation. 0◻

3 Analysis in the diffusion regime↩︎

After a first concentration phenomena near local minima of the loss function \(L\), the diffusion will allow the parameters of the optimization to jump from one local minimum of \(L\) to another one. This property of the diffusion is the main advantage of using the SGD as a learning process with respect to the classical Gradient Descent, since it allows to avoid local minima of \(L\) which could be attractors at first steps. Thus, the main question that we want to address in this section is what we formulated as Q2 in the introduction: How long does SGD take to escape from a local minimum?

To this end, we focus on how the diffusion affects the dynamics of the learning process. In particular, we want to estimate the mean exit time from a neighborhood of a local minimum of the loss function. Lower bounds for the mean exit time can be obtained without further assumption. On the other hand, some non-degeneracy condition on the diffusion is needed in order to ensure that the mean exit time is finite. Indeed, as a counterexample, consider the case with no diffusion at all, presented in 12 , which leads to the deterministic flow map \[\begin{cases} \frac{\text{d}}{\,{\rm d}t}\Phi_t(x)=-\nabla L(\Phi_t(x))\\ \Phi_0(x)=x\,. \end{cases}\] In this case, all the mass splits and concentrates in the local minima of \(L\), with no jump from one minimum to another one, i.e. the mean exit time is infinite.

Recall that the stochastic process describing the dynamics of the parameters when learning with SGD satisfies 2 , that is, \[\begin{cases} \text{d}X_t&=-\nabla L(X_t)\,{\rm d}t+\sqrt{2\varepsilon^2 Q(X_t)}\text{d}W_t\\ X_0&=x\,, \end{cases}\] with the degenerate but nonnegative matrix \(Q(x)\ge 0\), i.e. \(v^TQ(x)v\ge 0\) for every \(v\in\mathbb{R}^d\) and every \(x\in\mathbb{R}^d\). Due to the degeneracy of \(Q\), obtaining estimates on the Mean Exit Time (MET) from a basin of attraction of \(L\) is not obvious. Recall also that the MET \(u(x)=\mathbb{E}[\tau_\Omega^x]\) from an open set \(\Omega \subset \mathbb{\mathbb{R}}^d\) given an initial point \(x \in \Omega\) satisfies 5 , that is \[\begin{cases} -\mathcal{A}u(x) = 1, & in\Omega \\ u(x)=0, & on \partial \Omega \end{cases}\] where \[\mathcal{A}u(x) = \varepsilon^2tr\left(Q(x)D^2 u(x)\right)-\nabla L(x)\cdot\nabla u(x).\]

This family of equations contains, as special cases, the non diffusive scenario (\(Q(x)\equiv0\), i.e. pure drift) and the non-degenerate diffusive scenario (\(Q(x)\ge \delta I_{d\times d}\)).

In the non-diffusive case, the mean exit time is \(u(x)=+\infty\) for all \(x\in\Omega\). Indeed, the flow map \(X_t\) is deterministic and describes the concentration of the parameters around the local minima of the loss function \(L\). Hence, there is no possibility for \(X_t\) to escape from a neighbourhood of a local minimum of \(L\) in any time.

In the non-degenerate diffusive case, by Kramers’ Law, \(u(x)<+\infty\) for every \(x\in\Omega\). Indeed, even if the drift term \(\nabla L\) pushes \(X_t\) into the local minimum, the diffusion in every direction ensures the boundedness, see Appendix 8.

In this section, we provide lower and upper bounds for the viscosity solution \(u\) to 5 under mild conditions on \(Q\) and \(L\), allowing for degeneracies.

Our strategy is based a variation of the comparison principle given in [10], for viscosity solutions of 5 . In order to prove it, we adapt the arguments of [10] to our setting, following the indications in [10]. Before stating the result, we rewrite 5 in the notation of Definition 2, by setting \[F\left(x,\nabla u(x),D^2u(x)\right) = -\mathcal{A}u(x) - 1 = -\varepsilon^2tr\left(Q(x)D^2u(x)\right)+\nabla L(x)\cdot\nabla u(x) - 1\,.\] The comparison principle that we need is the following.

In order to obtain an upper bound for the Mean Exit Time, some non-degeneracy in the diffusion is required. Indeed, it suffices to have a positive lower bound of \(Q\) in one direction and some smoothness condition in \(L\).

Thanks to the semigroup property of the Markovian process described by 2 , we can use the subsolution and supersolution of Theorems 14 and 15 starting at any time \(t_0>0\). This allows us to estimate the time that the learning process will spend after the concentration phenomena occurred in the drift regime.

We are now in the position to prove Theorem 4 and 5. Note that, in order to make use of the classical comparison principle [10] in the next proofs, we are requiring that the mean exit time is a viscosity solution of 5 .

By the degenerate comparison principle of Theorem 13, Part i), we obtain that \[u(x)= \mathbb{E}\left[\tau_{B_R(0)}^x\right]\ge w(x)=\frac{R^2 -|x|^2}{2\varepsilon^2\sigma d} \,,\] since \(u\) is a solution of 5 by construction and \(w\) is a subsolution of 5 by Theorem 14. Note that the Markovian property of the stochastic process 2 allows the time shifting of the initial datum. We conclude recalling that \(|x|\le r\). 0◻

By a similar argument we can obtain an upper bound for the mean exit time.

Proof of Theorem 5 Since \(u\) is a solution of 5 by construction and \(w\) is a supersolution of 5 by Theorem 15, by the degenerate comparison principle of Theorem 13, Part ii), we obtain that \[u(x)= \mathbb{E}\left[\tau_{\Omega}^x\right]\le w(x)=\frac{1}{\Lambda}\left(e^{\frac{\Lambda R_v^2}{2\beta\varepsilon^2}}-e^{\frac{\Lambda (v\cdot x)^2}{2\beta\varepsilon^2}}\right).\] To complete the proof, it only remains to estimate the constants \(c, \Lambda, R_v\). We recall that \(\Omega \subseteq B(0,R_\Omega)\), and that \(|v| = 1\). Let \(c = R_\Omega + 1\). Then \[\frac{|\nabla L(x)\cdot v|}{|v\cdot x + c|} \leq \frac{M}{\min_{x \in \Omega}|v\cdot x + c|} \leq M\] where the last inequality follows by the fact that \(|v\cdot x + c| \geq |c| - |x| \geq 1\). Thus, setting \(\Lambda = M\), and noting that \(|v\cdot x + c| \leq 2(R_\Omega + 1)\), we get \[ \mathbb{E}\left[\tau_{\Omega}^x\right] \le \frac{1}{M}\left(e^{\frac{2M(R_\Omega + 1)^2}{\beta\varepsilon^2}} - 1\right).\qed\]

4 Asymptotic behaviour of the SGD↩︎

In this section, we analyze the asymptotic behaviour of the distribution of parameters optimized by SGD. The main motivation is neural network training: starting from an initial parameter configuration and a dataset, we aim to determine the final parameter probability distribution after prolonged SGD implementation. We will study this problem within the framework of the continuous stochastic process that approximates the SGD and its associated probability measure. Here we will address question Q3 from the introduction: Does the distribution of neural network parameters converge using SGD?

To answer this question is delicate, and we propose two different methods: duality and entropy methods. Both methods provide partial results about the convergence of the transition probability of the stochastic approximation to an invariant probability measure.

The duality method for a wide class of nondegenerate Fokker-Planck equations was introduced by Porretta in 2024 [4]. The first step consists of proving the time decay of oscillations of the solutions of the dual equation, often referred to as Ornstein-Uhlenbeck or backward Kolmogorov. As a second step, this decay can be translated into a decay of the moments of the probability measure that solves the original Fokker-Planck equation. Finally, the third step consists of showing that the convergence of the moments implies the existence of an invariant probability measure and the convergence to it. The main drawback of this method is that, to the best of our knowledge, it only works with nondegenerate diffusion, i.e. \(0<\delta I\le Q\le \sigma I\) with \(\sigma\geq\delta>0\). In order to adapt this method to the present setting, we propose a variant of the SGD adding some random noise at each step of the learning process.

A nowadays more standard approach to study the asymptotic behavior of Fokker-Planck equations is the entropy method [17], [56], also known as the Bakry-Emery method. The relative entropy - a special Lyapunov function - quantifies the distance between the probability at time \(t\) and the invariant probability measure. The first stability property, i.e. convergence to an invariant probability measure, can be proven by showing that the relative entropy decays to 0 as time goes to infinity. A more delicate analysis is then required to quantify the possible decay rates of the entropy. In general, even for the standard Fokker-Planck equation (\(Q = I\) and \(L(x) = |x|^2\)), exponential decay holds only under additional assumptions, typically the boundedness of the first and second moment for the initial data (see e.g. [57]). A quantitative exponential decay of the relative entropy can be proven by means of a weighted Poincaré type inequality which in turn is equivalent to the entropy-entropy production inequality.

In the present setting, the exponential decay follows by establishing a suitable Poincaré inequality obtained by using the results of [5] and [58]. The main limitations of this approach are that either the diffusion matrix must be non-degenerate with \(Q\equiv \sigma I\) or, when \(Q\) is allowed to be degenerate, the loss function \(L\) must be quadratic.

4.1 Duality method for Noisy SGD↩︎

Recall that, for a data set \(\lbrace (x_i,y_i) \rbrace_{i=1}^N\) and a total loss function \[L(\theta)=\frac{1}{N}\sum_{i=1}^{N}L_i(\theta)\,,\] where each \(L_i(\theta)\) for \(i=1,\dots,N\) is the partial loss function associated to the datum \((x_i,y_i)\), the SGD with constant learning rate \(\eta>0\) and batch size \(|B_n|=\mathcalboondox{b}>0\) reads as in 1 \[\theta_{n+1}=\theta_n-\frac{\eta}{\mathcalboondox{b}}\sum_{b_i\in B_n}\nabla L_{b_i}(\theta_n)\qquad\forall n\ge 0\,,\] with \(B_n=\lbrace b_i\rbrace_{i=1}^{\mathcalboondox{b}}\subset\lbrace 1,2,\dots,N\rbrace\) being a batch of indexes chosen with uniform distribution at each step \(n\). Its SDE counterpart, from [3], as defined by 2 , reads \[\begin{align} \text{d}X_t=-\nabla L(X_t)\,{\rm d}t+\sqrt{\frac{\eta}{\mathcalboondox{b}}}\Sigma(X_t)\text{d}W_t\,, \end{align}\] where the noise is always degenerate, because rank\((\Sigma(x)) \le N -1\) (and, in overparameterized scenarios, the degeneracy is more severe).

In order to obtain a nondegenerate SDE, we propose another learning iteration with a suitable extra noise \(\lbrace Z_n\rbrace_{n\ge 0}\), which we have called Noisy Stochastic Gradient Descent in 6 . This is \[\theta_{n+1}=\theta_n-\frac{\eta}{\mathcalboondox{b}}\sum_{b_i\in B_n}\nabla L_{b_i}(\theta_n)+\eta Z_n.\] In what follows, our choice for \(\lbrace Z_n\rbrace_{n\ge 0}\) is that of a family of i.i.d. Gaussian variables \[Z_n\sim\mathcal{N}(0, \delta I_{d\times d})\qquad\forall n\ge 0 \, .\] This extra noise does not depend on the characteristics of the learning process, and a priori may provide an invariant probability measure that is independent of the data set and the loss function. However, more sophisticated choices of noise could be considered, based on the loss function. We have chosen the simplest noise that allows to apply the results of [4], in the form of Theorem 17 below, indeed, our choice of noise leads us to the following nondegenerate SDE.

By standard theory, we can associate to ?? the following Ornstein-Uhlenbeck type equation for \(u(t,x) = \mathbb{E}[X_t]\) \[\label{O-U} \begin{cases} \partial_t u=\frac{\eta}{2}tr(Q_\delta(x)D^2u)-\nabla L(x)\cdot\nabla u=:-\mathcal{L}_\delta u\\ u(0,x)=u_0(x)\,, \end{cases}\tag{20}\] Analogously, the probability density function of \(X_t\) satisfies the \(L^2(\mathbb{R}^d)\)-dual equation of 20 , that is, the Fokker-Planck equation \[\label{F-P} \begin{cases} \partial_t \rho =\nabla\cdot\left(\frac{\eta}{2}\nabla\cdot(\rho\, Q_\delta(x))+\rho\nabla L(x)\right)=:-\mathcal{L}^*_\delta \rho\\ \rho(0,x)=\rho_0(x)\,, \end{cases}\tag{21}\] where \(\nabla\cdot(\rho \,Q_\delta(x))\) is the divergence taken column-wise.

The goal is to prove convergence results for the distribution of the parameters \(\theta\) in 6 . In order to do that, we prove that there exists a stationary state of equation 21 and the solution \(\rho\) converges to it as time goes to infinity. We are in the position to use the theory of [4] and establish the asymptotic behavior of the above non-degenerate Ornstein-Uhlenbeck and Fokker-Planck type equations, by exploiting the duality between 20 and 21 , and proving that the weighted oscillation of \(u\) implies convergence of the moments of \(m\). The weighted oscillation that we will consider is the following \[[u]_{\langle\cdot\rangle^k} = \sup\limits_{x,y\in\mathbb{R}^d}\frac{|u(x)-u(y)|}{\langle x\rangle^k+\langle y\rangle^k}\,,\] with \(\langle x\rangle^k=(1+|x|^2)^{k/2}\). This turns out to be a Lyapunov function of \(\mathcal{L}_\delta\) (see [4]).

The above theorem is written for zero mean valued solutions. As shown in [4], it can be easily extended to the case of non-zero mean valued solutions with finite \(k\)-moments: the linearity of the equation allows to show both existence of a stationary solution \(\rho_\infty^\delta\), and to show the convergence \(\rho(t)\xrightarrow[]{t\to\infty}\rho_\infty^\delta\) in appropriate topologies, noticing that the mass of both \(\rho\) and \(\rho^\delta_\infty\) need to be the same.

We conclude this section by recalling a technical lemma of [4], used in the proof of the above results, that we will need in what follows.

4.2 Existence of steady states: proof of Theorem 6↩︎

In this section we prove, under mild condition, the existence of at least one stationary measure for our Fokker-Planck type equations, as stated in Theorem 6, whose statement we recall here.

Consider the degenerate Fokker-Planck equation \[\partial_t \rho=\nabla\cdot\bigg(\frac{\eta}{2\mathcalboondox{b}}\nabla\cdot(\rho Q(x))+\rho\nabla L(x)\bigg)=:-\mathcal{L}^*\rho\,.\] If ?? , ?? and ?? hold, then there exists an invariant probability measure \(\rho_\infty\in\mathcal{P}(\mathbb{R}^d)\) such that \[\label{Stat46Eq} \int_{\mathbb{R}^d}\mathcal{L}\varphi(x)\text{d}\rho_\infty(x)=0\qquad\quad \forall\varphi\in C_b(\Omega)\,.\qquad{(20)}\] The idea of the proof is to consider the invariant measures of the nondegenerate equation with diffusion matrix \(Q_\delta(x)=\mathcalboondox{b}^{-1}Q(x)+\delta I_{d\times d}\) and apply Prokhorov’s Theorem. By Corollary [Asymptotic32convergence32nondegenerate], we know that for each \(\delta>0\) there exists \(\rho^\delta_{\infty}\in\mathcal{P}(\mathbb{R}^d)\) such that \[-\mathcal{L}_\delta^*\rho_\infty^\delta=\nabla\cdot\left(\frac{\eta}{2}\nabla\cdot(\rho^\delta_\infty Q_\delta(x))+\rho^\delta_\infty \nabla L(x)\right)= 0\,.\] In order to prove that the sequence \(\lbrace \rho^\delta_\infty\rbrace_{\delta>0}\) is tight we will use \(\langle x\rangle^k\) as a Lyapunov function. Notice that for any \(\varepsilon\in(0,1)\), by Lemma 19, it holds that \[\mathcal{L}_\delta\langle x\rangle^k\ge k\langle x\rangle^{k-2}(\alpha|x|^\gamma-\bar K)\,,\] with \(\bar K:=d(\sigma+1)(d+k-2)\). Hence, we can find two nonnegative functions \(\phi,\chi\in C^2(\mathbb{R}^2)\) such that \(\lim\limits_{|x|\rightarrow\infty}\phi(x)=+\infty\), \(\chi\) is compactly supported and \[\label{epsilon32invariant32measure} \mathcal{L}_\delta\langle x\rangle^k\ge\phi(x)-\chi(x)\qquad\quad\forall\delta\in(0,1)\,.\tag{22}\] Now, let us prove that \(\lbrace \rho_\infty^\delta\rbrace_{0<\delta<1}\) is tight, that is, for any \(\varepsilon>0\) there exists \(U_\varepsilon\subset\mathbb{R}^d\) such that \(\int_{U_\varepsilon}\text{d}\rho_\infty^\delta(x)>1-\varepsilon\) for every \(\delta\in(0,1)\). Using \(\mathcal{L}_\delta^*\rho_\infty^\delta=0\) and inequality 22 , we obtain \[\begin{align} 0&=\int_{\mathbb{R}^d}\langle x\rangle^k\mathcal{L}_\delta^*\rho^\delta_\infty\,{\rm d}x=\int_{\mathbb{R}^d}\mathcal{L}_\delta\langle x\rangle^k \text{d}\rho^\delta_\infty(x)\ge\int_{\mathbb{R}^d}(\phi(x)-\chi(x))\text{d}\rho^\delta_\infty(x)\,, \end{align}\] and hence \[0\le\int_{\mathbb{R}^d}\phi(x)\text{d}\rho^\delta_\infty(x)\le \int_{\rm supp(\chi)}\chi(x)\text{d}\rho_\infty^\delta(x)\le \|\chi\|_\infty\qquad\forall\delta\in(0,1)\,.\] On the other hand, since \(\phi(x)\rightarrow+\infty\) when \(|x|\rightarrow\infty\), for any \(n>0\) there exists \(R_n>0\) such that \(\phi(x)\ge n\) for \(x\in B_{R_n}\). Thus, \[\begin{align} \int_{B^c_{R_n}}n\,\text{d}\rho_\infty^\delta\le\int_{B^c_{R_n}}\phi(x)\,\text{d}\rho_\infty^\delta\le\|\chi\|_\infty\qquad\Rightarrow\qquad\int_{B^c_{R_n}}\text{d}\rho_\infty^\delta\le \frac{\|\chi\|_\infty}{n}\,, \end{align}\] which implies \[\begin{align} \int_{B_{R_n}}\text{d}\rho_\infty^\delta>1-\frac{\|\chi\|_\infty}{n}\qquad\quad\forall n>0\,. \end{align}\] Therefore, we can apply Prokhorov’s Theorem to obtain that there exists \(\rho_\infty\in\mathcal{P}(\mathbb{R}^d)\) such that \(\rho_\infty^\delta\rightharpoonup\rho_\infty\) in \(\mathcal{P}(\mathbb{R}^d)\) up to a subsequence.

It remains to prove that \(\rho_\infty\) solves the equation. For any \(\varphi\in C_c^\infty(\mathbb{R}^d)\), it holds that \[\begin{align} 0&\le\left|\int_{\mathbb{R}^d}\mathcal{L}\varphi(x)\text{d}\rho_\infty(x)\right|=\left|\int_{\mathbb{R}^d}\mathcal{L}\varphi(x)\text{d}\rho_\infty(x) -\int_{\mathbb{R}^d}\mathcal{L}_\delta\varphi(x)\text{d}\rho_\infty^\delta(x)\right|\\ &\le\left|\int_{\mathbb{R}^d}\mathcal{L}\varphi(x)\;\text{d}(\rho_\infty-\rho_\infty^\delta)(x)\right| +\left|\int_{\mathbb{R}^d}(\mathcal{L}\varphi(x)-\mathcal{L}_\delta\varphi(x))\text{d}\rho_\infty^\delta(x)\right|\\ &\le\langle\mathcal{L}\varphi,\rho_\infty-\rho_\infty^\delta\rangle_{C_b(\mathbb{R}^d),\mathcal{M}(\mathbb{R})}+\|\mathcal{L}\varphi-\mathcal{L}_\delta\varphi\|_{C_b(\mathbb{R}^d)} \stackrel{\delta\rightarrow0^+}{\longrightarrow}0\,, \end{align}\] which concludes the proof, since \(\|\mathcal{L}\varphi-\mathcal{L}_\delta\varphi\|_{C_b(\mathbb{R}^d)}\le C\,\delta\|\varphi\|_{C^2(\mathbb{R}^d)}\) , for some \(C>0\) . 0◻

Existence of stationary states -possibly measures- for quite general Fokker-Plank equations nowadays is a classical topic, we refer for instance to Chapter 2 of [9] and references therein. In particular, similar results, which do not apply directly to our specific case, have been obtained in [9], with slightly different assumptions and proofs. Uniqueness is in general a delicate issue and in general may turn out to be false.

4.3 The question of convergence to stationary measures↩︎

Once existence of stationary measures is established, the next natural question would be which solutions converge to it (for large times), and this will be explored in the next section. In general we cannot expect better than convergence in a topology compatible with measures, since stationary solutions may happen to be purely measures (not functions), in view of the Examples 6 and 7 below, see also Example 8 in the next section. Also, since the diffusion matrix can be highly degenerate, so that, heuristically, part of the flow can be driven by “pure transport”, in which case, the stationary solution can be a “pure measure”. This motivates the following example.

Let us consider the following transport equation \[\label{Example32Transport} \begin{cases} \partial_t\rho(t,x)\!\!\!\!!&=\nabla\cdot\left(Cx\rho(t,x)\right)\qquadin\quad(0,\infty)\times\mathbb{R}^d\,,\\ \rho(0,x)&=\rho_0(x) \qquad\quad\qquad\qquad\;\,in\quad\mathbb{R}^d\,, \end{cases}\tag{23}\] where \(C\in \mathbb{R}^{d\times d}\) with \(C>0\). Notice that \(\rho_\infty(x)=\delta_0(x)\) is a stationary measure since if we multiply by a test function \(\varphi\in C_c^1(\mathbb{R}^d)\) and we integrate, it holds that \[\int_{\mathbb{R}^d}x\cdot\nabla\varphi(x)\,\text{d}\delta_0(x)=x\cdot\nabla\varphi(x)\bigg|_{x=0}=0\qquad\qquad\forall\varphi\in C_c^1(\mathbb{R}^d)\,.\] The equation above has the following explicit solution \(\rho(t,x)=e^{ tr(C)t}\rho_0(e^{Ct}x)\), since \[\partial_t\rho(t,x)=tr(C)e^{ tr(C)t}\rho_0(e^{Ct}x)+e^{ tr(C)t}\left(\left(Cx\right)\cdot \nabla \rho_0(e^{Ct}x)\right)=\nabla\cdot\left(Cx\rho(t,x)\right)\,.\] Notice that the solution \(\rho(t,x)\) preserves the mass: using the change of variable \[y=e^{Ct}x\qquad and\qquad \,{\rm d}y=\left|\det{e^{Ct}}\right|\,{\rm d}x=e^{ tr(C)t}\,{\rm d}x\,,\] we obtain \[\int_{\mathbb{R}^d}\rho(t,x)\,{\rm d}x=\int_{\mathbb{R}^d}e^{ tr(C)t}\rho_0(e^{Ct}x)\,{\rm d}x=\int_{\mathbb{R}^d}\rho_0(y)\,{\rm d}y\qquad\forall t>0\,.\]

Let us show the convergence of the solutions of 23 to \(\delta_0\) whenever \(\rho_0\in\mathcal{P}_2(\mathbb{R}^d)\). Let us consider the equation above as a gradient flow for the second moments in the Wasserstein space, namely, \[\begin{cases} \partial_t\rho(t,x)\!\!\!\!!&=-grad_{W_2}\mathcal{E}[\rho]:=\nabla\cdot\left(\rho(t,x)\nabla\left(\frac{\delta \mathcal{E}}{\delta \rho}[\rho]\right)\right)\\ \rho(0,x)&=\rho_0(x)\in\mathcal{P}_2(\mathbb{R}^d)\,. \end{cases}\] with \[\mathcal{E}[\rho]:=\frac{1}{2}\int_{\mathbb{R}^d}x^T\,C^TC\,x\;\text{d}\rho(x)\qquad\quad\&\quad\qquad\frac{\delta \mathcal{E}}{\delta \rho}[\rho]=Cx\,.\] Notice that the global minimum of the energy is attained at \(\mathcal{E}[\delta_0]=0\). As it can be seen in [59] since \(\mathcal{E}\) is \(\lambda\)-convex with some \(\lambda>0\) (in particular, for any \(\lambda\in(0,\lambda_1(C^TC)\)), it holds that for every \(\rho\in \mathcal{P}_2(\mathbb{R}^d)\) \[\frac{\lambda}{2}W^2_2(\rho,\delta_0)\le\mathcal{E}[\rho]-\mathcal{E}[\delta_0]\le\frac{1}{2\lambda}\langlegrad_{W_2}\mathcal{E}[\rho],grad_{W_2}\mathcal{E}[\rho]\rangle_{\rho}\,.\] This estimates above for the \(\lambda\)-convex energy are the analogous of \[\varphi(x)-\varphi(x_0)\ge \frac{\lambda}{2}|x-x_0|^2\qquad\&\qquad|\nabla \varphi(x)|^2\ge2\lambda\left(\varphi(x)-\varphi(x_0)\right)\,,\] whenever \(x_0\in\mathbb{R}^d\) is the unique minimum of a \(\lambda\)-convex function \(\varphi\in C^1(\mathbb{R}^d)\).

Hence, following equation (4.18) of [59] we obtain \[\frac{\text{d}}{\,{\rm d}t}\mathcal{E}[\rho]=\langlegrad_{W_2}\mathcal{E}[\rho],\partial_t\rho\rangle_{W_2} =-\langlegrad_{W_2}\mathcal{E}[\rho],grad_{W_2}\mathcal{E}[\rho]\rangle_{\rho}\le -2\lambda\mathcal{E}[\rho]\,,\] which implies the decay of the energy \[\mathcal{E}[\rho]\le e^{-2\lambda t}\mathcal{E}[\rho_0]\,.\] Indeed, since \(\mathcal{E}[\delta_0]=0\), we have \[\frac{\lambda}{2}W^2_2(\rho,\delta_0)\le\mathcal{E}[\rho]-\mathcal{E}[\delta_0]=\mathcal{E}[\rho]\le e^{-2\lambda t}\mathcal{E}[\rho_0]\,.\] Thus, there is exponential Wasserstein convergence to the Dirac’s delta at zero. In conclusion, if we assume a condition on the decay of the initial datum \(\rho_0(x)\) when \(|x|\rightarrow\infty\), such as finite second moments, then the solution converge to \(\delta_0\).

A similar phenomena can also happen in the case of non-zero diffusion matrix (still very degenerate), as the next example shows. See also Example 8 in the next section for a generalization to higher dimensions.

Let us analyze the fundamental solution of the following degenerate Fokker-Planck equation: \[\label{Nonsmooth32fundamental32solution} \partial_t u=\nabla\cdot \left(\begin{pmatrix} 1 & 0 \\ 0 & 0 \end{pmatrix} \nabla u +u\begin{pmatrix} 1 & 0 \\ 0 & 1 \end{pmatrix} \begin{pmatrix} x \\ y \end{pmatrix} \right)\qquadin\quad(0,\infty)\times\mathbb{R}^2\,.\tag{24}\] In this case, condition [II] is not satisfied since \[\underbrace{\begin{pmatrix} 1 & 0 \\ 0 & 0 \end{pmatrix}}\limits_{Q_0} \begin{pmatrix} 0 \\ 0 \end{pmatrix} = \begin{pmatrix} 0 \\ 1 \end{pmatrix} \qquadand\qquad \underbrace{\begin{pmatrix} 1 & 0 \\ 0 & 1 \end{pmatrix}}\limits_{C} \begin{pmatrix} 0 \\ 1 \end{pmatrix} = \begin{pmatrix} 0 \\ 1 \end{pmatrix}\,.\] This orthogonality between the diffusion direction and the pure drift direction suggest that the fundamental solution should be of the form \[H(t,x)=g(t,x)h(t,y)\,,\] where \(g\) is the fundamental solution for the one dimensional problem \[\label{Example32diffusion32term} \partial_t g=\partial_{x_1}\left(\partial_{x}g+xg\right)\qquadin\quad(0,\infty)\times\mathbb{R}\,,\tag{25}\] and \(h\) solves the one dimensional transport equation \[\label{Example32transport32term} \partial_t h=\partial_{y}\left(y h\right)\qquadin\quad(0,\infty)\times\mathbb{R}\,.\tag{26}\] Using the explicit solutions of 25 and 26 , we conclude that the fundamental solution of 24 is \[H(t,x,y)=\underbrace{\frac{2}{\sqrt{2\pi}(1-e^{-2t})}\, e^{-\frac{2\,x^2}{(1-e^{-2t})}}}\limits_{g(t,x)}\;\underbrace{e^t\delta_0(e^ty)}\limits_{h(t,y)}\,.\] Thus, if we want to obtain the solution of the Cauchy problem associated to 24 with initial datum \(u_0\) regular enough, we have just to convolve \(u_0\) with the fundamental solution: \[\begin{align} u(t,x,y)&=\big(H(t,\cdot,\cdot)*u_0\big)(x,y)=\int_{\mathbb{R}}\int_{\mathbb{R}}g(t,x-\tilde{x})h(t,y-\tilde{y})u_0(\tilde{x},\tilde{y})\text{d}\tilde{x}\text{d}\tilde{y}\\ &=e^t\int_{\mathbb{R}}g(t,x-\tilde{x})u_0(\tilde{x},e^ty)\text{d}\tilde{x} \end{align}\] Notice that for the asymptotic behaviour is not clear what should be the steady states. On the one hand, fixed \(y\in\mathbb{R}\), \(\big(g(t,\cdot)*u_0(\cdot,y)\big)(x)\) converge to the gaussian \[g_\infty(x)=\frac{1}{\sqrt{2\pi}}e^{-\frac{x^2}{2}}\,.\] But on the other hand, if we fix \(x\in\mathbb{R}\), then \(\big(h(t,\cdot)*u_0(x,\cdot)\big)(y)\) converge to \(\delta_0(y)\) as a measure. Therefore, our educated guess for the steady state is

4.4 Entropy method↩︎

Once we have established the existence of a stationary measure (in the general case) we address the question of convergence towards it, possibly with rates. As already mentioned, we shall use entropy methods, and for this reason we shall assume that the stationary measure \(\rho_\infty\) is indeed an \(L^1\) function, that we fix throughout this section. We explore possible adaptations to our problem of the classical entropy method introduced by Bakry and Émery in [17], see also [56]. Recall the stochastic approximation of order 1 of the SGD defined by \[\text{d}X_t=-\nabla L(X_t)\,{\rm d}t+\sqrt{\frac{\eta}{\mathcalboondox{b}}Q(X_t)} \, \textrm{d} W_t\,,\] together with the following Fokker-Planck equation for the associated probability density function \(\rho\) : \[\label{FP32variant} \begin{cases} \partial_t\rho&=\nabla\cdot\bigg(\varepsilon^2\nabla\cdot\left(Q(x)\rho\right) +\rho\nabla L(x)\bigg)\\ \rho(0,\cdot)&=\rho_0\in L^2(\mathbb{R}^d, \rho_\infty^{-1}\,{\rm d}x). \end{cases}\tag{27}\] where \(\varepsilon^2= \frac{\eta}{2\mathcalboondox{b}}\). Entropy methods aim at obtaining differential inequalities between the entropy functional and its time derivative (entropy production), by means of suitable functional inequalities, typically of weighted Poincaré type. There may be several possible entropies for the above equation, however the following choice seems the most appropriate.

In order to establish a differential inequality for the entropy, we need a relationship between the entropy and its entropy production, which typically takes the form of a weighted Poincaré inequality associated to 27 .

Indeed, letting \(f=\mu/\rho_\infty\) with \(\int_{\mathbb{R}^d}\mu(x)\,{\rm d}x=1\), the above inequality reads as the entropy-entropy production inequality \[\label{Poinc46E-EP} \mathcal{E}(\mu|\rho_\infty)\le \frac{\varepsilon^2}{\lambda}\mathcal{I}\left(\mu\,|\,\rho_\infty\right).\tag{29}\] Note that \(f\in L^2(\mathbb{R}^d,\rho_\infty\,{\rm d}x)\) if and only if \(\mu\in L^2(\mathbb{R}^d,\rho_\infty^{-1}\,{\rm d}x)\) .

The above Poincaré inequality is relevant because it allows to derive exponential decay of the entropy, and show convergence towards \(\rho_\infty\) with precise rates, given by \(\lambda>0\).

Proving weighted Poincaré inequalities for a general nonnegative matrix \(Q\) and a general loss function \(L\) is a difficult challenge, in particular since the form of \(\rho_\infty\) is not explicit (except some special cases). To the best of our knowledge, only very partial results, under quite strong conditions, appear in the literature. We collect them hereafter, adapted to our notations.

In the isotropic framework presented in Theorem 24, we can use Proposition 23 to conclude the following entropy convergence.

If we want to consider more general diffusion matrices \(Q\) which can be degenerate, only few results are available. In this context of entropy methods, Fokker Planck equations with constant yet degenerate diffusion matrix \(Q_0\) and quadratic loss function \(L(x)=\tfrac{1}{2}x^T C x\) are studied by Arnold and Erb in [5]. Namely, \[\label{Simple32Degenerate32FP} \begin{cases} \quad\;\,\partial_t u=\nabla\cdot\left(Q_0\nabla u+u\,Cx\right)\qquadin\quad (0,\infty)\times\mathbb{R}^d \,,\\ u(0,x)=u_0(x)in \quad \mathbb{R}^d \,, \end{cases}\tag{30}\] For related Fokker Planck equations see [61], [62]. In [5], the authors adapted the entropy method for equation 30 under following assumptions:

1. Confinement potential: \(C\) is positive definite and symmetric³.

2. Hörmander condition: There are no eigenvectors of \(C\) in ker\(\,Q_0\).

Assumption [I] ensures that the mass is not escaping at infinity along the flow. However, the idea behind Hörmander condition [II] in the equation 30 is that if there is a point where diffusion is not acting, then the drift term will push it to the diffusion regime, see the discussion of 30 in [63].

According to [5], conditions [I] and [II] are equivalent to the existence and uniqueness of a steady state in \(L^1(\mathbb{R}^d)\) and imply the convergence of the solution of 30 to it.

We will denote by \(\lambda_{\max}(C)\) and \(\lambda_{\min}(C)\) the maximum and minimum eigenvalues of a nonnegative definite matrix \(C\), respectively.

There are interesting cases that fall out of the hypothesis of Theorem 27. Let us consider the following example where Hörmander’s condition [II] does not hold on a given subspace. In this case there is still a steady state, but it will not be smooth anymore. As a consequence, the convergence to equilibrium can only happen in a suitable weak sense.

Let us consider the equation for the function \(u:(0,\infty)\times\mathbb{R}^n\times\mathbb{R}^{d-n}\) with \(x\in\mathbb{R}^n\), \(y\in\mathbb{R}^{d-n}\) and matrices \(Q_0,C_0,C_1,C_2,C_3\) with the corresponding dimensions: \[\label{Generalized32simple32FP} \begin{cases} \partial_t u&=\nabla_{x,y}\cdot\left[ \begin{pmatrix} Q_0 & 0 \\ 0 & 0 \end{pmatrix} \begin{pmatrix} \nabla_x u \\ \nabla_y u \end{pmatrix} +u \begin{pmatrix} C_0 & C_1 \\ C_2 & C_3 \end{pmatrix} \begin{pmatrix} x \\ y \end{pmatrix} \right]\qquadin \;(0,\infty)\times\mathbb{R}^n\times\mathbb{R}^{d-n}\\ u(0)&=u_0on \;\mathbb{R}^n\times\mathbb{R}^{d-n} \end{cases}\tag{31}\] In order to construct the solution of 31 , we need the fundamental solution \(g\) defined in \((0,\infty)\times\mathbb{R}^n\) of the equation \[\partial_t g=\nabla\cdot\left(Q_0\nabla g+gC_0 x\right)\qquadin \;(0,\infty)\times\mathbb{R}^n\,.\] If \(Q_0\) and \(C_0\) satisfy assumption [II], then (see e.g. [5]) we have the explicit expression for \(g\) given by \[g(t,x)=\frac{1}{\left(2\pi\right)^{d/2}\det(W(t))}\exp\left(-x^TW^{-1}(t)x\right)\,,\] where \[W(t)=\int_{0}^{t}e^{C_0(s-t)}Q_0e^{C^T_0(s-t)}\,{\rm d}s\] is a positive definite matrix for all \(t>0\). In addition, if also condition [I] is satisfied there exists a unique steady state \(g_\infty\) as in Theorem 26.

5 Some conclusions and open questions↩︎

In the previous section we have seen that in some cases, the convergence of the distribution of parameters towards a steady state can also be exponentially fast: this holds under quite restrictive conditions on \(Q\) and \(L\). However, we think that this special situation is what happens to the distribution of the parameters near local minima of the loss function \(L\). The heuristic reason is that \(L\) should be \(\lambda\)-convex close to a minimum, and also \(Q\) should be (practically) constant there. Let us be more precise: assuming all the necessary regularity for \(\rho\), the behaviour of the stochastic process with probability density function \(\rho\) near the critical points of \(L\) can be understood as follow. Recall that \(\rho\) satisfies \[\label{Fokker-Planck} \begin{cases} \partial_t\rho=\nabla\cdot\bigg(\varepsilon^2Q(x)\nabla\rho+\rho\left(\nabla L(x)+\varepsilon^2\nabla\cdot Q(x)\right)\bigg) & on\quad(0,\infty)\times\mathbb{R}^d\,,\\[4mm] \rho(0,x)=\rho_0(x) &in\quad\mathbb{R}^d\,, \end{cases}\tag{32}\] with \(\nabla L(x)=\frac{1}{N}\sum_{i=1}^{N}\nabla L_i(x)\) and \[Q(x)=\frac{1}{N}\sum_{i=1}^{N}\nabla L_i(x)\otimes\nabla L_i(x)-\nabla L(x)\otimes\nabla L(x)\,.\] Suppose that \(x_0\in\mathbb{R}^d\) is a local minimum of \(L\), so that \(\nabla L(x_0)=0\). The goal is to obtain an asymptotic expansion of 32 in a neighbourhood of \(x_0\). We consider, for small \(h \in (-1,1)\), and for \(z \in \mathbb{R}^d\), \[x=x_0+ h z\,,\] and assuming enough regularity of the loss functions \(L_i\) for \(i=1,\dots,N\) we can compute the Taylor expansion of \(\nabla L\) and \(Q\) around \(x_0\): \[\begin{align} \nabla L(x)&={\nabla L(x_0)}+h D^2 L(x_0)z+h^2 D^3L(x_0)[z,z] +o\left(|hz|^2\right)\\ Q(x)&=\frac{1}{N}\sum_{i=1}^{N}\nabla L_i(x_0)\otimes\nabla L_i(x_0)-\nabla L(x_0)\otimes\nabla L(x_0)\\ &\;+\frac{h^2}{N}\sum_{i=1}^{N}\bigg[\!\left(D^2L_i(x_0)z\right)\otimes\left(D^2 L_i(x_0)z\right)-\left(D^2 L(x_0) z\right)\otimes\left( D^2 L(x_0)z\right)\!\bigg]\!+o(|hz|^2)\,. \end{align}\] Since \(\nabla L(x_0)=0\), the first order expansion reads as follows: \[\begin{align} \nabla L(x) &=h D^2L(x_0) z+o(|hz|)\\ Q(x)&=\frac{1}{N}\sum_{i=1}^{N}\nabla L_i(x_0)\otimes\nabla L_i(x_0)+o(|hz|)=Q(x_0)+o(|hz|)\,. \end{align}\] Now, let us consider the function \(\tilde{u}(t,z):=\rho(t,x(z))=\rho(t,x_0+h z)\). Assuming enough regularity for \(\tilde{u}\), we have \[\begin{align} \nabla_z \tilde{u}(t,y) & =h \nabla_x\rho(t,x(z))\,, \\ \nabla_z\cdot\left(A\;\nabla_z \tilde{u}(t,z)\right) & =h^2\nabla_x\cdot\left(A\;\nabla_x\rho(t,x(z))\right)\,, \end{align}\] for any matrix \(A\in \mathbb{R}^{d\times d}\). As a consequence we can write the equation for \(\tilde{u}\) as follows: \[\begin{align} \label{local32FP32with32error} \partial_t \tilde{u} & = \nabla_z\cdot \left(\varepsilon^2 Q(x_0)\nabla_z \tilde{u}+\tilde{u}\;D^2L(x_0)\,z\right)+h^2\;\nabla_z\cdot\left(O(1)\nabla_z \tilde{u}+\tilde{u} \,O(1)\right)\,, \end{align}\tag{33}\] This suggests that the behaviour of \(\rho\) near a critical point \(x_0\) of \(L\) is governed by the following equation for \(u_h = u_h(t,z)\): \[\label{localize32FP} \begin{cases} \quad\;\,\partial_t u_h =\nabla_z\cdot\left(\varepsilon^2 Q_0\nabla_z u_h+u_h\,Cz\right)\qquadin\quad (0,\infty)\times\mathbb{R}^d \,,\\ u_h(0,z)=\rho_0(x_0+h z)\qquad\qquad\qquad\qquad\quad\;\;in \quad \mathbb{R}^d \,, \end{cases}\tag{34}\] where \(Q_0 = Q(x_0) \ge 0\) is constant and symmetric, and \(C=D^2L(x_0)\in \mathbb{R}^{d\times d}\). Also, since \(x_0\) is a local minimum we have \(D^2 L(x_0)\ge 0\). This allows us to apply the convergence results of Theorem 28, and get exponential decay towards a steady state.

The conclusion that we draw from the above discussion is that the behaviour of the parameters during the learning process and their final distribution, can be understood through equation 34 , more precisely, when Theorem 28 applies we obtain that there exists a steady state \(u_\infty\) so that weak convergence in measure is exponentially fast, that is ?? holds, and the second moments converge strongly and exponentially fast, namely ?? holds, i.e., \[\left|\int_{\mathbb{R}^n}\int_{\mathbb{R}^{d-n}}(|x|^2+|y|^2)(u(t,x,y)-u_\infty(x,y))\,{\rm d}x\,{\rm d}y \right|\le \kappa_2 e^{-\lambda t}\,,\] for suitable \(\kappa_2,\lambda>0\).

Open Questions. We shall present a number of technical issues that we do not prove in this paper and would eventually lead to a rigorous proof of the above heuristic result.

i) Regularity. The first issue is undoubtedly the regularity of the solution \(\rho\) to problem 32 , that can be obtained through the regularity of the “heat” kernel or fundamental solution to 32 . We expect partial regularity in the “non degenerate” variables and we conjecture this to be enough to justify rigorously the above expansions.

ii) Localization error estimates. The next step would be to quantify the error between the original solution \(\tilde{u}\) of 33 and the localized solution \(u\) of 34 (in each neighbourhood of the local minima) when \(h>0\) is small (and fixed momentarily). We need suitable norm estimates for all \(t>0\), of the form \[\|\tilde{u}(t)-u(t)\|\le c(h,t)\,.\] The main issue would be to show that the error term \(c(h,t)\xrightarrow[]{h\rightarrow 0^+}0\) possibly with quantitative estimates of its behaviour.

iii) Building a global approximated solution. Assuming that the error term above can be controlled at each minima, it seems reasonable to approximate the solution \(\rho\) of 32 by a combination of the localized solutions of 34 around each local minimum.

Let us be more precise: let \(\lbrace x_1,\dots,x_M\rbrace\) denote the local minima of \(L\), and \(u_i\) the corresponding solution to the localized equation 34 around the local minimum \(x_i\) (for any \(1\le i\le M\)). We conjecture that the global behaviour of \(\rho\) should be given by \[\label{approximated32solution} \rho(t,x)\approx\sum_{i=1}^{M}m_i(t)\;u_i\left(t,\frac{x-x_i}{h}\right)\qquad\forall t>0\;\;\forall x\in\mathbb{R}^d\,,\tag{35}\] where \(m_i(t)\) represents the mass concentration/splitting around each local minimum \(x_i\) at time \(t>0\) and satisfy \(\sum_{i=1}^{M}m_i(t)=1\) for all \(t>0\).

Another challenge would be to show that the above approximation holds for all (sufficiently large) times. This would give information about the invariant probability measure (i.e., the global steady state), that should be well approximated by \[ \rho_\infty(x)\approx\sum_{i=1}^{M}m_i(\infty)\;u_{i,\infty}\left(\frac{x-x_i}{h}\right)\qquad\forall x\in\mathbb{R}^d\,,\] where \(u_{i,\infty}\) are the stationary solutions as in Theorem 28, with diffusion matrix \(Q(x_i)\) and drift matrix \(D^2L(x_i)\).

iv) Sharp mass displacement and splitting. A crucial question in the above approximation would be to obtain sharp information about the portion of mass \(m_i(t)\) that gradually concentrates around the minima \(x_i\) of \(L\), together with its dependence on the initial distribution \(\rho_0\). The local masses \(m_i\) represent the probability of the parameters of the optimization to be \(x_i\in\mathbb{R}^d\) after training with SGD. As we have shown in Example 6, when \(Q(x)=0\) for every \(x\in\mathbb{R}^d\), Equation 32 becomes a “pure transport” equation and the initial datum \(\rho_0\) will determine completely the asymptotic behaviour of \(\rho(t)\). Given a datum \(\rho_0\), a delicate issue is to be able to estimate in the sharpest possible way which portion of its mass concentrate around each of the \(x_i\), after we have entered the asymptotic regime. Another crucial issue would be to relate the asymptotic regime, i.e. the smallest time for which the approximation of point \(iii)\) is effective, with the mean exit times studied in Section 3.

Conflict of interest statement. On behalf of all authors, the corresponding author states that there is no conflict of interest.

Data availability statement. All data generated or analysed during this study are included in the published article.

References↩︎

---

1. Note that, due to the monotonicity of \(F\), this applies to the Laplace equation \(F(x,r,p,X) = - tr(X)\) means \(\Delta u \geq 0\)↩︎

2. The simplest case in which a Poincaré inequality in the sense of Definition 22 does not hold is when \(Q\) is a diagonal constant coefficient matrix with all eigenvalues zero except one.↩︎

3. Note that the authors in [5] address a slightly more general problem where \(C\) is not necessarily symmetric and hence they assume that \(C\) is positive stable, i.e. all eigenvalues have positive real part.↩︎