ENCE ON LARGE VOCABULARY NEURAL NETWORKS
Year: 2019

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 (Zhang et al., 2018) only achieves 6.7x speedup with 98.7% precision@1 and 98.1% precision@5 for the same task.

heading: 1 INTRODUCTION
text: Neural networks have been widely used in many natural language processing (NLP) tasks, including neural machine translation (Sutskever et al., 2014), text summarization (Rush et al., 2015) and dialogue systems (Li et al., 2016). In these applications, a neural network (e.g. LSTM) summarizes current state by a context vector, and a softmax layer is used to predict the next output word based on this context vector. The softmax layer first computes the “logit” of each word in the vocabulary, defined by the inner product of context vector and weight vector, and then a softmax function is used to transform logits into probabilities. For most applications, only top-k candidates are needed, for example in neural machine translation where k corresponds to the search beam size. In this procedure, the computational complexity of softmax layer is linear in the vocabulary size, which can easily go beyond 10K. Therefore, the softmax layer has become the computational bottleneck in many NLP applications at inference time.
Our goal is to speed up the prediction time of softmax layer. In fact, computing top-k predictions in softmax layer is equivalent to the classical Maximum Inner Product Search (MIPS) problem—given a query, finding k vectors in a database that have the largest inner product values with the query. In neural language model prediction, context vectors are equivalent to queries, and weight vectors are equivalent to the database. MIPS is an important operation in the prediction phase of many machine learning models, and many algorithms have been developed (Bachrach et al., 2014; Shrivastava & Li, 2014; Neyshabur & Srebro, 2015; Yu et al., 2017; Guo et al., 2016). Surprisingly, when we apply
∗This work is conducted during Pei-Hung Chen’s internship in Google Research.
recent MIPS algorithms to LSTM language model prediction, there’s not much speedup if we need to achieve > 98% precision (see experimental section for more details). This motivates our work to develop a new algorithm for fast neural language model prediction.
In natural language, some combinations of words appear very frequently, and when some specific combination appears, it is almost-sure that the prediction should only be within a small subset of vocabulary. This observation leads to the following question: Can we learn a faster “screening” model that identifies a smaller subset of potential predictions based on a query vector? In order to achieve this goal, we need to design a learning algorithm to exploit the distribution of context vectors (queries). This is quite unique compared with previous MIPS algorithms, where most of them only exploit the structure of database (e.g., KD-tree, PCA-tree, or small world graph) instead of utilizing the query distribution.
We propose a novel algorithm (L2S: learning to screen) to exploit the distribution of both context embeddings (queries) and word embeddings (database) to speed up the inference time in softmax layer. To narrow down the search space, we first develop a light-weight screening model to predict the subset of words that are more likely to belong to top-k candidates, and then conduct an exact softmax only within the subset. The algorithm can be illustrated in Figure 1. Our contribution is four folds:
• We propose a screening model to exploit the clustering structure of context features. All the previous neural language models only consider partitioning the embedding matrix to exploit the clustering structure of the word embedding to achieve prediction speedup. • To make prediction for a context embedding, after obtaining cluster assignment from screening model, L2S only needs to evaluate a small set of vocabulary in that cluster. Therefore, L2S can significantly reduce the inference time complexity from O(Ld) to O((r + L̄)d) with L̄ L and r L where d is the context vector’ dimension; L is the vocabulary size, r is the number of clusters, and L̄ is the average word/candidate size inside clusters. • We propose to form a joint optimization problem to learn both screening model for clustering as well as the candidate label set inside each cluster simultaneously. Using the Gumbel trick (Jang et al., 2017), we are able to train the screening network end-to-end on the training data. • We show in our experiment that L2S can quickly identify the top-k prediction words in the vocabulary an order of magnitude faster than original softmax layer inference for machine translation and next words prediction tasks.

