Searching for Efficient Transformers for Language Modeling

1 Introduction . Transformers [ 1 ] have been used extensively in many NLP advances over the past few years ( e.g. , [ 2 , 3 , 4 , 5 , 6 , 7 ] ) . With scaling , Transformers have produced increasingly better performance [ 3 , 7 , 8 , 9 ] , but the costs of training larger models have become prohibitively expensive . In this paper , we aim to reduce the training costs of Transformer language models . To this end , we propose searching for more efficient alternatives to Transformer by modifying its TensorFlow computation graph [ 10 ] . Given a search space of TensorFlow programs , we use evolution [ 11 , 12 , 13 , 14 , 15 , 16 , 17 ] to search for models that achieve as low of a validation loss as possible given a fixed amount of training compute . An advantage of using TensorFlow programs as the search space is that it is easier to find simple low-level improvements to optimize Transformers . We focus on decoder-only auto-regressive language modeling ( LM ) , because of its generality and success [ 18 , 7 , 19 , 20 , 21 ] .2 1https : //github.com/google-research/google-research/tree/master/primer 2We provide details of our primitives search in TensorFlow , but the same approach can also be applied to other deep learning libraries . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) , virtual . The discovered model , named Primer ( PRIMitives searched transformER ) , exhibits strong performance improvements over common Transformer variants on auto-regressive language modeling . Our experiments show that Primer has the benefits of ( 1 ) achieving a target quality using a smaller training cost , ( 2 ) achieving higher quality given a fixed training cost , and ( 3 ) achieving a target quality using a smaller inference cost . These benefits are robust and hold across model sizes ( 20M to 1.9B parameters ) , across compute scales ( 10 to 105 accelerator hours ) , across datasets ( LM1B , C4 , PG19 [ 22 ] ) , across hardware platforms ( TPUv2 , TPUv3 , TPUv4 and V100 ) , across multiple Transformer codebases using default configurations ( Tensor2Tensor , Lingvo , and T5 ) and across multiple model families ( dense Transformers [ 1 ] , sparse mixture-of-experts Switch Transformers [ 8 ] , and Synthesizers [ 23 ] ) . We open source these comparisons to help with the reproducibility of our results.1 Our main finding is that the compute savings of Primer over Transformers increase as training cost grows , when controlling for model size and quality . These savings follow a power law with respect to quality when using optimally sized models . To demonstrate Primer ’ s savings in an established training setup , we compare 500M parameter Primer to the original T5 architecture , using the exact configuration used by Raffel et al . [ 5 ] applied to auto-regressive language modeling . In this setting , Primer achieves an improvement of 0.9 perplexity given the same training cost , and reaches quality parity with the T5 baseline model using 4.2X less compute . We further demonstrate that Primer ’ s savings transfer to one-shot evaluations by comparing Primer to Transformer at 1.9B parameters in a setup similar to GPT-3 XL [ 7 ] . There , using 3X less training compute , Primer achieves similar performance to Transformer on both pretraining perplexity and downstream one-shot tasks . Our analysis shows that the improvements of Primer over Transformer can be mostly attributed to two main modifications : squaring ReLU activations and adding a depthwise convolution layer after each Q , K , and V projection in self-attention . These two modifications are simple and can be dropped into existing Transformer codebases to obtain significant gains for auto-regressive language modeling . 2 Search Space and Search Method . Searching Over TensorFlow Programs : To construct a search space for Transformer alternatives , we use operations from TensorFlow ( TF ) . In this search space , each program defines the stackable decoder block of an auto-regressive language model . Given input tensors X ∈ Rn×d that represent sequences of length n with embedding length d , our programs return tensors of the same shape . When stacked , their outputs represent next-token prediction embeddings at each sequence position . Our programs only specify model architectures and nothing else . In other words , the input and output embedding matrices themselves , as well as input preprocessing and weight optimization are not within the scope of our programs . Figure 1 shows how programs are constructed in our search space . Each program is built from an evolutionary search DNA , which is an indexed collection of subprograms . SUBPROGRAM 0 is the MAIN ( ) function that is the execution entry point , and the other subprograms are part of the DNA ’ s subprogram bank . Each subprogram is an indexed array of instructions with no length constraints . An instruction is an operation with a set of input arguments . The operation denotes the function that the instruction executes . Each operation maps to either a TF function from the primitives vocabulary or another subprogram in the DNA subprogram bank . The primitives vocabulary is comprised of simple primitive TF functions , such as ADD , LOG , and MATMUL ( see Appendix A.1 for details ) . It is worth emphasizing that high-level building blocks such as self-attention are not operations in the search space , but can be constructed from our low-level operations . The DNA ’ s subprogram bank is comprised of additional programs that can be executed as functions by instructions . Each subprogram can only call subprograms with a higher index in the subprogram bank , which removes the possibility of cycles . Each instruction ’ s argument set contains a list of potential argument values for each instruction operation . The set of argument fields represents the union of fields that all the operation primitives use : • Input 1 : The index of the hidden state that will be used as the first tensor input . The index of each hidden state is the index of the instruction that produced it , with the subprogram ’ s input states at indexes 0 and 1 . An example of an operation that uses this is SIN . • Input 2 : The index of the second tensor input . This is only used by operations that are binary with respect to tensor inputs . An example of an operation that uses this is ADD . • Constant : A real valued constant . An example of an operation that uses this is MAX ; tf.math.maximum ( x , C ) for C = 0 is how we express the Transformer ’ s ReLU activation . • Dimension Size : An integer representing the output dimension size for transformations that utilize weight matrices . An example of an operation that uses this is CONV 1X1 , the dense projection used by the Transformer ’ s attention projections and feed forward portions . See Appendix A.2 for how we employ relative dimensions [ 13 ] to resize our models . Our search subprograms are converted to TF programs by converting each subprogram instruction to a corresponding line of TF code , one at a time in indexing order . To create the TF line , the instruction operation is mapped to the corresponding TF primitive function or DNA subprogram , and any relevant arguments are plugged in ( see Appendix A.1 for the full TF primitives vocabulary , including argument mappings ) ; the other arguments are ignored . The TF tensor that is generated by the final instruction is taken as the subprogram output . We do not use TF Eager and so a useful property of the constructed programs is that irrelevant nodes that do not contribute to the programs ’ outputs are ignored as per TF ’ s original deferred execution design [ 10 ] . See Figure 2 for an illustration of how subprograms are converted to TF graphs and see Appendix A.2 for more details on how TF graphs are constructed , including how we handle causal masking . Evolutionary Search : The goal of our evolutionary search is to find the most training efficient architecture in the search space . To do this , we give each model a fixed training budget ( 24 TPUv2 hours ) and define its fitness as its perplexity on the One Billion Words Benchmark ( LM1B ) [ 24 ] in Tensor2Tensor [ 25 ] . This approach , which we call an implicit efficiency objective by fixed training budget , contrasts previous architecture search works that explicitly aim to reduce training or inference step time when optimizing for efficiency [ 26 , 27 , 28 , 29 ] . Our objective is different in that the trade-off between step time and sample efficiency is implicit . For instance , a modification that doubles step time , but triples sample efficiency is a good modification in our search , as it ultimately makes the architecture more compute efficient . Indeed , the modifications we find to be most beneficial , squaring ReLUs and adding depthwise convolutions to attention , increase training step time . However , they improve the sample efficiency of the model so much that they decrease the total compute needed to reach a target quality , by drastically reducing the number of training steps needed to get there . The search algorithm we use is Regularized Evolution [ 30 ] with hurdles [ 13 ] . We configure our hurdles using a 50th percentile passing bar and space them such that equal compute is invested in each hurdle band ; this reduces the search cost by a factor of 6.25X compared to the same experiment with full model evaluations ( see Appendix A.3 for more details ) . Additionally , we use 7 training hours as a proxy for a full day ’ s training because a vanilla Transformer comes within 90 % of its 24 hour training perplexity with just 7 hours of training . This reduces the search cost further by a factor of 3.43X , for a total compute reduction factor of 21.43X . So , although our target is to improve 24 hour performance , it only takes about 1.1 hours to evaluate an individual on average ( see Appendix A.4 for more search specifics , including mutation details and hyperparameters ) . We run our search for ∼25K individuals and retrain the top 100 individuals on the search task to select the best one . Our search space is different from previous search spaces ( see architecture search survey by [ 31 ] ) , which are often heavily biased such that random search performs well ( see analysis by [ 32 , 33 , 34 ] ) . As our search space does not have this bias , 78 % of random programs in our space with length equal to a Transformer program can not train more than five minutes , due to numerical instability . Because of this open-endedness and abundance of degenerate programs , it is necessary to initialize the search population with copies of the Transformer [ 13 ] ( input embedding size dmodel = 512 , feed forward upwards projection size dff = 2048 , and number of layers L = 6 ) ( Figure 3 ) . To apply this initialization to our search space , we must determine how to divide the Transformer program into subprograms . To do this , we divide along the lines of the machine learning concepts that constitute it . For instance , we create one subprogram each for self-attention , ReLU and layer norm , using commonly used implementations ( see Appendix A.5 for the complete list ) . We call this method conceptual initialization because it introduces a bias to the search through initialization , while leaving the search space for evolution and the action space for mutations open-ended . This contrasts the large amount of previous works that introduce bias through the search space . Although some works have also explored searching spaces that are open-ended like ours on miniature tasks [ 35 ] , we demonstrate that our techniques can scale to full sized deep learning regimes ( see Section 4 ) .

Summary: This paper studies how to improve the transformer architecture for the task of _autoregressive_ language modeling. To do so, they start from the transformer architecture and do low-level architecture search. They find an architecture, Primer, which provides better training efficiency / scaling laws over the original transformer. Most interestingly, they show that a simple variant of the transformer Primer-EZ, which simply adds a convolution layer and a squaring operation to the ReLU activations, can provide similar or better improvements than the full Primer. They test their implementation of Primer-EZ in multiple codebases and datasets, and show that the gains transfer.