Learning to Screen for Fast Softmax  Inference on Large Vocabulary Neural Networks

1 INTRODUCTION . 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 . 2 RELATED WORK . 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 can not 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 . 3 ALGORITHM . 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 can not 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

Summary: 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.