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 describes a technique for evaluating GRU networks based on the multigrid reduction in time (MGRIT) technique. These techniques are not new, in general or to neural network training, but the contribution here is their application to GRU layers. After presenting some of the theory behind ordinary differential equation (ODE) representations of the GRU and laying out the mathematical framework for the MGRIT method, the method is evaluated on two datasets. The results show impressive scalability up to 32 processes (CPU-only) at the cost of a slight loss of accuracy compared to the traditional, sequential GRU