heading: 2 RELATED WORK
text: We summarize previous works on speeding up the softmax layer computation.
Algorithms for speeding up softmax in the training phase. Many approaches have been proposed for speeding up softmax training. Jean et al. (2014); Mnih & Teh (2012) proposed importance sampling techniques to select only a small subset as “hard negative samples” to conduct the updates. The hierarchical softmax-based methods (Morin & Bengio, 2005; Grave et al., 2017) use the tree
structure for decomposition of the conditional probabilities, constructed based on external word semantic hierarchy or by word frequency. Most hierarchical softmax methods cannot be used to speed up inference time since they only provide a faster way to compute probability for a target word, but not for choosing top-k predictions as they still need to compute the logits for all the words for inference. One exception is the recent work by Grave et al. (2017), which constructs the tree structure by putting frequent words in the first layer—so in the prediction phase, if top-k words are found in the first layer, they do not need to go down the tree. We provide comparison with this approach in our experiments.
Algorithms for Maximum Inner Product Search (MIPS). Given a query vector and a database with n candidate vectors, MIPS aims to identify a subset of vectors in the database that have top-k inner product values with the query. Top-k softmax can be naturally approximated by conducting MIPS. Here we summarize existing MIPS algorithms:
• Hashing: (Shrivastava & Li, 2014; Neyshabur & Srebro, 2015) proposed to reduce MIPS to nearest neighbor search (NNS) and then solve NNS by Locality Sensitive Hashing (LSH) (Indyk & Motwani, 1998). • Database partitioning: PCA tree (Sproull, 1991) partitions the space according to the directions of principal components and shows better performance in practice. Bachrach et al. (2014) shows tree-based approaches can be used for solving MIPS but the performance is poor for high dimensional data. • Graph-based algorithm: Malkov et al. (2014); Malkov & Yashunin (2016a) recently developed an NNS algorithm based on small world graph. The main idea is to form a graph with candidate vectors as nodes and edges are formed between nearby candidate vectors. The query stage can then done by navigating in this graph. Zhang et al. (2018) applies the MIPS-to-NNS reduction and shows graph-based approach performs well on neural language model prediction. • Direct solvers for MIPS: Some algorithms are proposed to directly tackle MIPS problem instead of transforming to NNS. Guo et al. (2016); Wu et al. (2017) use quantization-based approach to approximate candidate set. Another Greedy MIPS algorithm is recently proposed in (Yu et al., 2017), showing significant improvement over LSH and tree-based approaches.
Algorithms for speeding up softmax in inference time. MIPS algorithms can be used to speed up the prediction phase of softmax layer, since we can view context vectors as query vectors and weight vectors as database. In the experiments, we also include the comparisons with hashing-based approach (LSH) (Indyk & Motwani, 1998), partition-based approach (PCA-tree (Sproull, 1991)) and Greedy approach (Yu et al., 2017). The results show that they perform worse than graph-based approach (Zhang et al., 2018) and are not efficient if we want to keep a high precision.
For NLP tasks, there are two previous attempts to speed up softmax layer prediction time. (Shim et al., 2017) proposed to approximate the weight matrix in the softmax layer with singular value decomposition, find a smaller candidate set based on the approximate logits, and then do a finegrained search within the subset. (Zhang et al., 2018) transformed MIPS to NNS and applied graphbased NNS algorithm to speed up softmax. In the experiments, we show our algorithm is faster and more accurate than all these previous algorithms. Although they also have a screening component to select an important subset, our algorithm is able to learn the screening component using training data in an end-to-end manner to achieve better performance.

