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 proposes a variant of the Transformer layer that includes sparsified dense and attention layers. For FFN layers, the author introduces a controller layer that selects only a subset of hidden dimensions to be used in the computation. For the attention layer, The author proposes a multiplicative transformation to simulate the permutation of the input hidden states and applies a convolution layer on top of it. The author shows no accuracy drop with the proposed method and speedup in the decoding phase.