[논문 리뷰] Parallelizing Legendre Memory Unit Training

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

Topic Modeling참고 문헌 34인용 수 8

한 줄 요약

이 논문은 RNN 추론 능력을 유지하면서도 선형 시간 불변(LTI) 메모리 구성요소를 재구성함으로써 효율적인 GPU 학습을 가능하게 하는 레지온드 메모리 유닛(LMU)의 병렬 처리 가능한 변종을 소개한다. 이 방법은 기존 LMU 대비 최대 200배 빠른 학습 속도를 달성하며, psMNIST 및 IMDB 감성 분석을 포함한 여러 벤치마크에서 더 적은 파라미터로 최신 기술(SoTA) 성능을 달성한다.

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.

연구 동기 및 목표

• GPU 활용을 저해하는 RNN 학습의 근본적 한계인 순차적 계산을 해결하기 위해.
• LMU의 LTI 메모리 구성요소를 병렬화함으로써 순환 네트워크의 확장성 장벽을 극복하기 위해.
• 복잡한 RNN과 Transformers에 비해 성능이 유사하거나 이를 초월하는 단순한 선형만을 포함하는 순환 아키텍처가 다양한 시퀀스 모델링 작업에서 성능을 달성할 수 있음을 입증하기 위해.
• 장수열 작업 및 전이 학습 시나리오를 포함한 다양한 벤치마크에서 모델의 효과성을 검증하기 위해.
• 선형 순환 구성요소를 가진 어떤 심층 네트워크라도 병렬화할 수 있는 일반화 가능한 프레임워크를 구축하기 위해.

제안 방법

• 선형 재귀 관계의 닫힌 형식 해를 사용하여, RNN 순차적 계산을 피하고 학습 중 비순차적 병렬 계산이 가능하도록 레지온드 메모리 유닛의 LTI 메모리 구성요소를 재구성한다.
• 기존의 순차적 RNN 재귀를, 매트릭스 지수와 레지온드 기저 투영을 사용해 전체 은닉 상태 궤적을 한 번에 계산하는 피드포워드 계산으로 대체한다.
• LTI 시스템의 상태 전이 매트릭스를 사용해 모든 은닉 상태를 동시에 계산함으로써 backpropagation 중 반복적 재귀를 피한다.
• 테스트 시 순차적 RNN 모드로 전환하여 원래 LMU의 추론 동작을 유지함으로써 낮은 메모리 사용량과 지연을 보존한다.
• 재귀를 선형 구성요소로 제한한 단순화된 LMU 변종에 이 병렬화 기법을 적용하여 완전한 GPU 병렬 처리를 가능하게 한다.
• 선형 시간 불변 시스템의 수학적 구조를 활용해 은닉 상태 진화에 대한 닫힌 형식 표현을 유도함으로써 배치 기반 병렬 학습을 가능하게 한다.

실험 결과

연구 질문

• RQ1 LMU의 선형 메모리 구성요소를 재구성하여 순차적 계산 없이도 학습을 병렬화하면서 추론 효율성을 유지할 수 있는가?
• RQ2 LTI 구성요소의 병렬화가 시퀀스 모델의 학습 시간을 얼마나 줄일 수 있는가?
• RQ3 단순한 선형만의 순환 아키텍처가 다양한 시퀀스 작업에서 원래 LMU 및 LSTMs, Transformers와 같은 최신 기술 모델을 능가하는가?
• RQ4 이 병렬화된 아키텍처가 기존 모델보다 더 적은 파라미터로 NLP 벤치마크에서 더 높은 성능을 달성할 수 있는가?
• RQ5 예를 들어 언어 모델을 하류 NLP 작업에 피지컬러닝하는 전이 학습 설정에서 모델의 성능은 어떠한가?

주요 결과

• 제안된 병렬 처리 가능한 LMU 변종은 매크시-글라스 작업과 같이 장수열에서 원래 LMU 대비 최대 200배 빠른 학습 속도를 달성한다.
• psMNIST에서 모델은 이전 최신 기술(SoTA) 모델보다 파라미터 수의 절반으로 새로운 최신 기술 성능을 달성한다.
• IMDB 감성 분석에서 모델은 LSTM 및 DistilBERT 모델을 모두 능가하지만, 파라미터 수는 최대 650배 적게 사용한다.
• IWSLT’15 En-Vi 번역 작업에서 BLEU 점수 26.2를 기록하여 LSTM 기준선보다 2.3 BLEU 포인트 높은 성능을 보였다(대소문자 구분 없음).
• LTI 모델은 시퀀스 길이에 따라 학습 시간이 선형적으로 증가하는 반면, 병렬 버전은 거의 일정한 학습 시간을 유지하여 효과적인 병렬 처리를 입증한다.
• 원래 LMU와 유사한 수준의 텍스트8 언어 모델링 점수 1.61을 기록하면서도, 극적으로 빠른 학습을 가능하게 한다.