Fluctuation-dissipation relations for stochastic gradient descent

1 INTRODUCTION . Equilibration rules the long-term fate of many macroscopic dynamical systems . For instance , as we pour water into a glass and let it be , the stationary state of tranquility is eventually attained . Zooming into the tranquil water with a microscope would reveal , however , a turmoil of stochastic fluctuations that maintain the apparent stationarity in balance . This is vividly exemplified by the Brownian motion ( Brown , 1828 ) : a pollen immersed in water is constantly bombarded by jittery molecular movements , resulting in the macroscopically observable diffusive motion of the solute . Out of the effort in bridging microscopic and macroscopic realms through the Brownian movement came a prototype of fluctuation-dissipation relations ( Einstein , 1905 ; Von Smoluchowski , 1906 ) . These relations quantitatively link degrees of noisy microscopic fluctuations to smooth macroscopic dissipative phenomena and have since been codified in the linear response theory for physical systems ( Onsager , 1931 ; Green , 1954 ; Kubo , 1957 ) , a cornerstone of statistical mechanics . Machine learning begets another form of equilibration . As a model learns patterns in data , its performance first improves and then plateaus , again reaching apparent stationarity . This dynamical process naturally comes equipped with stochastic fluctuations as well : often given data too gigantic to consume at once , training proceeds in small batches and random selections of these mini-batches consequently give rise to the noisy dynamical excursion of the model parameters in the loss-function landscape , reminiscent of the Brownian motion . It is thus natural to wonder if there exist analogous fluctuation-dissipation relations that quantitatively link the noise in mini-batched data to the observable evolution of the model performance and that in turn facilitate the learning process . Here , we derive such fluctuation-dissipation relations for the stochastic gradient descent algorithm . The only assumption made is stationarity of the probability distribution that governs the model parameters at sufficiently long time . Our results thus apply to generic cases with non-Gaussian mini-batch noises and nonconvex loss-function landscapes . Practically , the first relation ( FDR1 ) offers the metric for assessing equilibration and yields an adaptive algorithm that sets learning-rate schedule on the fly . The second relation ( FDR2 ) further helps us determine the properties of the lossfunction landscape , including the strength of its Hessian and the degree of anharmonicity , i.e. , the deviation from the idealized harmonic limit of a quadratic loss surface and a constant noise matrix . Our approach should be contrasted with recent attempts to import the machinery of stochastic differential calculus into the study of the stochastic gradient descent algorithm ( Mandt et al. , 2015 ; Li et al. , 2015 ; Mandt et al. , 2017 ; Li et al. , 2017 ; Smith & Le , 2018 ; Chaudhari & Soatto , 2017 ; Jastrzebski et al. , 2017 ; Zhu et al. , 2018 ; An et al. , 2018 ) . This line of work all assumes Gaussian noises and sometimes additionally employs the quadratic harmonic approximation for loss-function landscapes . The more severe drawback , however , is the usage of the analogy with continuous-time stochastic differential equations , which is inconsistent in general ( see Section 2.3.3 ) . Instead , the stochastic gradient descent algorithm can be properly treated within the framework of the KramersMoyal expansion ( Van Kampen , 1992 ; Gardiner , 2009 ; Risken , 1984 ; Radons et al. , 1990 ; Leen & Moody , 1993 ) . The paper is organized as follows . In Section 2 , after setting up notations and deriving a stationary fluctuation-dissipation theorem ( FDT ) , we derive two specific fluctuation-dissipation relations . The first relation ( FDR1 ) can be used to check stationarity and the second relation ( FDR2 ) to delineate the shape of the loss-function landscape , as empirically borne out in Section 3 . An adaptive scheduling method is proposed and tested in Section 3.3 . We conclude in Section 4 with future outlooks . 2 FLUCTUATION-DISSIPATION RELATIONS . A model is parametrized by a weight coordinate , θ = { θi } i=1 , ... , P . The training set of Ns examples is utilized by the model to learn patterns in the data and the model ’ s overall performance is evaluated by a full-batch loss function , f ( θ ) ≡ 1Ns ∑Ns α=1 fα ( θ ) , with fα ( θ ) quantifying the performance of the model on a particular sample α : the smaller the loss is , the better the model is expected to perform . The learning process can thus be cast as an optimization problem of minimizing the loss function . One of the most commonly used optimization schemes is the stochastic gradient descent ( SGD ) algorithm ( Robbins & Monro , 1951 ) in which a mini-batch B ⊂ { 1 , 2 , . . . , Ns } of size |B| is stochastically chosen for training at each time step . Specifically , the update equation is given by θ ( t+ 1 ) = θ ( t ) − η∇fB [ θ ( t ) ] , ( 1 ) where η > 0 is a learning rate and a mini-batch loss fB ( θ ) ≡ 1|B| ∑ α∈B fα ( θ ) . Note that q ∇fB ( θ ) y m.b . = ∇f ( θ ) , ( 2 ) with J. . .Km.b . denoting the average over mini-batch realizations . For later purposes , it is convenient to define a full two-point noise matrix C̃ through1 C̃i , j ( θ ) ≡ q [ ∂if B ( θ ) ] [ ∂jf B ( θ ) ] y m.b . ( 3 ) and , more generally , higher-point noise tensors C̃i1 , i2 , ... , ik ( θ ) ≡ q [ ∂i1f B ( θ ) ] [ ∂i2f B ( θ ) ] · · · [ ∂ikf B ( θ ) ] y m.b. . ( 4 ) Below , we shall not make any assumptions on the distribution of the noise vector ∇fB – other than that a mini-batch is independent and identically distributed from theNs training samples at each time step – and the noise distribution is therefore allowed to have nontrivial higher connected moments indicative of non-Gaussianity . It is empirically often observed that the performance of the model plateaus after some training through SGD . It is thus natural to hypothesize the existence of a stationary-state distribution , pss ( θ ) , that dictates the SGD sampling at long time ( see Section 2.3.4 for discussion on this assumption ) . For any observable quantity , O ( θ ) , – something that can be measured during training such as θ2 and f ( θ ) – its stationary-state average is then defined as 〈O ( θ ) 〉 ≡ ∫ dθ pss ( θ ) O ( θ ) . ( 5 ) 1A connected noise covariant matrix , Ci , j ( θ ) ≡ C̃i , j ( θ ) − [ ∂if ( θ ) ] [ ∂jf ( θ ) ] , will not appear in fluctuation-dissipation relations below but scales nicely with mini-batch sizes as ∝ 1|B| ( 1− |B| Ns ) ( Li et al. , 2017 ) . In general the probability distribution of the model parameters evolves as p ( θ , t + 1 ) =q∫ dθ′ p ( θ′ , t ) δ { θ − [ θ′ − η∇fB ( θ′ ) ] } y m.b . and in particular for the stationary state∫ dθ pss ( θ , t ) O ( θ ) = ∫ dθ pss ( θ , t+ 1 ) O ( θ ) ( 6 ) = s∫ dθ ∫ dθ′ pss ( θ ′ , t ) δ { θ − [ θ′ − η∇fB ( θ′ ) ] } O ( θ ) { m.b . = ∫ dθ′ pss ( θ ′ ) q O [ θ′ − η∇fB ( θ′ ) ] y m.b. . Thus follows the master equation 〈O ( θ ) 〉 = 〈q O [ θ − η∇fB ( θ ) ] y m.b . 〉 . ( FDT ) In the next two subsections , we apply this general formula to simple observables in order to derive various stationary fluctuation-dissipation relations . Incidentally , the discrete version of the FokkerPlanck equation can be derived through the Kramers-Moyal expansion , considering the more general nonstationary version of the above equation and performing the Taylor expansion in η and repeated integrations by parts ( Van Kampen , 1992 ; Gardiner , 2009 ; Risken , 1984 ; Radons et al. , 1990 ; Leen & Moody , 1993 ) . 2.1 FIRST FLUCTUATION-DISSIPATION RELATION . Applying the master equation ( FDT ) to the linear observable , 〈θ〉 = 〈q θ − η∇fB ( θ ) y m.b . 〉 = 〈θ〉 − η 〈∇f ( θ ) 〉 . ( 7 ) We thus have 〈∇f〉 = 0 . ( 8 ) This is natural because there is no particular direction that the gradient picks on average as the model parameter stochastically bounces around the local minimum or , more generally , wanders around the loss-function landscape according to the stationary distribution . Performing similar algebra for the quadratic observable 〈θiθj〉 yields 〈θi ( ∂jf ) 〉+ 〈 ( ∂if ) θj〉 = η 〈 C̃i , j 〉 . ( 9 ) In particular , taking the trace of this matrix-form relation , we obtain 〈θ · ( ∇f ) 〉 = 1 2 η 〈 Tr C̃ 〉 . ( FDR1 ) More generally , in the case of SGD with momentum µ and dampening ν , whose update equation is given by v ( t+ 1 ) = µv ( t ) − ( 1− ν ) ∇fB [ θ ( t ) ] , ( 10 ) θ ( t+ 1 ) = θ ( t ) + ηv ( t+ 1 ) , ( 11 ) a similar derivation yields ( see Appendix A ) 〈θ · ( ∇f ) 〉 = ( 1 + µ ) 2 ( 1− ν ) η 〈 v2 〉 . ( FDR1 ’ ) The last equation reduces to the equation ( FDR1 ) when µ = ν = 0 with v = −∇fB . Also note that 〈θ · ( ∇f ) 〉 = 〈 ( θ − θc ) · ( ∇f ) 〉 for an arbitrary constant vector θc because of the equation ( 8 ) . This first fluctuation-dissipation relation is easy to evaluate on the fly during training , exactly holds without any approximation if sampled well from the stationary distribution , and can thus be used as the standard metric to check if learning has plateaued , just as similar relations can be used to check equilibration in Monte Carlo simulations of physical systems ( Santen & Krauth , 2000 ) . [ It should be cautioned , however , that the fluctuation-dissipation relations are necessary but not sufficient to ensure stationarity ( Odriozola & Berthier , 2011 ) . ] Such a metric can in turn be used to schedule changes in hyperparameters , as shall be demonstrated in Section 3.3 . 2.2 SECOND FLUCTUATION-DISSIPATION RELATION . Applying the master equation ( FDT ) on the full-batch loss function and Taylor-expanding it in the learning rate η yields the closed-form expression 〈f ( θ ) 〉 = 〈q f [ θ − η∇fB ( θ ) ] y m.b . 〉 ( 12 ) = 〈 f + ∞∑ k=1 ( −η ) k k ! P∑ i1 , i2 , ... , ik=1 ( ∂i1∂i2 · · · ∂ikf ) q ( ∂i1f B ) ( ∂i2fB ) · · · ( ∂ikfB ) ym.b . 〉 = 〈f〉 − η 〈 ( ∇f ) 2 〉 + ∞∑ k=2 ( −η ) k k ! 〈 P∑ i1 , i2 , ... , ik=1 Fi1 , i2 , ... , ikC̃i1 , i2 , ... , ik 〉 where we recalled the equation ( 4 ) and introduced Fi1 , i2 , ... , ik ( θ ) ≡ ∂i1∂i2 · · · ∂ikf ( θ ) . ( 13 ) In particular , Hi , j ( θ ) ≡ Fi , j ( θ ) is the Hessian matrix . Reorganizing terms , we obtain 〈 ( ∇f ) 2 〉 = η 2 〈 Tr ( HC̃ ) 〉 − η2 ∞∑ k=3 ( −η ) k−3 k ! 〈 P∑ i1 , i2 , ... , ik=1 Fi1 , i2 , ... , ikC̃i1 , i2 , ... , ik 〉 . ( FDR2 ) In the case of SGD with momentum and dampening , the left-hand side is replaced by ( 1 − ν ) 〈 ( ∇f ) 2 〉 − µ 〈v ·∇f〉 and C̃i1 , i2 , ... , ik by more hideous expressions ( see Appendix A ) . We can extract at least two types of information on the loss-function landscape by evaluating the dependence of the left-hand side , G ( η ) ≡ 〈 ( ∇f ) 2 〉 , on the learning rate η . First , in the small learning rate regime , the value of 2G ( η ) /η approximates Tr ( HC̃ ) around a local ravine . Second , nonlinearity ofG ( η ) at higher η indicates discernible effects of anharmonicity . In such a regime , the Hessian matrixH can not be approximated as constant ( which also implies that { Fi1 , i2 , ... , ik } k > 2 are nontrivial ) and/or the noise two-point matrix C̃ can not be regarded as constant . Such nonlinearity especially indicates the breakdown of the harmonic approximation , that is , the quadratic truncation of the loss-function landscape , often used to analyze the regime explored at small learning rates .

Summary: Understanding the stationary equilibrium helps to understand the practical performance of stochastic gradient descent. In this paper, the authors propose two fluctuation-dissipation relation to link measurable quantities and hyperparameters in the stochastic gradient descent algorithm. An advantage over the existing study is that the results here hold for any stationary state and do not need the analogy with continuous-time differential equations. Empirical results are also reported to verify these fluctuation relation.