Source-linked AI summary
Simple Recurrent Units for Highly Parallelizable Recurrence
Tao Lei, Yu Zhang, Sida I. Wang, Hui Dai, Yoav Artzi
TL;DR
Common recurrent architectures are difficult to scale because sequential state dependencies impede parallel computation. The paper introduces SRU, a light recurrent unit that separates recurrent state computation across dimensions and uses architectural and initialization choices for scalable deep modeling. Across NLP tasks, SRU improves results while achieving 5–9x speed-up over cuDNN LSTM on classification and question answering, and adds an average 0.7 BLEU improvement when incorporated into Transformer translation models.
Problem
Sequential state dependencies make common recurrent architectures difficult to parallelize and scale.
Method
SRU combines light recurrence, independent state dimensions, highway connections, and tailored initialization for highly parallelizable sequence modeling.
Results
Across NLP tasks, SRU outperforms common recurrent and non-recurrent architectures, achieves 5–9x speed-up over cuDNN LSTM, and improves Transformer translation by an average of 0.7 BLEU.
Takeaways & Limitations
SRU offers a trade-off between scalability and representational power, with deeper stacking retaining modeling capacity despite simplified recurrence.
Abstract
from arXiv · showhide
Common recurrent neural architectures scale poorly due to the intrinsic difficulty in parallelizing their state computations. In this work, we propose the Simple Recurrent Unit (SRU), a light recurrent unit that balances model capacity and scalability. SRU is designed to provide expressive recurrence, enable highly parallelized implementation, and comes with careful initialization to facilitate training of deep models. We demonstrate the effectiveness of SRU on multiple NLP tasks. SRU achieves 5--9x speed-up over cuDNN-optimized LSTM on classification and question answering datasets, and delivers stronger results than LSTM and convolutional models. We also obtain an average of 0.7 BLEU improvement over the Transformer model on translation by incorporating SRU into the architecture.
1 Introduction
SRU addresses the sequential state-computation bottleneck in recurrent networks with light recurrence designed for high parallelization and sequence modeling. Across NLP tasks, it improves performance while substantially reducing runtime relative to cuDNN LSTM.
- RNN scalability is limited because each state computation waits for the complete execution of the previous step.
- SRU combines light recurrence with sequence modeling capacity and high parallelization.
- SRU makes state dimensions independent, enabling CUDA-level parallelization across hidden dimensions and time steps.
- SRU replaces QRNN- and KNN-style convolutions with more recurrent connections while using less computation and fewer hyper-parameters.
- SRU uses highway connections and tailored parameter initialization to improve training of deep recurrent models.
- 5–9x speed-up over cuDNN LSTM was achieved on classification and question answering datasets, while SRU outperformed recurrent and non-recurrent architectures.
2 Related Work
Related work seeks faster sequence processing through simplified recurrence, convolutions, hardware strategies, conditional computation, and factorization. SRU builds on these directions while emphasizing parallelization and retained modeling capacity through depth.
- Quasi-RNN alternates word-level convolutions with minimalist recurrent pooling and achieves significant speed-up over LSTM.
- Prior results suggest simplified recurrence can retain strong modeling capacity through layer stacking, motivating SRU’s design.
- SRU uses point-wise hidden connections equivalent to a diagonal weight matrix, relating it to identity-diagonal initialization in IRNN.
- Other scaling strategies include caching RNN parameters, conditional computing, and matrix factorization to speed recurrent-network computation.
3 Simple Recurrent Unit
SRU separates light recurrence from a highway network, then reorganizes and fuses computation to parallelize it efficiently. Its point-wise hidden connections support linear per-layer point-wise work, while scaling correction and initialization target deep-model training.
- A single SRU layer consists of light recurrence and a highway network.The light recurrence computes sequential states, while the highway network combines inputs and states through a reset gate.
- Point-wise multiplication of the previous state makes each state dimension independent and parallelizable, unlike full hidden-to-hidden matrix multiplication.
- The highway component uses a reset gate to combine the input and recurrent state, with a skip connection for direct gradient propagation.
- The two components together make SRU simple, expressive, and easier to scale through enhanced parallelization and gradient propagation.
- 3.1 Parallelized Implementation: SRU batches matrix multiplications across all time steps and fuses point-wise operations into one CUDA kernel.
- 3.1 Parallelized Implementation: O(L · B · d) per layer is the complexity of SRU’s point-wise step, versus O(L · B · d2) for LSTM’s hidden-to-hidden multiplications.
- 3.2 Initialization: Scaling correction offsets hidden-representation variance reduction so Var[ht] ≈ Var[xt] at initialization.
- 3.2 Initialization: Figure 2 compares training with and without scaling correction, which improves progress especially in deeper models with many stacked layers.
4 Experiments
The experiments evaluate SRU across classification, question answering, translation, and character-level language modeling, focusing on its performance–speed trade-off. Across these tasks, SRU generally combines competitive or stronger results with substantially faster training.
- Experimental scope: SRU is evaluated on classification, question answering, machine translation, and character-level language modeling.The benchmarks span computation times from minutes to days.
- Text Classification: SRU outperforms most previous classification methods while achieving comparable results to a more sophisticated state-of-the-art model.The comparison uses six sentence-classification benchmarks and a simple setup without fine-tuning word embeddings.
- Text Classification: 5–9 times faster than cuDNN LSTM, SRU also runs 6–40% faster than Kim’s CNN model on classification training.On MR, SRU completes 100 epochs within 40 seconds, whereas LSTM takes over 320 seconds.
- Question Answering: 71.4% exact match and 80.2% F1, SRU outperforms bidirectional LSTM on SQuAD while achieving over 5x speed-up.SRU improves over QRNN by 0.8% exact match and 0.6% F1, and runs 60% faster.
- Machine Translation: 0.7 average test BLEU improvement over Transformer base is achieved by the 5-layer SRU translation model.The 4-layer model is 10% faster, while the 5-layer model is only 5% slower than the base model.
- Character-level Language Modeling: 1.19 BPC is reached by the 12-layer SRU with 256-character context, improving over the 8-layer model’s 1.21 BPC.The 8-layer SRU already outperforms previously reported LSTM, QRNN, and RHN results.
5 Discussion
SRU is presented as a scalable recurrent architecture whose simplified dependencies support fast computation while preserving modeling capacity through depth and design components.
- SRU operates as fast as feed-forward and convolutional units, according to the paper’s conclusion.
- Ablation results on classification and question answering show that removing SRU components worsens classification accuracies and exact match scores.The analyzed components include the v ⊙ c_t−1 multiplication and highway connection.
- SRU balances high parallelization with representational power by simplifying hidden-to-hidden dependencies and increasing model depth.The paper argues that additional layers retain modeling capacity despite the reduced capacity of a single layer.
A Parameter Initialization Derivation
The initialization derivation assumes small, zero-mean inputs and weight variance 1/d to preserve value scale after matrix multiplication.
- The derivation assumes input-vector entries are independent, zero-mean, and have small variance.
- Each weight matrix is initialized with zero mean and variance 1/d.
- The resulting matrix multiplication preserves the scale of the values.
A.1 Computing Var[ct]
The variance of SRU’s cell state depends on input correlation, with empirical measurements showing the variance ratio rises toward 1 in deep layers.
- The cell-state value c_t,i is computed from the current input, recurrent state, and forget-gate terms.
- The variance of c_t,i depends on correlations between input vectors.Independent and perfectly correlated inputs provide different variance regimes.
- For independent inputs, the derived cell-state variance is Var[x]/3.
- The derivation ignores correlations between forget-gate entries because their variance is small.
- For perfectly correlated inputs, the derived cell-state variance is Var[x].
- In stacked SRU networks, the actual variance ratio is expected to lie between these bounds and converge toward 1 in deep layers.Figure 6 empirically confirms this expectation.
A.2 Computing Var[ht]
The hidden-state variance is derived from the cell-state variance bounds and the highway reset gate, yielding an explicit bound under the stated initialization conditions.
- The derivation computes Var[h_t] after obtaining the variance of the cell state.
- The highway reset gate uses a bias b that is not necessarily initialized to zero.
- A Taylor expansion around u = 0 approximates the highway reset-gate value.The approximation uses E[u] = 0 and small Var[u].
- The u^2 term is ignored because Var[u] ≪1.
- Using Var[x]/3 ≤ Var[c] ≤ Var[x], the hidden-state variance ratio is bounded by the expressions involving b.The stated bound is evaluated under b = 0.
A.3 Computing the Scaling Constant α
The section derives the scaling constant α by matching the variance of the hidden state to the input variance under the stated recurrence assumptions.
- A.3 Computing the Scaling Constant α: The variance-matching condition yields e_b^2 + α = (1 + e_b)^2.This follows from the displayed variance expression and the requirement that hidden-state variance remain near input variance.
- A.3 Computing the Scaling Constant α: α is chosen so that Var[h] ≈ Var[x], using the relation Var[c] → Var[x].The derivation uses Equation (5) and the stated variance condition.
- A.3 Computing the Scaling Constant α: When b = 0, the resulting expression for α evaluates to 3.The supplied derivation explicitly states this special case.
B Experimental Details
The experiments use additional setup choices for optimization, regularization, preprocessing, and stacked recurrent architectures.
- B Experimental Details: Additional experiments and results are included in this section.The passage introduces the supplementary experimental material without specifying its individual outcomes.
- B Experimental Details: The data and preprocessing code come from the Harvard NLP code repository.The cited repository is identified as the source of the data and preprocessing implementation.
- B Experimental Details: All models use batch size 32 and dropout probability 0.5, with dropout increased to 0.55 or 0.6 for the 8-layer SRU.Out-of-vocabulary embedding vectors are randomly initialized from [−0.25, 0.25].
- B Experimental Details: With highway connections, SRU and Q-RNN pass only the last layer’s output to subsequent layers.The passage contrasts this with bi-LSTM concatenation of every layer’s output.
B.3 Machine Translation
The translation experiments evaluate SRU-based architectures with specified decoding and training configurations, reporting stronger BLEU results and faster-progressing SRU language-model curves.
- B.3 Machine Translation: The translation evaluation uses beam size 5 and length penalty 0.6 with the OpenNMT PyTorch implementation.Table 6 contains the training configuration options.
- B.3 Machine Translation: Table 7 reports averaged BLEU scores from the 20th through 40th epoch, with improvement over the Transformer base model consistent across epochs.The table is explicitly described as averaging scores over this epoch range.
- B.3 Machine Translation: The SRU models show much faster training progress and lower training perplexity than the compared models.With dropout 0.2, the 5-layer SRU also has consistently lower validation perplexity than the base and 4-layer models.
- B.3 Machine Translation: The projection trick reparameterizes stacked SRU matrices with low-rank factors Q and P, using projection dimension d′ much smaller than input and output dimensions.Q ∈ R^d_in×d′ and P ∈ R^3d_out×d′ are learned parameter matrices.