[Paper Review] Parallelizing Linear Recurrent Neural Nets Over Sequence Length.

Éric Martin, Chris Cundy|arXiv (Cornell University)|Feb 15, 2018

Neural Networks and Applications12 citations

TL;DR

This paper proposes a parallel scan algorithm to enable training of RNNs with linear sequential dependencies over sequence length, achieving up to 9x speedup in training and inference. It introduces the GILR-LSTM, a linear surrogate model that leverages parallel linear recurrence to train on sequences up to one million timesteps—previously infeasible with standard RNNs.

ABSTRACT

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.

Motivation & Objective

• Address the fundamental limitation of RNNs in parallelizing training over sequence length due to non-linear dependencies.
• Enable efficient training and inference of RNNs on extremely long sequences, previously computationally prohibitive.
• Develop a framework for linear surrogate RNNs that preserves the representational capacity of standard RNNs while enabling parallelization.
• Extend the practical reach of sequence modeling to sequences with one million timesteps using a novel GILR-LSTM architecture.

Proposed method

• Apply the parallel scan algorithm to RNNs with only linear sequential dependencies, enabling data-parallel training across sequence length.
• Implement a CUDA kernel for parallel linear recurrence that efficiently computes hidden states in parallel across timesteps.
• Abstract recent linear RNN work into a general framework of linear surrogate RNNs, enabling modular design of linearized RNN units.
• Design the GILR-LSTM, a linear surrogate of the LSTM unit, using parallel linear recurrence to maintain long-term dependency modeling.
• Use the parallel scan kernel to accelerate both training and inference in the GILR-LSTM, reducing computation time significantly.
• Train the GILR-LSTM on a synthetic sequence classification task with one million timesteps to validate scalability.

Experimental results

Research questions

• RQ1 Can RNN training with linear dependencies be parallelized over sequence length to enable faster training on long sequences?
• RQ2 To what extent can the parallel scan algorithm accelerate training and inference in state-of-the-art RNN architectures?
• RQ3 Can a linear surrogate model like GILR-LSTM maintain performance on long sequences while enabling parallelization?
• RQ4 Is it feasible to train RNNs on sequences with one million timesteps using a linear recurrence-based approach?

Key findings

• The parallel scan algorithm enables full data-parallelization of RNN training over sequence length when dependencies are linear.
• The proposed CUDA kernel achieves up to 9x speedup in training and inference across multiple state-of-the-art RNN architectures.
• The GILR-LSTM model successfully learns long-range dependencies in a synthetic sequence classification task with one million timesteps.
• The linear surrogate framework allows for efficient, scalable training of RNNs on sequences far exceeding previous practical limits.