Stochastic Gradient Descent Learns State Equations with Nonlinear Activations

1 INTRODUCTION . A wide range of problems involve sequential data with a natural temporal ordering . Examples include natural language processing , time series prediction , system identification , and control design , among others . State-of-the-art algorithms for sequential problems often stem from dynamical systems theory and are tailored to learn from temporally dependent data . Linear models and algorithms ; such as Kalman filter , PID controller , and linear dynamical systems , have a long history and are utilized in control theory since 1960 ’ s with great success ( Brown et al . ( 1992 ) ; Ho & Kalman ( 1966 ) ; Åström & Hägglund ( 1995 ) ) . More recently , nonlinear models such as recurrent neural networks ( RNN ) found applications in complex tasks such as machine translation and speech recognition ( Bahdanau et al . ( 2014 ) ; Graves et al . ( 2013 ) ; Hochreiter & Schmidhuber ( 1997 ) ) . Unlike feedforward neural networks , RNNs are dynamical systems that use their internal state to process inputs . The goal of this work is to shed light on the inner workings of RNNs from a theoretical point of view . In particular , we focus on the RNN state equation which is characterized by a nonlinear activation function φ , state weight matrixA , and input weight matrixB as follows ht+1 = φ ( Aht +But ) , ( 1.1 ) Here ht is the state vector and ut is the input data at timestamp t. This equation is the source of dynamic behavior of RNNs and distinguishes RNN from feedforward networks . The weight matrices A and B govern the dynamics of the state equation and are inferred from data . We will explore the statistical and computational efficiency of stochastic gradient descent ( SGD ) for learning these weight matrices . Contributions : Suppose we are given a finite trajectory of input/state pairs ( ut , ht ) Nt=0 generated from the state equation ( 1.1 ) . We consider a least-squares regression obtained from N equations ; with inputs ( ut , ht ) Nt=1 and outputs ( ht+1 ) N t=1 . For a class of activation functions including leaky ReLU and for stable systems1 , we show that SGD linearly converges to the ground truth weight matrices while requiring near-optimal trajectory length N . In particular , the required sample size is O ( n+ p ) where n and p are the dimensions of the state and input vectors respectively . The results are extended to unstable systems when the samples are collected from multiple independent RNN trajectories rather than a single trajectory . Our theory applies to increasing activation functions 1Throughout this work , a system is called stable if the spectral norm of the state matrixA is less than 1. whose derivatives are bounded away from zero , which includes leaky ReLU , and Gaussian input data . Numerical experiments on ReLU and leaky ReLU corroborate our theory and demonstrate that SGD converges faster as the activation slope increases . To obtain our results , we i ) characterize the statistical properties of the state vector ( e.g . well-conditioned covariance ) and ii ) derive a novel SGD convergence result with nonlinear activations ; which may be of independent interest . As a whole , this paper provides a step towards foundational understanding of RNN training via SGD . 1.1 RELATED WORK . Our work is related to the recent optimization and statistics literature on linear dynamical systems ( LDS ) and neural networks . Linear dynamical systems : The state-equation ( 1.1 ) reduces to a LDS when φ is the linear activation ( φ ( x ) = x ) . Identifying the weight matrices is a core problem in linear system identification and is related to the optimal control problem ( e.g . linear quadratic regulator ) with unknown system dynamics . While these problems are studied since 1950 ’ s ( Ljung ( 1998 ; 1987 ) ; Åström & Eykhoff ( 1971 ) ) , our work is closer to the recent literature that provides data dependent bounds and characterize the non-asymptotic learning performance . Recht and coauthors have a series of papers exploring optimal control problem ( Simchowitz et al . ( 2018 ) ; Tu et al . ( 2018 ; 2017 ) ) . In particular , Hardt et al . ( 2016 ) shows gradient descent learns single-input-single-output ( SISO ) LDS with polynomial guarantees . Oymak & Ozay ( 2018 ) and Faradonbeh et al . ( 2018 ) provide sample complexity bounds for learning LDS . Sanandaji et al . ( 2011b ; a ) ; Pereira et al . ( 2010 ) study the identification of sparse systems . Neural networks : There is a growing literature on the theoretical aspects of deep learning and provable algorithms for training neural networks . Most of the existing results are concerned with feedforward networks . Ge et al . ( 2017 ) ; Li & Yuan ( 2017 ) ; Mei et al . ( 2018b ) ; Soltanolkotabi ( 2017 ) ; Janzamin et al . ( 2015 ) ; Zhong et al . ( 2017b ) consider learning fully-connected shallow networks with gradient descent . Mei et al . ( 2018a ) ; Soltanolkotabi et al . ( 2017 ) ; Foster et al . ( 2018 ) analyze empirical landscape of related nonlinear learning problems . Brutzkus & Globerson ( 2017 ) ; Zhong et al . ( 2017a ) ; Du et al . ( 2017 ) ; Goel et al . ( 2018 ) address convolutional neural networks ; which is an efficient weight-sharing architecture . Brutzkus et al . ( 2017 ) ; Wang et al . ( 2018 ) studies over-parameterized networks when data is linearly separable . Janzamin et al . ( 2015 ) ; Oymak & Soltanolkotabi ( 2018 ) utilize tensor decomposition techniques for learning feedforward neural nets . For recurrent networks , Sedghi & Anandkumar ( 2016 ) proposed tensor algorithms with polynomial guarantees and Khrulkov et al . ( 2017 ) studied their expressive power . More recently , Miller & Hardt ( 2018 ) showed that stable RNNs can be approximated by feed-forward networks . 2 PROBLEM SETUP . We first introduce the notation . ‖ · ‖ returns the spectral norm of a matrix and smin ( · ) returns the minimum singular value . The activation φ : R → R applies entry-wise if its input is a vector . Throughout , φ is assumed to be a 1-Lipschitz function . With proper scaling of its parameters , the system ( 1.1 ) with a Lipschitz activation can be transformed into a system with 1-Lipschitz activation . The functions Σ [ · ] and var [ · ] return the covariance of a random vector and variance of a random variable respectively . In is the identity matrix of size n × n. Normal distribution with mean µ and covariance Σ is denoted by N ( µ , Σ ) . Throughout , c , C , c0 , c1 , . . . denote positive absolute constants . Setup : We consider the dynamical system parametrized by an activation function φ ( · ) and weight matricesA ∈ Rn×n , B ∈ Rn×p as described in ( 1.1 ) . Here , ht is the n dimensional state-vector and ut is the p dimensional input to the system at time t. As mentioned previously , ( 1.1 ) corresponds to the state equation of a recurrent neural network . For most RNNs of interest , the state ht is hidden and we only get to interact with ht via an additional output equation . For Elman networks Elman ( 1990 ) , this equation is characterized by some output activation φy and output weights C , D as follows yt = φy ( Cht +Dut ) . ( 2.1 ) In this work , our attention is restricted to the state equation ( 1.1 ) ; which corresponds to setting yt = ht+1 in the output equation . To analyze ( 1.1 ) in a non-asymptotic data-dependent setup , we assume a finite input/state trajectory of { ut , ht } Nt=0 generated by some ground truth weight matrices Algorithm 1 Learning state equations with nonlinear activations 1 : Inputs : ( yt , ht , ut ) Nt=1 sampled from a trajectory . Scaling µ , learning rate η. Initialization A0 , B0 . 2 : Outputs : Estimates Â , B̂ of the weight matricesA , B . 3 : xt ← [ µhTt uTt ] T for 1 ≤ t ≤ N . 4 : Θ0 ← [ µ−1A0 B0 ] 5 : for τ from 1 to END do 6 : Pick γτ from { 1 , 2 , . . . , N } uniformly at random . 7 : Θτ ← Θτ−1 − η∇Lγτ ( Θτ−1 ) 8 : end for 9 : return [ Â B̂ ] ← ΘEND [ µIn 0 0 Ip ] . ( A , B ) . Our goal is learning the unknown weightsA andB in a data and computationally efficient way . In essence , we will show that , if the trajectory length satisfies N & n+ p , SGD can quickly and provably accomplish this goal using a constant step size . Appoach : Our approach is described in Algorithm 1 . It takes two hyperparameters ; the scaling factor µ and learning rate η . Using the RNN trajectory , we constructN triples of the form { ut , ht , ht+1 } Nt=1 . We formulate a regression problem by defining the output vector yt , input vector xt , and the target parameter C as follows yt = ht+1 , xt = [ µht ut ] ∈ Rn+p , C = [ µ−1AB ] ∈ Rn× ( n+p ) . ( 2.2 ) With this reparameterization , we find the input/output identity yt = φ ( Cxt ) . We will consider the least-squares regression given by L ( Θ ) = 1 N N∑ t=1 Lt ( Θ ) where Lt ( Θ ) = 1 2 ‖yt − φ ( Θxt ) ‖2 ` 2 . ( 2.3 ) For learning the ground truth parameter C , we utilize SGD on the loss function ( 2.3 ) with a constant learning rate η . Starting from an initial point Θ0 , after END SGD iterations , Algrorithm 1 returns an estimate Ĉ = ΘEND . Estimates of A and B are decoded from the left and right submatrices of Ĉ respectively . 3 MAIN RESULTS . 3.1 PRELIMINARIES . The analysis of the state equation naturally depends on the choice of the activation function ; which is the source of nonlinearity . We first define a class of Lipschitz and increasing activation functions . Definition 3.1 ( β-increasing activation ) . Given 1 ≥ β ≥ 0 , the activation function φ satisfies φ ( 0 ) = 0 and 1 ≥ φ′ ( x ) ≥ β for all x ∈ R. Our results will apply to strictly increasing activations where φ is β-increasing for some β > 0 . Observe that , this excludes ReLU activation which has zero derivative for negative values . However , it includes Leaky ReLU which is a generalization of ReLU . Parameterized by 1 ≥ β ≥ 0 , Leaky ReLU is a β-increasing function given by LReLU ( x ) = max ( βx , x ) . ( 3.1 ) In general , given an increasing and 1-Lipschitz activation φ , a β-increasing function φβ can be obtained by blending φ with the linear activation , i.e . φβ ( x ) = ( 1− β ) φ ( x ) + βx . A critical property that enables SGD is that the state-vector covariance Σ [ ht ] is well-conditioned under proper assumptions . The lemma below provides upper and lower bounds on this covariance matrix in terms of problem variables . Lemma 3.2 ( State vector covariance ) . Consider the state equation ( 1.1 ) where h0 = 0 and ut i.i.d.∼ N ( 0 , Ip ) . Define the upper bound term Bt as Bt = ‖B‖ √ 1− ‖A‖2t 1− ‖A‖2 . ( 3.2 ) • Suppose φ is 1-Lipschitz and φ ( 0 ) = 0 . Then , for all t ≥ 0 , Σ [ ht ] B2t In . • Suppose φ is a β-increasing function and p ≥ n. Then , Σ [ ht ] β2smin ( B ) 2In . As a natural extension from linear dynamical systems , we will say the system is stable if ‖A‖ < 1 and unstable otherwise . For activations we consider , stability implies that if the input is set to 0 , state vector ht will exponentially converge to 0 i.e . the system forgets the past states quickly . This is also the reason ( Bt ) t≥0 sequence converges for stable systems and diverges otherwise . The condition number of the covariance will play a critical role in our analysis . Using Lemma 3.2 , this number can be upper bounded by ρ defined as ρ = ( B∞ βsmin ( B ) ) 2 = ( ‖B‖ smin ( B ) ) 2 1 β2 ( 1− ‖A‖2 ) . ( 3.3 ) Observe that , the condition number ofB appears inside the ρ term .

Summary: This work considers the problem of learning a non-linear dynamical system in which the output equals the state.  Under several assumptions (input is Gaussian, non-linear activation is strictly increasing, stable system) it is shown that SGD converges linearly to the ground truth system with near-optimal sample complexity. The proof idea is to reduce this problem to the problem of learning a single non-linear neuron in the case that the covariance matrix of the data is well-conditioned. The main challenge is to show the covariance is well-conditioned under the reduction. In a nutshell, this is done by splitting the trajectory to sub-trajectories with independent states and using results from random matrix theory on matrices with independent rows.