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 aims at providing insights from scaling Transformers for pre-training and finetuning. Based on extensive experiments involving pre-training and fine-tuning over different transformer configurations, the authors find that the model shape matters a lot besides model size when considering the downstream performance, and scaling strategies differ at different compute regions. With these new findings, the authors propose the DeepNarrow scaling strategy and verify it on additional experiments including tasks in different domains (language and vision).