Resonance in Weight Space: Covariate Shift Can Drive Divergence of SGD with Momentum

1 INTRODUCTION . Stochastic gradient descent ( SGD ) ( Robbins & Monro , 1951 ) – and its variants such as Adagrad ( Duchi et al. , 2011 ) , ADAM ( Kingma & Ba , 2014 ) and RMSprop ( Hinton et al. , 2012 ) – are very widely used optimization algorithms across machine learning . SGD is conceptually straightforward , easy to implement , and often performs well in practice . Among the variants of SGD , accelerated versions based on Polyak ’ s or Nesterov ’ s acceleration ( Polyak , 1964 ; Nesterov , 1983 ) , known generally as Stochastic Gradient Descent with Momentum ( SGDm ) , are used widely due to the improvements in convergence rate they offer . SGDm can give up to a quadratic speedup to SGD on many functions , and is in fact optimal among all methods having only information about the gradient at consecutive iterates for convex and Lipschitz continuous optimization problems ( Nesterov , 2004 ; Goh , 2017 ) . SGDm has the same computational complexity as SGD , but exhibits superior convergence rates under reasonable assumptions ( Su et al. , 2014 ) . These convergence results for SGDm , however , rely on independent and identically distributed ( iid ) sampling . Little is known about the convergence properties of SGDm under non-iid sampling , yet the non-iid setting is critical . In many machine learning problems it is expensive or impossible to obtain iid samples . In online learning ( Rakhlin et al. , 2010 ) and reinforcement learning ( RL ) ( Sutton & Barto , 2018 ) , the data becomes available in a sequential order and there is a temporal dependence among the samples . There is a particularly strong temporal dependence in RL , where observed states are sampled according to the transition dynamics of the underlying Markov decision process ( MDP ) . Federated learning ( Hsieh et al. , 2020 ) and time-series learning ( Kuznetsov & Mohri , 2014 ) provide further examples of when non-iid sampling is essential to the learning problem . Without momentum , SGD ’ s convergence rate has been examined under non-iid sampling with the stochastic approximation framework in ( Benveniste et al. , 2012 ; Kushner & Yin , 2003 ) , and more recently under specific assumptions of ergodicity ( Duchi et al. , 2012 ) or Markovian sampling ( Nagaraj et al. , 2020 ; Doan et al. , 2020b ; Sun et al. , 2018 ) . Convergence rates under Markovian sampling are also known for ADAM-type algorithms when applied to policy gradient and temporal difference learning ( Xiong et al. , 2020 ) . To our knowledge , however , there has been little work on providing convergence rates or guarantees for SGDm under non-iid sampling . In ( Doan et al. , 2020a ) , a progress bound is provided for SGDm under Markovian sampling based on mixing time—the time required for a distribution ’ s convergence toward its stationary distribution—along with a convergence rate guarantee under decaying step-sizes . In this work , we assume a fixed step-size , since it is a common choice in the online setting , especially when the practitioner is unsure of mixing rate or stationarity . There is a broad literature using ordinary differential equations ( ODEs ) to analyze gradient descent methods by approximating descent updates as continuous-time flows . Early work is comprehensively discussed in ( Kushner & Yin , 2003 ; Benveniste et al. , 2012 ) . Despite the age and establishment of the linear setting , the gradient flow lens continues to reveal new insights ( e.g . implicit rank reduction ( Arora et al. , 2019 ) ) and new perspectives on old insights ( e.g . regularization of early stopping ( Ali et al. , 2019 ) ) . Recent work has paid particular attention to the flow induced by momentum accelerated methods ( Su et al. , 2014 ; Wibisono et al. , 2016 ; Wilson et al. , 2016 ; Scieur et al. , 2017 ; Muehlebach & Jordan , 2021 ; Diakonikolas & Orecchia , 2019 ; Muehlebach & Jordan , 2019 ; Li et al. , 2017 ; Simsekli et al. , 2020 ; Betancourt et al. , 2018 ; Kovachki & Stuart , 2019 ; Attouch et al. , 2018 ; Shi et al. , 2019 ; Siegel , 2019 ; Zhang et al. , 2018 ; Berthier et al. , 2021 ; Kovachki & Stuart , 2021 ) . In these works , it is demonstrated that linear regression under iid sampling with SGDm can be represented as an ODE resembling a harmonic oscillator , sometimes with a time-decaying damping coefficient . Contributions : In this work , we show that non-iid sampling , due to covariate shift , induces a related system : the parametric oscillator , a harmonic oscillator having coefficients which vary over time in a manner capable of exponentially exciting the system ( Mumford , 1960 ; Csörgő & Hatvani , 2010 ; Halanay , 1966 ) . We characterize the underlying time-varying ODE , which requires a novel approach to incorporate covariate shift , and develop specific conditions on covariate shift that lead to divergence in SGDm . We provide an empirical design , with synthetic data , to systematically test the response of the learning system to different covariate shift frequencies . We empirically validate that resonancedriven divergence occurs in a learning system which aligns well with theoretical assumptions . We follow with similar demonstrations on learning systems which progressively relax further and further away from our theoretical assumptions , including non-periodic covariate shift , learning with neural networks and optimizers other than SDGm . The implications of this work are that resonance can occur in learning systems with momentum , causing instability or poor accuracy , simply due to certain temporal correlations in the input sequence . Our results provide a novel direction to better understand the properties of our learning systems for time varying inputs—namely beyond iid sampling—and point to a gap in the robustness of our algorithms that merits further investigation . 2 PROBLEM SETTING . We investigate the effect of non-iid sampling on a model optimized using SGDm . We assume labelled training data sampled from discreteindexed , real-valued stochastic processes { Xk } k∈N and { Yk } k∈N such that { Xk } k∈N has a unique stationary distribution Π. Yk is a function of Xk with zero-mean observation noise , Yk = f ( Xk ) + k. The learning algorithm does not have access to iid samples from Π . Instead , the learning algorithm may only update parame- ters θk at time k using samples zk = ( xk , yk ) , where Xk ’ s marginal over time converges to Π . The goal is to train a model with parameters θ to minimize an objective function L ( z ; θ ) with respect to the stationary distribution Π : θ∗ = arg minθ∈Rd EΠ [ L ( z ; θ ) ] = ∫ R L ( z ; θ ) dΠ ( x ) . An example of this problem setting is illustrated in Figure 1 . This setting is identical to that explored in prior work on Markovian sampling ( Sun et al. , 2018 ; Xiong et al. , 2020 ; Nagaraj et al. , 2020 ) , but we do not require that our stochastic processes adhere to the Markov property . Note , also , that the underlying functional relationship between the inputs , Xk , and the targets , Yk , does not change over time . That is , the non-iid sampling is a result of covariate shift rather than a changing relationship between the inputs and targets . We consider Polyak ’ s Heavy Ball method ( Polyak , 1964 ) with the formulation from Sutskever et al . ( 2013 ) , namely SGD with momentum ( SGDm ) vk+1 = µvk − η∇θL ( zk ; θk ) θk+1 = θk + vk+1 ( 1 ) where η ≥ 0 is the learning rate and µ ∈ ( 0 , 1 ) the momentum coefficient . Our goal is to understand the behavior of SGDm when learning on this non-iid sequence zk . 3 COVARIATE SHIFT AS A DRIVING FORCE . In this section we characterize SGDm as a discretization of a particular parametric oscillator in continuous time . A parametric oscillator resembles a harmonic oscillator ODE , but has time-varying system coefficients which are capable of driving the system . It is well understood that parametric oscillators can suffer from global solution instability due to coefficients oscillating at particular frequencies ( Halanay , 1966 ) , a condition known as parametric resonance . We show parametric resonance conditions sufficient to induce exponential divergence in SGDm . We provide proof sketches and defer complete proofs to Appendix A.3 . Notation : Capital letters are matrices and random variables . { Xk } is a stochastic process , shorthand for { Xk } k∈N . k and t index discrete and continuous time , respectively . When a discrete sequence and a function approximate each other , they share the same symbol and are differentiated by subscript k and function argument t , e.g . the sequence { θk } and the function θ ( t ) . Pk is the joint distribution of Xk and Yk . For loss function L ( z ; θ ) on training pair z = ( x , y ) and weights θ , we denote the time-varying expected gradient gk ( θ ) : = EPk [ ∇θL ( z ; θ ) ] 3.1 COVARIATE SHIFT INDUCES A LINEAR TIME-VARYING EXPECTED GRADIENT . We begin by showing that linear least squares regression with covariate shift and a fixed target induces a linear time-varying expected loss gradient . That is , the gradient function gk ( θ ) = Bk ( θ − θ∗ ) for some matrix Bk that varies with time , due to covariate shift . Assumption 1 . The covariate generating process { Xk } k∈N is not identically distributed : Corr ( Xk1 , Xk1 ) 6= Corr ( Xk2 , Xk2 ) for some k1 , k2 ∈ N. Assumption 2 . The targets Yk are a fixed linear function of Xk with iid zero-mean observation noise , k such that E [ k ] = 0 . That is , Yk = 〈θ∗ , Xk〉+ k where θ∗ ∈ Rd is fixed for all k. Proposition 1 . Under Assumptions 1 , 2 , a linear model with weights θk making predictions Ŷk = 〈θk , Xk〉 with a mean squared error ( MSE ) objective will induce time-varying linear expected loss gradients gk ( θk ) = Bk ( θk − θ∗ ) for all k ∈ N , where Bk ∈ Rd×d and Bk1 6= Bk2 for some k1 , k2 . Proof Sketch : We fix a time step k and take the expectation of the MSE loss over observation noise , and over the distribution of Xk . A time-varying expected loss surface remains , and the linear setting implies the loss gradient is a linear function for all steps k , with Bk depending on Xk : Bk = 2Corr ( Xk , Xk ) = 2Cov ( Xk , Xk ) + 2EPk [ Xk ] EPk [ Xk ] T ( 2 ) 3.2 ODE CORRESPONDENCE . Next , we show that the iterates { θk } generated by SGDm are a first order numerical integration of a particular ODE . The procedure is similar to ( Muehlebach & Jordan , 2021 ) , but due to the time-varying loss gradient , we must pay specific attention to the conditions on the ODE necessary to have integration consistency , . Let B ( t ) be a matrix-valued function , Lipschitz continuous in t , such that { Bk } are samples spaced √ η apart , i.e . Bk = B ( √ ηk ) . The existence and uniqueness of this B ( t ) are ensured if we assume the { Xk } are sampled from an underlying continuous-time X ( t ) . Proposition 2 . The SGDm iterates { θk } numerically integrate the ODE system in equation 3 with integration step √ η and first order consistency . θ̈ ( t ) + 1− µ √ η θ̇ ( t ) +B ( t ) ( θ ( t ) − θ∗ ) = 0 ( 3 ) Proof Sketch : First , the ODE in equation 3 is converted to first order linear form , and split into a sum of two separate systems . The sum ’ s terms have implicit and explicit Euler integrations ; using operator splitting , these can be composed into a consistent numerical integrator for linear time-varying systems . Then , it is shown that the composed integrator is precisely equivalent to SGDm when the integration time step is √ η . The first order consistency guarantee in Proposition 2 means that when the step-size η is small and initial conditions agree between continuous and discrete time , i.e . θ0 = θ ( 0 ) , the continuous and discrete time trajectories approximate each other as θk ≈ θ ( k √ η ) , where the difference between them depends on the integration time step h = √ η ( Hairer et al. , 2006 ) . Therefore , assuming θ0 = θ ( 0 ) , the difference accrued in one step of k ( local error ) is θ1 = θ ( √ η ) +O ( η ) and over arbitrarily many steps ( global error ) is θk = θ ( √ ηk ) +O ( ηk ) .

Summary: This paper analyses the phenomenon of resonance in momentum SGD (SGDm) under a time-dependent covariate shift. This setting is useful to depart from iid-sampling-based theory, and could have useful implications for continual learning and reinforcement learning. In the paper, it is shown that such a setting corresponds to a _parametric_ oscillator of the parameters, which explains instability and divergence in a non-iid SGDm setting. The authors then test their hypothesis empirically, first in a simple setting that respects all the assumptions of the theorem, and then progressively getting rid of more and more assumptions, which ends up in finding a similar (albeit heavily dampened) phenomenon in non-linear ReLU-MLPs using Adam.