Scale Efficiently: Insights from Pretraining and Finetuning Transformers

1 INTRODUCTION . Training Transformers incurs both financial and environmental costs ( Schwartz et al. , 2019 ; Patterson et al. , 2021 ) . To this end , researchers and practitioners often have to work around fixed compute budgets and figure out the best ways to train their models . In lieu of the rising computation demand for training state-of-the-art Transformer ( Vaswani et al. , 2017 ; Devlin et al. , 2018 ; Raffel et al. , 2019 ; Brown et al. , 2020 ; Fedus et al. , 2021 ) models , the goal of this paper is to present insights and lessons from scaling Transformers and making them efficient and effective for transfer learning on downstream tasks . Despite the insights offered in scaling laws research ( Kaplan et al. , 2020 ; Hernandez et al. , 2021 ) there remain unresolved questions : Should one follow fixed scaling ratios ? If not , should one scale by depth ? Or by width ? Will scaling experiments on upstream pre-training generalize for downstream transfer ? Do scaling protocols for small models generalize to larger models ? Are scaling behaviours similar in all compute regions ? We hope the insights presented in this paper can be useful to both practitioners and researchers in informing their scaling decisions . Neural scaling laws ( Kaplan et al. , 2020 ) is a common resource that many look to for advice on scaling Transformer architectures . However , this paper limited its scope to an exhaustive study of upstream cross entropy on language modeling tasks . It is furthermore unclear if findings from ( Kaplan et al. , 2020 ) will transfer to downstream applications . Specifically , Kaplan et al . ( 2020 ) proposed that the performance of a Transformer language model strongly depends on model size and only weakly on its shape . They also argue that many model configurations with the same number of parameters perform similarly regardless of architectural details . Our work empirically confirms this on upstream training but finds a distinct discrepancy when considering practical downstream performance – a key insight that we believe is highly important . To this end , we conduct extensive experiments involving pre-training and fine-tuning over 200 transformer configurations ranging from 5M to 30B parameters . To the best of our knowledge , this is the largest empirical study of practical scaling of transformer to date that considers both upstream and practical downstream transfer . While there have been many proposed scaling protocols for ConvNets ( Tan and Le , 2019 ; Bello et al. , 2021 ) , there is still limited advice on scaling of transformer architectures , apart from ( Kaplan et al. , 2020 ; Li et al. , 2020 ; Levine et al. , 2020 ) . Hence , the key goal of this paper is to distill our experiences and insights with scaling Transformer architectures and share them with the broader community . The overall findings and insights of the paper can be summarized as follows : • We find that scaling laws may differ in upstream and downstream setups . Specifically , contrary to Kaplan et al . ( 2020 ) , we find that downstream performance strongly depends on shape and not only on model size . Hence , pretraining performance may not necessarily transfer to downstream applications . ( Figure 1 ) . • Our findings show that pre-training perplexity can often be a deceiving indicator of downstream quality and therefore model building based on upstream perplexity can be challenging . Scaling laws can differ substantially when considering metrics on actual downstream finetuning . ( Figure 1 ) • Given that empirical scaling laws differ when considering quality on the downstream , our work investigates the pareto-frontier of transformer configurations in this setup . We find that the canonical model configurations such as T5-Base and T5-Large sizes ( Raffel et al. , 2019 ) are relatively inefficient ( Figure 2 ) . Note that these sizes are based off the canonical BERT ( Devlin et al. , 2018 ) base and large sizes . • We find that scaling strategies differ at different compute regions , i.e. , applying same strategies at different compute regions ( small vs large ) has a different effect on model quality . This has practical implications since finding strategies at small scale might not necessarily transfer or generalize to higher compute regions ( section 4.2 ) . • After extensive empirical exploration of the pareto-frontier of transformer models , we propose a simple but effective scaling strategy which we call the DeepNarrow strategy . We show that we are able to obtain model quality on par or better than canonical model sizes ( e.g. , base ) with 50 % less parameters and being 40 % faster . While we highlight the limitations of this strategy , we also show that this DeepNarrow strategy is applicable to all model sizes . ( Table 4 ) . • To consider how generalized these scaling strategies are , we conduct additional experiments on Vision Transformers ( ViT ; Dosovitskiy et al. , 2020 ) to verify them in the vision domain . Moreover , on top of the 17 GLUE ( Wang et al. , 2018 ) / SuperGLUE ( Wang et al. , 2019 ) and SQuAD ( Rajpurkar et al. , 2016 ) tasks we employed in our extensive study , we verify our findings via additional downstream experiments across 12 diverse language tasks ( section A.1 ) . • We will release ( 1 ) the pre-trained checkpoints for our T5 models with improved scaling protocols and ( 2 ) all 100+ model checkpoints , including intermediate training checkpoints to the research community . We believe that this is a treasure trove of data to study the behaviour of large LM pretraining and finetuning especially pertaining to scaling laws . The checkpoints and code will be released at anonymous_link_withheld . 2 RELATED WORK . Transformers ( Vaswani et al. , 2017 ) have become ubiquitous in the modern deep learning stack and have seen widespread impact across not only language ( Devlin et al. , 2018 ; Raffel et al. , 2019 ; Brown et al. , 2020 ) but also computer vision ( Dosovitskiy et al. , 2020 ; Arnab et al. , 2021 ) , reinforcement learning ( Parisotto et al. , 2020 ) and computational biology ( Senior et al. , 2020 ) . To this end , discovering empirical scaling laws of these models is a research area that has garnered considerable interest ( Kaplan et al. , 2020 ; Henighan et al. , 2020 ; Hernandez et al. , 2021 ; Bahri et al. , 2021 ; Levine et al. , 2020 ) . Discovering empirical scaling laws that govern neural language model scaling has been a recent subject of keen interest ( Kaplan et al. , 2020 ; Hernandez et al. , 2021 ; Bahri et al. , 2021 ) . Many of these works present scaling laws across a variety of axis such as model size , compute and/or dataset size . It is worth to note that many of these works evaluate on autoregressive language modeling and use cross entropy loss to measure performance ( Kaplan et al. , 2020 ; Hernandez et al. , 2021 ) . There Table 1 : Table of model configurations . NL is the number of layers , dff is the size of the MLP , dmodel is the hidden size of the model . dkv is the size of each key-value vector . NH is the number of heads . P is the default model parallelism . Model NL dff dmodel dkv NH P # Params Tiny 4/4 1024 256 32 4 1 16M Mini 4/4 1536 384 32 8 1 31M Small 6/6 2048 512 32 8 1 60M Base 12/12 3072 768 64 12 1 220M Large 24/24 4096 1024 64 16 2 738M XL 24/24 16384 1024 128 32 8 3B XXL 24/24 65536 1024 128 128 32 11B XXXL 28/28 131072 1280 128 256 32 30B Table 2 : Description of different knobs used in the paper to define scaling operations . Scaling Op Description NL Num . layers EL Num enc . layers DL Num . dec. layers DM dmodel KV dKV NH Num . of heads FF dff SH Shared heads SKV Tied key-values are a multitude of interesting findings presented ( Kaplan et al. , 2020 ) whereby the authors show that performance ( loss ) scales as a power-law with model size and dataset size . However , one notable claim is that architectural details ( e.g. , network depth and width ) have minimal effects . Subsequently , Hernandez et al . ( 2021 ) builds upon the work of Kaplan et al . ( 2020 ) , evaluating scaling laws for ‘ transfer ’ . To this end , the authors study the effect of dataset scaling on unsupervised transfer learning and finetuning . That said , the experiments of Hernandez et al . ( 2021 ) are mainly targeted at dataset transfer between two different distributions ( language and code ) and make the same assumptions as Kaplan et al . ( 2020 ) about model scaling . In a similar vein , Henighan et al . ( 2020 ) studied empirical scaling laws for different domains for generative modeling in vision , video and multimodal setups . There have been increasing demand for training and scaling Transformers ( Shoeybi et al. , 2019 ; Raffel et al. , 2019 ; Fedus et al. , 2021 ; Conneau et al. , 2019 ; Lin et al. , 2021 ) . Despite the benefits from improved performance , there are financial considerations and environmental costs ( Schwartz et al. , 2019 ; Patterson et al. , 2021 ) to training these models . Given that every moment spent on hardware accelerators is a cost incurring activity , we believe that research in distilling practical scaling insights and recommendations to be highly crucial ( Li et al. , 2020 ; Kaplan et al. , 2020 ; Bello et al. , 2021 ) . Notably , the research problem of making transformers efficient have also been tackled from an extensive number of angles such as ( but not limited to ) distillation ( Hinton et al. , 2015 ) , compression ( Zafrir et al. , 2019 ) , parameter sharing ( Lan et al. , 2019 ; Tay et al. , 2019 ; Zhang et al. , 2021 ) , efficient attention ( Tay et al. , 2020c ; Kitaev et al. , 2020 ; Choromanski et al. , 2020 ; Tay et al. , 2020b ; Ainslie et al. , 2020 ; Jaegle et al. , 2021 ) , architecture search ( So et al. , 2019 ) , alternative non Transformer-based architectures ( Tolstikhin et al. , 2021 ; Tay et al. , 2021 ; 2020a ; Lee-Thorp et al. , 2021 ) . With so much extensive research into novel techniques to improving the efficiency of transformers , it is surprising that the standard configurations ( e.g. , base , large ) of transformers in BERT ( Devlin et al. , 2018 ) or T5 ( Raffel et al. , 2019 ) have not been rethought . 3 METHODS . This section describes our main experimental setup . Architecture We study a Transformer encoder-decoder architecture that uses relative attention based of the T5 architecture ( Raffel et al. , 2019 ) . The choice of adopting Seq2Seq architectures ( Sutskever et al. , 2014 ) is mainly due to their universality and ability to both subsume encoder ( Bert-like ) and decoder ( language ) models within an identical framework . Moreover , the universality of Seq2Seq architectures also allow us to fine-tune across a broad range of tasks . Our implementation and experiments are performed in Mesh Tensorflow1 ( Shazeer et al. , 2018 ) using the default T5 library2 . Model Configurations We first define eight Transformer sizes , i.e. , tiny , mini , small , base , large , XL , XXL and XXXL . The small , base , large , XL and XXL corresponds to the canonical T5 sizes that 1https : //github.com/tensorflow/mesh 2https : //github.com/google-research/text-to-text-transfer-transformer are released in Raffel et al . ( 2019 ) . We use three other sizes as starting points , e.g. , tiny and mini since there is a lack of representation of transformers at lower compute regions . Pretraining We pretrain on the Colossal Cleaned Common Crawl Corpus ( C4 ; Raffel et al. , 2019 ) . We pre-train encoder-decoder models using the span-based masked language modeling ( MLM ) objective ( Fedus et al. , 2018 ; Devlin et al. , 2018 ) . We pretrain all our models for 219 steps using 16 TPU-v3 chips . For larger models , we run our models with 64 TPU-v3 chips . We use 219 steps since majority of the experiments in ( Raffel et al. , 2019 ) were conducted in the same fashion . We would also like to emphasize that the official released T5 checkpoints were pretrained on 1T tokens ( 1 million steps with a batch size of 2048 ) . Given that this extended long pretraining setup is infeasible given the number of experiments we would have to run , we opt to follow the standard ablation setup in ( Raffel et al. , 2019 ) which pretrains on more manageable number of tokens . Downstream Tasks We consider a myriad of downstream tasks . In total , we consider 17 tasks . We finetune on a mixture of GLUE ( Wang et al. , 2018 ) , SuperGLUE ( Wang et al. , 2019 ) , SQuAD ( Rajpurkar et al. , 2016 ) for the key downstream experiment results and report aggregate GLUE/SuperGLUE scores . We believe that an aggregate of 17 tasks in natural language understanding that conmprises of both high-resource and low-resource tasks gives us a good overview of a model ’ s downstream performance . Finetuning is typically performed with 16 TPU-v3 chips . Notation for Scaling Operators For the remainder of the paper , we use a shortform code for each scaling operator applied on a standard transformer size . For example , NL32-SM refers to scaling small ( SM ) transformers to 32 layers ( NL32 ) . We use EL , DL to represent scaling encoder and decoder layers independently , KV to represent scaling each key-value size , DM to represent scaling dmodel . NH to represent modifying the number of heads and FF to represent scaling dFF . The initial/standard model sizes are tiny , mini , small , base , large , XL , XXL and XXXL . This is described in Table 2 . Convention With the exception of Figure 1 , all charts are plotted with FLOPS as the main compute metric . We use number of params for Figure 1 to align with Kaplan et al . ( 2020 ) . All of the downstream results are plot with SuperGLUE accuracy ( Wang et al. , 2019 ) as the Y-axis . Due to the lack of space , we report charts/plots of other metrics ( params of speed ) and other tasks ( GLUE or SQuAD ) in the supplementary material . All parameter counts also include the embedding parameters . We re-emphasize that it is critical to take into account multiple facets of efficiency and therefore report all three key metrics ( FLOPs , throughput/speed and parameter count ) in the supplementary material . Model and Data Parallelism By default , our models are trained across multiple workers via data parallelism . As per convention in the T5 paper , our larger models use the default model parallelism . Specifically , this is set to 2 for large models , 8 for XL models and 32 for XXL models .

Summary: This paper provides various ablation studies to show the effect of scaling various architecture parameters on the performance of a downstream task. The studies are performed on encoder-decoder transformer architectures that follow the T5 architecture as the backend. As a result of the ablation studies, the authors provide a heuristic on how to design architectures that lie on a better Pareto front compared to the previously proposed T5 architectures.