P ARALLELIZING L INEAR R ECURRENT N EURAL N ETS O VER S EQUENCE L ENGTH
A BSTRACT
Recurrent neural networks (RNNs) are widely used to model sequential data but their non-linear dependencies between sequence elements prevent parallelizing training over sequence length.
We show the training of RNNs with only linear sequential dependencies can be parallelized over the sequence length using the parallel scan algorithm, leading to rapid training on long sequences even with small minibatch size.
We develop a parallel linear recurrence CUDA kernel and show that it can be applied to immediately speed up training and inference of several state of the art RNN architectures by up to 9x.
We abstract recent work on linear RNNs into a new framework of linear surrogate RNNs and develop a linear surrogate model for the long short-term memory unit, the GILR-LSTM, that utilizes parallel linear recurrence.
We extend sequence learning to new extremely long sequence regimes that were previously out of reach by successfully training a GILR-LSTM on a synthetic sequence classiﬁcation task with a one million timestep dependency.
I NTRODUCTION
Recurrent neural networks (RNNs) are widely used for sequence modelling tasks in domains such as natural language processing (Sutskever et al., 2014),
speech recognition (Amodei et al., 2015), and reinforcement learning (Hausknecht and Stone, 2015).
Most RNNs, including popular variants such as long short-term memories (LSTMs), introduced by Hochreiter and Schmidhuber (1997) and gated recurrent units (GRUs), introduced by Cho et al.
(2014), contain a non-linear dependency between sequential inputs.
These non-linear dependencies create a very ﬂexible class of models but limit the feasibility of training RNNs on long sequences as each sequence element must be processed sequentially.
Modelling sequences of thousands to millions of elements is important to domains such as robotics, remote sensing , control systems, speech recognition, medicine, and ﬁnance
The RNN serial evaluation inefﬁciency problem is usually mitigated by parallelizing the forward and backward pass over a minibatch of inputs.
Without minibatches, RNN evaluation is a sequence of matrix-vector multiplications.
Minibatches transform RNN computation into a sequence of more efﬁcient matrix-matrix multiplications, but this speed-up brings several disadvantages.
RNN model size is often limited by GPU memory size, and running a forward and backward pass on a minibatch requires memory linear in the minibatch size.
Grouping data into minibatches increases the latency of each pass and reduces the rate of optimization steps.
Finally, training with larger minibatches damages generalization ability (Keskar et al., 2017).
Given these effects, it is desirable to obtain high training throughput with small minibatches.
Persistent RNNs (Diamos et al., 2016) use a novel implementation that can achieve high GPU utilization with very small minibatch sizes when the recurrent state is larger than 500 elements, but even persistent RNNs become limited by the serial evaluation inefﬁciency at smaller hidden sizes.
Numerous prior works have shown strong performance from neural sequential models with only linear dependence on earlier sequence elements.
Balduzzi and Ghifary (2016) investigated RNNs with only elementwise linear recurrence relations h t =
α t (cid:12) h t − 1 + (1 − α t ) (cid:12) x t and developed linear variants of LSTM and GRU that perform similarly to standard non-linear RNNs on text generation tasks.
Bradbury et al.
(2017), Kalchbrenner et al.
Gehring et al.
(2017), and van den Oord et al.
(2016) have successfully applied networks of convolutions over sequences for tasks such as machine translation, language modelling, and audio generation.
These works have observed up to an order of magnitude increase in training throughput compared to RNN alternatives.
Convolutional sequence models typically rely on either an attention mechanism or a (possibly linear) recurrent layer to integrate information at scales larger than the ﬁlter width.
Introduction of a recurrent layer prevents full parallelization over the sequence length while attention mechanisms are expensive to apply on long sequences in online inference use cases.
A linear recurrence is a speciﬁc instance of a general form of computation known as a scan.
Scans and reductions are computations involving repeated application of a binary operator ⊕ over an array of data.
Computing the sum or maximum of an array is an example of a reduction, while a cumulative sum is a common example of a scan operation.
Throughout this work, the scan of ⊕ with initial value b is deﬁned as
[Equation]
The reduction of ⊕ over array A and initial value b is denoted REDUCE ( ⊕ , A, b ) and is the ﬁnal element of SCAN ( ⊕ , A, b ) .
Despite their dependent computation graph, algorithms exist to parallelize scans and reductions when ⊕ is associative (Ladner and Fischer, 1980).
Blelloch (1990) shows that ﬁrst order recurrences of the form h t
= (Λ t ⊗ h t − 1 ) ⊕ x t can be parallelized with the parallel scan algorithm if three conditions are met:
⊕ is associative: ( a ⊕ b ) ⊕ c = a ⊕ ( b ⊕ c ) 2.
⊗ is semiassociative: there exists a binary associative operator (cid:12) such that a ⊗ ( b ⊗ c ) =
( a (cid:12) b ) ⊗ c 3.
⊗ distributes over ⊕ : a ⊗ ( b ⊕ c )
= ( a ⊗ b ) ⊕ ( a ⊗ c )
Considering the familiar operations in linear algebra, we see that the associative operation of vector addition ( x ⊕ y =
x + y ), the semiassociative operation of matrix-vector multiplication ( A ⊗ x = Ax ) and the associative operation of matrix-matrix multiplication ( A (cid:12) B = AB ) satisfy Blelloch’s three conditions, allowing h t = Λ t h t − 1 + x t to be evaluated in parallel over time steps t for vectors x t and square matrices Λ t .
We investigate this idea further and deliver the following contributions:
• We classify RNNs which satisfy the conditions above, and show that many RNNs used in practice such as the Quasi-RNNs (QRNNs) introduced by Bradbury et al.
(2017) are contained in this class.
• We provide an implementation of the parallel linear recurrence algorithm as a CUDA kernel, and show that it speeds up training of QRNN and Lei and Zhang (2017)’s
Simple Recurrent Unit (SRU) architectures by factors of up to 9x.
• We describe how several recent linear RNNs can be described as linear surrogates for non-linear architectures.
We introduce a linear surrogate for the LSTM and show that we are able to train it with a speedup of 5-10x compared to the CuDNN LSTM when we use the parallel linear recurrence algorithm.
P ARALLEL LINEAR RECURRENCE
As the method is essential to this work, Algorithm 1 presents the parallel linear recurrence algorithm for the interested reader.
2. T HEORETICAL PERFORMANCE
The cost of a serial scan over a sequence of length T is C sscan ∈
O (( C ⊗ + C ⊕ ) T ) , compared to the parallel scan cost C pscan ∈
O (2( C (cid:12) + C ⊗ + C ⊕ )( T/p + lg p )) on p processors (Blelloch, 1990).
If h t is a vector of dimension n then C (cid:12) ∈ O ( n 3 ) , C ⊗ ∈ O ( n 2 ) , C ⊕ ∈ O ( n ) giving
1: Let y =
[(Λ 1 , x 1 ) , (Λ 2 , x 2 ) , ..., (Λ T , x T )] 2: Let binary operator • act as (Λ , x ) • h = Λ h + x 3: Let S
= 1 , S
i < E i ,
E i + 1 = S
i +1 ,
E p − 1 = T for i in 0 , p −4: 5: parfor i ← 0 , p − 1 do 6:
P i = REDUCE ( (cid:12) , Λ S i : E i , I ) 7:
R i = REDUCE ( • , y S i : E i , 0) 8: end parfor 9: 10: Let z =
[( P 0 , R 0 ) , ( P 1 , R 1 ) , ..., ( P p , R p )] .
11: C = SCAN ( • , z, h 0 ) (cid:46) compute C
i
=
P
i C
i −
+ R i with C − 1 = h12: 13: parfor i ← 0 , p − 1 do 14: h S i : E i =
SCAN ( • , y S
i : E i , C i − 1 ) 15: end parfor 16: return h
C pscan ∈ O (2( n 3 + n 2 + n )( T/p + lg p )) and C sscan ∈ O (( n 2 + n ) T ) .
The O ( n 3 ) cost of the matrix multiplication in the parallel algorithm can counter-act any parallel speedups for sufﬁciently large hidden states and lead to a slower algorithm overall.
To avoid this problem, we will only consider diagonal matrices Λ t , in which case both matrix-matrix and matrix-vector multiplication have cost proportional to n and C pscan ∈
O (6 n ( T/p + lg p )) and
C sscan ∈ O (2 nT ) .
This gives a parallel speedup factor of pT/ 3( T + lg p ) .
Assuming p (cid:28) T , then
C pscan ≤ C sscan when p ≥ 3 .
As we are only considering diagonal matrices, we write the linear recurrence as h t =
λ t (cid:12) h t − 1 + x t where (cid:12) indicates elementwise multiplication.
Limiting Λ t to be diagonal may seem like a severe constraint but there are several reasons to do so beyond the favorable parallelization performance.
Relatively few neural network models use separate recurrent matrices for each sequence element and using these separate matrices would require potentially prohibitive n 2 T memory.
Applying the same matrix Λ to each sequence element is also unappealing considering that a matrix multiplication can be thought of as a rotation and a scaling.
The same rotation at every element seems unlikely to be useful, and the scaling is exactly what’s captured in diagonal vectors λ t .
Recurrent coefﬁcient vectors λ t provide enough ﬂexibility to implement schemes such as exponential moving averages or a gating mechanism.
2. B ACKPROPAGATION
[Equation]
The backpropagation equations center around a linear recurrence over ∂L∂h t in the reverse order of the original sequence.
This allows for parallelizing both the forwards and backwards pass of a linear
RNN over the sequence length.
2. I MPLEMENTATION
GPUs commonly used for deep learning in 2017 consist of between 640 and 3200 parallel processors known as warps.
Each warp operates on 32 single precision ﬂoating point numbers in parallel.
This work implemented parallel linear recurrence as a CUDA kernel with bindings into the TensorFlow (Abadi et al., 2016) framework.
Each warp acts as a processor, which means the algorithmic p is up to 3200 and the theoretical parallelization speedup factor is up to several hundred.
The 32 lanes of each warp work on different elements of the recurrence vector in parallel.
These implementation details mean that peak performance is only obtained on sequences of at least several thousand steps on at least a 32 element vector.
The parallel linear recurrence CUDA kernel and TensorFlow bindings are available at https: //github.com/eamartin/parallelizing_linear_rnns .
M ODELS
Parallel linear recurrence can be used to construct a wide variety of differentiable modules that can be evaluated in parallel.
Common applications of linear recurrence include gating schemes and exponential moving averages.
Although linear recurrence values can depend only linearly on previous elements, the stacking of linear recurrent layers separated by non-linearities allows for a non-linear dependence on the past.
In this sense the non-linear depth of a linear recurrent network is the number of layers and not the sequence length.
3. G ATED IMPULSE LINEAR RECURRENT LAYER
A gated impulse linear recurrent (GILR) layer transforms its m dimensional inputs x t into a sequence of n dimensional hidden states h t :
[Equation]
A GILR layer applies the same non-linear transform to each sequence element and then accumulates the sequence elements with a non-linear gating mechanism.
Gate g t uses the sigmoid activation function to give values in [0,1] for reasonable gating semantics, while impulse i t can use any activation function τ .
Stacking GILR layers allows for rich non-linear dependence on previous events while still taking advantage of fast parallel sequence evaluation.
3.1. I MPACT ON EFFECTIVE " BATCH SIZE "
Consider evaluating an RNN with recurrence h t = σ ( Uh t − 1 + V x t + b ) from m inputs to n hidden units on a sequence of length T with minibatch size b using a serial evaluation strategy.
At each of T iterations, the naive approach performs two ( b, m ) × ( m, n ) matrix multiplications.
Larger matrix multiplications achieve higher throughput due to less IO overhead, so the better approach computes
V x t for all t ahead of time in a single ( bT, m ) × ( m, n ) matrix multiply.
The non-linear recurrence forces even the better approach to perform T potentially small ( b, m ) × ( m, n ) matrix multiplications in serial.
This makes serial RNN performance heavily dependent on minibatch size.
Now consider the GILR, noting that it has the same two matrix-vector multiplications per iteration as the above RNN.
The intermediate variables g
and i can be evaluated for all t with a single ( bT, m ) ×
( m, n ) matrix multiplication each.
Given g and i , h can be computed using a parallel linear recurrence over T vectors each of bn elements.
Rather than T small operations, the GILR can be evaluated over all sequence elements with two large matrix multiplications and a parallel linear recurrence.
GILR performance is much less dependent on batch size as the matrix multiplication kernel sees an "effective batch size" of bT and T is typically large.
3. L INEAR SURROGATE RNN S
RNNs learn a transition function s t = f
( s t − 1 , x t ) which combines previous state s t − 1 with input x t to compute current state s t .
Non-linear f prevents application of the parallel linear recurrence algorithm and forces slow serial evaluation.
To work around this inefﬁciency, note that s t serves dual purposes.
In s t = f ( s t − 1 , x t ) , s t − 1 serves as an input to f summarizing the previous inputs while s t serves as the output of f to be passed to other layers of the network.
We can decouple these uses and introduce independent variables for each purpose: s t is passed onto other layers of the network and we introduce the linear surrogate ˜ s t which is passed onto the next state, with s t = f
(˜ s t − 1 , x t ) .
We are still able to choose a non-linear f , our only limitation being that ˜ s t must be linearly computable.
We refer to this class of model as a linear surrogate RNN (LS-RNN).
QRNNs (Bradbury et al., 2017) are LS-RNNs using ˜ h t − 1 =
W k x t − k
+ ...
W 1 x t − 1 and strongly typed RNNs (Balduzzi and Ghifary, 2016) are LS-RNNs with ˜ h t
= x t − 1 .
Although not a rule, LS-RNNs can often be parallelized over sequence length with either convolution or linear recurrence.
Consider an LSTM:
[Equation]
An LSTM has state s t
=
( h t , c t ) .
Since c t depends only linearly on c t − 1 , no surrogate is needed for c t .
h t has a non-linear dependence on h t − 1 , so h t needs a linear surrogate.
Introducing a GILR layer as the surrogate, we obtain the GILR-LSTM:
[Equation]
j t = τ ( V j x t + b j ) ˜ h t =
g t (cid:12) ˜
h t − 1 + (1 − g t ) (cid:12) j t f t , i t , o t = σ ( U f,i,o ˜ h t − 1 + V f,i,o x t + b f,i,o )
z t = τ
( U z ˜ h t − 1 + V z x t + b z ) c t = f t (cid:12) c t − 1 + i t (cid:12) z t h t =
o t (cid:12) c t For m inputs and hidden size n , a GILR-LSTM contains 2 n ( n + m ) more parameters than the equivalently sized LSTM to handle the mapping from x to ˜ h .
More generally, a LS-RNN contains all of the same parameters as the underlying RNN as well as some additional parameters to compute the linear surrogate.
E XPERIMENTS
We perform several experiments.
First we ﬁnd that our parallel linear recurrence kernel is able to achieve up to 40x higher throughput than a serial implementation when applied to long sequences.
Secondly, we conﬁrm that this kernel speedup translates to up to a 9x speedup to LS-RNNs such as
QRNNs.
In order to illustrate that the linearization does not necessarily come at the cost of expressibility, we show that the GILR-LSTM architecture computed with the parallel linear recurrence algorithm is able to train signiﬁcantly faster than an optimized LSTM implementation on a pathological long-term dependency problem from the original LSTM paper (Hochreiter and Schmidhuber, 1997).
4. T HROUGHPUT BENCHMARKS 4.1. K ERNEL PERFORMANCE
We ﬁrst illustrate the throughput advantage of the parallel scan algorithm for evaluating the linear recurrence.
For a minibatch comprised of b sequences of length T , we deﬁne the number of events as bT and the throughput as the number of events processed per second. We implement two CUDA [Table]
[Table] kernels, one which evaluates the parallel linear recurrence described in algorithm [Table] and one which evaluates the same linear recurrence on GPU in serial over sequence length and in parallel over features and minibatch.
The performance of each kernel depends on two factors: the sequence length and the product of number of features and minibatch size.
The performance measurements for this experiment are made directly at the kernel level, avoiding any overhead from TensorFlow.
We ﬁnd that the parallel kernel has a distinct advantage at long sequence lengths with a speedup factor of up to 40x, as shown in table 1.
The parallel kernel does not perform well at short sequence lengths due to the overhead of multiple passes over data and communication between processors.
4.1. A CCELERATING EXISTING RNN ARCHITECTURES
Several recently introduced LS-RNNs can be accelerated with the parallel linear recurrence algorithm.
We implemented SRUs, QRNNs (with ﬁlter width 2 and 10) [Table] and GILR-LSTMs that can be computed with either the standard serial linear recurrence algorithm or parallel linear recurrence.
Both methods compute an identical recurrence, so switching from a serial to parallel implementation does not cause any numerical changes and takes only a single line of code changes.
Notably, both SRUs and QRNNs claim an order of magnitude speedup compared to CuDNN LSTM when implemented with serial linear recurrence.
Any further speedup from parallel linear recurrence applies on top of the existing speedup.
We timed train throughput (forwards and backwards propagation), but the linear time of each pass also makes the results applicable to forwards (inference) performance.
However, parallel linear recurrence can only accelerate inference in scenarios where the entire input sequence is known at the start of the inference phase.
We controlled for GPU memory usage within these experiments by ﬁxing bT = 65 , 536 for minibatch size b and sequence length T , and chose a popular architecture consisting of two stacked RNN layers with 256 hidden units and an input size of 4.
Table 2 shows that the throughput advantage from using parallel linear recurrence compared to serial linear recurrence reaches up to 9x.
Simpler architectures (for which the linear recurrence is a higher proportion of the total computational load) are more affected by the switch to the parallel kernel.
This is particularly clear in the case of the QRNN, where including wider convolutional ﬁlters results in more time spent outside of the linear recurrence and therefore reduces the speedup from linear recurrence parallelization.
4. S YNTHETIC E XPERIMENT
One of the key strengths of the LSTM is that it is capable of dealing with long-term dependencies.
In order to demonstrate that the GILR-LSTM is also able to handle long-term dependencies we tackle a canonical example of inference over many time steps from Hochreiter and Schmidhuber (1997).
We show that in fact the GILR-LSTM is able to outperform the CuDNN LSTM and extend to sequence [Table] lengths orders of magnitude longer than dealt with previously.
The input consists of sequences of length n where for n > 0 each element is a randomly chosen one-hot vector x in p -dimensional space.
The ﬁrst vector in each sequence, x 0 , is always either (1 , 0 , . .
. , 0) or ( − 1 , 0 , . .
The sequential model must read in an entire sequence and then output the sign of the ﬁrst sequence element.
This sequence classiﬁcation problem requires remembering the ﬁrst element over the length of the sequence, and early RNNs struggled with this for p as small as a few dozen.
In the original formulation of the problem (dealing in the regime with around one hundred timesteps), the dimensionality of the input p is set equal to n .
Since this would make the size of the input data grow impractically large as O ( n 2 ) for long sequences, we ﬁx p = 128 as we vary n .
We generated sequences for n equal to 1,024, 8,192, and 1,048,576.
For each of these we compared a two layer GILR-LSTM with 512 hidden units to a two layer LSTM with 512 hidden units 1 per layer implemented by CuDNN.
We ran all experiments on a NVIDIA K80 GPU, with ﬁve runs per conﬁguration allowing us to ﬁnd the average and standard deviation of the time and number of iterations to convergence.
We continually generated random sequences to serve as input data.
A brief search over learning rate and batch size was carried out to ﬁnd the parameters which allow the network to converge most rapidly for all runs.
The criterion for convergence was ﬁve consecutive minibatches giving 100% accuracy.
The learning curves in ﬁgure 2 give support to this being a reasonable convergence criteria.
For the longest sequence length, we did not observe the CuDNN LSTM converging, even after several days’ training.
The results as show in table 4.2 illustrate that the GILR-LSTM is able to converge between 6 andtimes faster than the CuDNN LSTM.
This is somewhat surprising given the LSTM was speciﬁcally constructed for problems of this sort, and the CuDNN LSTM implementation is highly optimized (to the extent that the monolithic interface it exposes is difﬁcult to modify or extend).
The GILR-LSTM is implemented entirely in standard TensorFlow with the exception of using the new linear recurrence op instead of a TensorFlow symbolic loop.
Convergence of the GILR-LSTM models leads to the conclusion that the non-linearities present in LSTM are not necessary for solving this instance of the long-term dependency problem.
The time to convergence further leads to the conclusion that inclusion of a non-linearity at every step incurs a signiﬁcant training time slowdown.
Furthermore, the GILR-LSTM is able to learn to carry dependencies over a one million element sequence.
As far as we know, this one million step sequence experiment is the longest sequential learning problem to be handled by neural networks to date.
C ONCLUSION
A signiﬁcant portion of the success of deep learning can be attributed to access to massive amounts of computation.
Most of this computation is accessed through two highly efﬁcient and parallelizable building blocks: matrix multiplication and convolution.
Recent research has demonstrated that linear
RNNs can achieve similar prediction accuracy to non-linear RNNs on a wide variety of tasks in a fraction of the training time.
We propose the framework of LS-RNNs as a way to tame the growing zoo of sequential neural nets.
We identify linear recurrence as another parallelizable building block for current and future sequential models and we use it to obtain signiﬁcant speedups on already fast models.
With the power of parallel linear recurrence we are able to solve a sequential dependency problem multiple orders of magnitude larger than anything done prior.
Future applications of parallel linear recurrence within neural nets could include parallel training of memory augmented models or providing a new sort of image ﬁlter on very high resolution images.
We hope that parallel linear recurrence can be to large scale sequence modelling what fast convolution algorithms are to image recognition.
A CKNOWLEDGMENTS
We would like to acknowledge Kevin Bowers, Alex Meiburg, JD Co-Reyes, Carson McNeil, Andy Palan, Sören Mindermann, and several others for fruitful conversations and guidance.