Source-linked AI summary
Parallelizing Linear Recurrent Neural Nets Over Sequence Length
Eric Martin, Chris Cundy
TL;DR
Long-sequence RNN training is limited by nonlinear sequential dependencies and serial evaluation, motivating a way to parallelize linear recurrence over sequence length. The paper develops parallel-scan algorithms and CUDA kernels, frames linear models as linear surrogate RNNs, and reports speedups up to 9x while training GILR-LSTM on a one-million-element dependency.
Problem
Nonlinear dependencies in RNNs prevent parallelizing training over sequence length, while serial evaluation limits efficient training on long sequences and with small minibatches.
Method
The paper parallelizes linear recurrences with scan algorithms, implements a CUDA kernel, and introduces linear surrogate RNNs including a surrogate LSTM.
Results
Up to 9x speedup is reported for training and inference of several RNN architectures, while GILR-LSTM converges 6–10 times faster than CuDNN LSTM and learns a one-million-element dependency.
Takeaways & Limitations
Parallel linear recurrence provides a parallelizable building block for sequential models and extends neural sequence learning to a one-million-timestep dependency.
Takeaways & Limitations
For general hidden states, the O(n^3) matrix-multiplication cost can offset parallel speedups, motivating diagonal recurrence matrices.
Abstract
from arXiv · showhide
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 classification task with a one million timestep dependency.
1 INTRODUCTION
RNNs offer flexible sequence modeling but nonlinear sequential dependencies hinder training over long sequences. The paper shows that associative linear recurrences can be parallelized with scans, motivating practical speedups and linear-surrogate architectures.
- Motivation: Nonlinear dependencies in common RNNs, including LSTMs and GRUs, require sequential processing and limit training on long sequences.The motivation includes sequence lengths ranging from thousands to millions of elements.
- Motivation: Small minibatches can improve optimization behavior but leave serial recurrence as a throughput bottleneck, especially for smaller hidden states.Larger minibatches also increase memory and latency and can damage generalization.
- Parallelization principle: Linear recurrence is a scan operation, and associative operators permit scans and reductions to be parallelized.A scan repeatedly applies a binary operator, with cumulative sums as a familiar example.
- Parallelization principle: The recurrence ht = Λtht−1 + xt satisfies the scan conditions using vector addition, matrix-vector multiplication, and matrix-matrix multiplication.The required operations correspond to ⊕, ⊗, and ⊙ in the scan formulation.
- Contributions: The paper classifies practical RNNs with parallelizable linear recurrences, implements a CUDA kernel, and introduces a linear surrogate for LSTM.The contributions include applications to QRNNs and SRUs and the GILR-LSTM model.
2 PARALLEL LINEAR RECURRENCE
Parallel linear recurrence divides sequence computation across processors using scans and reductions, enabling parallel forward and backward passes. Its efficiency depends on recurrence structure: unrestricted matrices can make parallelism costly, whereas diagonal matrices reduce the per-step operations and support GPU execution.
- Algorithm: The parallel linear recurrence algorithm partitions the sequence, computes per-partition reductions, scans partition summaries, and then scans each partition from its propagated state.The algorithm uses matrix products and the binary operator (Λ, x) • h = Λh + x.
- Theoretical performance: For general n-dimensional states, matrix-matrix multiplication adds an O(n^3) cost that can make the parallel algorithm slower than serial evaluation.The stated costs are Cpscan ∈O(2(n^3 + n^2 + n)(T/p + lg p)) and Csscan ∈O((n^2 + n)T).
- Theoretical performance: Restricting recurrence matrices to diagonal form makes matrix operations proportional to n and yields a parallel speedup factor of pT/3(T + lg p).Under p ≪T, the passage states that parallel cost is no greater than serial cost when p ≥3.
- Forward and backward computation: Parallel scans can process both the forward recurrence and the reverse-order backpropagation recurrence over sequence length.The backward equations form a linear recurrence over gradients.
- Implementation: The CUDA implementation maps processors to GPU warps and reaches peak performance only on sequences of at least several thousand steps.Each warp processes different recurrence-vector elements across its 32 lanes.
3 MODELS
The models replace or augment nonlinear recurrent dependencies with linear recurrences that can be evaluated in parallel, while retaining nonlinear transformations and gating. This framework includes GILR layers, linear surrogate RNNs, and the GILR-LSTM.
- Gated impulse linear recurrence: Linear recurrent layers apply nonlinear transforms to each input, then accumulate sequence elements through gated linear recurrence.GILR gates use sigmoid values in [0,1], while the impulse transform can use any activation function.
- Gated impulse linear recurrence: GILR evaluation uses two large matrix multiplications followed by a parallel linear recurrence over the sequence.Its matrix multiplications see an effective batch size of bT, reducing dependence on minibatch size.
- Linear surrogate RNNs: Linear surrogate RNNs decouple the state passed to later layers from the state used to compute the next transition, requiring only the surrogate state to be linearly computable.QRNNs and strongly typed RNNs are examples, and such models can often be parallelized with convolution or linear recurrence.
- GILR-LSTM: The GILR-LSTM uses a GILR layer as a surrogate for h_t, while retaining the LSTM cell state c_t's linear dependence on c_{t-1}.The resulting recurrence combines the surrogate hidden state with the LSTM cell and output equations.
- GILR-LSTM: A GILR-LSTM adds 2n(n + m) parameters beyond an equivalently sized LSTM to compute the mapping from x to the surrogate hidden state.More generally, a linear surrogate RNN retains the underlying RNN's parameters and adds parameters for its linear surrogate.
4 EXPERIMENTS
Experiments show that parallel linear recurrence substantially improves throughput for linear-surrogate RNNs and enables GILR-LSTM training on million-step dependencies. The kernel provides large speedups over serial recurrence, while GILR-LSTM converges faster than CuDNN LSTM on the long-term dependency task.
- Kernel performance: Up to 40x higher throughput than serial evaluation was achieved by the parallel linear recurrence kernel on long sequences.The comparison was made directly at the kernel level, avoiding TensorFlow overhead.
- Accelerating existing RNN architectures: Up to 9x throughput speedup was obtained for linear-surrogate RNN architectures including SRU, QRNN, and GILR-LSTM.The experiments used two stacked RNN layers with 256 hidden units and controlled GPU memory by fixing bT = 65,536.
- Synthetic experiment: GILR-LSTM was evaluated against CuDNN LSTM on synthetic sequences requiring the model to remember the first element until sequence end.The task used sequence lengths of 1,024, 8,192, and 1,048,576, with input dimensionality fixed at p = 128.
- Synthetic experiment: GILR-LSTM converged between 6 and 10 times faster than CuDNN LSTM on the long-term dependency problem.For the longest sequence, hidden units were reduced to 64 for both architectures so the networks fit in memory.
- Synthetic experiment: GILR-LSTM learned dependencies across a one-million-element sequence, whereas CuDNN LSTM did not converge after several days at that length.The authors describe this as the longest sequential learning problem handled by neural networks to that date.
5 CONCLUSION
The paper presents linear recurrence as a parallelizable building block for sequential models, enabling substantial speedups and sequence dependencies far beyond prior work.
- Linear recurrence provides another parallelizable building block for current and future sequential neural models.The paper positions it alongside matrix multiplication and convolution as an efficient computational primitive.
- The LS-RNN framework organizes recent linear recurrent neural networks into a common modeling framework.
- Parallel linear recurrence achieves significant speedups on already fast sequential models.
- Parallel linear recurrence enables solving a sequential dependency problem multiple orders of magnitude larger than prior work.