Primer: Searching for Efficient Transformers for Language Modeling
arXiv:2109.08668
Abstract
Large Transformer models have been central to recent advances in natural language processing. The training and inference costs of these models, however, have grown rapidly and become prohibitively expensive. Here we aim to reduce the costs of Transformers by searching for a more efficient variant. Compared to previous approaches, our search is performed at a lower level, over the primitives that define a Transformer TensorFlow program. We identify an architecture, named Primer, that has a smaller training cost than the original Transformer and other variants for auto-regressive language modeling. Primer's improvements can be mostly attributed to two simple modifications: squaring ReLU activations and adding a depthwise convolution layer after each Q, K, and V projection in self-attention. Experiments show Primer's gains over Transformer increase as compute scale grows and follow a power law with respect to quality at optimal model sizes. We also verify empirically that Primer can be dropped into different codebases to significantly speed up training without additional tuning. For example, at a 500M parameter size, Primer improves the original T5 architecture on C4 auto-regressive language modeling, reducing the training cost by 4X. Furthermore, the reduced training cost means Primer needs much less compute to reach a target one-shot performance. For instance, in a 1.9B parameter configuration similar to GPT-3 XL, Primer uses 1/3 of the training compute to achieve the same one-shot performance as Transformer. We open source our models and several comparisons in T5 to help with reproducibility.
arXiv comment"Primer: Searching for Efficient Transformers for Language Modeling" NeurIPS camera ready. 34 pages
References in corpus (21)
- Sequence to Sequence Learning with Neural Networks
- TensorFlow: A system for large-scale machine learning
- Neural Architecture Search: A Survey
- Scaling Laws for Neural Language Models
- Searching for Activation Functions
- Semi-supervised Sequence Learning
- Conformer: Convolution-augmented Transformer for Speech Recognition
- Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
- Towards a Human-like Open-Domain Chatbot
- The Evolved Transformer
- Synthesizer: Rethinking Self-Attention in Transformer Models
- Lingvo: a Modular and Scalable Framework for Sequence-to-Sequence Modeling
- CvT: Introducing Convolutions to Vision Transformers
- Adafactor: Adaptive Learning Rates with Sublinear Memory Cost
- Making Pre-trained Language Models Better Few-shot Learners
- AutoML-Zero: Evolving Machine Learning Algorithms From Scratch
- Carbon Emissions and Large Neural Network Training
- Entailment as Few-Shot Learner
- Adaptive Input Representations for Neural Language Modeling
- GLU Variants Improve Transformer
- Can weight sharing outperform random architecture search? An investigation with TuNAS