Parallel Training of GRU Networks with a Multi-Grid Solver for Long Sequences
Authors: Euhyun Moon
Eric C Cyr
Venue: ICLR 2022 Poster (2022)
Keywords: GRU
MGRIT
parallel-in-time
distributed machine learning

Abstract: Parallelizing Gated Recurrent Unit (GRU) is a challenging task, as the training procedure of GRU is inherently sequential. Prior efforts to parallelize GRU have largely focused on conventional parallelization strategies such as data-parallel and model-parallel training algorithms. However, when the given sequences are very long, existing approaches are still inevitably performance limited in terms of both training time and model accuracy. In this paper, we present a novel parallel training scheme (called  parallel-in-time) for GRU based on a multigrid reduction in time (MGRIT) solver. MGRIT partitions a sequence into multiple shorter sub-sequences and trains the sub-sequences on different processors in parallel. The key to achieving speedup is a hierarchical correction of the hidden state to accelerate end-to-end communication in both the forward and backward propagation phases of gradient descent. Experimental results on the HMDB51 dataset, where each video is an image sequence, demonstrate that a new parallel training scheme of GRU achieves up to $6.5 \times$ speedup over a serial approach. As efficiency of our new parallelization strategy is associated with the sequence length, our parallel GRU algorithm achieves significant performance improvement as the length of sequence increases. Further, the proposed approach can be applied simultaneously with batch and other forms of model parallelism.

Published as a conference paper at ICLR 2022
PARALLEL TRAINING OF GRU NETWORKS WITH A MULTI-GRID SOLVER FOR LONG SEQUENCES
Gordon Eubyun Moon Korea Aerospace University Chmoon@kau.c aac.kack kr
Eric C. Cyr
Sandia National Laboratories eccyr & sandia gov
ABSTRACT
Parallelizing Gated Recurrent Unit (GRU) networks is challenging task, the training procedure of GRU is inherently sequential. Prior efforts to parallelize GRU have largely focused on conventional parallclization strategies such as data- parallel and model-parallel training algorithms. However, when the given se- quences are very long, existing approaches are still inevitably performance limited in terms of training time. In this paper, we present a novel parallel training scheme (called parallel-in-tin
for GRU based on a multigrid reduction in time (MGRIT) solver. MGRIT partitions sequence into multiple shorter sub-sequences and trains the sub-sequences on different processors in parallel. The key to achieving speedup is hierarchical correction of the hidden state accelerate end-to-end communication in both the forward and backward propagation phases of gradi- ent descent. Experimental results the HMDB51 dataset, where each video is an image sequence, demonstrate that the new parallel training scheme achieves up 6.5x speedup over a serial approach. As efficiency of our new paralleliza- tion strategy is associated with the sequence length, our parallel GRU algorithm achieves significant performance improvement as the sequence length increases.
INTRODUCTION
Recently, the model complexity of Deep Neural Networks (DNN) has 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 2011.. Sequence modeling one learning task that suffers from substantial computational time due to data size. The goal to capture temporal dependencies within a sequence. Several different deep learning approaches addressing sequence tasks have been proposed, including 3D-CNNs (Carreira & Zisserman, (017), LSTM (Srivastava et al. 2015), and Transformers (Girdhar al. (019). Focusing sequence classification, this paper proposes a parallelization strategy for a gated recurrent unit (GRU) which special type of recurrent neural network (RNN) architecture.
Parallel algorithms for training neural networks are generally broken into two classes; data- parallelism distributes batch across processors, and model-parallelism distributes the architec- ture (Ben-Nun & Hoefler (201)). Data-parallelism inherent in mini-batch stochastic gradient de- scent (SGD) is the primary strategy for parallelizing RNN architecture (You et al. 2019 Huang jet al. 2013). However, the accuracy of mini-batch SGD degrades as the number of sequences in each mini-bateh 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 RNN 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 large-scale RNN such as the Google Neural Machine Translation (GNMT) developed by Wu ct al. 2016) have only 8 LSTM layers. Therefore, the scalability of layer-based model-parallelism is limited. A common problem data-parallel and model
algorithms for GRU that the training execution time increases with the sequence length. At the heart of this difficulty is fun- damental need to process sequences in serial fashion, yielding forward and backward propagation algorithms that inherently sequential, and thus hard to parallelize.

---

Published as a conference paper at ICLR 2022
We propose parallel-in-timg (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, resurgence in PinT has paralleled increasing computational resources (Gander 2015 Ong & Schroder 2020). These efforts have even been ap- plied to the training of neural Ordinary Differential Equations (ODEs) (Schroder [017] Gunther et al. 2020; Kirby et al. 2020 Cyr et al., 2019), where inexact forward and back propagation is exploited achieve parallel speedups. To date, these techniques have not been applied to RNN.
Following Jordan et al. (222), GRU networks can be written as an ODE to facilitate application of MGRIT. Different from existing parallel training algorithms for RNN, our MGRIT parallel training scheme partitions 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 in- put 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 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, 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 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 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., (201]). As the size of each image in the video is too large to directly use it ate dindenesiional input features for the GRU network. This approach captures spatio-temporal information in videos (Yue-Hei Ng et al. 2015 Donahue et al. 2013) Wang et al., (0166). In experiments, extracted image features from the ResNet pre-trained on ImageNet are used as input se- quences for GRU networks. Using this architecture on the HMDB51 dataset, the MGRIT algorithm demonstrates a 6.5x speedup in training over the serial approach while maintaining accuracy.
input for a RNN-based model, we use pre-trained CNN model to gener-
The paper is organized as follows. Section 2
presents background 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.
GATED RECURRENT UNIT
The GRU network, proposed by Cho et al. (2014), is a type of RNN and a state-of-the- sequential model widely used in applications such as natural language processing and video classification. Each GRU cell time takes input data x, and the previous hidden state 14-1, and computes new hidden state hy. The cell resets or updates information carefully regulated by gating units in response to prior hidden state and input sequence data stimulation as follows:
Ft = a (W e + bar + W + bhr)
oWii + biz + Whahh-1 + bbz))
(I)
ny Wn + b + re @ (Wanhe-1 + bbn))
h4 O h + 2)) @ n4-
We refer to this method as Classic GRU in the text.
GRUs to ODEs
Recent work has shown how RNN can be modified to obtain an ODE descrip-
tion (Chang et al. 2019 Habiba & Pearlmutter 2020). Of particular interest here is those works for
2

---

