Riemannian Stochastic Gradient Descent for Tensor-Train Recurrent Neural Networks

1 INTRODUCTION . Recurrent Neural Networks ( RNNs ) are typically composed of large weight matrices of fullyconnected and recurrent layers , thus massive training data as well as exhaustive computational resources are required . The Tensor-Train factorization ( TTF ) aims to reduce the redundancy of RNN parameters by reshaping large weight matrices into high-dimensional tensors before factorizing them in a Tensor-Train format Oseledets ( 2011 ) . The notation of Tensor-Train usually suggests that TTF is applied for the tensor representation of model parameters . Tensor-Train was initially applied to fully-connected layers Novikov et al . ( 2015 ) , and it has been recently generalized to recurrent layers in RNNs such as LSTM and GRU Yu et al . ( 2017 ) . Compared with other tensor decomposition techniques like the CANDECOMP/PARAFAC decomposition Kolda & Bader ( 2009 ) and Tucker decomposition Kim & Choi ( 2007 ) , Tensor-Train can be easily scaled to arbitrarily high dimensions and have the advantage of computational tractability to significantly large weight matrices . Given a vector of Tensor-Train ranks r = ( r1 , r2 , · · · , rd+1 ) , TTF decomposes a d-dimensional tensor W ∈ R ( m1·n1 ) × ( m2·n2 ) ×···× ( md·nd ) into a multiplication of core tensors according to ( 1 ) , where the k-th core tensor C [ k ] ∈ Rrk×mk×nk×rk+1 , and any index pair ( ik , jk ) satisfies 1 ≤ ik ≤ mk , 1 ≤ jk ≤ nk . Additionally , the ranks r1 and rd+1 are fixed to 1 . W ( ( i1 , j1 ) , ( i2 , j2 ) , ... , ( id , jd ) ) = C [ 1 ] ( r1 , i1 , j1 , r2 ) C [ 2 ] ( r2 , i2 , j2 , r3 ) ···C [ d ] ( rd , id , jd , rd+1 ) ( 1 ) Thus , when the TTF technique is applied to the fully-connected ( FC ) layer with feed-forward weight matrix W , a tensor W is firstly converted from W and is then decomposed to a multiplication of the core tensors as shown in ( 2 ) , where C [ t ] is the t-th core tensor , X denotes a tensor of input , B refers to a tensor of bias , the tensor of outputs Ŷ ∈ Rn1×n2×···×nd , and σ is a sigmoid function . For clarity , the notation TTL ( W , X ) is used to simplify the representation of a Tensor-Train fully-connected layer , which is shown in ( 3 ) . Ŷ ( j1 , j2 , ... , jd ) = σ ( m1∑ i1=1 · · · md∑ id=1 d∏ t=1 C [ t ] ( rt , it , jt , rt+1 ) X ( i1 , i2 , ... , id ) + B ( j1 , j2 , ... , jd ) ) ( 2 ) Ŷ = σ ( TTL ( W , X ) ) . ( 3 ) Likely , we use ( 4 ) to represent an RNN with a feed-forward weight matrix W and a recurrent weight matrix U . In ( 4 ) , Xt is an input matrix at time t , ht−1 and ht separately denote the hidden vectors of time t− 1 and t , and σ refers to the sigmoid function . ht = σ ( TTL ( W , Xt ) + TTL ( U , ht−1 ) ) . ( 4 ) The largest benefit from Tensor-Train models is the capability to reduce the model parameters tremendously . For example , a Tensor-Train FC layer needs ∑ imi · ni · ri · ri+1 parameters in total . In comparison , the total number of parameters of an FC layer is about O ( ∏ imi · ∏ i ni ) , which is much larger than the associated Tensor-Train FC one . The Tensor-Train models are found widespread . For example , Yang et al . ( 2017 ) set up Tensor-Train Recurrent Neural Networks for video classification , and Yu et al . ( 2017 ) applied the Tensor-Train as an End-to-End dynamic model for multi-variate forecasting environmental data . However , those applications focused on simple deep learning architectures , and a vector of Tensor-Train ranks r is element-wise fixed . When more complex deep models with numerous Tensor-Train layers are involved , it is not easy to find appropriate Tensor-Train ranks for each core tensor . However , when applying a vector of shared and high Tensor-Train ranks to all parameters of Tensor-Train layers , each of them has a vector of lower Tensor-Train ranks , which results in an additional decrease in the number of parameters . This work applies Riemannian Stochastic Gradient Descent ( RSGD ) to the RNN Tensor-Train layers . Unlike Stochastic Gradient Descent ( SGD ) , which conducts the update of parameters in the Euclidean space , RSGD finds optimal core tensors of parameters associated with a vector of lower Tensor-Train ranks in Riemannian Manifold , thereby decreasing the number of total parameters of Tensor-Train RNNs . Our work is inspired by Bonnabel ( 2013 ) and Lubich et al . ( 2015 ) , which are two excellent introductions to the theory of Stochastic Gradient Descent on Riemannian Manifolds . Additionally , Zhang et al . ( 2016 ) was a recent work that proposed a fast stochastic optimization on Riemannian Manifolds . Besides , Absil et al . ( 2009 ) is a useful reference of Riemannian optimization . Contributions . We summarize the key contributions of this paper as follows : • This work applies RSGD to iteratively find vectors of lower Tensor-Train ranks with the update of parameters in the training process . The RSGD algorithm and the related theoretical analysis are also presented . • We design Bi-directional Tensor-Train GRU/LSTM Chung et al . ( 2014 ) , and EncoderDecoder Tensor-Train RNNs with a Tensor-Train Attention mechanism Luong et al . ( 2015 ) . We apply the RSGD algorithm to the Tensor-Train RNN models on the digit recognition and machine translation tasks . To the best of our knowledge , this is the first work that applies RSGD to train Tensor-Train RNNs to find the optimal core tensors of parameters with vectors of lower Tensor-Train ranks . Moreover , this is the first work that builds Tensor-Train RNNs with complex architectures for natural language processing ( NLP ) tasks . 2 RIEMANNIAN STOCHASTIC GRADIENT DESCENT . The optimization problem of Tensor-Train RNNs can be formulated as a Riemannian optimization problem as shown in ( 5 ) , where { X , Y } is a data sequence with length T , W represents a tensor of parameters which lies in a d-dimensional Riemannian Manifold ( M , µ ) with a Riemannian measure µ , and the Tensor-Train ranks for core tensors of W must be element-wise no higher than the vector r = ( r1 , r2 , ... , rd+1 ) as shown in ( 5 ) . min W∈M f ( W ; X , Y ) s.t. , tt rank ( W ) ≤ r. ( 5 ) Besides , the Riemannian measure µ induces an inner product structure in each tangent space TxM associated with a tensor x ∈M . Specifically , ∀u , v ∈ TxM , the inner product < u , v > = µx ( u , v ) . Algorithm 1 Riemannian Stochastic Gradient Descent 1 . Given the labeled input data ( X , Y ) with sequence length T , and the learning rate η . 2 . W ( 0 ) ← the randomly initialized core tensors { C [ 1 ] , C [ 2 ] , ... , C [ d ] } . 3 . For t = 1 , 2 , ... , T : 4 . For i = 1 , 2 , ... , d : 5 . Choose a gradient gC [ i ] = ∇C [ i ] f ( W ( t ) ; X , Y ) in the tangent space TC [ i ] M . 6 . Âi = C [ i ] − ηgC [ i ] . 7 . Ĉ [ i ] ← ExpC [ i ] ( Âi ) . 8 . ( Ĉ [ i ] , r̂i ) ← rounding ( Ĉ [ i ] ) . 9 . W ( t ) ← { Ĉ [ 1 ] , Ĉ [ 2 ] , · · ·Ĉ [ d ] } . 10. r← ( r̂1 , r̂2 , ... r̂d ) . 11 . Reshape the core tensors { C [ 1 ] , C [ 2 ] , ... , C [ d ] } based on the updated r. 12 . Return W ( T+1 ) = { C [ 1 ] , C [ 2 ] , · · · , C [ d ] } . Similarly , µ induces the norm of u ∈ TxM as ||u|| = √ µx ( u , u ) ≥ 0 . In addition , the µ induced inner product and the norm preserve the basic properties like definiteness , homogeneity and triangle inequality . Algorithm 1 presents the RSGD Algorithm . The algorithm mainly consists of two main procedures : one is the update of parameters in the tangent space and conducting an exponential mapping , and the second one is the rounding to lower Tensor-Train ranks . As illustrated in Figure 1 , step 5 firstly obtains a gradient gC [ i ] on a tangent space TCiM at the core tensor C [ i ] in Riemannian Manifold ( M , µ ) , and step 6 conducts a gradient descent on the tangent space to generate a new tensor Âi . Step 7 projects Âi back to Ĉ [ i ] in Riemannian Manifold ( M , µ ) by an exponential mapping . Finally , as shown in Figure 2 , the rounding function in step 8 transforms Ĉ [ i ] in the submanifold Sr to the core tensor Ĉ [ i ] with a vector of lower Tensor-Train ranks r̂ in a new submanifold Sr̂ . Note that the vectors of Tensor-Train ranks r and r̂ span two submanifolds Sr ⊂ M and Sr̂ ⊂ M respectively . After that , the next iteration of the the parameter update is conducted in Sr̂ . The exponential mapping in Algorithm 1 is formulated in ( 6 ) . Unfortunately , it is not easy to solve the problem because we have to deal with the calculus of variations , or we have to know the Christoffel symbols Lubich et al . ( 2015 ) . Therefore , a fast and straightforward retraction method is applied as a first-order approximation to the exponential mapping as shown in Algorithm 2. min Y∈Sr⊂Md ||W− Y||F s.t. , tt rank ( Y ) = r. ( 6 ) Algorithm 2 The Retraction Algorithm 1 . Given the core tensors { C [ 1 ] , C [ 2 ] , · · · , C [ d ] } and the Tensor-Train rank r̂ = { r1 , r2 , · · · , rd+1 } . 2 . For i = 1 to d : 3 . Ai ← reshape C [ i ] by the shape of ( ri · ni ·mi ) × ri+1 . 4 . For i = 1 to d : 5 . Âi , Σi ← QR decomposition ( Ai ) . 6 . C [ i ] ← reshape Âi by ( ri , ni , mi , ri+1 ) . 7 . Ai+1 ← multiply ( Σi , Ai+1 ) , if i < d. 8 . Return { C [ 1 ] , C [ 2 ] , · · · , C [ d ] } . Algorithm 3 The Rounding Algorithm 1 . Given the core tensors { C [ 1 ] , C [ 2 ] , · · · , C [ d ] } and a constant maximum rank rmax . 2 . Initialize r̂ = { 1 , 1 , ... , 1 } . 3 . For i = 1 to d : 4 . Ai ← reshape C [ i ] by the shape of ( ri · ni ·mi ) × ri+1 . 5 . For i = 1 to d : 6. col num ( C [ i ] ) = ni ·mi · r̂i+1 . 7. row num ( C [ i ] ) = ri · ri+1 r̂i+1 . 8. r̂i+1 = min ( rmax , col num ( C [ i ] ) , row num ( C [ i ] ) ) . 9 . ( Ui , Si , Vi ) = SVD ( Ai ) . 10 . Ûi = Ui [ : , 1 : r̂i+1 ] , Ŝi = Si [ 1 : r̂i+1 , 1 : r̂i+1 ] , V̂i = Vi [ 1 : r̂i+1 , : ] . 11 . Âi = ÛiŜiV̂i . 12 . C [ i ] ← reshape Âi as Rr̂i×ni×mi×r̂i+1 . 13 . Return the updated core tensors { C [ 1 ] , C [ 2 ] , · · · , C [ d ] } . Algorithm 2 presents the retraction algorithm . The main idea of the retraction algorithm is to orthogonalize the core tensors { C [ 1 ] , C [ 2 ] , · · · , C [ d ] } in a left-to-right order by the QR decomposition . The rounding algorithm is shown in Algorithm 3 . Similar to the retraction algorithm , a left-to-right tensor matricization is firstly initialized . Then , the Tensor-Train rank is updated before conducting an SVD computation . The returned core tensors { C [ 1 ] , C [ 2 ] , · · · , C [ d ] } are based on the updated vector of lower Tensor-Train ranks . Furthermore , the rounding procedure has the property presented in Proposition 1 . Proposition 1 . The rounding procedure of Algorithm 3 does not change values of the objective function f for the Tensor-Train RNNs . That is , for a tensor x ∈ Sr , the rounding tensor x̂ ∈ Sr̂ , we have f ( x̂ ) = f ( x ) . Proof . Given the weight matrix W with core tensors { C [ 1 ] , C [ 2 ] , · · · , C [ d ] } , for an input tensor X ∈M , we obtain ( 7 ) and ( 8 ) according to ( 1 ) . Y ( j1 , j2 , ... , jd ) = m1∑ i1=1 · · · md∑ id=1 d∏ t=1 C [ t ] ( rt , it , jt , rt+1 ) X ( i1 , i2 , ... , id ) + B ( j1 , j2 , ... , jd ) ( 7 ) = m1∑ i1=1 · · · md∑ id=1 d∏ t=1 C [ t ] ( r̂t , it , jt , r̂t+1 ) X ( i1 , i2 , ... , id ) + B ( j1 , j2 , ... , jd ) , ( 8 ) which suggests that a vector of Tensor-Train ranks determines a submanifold for generating core tensors of the tensor W , but the values of the objective functions are invariant to the change of the vector of Tensor-Train ranks , obtaining Proposition 1 .

Summary: In this paper, the authors proposed a new method to update the weights in RNN by SGD on Riemannian manifold.  Due to the properties of manifold learning, the updated weights in each iteration are contracted with a low-rank structure, such that the number of the parameters of TT can be automatically decreased during the training procedure. By using the new algorithm, the authors modified two types of sophisticated RNNs, i.e., bi-directional GRU/LSTM and Encoder-Decoder RNN. The experimental results validate effectiveness of the proposed method. How to determine the rank of the tensor networks in weight compression problem is indeed an important and urgent task, this paper does not clearly illustrate how RSGD can efficiently solve this problem.