Hash Layers For Large Sparse Models

We investigate the training of sparse layers that use different parameters for different inputs based on hashing in large Transformer models . Specifically , we modify the feedforward layer to hash to different sets of weights depending on the current token , over all tokens in the sequence . We show that this procedure either outperforms or is competitive with learning-to-route mixture-of-expert methods such as Switch Transformers and BASE Layers , while requiring no routing parameters or extra terms in the objective function such as a load balancing loss , and no sophisticated assignment algorithm . We study the performance of different hashing techniques , hash sizes and input features , and show that balanced and random hashes focused on the most local features work best , compared to either learning clusters or using longer-range context . We show our approach works well both on large language modeling and dialogue tasks , and on downstream fine-tuning tasks . 1 Introduction . Recent studies of Transformer models have shown a clear trend towards improvements with scale in data and model size [ 1 ] , mirroring the same trend in Machine Learning more generally . However , when architected naively , larger ( in terms of parameter count ) models are slower to train and to evaluate ; and at extreme scale , with current computer systems , necessitate complex engineering to facilitate communication between workers . To address these challenges , researchers have studied Mixtures-of-Experts ( MoE ) models [ 2 , 3 , 4 , 5 , 6 , 7 , 8 ] , where a “ gater ” routes computation through a sparse subset of the weights of the model ( the “ expert modules ” ) . Specifically in the setting of Transformers for Natural Language Processing ( NLP ) , recent approaches have led to state of the art performance in language modeling [ 8 ] . MoE models allow increasing the number of parameters in the model while holding steady the number of computations that affect a given sample . A key component to a MoE model is the routing ( gating ) strategy . While MoE models can be computationally advantageous per parameter compared to a dense model , they might be functionally less powerful per parameter . A poor routing strategy might lead to expert modules that are not properly specialized ( essentially making a stochastic ensemble model ) ; or overly specialized , using the data assignment function to overfit . Meanwhile , the routing strategy itself must be efficient . A standard approach is to train a layer of weights that makes the routing decision based upon the input to the layer to be routed . Classically , this may have been implemented with a softmax over the choice of expert modules , and fitted via backpropagation . However , a dense softmax requires all expert modules to run on all data points at train time , which negates the computational savings . Several works have shown that sparsity can be maintained during training , e.g . [ 9 , 7 , 8 , 10 ] . In particular , Switch Transformers [ 8 ] select the top expert per token using a softmax over the token ’ s hidden state , but require a load balancing term in the objective function or they can become imbalanced or degenerate , giving poor results . BASE Layers [ 10 ] employ a linear assignment algorithm to try to resolve the same problem . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . In this work , we describe a simple , sparse , efficient routing strategy based on hashing input tokens that is effective in the Transformers-for-NLP setting . We show this approach is effective on a number of datasets , comparing favorably to both Switch Transformers and BASE Layers . As the routing strategy requires no extra parameters , no change to the objective function or assignment algorithm , its simplicity means it is robust , fast and easy to implement . We provide detailed analysis to explain why our method works , and in which conditions . Given that when training very large models one may typically have only one shot given the required compute budget , and experimenters will be unable to try many parameter choices , we hence advocate our approach as a strong candidate for such a setting . 2 Background . Let us first introduce the Mixture-of-Experts setting where we apply our hash-based routing strategy . We use the same setting as [ 11 , 8 , 10 ] where a feedforward network ( FFN ) in a Transformer is replaced by its MoE version . Given a tokenized input sequence { x1 , x2 , . . . , xT } of T tokens , a representation for each token is computed in parallel by a standard Transformer [ 12 ] hL1 , h L 2 , . . . , h L T = TRANSFORMER ( x1 , x2 , . . . , xT ) . ( 1 ) The Transformer consists of L layers that computes final hidden states for each token , and each layer is composed of self-attention and FFN sublayers , where FFNs are two-layer fully connected networks h̄lt = SelfAttn ( h l−1 t ) h l t = FFN ( h̄ l t ) . ( 2 ) Here we omit skip-connections and normalization for brevity . We can then replace one or more of the FFN sublayers with expert modules . Replacing the FNN at layer l with K expert FFNs , their output is then mixed with some gating function g ( · ) : hlt = FFN ( h̄ l t ) → hlt = K∑ i=1 gi ( h̄ l t ) FFNi ( h̄ l t ) , t = 1 , . . . , T , ( 3 ) where importantly each token is routed to a different mixture of experts , as the gating function depends on the token ’ s specific hidden state h̄lt . Sparse MoE methods assume gating values gi are often zero , so only a few experts need to be computed for better efficiency . As expert FFNs do not share parameters , the number of parameters increases with K while the amount of computations per input token stays the same if the MoE FFN only routes to a single expert , and computation of gi is cheap . While this allows training of large capacity models with small compute budget , optimizing gi in the sparse setting can be tricky . 3 Method . In this paper we propose a simple gating mechanism that is especially efficient because only one expert is active , and it has no routing network parameters to be learnt . Recent work [ 11 , 8 , 10 ] has to learn parameters that determine the routing to expert modules based on hidden states , which have to be optimized in tandem with the expert weights themselves . This can potentially cause difficulty because during training membership for each expert is changing while it is trying to learn the mapping for those members . We instead advocate for a fixed mapping to experts . Namely , by hashing the tokens into a fixed number of buckets , each bucket corresponding to an expert : hlt = FFNhash ( xt ) ( h̄ l t ) , t = 1 , . . . , T. ( 4 ) While the FFN still takes the hidden state h̄lt as input , our routing function uses the original input token xt rather than the hidden state , see Figure 1 for a graphical depiction . We are free to choose from various possible hash functions , which we will consider below . However , for training purposes , the hash function is fixed in advance , and in this way , our routing mechanism requires no training and has no adjustable parameters . 3.1 Hash Functions . Hash functions have long been employed throughout Computer Science [ 13 ] , and can take a variety of forms . In our work , we generally employ pre-computed hash functions , which use a lookup table during learning – precomputed in advance – to map tokens to expert modules . We consider several kinds of hash functions as possible choices for routing tokens to expert modules . The simplest is Random Hash , wherein we assign every token to a fixed , random expert at initialization . Due to the Zipfian distribution of token frequency , this naturally produces imbalance across the different expert modules . As balancing has been previously shown to be important for training MoE models [ 8 , 10 ] , we also consider Balanced assignment . In this method , we build the lookup table before training the model using the training data distribution by greedily assigning the most frequent tokens to the emptiest buckets . The resulting assignment structure is significantly more balanced than Random Hashing , but not perfect , as the frequency of some tokens exceeds the ideal distribution . Random and Balanced hashing exploit the inductive bias of auto-regressive models and hash on the input token , but we also consider other possibilities : Bigram Hash uses the current and previous token ( xt−1 , xt ) rather than only the current token , while Previous Token Hash uses the previous token xt−1 , ignoring the current input . We also consider a sanity check which hashes based on the Position in the sequence , which we expect to have little impact , as absolute positions carry little information in natural language . Each of these hash functions is used to assess the value of the information being routed-on in our subsequent experimental analysis . As an upper baseline , we also evaluate using an Oracle Future Hash , which hashes based on the output token xt+1 , rather than input token . This Oracle Hash checks how powerful routing decisions can be in solving a task . Similarly , we also consider Predicted Future Token Hash , which utilizes a baseline Transformer to make a prediction of the output token , and then hashes over this prediction . Clustered Hashes Based on the intuition that similar tokens may want to be routed to the same expert , we also experiment with Clustered Hashes . We obtain clusters by performing k-means clustering with a fixed number of clusters using token embeddings from a baseline Transformer model . Each expert is assigned a centroid , and tokens are assigned to their closest cluster . Dispersed Hashes We also consider the opposite hypothesis : that similar-tokens should be placed in different buckets , where the assumption is that very similar tokens need fine distinctions which requires more model capacity ( hence assigning to different experts ) . To do this , we use the same k-means clusters as before , but distribute all tokens within each cluster equally across all buckets . 3.2 MultiHash Layers . In the standard FFN MoE approach , all K expert modules have independent parameters , but here we consider another option . It is known in the hashing literature that multiple hashes can provide better allocations in many contexts [ 14 ] . We consider such schemes in the context of sparse routing . Let us assume we are given N different hashing functions , and for a given input token x we compute these hashes , denoted as km = hashm ( x ) , m = 1 , . . . , N . Assuming the usual expert FFN is a function B ( relu ( A ( h ) ) ) where A : Rd → RD and B : RD → Rd , we split the linear layers into N segments , Am : Rd → RD/N and Bm : RD → Rd/N . Then we compute : v = relu ( [ Ak1 ( h ) , . . . , AkN ( h ) ] ) FFNMH ( h ) = [ Bk1 ( v ) , . . . , BkN ( v ) ] . That is , use hashing to select the parameters we are going to use for each segment , and then concatenate them together . The advantage is that we are now no longer reliant on the quality of a single hash function , but have multiple chances to produce good quality partitions . This perhaps can also be seen as analogous to the multi-head attention process already used in Transformers .

Summary: The authors propose using hashing layer to replace the routing mechanism in the methods of mixture-of-expert such as Switch Transformers and BASE Layers. The main benefit of the proposed method is that no routing parameters need to learn and achieve quite good performance on the language modeling tasks. The authors also have a wide exploration of different hashing mechanisms, such as random hash, balanced hash, bigram hash, clustered hash, etc.