Long Expressive Memory for Sequence Modeling

1 INTRODUCTION . Learning tasks with sequential data as inputs ( and possibly outputs ) arise in a wide variety of contexts , including computer vision , text and speech recognition , natural language processing , and time series analysis in the sciences and engineering . While recurrent gradient-based models have been successfully used in processing sequential data sets , it is well-known that training these models to process ( very ) long sequential inputs is extremely challenging on account of the so-called exploding and vanishing gradients problem ( Pascanu et al. , 2013 ) . This arises as calculating hidden state gradients entails the computation of an iterative product of gradients over a large number of steps . Consequently , this ( long ) product can easily grow or decay exponentially in the number of recurrent interactions . Mitigation of the exploding and vanishing gradients problem has received considerable attention in the literature . A classical approach , used in Long Short-Term Memory ( LSTM ) ( Hochreiter & Schmidhuber , 1997 ) and Gated Recurrent Units ( GRUs ) ( Cho et al. , 2014 ) , relies on gating mechanisms and leverages the resulting additive structure to ensure that gradients do not vanish . However , gradients might still explode , and learning very long-term dependencies remains a challenge for these architectures ( Li et al. , 2018 ) . An alternative approach imposes constraints on the structure of the hidden weight matrices of the underlying recurrent neural networks ( RNNs ) , for instance by requiring these matrices to be unitary or orthogonal ( Henaff et al. , 2016 ; Arjovsky et al. , 2016 ; Wisdom et al. , 2016 ; Kerg et al. , 2019 ) . However , constraining the structure of these matrices might lead to significantly reduced expressivity , i.e. , the ability of the model to learn complicated inputoutput maps . Yet another approach relies on enforcing the hidden weights to lie within pre-specified bounds , leading to control on gradient norms . Examples include Li et al . ( 2018 ) , based on independent neurons in each layer , and Rusch & Mishra ( 2021a ) , based on a network of coupled oscillators . Imposing such restrictions on weights might be difficult to enforce , and weight clipping could reduce expressivity significantly . This brief survey highlights the challenge of designing recurrent gradient-based methods for sequence modeling which can mitigate the exploding and vanishing gradients problem , while at the same time being sufficiently expressive and possessing the ability to learn complicated input-output maps efficiently . We seek to address this challenge by proposing a novel gradient-based method . The starting point for our method is the observation that realistic sequential data sets often contain information arranged according to multiple ( time , length , etc. , depending on the data and task ) scales . Indeed , if there were only one or two scales over which information correlated , then a simple model with a parameter chosen to correspond to that scale ( or , e.g. , scale difference ) should be able to model the data well . Thus , it is reasonable to expect that a multiscale model should be considered to process efficiently such multiscale data . To this end , we propose a novel gradientbased architecture , Long Expressive Memory ( LEM ) , that is based on a suitable time-discretization of a set of multiscale ordinary differential equations ( ODEs ) . For this novel gradient-based method ( proposed in Section 2 ) : • we derive bounds on the hidden state gradients to prove that LEM mitigates the exploding and vanishing gradients problem ( Section 4 ) ; • we rigorously prove that LEM can approximate a very large class of ( multiscale ) dynamical systems to arbitrary accuracy ( Section 4 ) ; and • we provide an extensive empirical evaluation of LEM on a wide variey of data sets , including image and sequence classification , dynamical systems prediction , keyword spotting , and language modeling , thereby demonstrating that LEM outperforms or is comparable to state-of-the-art RNNs , GRUs and LSTMs in each task ( Section 5 ) . We also discuss a small portion of the large body of related work ( Section 3 ) , and we provide a brief discussion of our results in a broader context ( Section 6 ) . Much of the technical portion of our work is deferred to Supplementary Materials . 2 LONG EXPRESSIVE MEMORY . We start with the simplest example of a system of two-scale ODEs , dy dt = τy ( σ ( Wyz + Vyu + by ) − y ) , dz dt = τz ( σ ( Wzy + Vzu + bz ) − z ) . ( 1 ) Here , t ∈ [ 0 , T ] is the continuous time , 0 < τy ≤ τz ≤ 1 are the two time scales , y ( t ) ∈ Rdy , z ( t ) ∈ Rdz are the vectors of slow and fast variables and u = u ( t ) ∈ Rm is the input signal . For simplicity , we set dy = dz = d. The dynamic interactions between the neurons are modulated by weight matrices Wy , z , Vy , z , bias vectors by , z and a nonlinear tanh activation function σ ( u ) = tanh ( u ) . Note that refers to the componentwise product of vectors . However , two scales ( one fast and one slow ) , may not suffice in representing a large number of scales that could be present in realistic sequential data sets . Hence , we need to generalize ( 1 ) to a multiscale version . One such generalization is provided by the following set of ODEs , dy dt = σ̂ ( W2y + V2u + b2 ) ( σ ( Wyz + Vyu + by ) − y ) , dz dt = σ̂ ( W1y + V1u + b1 ) ( σ ( Wzy + Vzu + bz ) − z ) . ( 2 ) In addition to previously defined quantities , we need additional weight matrices W1,2 , V1,2 , bias vectors by , z and sigmoid activation function σ̂ ( u ) = 0.5 ( 1 + tanh ( u/2 ) ) . As σ̂ is monotone , we can set W1,2 = V1,2 ≡ 0 and ( b1 ) j = by , ( b2 ) j = bz , for all 1 ≤ j ≤ d , with σ̂ ( by , z ) = τy , z to observe that the two-scale system ( 1 ) is a special case of ( 2 ) . One can readily generalize this construction to obtain many different scales in ( 2 ) . Thus , we can interpret ( τ z ( y , t ) , τ y ( y , t ) ) = ( σ̂ ( W1y + V1u + b1 ) , σ̂ ( W2y + V2u + b2 ) ) in ( 2 ) as input and state dependent gating functions , which endow ODE ( 2 ) with multiple time scales . These scales can be learned adaptively ( with respect to states ) and dynamically ( in time ) . Moreover , it turns out that the multiscale ODE system ( 2 ) is of the same general form ( see SM§C ) as the well-known HodgkinHuxley equations modeling the dynamics of the action potential for voltage-gated ion-channels in biological neurons ( Hodgkin & Huxley , 1952 ) . Next , we propose a time-discretization of the multiscale ODE system ( 2 ) , providing a circuit to our sequential model architecture . As is common with numerical discretizations of ODEs , doing so properly is important to preserve desirable properties . To this end , we fix ∆t > 0 , and we discretize ( 2 ) with the following implicit-explicit ( IMEX ) time-stepping scheme to arrive at LEM , written in compact form as , ∆tn = ∆tσ̂ ( W1yn−1 + V1un + b1 ) , ∆tn = ∆tσ̂ ( W2yn−1 + V2un + b2 ) , zn = ( 1−∆tn ) zn−1 + ∆tn σ ( Wzyn−1 + Vzun + bz ) , yn = ( 1−∆tn ) yn−1 + ∆tn σ ( Wyzn + Vyun + by ) . ( 3 ) For steps 1 ≤ n ≤ N , the hidden states in LEM ( 3 ) are yn , zn ∈ Rd , with input state un ∈ Rm . The weight matrices are W1,2 , z , y ∈ Rd×d and V1,2 , z , y ∈ Rd×m and the bias vectors are b1,2 , z , y ∈ Rd . We also augment LEM ( 3 ) with a linear output state ωn ∈ Ro with ωn = Wyyn , and Wy ∈ Ro×d . 3 RELATED WORK . We start by comparing our proposed model , LEM ( 3 ) , to the widely used LSTM of Hochreiter & Schmidhuber ( 1997 ) . Observe that ∆tn , ∆tn in ( 3 ) are similar in form to the input , forget and output gates in an LSTM ( see SM§D ) , and that LEM ( 3 ) has exactly the same number of parameters ( weights and biases ) as an LSTM , for the same number of hidden units . Moreover , as detailed in SM§D , we show that by choosing very specific values of the LSTM gates and the ∆tn , ∆tn terms in LEM ( 3 ) , the two models are equivalent . However , this analysis also reveals key differences between LEM ( 3 ) and LSTMs , as they are equivalent only under very stringent assumptions . In general , as the different gates in both LSTM and LEM ( 3 ) are learned from data , one can expect them to behave differently . Moreover in contrast to LSTM , LEM stems from a discretized ODE system ( 2 ) , which endows it with ( gradient ) stable dynamics . The use of multiscale neural network architectures in machine learning has a long history . An early example was provided in Hinton & Plaut ( 1987 ) , who proposed a neural network with each connection having a fast changing weight for temporary memory and a slow changing weight for long-term learning . More recently , one can view convolutional neural networks as multiscale architectures for processing multiple spatial scales in data ( Bai et al. , 2020 ) . The use of ODE-based learning architectures has also received considerable attention in recent years with examples such as continuous-time neural ODEs ( Chen et al. , 2018 ; Queiruga et al. , 2020 ; 2021 ) and their recurrent extensions ODE-RNNs ( Rubanova et al. , 2019 ) , as well as RNNs based on discretizations of ODEs ( Chang et al. , 2018 ; Erichson et al. , 2021 ; Chen et al. , 2020 ; Lim et al. , 2021 ; Rusch & Mishra , 2021a ; b ) . In addition to the specific details of our archiecture , we differ from other discretized ODE-based RNNs in the explicit use of multiple ( learned ) scales in LEM . 4 RIGOROUS ANALYSIS OF LEM . Bounds on hidden states . The structure of LEM ( 3 ) allows us to prove ( in SM§E.1 ) that its hidden states satisfy the following pointwise bound . Proposition 4.1 . Denote tn = n∆t and assume that ∆t ≤ 1 . Further assume that the initial hidden states are z0 = y0 ≡ 0 . Then , the hidden states zn , yn of LEM ( 3 ) are bounded pointwise as , max 1≤i≤d max { |zin| , |yin| } ≤ min ( 1 , ∆ √ tn ) , ∀1 ≤ n , with ∆ = 1 + ∆t√ 2−∆t . ( 4 ) On the exploding and vanishing gradient problem . For any 1 ≤ n ≤ N , let Xn ∈ R2d , denoted the combined hidden state , given by Xn = [ z1n , y 1 n , . . . . . . , z i n , y i n , . . . . . . , z d n , y d n ] . For simplicity of the exposition , we consider a loss function : En = 12‖yn − yn‖2 , with yn being the underlying ground truth . The training of our proposed model ( 3 ) entails computing gradients of the above loss function with respect to its underlying weights and biases θ ∈ Θ = [ W1,2 , y , z , V1,2 , y , z , b1,2 , y , z ] , at every step of the gradient descent procedure . Following Pascanu et al . ( 2013 ) , one uses chain rule to show , ∂En ∂θ = ∑ 1≤k≤n ∂E ( k ) n ∂θ , ∂E ( k ) n ∂θ = ∂En ∂Xn ∂Xn ∂Xk ∂+Xk ∂θ ( 5 ) In general , for recurrent models , the partial gradient ∂E ( k ) n ∂θ , which measures the contribution to the hidden state gradient at step n arising from step k of the model , can behave as ∂E ( k ) n ∂θ ∼ γn−k , for some γ > 0 Pascanu et al . ( 2013 ) . If γ > 1 , then the partial gradient grows exponentially in sequence length , for long-term dependencies k < < n , leading to the exploding gradient problem . On the other hand , if γ < 1 , then partial gradients decays exponentially for k < < n , leading to the vanishing gradient problem . Thus , mitigation of the exploding and vanishing gradient problem entails deriving bounds on the gradients . We start with the following upper bound ( proved in SM§E.2 ) , Proposition 4.2 . Let zn , yn be the hidden states generated by LEM ( 3 ) . We assume that ∆t < < 1 is chosen to be sufficiently small . Then , the gradient of the loss function En with respect to any parameter θ ∈ Θ is bounded as ∣∣∣∣∂En∂θ ∣∣∣∣ ≤ ( 1 + Ŷ ) tn + ( 1 + Ŷ ) Γt2n , Ŷ = ‖yn‖∞ , η = max { ‖W1‖∞ , ‖W2‖∞ , ‖Wz‖∞ , ‖Wy‖∞ } , Γ = 2 ( 1 + η ) ( 1 + 3η ) ( 6 ) If we choose the hyperparameter ∆t = O ( n−1 ) ( see SM ( 17 ) for the order-notation ) , then one readily observes from ( 6 ) that the gradient ∂θEn is uniformly bounded for any sequence length n and the exploding gradient problem is clearly mitigated for LEM ( 3 ) . Even if one chooses ∆t = O ( n−s ) , for some 0 ≤ s ≤ 1 , we show in SM Remark E.1 that the gradient can only grow polynomially ( e.g . as O ( n ) for s = 1/2 ) , still mitigating the exploding gradient problem . Following Pascanu et al . ( 2013 ) , one needs a more precise characterization of the partial gradient ∂θE ( k ) n , for long-term dependencies , i.e. , k < < n , to show mitigation of the vanishing gradient problem . In SM§E.3 , we state and prove proposition E.2 , which provides a precise formula for the asymptotics of the partial gradient . Here , we illustrate this formula in a special case as a corollary , Proposition 4.3 . Let yn , zn be the hidden states generated by LEM ( 3 ) and the ground truth satisfy yn ∼ O ( 1 ) . Then , for any k < < n ( long-term dependencies ) we have , ∂E ( k ) n ∂θ = O ( ∆t 3 2 ) . ( 7 ) Here , constants in O ( ∆t 32 ) depend on only on η ( 6 ) and η = ‖W2‖1 and are independent of n , k. This formula ( 7 ) shows that although the partial gradient can be small , i.e. , O ( ∆t 32 ) , it is in fact independent of k , ensuring that long-term dependencies contribute to gradients at much later steps and mitigating the vanishing gradient problem . Universal approximation of general dynamical systems . The above bounds on hidden state gradients show that the proposed model LEM ( 3 ) mitigates the exploding and vanishing gradients problem . However , this by itself , does not guarantee that it can learn complicated and realistic input-output maps between sequences . To investigate the expressivity of the proposed LEM , we will show in the following proposition that it can approximate any dynamical system , mapping an input sequence un to an output sequence on , of the ( very ) general form , φn = f ( φn−1 , un ) , on = o ( φn ) , ∀ 1 ≤ n ≤ N , ( 8 ) with φn ∈ Rdh , on ∈ Rdo denoting the hidden and output states , respectively . The input signal is un ∈ Rdu and maps f : Rdh × Rdu 7→ Rdh and o : Rdh 7→ Rdo are Lipschitz continuous . For simplicity , we set the initial state φ0 = 0 . Proposition 4.4 . For all 1 ≤ n ≤ N , let φn , on be given by the dynamical system ( 8 ) with input signal un . Under the assumption that there exists a R > 0 such that max { ‖φn‖ , ‖un‖ } < R , for all 1 ≤ n ≤ N , then for any given > 0 there exists a LEM of the form ( 3 ) , with hidden states yn , zn ∈ Rdy and output state ωn = Wyyn ∈ Rdo , for some dy such that the following holds , ‖on − ωn‖ ≤ , ∀1 ≤ n ≤ N. ( 9 ) From this proposition , proved in SM§E.4 , we conclude that , in principle , the proposed LEM ( 3 ) can approximate a very large class of dynamical systems . Universal approximation of multiscale dynamical systems . While expressing a general form of input-output maps between sequences , the dynamical system ( 8 ) does not explicitly model dynamics at multiple scales . Instead , here we consider the following two-scale fast-slow dynamical system of the general form , φn = f ( φn−1 , ψn−1 , un ) , ψn = τg ( φn , ψn−1 , un ) , on = o ( ψn ) . ( 10 ) Here , 0 < τ < < 1 and 1 are the slow and fast time scales , respectively . The underlying maps ( f , g ) : Rdh×dh×du 7→ Rdh are Lipschitz continuous . In the following proposition , proved in SM§E.5 , we show that LEM ( 3 ) can approximate ( 10 ) to desired accuracy . Proposition 4.5 . For any 0 < τ < < 1 , and for all 1 ≤ n ≤ N , let φn , ψn , on be given by the two-scale dynamical system ( 10 ) with input signal un . Under the assumption that there exists a R > 0 such that max { ‖φn‖ , ‖ψn‖ , ‖un‖ } < R , for all 1 ≤ n ≤ N , then for any given > 0 , there exists a LEM of the form ( 3 ) , with hidden states yn , zn ∈ Rdy and output state ωn ∈ Rdo with ωn = Wyn such that the following holds , ‖on − ωn‖ ≤ , ∀1 ≤ n ≤ N. ( 11 ) Moreover , the weights , biases and size ( number of neurons ) of the underlying LEM ( 3 ) are independent of the time-scale τ . This argument can be readily generalized to more than two time scales ( see SM Proposition E.4 ) . Hence , we show that , in principle , the proposed model LEM ( 3 ) can approximate multiscale dynamical systems , with model size being independent of the underlying timescales . These theoretical results for LEM ( 3 ) point to the ability of this architecture to learn complicated multiscale inputoutput maps between sequences , while mitigating the exploding and vanishing gradients problem . Although useful prerequisities , these theoretical properties are certainly not sufficient to demonstrate that LEM ( 3 ) is efficient in practice . To do this , we perform several benchmark evaluations , and we report the results below .

Summary: This paper introduces a model class of neural networks whose architecture is defined by a circuit that computes a discretized step of a pair of multi-scale ODEs. The key novelty is on the multi-scale aspect of this system. In terms of representation capability, this class can represent LSTMs (and vice-versa) and can also approximate any dynamical system from a very broad class of dynamical systems (as well as multi-scale dynamical systems). The authors also prove upper bounds in l_infinity on both the hidden state values and the gradients, and argue this addresses the exploding gradient problem. They also address the vanishing gradient problem by arguing that the magnitude of the partial gradient corresponding to the contribution that depends on step k of the circuit does not depend on k.  They then present a variety of experiments using this new architectural component for modeling several sequence problems. They include a small study on the multi-scale nature of the datasets they study as well.