LEARNING TO SCREEN FOR FAST SOFTMAX INFER- ENCE ON LARGE VOCABULARY NEURAL NETWORKS
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 (Zhang et al., 2018) only achieves 6.7x speedup with 98.7% precision@1 and 98.1% precision@5 for the same task.
This paper proposes a novel method to speedup softmax computation at test time. Their approach is to partition the large vocabulary set into several discrete clusters, select the cluster first, and then do a small scale exact softmax in the selected cluster. Training is done by utilizing the Gumbel softmax trick.
Pros:
1. The method provides another way that allows the model to learn an adaptive clustering of vocabulary. And the whole model is made differentiable by the Gumbel softmax trick. 
2. The experimental results, in terms of precision, is quite strong. The proposed method is significantly better than baseline methods, which is a really exciting thing to see. 
3. The paper is written clearly and the method is simple and easily understandable. 
Cons:
1. I’d be really expecting to see how the model will perform if it is trained from scratch in NMT tasks. And I have reasons for this. Since the model is proposed for large vocabularies, the vocabulary of PTB (10K) is by no terms large. However, the vocabulary size in NMT could easily reach 30K, which would be a more suitable testbed for showing the advantage of the proposed method.  
2. Apart from the nice precision results, the performance margin in terms of perplexity seems not as big as that of precision. And according to earlier discussions in the thread, the author confirmed that they are comparing the precision w.r.t. original softmax, not the true next words. This could raise a possible assumption that the model doesn’t really get the probabilities correct, but somehow only fits on the rank of the words that was predicted by the original softmax. Maybe that is related to the loss? However, I believe sorting this problem out is kind of beyond the scope of this paper.  
3. In another scenario, I think adding some qualitative analysis could better present the work. For example, visualize the words that got clustered into the same cluster, etc. 
In general, I am satisfied with the content and enjoys reading the paper.