Primer: Searching for Efficient Transformers for Language Modeling

By David R. So, Wojciech Ma'nke, Hanxiao Liu, Zihang Dai, Noam Shazeer, Quoc V. Le. 2021

224 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Model Architectures

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Model Architectures

Citations in the corpus

Scaling Instruction-Finetuned Language Models (2022, background): “compute includes better architectures (So et al., 2021), improved training objectives (Tay et al.”
Flamingo: a Visual Language Model for Few-Shot Learning (2022, methodology): “[104] David R So, Wojciech Mańke, Hanxiao Liu, Zihang Dai, Noam Shazeer, and Quoc V.”
Fast Inference from Transformers via Speculative Decoding (2022, background): “Numerous techniques try to make inference more efﬁcient for all tokens, e.g. distillation (Hinton et al., 2015), sparciﬁcation (Jaszczur et al., 2021), quantization (Hubara et al., 2016), and architecture modiﬁcation (So et al., 2021; Shazeer, 2019).”
The Efficiency Misnomer (2021, background): “Speed: Speed is one the most informative indicator for comparing the efﬁciency of different models (So et al., 2021; Dosovitskiy et al., 2020; Arnab et al., 2021; Tay et al., 2021b; Kim et al., 2021; He et al., 2021a; Narayanan et al., 2021; Lagunas et al., 2021; Liu et al., 2021; Tay et al.,…”
General-purpose, long-context autoregressive modeling with Perceiver AR (2022, methodology): “, 2020) and squared-ReLU nonlinearities (So et al., 2021).”
Big Little Transformer Decoder (2023, methodology): “Well-known approaches include efficient architecture design [25, 31, 34, 57, 64, 72], quantization [8, 29, 51, 71, 76, 78, 79], pruning [11, 13, 32, 33, 39, 46, 61], and neural architecture search [5, 53, 54, 63, 73, 77].”
NormFormer: Improved Transformer Pretraining with Extra Normalization (2021, methodology): “7B with Relu 2 activations (So et al., 2021), but use slightly fewer training steps (20% less) for compute efﬁciency.”
The Carbon Footprint of Machine Learning Training Will Plateau, Then Shrink (2022, methodology): “In 2021, a different NAS found the Primer model that again matched the quality scores but was 4.2x faster than the original Transformer [9].”
Zero-Shot Video Question Answering via Frozen Bidirectional Language Models (2022, background): “[92] David R So, Wojciech Mańke, Hanxiao Liu, Zihang Dai, Noam Shazeer, and Quoc V Le.”
Transformer Quality in Linear Time (2022, methodology): “To get the best performance, we use the rowmajor-axial variant of Combiner (Combiner-Axial) and the ReLU-kernel variant of Performer.”
RWKV: Reinventing RNNs for the Transformer Era (2023, methodology): “Further, the channel-mixing block is given by: where we adopt squared ReLU activation (So et al., 2021).”
OtterHD: A High-Resolution Multi-modality Model (2023, methodology): “Perssimon-8B [19] Persimmon-8B is a decoder-only transformer with modifications like squared ReLU activation [49], rotary positional encodings [50], and decoupled input \ output embeddings.”
ST-MoE: Designing Stable and Transferable Sparse Expert Models (2022, background): “Recently depthwise convolutions, which only accumulate 3-5 elements, have also been shown to greatly improve Transformer performance (So et al., 2021).”
SqueezeLLM: Dense-and-Sparse Quantization (2023, unknown): “This can be broadly categorized into efficient architecture design [32, 38, 42, 60, 67, 73], pruning [18, 20, 40, 41, 48, 52, 65, 79], neural architecture search [7, 58, 59, 66, 74, 78], and quantization [35, 55, 81, 82].”
Designing Effective Sparse Expert Models (2022, background): “Recently depthwise convolutions, which only accumulate 3-5 elements, have also been shown to greatly improve Transformer performance (So et al., 2021).”
MetaFormer Baselines for Vision (2022, background): “(6) This activation is called Squared ReLU [62], only costing 2 FLOPs for each input unit.”