Parallel Training of GRU Networks with a Multi-Grid Solver for Long Sequences

1 INTRODUCTION . Recently , the model complexity of Deep Neural Networks ( DNN ) have grown to keep pace with the increasing scale of data . Hence , parallelization of DNNs is becoming important to achieve high quality models in tractable computational run times ( Ben-Nun & Hoefler , 2019 ) . Sequence modeling is one learning task that suffers from substantial computational time due to data size . The goal is to capture temporal dependencies within a sequence . Several different deep learning approaches addressing sequence tasks have been proposed , including 3D-CNNs ( Carreira & Zisserman , 2017 ) , LSTM ( Srivastava et al. , 2015 ) , and Transformers ( Girdhar et al. , 2019 ) . Focusing on sequence classification , this paper proposes a parallelization strategy for a recurrent neural network ( RNN ) architecture known as the gated recurrent unit ( GRU ) developed by Cho et al . ( 2014 ) . Parallel algorithms for training neural networks are generally broken into two classes ; dataparallelism distributes a batch across processors , and model-parallelism distributes the architecture ( Ben-Nun & Hoefler , 2019 ) . Data-parallelism inherent in mini-batch stochastic gradient descent ( SGD ) is the primary strategy for parallelizing GRU ( You et al. , 2019 ; Huang et al. , 2013 ) . However , the accuracy of mini-batch SGD degrades as the number of sequences in each mini-batch increases . Moreover , although multiple mini-batches are processed in parallel , the training procedure of each sequence ( within a mini-batch ) is invariably sequential . Alternatively , the GRU model can be parallelized across hidden layers ( Wu et al. , 2016 ; MXNet , 2018 ) . This model-parallelism approach is beneficial when the network requires a large number of layers . However , even largescale RNNs such as the Google Neural Machine Translation ( GNMT ) developed by Wu et al . ( 2016 ) have only 8 LSTM layers . Therefore , the scalability of layer-based model-parallelism is limited . A common problem in data-parallel and model-parallel algorithms for GRU is that the training execution time increases with the sequence length . At the heart of this difficulty is a fundamental need to process sequences in a serial fashion , yielding forward and backward propagation algorithms that are inherently sequential , and thus hard to parallelize . We propose a parallel-in-time ( PinT ) training method for GRU networks with long sequences . To achieve this , we adapt a multigrid reduction in time ( MGRIT ) solver for forward and backward prop- agation . Within the numerical methods community , a resurgence in PinT has paralleled increasing computational resources ( Gander , 2015 ; Ong & Schroder , 2020 ) . These efforts have even been applied to the training of neural Ordinary Differential Equations ( ODEs ) ( Schroder , 2017 ; Gunther et al. , 2020 ; Kirby et al. , 2020 ; Cyr et al. , 2019 ) , where inexact forward and back propagation is exploited to achieve parallel speedups . To date , these techniques have not been applied to RNNs . Following Jordan et al . ( 2021 ) , GRU networks can be written as an ODE to facilitate application of MGRIT . Different from existing parallel training algorithms for RNNs , our MGRIT parallel training scheme partitions a sequence into multiple shorter sub-sequences and distributes each sub-sequence to different processors . By itself , this provides local improvement of the hidden states , yet global errors remain . To correct these errors , propagation is computed on a coarse representation of the input sequence requiring less computational work while still providing an improvement to the original hidden states . This process is iterated to the accuracy required for training . Applying this algorithm to the classic GRU networks will achieve parallelism . However , this training algorithm will not lead to accurate networks due to stability problems on the coarse grid . This emphasizes the challenges of choosing a proper coarse grid model for neural networks , and multigrid algorithms in general . To alleviate this , we develop a new GRU architecture , or discretization of the ODE , which we refer to as Implicit GRU that handles stiff modes in the ODE . This is required for application of the MGRIT algorithm where coarse sub-sequences providing the correction , correspond to discretizations with more strict numerical stability restrictions . We also compare the accuracy of serial versus parallel inference . This ensures that the network is not compensating for the error introduced by the MGRIT procedure , and is a study that has not been considered in prior work . Two datasets are chosen for evaluation of the PinT algorithm . The first is the UCI-HAR dataset for the human activity recognition using smartphones ( Anguita et al. , 2013 ) . This small dataset is used to demonstrate the algorithm can obtain similar accuracy to classical training approaches , while also demonstrating speedups . The second learning problem uses the HMDB51 dataset for human activity classification in videos ( Kuehne et al. , 2011 ) . As the size of each image in the video is too large to directly use it as an input for a RNN-based model , we use a pre-trained CNN model to generate low-dimensional input features for the GRU network . This approach has the advantage of fully capturing spatio-temporal information in videos ( Yue-Hei Ng et al. , 2015 ; Donahue et al. , 2015 ; Wang et al. , 2016 ) . In our experiments , extracted image features from the ResNets pre-trained on ImageNet are used as input sequences for GRU networks . Using this architecture on the HMDB51 dataset , the MGRIT algorithm demonstrates a 6.5× speedup in training over the serial approach while maintaining accuracy . The paper is organized as follows . Section 2 presents background on GRUs and their formulation as ODEs . Section 3 details PinT training for GRU using the MGRIT method . Section 4 compares the performance of the proposed approach on sequence classification tasks . We conclude in Section 5 . The appendix provides greater detail on the MGRIT method including a serial implementation and parameters used in the experiments . 2 GATED RECURRENT UNIT . The GRU network , proposed by Cho et al . ( 2014 ) , is a type of RNN and a state-of-the-art sequential model widely used in applications such as natural language processing and video classification . Each GRU cell at time t takes input data xt and the previous hidden state ht−1 , and computes a new hidden state ht . The cell resets or updates information carefully regulated by gating units in response to prior hidden state and input sequence data stimulation as follows : rt = σ ( Wirxt + bir +Whrht−1 + bhr ) zt = σ ( Wizxt + biz +Whzht−1 + bhz ) nt = ϕ ( Winxt + bin + rt ( Whnht−1 + bhn ) ) ht = zt ht−1 + ( 1− zt ) nt . ( 1 ) We refer to this method as Classic GRU in the text . GRUs to ODEs Recent work has shown how RNNs can be modified to obtain an ODE description ( Chang et al. , 2019 ; Habiba & Pearlmutter , 2020 ) . Of particular interest here is those works for GRU networks ( De Brouwer et al. , 2019 ; Jordan et al. , 2021 ) where the ODE is ∂tht = − ( 1− zt ) ht︸ ︷︷ ︸ stiff mode + ( 1− zt ) nt . ( 2 ) Jordan et al . ( 2021 ) present a detailed description of the network behavior . For our purposes , note that the first term can generate stiff dynamics that drive the hidden state h towards a manifold specified by the data . The second term injects changes from the data . The term zt controls the decay rate towards the manifold for each component of the hidden state . If zt is near 1 , then the rate is slow . If zt is close to zero the decay rate is rapid and the first term can be stiff , constraining the choice of time step ( for numerical treatment of stiff ODE ’ s see Wanner & Hairer ( 1996 ) ) . The classic GRU method in Eq . 1 is rederived using an explicit Euler scheme with a unit time step to integrate Eq . 2 . Adding and subtracting ht−1 on the right hand side of the hidden state update , and for notational purposes multiplying by γ = 1 , yields ht = Φγ ( xt , ht−1 ) : = ht−1 + γ ( − ( 1− zt ) ht−1 + ( 1− zt ) nt ) . ( 3 ) The notation Φγ , introduced here , denotes the classic GRU hidden state update if γ = 1 . 3 PARALLEL-IN-TIME TRAINING . Training GRU networks is typically done with the BPTT algorithm ( Sutskever , 2013 ) . BPTT performs a forward evaluation of the neural network , followed by backpropagation to compute the gradient . Both forward and backward propagation are computationally expensive algorithms that do not naturally lend themselves to parallelization in the time domain . Despite this , we develop a methodology achieving the parallel distribution shown in Fig . 1 . This distribution is enabled by computing the forward and backward propagation inexactly , trading inexactness for parallelism and ultimately run time speedups . MGRIT , a multigrid method , is used to compute the hidden states for both forward and backward propagation . Multigrid methods are used in the iterative solution of partial differential equations ( Brandt , 1977 ; Douglas , 1996 ; Briggs et al. , 2000 ; Trottenberg et al. , 2000 ) . Recent multigrid advances have yielded the MGRIT method achieving parallel-in-time evolution of ODEs like Eq . 2 ( see for instance Falgout et al . ( 2014 ) ; Dobrev et al . ( 2017 ) ; Gunther et al . ( 2020 ) ) . Thus , multigrid and MGRIT methods are reasonably well developed ; however , they remain quite complex . Here , we briefly outline the algorithm for forward propagation , and present a brief motivation for the method in Appendix A.1 . In Section 3.3 , an implicit formulation for GRU that handles the stiff modes is proposed . This is the key advance that enables the use of MGRIT for training GRU networks . 3.1 NOTATION AND PRELIMINARIES . Denote a sequence by ~u = { ut } Tt=0 where the length T will be clear from context , and u0 = 0 by assumption . This notation is used for the input data sequence ~x and the sequence of hidden states ~h . The GRU forward propagator in Eq . 3 can be written as the system of nonlinear algebraic equations f ( ~h ) = { ht − Φ1 ( xt , ht−1 ) } Tt=1 . ( 4 ) When f ( ~h ) = 0 , then ~h is the hidden state for the GRU networks with input sequence ~x . A two-level multigrid algorithm uses a fine sequence on level l = 0 , corresponding to a data input of length T . The coarse level , defined on level l = 1 , has data of length T/cf where cf is an integer coarsening factor . Sequences are transferred from the fine to coarse level using a restriction operator ~U = R ( ~u ) : = { ucf t } T/cf t=1 . ( 5 ) Applying R to ~x gives a coarse sub-sequence of the original . Time points on the fine level that are included in the coarse level are referred to as coarse points ( see Fig . 2 ) . Transfers from coarse to fine grids are denoted by the prolongation operator ~u = P ( ~U ) : = { Ubt/cfc } T t=1 . ( 6 ) Prolongation fills intermediate values using replication of the prior coarse point . The nonlinear system from Eq . 4 is modified using ∆t = cf , yielding the residual for the coarse network : fc ( ~H ) = { Ht − Φcf ( Xt , Ht−1 ) } T/cf t=1 . ( 7 ) Notice that in addition to a larger time step , the length of sequence ~H is smaller than the fine grid ’ s . Finally , we introduce the notion of relaxation of the hidden state . Here , we define this formally , while in Fig . 2 a schematic view of these algorithms is taken . Relaxation applies propagation to an initial condition for the hidden state in a parallel fashion . F -relaxation is defined in Algorithm 1 as ΨF ( cf , ~x , ~h ) and demonstrated in the top of Fig . 2 . F -relaxation uses the hidden state at the coarse points , and fills in the fine points up to the next coarse point ( exclusive ) . FC-relaxation , ΨFC , is identical to F -relaxation except the next coarse point is included ( the bottom row in Fig . 2 , where closed circles fall on the coarse grid point ) . Combining F - and FC-relaxation yields FCF - relaxation used in MGRIT to ensure sequence length independent convergence ~h′ = ΨFCF ( cf , ~x , ~h ) : = ΨF ( cf , ~x , ΨFC ( cf , ~x , ~h ) ) . ( 8 ) The available parallelism for all relaxation methods is equal to the number of coarse points . Algorithm 1 : ΨF ( γ , ~x , ~h ) - F -relaxation 1. parallel for tc ← 1 to T/γ 2. t = γtc 3. h′t = ht 4. for τ ← t to t+ γ − 1 5. h′τ+1 = Φγ ( xτ , h ′ τ ) 6. return ~h′

Summary: The paper proposed to parallelize the inference and training of GRU networks (a type of recurrent neural networks) at the `time` dimension.  The main contribution is the application of multigrid reduction in time (MGRIT) solver, and a new GRU architecture (Implicit GRU) that handles the stiffness in the architecture. As a result, the evaluation shows 6.5 times faster training time.