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 studies the convergence of SGD with momentum under linear regression with covariate shift. Data comes from $Y_i = <\theta^*, X_i> + \epsilon_i$ for a fixed $\theta^*$ and iid zero-mean noise $\epsilon_i$, but the distribution of $X_i$ varies over time ("covariate shift"), leading to non-iid samples. In this setting, the paper shows that the expected progress of SGD with momentum is a discretization of a second-order ODE. Then, the paper proceeds to show that if the covariate shift is periodic, the convergence/divergence of the ODE can be determined by looking at the spectral radius of a matrix called the monodromy matrix. The paper moves on to show experiments that test the theoretical characterizations and also investigate if different settings under relaxed assumptions show similar "resonance" behavior.  The keyword "resonance-driven divergence" in this paper can be understood as follows: divergence of SGD with momentum happens only at some specific frequency of covariate shift. Metaphorically, this is similar to breaking a wine glass with a sound wave tuned to the right frequency.