[論文レビュー] Parallelizing Legendre Memory Unit Training

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

Topic Modeling参考文献 34被引用数 8

ひとこと要約

本論文は、RNN推論能力を維持しつつ、線形時不変（LTI）メモリ部の再定式化により、効率的なGPU学習を可能にする、Legendre Memory Unit（LMU）の並列処理可能な変種を提案する。この手法は、元の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変種に、この並列化技術を適用し、完全なGPU並列化を実現する。
• 線形時不変系の数学的構造を活用して、隠れ状態の進化に関する閉形式表現を導出し、バッチ化・並列学習を可能にする。

実験結果

リサーチクエスチョン

• RQ1 LMUの線形メモリ部を再定式化することで、推論効率を損なわずに並列学習を可能にできるか？
• RQ2 LTI部の並列化によって、系列モデルの学習時間をどの程度短縮できるか？
• RQ3 簡素化された線形のみの再帰的アーキテクチャが、元のLMUやLSTM・Transformerといった最先端モデルを上回る性能を示せるか？
• RQ4 この並列化アーキテクチャは、既存のモデルよりも少ないパラメータ数で、NLPベンチマークでより優れた性能を達成できるか？
• RQ5 トランスファーラーニングのシナリオ、例えば言語モデルを下流のNLPタスクに微調整する際、モデルの性能はどの程度か？

主な発見

• 提案された並列化LMU変種は、Mackey-Glassタスクのような長系列において、元のLMUに比べ最大200倍の高速化を達成した。
• psMNISTでは、パラメータ数を半分に抑えながら、過去の最先端モデルを上回る新たな最先端結果を樹立した。
• IMDBセンチメント分析では、LSTMおよびDistilBERTモデルを上回ったが、パラメータ数は最大650倍も少なかった。
• IWSLT’15 En-Vi翻訳タスクでは、BLEUスコア26.2（小文字無視）を達成し、LSTMベースラインより2.3ポイント高いスコアを記録した。
• LTI版のモデルは、系列長に比例して学習時間が線形に増加する一方、並列版はほぼ一定の学習時間を維持しており、効果的な並列化が実証された。
• テキスト8言語モデリングタスクでは、元のLMUと同等の1.61のスコアを達成しながら、著しく高速化された学習を可能にした。