[論文レビュー] Learning to Screen for Fast Softmax Inference on Large Vocabulary Neural Networks

Patrick H. Chen, Si Si+3 more|arXiv (Cornell University)|Oct 29, 2018

Topic Modeling参考文献 13被引用数 10

ひとこと要約

本稿では、文脈ベクトルのクラスタリングを用いて、高 Vocabulary のニューラル言語モデルにおけるソフトマックス推論を高速化する L2S（Learning to Screen）を提案する。軽量なスクリーニングモデルを訓練し、類似する上位k語のサブセットを予測することで、Gumbelソフトマックスのテクニックを用いてエンド・ツー・エンド微分可能に訓練することで、正確なソフトマックスと比較して最大20.4倍の高速化を達成し、同程度の精度（98.9% precision@1）を維持する。従来の手法（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.

研究の動機と目的

• 特にモバイルデバイス上で顕著な計算ボトル neck を解消するため、大規模語彙のニューラル言語モデルにおけるソフトマックス推論の計算コストを軽減すること。
• 予測精度を損なわずに上位k語の予測にかかる推論時間を短縮すること。
• 文脈ベクトルと単語埋め込みのクラスタ構造を活用し、候補検索空間を狭小化すること。
• クラスタ割り当てと候補セット選択を同時に最適化するエンド・ツー・エンドで訓練可能なスクリーニングモデルを開発すること。
• 機械翻訳や次単語予測などの NLP タスクにおいて、従来の MIPS や近似手法よりも優れたスピード・アキュラシーのトレードオフを達成すること。

提案手法

• 2段階の推論パイプラインを提案：まず、軽量なスクリーニングモデルが文脈ベクトルに基づいて候補語の小さなクラスタを予測し、次にそのサブセット内でのみ正確なソフトマックスを計算する。
• 離散的クラスタ割り当てを微分可能にするために Gumbel-Softmax テクニックを用い、ディスクリートなクラスタ割り当てを経由したエンド・ツー・エンドバックプロパゲーションを可能にする。
• スクリーニングモデルと各クラスタごとの候補語セットを同時に学習するための共同最適化目的関数を定式化する。
• 文脈ベクトルにおけるデータ駆動型のクラスタリングパターンを活用するため、訓練分布全体でスクリーニングモデルを訓練する。
• スクリーニングモデルをビームサーチおよび次単語予測パイプラインに統合し、ソフトマックス計算を候補サブセットに限定する。
• 損失関数に時間予算制約を組み込み、実用的な文脈でスピードと精度の調整可能なトレードオフを実現する。

実験結果

リサーチクエスチョン

• RQ1 文脈ベクトルのクラスタ構造を活用する学習可能なスクリーニングモデルは、大規模語彙言語モデルにおけるソフトマックス推論時間を顕著に短縮できるか？
• RQ2 離散的クラスタ割り当てをニューラルネットワークのエンド・ツー・エンド学習に適応可能にするにはどうすればよいか？
• RQ3 従来の MIPS アルゴリズムと比較して、学習可能なスクリーニングを用いる場合の推論スピードアップと予測精度のトレードオフはいかなるものか？
• RQ4 実際の NLP タスク（例：機械翻訳、次単語予測）において、オーダー・オブ・マグニチュードのスピードアップを達成しつつ、高精度（例：98%以上）のprecision@kを維持できるか？
• RQ5 クラスタ数や時間予算制約の変動に対して、本手法はどれほど頑健か？

主な発見

• ドイツ語→英語翻訳タスクにおいて、L2S は 20.4倍の高速化を達成し、precision@1 が 98.9%、precision@5 が 99.3% に達する。Zhang et al. (2018) は 6.7倍の高速化で precision@1 が 98.7% にとどまるのに対し、L2S は顕著に優れている。
• EN-VE 翻訳タスクでは、ビームサイズ 5 を用いた場合、L2S は 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 対 30.33）を維持しながら、DE-EN で最大 13.4 倍の推論時間短縮を達成（ビームサイズ 5）。
• Gumbel-Softmax を用いた共同最適化により、エンド・ツー・エンドの訓練が有効に可能となり、データ駆動型のクラスタリングパターンを学習することで、スピードと精度の両方を向上させられる。