Implicit Regularization of SGD via Thermophoresis

A central ingredient in the impressive predictive performance of deep neural networks is optimization via stochastic gradient descent ( SGD ) . While some theoretical progress has been made , the effect of SGD in neural networks is still unclear , especially during the early phase of training . Here we generalize the theory of thermophoresis from statistical mechanics and show that there exists an effective force from SGD that pushes to reduce the gradient variance in certain parameter subspaces . We study this effect in detail in a simple two-layer model , where the thermophoretic force functions to decreases the weight norm and activation rate of the units . The strength of this effect is proportional to squared learning rate and inverse batch size , and is more effective during the early phase of training when the model ’ s predictions are poor . Lastly we test our quantitative predictions with experiments on various models and datasets . 1 INTRODUCTION . Deep neural networks have achieved remarkable success in the past decade on tasks that were out of reach prior to the era of deep learning . Yet fundamental questions remain regarding the strong performance of over-parameterized models and optimization schemes that typically involve only first-order information , such as stochastic gradient descent ( SGD ) and its variants . In particular , optimization via SGD is known in many cases to result in models that generalize better than those trained with full-batch optimization . To explain this , much work has focused on how SGD navigates towards so-called flat minima , which tend to generalize better than sharp minima ( Hochreiter & Schmidhuber , 1997 ; Keskar et al. , 2017 ) . This has been argued by nonvacuous PACBayes bounds ( Dziugaite & Roy , 2017 ) and Bayesian evidence ( Smith & Le , 2018 ) . More recently , Wei & Schwab ( 2019 ) discuss how optimization via SGD pushes models to flatter regions within a minimal valley by decreasing the trace of the Hessian . However , these perspectives apply to models towards the end of training , whereas it is known that proper treatment of hyperparameters during the early phase is vital . In particular , when training a deep network one typically starts with a large learning rate and small batch size if possible . After training has progressed , the learning rate is annealed and decreased so that the model can be further trained to better fit the training set ( Krizhevsky et al. , 2012 ; Simonyan & Zisserman , 2015 ; He et al. , 2016b ; a ; You et al. , 2017 ; Vaswani et al. , 2017 ) . Crucially , using a small learning rate during the first phase of training usually leads to poor generalization and also result in large gradient variance practically ( Jastrzebski et al. , 2020 ; Faghri et al. , 2020 ) . However , limited theoretical work has been done to understand the effect of SGD on the early phase of training . Jastrzebski et al . ( 2020 ) argue for the existence of a “ break-even ” point on an SGD trajectory . This point depends strongly on the hyperparameter settings . They argue that the breakeven point with large learning rate and small batch size tends to have a smaller leading eigenvalue of the Hessian spectrum , and this eigenvalue sets an upper bound for the leading eigenvalue beyond this point . They also present experiments showing that large learning rate SGD will reduce the variance of the gradient . However their analysis focuses only on the leading eigenvalue of the Hessian spectrum and requires the strong assumption that the loss function in the leading eigensubspace is quadratic . Meanwhile Li et al . ( 2020 ) studied the simple setting of two-layer neural networks . They demonstrate that in this model , training with large learning rate in the early phase tends to result in better generalization than training with small learning rate . To explain this , they hypothesize a separation of features in the data : easy-to-generalize yet hard-to-fit features , and hard-to-generalize , easierto-fit features . They argue that a model trained with small learning rate will memorize easy-togeneralize , hard-to-fit patterns during phase one , and then generalize worse on hard-to-generalize , easier-to-fit patterns , while the opposite scenario occurs when training with large learning rate . However , this work relies heavily on the existence of these two distinct types of features in the data and the specific network architecture . Moreover , their analysis focuses mainly on learning rate instead of the effect of SGD . In this paper , we study the dynamics of model parameter motion during SGD training by borrowing and generalizing the theory of thermophoresis from physics . With this framework , we show that during SGD optimization , especially during the early phase of training , the activation rate of hidden nodes is reduced as is the growth of parameter weight norm . This effect is proportional to squared learning rate and inverse batch size . Thus , thermophoresis in deep learning acts as an implicit regularization that may improve the model ’ s ability to generalize . We first give a brief overview of the theory of thermophoresis in physics in the next section . Then we generalize this theory to models beyond physics and derive particle mass flow dynamics microscopically , demonstrating the existence of thermophoresis and its relation to relevant hyperparameters . Then we focus on a simple two-layer model to study the effect of thermophoresis in detail . Notably , we find the thermophoretic force is strongest during the early phase of training . Finally , we test our theoretical predictions with a number of experiments , finding strong agreement with the theory . 2 THERMOPHORESIS IN PHYSICS . Thermophoresis , also known as the Soret effect , describes particle mass flow in response to both diffusion and a temperature gradient . The effect was first discovered in electrolyte solutions ( Ludwig , 1859 ; Soret , 1897 ; Chipman , 1926 ) . However it was discovered in other systems such as gases , colloids , and biological fluids and solid ( Janek et al. , 2002 ; Köhler & Morozov , 2016 ) . Thermophoresis typically refers to particle diffusion in a continuum with a temperature gradient . In one method of analysis , the non-uniform steady-state density ρ is given by the ” Soret Equilibrium ” ( Eastman , 1926 ; Tyrell & Colledge , 1954 ; Wurger , 2014 ) , ∇ρ+ ρST∇T = 0 , ( 1 ) where T is temperature and ST is called the Soret coefficient . In other work by de Groot & Mazur ( 1962 ) , mass flow was calculated by non-equilibrium theory . They considered two types of processes for entropy balance : a reversible process stands for the entropy transfer and an irreversible process corresponds to the entropy production , or dissipation . The resulting mass flow induced by diffusion and temperature gradient was found to be J = −D∇ρ− ρDT∇T , ( 2 ) where D is the Einstein diffusion coefficient and DT is defined as thermal diffusion coefficient . Comparing the steady state in 1 and setting the flow to be zero , the Soret coefficient is simply ST = DT D . ( 3 ) The Soret coefficient can be calculated from molecular interaction potentials based on specific molecular models ( Wurger , 2014 ) . 3 THERMOPHORESIS IN GENERAL . In this section , we first study a kind of generalized random walk that has evolution equations for a particle state with coordinate q = { qi } i=1 , ... , n as qt+1 = qt − ηγf ( qt , ξ ) , ( 4 ) where f is a vector function , γ and ξ are random variables , and η is a small number controlling the step size . Notice that this is a generalized inhomogeneous random walk for the particle . Before further analysis , it is noted that the evolution equations 4 is similar to SGD updates in machine learning and we will show this in the next section . To isolate the effect of thermophoresis , we assume the random walk is unbiased , in which case P ( γf ( q , ξ ) = a ) = P ( γf ( q , ξ ) = −a ) , ( 5 ) for an arbitrary vector a . Thus there is no explicit force exerted on the particle . This simplification was used to demonstrate a residual thermophoretic force in the absence of a gradient . Including gradients is straightforward and corresponds to an external field that creates a bias term . We also denote the probability density , which we also call the mass density , as ρ ( q ) and gi ( q ) : = √∫ γ2f2i ( q , ξ ) dµ ( γ , ξ ) , ( 6 ) so that ηgi ( q ) is the standard deviation of the random walk in the ith direction . From a position q , we consider a subset of coordinate indices , U ⊆ { 1 , . . . , n } , wherein sign ( fi ( q , x ) ) = sign ( fj ( q , x ) ) and ∂igj ( q ) ≥ 0 ( 7 ) for all i , j ∈ U . We note here that indices will correspond to parameters when we study learning dynamics . The first property is necessary for our derivation . The second condition will be used at the end to conclude that each gi decreases . In order to study the dynamics of the particle and its density function , we focus on the probability mass flow induced by the inhomogeneous random walk . We will show that there is always a flow from regions with larger gi ( q ) to those with smaller gi ( q ) for i ∈ U , which is a generalization of thermophoresis in physics . Since η 1 , the movement of the particle will have a mean free path of gi ( q ) in ith direction . Therefore the random walk equation 4 becomes qi = qi − ηgi ( q ) ζi , ( 8 ) where i = 1 , . . . , n and ζi is a binary random variable with P ( ζi = −1 ) = P ( ζi = 1 ) = 0.5 . Moreover , from Eq . 7 , we also have that ζi = ζj for all i and j ∈ U . Next we will show that the flow projecting on the subspace U is always toward smaller gi ( q ) . Notice that although U can be multi-dimensional , the degree of freedom of the particle dynamics is 1 within U due to the sharing of the ζs , and therefore the mass flow projecting on it is also 1-dimensional . For each i ∈ U , we define the average flow in this dimension to be the mass that enters qi from q−i minus the mass from the opposite direction q+i . From Eq . 8 and the assumption that η 1 , only mass close to qi will move across qi at each step . We let the farthest mass that will flow across qi in step i be qi + ∆+i and qi − ∆ − i , where ∆ + i and ∆ + i are positive . ∆ + i and ∆ − i are thus defined implicitly by the equations : ∆+i = ηgi ( q+∆ + ) and ∆−i = ηgi ( q−∆− ) , respectively . Notice that if the random walk were homogeneous , we would have ∆+i = ∆ − i . In our inhomogeneous case , we have ∆+i ∼ ∆ − i ∼ ηgi ( q ) up to leading order of η , and the next to leading order will be calculated in order to compute the difference between ∆+i and ∆ − i . Now we are ready to calculate the mass flow through q . The mass flow projecting onto the subspace U is calculated by the mass through q from q + ∆+ minus the mass from q−∆− where ∆+i and ∆−i are as above for i ∈ U . It is straightforward to show that1 ∆+i −∆ − i = 2η 2 ∑ j∈U gj ( q ) ∂jgi ( q ) +O ( η 3 ) . ( 9 ) With this , we can compute the flow density , J , through q , finding J = −η2 √∑ i∈U g2i ( q ) ∑ i∈U gi ( q ) ∂iρ ( q ) − η2 ∑ i , j∈U gi ( q ) gj ( q ) ∂jgi ( q ) √∑ i∈U g 2 i ( q ) ρ ( q ) +O ( η3 ) , ( 10 ) where the derivation can be found in Appendix A.3 . This can be understood as described in Diagram 1 . Notice that this probability mass flow consists of two terms at order η2 . The first represents diffu- sion and the second corresponds to our goal in this section , namely the flow due to thermophoresis . By the second property of the gi in Eq . 7 , we find that the coefficient of thermophoresis ( Soret coefficient ) , which is defined as c : = −η2 ∑ i , j∈U gi ( q ) gj ( q ) ∂jgi ( q ) 2 √∑ i∈U g 2 i ( q ) ( 11 ) ≤ 0 , ( 12 ) is negative . This means that there is an effective force exerted on a particle at position q towards the smaller variance regime ( by analogy , the colder area ) . The coefficient is proportional to η2 .

Summary: This paper proposes that SGD has the implicitly bias of reducing gradient variance via the phenomenon of thermophoresis that masses tend to flow from regions with higher temperature / variance of random walk to regions with lower temperature / variance of random walk. In the setup of two-layer neural networks trained by SGD for binary classification, the authors show the analogous phenomenon that the model is biased towards smaller activation rate and the norm of its second layer weight. The dependence of the rate at which this phenomenon happens on the learning rate and batch size are verified in the experiments.