[논문 리뷰] Learning to Screen for Fast Softmax Inference on Large Vocabulary Neural Networks

Patrick H. Chen, Si Si외 3명|arXiv (Cornell University)|2018. 10. 29.

Topic Modeling참고 문헌 13인용 수 10

한 줄 요약

이 논문은 맥락 벡터의 군집화를 이용해 상위-k 단어의 후보 집합을 예측하는 경량 스크리닝 모델을 훈련시켜 대규모 어휘를 가진 신경 언어 모델에서 소프트맥스 추론을 가속화하는 L2S(Learning to Screen)를 제안한다. Gumbel 소프트맥스 기법을 적용해 엔드 투 엔드 미분 가능 훈련을 구현함으로써, 정확도는 거의 동일한 98.9% precision@1을 유지하면서도 정확한 소프트맥스 대비 최대 20.4배의 속도 향상을 달성하였으며, 기존의 FGD 및 Zhang et al. (2018)와 같은 방법들을 크게 능가한다.

ABSTRACT

Neural language models have been widely used in various NLP tasks, including machine translation, next word prediction and conversational agents. However, it is challenging to deploy these models on mobile devices due to their slow prediction speed, where the bottleneck is to compute top candidates in the softmax layer. In this paper, we introduce a novel softmax layer approximation algorithm by exploiting the clustering structure of context vectors. Our algorithm uses a light-weight screening model to predict a much smaller set of candidate words based on the given context, and then conducts an exact softmax only within that subset. Training such a procedure end-to-end is challenging as traditional clustering methods are discrete and non-differentiable, and thus unable to be used with back-propagation in the training process. Using the Gumbel softmax, we are able to train the screening model end-to-end on the training set to exploit data distribution. The algorithm achieves an order of magnitude faster inference than the original softmax layer for predicting top-$k$ words in various tasks such as beam search in machine translation or next words prediction. For example, for machine translation task on German to English dataset with around 25K vocabulary, we can achieve 20.4 times speed up with 98.9\% precision@1 and 99.3\% precision@5 with the original softmax layer prediction, while state-of-the-art ~\citep{MSRprediction} only achieves 6.7x speedup with 98.7\% precision@1 and 98.1\% precision@5 for the same task.

연구 동기 및 목표

• 모바일 디바이스에서 특히 중요한 대규모 어휘 신경 언어 모델의 소프트맥스 추론에서 발생하는 계산적 병목 현상을 해결하기 위해.
• 예측 정확도를 훼손하지 않으면서 상위-k 단어 예측의 추론 시간을 단축시키기 위해.
• 맥락 벡터와 단어 임베딩의 군집화 구조를 활용하여 후보 검색 공간을 좁히기 위해.
• 클러스터 할당과 후보 집합 선택을 동시에 최적화하는 엔드 투 엔드 훈련 가능한 스크리닝 모델을 개발하기 위해.
• 기계 번역 및 다음 단어 예측과 같은 NLP 작업에서 기존의 MIPS 기반 및 근사 기반 방법들보다 속도-정확도 트레이드오프 측면에서 뛰어나게 성능을 높이기 위해.

제안 방법

• 두 단계 추론 파이프라인 제안: 첫 번째로, 경량 스크리닝 모델이 맥락 벡터를 기반으로 후보 단어의 소규모 군집을 예측하고, 두 번째로, 정확한 소프트맥스는 해당 하위집합 내에서만 계산함.
• 이산적 클러스터 할당을 위해 Gumbel-Softmax 기법을 활용한 미분 가능한 군집화 메커니즘을 도입하여 신경망 내에서 엔드 투 엔드 역전파를 가능하게 함.
• 스크리닝 모델과 각 군집별 후보 단어 집합을 동시에 학습하는 공동 최적화 목적 함수를 수립함.
• 맥락 벡터 내 데이터 기반 군집 패턴을 활용하기 위해 전체 훈련 분포에서 스크리닝 모델을 훈련함.
• 스크리닝 모델을 빔 서치 및 다음 단어 예측 파이프라인에 통합하여 소프트맥스 계산을 후보 집합으로만 제한함.
• 속도와 정확도의 균형을 맞추기 위해 손실 함수에 시간 예산 제약 조건을 설정하여 실무에서 조정 가능한 트레이드오프를 가능하게 함.

실험 결과

연구 질문

• RQ1 맥락 벡터의 군집화를 활용하는 학습 가능한 스크리닝 모델이 대규모 어휘 언어 모델에서 소프트맥스 추론 시간을 크게 줄일 수 있는가?
• RQ2 이산적 군집 할당을 어떻게 신경망 내에서 미분 가능하게 만들 수 있는가?
• RQ3 기존의 MIPS 알고리즘 대비 학습된 스크리닝을 사용할 경우 속도 향상과 예측 정확도 사이의 트레이드오프는 어떻게 되는가?
• RQ4 실제 NLP 작업에서 소프트맥스 추론을 순서 수준으로 가속화하면서도 높은 precision@k(예: 98% 이상)를 유지할 수 있는가?
• RQ5 클러스터 수와 시간 예산 제약 조건의 변화에 대해 이 방법은 얼마나 강인한가?

주요 결과

• 독일어-영어 번역 작업에서 L2S는 20.4배의 속도 향상을 달성했으며, precision@1은 98.9%로 정확도는 거의 동일했고, Zhang et al. (2018)는 6.7배의 속도 향상과 98.7% precision@1을 기록함.
• EN-VE 번역 작업에서 L2S는 빔 크기 5일 때 20배의 속도 향상을 기록했으며 BLEU 점수 하락 폭은 0.08에 불과했고, FGD는 유사한 정확도로도 6배 이하의 속도 향상에 그침.
• PTB-Small 및 PTB-Large 언어 모델링 작업에서 L2S는 전체 소프트맥스 대비 퍼플렉서피 증가 폭이 5% 이내로 5배 이상의 속도 향상을 달성함.
• 클러스터 수에 대해 강인함: 100~200개의 클러스터에서 고정밀도(P@1 > 0.99)와 저소요 추론 시간(0.14~0.17ms)을 유지함.
• 빔 서치에서 L2S는 BLEU 점수를 거의 동일하게 유지함(예: 30.19 vs. 30.33)과 동시에 DE-EN에서 빔 크기 5일 때 추론 시간을 최대 13.4배까지 단축함.
• Gumbel-Softmax를 통한 공동 최적화는 효과적인 엔드 투 엔드 훈련을 가능하게 하여, 데이터 기반 군집 패턴을 학습함으로써 속도와 정확도를 동시에 향상시킴.