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 the 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 the translation by incorporating SRU into the architecture. The experimental code and SRU++ implementation are available on the dev branch which will be merged into master later.
Features
- SRU is a recurrent unit that can run over 10 times faster than cuDNN LSTM
- Without loss of accuracy tested on many tasks
- Simple Recurrent Units for Highly Parallelizable Recurrence
- SRU can be installed as a regular package via python setup.py install or pip install
- Directly use the source without installation
- The usage of SRU is similar to nn.LSTM