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