heading: 3 ALGORITHM
text: Softmax layer is the main bottleneck when making prediction in neural language models. We assume L is the number of output tokens, W ∈ Rd×L is the weight matrix of the softmax layer, and b ∈ RL is the bias vector. For a given context vector h ∈ Rd (such as output of LSTM), softmax layer first computes the logits
xs = w T s h + bs for s = 1, · · · , L (1)
where ws is the s-th column of W and bs is the s-th entry of b, and then transform logits into probabilities ps = e
xs∑L l=1 e xl for s = 1, · · · , L. Finally it outputs the top-k candidate set by sorting
the probabilities [p1, · · · , pL], and uses this information to perform beam search in translation or predict next word in language model.
To speedup the computation of top-k candidates, all the previous algorithms try to exploit the structure of {ws}Ls=1 vectors, such as low-rank, tree partitioning or small world graphs (Zhang et al., 2018; Shim et al., 2017; Grave et al., 2017). However, in NLP applications, there exists strong structure of context vectors {h} that has not been exploited in previous work. In natural language, some combinations of words appear very frequently, and when some specific combinations appear, the next word should only be within a small subset of vocabulary. Intuitively, if two context vectors hi and hj are similar, meaning similar context, then their candidate label sets C(hi) and C(hj) can be shared. In other words, suppose we already know the candidate set of C(hi) for hi, to find the candidate set for hj , instead of computing the logits for all L tokens in the vocabulary, we can narrow down the candidate sets to be C(hi), and only compute the logits in C(hi) to find top-k prediction for hj .
The Prediction Process. Suppose the context vectors are partitioned into r disjoint clusters and similar ones are grouped in the same partition/cluster, if a vector h falls into one of the cluster, we will narrow down to that cluster’s label sets and only compute the logits of that label set. This screening model is parameterized by clustering weights v1, . . . ,vr ∈ Rd and label candidate set for each cluster c1, . . . , cr ∈ {0, 1}L. To predict a hidden state h, our algorithm first computes the cluster indicator
z(h) = arg max t
vTt h, (2)
and then narrows down the search space to C(h) := {s | cz(h),s = 1}. The exact softmax is then computed within the subset C(h) to find the top-k predictions (used in language model) or compute probabilities used for beam search in neural machine translation. As we can see the prediction time includes two steps. The first step has r inner product operations to find the cluster which takesO(rd) time. The second step computes softmax over a subset, which takes O(L̄d) time where L̄ (L̄ L) is the average number of labels in the subsets. Overall the prediction time for a context embedding h is O((r + L̄)d), which is much smaller than the O(Ld) complexity using the vanilla softmax layer. Figure 1 illustrates the overall prediction process.
However, how to learn the clustering parameter {vt}rt=1 and the candidate sets {ct}rt=1? We found that running spherical kmeans on all the context vectors in the training set can lead to reasonable results (as shown in the appendix), but can we learn even parameters to minimize the prediction error? In the following, we propose an end-to-end procedure to learn both context clusters and candidate subsets simultaneously to maximize the performance.
Learning the clustering. Traditional clustering algorithms such as kmeans on Euclidean space or cosine similarity have two drawbacks. First, they are discrete and non-differentiable, and thus hard to use with back-propagation in the end-to-end training process. Second, they only consider clustering on {hi}Ni=1, without taking the predicted label space into account. In this paper, we consider learning the partition through Gumbel-softmax trick. We will briefly summarize the technique and direct the reader to (Jang et al., 2017) for further details on these techniques. In Table 5 in the appendix, we compare our proposed method to traditional spherical-kmeans to show that it can further improve the performance.
First, we turn the deterministic clustering in Eq(2) into a stochastic process: the probability that h belongs to cluster t is modeled as
P (t|h) = exp(v T t h)∑r
j=1 exp(v T j h)
, ∀t, and z(h) = arg max t P (t|h). (3)
However, since argmax is a discrete operation, we cannot combine this operation with final objective function to find out better clustering weight vectors. To overcome this, we can re-parameterize Eq(3) using Gumbel trick. Gumbel trick provides an efficient way to draw samples z from the categorical distribution calculated in Eq(3):
m(h) = one hot(arg max j [gj + logP (j|h)]), (4)
where each gj is an i.i.d sample drawn from Gumbel(0, 1). We then use the Gumbel softmax with temperature = 1 as a continuous, differentiable approximation to argmax, and generate rdimensional sample vectors p = [p1, · · · , pr] which is approximately one-hot m(h) with
pt = exp(log(P (t|h)) + gt)∑r
j=1 exp(log(P (j|h)) + gj) ,∀t ∈ {1, . . . , r}. (5)
Using the Straight-Through (ST) technique proposed in (Jang et al., 2017), we denote p̄ = p + (one hot(arg maxj pj) − p) as the one-hot representation of p and assume back-propagation only goes through the first term. This enables end-to-end training with the loss function defined in the following section. We also use p̄(h) to denote the one-hot entry of p̄ (i.e., the position of the “one” entry of p̄).
Learning the candidate set for each cluster. For a context vector hi, after getting into the partition t, we will narrow down the search space of labels to a smaller subset. Let ct be the label vector for t-th cluster, we define the following loss to penalize the mis-match between correct predictions and the candidate set:
`(hi,yi) = ∑
s:yis=1
(1− cts)2 + λ ∑
s:yis=0
(cts) 2, (6)
where yi ∈ {0, 1}L is the ’ground truth’ label vector for hi that is computed from the exact softmax. We set y to be the label vector from full softmax because our goal is to approximate full softmax prediction results while having faster inference (same setting with (Shim et al., 2017; Zhang et al., 2018)). The loss is designed based on the following intuition: when we narrow down the candidate set, there are two types of loss: 1) When a candidate s (yis = 1) is a correct prediction but not in the candidate set (cts = 0), then our algorithm will miss this label. 2) When a candidate j (yis = 0) is not a correct prediction but it’s in the candidate set (cts = 1), then our algorithm will waste the computation of one vector product. Intuitively, 1) is much worse than 2), so we put a much smaller weight λ ∈ (0, 1) on the second term. The choice of true candidate set in y can be set according to the application. Throughout this paper, we set y to be the correct top-5 prediction (i.e., positions of 5-largest xs in Eq(1). yis = 1 means s is within the correct top-5 prediction of hi, while yis = 0 means it’s outside the top-5 prediction.
Final objective function: We propose to learn the partition function (parameterized by {vt}rt=1) and the candidate sets ({ct}rt=1) simultaneously. The joint objective function will be:
minimize v1,··· ,vr c1,··· ,cr N∑ i=1 ( ∑ s:yis=1 (1− cp̄(hi),s) 2 + λ ∑ s:yis=0 (cp̄(hi),s) 2) (7)
s.t. ct ∈ {0, 1}L ∀t = 1, . . . , r L̄ ≤ B ∀i = 1, . . . , N
where N is the number of samples, L̄ is the average label size defined as L̄ = ( ∑N
i=1 ∑L s=1 cp̄(hi),s)/N , p̄(hi) is the index for where p̄t = 1 for t = 1, · · · , r; B is the desired average label/candidate size across different clusters which could be thought as prediction time budget. Since L̄ is related to the computation time of proposed method, by enforcing L̄ ≤ B we can make sure label sets won’t grow too large and desired speed-up rate could be achieved. Note that p̄(hi) is for clustering assignment, and thus a function of clustering parameters v1, · · · ,vr as shown in Eq(3).
Optimization. To solve the optimization problem in Eq (7), we apply alternating minimization. First, when fixing the clustering (parameters {vt}) to update the candidate sets (parameters {ct}), the problem is identical to the classic “Knapsack” problem—each ct,s is an item, with weight proportional to number of samples belonging to this cluster, and value defined by the loss function of Eq(7), and the goal is to maximize the value within weight capacity B. There is no polynomial time solution with respect to r, so we apply a greedy approach to solve it. We sort items by the value-capacity ratio and add them one-by-one until reaching the upper capacity B.
When fixing {ct} and learning {vt}, we convert the cluster size constraint to objective function by Lagrange-Multiplier:
minimize v1,··· ,vr N∑ i=1 ( ∑
j:yis=1
(1− cp̄(hi)s) 2 + λ ∑ j:yis=0 (cp̄(hi)s) 2) + γmax(0, L̄−B) (8)
s.t. ct ∈ {0, 1}L ∀t = 1, . . . , r, and simply use SGD since back-propagation is available after applying Gumbel trick. To deal with L̄ in the mini-batch setting, we replace it by the moving-average, updated at each iteration when we go through a batch of samples. The overall learning algorithm is given in Algorithm 1.
Models vocabulary size dimension
PTB-Small 10000 200 PTB-Large 10000 1500 WikiText-103 80000 400 NMT: DE-EN 24725 500 NMT: EN-VE 22749 200

heading: 4 EXPERIMENTS
text: We evaluate our method on two tasks: Language Modeling (LM) and Neural Machine Translation (NMT). For LM, we use the Penn Treebank Bank (PTB) dataset (Marcus et al., 1993). For NMT, we use the IWSLT 2014 German-to-English translation task (Cettolo et al., 2014) and IWSLT 2015 English-Vietnamese data (Luong & Manning, 2015). All the models use a 2-layer LSTM neural network structure. For IWSLT-14 DE-EN task, we use the PyTorch checkpoint provided by OpenNMT (Klein et al., 2017). For IESLT-15 EN-VE task, we set the dimension of hidden size to be 200, and the rest follows the default training hyperparameters of OpenNMT. For PTB, we train a 2-layer LSTM-based language model on PTB from scratch with two setups: PTB-Small and PTB-Large. The LSTM hidden state sizes are 200 for PTB-Small and 1500 for PTB-Large, so are their embedding sizes. Vocabulary sizes and the hidden dimension used in our experiments are summarized in Table 1.
We verified that all these models achieved benchmark performance on the corresponding datasets as reported in the literature. We then apply our method to accelerate the inference of these benchmark models.

heading: 4.1 COMPETING ALGORITHMS
text: We include the following algorithms in our comparisons:
• L2S (Our proposed algorithm): the proposed learning-to-screen method. Number of clusters and average label size across clusters will be the main hyperparameters affecting computational time. We could control the tradeoff of time and accuracy by fixing the number of clusters and varying the size constraint B. For all the experiments we set parameters λ = 0.0003 and γ = 10. We will show later that L2S is robust to different numbers of clusters. • FGD (Zhang et al., 2018): transform the softmax inference problem into nearest neighbor search (NNS) and solve it by a graph-based NNS algorithm.
• SVD-softmax (Shim et al., 2017): a low-rank approximation approach for fast softmax computation. We vary the rank of SVD to control the tradeoff between prediction speed and accuracy. • Adaptive-softmax (Grave et al., 2017): a variant of hierarchical softmax that was mainly developed for fast training on GPUs. However, this algorithm can also be used to speedup prediction time (as discussed in Section 2), so we include it in our comparison. The tradeoff is controlled by varying the number of frequent words in the top level in the algorithm. • Greedy-MIPS (Yu et al., 2017): the greedy algorithm for solving MIPS problem. The tradeoff is controlled by varying the budget parameter in the algorithm. • PCA-MIPS (Bachrach et al., 2014): transform MIPS into Nearest Neighbor Search (NNS) and then solve NNS by PCA-tree. The tradeoff is controlled by varying the tree depth. • LSH-MIPS (Neyshabur & Srebro, 2015): transform MIPS into NNS and then solve NNS by Locality Sensitive Hashing (LSH). The tradeoff is controlled by varying number of hash functions.
We implement L2S, SVD-softmax and Adaptive-softmax in numpy. For FGD, we use the C++ library implemented in (Malkov & Yashunin, 2016b; Boytsov & Naidan, 2013) for the core NNS operations. The last three algorithms (Greedy-MIPS, PCA-MIPS and LSH-MIPS) have not been used to speed up softmax prediction in the literature and they do not perform well in these NLP tasks, but we still include them in the experiments for completeness. We use the C++ code by (Yu et al., 2017) to run experiments for these three MIPS algorithms.
Since our focus is to speedup the softmax layer which is known to be the bottleneck of NLP tasks with large vocabulary, we only report the prediction time results for the softmax layer in all the experiments. To compare under the same amount of hardware resource, all the experiments were conducted on an Intel Xeon E5-2620 CPU using a single thread.

heading: 4.2 PERFORMANCE COMPARISONS
text: To measure the quality of top-k approximate softmax, we compute Precision@k (P@k) defined by |Ak ∩ Sk|/k, where Ak is the top-k candidates computed by the approximate algorithm and Sk is
the top-k candidates computed by exact softmax. We present the results for k = 1, 5. This measures the accuracy of next-word-prediction in LM and NMT. To measure the speed of each algorithm, we report the speedup defined by the ratio of wall clock time of the exact softmax to find top-k words divided by the wall clock time of the approximate algorithm.
For each algorithm, we show the prediction accuracy vs speedup over the exact softmax in Figure 2, 3, 4, 5, 6, 7 (The last three are in the appendix). We do not show the results for PCA-MIPS and LSHMIPS in the figures as their curves run outside the range of the figures. Some represented results are reported in Table 2. These results indicate that the proposed algorithm significantly outperforms all the previous algorithms for predicting top-k words/tokens on both language model (next word prediction) and neural machine translation.
Next, we measure the BLEU score of the NMT tasks when incorporating the proposed algorithm with beam search. We consider the common settings with beam size = 1 or 5, and report the wall clock time of each algorithm excluding the LSTM part. We only calculate log-softmax values on reduced search space and leave probability of other vocabularies not in the reduced search space to be 0. From the precision comparison, since FGD shows better performance than other completing methods in Table 2, we only compare our method with state-of-the-art algorithm FGD in Table 3 in terms of BLEU score. Our method can achieve more than 13 times speed up with only 0.14 loss in BLEU score in DE-EN task with beam size 5. Similarly, our method can achieve 20 times speed up in EN-VE task with only 0.08 loss in BLEU score. In comparison, FGD can only achieve less than 3-6 times speed up over exact softmax to achieve a similar BLEU score. We also compare our algorithm with other methods using perplexity as a metric in PTB-Small and PTB-Large as shown in Table 6 in the appendix. We observe more than 5 times speedup over using full softmax without losing much perplexity (less than 5% difference). More details can be found in the appendix.
In addition, we also show some qualitative results of our proposed method on DE-EN translation task in Table 8 to demonstrate that our algorithm can provide similar translation results but with faster inference time.

heading: 4.3 SELECTION OF THE NUMBER OF CLUSTERS
text: Finally, we show the performance of our method with different number of clusters in Table 4. When varying number of clusters, we also vary the time budget B so that the prediction time including finding the correct cluster and computing the softmax in the candidate set are similar. The results indicate that our method is quite robust to number of clusters. Therefore, in practice we suggest to just choose the number of clusters to be 100 or 200 and tune the “time budget” in our loss function to get the desired speed-accuracy tradeoff.

heading: 5 CONCLUSION
text: In this paper, we proposed a new algorithm for fast softmax inference on large vocabulary neural language models. The main idea is to use a light-weight screening model to predict a smaller subset of candidates, and then conduct exact search within that subset. By forming a joint optimization problem, we are able to learn the screening network end-to-end using the Gumbel trick. In the experiment, we show that the proposed algorithm achieves much better inference speedup than stateof-the-art algorithms for language model and machine translation tasks.

heading: 6 ACKNOWLEDGEMENT
text: We are grateful to Ciprian Chelba for the fruitful comments, corrections and inspiration. CJH acknowledges the support of NSF via IIS-1719097, Intel faculty award, Google Cloud and Nvidia.

heading: 7 APPENDIX

heading: 7.1 PRECISION@5 RESULTS

heading: 7.2 COMPARISON TO SPHERICAL-KMEANS INITIALIZATION
text: Since we firstly initialize parameters in our method by Shperical-KMEANS, we also show in Table 5 that L2S can further improve over the baseline clustering methods. Notice that even the basic
Spherical-KMEANS can outperform state-of-the-art methods. This shows that clustering structure of context features is a key to perform fast prediction.

heading: 7.3 PERPLEXITY RESULTS
text: Finally, we go beyond top-k prediction and apply our algorithm to speed up the perplexity computation for language models. To get perplexity, we need to compute the probability of each token appeared in the dataset, which may not be within top-k softmax predictions. In order to apply a top-k approximate softmax algorithm for this task, we adopt the low-rank approximation idea proposed in (Shim et al., 2017). For tokens within the candidate set, we compute the logits using exact inner product computation, while for tokens outside the set we approximate the logits by W̃h where W̃ is a low-rank approximation of the original weight matrix in the softmax layer. The probability can then be computed using these logits. For all the algorithms, we set the rank of W̃ to be 20 for PTB-Small and 200 for PTB-Large. The results are presented in Table 6. We observe that our method outperforms previous fast softmax approximation methods for computing perplexity on both PTB-small and PTB-large language models.

heading: 7.4 RESULTS ON LARGER VOCABULARY SIZE
text: In this section we validate that proposed L2S method can still work effectively on models with larger vocabulary sizes. We demonstrate this on Wikitext-103 dataset Merity et al. (2016) by choosing vocabulary size to be 80k. We follow the model and training procedure as in Merity et al. (2016). Precision@1 and Prevision@5 results of L2S are shown in Fig 9. We can see that L2S could speed-up over 15 times without losing much accuracy.

heading: 7.5 SPEED-UP RATE VERSES VOCABULARY SIZE.
text: We conduct experiments on PTB-Small setup to study the effect of vocabulary size to speedup. We compress the size of PTB dataset by choosing most frequent words and set the rest of other words to be ”unknown” and follow the same training procedure. The cluster size is set to 200 for all setups, and the results are shown in Fig 8. Basically, the results show that the compression rate could be made higher when vocabulary size is larger. This is reasonable as the occurrence of the words follow zipf’s distribution in natural languages. So when vocabulary size goes large, we could imagine many words only need to be searched a few times so full softmax computation will not be very efficient. Thus the speed-up rate could be improved further.

heading: 7.6 QUALITATIVE RESULTS ON CLUSTERING
text: We selected few clusters and show the words contained in that cluster to demonstrate qualitative results of L2S. Examples are listed in Table 7. Basically, some frequent words will appear in almost all clusters in order to make sure the predicted accuracy is high. For example, ”the”, ”in”, and ”a”
those tokens appear in all clusters shown in the table. On the other hand, each cluster might contain some
In the cluster 1, we can observe that all words related to timing is included. In the second cluster, some terms related to financial investment are included. More interestingly, from the term ”overthe-counter” in third cluster we know that it’s about finance. And we also observe that important cities of finance world such as Tokyo, ”Hong” Kong and London are included. Since these trading trend might be seasonal, we also see that many monthes are included in this cluster.

heading: 7.7 QUALITATIVE RESULTS ON NEURAL MACHINE TRANSLATION
text: We select some translated sentences of DE-EN task shown in Table 8 to demonstrate that our algorithm can provide similar translations but with faster inference time.

Review summaries:
This paper presents an approximation to the softmax function to reduce the computational cost at inference time and the proposed approach is evaluated on language modeling and machine translation tasks. The main idea of the proposed approach is to pick a subset of the most probable outputs on which exact softmax is performed to sample top-k targets. The proposed method, namely Learning to Screen (L2S), learns jointly context vector clustering and candidate subsets in an end-to-end fashion, so that it enables to achieve competitive performance.

The paper proposes a way to speed up softmax at test time, especially when top-k words are needed. The idea is clustering inputs so that we need only to pick up words from a learn cluster corresponding to the input. The experimental results show that the model looses a little bit accuracy in return of much faster inference at test time.