[论文解读] Parallelizing Legendre Memory Unit Training

Narsimha Chilkuri, Chris Eliasmith|arXiv (Cornell University)|Feb 22, 2021

Topic Modeling参考文献 34被引用 8

一句话总结

本文提出了一种可并行化的勒让德记忆单元（Legendre Memory Unit, LMU）变体，该变体在保持RNN推理能力的同时，通过重新表述其线性时不变（LTI）记忆组件，实现了高效的GPU训练。该方法将原始LMU的训练速度提升了高达200倍，并在多个基准测试中取得了新的最先进结果，包括psMNIST和IMDB情感分析任务，且参数量少于同类的LSTM和Transformer模型。

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), thus overcoming a well known limitation of training RNNs on GPUs. We show that this reformulation that aids parallelizing, which can be applied generally to any deep network whose recurrent components are linear, makes training up to 200 times faster. Second, to validate its utility, we compare its performance against the original LMU and a variety of published LSTM and transformer networks on seven benchmarks, ranging from psMNIST to sentiment analysis to machine translation. We demonstrate that our models exhibit superior performance on all datasets, often using fewer parameters. 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.

研究动机与目标

• 解决RNN训练的根本局限——顺序计算——导致无法高效利用GPU的问题。
• 通过使LMU中的LTI记忆组件实现并行化，突破循环网络的可扩展性瓶颈。
• 证明简化版的纯线性循环架构可在多种序列建模任务中达到甚至超越复杂RNN和Transformer模型的性能。
• 在多个基准测试中验证模型的有效性，包括长序列任务和迁移学习场景。
• 建立一个可泛化的框架，用于并行化任何具有线性循环组件的深度网络。

提出的方法

• 通过线性递推关系的闭式解，重新表述LMU的LTI记忆组件，使训练期间可进行非顺序的并行计算。
• 用一次计算整个隐藏状态轨迹的前馈计算，替代原始的顺序RNN递推，该计算使用矩阵指数和勒让德基投影实现。
• 利用LTI系统的状态转移矩阵，同时计算所有隐藏状态，避免反向传播过程中的迭代递推。
• 在测试时切换回顺序RNN模式，以保持原始LMU的推理行为，从而维持低内存占用和低延迟。
• 将并行化技术应用于简化版LMU变体，其中循环仅限于线性组件，从而实现完整的GPU并行化。
• 利用线性时不变系统数学结构，推导出隐藏状态演化的闭式表达式，实现批量并行训练。

实验结果

研究问题

• RQ1 LMU的线性记忆组件能否被重新表述，以实现在不牺牲推理效率的前提下实现并行训练？
• RQ2 对LTI组件的并行化在多大程度上能减少序列模型的训练时间？
• RQ3 简化版的纯线性循环架构是否能在多种序列任务中超越原始LMU及其他SotA模型（如LSTM和Transformer）？
• RQ4 该并行化架构是否能在NLP基准测试中以更少的参数实现更优性能？
• RQ5 在迁移学习场景中，例如在下游NLP任务上微调语言模型，该模型表现如何？

主要发现

• 所提出的并行化LMU变体在长序列任务（如Mackey-Glass任务）中，训练速度相比原始LMU最高提升了200倍。
• 在psMNIST任务上，该模型取得了新的最先进结果，性能优于以往SotA模型，且参数量仅为一半。
• 在IMDB情感分析任务中，该模型优于LSTM和DistilBERT模型，且参数量最多减少650倍。
• 在IWSLT’15 En-Vi翻译任务上，该模型取得了26.2的BLEU分数（不区分大小写），比LSTM基线高出2.3个BLEU点。
• LTI版本的模型训练时间随序列长度呈线性增长，而并行版本的训练时间几乎保持恒定，证明了有效的并行化能力。
• 该模型在text8语言建模任务上取得了1.61的得分，与原始LMU相当，同时实现了极快的训练速度。