Doubly Sparse: Sparse Mixture of Sparse Experts for Efficient Softmax Inference

Computations for the softmax function in neural network models are expensive when the number of output classes is large . This can become a significant issue in both training and inference for such models . In this paper , we present Doubly Sparse Softmax ( DS-Softmax ) , Sparse Mixture of Sparse Experts , to improve the efficiency for softmax inference . During training , our method learns a two-level class hierarchy by dividing entire output class space into several partially overlapping experts . Each expert is responsible for a learned subset of the output class space and each output class only belongs to a small number of those experts . During inference , our method quickly locates the most probable expert to compute small-scale softmax . Our method is learning-based and requires no knowledge of the output class partition space a priori . We empirically evaluate our method on several real-world tasks and demonstrate that we can achieve significant computation reductions without loss of performance . 1 INTRODUCTION . Deep learning models have demonstrated impressive performance in many classification problems ( LeCun et al. , 2015 ) . In many of these models , the softmax function/layer is commonly used to produce categorical distributions over the output space . Due to its linear complexity , the computation for the softmax layer can become a bottleneck with large output dimensions , such as language modeling ( Bengio et al. , 2003 ) , neural machine translation ( Bahdanau et al. , 2014 ) and face recognition ( Sun et al. , 2014 ) . In some models , softmax contributes to more than 95 % computation . This becomes more of an issue when the computational resource is limited , like mobile devices ( Howard et al. , 2017 ) . Many methods have been proposed to reduce softmax complexity for both training and inference phases . For training , the goal is to reduce the training time . Sampling based ( Gutmann & Hyvärinen , 2012 ) and hierarchical based methods ( Goodman , 2001 ; Morin & Bengio , 2005 ) were introduced . D-Softmax ( Chen et al. , 2015 ) and Adaptive-Softmax ( Grave et al. , 2016 ) , construct two levelhierarchies for the output classes based on the unbalanced word distribution for training speedup . The hierarchies used in these methods are either pre-defined or constructed manually , which can be unavailable or sub-optimal . Unlike training , in inference , our goal is not to computing the exact categorical distribution over the whole vocabulary , but rather to search for top-K classes accurately and efficiently . Existing work ( Shrivastava & Li , 2014 ; Shim et al. , 2017 ; Zhang et al. , 2018 ) on this direction focus on designing efficient approximation techniques to find the top-K classes given a trained model . Detailed discussions of related works are to be found in Section 4 . Our work aims to improve the inference efficiency of the softmax layer . We propose a novel Doubly Sparse softmax ( DS-Softmax ) layer . The proposed method is motivated by ( Shazeer et al. , 2017 ) , and it learns a two-level overlapping hierarchy using sparse mixture of sparse experts . Each expert is trained to only contain a small subset of entire output class space , while each class is permitted to belong to more than one expert . Given a set of experts and an input vector , the DS-Softmax first selects the top expert that is most related to the input ( in contrast to a dense mixture of experts ) , and then the chosen expert could return a scored list of most probable classes in it sparse subset . This method can reduce the linear complexity in original softmax significantly since it does not need to consider the whole vocabulary . We conduct experiments in different real tasks , ranging from language modeling to neural machine translation . We demonstrate our method can reduce softmax computation dramatically without loss of prediction performance . For example , we achieved more than 23x speedup in language modeling and 15x speedup in translation with similar performances . Qualitatively , we demonstrate learned two-level overlapping hierarchy is semantically meaningful on natural language modeling tasks . 2 DS-SOFTMAX : SPARSE MIXTURE OF SPARSE EXPERTS . 2.1 BACKGROUND . Before introducing our method , we first provide an overview of the background . Hierarchical softmax . Hierarchical softmax uses a tree to organize output space where a path represents a class ( Morin & Bengio , 2005 ) . There are a few ways to construct such hierarchies . Previous work ( Morin & Bengio , 2005 ; Chen et al. , 2015 ; Grave et al. , 2016 ) focus on building hierarchies with prior knowledge . Other approaches , like Mnih & Hinton ( 2009 ) , performed clustering on embeddings to construct a hierarchy . Our work aims to learn a two-level hierarchy while the major difference is that we allow overlapping in the learned hierarchy . Sparsely-gated mixture-of-experts . Shazeer et al . ( 2017 ) designed a sparsely gated mixture of experts model so that outrageously large networks can achieve significantly better performance in language modeling and translation . They borrowed conditional computation idea ( Bengio et al. , 2015 ) to keep similar computation even though the number of parameters increases dramatically . Their proposed sparsely-gated Mixture of Experts ( MoE ) only use a few experts selected by the sparsely gating network for computation on each example . The original MoE can not speedup softmax computation but serves as an inspiration for our model design . Group lasso . Group lasso has been commonly used to reduce effective features in linear model ( Friedman et al. , 2010 ; Meier et al. , 2008 ) . Recently , it has been applied in a neural network for regularization ( Scardapane et al. , 2017 ) and convolutional deep neural network speedup ( Wen et al. , 2016 ) . It has been demonstrated as an effective method to reduce the number of nodes in the neural network . In this work , we use group lasso to sparsify the experts . 2.2 OUR METHOD . Goodman ( 2001 ) studied a two-level hierarchy for language modeling , where each word belongs to a unique cluster . ( A “ cluster ” here refers to a cluster of words . ) From this perspective , our method can be thought of as an extension of their method to allow overlapping hierarchy . This is because , in language modeling , it is often difficult to exactly assign a word to a single cluster . For example , if we want to predict the next word of “ I want to eat ” and one possible correct answer is ” cookie ” , we can quickly notice that possible answer belongs to something eatable . If we only search for the right answer inside words with the eatable property , we can dramatically increase the efficiency . Even though words like ” cookie ” are one of the correct answers , it might also appear under some non-edible context such as “ a piece of data ” in computer science literature . Thus , a two-level overlapping hierarchy can naturally accommodate word homonyms like these by allowing each word to belong to more than one cluster . We believe that this observation is likely to be true in other applications besides language modeling . Doubly Sparse softmax ( DS-Softmax ) is designed to capture such overlapped two-level hierarchy among output classes . In DS-Softmax , the first level is the sparse mixture and second level contains several sparse experts . ( Here an expert can be thought as a similar concept as a cluster . ) The sparse mixture chooses the right expert/cluster while sparse experts are responsible to separate the full output space into multiple , overlapped and small class clusters . The design of mixture gating is inspired by Shazeer et al . ( 2017 ) , where each expert in their model needs to search whole output space . In contrast , DS-Softmax only searches a small subset of the output space and inference becomes significantly faster than full softmax . Sparse gating . The first level of sparsification is a sparse gating mechanism inspired by Shazeer et al . ( 2017 ) , which is designed to choose the right experts . The sparse gating outputs a sparse activation over a set of experts . For faster inference purpose , only the top-one expert is chosen . One major difference comparing to Shazeer et al . ( 2017 ) is described as follows . Suppose we have K experts . Given input activation vector h ∈ Rd and gating network weight W g ∈ RK×d , gating values Gk ( h ) , k = 1 , ... , K , are calculated and normalized prior to the selection as shown in Eq . 1 and then we choose the gate with the largest value gk = maxiGi ( h ) and set all other gates to be zero , which means the corresponding k-th expert is chosen . Gk ( h ) = exp ( W gk h ) ∑ k′ exp ( W g k′h ) , gk = { Gk ( h ) , if k = argmaxiGi ( h ) , 0 , otherwise . ( 1 ) Where W g ∈ RK×d is the weighting matrix for group selection . We should mention that despite only the top-1 expert is selected , Eq . 1 still allows the gradient to be back-propagated to whole W g. Given the sparse gate , we can compute the probability of class c under the context h as : O ( h ) = p ( c|h ) = exp ( ∑ k gkW e ( c , k ) h ) ∑ c′ exp ( ∑ k gkW e ( c′ , k ) h ) , ( 2 ) where W e ( c , k ) ∈ R d is softmax embedding weight vector for class c in expert k , and only one gk ( the chosen expert ) is nonzero in the formulation above . Note that our design of Eq . 2 is different from Shazeer et al . ( 2017 ) , in that the gating values gk are put inside the softmax . This choice is critical as selecting top-1 expert under Shazeer et al . ( 2017 ) , where normalization is done after top-K experts are selected , results in a constant that carries no gradient information . Also , the gating values can be interpreted as an inverse temperature term for final categorical distribution produced by the chosen expert k Hinton et al . ( 2015 ) , shown in Eq . 2 . A smaller gk gives a more uniform distribution and larger gk corresponds to a sharper one . Sparse experts with group lasso . The second level sparsification is the sparse experts that output a categorical distribution for only a subset output classes . To sparsify each expert , we apply group lasso loss to restrain the W e ( c , k ) , shown in Eq . 3 . Then , pruning is carried out for W e ( c , k ) during training with γ , a lasso threshold according to Eq . 4. expert 1 expert N expert 2 Noisy Cloning Training & Pruning expert 1+N expert 2N expert 2+N expert 1 expert N expert 2 expert 1 expert N expert 2 expert 1+N expert 2N expert 2+N Iteration : i Iteration : i + 1 Loading Balance . We denote the sparsity percentage out of full softmax in k-th expert as sparsityk and proportion of k-th expert activated as utilizationk . The sparsity percentage indicates the percent of pruned output classes . Then , the overall speedup compared to the full softmax can be calculated as 1/ ∑ k ( utilizationk ∗ sparsityk ) . Thus , better utilization is essential for speedup as well . For example , there is no speedup if the expert with full output space is always chosen . We borrow a similar loading balance function from Shazeer et al . ( 2017 ) in Eq . 5 . It encourages the utilization percentage of each expert to be balanced by maximizing the coefficient of variation ( CV ) for gating outputs . In addition , to encourage the exclusiveness of classes , we incorporate group lasso loss on the expert level where each class should only exist in only one expert as shown in Eq . 6 . Lload = −CV ∑ h∈H ( x ) G ( h ) , ( 5 ) Lexpert = ∑ k √∑ c ‖W e ( c , k ) ‖ 2 2 . ( 6 ) Mitosis training . Memory might become a bottleneck during training if we initialize all experts with full softmax . Therefore , we design a training scheme , called mitosis training , to reduce the memory requirement . The method is to initialize with a smaller model ( fewer number of experts ) and gradually breed to a bigger number after noisy cloning showed in Fig . 2 . For each cloning , the sparsity is inherited so that less memory is required . For example , in one of our experiments , we only need 3.25x memory with 64 experts compared to a full softmax implementation . The final training algorithm . Our final training objective , Lall , consists of a combination of the related contributions discussed above . We describe our training procedure in Algorithm 1 .

Summary: In this paper the authors introduce a new technique for softmax inference. In a multiclass setting, the idea is to take the output of a NN and turn it into a gating function to choose one expert. Then, given the expert, output a particular category. The first level of sparsity comes from the first expert. The second level of sparsity comes from every expert only outputting a limited set of output categories.