Published as a conference paper at ICLR 2022
Processor
Processor 2
Processor P-1
Processor P
hg
T
17
Figure I: Parallel distribution of work in parallel-in-time training.
GRU networks De Brouwer et al.
2019 0,h,
Jordan et al.
2021 where the ODE is
2+) h2
+ 24) T4--
stiff mode
Jordan et al. (202I) 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 2; controls the decay rate towards the manifold for each component of the hidden state. If 2 is near 1, then the rate is slow. If 2 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. 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 y = 1, yields
h = (2,, h2-1) == h4-1 + 7 (-(1- 24) O h4-1 + (1 21) O n4).
(3)
The notation 4, introduced here, denotes the classic GRU hidden state update if 1.
3 PARALLEL-IN-TIMAE TRAINING
Training GRU networks is typically done with the backpropagation through time (BPTT) algo- rithm (Sutskever) 2013). BPTT performs a forward evaluation of the neural network, followed by backpropagation to compute the gradient. Both forward and backward propagation are computa- tionally 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, D This distribution is enabled by computing the forward and backward propagation inexactly, trading inexactness for parallelism and ultimately run time speedups. MGRIT, 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 (Brandi 1977; Douglas (996//Brige 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 Fal- gout et al. 2014); Dobrev et al. (2017); Gunther et al. (2020). Thus, multigrid and MGRIT methods 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 [ In Section 3.3 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 sequence by 77 {wr) =o where the length T will be clear from context, and uo by assumption. This notation is used for the input data sequence and the sequence of hidden states h. The GRU forward propagator in Eq. 3 can be written the system of nonlinear algebraic equations
f(h) (h+ (1 h--1)
(4)
When f(h 0, then h is the hidden state for the GRU networks with input sequence 7.
A two-level multigrid algorithm uses a fine sequence on level / = O, corresponding to a data input of length T. The coarse level, defined level / = 1, has data of length T/cs where Cy is an integer coarsening factor. Sequences are transferred from the fine to coarse level using a restriction operator
Ú R( (Mert
T/o5
(5)
3

---

Published as a conference paper at ICLR 2022
F-Relaxation
4
4
4
47
17
11
17
17
17
12
K1
K1
K1
FC-Relaxation
Figure 2: For Cy = with T = 16, this image shows F- and FC-reelxation. In both cases, forward propagation performed serially starting at the open green circle and ending at the closed one. Each propagation step is performed in parallel, with four-way parallelism depicted in the image.
Applying R. i gives coarse sub-sequence of the original. Time points the fine level that 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
7 = P(ū) == (UW+e;])]--
Prolongation fills intermediate values using replication of the prior coarse point. The nonlinear system from Eq. modified using At = cf. yielding the residual for the coarse network:
fc(Ht = H, - D.y (X,, H1-1) 1/er
(7)
Notice that in addition a larger time step. the length of sequence H smaller than the fine grid's s.
Finally, we introduce the notion of relaxation of the hidden state. Here, we define this formally, while Fig. D 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 p ,,,y, k) 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 in Fig. 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
F' WFFF H, h) = v F I. V pcles, h)).
(8)
The available parallelism for all relaxation methods is equal the number of coarse points.
Algorithm 1: V, 7)) F-relxxaioon
1.
parallel_fom for l. to T/7
2.
ytc 3. I1 = ht
4.
for f- t to +
5.
Mr+11 = (2,((xT,)) N,)
6.
return N
3.2 MULTI-GRID IN TIME
Algorithm 2 presents pseudocode for the two-level MGRIT method for forward propagation. A related schematic is presented Fig. where line numbers are included in parentheses. Fig, 3 the fine grid on the top corresponds to input sequence of length 16. It is depicted being sub-sequence of length 4, its
distributed across 4 processors. The coarse grid representation single processor. Red anoows in Fig, iindicate grid transfers, while green arrows indicate relaxation and propagation through the network. Appendix A..] provides a brief motivation for MGRIT, and serial implementation is included in the supplementary materials.
propagation is limited

---

Published as a conference paper at ICLR 2022
Algorithm 2: MGProp(ž,, h) Forward propagation using multigrid
// Relax the initial guess
1. N = Wrcr(E,
// Transfer approximate hidden state to the coarse grid
2. F. = R(f(H()
3. M. -R(R')
// Compute coarse grid correction, return to fine grid
4. Solve for h fa(ht) = fe(k))
5. h" = h' + P(h - h()
// Update the fine grid solution, and relax on fine level
6. return V (F(
Relsxation (186)
Processor
Processor
Processor
Processor
Fine Grid Relacation AAprrwimmat/Paraak
Coarse Correction (6)
Comectionn
Approx. (2-3)
Coarse
Processor
Digger Time Steps
(Exast-SensiChhettg
Coarse Propagation (4)
Figure 3: The two level MGRIT method. Green arrows represent relaxation, and forward propaga- tion. Grid transfers are represented by red arrows. Numbers in parentheses are lines in Algorithm 2
Line I Algorithm relaxes the initial guess, and improves the fine grid hidden state. FCF- relaxation is typically used to obtain sequence length independent convergence, though often on the fine level F-relaxation used. Regardless, relaxation applied parallel yielding a processor local improved hidden state. The fine level residual given the relaxed hidden state is
== f(h)-
Let f be the forward propagated hidden state that satisfies [(īi²) = 0. This is easily computed by serial forward propagation. Rewriting F in terms of h' and A* gives
== f(') f(h²).
(10)
Coarsening the residual using the restriction operator from Eq, we arrive at an equation for the coarse grid approximation of h
R(F).
(11)
The coarse grid approximation computed using Eq. II in Line of the algorithm. This compu- tation has reduced time relative to the fine grid because the input sequence is shorter. Line transfers a coarse approximation to the fine grid. The final step in the algorithm is F F-relaxation on 6iht propagates the corrected solution from the coarse points the fine points.
Remark (Multiple Iterations). MGRIT is an iterative algorithm. One pass through Algorithm 2 may not provide sufficient accuracy in the hidden state for use in neural network training. However, in Gunther et al. (2020), relatively few iterations were needed for successful training.
Remark 2 (Beyond Two-Level). As presented, this is a two level algorithm. The coarse grid correc- tion in line 4 provides an end-to-end correction at a reduced run-time of T/c.. However, for large T, this may be prohibitively expensive. The solution is recognize that the coarse grid problem it- self can be solved using MGRIT in a recursive fashion. Thus, new levels are added to reduce the run time of the coarse problem. The run time of level with T time steps scales c(l) T T/P +c(++l) where P is the number of processors. Therefore, the total run time for levels scales as
L-1
T
T Cf
c(0) = Σ
<
(12)
Pely
P 1=0
c
11
1-0
5

---

Published as a conference paper at ICLR 2022
3.3 COARSE GRID PROPAGATORS
The discussion above assumes that all levels use the classic GRU propagator in Eq. B (forward Euler with a time step of A = 1). This presents a problem for coarsening. Consider only the stiff term in the dynamics from Eq. 2 and assume that [0, is constant for all time. This results in the model
ODE 0,h = -1 2) o h. Discretizing with forward Euler, stable time steps satisfy
|| - M1(1 2)1 < - 0 < A S 2(1-2)-1
(13)
The model ODE shares the stiff character of the GRU formulation, and the time step is restricted to be At < 2 assuming the worst case 0 (z implies there is no time step restriction). On the fine grid with At = 1, forward Euler is not unstable. However, on coarse grids, where the time step is equal to cÿ. stability of forward and back propagation is a concern. In practice, we found instabilities coarse grids materialize as poor training, and generalization.
To alleviate problems with stiffness, the traditional approach is to apply an implicit method. For instance, applying backward Euler discretization to the model ODE yields stability region of
[| + A1(1 - - S 1.
(14)
Hence, the time step unlimited for the model problem (for 1 - z < 0 the ODE is unstable).
A backward Euler discretization of Eq. requiring a nonlinear solver is possible; however, it is likely needlessly expensive. We hypothesize, considering the model problem, that the stiffness arises from the first term in Eq. 2 The second term incorporates the input sequence data, changing the hidden state relatively slowly to aggregate many terms over time. Finally, the nonlincarity in 2 may be problematic; however, this term primarily controls the decay rate of the hidden state.
As result of this discussion, we propose an implicit scheme to handle the term labeled "stiff mode" in Eq. 2 The discretization we use computes the updated hidden state to satisfy
(1 + At(1 - 21)) @ he = h + At(1 - 21) 3 me.
(15)
We refer to this method as Implicit GRU. Here, the term (1 21) on the left hand side and the right hand side are identical and evaluated at the current data point 2, and the old hidden state h,- Compared to classic GRU, this method inverts the left-hand side operator to compute the new hidden
state h4 This is easily achieved because (1 A1 2+)) is diagonal matrix. Thus, the evaluation
cost of Eq. 15 is not substantially different from Eq. 3 Yet the (linear) stiffness implied by the first
term in the ODE is handled by the implicit time integration scheme. To simplify notation, Eq. 15 is
rewritten by defining an update operator Φ, (similar to Eq. 3)
(It, 1i-1;) (1 + 7(1 -2) @ (hi-1 + 21) mm). h2 (21,
(16)
The subscript 7 is the time step size, and the learned parameters are identical to classic GRU. This propagator replaces the classic GRU propagator in all levels of the MGRIT algorithm.
3.4 TRAINING WITH MGRIT
Training using MGRIT follows the traditional approaches outlined in Goodfellow et al. (2016). In fact, the MGRIT algorithm has been implemented in PyTorch (Paszke et al., 2019), and the Adam implementation is used without modification. The difference comes in the forward propagation, and backward propagation steps. These are replaced with the MGRIT forward and backward propa- gation. We provide a more detailed description of this approach in Appendix A.2 However, the iterations used for these steps are much fewer than would be required to achieve zero error in the propagation. We found that using 2 MGRIT iterations for the forward propagation, and 1 iteration for back propagation was all that required to reduce the training loss.
Remark 3 (Serial vs. Parallel Inference). As training with MGRIT is inexact, a question arises: does inference with the network evaluated in parallel differ from when the network is evaluated in serial? Our results below indicate consistency between serial and parallel inference even when training is done with MGRIT. This suggests that error introduced by MGRIT's is a severe limitation.

---

Published as a conference paper at ICLR 2022
04
TIme Per Eoech en
Teer AXTUU
0.7
Inplitt DIU
Socid
Processors
Figure 4: For the UCI-HAR dataset: (left) Accuracy for classic, implicit and parallel training of a GRU network on the UCI-HAR dataset. (right) Run time per epoch for training.
EXPERIMENTAL EVALUATION
In this section, evaluate the performance of our parallel-in-time training approach on sequence classification tasks such as human activity recognition and video classification. Experiments are run multi-node supercomputer, named Attaway at Sandia National Laboratories, with each node comprised of 2 sockets with 2.3 GHz Intel Xeon processor, each processor has 18 cores yielding 36 cores per node. Inter-process communication is handled using the Message Passing Interface (MPI) (Gropp et al. 1999) which has been optimized for the Intel Omni-Path interconnect. For the video classification dataset, OpenMP threads were used per MPI rank, yielding 4 ranks per node.
We used two publicly available datasets UCI-HAR (Anguita et al. 2013) and HMDB51 (Kuehne jet af., 201T). The UCI-HAR dataset is a collection T0299 sequences of human activities labeled with 6 activities such as walking, walking upstairs, walking downstairs, sitting, standing, and laying. Each sequence consists of 128 consccutive input data samples size measuring three different types of motion recorded by smartphone accelerometers. The HMDB51 dataset contains 6726 video clips recording 51 human activities, To limit computational costs, in several of our experiments a subset of classes used (see Table the appendix). For consistency, each video clip is converted (truncated/padded) into 128 frames. Each image in the sequence of size 240 x 240 pixels, and passed through ResNet18, 34 or 50 networks pre-trained ImageNet to extract features for the GRU networks (Dng et al. 2009 He et al.) 2016). See Appendix A.3.2) for architectural details.
The software implementation of the MGRIT GRU propagation algorithm is combination of the C++ library XBraid and PyTorch. The combined capability is referred to as TorchBraid The implementation leverages performance enhancements developed for classic GRU within PyTorch. While Implicit GRU has been carefully implemented, further optimizations are certainly possible achieve greater computational performance. A GPU implementation of the MGRIT algorithm is under development. Finally, it is important to note that the time parallel implementation does not interfere substantially with traditional data and model forms of parallelism.
4.1 UCI-HAR DATASET
In Fig. 4 the left image shows the test accuracy variability over 16 training runs for the UCI-HAR dataset. The set of parameters detailed Table of the appendix. Solid lines are the mean accuracy, while the colored shadows indicate one standard deviation. The classic GRU (black) achieves 90% accuracy. Implicit GRU (blue line) does as well though with a more rapid increase with epochs. When training with MGRIT, three levels with coarsening rate of 4 are used for each sequence of length 128. This yields hicrarchy of grids of size 128, 32, and Here, we see that the test accuracy using parallel inference (red line) is modestly lower than the serial case. Inference with the serial algorithm suffers a little more. The plot on the right demonstrates a roughly 2x speedup over the classic GRU method in serial (black dashed line) when using the MGRIT algorithm. This dataset is very small, and likely suffers from the effects of Amdahls' law (see Sun & Ni (1993) for discussions of Amdahl's law). Thus, even achieving 2x speedup notable. Beyond 32 MPI ranks there is no further speedup due to communication swamping potential gains from parallelism.
Available http:://hhhaa..cmm////iteeeeeeee
7

---

Published as a conference paper at ICLR 2022
Convergance of MORI Propagation
4.44, Henheel8
Racovard
2058
405
872
187164
52288
64070
Figure 5: HMDB51 Subset: (left) MGRIT converges rapidly for cj = 2, 4, and 8. Sufficient accu- racy for training is achieved with forward iterations, and backward. (right) Parallel performance of training with different sequence lengths.
lest Aboout onyy Va. Epoch
26. harning Hura
5.70
2.05
4.00
2.00
0.50
2:00
0.50
p-10
p-22
0.4
Epoch
Epochs
Figure 6: HMDB51 Subset: The accuracy for GRU, and implicit GRU, using serial and parallel training. The right image highlights serial vs. parallel inference with the MGRIT trained network.
4.2 HMDB51 DATASET
Fig. 5 presents parameter studies elucidating the scalability and performance of the MGRIT algo- rithm. The left image in Fig. shows the residual norm function of MGRIT iterations for both forward (solid lines) and backward (dashed lines) propagation. The colors differentiate coarsening factors. This demonstrates that the MGRIT algorithm converges iterations marginally faster for smaller cf. Note that the deeper hierarchies implied by cj 2 can be more expensive, while ey 8 may not provide sufficient robustness. As a result, Cf 4 is used in the remainder. Based on the convergence plots and experimentation, only 2 iterations for forward propagation and iteration of backward are used for training (sufficient accuracy indicated in the figure). Further, iterations do not improve training; similar observation was made in Gunther et al. (2020) and Kirby et al.) (2020). The right image in Fig. 5 shows the parallel performance over multiple sequence lengths. While there is some degradation in the weak scaling, the run time for 128 timestep dataset is roughly only twice that of a dataset with sequences of length 8. Further, MGRIT performance improves for all lengths if the number of timesteps per MPI rank is or greater. When timesteps per rank are too few, the computational work is dominated by parallel communication yielding no speedup. Sequences of length 128 are used in the remaining studies in this section.
Fig. 6 (left) shows test accuracy achieved for a subset of the HMDB51 dataset (see dataset details in Table 2). The Adam optimizer is used with fixed learning rate of 10-3, The accuracy for the parallel algorithms are nearby the result obtained for classic GRU. Note that the use of different processor counts does not yield different algorithm. In the right image, the differences between the schemes are examined after training 16 times. Here, we see that the standard deviation (shaded regions) and mean (solid lines) for the different approaches overlap each other. In addition, both serial and parallel inference of the Implicit GRU yield similar results, supporting Remark
To further study the parallel performance of MGRIT training of GRU, we now consider the full HMDB51 dataset. Fig. 7 shows the run time per batch (with a batch size of 100) when different pre-trained ResNet architectures are used. In all cases, implicit GRU networks are trained with 32 MPI ranks and 9 OpenMP threads (using total of 288 CPUs). achieving roughly 6.5 speedup over the classic GRU algorithm trained with 9 OpenMP threads. The largest amount of speedup occurs after only MPI ranks. Increasing parallelism 64 MPI ranks did not result in any speedup. Note
8

---

Published as a conference paper at ICLR 2022
Resintt
Tiee Butch 11)
PM
(3)
334
<06
$75
30088
213
4:25
16:.5
20080
2014
4/1
14/644
46700
Figure 7: For HMDB51: Batch run times for training GRU network with ResNet18, 34, and 50.
3.5
3.0
Saridl s21000
Spacha
Figure 8: For HMDB51: (left) Per epoch breakdown of times with ResNet18 preprocessing images. (right) Test accuracy as function of epochs for a ResNet34 GRU network, trained on 32 MPI ranks with OpenMP threads per rank. The thin blue lines are different initializations for parallel training, and the red line provides a baseline for serial training.
that in this case, each rank is responsible for computing only 2 time steps. Based on Fig. 5 (right), hypothesize that at this scale the amount of computational work is too small to see benefits from increased parallelism. Fig. 8 (left) shows a breakdown of run times for the ResNet18 case. This log plot shows consistent gains with parallelism for the total (red), forward (green), backward (blue) and disk compute times (orange). The expense of the forward simulation is dominated by the evaluation of the ResNet each time step. The backward propagation cost is reduced because sensitivities with respect to the ResNet are not required. Finally, the disk time can be reduced by simply storing these in memory overlapping computation and disk read. Note that both of these options are available to the MGRIT training as well. The robustness of MGRIT training applied to implicit GRU using ResNet34 on 32 MPI ranks investigated with respect to initialization with different random seeds (thin blue lines) and compared to a serial baseline run (red line) in Fig. g (right). learning rate schedule starting from 10-3 and being reduced by 1/2 after 5 epochs is used to achieve roughly 65% accuracy for both parallel and serial training. The variation in parallel training relatively small, and the overall performance robust. For more details on these studies, including the hyper-parameters, see Appendix A.3.2
CONCLUSION
We presented parallel-in-time approach for forward and backward propagation of GRU network. Parallelism is achieved by admitting controlled errors in the computation of the hidden states by using the MGRIT algorithm. key development network based on the ODE formulation of the GRU methodology. Applying MGRIT to other types of RNN (LSTTM or those proposed in [Chang al. (2019) is possible with an implicit solver.
make this possible is a new Implicit GRU
The implementation of this approach leverages existing automatic differentiation frameworks for neural networks. PyTorch and its Adam optimizer were used in this work. Comparisons between PyTorch's GRU implementation and the parallel algorithm were made on CPUs using MPI and OpenMP demonstrating a 6.5x speedup for HMDB51 dataset. A GPU version of the MGRIT tech- nology currently underdevelopment, though related work has demonstrated potential on Neural ODES (Kirby et al., 2020).. Other forms of data and model parallelism are complementary to the proposed technique, though undoubtedly there are unexplored optimizations.

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGMENTS
This work was performed in part, at Sandia National Laboratories, with support from the U.S. De- partment of Energy, Office of Advanced Scientific Computing Research under the Early Career Research Program. Sandia National Laboratories is multimission laboratory managed and oper- ated by National Technology & Engineering Solutions of Sandia, LLC, wholly owned subsidiary of Honeywell International Inc., for the U.S. Department of Energy's National Nuclear Security Ad- ministration under contract DE-NA0003525. This paper describes objective technical results and analysis. Any subjective views opinions that might be expressed in the paper do not necessarily represent the views of the U.S. Department Number: SAND2022-0890 C. This work was supported by the National Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) (No. NRF-2002RIIGA11025595)
Energy or the United States Government. SAND
REFERENCES
Davide Anguita, Alessandro Ghio, Luca Oneto, Xavier Parra, Jorge Luis Reyes-Ortiz, et al. A public domain dataset for human activity recognition using smartphones. In Esann, volume 3. pp. 3. 2013. / icc.rrciiiiiiiii.eiiiii.ooo uci edw.dddu/mmi. /tt.c./ee//mmmaaatttttt////aeaa http:/ssstsssammaaaattttttttteee
Tal Ben-Nun and Torsten Hoefler. Demystifying parallel and distributed deep learning: An in-depth concurrency analysis. ACM Computing Surveys (CSUR), 52(4):1-43, 2019.
Achi Brandt. Multi-level adaptive solutions to boundary-value problems. Mathematics of computa- tion, 31(138):333-39, 1977.
William L Briggs, Van Emden Henson, and Steve F McCormick. A multigrid tutorial. SIAM, 2000.
Joao Carreira and Andrew Zisserman. Quo vadis, action recognition? new model and the kinetics dataset. In proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp- 6299-6308, 2017.
Bo Chang, Minmin Chen, Eldad Haber, and Ed H Chi. Antisymmetticrmm: A dynamical system view recurrent neural networks. arXiv preprint ar arXiv:1902.09689, 2019.
Tian Qi Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary diffcrential equations. arXiv preprint arXiv:1806.07366 2018.
Kyunghyun Cho, Bart Van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Hol- ger Schwenk, and Yoshua Bengio. Learning phrase representations using rnn encoder-decoder for statistical machine translation. arXiv preprint arXiv: 1078, 2014.
Eric C Cyr, Stefanie Günther, and Jacob B Schroder. Multilevel initialization for layer- aarer-laalllel deep neural network training. arXiv preprint arXiv:1912.08974 2019.
Edward De Brouwer, Simm, Ádám Arany, and Y Moreau. Gru-ode-bayes: Continuous modeling of sporadically-osserved time series. Proc. of Advances in Neural Information Processing Systems (NeurIPS), 32:7377-7388. 2019.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hi- erarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255. leee, 2009.
V. A. Dobrev, Tz. Kolev, N. A. Petersson, and J. B. Schroder. Two-level convergence theory for multigrid reduction in time (mgrit). SIAM Journal on Scientific Computing, (N5):5501-$521, jan 6m1 074096)
2017. doi: 10. 137/ 6m1074096. UtL.ttts::// //tpi::///o..oo www./or..1.......
Jeffrey Donahue, Lisa Anne Hendricks, Sergio Guadarrama, Marcus Rohrbach, Subhashini Venu- gopalan, Kate Saenko, and Trevor Darrell. Long-teme term recurrent convolutional networks for visual recognition and description. In Proceedings of the IEEE conference on computer vision and pat- tern recognition, pp. 2625-2634. 2015.
10

---

Published as a conference paper at ICLR 2022
Craig C Douglas. Multigrid methods in science and engineering. IEEE Computational Science and Engineering, 3(4):55-68 68, 1996.
Robert D Falgout, Stephanie Friedhoff, Tz V Kolev, Scott P MacLachlan, and Jacob B Schroder.
Parallel time integration with multigrid. SIAM Journal on Scientific Computing, 36(6):C635- C661, 2014.
Martin J Gander. 50 years of time parallel time integration. In Multiple shooting and time domain decomposition methods, pp. 69-113. Springer, 2015.
Rohit Girdhar, Joao Carreira, Carl Doersch, and Andrew Zisserman. Video action transformer net- work. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 344-253, 2019.
Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep learning. MIT press, 2016.
William Gropp, William D) Gropp, Ewing Lusk, Anthony Skjellum, and Argonne Distinguished Fellow Emeritus Ewing Lusk. Using MPI: portable parallel programming with the message- passing interface, volume 1. MIT press, 1999.
Stefanie Gunther, Lars Ruthotto, Jacob Schroder, Eric C Cyr, and Nicolas R Gauger. Layer- parallel training of deep residual neural networks. SIAM Journal on Mathematics of Data Science, 2(1):1-23, 2020.
Mansura Habiba and Barak A Pearlmutter. Neural ordinary differential equation based recurrent neural network model. In 2020 31st Irish Signals and Systems Conference (ISSC), pp. 1-6. IEEE, 2020.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Zhiheng Huang. Geoffrey Zweig, Michael Levit, Benoit Dumoulin, Barlas Oguz, and Shawn Chang. Accelerating recurrent neural network training via two stage classes and parallelization. In 2013 IEEE Workshop on Automatic Speech Recognition and Understanding, pp. 326-331. IEFE, 2013.
lan Jordan, Piotr Aleksander Sokót, and Il Memming Park. Gated recurrent units viewed through the lens of continuous time dynamical systems. Frontiers in computational neuroscience, pp. 67, 2021.
Andrew Kirby, Siddharth Samsi, Michael Jones, Albert Reuther, Jeremy Kepner, and Vijay Gade- pally. Layer-paralee training with gpu concurrency deep residual neural networks via nonlin- ear multigrid. In 2020 IEEE High Performance Extreme Computing Conference (HPEC). IEEE, sep 2020. doi: 10. wtt.hpeec4367442200999555
URL httpss///cci www.oorgo...... 1109%
2Fhpec43674 2020 9286180
Hildegard Kuehne, Hueihan Jhuang, Estibaliz Garrote, Tomaso Poggio, and Thomas Serre. Hmdb: large video database for human motion recognition. In 20/ International conference on com- puter vision, pp. 2556-2563. IEEE, 2011. URL //serre lab. clps wwwwn. eew. resource/ hmdlo a large human mot lon -at.iiitddddaaseee.t
Apache MXNet. Training with multiple gpus using model parallelism. 2018. URL //xpne. incubator apache http://www....aaaaaa.....oeaaaa. model 1 parallel 1ttm.htm.
Technical report,
Benjamin W Ong and Jacob B Schroder. Applications of time parallelization. Computing and Visualization in Science, 23(1):1-15, 2020.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary Vito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc,
11

---

Published as a conference paper at ICLR 2022
E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 8024 8035. Curran Associates, Inc., 2019. URL http://////////pppaaaaaaaaaaaaa.
9015 aat-tppptooo----teeaee http://ww...................o.aa
liatiiinnnn--irry..ooom library..
pdf
Jacob B Schroder. Parallelizing over artificial neural network training runs with multigrid. arXiv preprint arXiv:1708.0m 02276, 2017.
Nitish Srivastava, Elman Mansimov, and Ruslan Salakhudinov. Unsupervised learning video rep- resentations using Istms. In International conference machine learning, pp. 843-852. PMLR,
2015.
Xian-He Sun and Lionel M Ni. Scalable problems and memory-bounded speedup. Journal of Parallel and Distributed Computing, 19(1):27-77, 1993.
Iyy Sutskever. Training recurrent neural networks. University of Toronto Toronto, Canada, 2013.
UIrich Trottenberg, Cornelius W Oosterlee, and Anton Schuller. Multigrid. Elsevier, 2000.
Xiaolong Wang, Ali Farhadi, and Abhinav Gupta. Actions" transformations. In Proceedings of the IEEE conference Computer Vision and Pattern Recognition, pp.
2016.
Gerhard Wanner and Ernst Hairer. Solving ordinary differential equations II, volume 375. Springer Berlin Heidelberg, 1996.
Yonghui Wu, Mike Schuster, Zhifeng Chen, Quoe V Le, Mohammad Norouzi, Wolfgang Macherey, Maxim Krikun, Yuan Cao, Qin Gao, Klaus Macherey, et al. Google's neural machine trans- lation system: Bridging the gap between human and machine translation. arXiv preprint arXiv:1609.08144, 2016.
Yang You, Jonathan Hseu, Chris Ying, James Demmel, Kurt Keutzer, and Cho-Jui Hsich. Large- batch training for Istm and beyond. In Proceedings the International Conference for High Performance Computing, Networking, Storage and Analysis, pp- 1-16, 2019.
Joe Yue-Hei Ng, Matthew Hausknecht, Sudheendra Vijaynnarasimhan. Oriol Vinyals, Rajat Monga, and George Toderici. Beyond short snippets: Deep networks for video classification. In Pro- ceedings of the IEEE conference on computer vision and pattern recognition, pp. 4694-4702, 2015.
12

---

Published as a conference paper at ICLR 2022
A APPENDIX
A. WHY DOES MGRIT WORK?
To help the reader who may not be familiar with the details of the MGRIT algorithm, we include this brief discussion to provide some intuition about how it works. For brevity, we focus only on MGRIT forward propagation. As discussed above, there are a number of high quality introductory texts on multigrid in general that we recommend for a more in-depth exploration. The monograph "A Multigrid Tutorial" by BBriggs et al. (2000) is a really excellent starting place. For MGRIT specifically see (Falgout et al.) 2014) and Dobrev et al. (2017). Moreover, serial implementation of the MGRIT algorithm written in python using NumPy is included in the supplementary material (see Section A.4 for brief overview). The ODE used in this code is
a,h(t) - 2 - h(t) + (Ah(t)) + Bd(t) + b).
1
(17)
Here, h is the hidden state, d represents the user data, and A, B, b are randomly selected weights and biases independent of time. In the code, the width of the hidden state is user defined parameter, while in the discussion that follows E R This ODE is meant as a simplified representation of an RNN/GRU network. Finally, realize that in the code and in this discussion only forward propagation is considered, no training is performed.
MGRIT is divide and conquer algorithm where the problem of forward propagation is partitioned into reducing high-frequency error, and low frequency error separately. To this end, the error is written as decomposition
h hh C= high-freq -lowffre
(18)
where h is the exact solution, h is the approximation, ē' the error, and the remaining two terms are the high-frequency and low-frequency errors. The high-frequenc error is reduced in MGRIT by using Freeaaatttion. To demonstrate this we consider the ODE in Eq. 17 and choose our initial approximation by selecting random values at each time point to define h. This seeds the initial error with a number of the initial error plotted in time and the Fourier coefficients plotted as a function of wave number. Each color represents a different component of the hidden state. The high frequencies are large and significant, even if the low frequency mode is also substantial. Plots (C) and (D) show the error and its Fourier coefficients after 4 sweeps of FC'F-relxxatee has improved in early times in dramatic fashion. This is expected as FCF Frelaxxttioon propagates exactly from the initial condition. In addition, the error is reduced in magnitude throughout the time domain. Not so evident from the error plot is the reduction in high frequency errors. This is clearly demonstrated in plot (D), where the low frequency modes remain relatively large while the high frequency modes are substantially damped. Thus, FCF--reaxate can be used to reduce high-frequency errors throughout the time domain, and is the first component of the MGRIT divide and conquer strategy.
modes. This is evident in Fig. plots (A) and (B), which show
Clearly, from plot (C), the solution
The second part of the divide and conquer strategy is the coarse grid correction, where coarse representation of the ODE is used to provide a correction based on the residual from the fine grid. The residual implicitly contains the error in the current guess of the solution on the fine grid. The final pair of images, (E) and (F), shows how the error from the fine grid behaves after it is restricted to the coarse grid, in this case with a coarsening factor of 4. Here, the highest frequency modes from the fine grid are removed from the error, leaving only modes that are low-frequency on the fine grid. The goal of the coarse grid is then to compute a correction these low frequency modes. If this is the coarsest level, then this computation is a serial forward propagation, which exactly resolves the coarse grid error and the correction is transferred to the fine grid with the prolongation operator. For multi-level scheme, another round of the divide and conquer algorithm is initiated recursively (applying the MGRIT algorithm again, as discussed in Remark 2). Here, the key is to recognize that parts of the low-frequency error from the fine grid are aliased onto the coarse grid, and become, relative to the coarse grid, high high-frequency. This can be seen as a spreading of the Fourier spectrum image (F) relative to image (D). Again, the FCF-relaxatin algorithm will damp the high-frequency error modes on this level, and the coarse grid correction will correct the low frequency error modes. Regardless, when the recursion returns to the coarse level, the error in both
13

---

Published as a conference paper at ICLR 2022
Initial Error
Raniin Conficiens
-20
(A)
(B)
26
Time
Wave Number
Posl, FCF Sweep Error
Swoep Coefficients
635
1.0
Fourie Coufitianna
65
335
8.3
-0.35
(C)
-022
10
Number
Coarse Error
-22
-000
(E)
(F)
-010
10
20
Figure 9: This figure is a sequence of images representing the error in the hidden state for an ODE using 64 time steps. Each color different component of the hidden state (e.g.. l()) R for all time). The left column shows the error, while the right column shows the Fourier transform of the error. The first row (images and B) the initial error. The second row (images and is the error after 4 sweeps of FCF-relaxation The final row (images E and F) is the restriction of the relaxed error (in the second row) onto the coarse grid. Note that the scale in the Fourier coefficients, and wave numbers axes vary by row.
14

---

Published as a conference paper at ICLR 2022
the high-frequency modes and low-frequency modes will have been reduced. Thus, correction transferred to the fine grid will provide further error reduction.
This presentation has focused forward propagation. However, using adjoint ODE the motiva- tion and implementation are nearly identical. The most notable difference in the implementation is the reversal of the time evolution.
A.2 TRAINING WITH MGRIT
Incorporating the time-pane üime-parallll approach into standard stochastic gradient descent algorithm re- quires replacing the forward and backward propagation steps with the equivalent MGRIT compu- tations. Algorithm 3 presents pseudo-code demonstrating how that works. This algorithm requires additional notation for the parameters
neural network, here denoted by O. The structure the algorithm follows the batched gradient descent approach. In the first part, lines 3 the MGRIT al- gorithm applied to perform forward propagation. the second block, lines 6 MGRIT is called again, this time compute the backpropagation through the neural network. Line I0 computes the gradients of the loss with respect to the weights. Note that the sensitivities computed by backpropa- gation are used to compute the full gradient with respect the parameters in parallel fashion. The final line of the algorithm is simply the gradient descent step with specified learning rate.
Algorithm 3: MGRITTTaining((.. Batches, Forwardlters, Backwardlttrr)
// This is training Loop epochs, and siinh-bachess 1. for E 1... ....hs
2.
for each Batches
// Iterate with MGRIT to porform forward propagation
3.
h=0 0 for E 1... Forwardlters
4.
5.
= MGProp(ē, h)
Iterate with MGRIT to perform backward propagation
6.
W 0
7.
WT V for E 1... Backwardlters
8.
9.
w- = MGBackProp(ž,, w)
Conpute the gradient with respect to parameters 9: h; Vi
10.
11.
6 θ 4g
The algorithm specified shows MGRIT embedded within broader stochastic gradient descent al- gorithm. However, the code described is already relatively modular. Lines 3 are simply forward propagation. Lines 6 TO use back propagation to compute the gradient. In our own PyTorch implementation, we were able to encapsulate the forward propagation a torch. Norch.Modiin Module eval- uation call. Similarly, the gradient computation is implemented to be compatible with PyTorch's autograd capability. As result, utilization of PyTorch's optimizers, like Adam, is straightforward.
A final point needs to be made about the application of MGBackProp Algorithm This is meant to imply that the MGRIT algorithm is applied parallel-in in in-time fashion to the backpropagation algorithm. In opinion, the most transparent way do this is by considering the adjoint of Eq. 15 and its discretization. This the approach taken the Neural ODE paper (Chen et al. 2018).
A.3 COMPUTATIONAL EXPERIMENTS
A.3.1 UCI-HAR DATASET
The UCI-HAR dataset is used as an initial evaluation of the effectiveness the MGRIT training for GRU, The parameters used for this study listed in Table [ The subsets of labeled Training Parameters are values used for both the computational resources and algorithmic details of the op- timizers. Note that the number of threads used per MPI Rank is only this dataset is relatively
15

---

Published as a conference paper at ICLR 2022
Optimizer
Adam
Learning Rate
0.001
Training Parameters
Batch Size
100
OpenMP Threads per MPI Rank
1
Seq. Length
128
Hidden Size
100
GRU Parameters
Num. Layers
Time step (At)
1.0
T
128.0
Iters Fwd/Bwd
MGRIT Parameters
Cf
4
Levels
Table I: Parameters used for the UCI-HAR dataset learning problem in Section 4.1
Training Videos
Test Videos
Classes
Full Dataset Subset of Dataset
6053
673
51 classes
1157
129
chew, eat. jump, run, sit, walk
Table 2: The full HMDB51 and the subset of HMDB51 partitioned using the 90/10 split between training and test data. The subset contains only 6 of the 51 classes of the original dataset.
small. This implies, on our architecture that contains 36 cores per node, only one node is used for this simulation. Those labeled GRU Parameters are network architecture specific parameters. All of these apply to Classic GRU, serial Implicit GRU, and parallel Implicit GRU. Note that time step of At 1 is always chosen to ensure alignment with the Classic GRU architecture. Finally, the MGRIT parameters have been chosen by following the recommendations in prior works of Gunther ct al. (2020) and Kirby et al. (2020).
A.3.2 HMDB51 DATASET
The HMDB51 dataset is composed of 51 classes, and 672 different videos. This full dataset was used in the generation of the results in Fig. and Fig. The subset of the data was used generate Fig. [5] and Fig. [6] The breakdown of training to testing videos and the specific classes used are detailed in Table 2 The images in the video are cropped and padded so the image sizes are all 240x240 Further, to focus on the performance of the MGRIT algorithm, we truncate/pad (with blank frames) the video sequences to a specified length.
Table 3 shows the parameters used for the simulations. These follow the same grouping convention as the parameters for the UCI-HAR dataset. This dataset is substantially larger and more computa- used.
tionally expensive than the UCI-HAR dataset. As result, more computational resources In the table, this is demonstrated by the use of 9 OpenMP threads per MPI rank. On the platform used, this implies there thread). While we experimented with different configurations, we were unable to get more perfor- mance by increasing the number of threads per MPI rank. We suspect that this has do with the configuration of the memory, and software stack on our system. The table also shows the types of studies and parameter variations we have performed and documented the present work. Parameter names denoted with an asterisk are varied in our studies. For instance, Adam is used as the optimizer and At on the fine grid in all experiments, while the learning rate varied in the experiment producing Fig. 8
MPI ranks per node (each of the 36
node gets assigned
Images from the video sequence are preprocessed during training and inference into reduced features (size 1000, also specifying the GRU hidden size) using pre-trained ResNet18, 34 or 50 networks. In training. the weights and biases of the ResNet's remain fixed. Further, the MGRIT algorithm computes the application of the ResNet only once the fine grid. The coarse grids make use of
TThis differs from the documented number somewhat. At the time downloaded these videos, several unavailahle.
16

---

Published as a conference paper at ICLR 2022
Optimizer
Adam
Learning Rate"
0.001
Training Parameters
Batch Size
100
OpenMP Threads per MPI Rank
Seq. Length*
128
Hidden Size
1000
Num. Layers
GRU Parameters
Time step (At)
1.0
T"
128.0
Image Size
240
240
Iters Fwd/Bwd'
2/1
MGRIT Parameters
of
Levels*
Table 3: Baseline parameters used for the HMDB51 dataset learning problem in Section 4.2 These parameters serve as the defaults. However, those that are marked with an asterisk are varied during the study.
Loss spach
15 training Runs
1.5
3.4
1.2
Fecite p-32
2
0.5
0.4
0.2
00
30
25
95
16.5
19.1
15.0
17%
20
35
5.5
100
17%
15.
17%
30.0
Figure 10: For Subset of HMDB51: The training loss for the Classic GRU method and Implicit GRU using the parallel MGRIT propagation algorithm.
those precomputed features to save computational time. In the MGRIT backpropagation step, the features are also reused.
The experiments producing Fig. and Fig. used the smaller subset dataset. For the residual convergence study, the MGRIT coarsening rate was varied. For Cy 2, the number levels each level (from fine to coarse). For Cy 4
used was 6, with 128, 64, 32, 16, 8 and 4 timesteps (the baseline), the number of levels used was 3, with 128, 32 and timesteps on each level. Finally, for Cy 8, the number of levels used was 2, with 128 and 16 timesteps number of timesteps 4 on the coarse level was used that defined the depth of the MGRIT hierarchy. For the performance scaling study in the right image of Fig. 5 varying sequence lengths were used, and, as a result, the time domain parameter T also varied. The remaining parameters followed the baseline from Table notably using cy 4, and a pre-trained ResNet18.
each level. A minimum
For Fig. 6 ResNet18 was used and all parameters followed the baseline. In addition to the test accuracy, include analogous figure with the loss in Fig. [0 In all cases, the loss decreases steadily regardless of the number of processors. Varying training over multiple initial states yields a similar reduction in loss for the MGRIT algorithm.
Figures 7 and g use the full data set. All performance simulations use the baseline parameter set with sequence length of 128. The plot of the timing breakdown uses ResNet18. Finally, in the full dataset accuracy results using ResNet34 in Fig. 8 (right), the baseline parameters are used but the learning rate is varied. The learning starts from 10 and then reduced by 0.5 every 5 epochs, terminating at 1.25 x 10-1
A.4 SUPPLEMENTAL MATERIAL: SERIAL MGRIT CODE
To help the reader new to MGRIT, we have included a serial implementation of the two-level al- gorithm using python with minimal prerequisites (only NumPy). It has also been included in
17

---

Published as a conference paper at ICLR 2022
TorchBraid library example and teaching tool. This example does training, but produces only one forward solve of the ODE in Eq. [7 By default, the ODE has dimension 10, a = 1.0, and the number of steps is 128. The result of running this example is:
Two Level MG
error 1.6675e-01,
residual 1.1287e-01
error = 3.1409e-03, residual 1.9173e-03
error 5.4747e-05,
residual 3.9138e-05
error 3.5569e-07, residual 2.6001e-07
error = 3.2963e-09, residual 2.9624e-09
error 5.5450e-11, residual 3.6914e-11
error = 5.7361e-13, residual 3.9258e-13
error 5.1118e-15, residual 4.5498e-15
As the weights and biases are randomly selected, the specific values may be different depending on the seed. However, in general, the error reduction should be similar. The code itself commented. Further, the implementation of the two level algorithm is labeled with the line numbers used in Algorithm 2
18


Reviews:
[{"review_id": "zB-k-BlZ76P", "reviewer": "Reviewer_Y8yE", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Major comments:\n\n-) The written English can be improved. \n\n-) The paper is hard to follow as the authors assume that readers are familiar with MGRIT. I understand that this is a technical paper which is geared towards researchers that are seasoned in the topic, and there is also a page limitation. Nonetheless, the authors should put every possible effort to include details, even in the form of an appendix.\n\n-) Unfortunately, it is not clear at all what factors affect parallelism and scalability. \n\n-) Restricting all MPI processes within the same node does not really lead to interesting results. Large-scale networks will need way more resources than a 2-socket Intel node. In fact, I am not sure I understand why there is no benefit beyond 36 MPI processes, at least for the larger examples.\n\n-) MPI is undefined -- please do so and add a proper citation. The appendix should provide guidance as to how the experiments are organized and might be replicated. I understand that this might be a considerable effort but any additional information will be particularly \nhelpful. \n\n-) The main drawback of this paper is novelty. There is already some work done in the field (the authors cite it properly) and it is not \nclear to me what is the main differentiation factor. MGRIT is discussed for example in the paper by Kirby et al. (2020) and I am not \nsure that simply focusing on GRUs (together with the discussion in Section 3.3) warrants a publication at ICLR. \n\n-) The main positive part of this paper is the speedup and engineering of making MGRIT work for GRUs. Again, I do not think this is a \nmajor advancement in the field, especially for such a small number of MPI processes.  The authors claim that a GPU version is under \ndevelopment. This would be a very nice addition -- note though that the related paper cited by the authors uses MPI+GPU with experiments on up to several tens of GPUs.\n\n-) Are hyper-parameters used? I did not see any detailed study. How the algorithm performs for different combinations of T and c_f? \nMaybe the authors include some and I missed them.\n\n-) Figures 5, 8, and 9 are too small. Figures 8 and 9 could be combined together.\n\n\nMinor comments:\nSection 4.1: \n\n-)\"When training with MGRIT three levels are with a coarsening rate of 4 are used for the length 128 sequences.\" Please rephrase.\n-)\"Here we see that the the test accuracy.\" Remove \"the\".", "summary_of_the_paper": "The main focus/goal of the submitted paper is the parallelization of the Gated Recurrent Unit (GRU) (the authors focus on classification problems). The authors describe the incorporation of a multigrid reduction in time (MGRIT) solver to speed-up and better parallelize the \napplication of forward and back propagation of information. The proposed technique seems to provide a speedup of about an order of \nmagnitude (at most) when implemented on distributed/shared memory hybrid computing environments.", "summary_of_the_review": "I find the idea to apply MGRIT to GRUs interesting, however I am not sure that there is sufficient differentiation (and major advancement) between the methodology discussed in this paper and previous work cited by the authors (I am aware of some of this paper by my own reading). This is further pronounced by the fact that previous work also considers larger hardware installations of MGRIT; the amount of \nparallel resources considered by the authors does not allow for a trustworthy scaling analysis as inter-node communication is not currently present.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "kJFgDCIFdu", "reviewer": "Reviewer_2GTo", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\nIn this paper, authors propose a novel parallel-in-time (PinT) training method based on a MGRIT solver for GRU networks with long sequences. This method has been well motivated by illustrating the limitations and unsolved problems of the existing methods.\nThe proposed algorithm enables the accelerated parallel training of GRUs on long sequence, which is a unique capability that permits growth in this dimension.\nExperimental results provided on the HMDB51 dataset is impressive, which demonstrate the proposed parallel training scheme of GRU achieves up to 6.5× speedup over a serial approach. \nOverall, the paper is well written and clear, especially the part of methodology introduction.\n\nWeaknesses:\nCan you give a discussion about the different performance of proposed method in parallel inference versus serial inference?\nThere is a confusion that regarding the test accuracy, the experiment setting for UCI-HAR and HMDB51 dataset is a bit different. Can you show the results/plot in the same manner or provide an explanation to it?\nCan you provide the test accuracy of baseline in Figure 9 as well? Otherwise, the accuracy performance of MGRIT for full HMDB51 is unsure, which is necessary to demonstrate whether the remarkable speedup on HMDB51 dataset is meaningful.\nThe experiment results are not clear to demonstrate the statement in Remark 3, “Our results below indicate consistency between serial and parallel inference even when training is done with MGRIT. ” Consequently, confusion is introduced for the statement as follows, “This suggests that error introduced by MGRIT is not a severe limitation.” A more convincing illustration is expected. \n\nComments:\nGrammar errors(underline): Sec. 4.1 - “When training with MGRIT three levels are with a coarsening rate of 4 are used for the length 128 sequences. ”\nIn Sec. 4.1 - “Here we see that the the test accuracy using parallel inference (red line) is modestly lower then purely serial case. ”, “then” -- > “than” \nGrammar errors(underline): Sec. 4 – “In this section, we evaluate the performance our new parallel-in-time training approach on the sequence classification tasks such as human activity recognition and video classification. ”\nIn Conclusion – using “6.5\\times” would be better than “6.5x(letter x)”", "summary_of_the_paper": "This paper aims to address the limitations of existing approaches for training Gated Recurrent Unit (GRU) given long sequence in terms of both training time and model accuracy. To tackle this challenge, author propose a novel parallel training scheme (called parallel-in-time) for GRU based on a multigrid reduction in time (MGRIT) solver.  Specifically, the key to achieving speedup is a hierarchical correction of the hidden state to accelerate end-to-end communication in both the forward and backward propagation. Authors gives experimental results on two public datasets to demonstrate the performance improvement in the long sequence scenario. ", "summary_of_the_review": "N/A", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Wbid4SPhRZF", "reviewer": "Reviewer_bzSV", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "First of all, I have to admit that I am not very familiar with the multigrid reduction in time solver technique and its application to ordinary differential equations. My understanding of this paper is thus not very deep.\n\nStrengths:\n1. The paper provides a new view to parallelize the training task of machine learning models. Previously, progresses in parallelization techniques have lend great power to the efficiency of parallel training, and helped reduce idle time when many devices are used to train large models with large amount of training data. Each new idea in this domain is a great step forward.\n2. The paper is based on solid techniques (multigrid reduction in time solver) that have been used in other domains (solving ordinary differential equations). Pivoting a proven technique to new applications always lends more credibility to the proposed solution.\n\n\nWeakness:\n1. Since the parallelization technique in this paper gives best performance for recurrent neural networks with long sequences, the impact of this paper might not be that wide. \n2. The GRU architecture has to be adapted to achieve the result mentioned in the evaluation. This is a bit limiting because machine learning professional might have their reason to use a specific recurrent architecture, or might want \n\nQuestions:\nSince I am not very familiar with the details of the technique, I would rather use this reviewing opportunity to ask a few questions.\n1. If my understanding is not way-off, the technique trades off computation for parallelization. In other words, more computation is needed for the same number of iterations of training, so that the GRU can be parallelized at the `time` dimension. If this understanding is correct, I am curious where exactly the speed up (6.5 times) is from. Is it because we can utilize the multiple devices more efficiently (there is less idle time on the devices)? Could I say that for the same number of iterations of training, more FLOPs are actually needed, but less time is taken since the devices are running in parallel? If the same number of iterations of training is ran, should we expect the technique to have the same degree of model accuracy?\n\n2. Let's assume that the model is as simple as a 2 layer GRU. If we unroll the model, it just looks like a long repetitive model with lots of GEMMs and activations. There are still data dependencies between the repetitive units, just like the GRU data dependencies. For the unrolled model, there are published parallelization techniques (such as GPipe and PipeDream). I think GPipe might require the batch size to be large enough for good parallelization factors, but PipeDream seems to be very good even if the batch size is not big. PipeDream also has certain inaccuracy, in the sense that the gradients are always stale. I am just wondering how do the authors compare this technique with PipeDream?\n\nNits:\n1. in page 2, the second usage of ODE has explanation \"ordinary differential equation\", but it should be explained in the first use (end of line 2)\n2. in page 4 (around the middle) there is a \"the the hidden state\".\n3. in page 6 section 3.4, there is a \"fewer then would be required\" but I think it should be \"fewer than would be required\"\n\n ", "summary_of_the_paper": "The paper proposed to parallelize the inference and training of GRU networks (a type of recurrent neural networks) at the `time` dimension. \nThe 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.", "summary_of_the_review": "I think the paper provides a very interesting new approach of `time-parallelization`, which can benefit the machine learning community a lot. Thus I think this paper can be accepted.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "MQGQCPuslwM", "reviewer": "Reviewer_ix54", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper describes an iterative technique for evaluating GRU networks based on the multigrid reduction in time (MGRIT) technique. Key idea is partitioning of a long sequence into shorter sub sequences. These subsequences are trained in parallel, error corrections between grids are applied to correct for inexact computation. \n\nThe paper is well-organized, and the mathematical descriptions of the method are clear. The application of the method to GRU layers is novel and will likely be a useful strategy for overcoming the inherently sequential nature of the GRU. The authors describe the tradeoffs in play here well: more resources will be effectively utilized (to a limit noted in the discussion) in applying a fast and scalable, though inexact, method at the cost of a small amount of accuracy in the trained model relative to exact evaluation of the GRU. Obviously, a GPU implementation of this technique would be interesting as well, and the authors note that this is in development. Even without it, though, the results are compelling enough.\n\nRelated work is incomplete, recent work in sequence modeling are particularly missing.\n\nMore details on “Training with MGRIT” (section 3.4) are needed. In the absence of a reproducible (anonymized) source code, pseudo code or code snippet would provide more clarity (sufficient as supplementary materials). \n\nMore could be done in experimental evaluations e.g., quantify impacts of coarsening factor on model quality.\n\nThere were some minor grammatical issues; careful copy-editing is recommended. Figure 6 is missing legend.", "summary_of_the_paper": "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", "summary_of_the_review": "A paper describing a novel parallelization strategy for GRU networks that adds dramatic scalability with some minimal loss in accuracy. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents a way of using multigrid techniques to parallelize GRU networks across the time dimension. Reviewers are uniformly in favor of accepting the paper.  The main strength is that the paper provides a new perspective on dealing with long input sequences by parallelizing RNNs across time. The main weaknesses are around the experiments: only CPU experiments are run, and sequences are not very long (max 128 length). All-in-all, though, it provides an interesting perspective that should be valuable to the community.
Decision: Accept (Poster)