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: The authors propose a new RNN architecture, long expressive memory (LEM), motivated by a system of ODEs with multiple time constants. They prove that it can avoid the vanishing gradient problem while retaining the flexibility to approximate a broad class of dynamical systems. They report comparable or improved prediction performance of LEM-based sequence models across a very wide variety of tasks, as compared to several recent alternatives.