A Chaos Theory Approach to Understand Neural Network Optimization

Despite the complicated structure of modern deep neural network architectures , they are still optimized with algorithms based on Stochastic Gradient Descent ( SGD ) . However , the reason behind the effectiveness of SGD is not well understood , making its study an active research area . In this paper , we formulate deep neural network optimization as a dynamical system and show that the rigorous theory developed to study chaotic systems can be useful to understand SGD and its variants . In particular , we first observe that the inverse of the instability timescale of SGD optimization , represented by the largest Lyapunov exponent , corresponds to the most negative eigenvalue of the Hessian of the loss . This observation enables the introduction of an efficient method to estimate the largest eigenvalue of the Hessian . Then , we empirically show that for a large range of learning rates , SGD traverses the loss landscape across regions with largest eigenvalue of the Hessian similar to the inverse of the learning rate . This explains why effective learning rates can be found to be within a large range of values and shows that SGD implicitly uses the largest eigenvalue of the Hessian while traversing the loss landscape . This sheds some light on the effectiveness of SGD over more sophisticated second-order methods . We also propose a quasi-Newton method that dynamically estimates an optimal learning rate for the optimization of deep learning models . We demonstrate that our observations and methods are robust across different architectures and loss functions on CIFAR-10 dataset . 1 INTRODUCTION . An interesting observation from current deep learning research is that classification and regression accuracy gains seem to be achieved from the intricacy of the underlying models rather than the optimization algorithm used for their training . Actually , the de facto choice for the optimization algorithm is still the classic Stochastic Gradient Descent ( SGD ) algorithm ( Robbins & Monro , 1951 ) with minor modifications ( Duchi et al. , 2011 ; Sutskever et al. , 2013 ; Kingma & Ba , 2014 ) . Even though several sophisticated second-order and quasi-Newton methods ( Martens , 2010 ; Martens & Grosse , 2015 ; Berahas et al. , 2019 ) have been introduced , first-order methods remain popular and none of them seem to outperform SGD with a carefully tuned learning rate schedule ( Hardt et al. , 2016 ) . This indicates that SGD ( or in general first-order methods ) probably has some intrinsic properties that make it effective to optimize over-parametrized deep neural networks . Despite various attempts to explain such phenomenon ( Chaudhari & Soatto , 2018 ; Keskar et al. , 2016 ; Kleinberg et al. , 2018 ) , little is understood about the effectiveness of SGD over sophisticated second-order optimization methods . In this paper , we argue that chaos theory ( Sprott & Sprott , 2003 ) is a useful approach to understand the neural network optimization based on SGD . The basic idea is to view neural network optimization as a dynamical system where the SGD update equation maps from the space of learnable parameters to itself and describes the evolution of the system over time . Once the evolution is defined , the rich theory developed to study chaotic dynamical systems can be leveraged to analyze and understand SGD and its variants . In essence , chaos theory enables us to study the evolution of the learnable parameters ( i.e. , the optimization trajectory ) in order to understand the training behavior over large time scales ( i.e. , number of iterations ) . In particular , we focus on understanding the influence of the learning rate on the SGD optimization trajectory . First , by observing that the Lyapunov exponent of SGD is the most negative eigenvalue of the Hessian of the loss , we introduce an efficient and accurate method to estimate the loss curvature . Then , we empirically show that for a range of learning rate schedules , SGD traverses the optimization landscape across regions with largest eigenvalue of the Hessian similar to the inverse of the learning rate . This demonstrates that at a specific time step , performing SGD update is similar to performing a quasi-Newton step , considering only the largest eigenvalue of the Hessian of the loss . This , for the first time , sheds some light on the effectiveness of SGD over more sophisticated second-order methods and corroborates the observation that SGD robustly converges for a variety of learning rate schedules ( Sun , 2019 ) . Furthermore , as pointed out in ( LeCun et al. , 1993 ) , the inverse of the estimated curvature can be used as the learning rate when applying SGD to a new dataset or architecture . Hence , we can set up a “ feedback ” system where the quasi-Newton optimal learning rate is calculated dynamically based on the current largest eigenvalue of the Hessian ( curvature ) , and the learning rate is consequently adjusted during the training , allowing a “ parameter free ” stochastic gradient descent optimization . The experiments are conducted on CIFAR-10 dataset to demonstrate that our observations are robust across a variety of models , including a simple linear model regression and more modern deep neural network architectures , trained with both cross entropy and mean square error loss functions . 2 CHAOS THEORY FOR NEURAL NETWORK OPTIMIZATION . In recent years , several papers have used dynamical systems to study theoretical aspects of deep learning optimization ( Liu & Theodorou , 2019 ) . Essentially , this is achieved by defining the optimization of deep neural networks as the evolution of parameters over time . In particular , a dynamical system progresses according to a map function that describes how the system evolves in a specific time step . In the case of deep neural network optimization , this map function is defined from the space of parameters into itself . By describing the system evolution using such a map function , it is possible to leverage the mathematical machinery of dynamical systems . For instance , viewing SGD as a discrete approximation of a continuous stochastic differential equations , allowed Li et al . ( 2017 ) and An et al . ( 2018 ) to propose adaptive SGD algorithms . Furthermore , dynamical systems enabled LeCun et al . ( 1993 ) to relate learning rate with the inverse of the local Hessian in a quasiNewton optimization framework . Our paper also uses dynamical systems to study deep learning optimization , but differently from all methods above , we rely on chaos theory . Chaos theory ( Sprott & Sprott , 2003 ) studies the evolution of dynamical systems over large time scales and can categorize systems into chaotic or non chaotic . Under some simplifying but still general assumptions , chaotic systems are bounded and have strong dependence on the initial conditions . This means that chaotic systems evolving from different starting points that are within a relatively small region around a particular reference point , will diverge exponentially during the evolution process , where the amount of time taken for this divergence to happen is defined as the chaotic timescale . This chaotic timescale imposes a limit on our ability to predict the future state distribution of a dynamical system . In fact , the distribution of the future state , which have evolved for more than a few times the chaotic timescale , can not be distinguished from random distributions , even when the system is fully deterministic . We apply concepts from chaos theory to improve our current understanding of the optimization of deep neural networks . More specifically , we describe how to use standard chaos theory techniques to efficiently calculate the leading ( positive and negative ) eigenvalues of the Hessian of the loss function . With these eigenvalues we measure , in turn , the loss function curvature , which can be used to study the behavior of first-order optimization methods , such as SGD ( Robbins & Monro , 1951 ) . In particular , with this technique we formulate an explanation for the empirical robustness of SGD to the choice of learning rate and its scheduling function , and we investigate a method ( based on quasi-Newton second order method ) for dynamically finding the optimal learning rate during the optimization of deep neural networks . Such automated and dynamic estimation of optimal learning rate can lift a significant burden from the manual definition of learning rate schedules in deep learning optimization . 2.1 LYAPUNOV EXPONENTS . In chaos theory , the Lyapunov exponents define the divergence rate of infinitesimally close trajectories , and the inverse of the largest Lyapunov exponent is the timescale that corresponds to the onset of chaos into the system . Two arbitrarily close initial conditions generate two solutions that diverge with time . Under the assumption that the map function of the system is differentiable , if one observes this divergence for a short time window , it grows exponentially . If the initial divergence q ( 0 ) is made smaller , the time window can be made larger ( t→∞ ) . The largest Lyapunov exponent λ is a measure of the growth of the divergence q ( t ) in the direction q̂ ( 0 ) = q ( 0 ) /‖q ( 0 ) ‖ with the largest growth ( maxq̂ ( 0 ) ) along the trajectory , as in λ = max q̂ ( 0 ) lim t→∞ lim ‖q ( 0 ) ‖→0 1 t log ‖q ( t ) ‖ ‖q ( 0 ) ‖ . ( 1 ) In this paper , we rely on the local finite size Lyapunov exponent . In this context , local in time means that there is no limit to infinity for t in equation 1 – instead , it is an average over a constant time window t. Finite size means keeping the difference in parameter space fixed as a small constant with ‖q‖ = ∆q ( i.e. , no limit ‖q‖ → 0 in equation 1 ) . Using a finite size allows the study of the dynamic system at a specific spatial scale ( for a comprehensive review , see ( Cencini & Vulpiani , 2013 ) ) , corresponding to the eigenvalues of the Hessian of a spatially smoothed version of the loss ( or equivalently , to the numerical second derivative with a finite delta ) . When this analysis is used to study the Hessian , this is equivalent to calculating the local numerical second derivative . We found empirically that the results do not depend on the ∆q parameter within a large range of values . We will show in Sec . 3 that this timescale ( i.e. , the Lyapunov exponent ) corresponds to the most negative eigenvalue of the Hessian of the loss , when optimizing deep neural networks with SGD . Intuitively , a small difference in the initial condition will amplify exponentially in the directions with negative second derivatives and will dampen in directions with positive second derivatives . Empirically , we find that the chaotic timescales in effective training of deep neural networks are short ( in the order of tens of iterations ) when compared with the time of one epoch ( i.e. , the total number of iterations in one epoch ) . We also find that there are multiple unstable directions throughout the training , i.e. , the system is hyper-chaotic . 2.2 LYAPUNOV EXPONENTS FOR GD AND SGD . In this section we derive the formula to compute the largest Lyapunov exponents for Gradient Descent ( GD ) following ( Sprott & Sprott , 2003 ) . We first show that the largest Lyapunov exponent corresponds to the most negative eigenvalue of the Hessian of the loss and provide an algorithm to efficiently compute it . This will be later extended to calculate the largest ( or in general the top-k ) eigenvalue of the Hessian in section 3 . For simplicity of the exposition , in this section we initially consider the non-stochastic setting . Also , for the results of this section to hold , we assume that the Hessian of the loss does not change quickly through time , and it does not change quickly along the optimization trajectory compared to the chaotic time scale . These assumptions can easily be checked a posteriori , and we will show how to overcome this ( potential ) limitation in section 3 . Let θ be the vector of learnable parameters of the deep neural network , L ( · ) be the loss function , and α > 0 be the learning rate . The gradient descent step at iteration t is written as : θt+1 = θt − α dL ( θt ) dθ , ( 2 ) where the update step ∆θ = −αdL/dθ . In the limit of small steps the formulation is equivalent to a Partial Differential Equation ( PDE ) dθ dt = −α∂L ( θ ) ∂θ . ( 3 ) Integrating equation 3 gives the evolution of the system , which is equivalent to training the neural network . To compute the chaotic time scale ( i.e . the inverse of the Lyapunov exponent ) , one needs to analyze the difference in evolution of GD at two arbitrarily close initial points . To this end , we consider a small perturbation q0 added to the initial weights θ0 . For this perturbed starting point θ0 + q0 , the PDE becomes : d ( θ + q ) dt = −α∂L ( θ + q ) ∂θ . ( 4 ) In the limit of small q , considering the first order Taylor approximation of the above equation and subtracting equation 3 , we obtain : dq dt = ∂ ( −α∂L ( θ ) ∂θ ) ∂θ q . ( 5 ) Then , integrating equation 5 , we obtain the evolution of the perturbation under GD : q ( t ) = exp ( −α∂ 2L ( θ ) ∂θ2 t ) q0 . ( 6 ) This remains true as long as q ( t ) remains small , where the definition of small depends on the properties of L. We consider the decomposition of q0 as a sum of its projections on the eigenspace of the Hessian of the loss ( with the Hessian being represented at the exponent of the formula in equation 6 ) . In this space , the projection of q0 along the direction corresponding to the largest eigenvalue is the one growing the fastest . Starting with a random q0 , the direction of q that becomes dominant after sufficient time is aligned with the eigenvector of the largest eigenvalue of the matrix at the exponent , and the growth rate of q is equal to the corresponding eigenvalue . Measuring this growth rate provides a simple and linear ( in the number of parameters ) method to measure the leading eigenvalue . This procedure represents the calculation of the largest Lyapunov exponent , i.e. , the largest eigenvalue ( λ0 ) of the matrix −α∂2L/∂θ2 . Due to the minus sign , this corresponds to the smallest eigenvalue ( hN ) of the Hessian of the loss ( H = ∂2L/∂θ2 ) . More precisely , the smallest eigenvalue of the Hessian and the largest Lyapunov exponent are related as hN = −λ0/α . For non-convex losses , hN is the most negative eigenvalue and the matching eigenvector corresponds to the most unstable direction of the optimization of the loss . Once q ( t ) is aligned with the largest eigenvector , equation 6 becomes q ( t+ ∆t ) = exp ( λ0∆t ) q ( t ) . ( 7 ) The algorithm to calculate λ0 requires normalizing the length of q at each step to keep the increment “ small ” . This reference distance is equivalent to the choice of the step size for the calculation of finite difference based second derivative . In dynamical systems terminology this is called calculating the finite-size Lyapunov exponent . Now , the largest Lyapunov exponent is obtained by iterating the following two steps : λ0 ← log ( ‖q ( t+ ∆t ) ‖ ‖q ( t ) ‖ ) /∆t , ( 8 ) q ( t+ ∆t ) ← q ( t+ ∆t ) ‖q ( t ) ‖ ‖q ( t+ ∆t ) ‖ , where ‖·‖ denotes the L2 norm and ∆t denotes the time step . One could see that the computation of the largest Lyapunov exponent is analogous to the power method to compute the largest eigenvalue of a given matrix . This idea can be easily extended to compute the top-k Lyapunov exponents following the idea of Benettin et al . ( 1980 ) . Please refer to the appendix C. SGD can be described with the same approach , with the loss function replaced by L ( θ , ω ) where ω are random variables that describe which images are picked in each minibatch , the data augmentation used , and in principle any other random process engineered in the network . We note that chaos theory is fully applicable with equivalent results in such general stochastic setting ( Arnold , 1988 ) . In the subsequent analysis we will leverage this and work with SGD . Finally , we demonstrate how to extend the method explained in the current section to compute the Lyapunov exponent for SGD with momentum in Appendix B .

Summary: The authors use an insight from chaos theory to derive an efficient method of estimating the largest and smallest eigenvalues of the loss Hessian wrt the weights. To do that, they use nearby weight space positions, optimize for a bit (either gradient climbing or descending), check how quickly the points are departing from each other, and use that to estimate the extreme eigenvalues using a connection to Lyapunov coefficients in chaos theory. Then they use on the fly estimated largest eigenvalue to automatically tune the learning rate of SGD.