Sparse is Enough in Scaling Transformers

1 Introduction . The field of natural language processing has seen dramatic improvements in recent years due to large neural networks based on the Transformer architecture . The original Transformer [ 42 ] significantly advanced state-of-the-art in machine translation . BERT [ 7 ] surpassed all previous methods on question answering , language inference and other NLP tasks and was followed by a line of models like T5 [ 30 ] that further improved these results . The GPT line of models [ 29 , 3 ] elevated language generation to the point that GPT-2 was invited to write short passages for the Economist and GPT-3 created whole articles almost indistinguishable from human-written ones . The benefits of this progress are undercut by the huge costs such models incur . Strubell et al . [ 36 ] estimate that training a single base BERT model costs $ 4k- $ 12k and emits as much CO2 as one passenger ’ s share of a 4-hour flight and later Patterson et al . [ 27 ] estimate that training GPT-3 has three times as much tCO2e ( metric tons of CO2 equivalent ) emissions as a SF-NY round trip flight . Data and serving costs are also forbidding : a single training run of BERT , for example , processes 128B tokens , and Google Translate reportedly1 serves over 143B words per day . With the growing popularity and size of these models , it is increasingly valuable to make them scale efficiently . In this work we propose Scaling Transformers with a separate sparse mechanism for the query , key , value and output layers ( QKV layers for short ) and combine it with sparse feedforward blocks to get a fully sparse Transformer architecture . To quantify the computational complexity of inference in Transformer models , recall the architecture of a Transformer decoder block . It consists of three parts : a masked self-attention layer , an encoderdecoder attention layer and a feedforward block . The sizes of these layers are parameterized by dmodel and dff . The base BERT model sets dmodel = 768 , the large BERT has dmodel = 1024 , the largest ∗Work done while at Google Research . 1https : //cutt.ly/skkFJ7a 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . GPT-2 has dmodel = 1600 and GPT-3 reaches dmodel = 12288 . For both BERT and GPT models the authors use dff = 4 dmodel . While decoding a token , the self-attention layer needs to activate four matrices of size dmodel × dmodel : one each for the queries , keys and values input to the attention and one for merging the output . In the encoder-decoder attention , the keys and values may already be cached , so only two matrices of size dmodel × dmodel are activated . The feedforward block consists of two matrices of size dmodel × dff , omitting small additional contribution of biases . The total adds up to : 4 d 2model + 2 d 2 model + 2 dmodel dff . This sum describes both the number of trainable weights of a single block and approximates well the number of floating-point operations needed for decoding a single token , except for the attention operations ( discussed later ) . The complexity is quadratic in dmodel ; for example , as dmodel increases 16-fold from base BERT to GPT-3 , the complexity of a single block grows 256-fold . In comparison Scaling Transformers use only 2dmodel √ dmodel = 2d 1.5 model parameters in QKV layers and yield results as good as the baseline ( fully dense ) Transformer with the same number of parameters and complexity : 8 d 1.5model + 4 d 1.5 model + 4 d 1.5 model . We were surprised that the fully sparse Scaling Transformers are indeed enough to match the results of the baseline Transformer on the large C4 dataset [ 30 ] ( Figure 1 ) . The improvement in complexity holds not just asymptotically but yields over 2.6x speedup in wall-clock hed decoding time already for a model with 800M parameters and 20x improvement for a model with 17B parameters , as shown in Table 1 . To verify that Scaling Transformers can be used with other Transformer improvements on real tasks , we create Terraformer – a Transformer model that uses reversible layers for memory efficiency and sparse attention to handle long sequences . We pre-train Terraformer on the C4 dataset and fine-tune it on the challenging task of summarizing arxiv articles . Terraformer yields results competitive to the state-of-the-art BigBird-Pegasus without using the Pegasus loss in pre-training ( Table 5 ) . 2 Related Work . As discussed in the previous section , large Transformer models brings significant improvements in performance , as seen in models such as GPT-3 [ 3 , 17 ] or T5 [ 44 , 30 ] . Training and inference incur a high computational cost at the scale of hundreds of billions of parameters . Numerous techniques improve the efficiency of Transformer models , and Gupta and Agrawal [ 11 ] divide them into several classes , including pruning , knowledge distillation , quantization , parameter sharing , efficient attention , and efficient feedforward . Model compression . Model pruning [ 24 , 2 ] makes matrices smaller by removing unneeded weights after or during training , however , the gains in computational complexity on sparse matrices often do 2The 800M model has 24 layers of Encoder & Decoder , dmodel = 1024 , 16 attn heads , attention-sparsity = 16 , ff-sparsity = 64 . We scale up this model to approximately 17B parameters with dmodel = 9216 and get up to 20x speedup in decoding compared to baseline dense model . This 17B param model has six layers of Encoder & Decoder , 96 attn heads , attention-sparsity = 64 , ff-sparsity = 256. not result in inference speedups on actual hardware [ 9 ] . Structured pruning based approaches [ 47 , 22 , 43 ] account for this challenge by leveraging sparsity in hardware in CPU and GPU architectures [ 1 ] . Our paper is different from pruning approaches in that it relies on dynamic sparsity wherein the feedforward layer loads only a subset of weights in the layer for each token . Our approach is complementary to model quantization studies [ 35 , 38 , 28 ] that use fewer bits for the weights . Model distillation . Several natural language models used for mobile inference [ 13 , 39 ] rely on distillation [ 32 ] to speed up inference from the pretrained large models . For example , [ 18 ] pretrains a large model and uses knowledge distillation along with pruning to get more than 10x faster inference . Instead of distilling a large model , our approach speeds up inference by reducing the number of weights loaded in memory from the model . Sparse attention . Sparse attention-based approaches have made the attention layer more efficient , especially for long sequences , by incorporating additional combinatorial mechanisms , as in [ 40 ] , or selecting a subset of tokens this layer attends to [ 31 , 5 , 19 , 37 , 15 , 4 ] or other approaches [ 12 ] . Our work is complementary to these approaches for sparse attention and reuses the advances on SOTA therein . Inference speedups in the attention layers also use bottleneck layers [ 39 ] or grouped convolutions [ 13 ] . Our work extends beyond the idea of grouped convolutions approach because each attention head is limited to using only a fixed part of the embedding while our work is able to permute the embeddings to improve model quality ; see Section 3.2 for details . Tensor Decomposition . The approaches discussed above significantly improve Transformer speed and handling of long sequences , however none of them addresses the fundamental scaling issue : even if we distill into a smaller model , quantize it and prune a percentage of the weights , the complexity still grows quadratically with dmodel . The final approach , which does attack this scaling issue , is called tensor decompositions in [ 11 ] . Unluckily , as the authors there note , the approach is most effective in dealing with large input and output embedding matrices and tends to produce lower performance than unstructured models if used inside the decoder block . Sparse feedforward . Mixture of experts approaches have been shown to achieve computational efficiency in training [ 33 , 21 , 34 ] , scaling up to a trillion parameters [ 8 ] . The key idea is to partition the dff-sized dimension into parts ( called experts ) and retrieve only one part per token , which reduces the complexity of the feedforward block from 2dmodeldff to 2dmodeldff/nexperts . These speedups are mostly measured in training speed , and the method focuses on feedforward blocks . In contrast to prior methods , we train a full weight matrix and then only activate specific parts of it for each input token during decoding ; see Section 3.1 . 3 Sparse is Enough . We study how to sparsify every part of the Transformer model—otherwise the non-sparse parts dominate decoding time and become a bottleneck . This means we need sparse equivalents for the feedforward blocks , for the dense Q , K , V and output layers in attention , and for the final dense layer before the softmax and loss . 3.1 Sparse Feedforward Layer . In a baseline Transformer , decoding speed is dominated by the execution cost of the feedforward block . Recall that this block consists of two fully-connected ( dense ) layers with a ReLU nonlinearity in between . The dimensionality of activation vectors between these 2 layers is usually denoted by dff and is often 4 or 8 times larger than the dimensionality of the activations in other places ( dmodel ) . We make use of the structure of the feedforward block to sparsify it . One main observation is that the ReLU in the middle creates a lot of zeros2 . We impose a fixed structure on this middle activation vector : only one float in every block of N will be allowed to be non-zero . Prior techniques prune weights or blocks from weight matrices and can be referred to as static sparsity . Our proposed technique will train a full weight matrix but only activate specific parts of it for each input token during decoding . We call this dynamic sparsity , because the model dynamically selects only a fraction of its parameters , and the selection is independent for each token . 2GeLU is another non-linearity often used in the Transformer feedforward block . Table 1 in [ 26 ] shows the same final loss for ReLU and GeLU on the C4 dataset , though , so in this work for simplicity , we focus on ReLU . We train a controller to determine which activation in each block can be non-zero ; the rest will be set to zero . This can be represented as Ysparse = max ( 0 , xW1 + b1 ) Controller ( x ) SparseFFN ( x ) = YsparseW2 + b2 where is element-wise multiplication . Note that each activation in Ysparse corresponds to a single column in W1 and a single row in W2 . Therefore , if we compute Controller ( x ) output first , we don ’ t have to use any columns in W1 or any rows in W2 that correspond to an activation set to zero by the controller . This allows for much faster decoding , as we have to process only 1 in N columns in W1 and rows in W2 ( see Figure 2 ( a ) ) . To design the controller to be computationally inexpensive , we project the input using a low-rank bottleneck dense layer . Figure 2 ( b ) illustrates the controller which produces the output as follows Controller ( x ) = argmax ( Reshape ( xC1C2 , ( −1 , N ) ) ) where C1 ∈ Rdmodel×dlowrank and C2 ∈ Rdlowrank×dff , with dlowrank usually set to ( dmodel/N ) . During inference the controller uses a discrete argmax function , but during training the model uses a softmax to calculate and sample from a distribution . The model learns to select which row/column will be non-zero using the Gumbel-Softmax trick for discretization . To determine the active row/column in each block , we reparameterize sampling from a Bernoulli distribution by using the Gumbel-Softmax trick [ 25 ] . Instead of using the logits in each block to directly sample a binary value , we add independent noise from the Gumbel distribution to each of the logits , and then select the binary value with the highest logit ( i.e. , argmax ) as the sample z . The argmax operation is not differentiable , but it can be approximated by a softmax with annealing temperature . Therefore , on the forward pass , we use the argmax to obtain a binary one-hot vector for each block , while on the backward pass , we approximate it with softmax . This approach is known as the Straight-Through Gumbel-Softmax estimator [ 14 ] . Ablations . We investigate the impact of sparse FF on the model equivalent to T5-large with varying levels of sparsity , with dmodel = 1024 , dff = 4096 , and 16 attention heads . When we set the sparsity level to N ( for e.g . N = 64 ) then every block of size N has one non-zero value activated for inference . During training , the controller uses the bottleneck layer with dlowrank = 64 and temperature of Gumbel softmax estimator set to 0.1 . To improve training stability , the controller in the forward pass will use the output of argmax that is a binary one-hot vector for each block with a probability of 30 % and otherwise it uses the output of softmax . Table 2 and Figure 3 show the perplexity and the decoding time of this model with varying levels of sparsity in feedforward layer . As the level of sparsity increases from 0 to 128 , we observe a significant decrease in the decoding time , while the neg-log-perplexity of the model with N = 64 sparsity is comparable to the baseline . We also checked the performance of the feedforward block with Mixture-of-Experts [ 33 ] style sparsity . As expected , this technique achieved decoding time comparable to sparse FF – 0.11s instead of 0.09s – but with its lack of granularity it achieved log-perplexity of 1.64 , worse than both our method and the dense baseline .

Summary: The paper presents an approach to training large-scale sparse transformers by sparsifying the activations of linear layers in the network (feedforward and Q, K, V projections). The approach is particularly appealing because it doesn't appear to require any hardware-specific or low-level optimizations to be efficient since it sparsifies activations. A controller network produces a block-wise one-hot mask on the activations and weights for the subsequent layer are selected on the fly based on this mask. Gradients are propagated through this hard-selection operation via the straight-through Gumbel softmax estimator.  The authors are able to achieve fairly substantial inference time speedups via sparsification. They also apply the same approaches to sparsifying models that are tailored towards handling long sequences such as the Reformer with LSH attention.