Parallelizing Legendre Memory Unit Training

Proceedings of the 38th International Conference on Machine Learning, 2021

Narsimha Reddy Chilkuri, Chris Eliasmith

Abstract

Recently, a new recurrent neural network (RNN) named the Legendre Memory Unit (LMU) was proposed and shown to achieve state-of-the-art performance on several benchmark datasets. Here we leverage the linear time-invariant (LTI) memory component of the LMU to construct a simplified variant that can be parallelized during training (and yet executed as an RNN during inference), resulting in up to 200 times faster training. We note that our efficient parallelizing scheme is general and is applicable to any deep network whose recurrent components are linear dynamical systems. We demonstrate the improved accuracy of our new architecture compared to the original LMU and a variety of published LSTM and transformer networks across seven benchmarks. For instance, our LMU sets a new state-of-the-art result on psMNIST, and uses half the parameters while outperforming DistilBERT and LSTM models on IMDB sentiment analysis.

Full text links

 PDF

 External link

Conference Proceedings

Booktitle
Proceedings of the 38th International Conference on Machine Learning
Pages
1898–1907
Volume
139
Series
Proceedings of Machine Learning Research
Month
18–24 Jul
Publisher
PMLR
Editors
Marina Meila, Tong Zhang

Cite

Plain text

BibTeX