Searching for Efficient Transformers for Language Modeling
Authors: David So
Wojciech Mańke
Hanxiao Liu
Zihang Dai
Noam Shazeer
Quoc V Le
Venue: NeurIPS 2021 Poster (2021)
Keywords: language modeling
lm
transformer
automl
architecture search
nas
primer

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.

Primer: Searching for Efficient Transformers for Language Modeling
David So, Wojciech Manke, Hanxiao Liu, Zihang Dai, Noam Shazeer, Quoe V. Le Google Research, Brain Team {davidso, wojciechm, hanxiaol, zihangd, noam, qvl)@google.co com
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 aim to reduce the costs of Transformers by searching
a more efficient variant. Compared to previous approaches, our search is performed at lower level, over the primitives that define a Transformer TensorFlow program. We identify an architecture, named Primer, that has 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 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 significantly speed training without additional tuning. For example, at 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 reach 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 achieve the same one-shot performance as Transformer. We open source our models and several comparisons in T5 to help with reproducibility.
Transformer increase compute scale
Introduction
Transformers [1] have been used extensively in many NLP advances over the past few years (c.g., [2, 3, 4, 5, 6, 7J). 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 search space of TensorFlow programs, we use evolution [11, 12, 13, 14, 15, 16, 17] to search for models that achieve as low a validation loss as possible given fixed amount of training compute. An advantage of using TensorFlow programs as the search space is that it easier to find simple lew-level improvements optimize Transformers. We focus on decoder-only auto-regressive language modeling (LM), because of its generality and success [18, 7, 19, 20, 21).
'https:// //githb con/ "e provide details of primitives search in TensorFlow, but the same approach can also be applied other deep learning libraries.
cam//google-ae google research google wwwparrc/tteeeeetteteeeeeeeeeeee
35th Conference on Neural Information Processing Systems (NcurIPS 2021), virtual.

---

The discovered model, named Primer (PRIMitives searched transformER), exhibits strong perfor- mance improvements over common Transformer variants on rugo-regressive language modeling. Our experiments show that Primer has the benefits of (1) achieving a target quality using 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 10 accelerator hours), across datasets (LMIB, C4, PG19 [22]), across hardware platforms (TPUv2, TPUv3, TPUv4 and V100), across multiple Transformer codebases using default configurations (Tensor2Tensor, Lingvo, and T5) and across mul- tiple model families (dense Transformers [1]. sparse mixture-of cettuure-offeepeee experts Switch Transformers [8], and Synthesizers [23]). We open source these comparisons to help with the reproducibility of our results."
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 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 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 E Rnxd that represent sequences of length 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.
DNA (Model)
Subprogram
Instruction
Susprograre am (Marn))
Instudter
Operation: CONY 1X1
Sutprogram Bark
Punction Asguments Input nd
Subprogrom
Constant: 9.787
Susprogrum
Inaduction N
Rid0on literssoon: 612
Susprogrom N
Instruction functions come from primiives vocab or sutprogram bank.
TF Primitives Vocab
Generated TF Code
NLC: if (rrt
of Joyers.darse
(nlten612)
nath
Figure 1: Overview of DNAs that define decoder model program (i.e., auto-regressive language model). Fach DNA has a collection of subprograms, where SUBPROGRAM AM 0 is the MAIN() function entry point. Each subprogram is comprised of instructions, which are converted to lines of TensorFlow code. Instruction operations map to either basic TensorFlow library functions from the primitives vocabulary or one of the parent DNA's subprograms. The operation's 's arguments are filled using the parent instruction's 's argument set, which contains values for all potential operation arguments; arguments that are not used by particular operation are simply ignored.

---

Figure shows how programs are constructed in our search space. Each program is built from an evo- lutionary search DNA, which an indexed collection of subprograms. SUBPROGRAM 0 the MAIN( function that is the execution entry point, and the other subprograms are part of the DNA's subprogram bank. Each subprogram an indexed array of instructions with no length constraints. An instruction is an operation with set of input arguments. The operation denotes the function that the instruction executes. Each operation maps to either TF function from the primitives vocabulary or another sub- program in the DNA subprogram bank. The primitives vocabulary is comprised of simple primitive TF functions, such as ADD, LOG, and MATMUL (see Appendix A. I. for details). It is worth empha- sizing that high-level building blocks such 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 only call
additional programs that can be executed functions by instructions. Each subprogram subprograms with a higher index in the subprogram bank, which removes the possibility of cycles.
Each instruction's is argument set contains 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 the first tensor input. The index 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.
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.
Input 2:
Constant: A real valued constant. An example of an operation that uses this is MAX; maximmui, 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, 1, the dense projection used by the Transformer's '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 corresponding line of TF code, one at 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 the subprogram output. We do not use TF Eager and 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 for an illustration of how subprograms converted to TF graphs and see Appendix A.2 for more details on how TF graphs are constructed, including how we handle causal masking.
useful property of the constructed
Subprogram
(10) IN
(h1)IN 1
(h2) CONV 1X1
Snpat 21 Sghare
Covscentt Igsore
def subprogras__ 1 (b0, h1]:
1o 1: In boxt Ignora Speore
19 tf. layera dhnsaa 512]
(1)) CONY
(h3)
(13)
hg tf maach sin(h1)
Coverann: Iazara
h4 tr www.mmhh.uitttee.. ruitspiyyyyp2 (22 h3)
hs tr aach (hê, 82)
1ast state output
MUL
1roat_11 to Bie: Tptore Satore
output
return output
ADD
(h4) MUL
(1S) ADD
Suavl ld dgsore
QUT
Coratant:
Figure 2: Example of program converted into its corresponding TensorFlow graph. Nodes that do not contribute the program output are not executed thanks to TensorFlow's deferred execution design.
3

---

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 fixed training budget (24 TPUv2 hours) and define its fitness as its perplexity on the One Billion Words Benchmark (LMIB) [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
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 training hours as proxy for 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 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 search for 255 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 [3I], which are often heavily bi- ased such that random search performs well (see analysis by [32, 33, 34J). As our search space does not have this bias, 78% of random programs in our space with length equal to Transformer program cannot 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 popu- lation with copies of the Transformer [13] (input embedding size dmodel = 512, feed forward upwards projection size di = 2048, and number of layers L 6) (Figure 3). To apply this initializa- tion 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).
-10866
Max Search Fitness
-100.70
Random
Concept
-10974
500
1000
Num Individuals
Figure 3: Small scale searches (10 hours on 10 TPUv2 chips) comparing conceptual initializa- tion to random initialization. Our search space is open-ended enough that it is infeasible to search without strong initializa- tion.
3 Primer
Primer: We name the discovered model Primer, which stands for PRIMitives searched transformER (See Appendix Figure 23 for the full program). Primer shows significant improvement when retrained on the search task, requiring less than half the compute of Transformer reach the same quality (Figure 6). In Section 4, we additionally show that Primer makes equally large gains when transferred to other codebases, training regimes, datasets, and downstream one-shot tasks.
Primer-FZ: A core motivation of this work is to develop simple techniques that can be easily adopted by language modeling practitioners. To accomplish this, perform ablation tests across two

---

codebases (T5 [5] and Tensor2Tensor [25]) and determine which Primer modifications are generally useful (Appendix Figure 26). The two that produce the most robust improvements are squaring feed forward ReLUs and adding depthwise convolution to attention multi-head projections (Figure 4). We refer to a Transformer with just these two easy modifications as Primer-FZ; this our recommended starting point for language modeling practitioners interested in using Primer. We now explain these modifications and then measure their empirical effectiveness.
MDHA Projection Pseudo-code
Mutt-DConv-Head Attention (MDHA)
Squared ReLU in Feed Forward Block
Use to create each K, Q, and V head of size 'hs',
Output
Output
def aohha_proecctoont: ha):
Create head.
Downwards Proj
proj (x,
Multi Head Self-Altention
head_size=he,
axis="channee1")
Square
Spalial D-Canv
Spalial D-Conv 3xt
Spelial D-Conw 3xt
# Apply D-Conv to head.
RaLU
x d_ (x
width-3, head size=hs,
Q Head Projection
K Head Projection
V Head Projection
Upwards Proj
axis- axis="spatial"
mask- causal)
Input
Input
X Num Heads
return
Figure 4: The two main modifications that give Primer most of its gains: depthwise convolution added to attention multi-head projections and squared ReLU activations. These modifications are easy to implement and transfer well across codebases. We call the model with just these two modifications Primer-EZ. Blue indicates portions of the original Transformer and red signifies one of our proposed modifications.
The most effective modification is the improvement from a ReLU activation to squared ReLU activation in the Transformer's s feed forward block. Rectified polynomials of varying degrees have been studied in the context of neural network activation functions [36], but are not commonly used; to the best of our knowledge, this is the first time such rectified polynomial activations are demonstrated to be useful in Transformers. Interestingly, the effectiveness of higher order polynomials [37] can also be observed in other effective Transformer nonlinearities, such as GLU [38] variants like ReGLU [39] (y = U V z, 0) where @ is an weementtwwiee product) and point-wise activations like approximate GELU [40] (y- = 0.5x(11 + tanh( /2/7(x
Squared ReLU:
(x + ..044715=²)))) However, squared ReLU has drastically different asymptotics as x + Do compared to the most commonly used activation functions: ReLU, GELU and Swish (Figure 5 left side). Squared ReLU does have significant overlap with ReGLU and in fact is equivalent when ReGLU's U and V weight matrices are the same and squared ReLU is immediately preceded by a linear transformation with weight matrix U. This leads us believe that squared ReLUs capture the benefits of these GLU variants, while being simpler, without additional parameters, and delivering better quality (Figure 5 right side).
T5 C4 LM 110M Params -19.5
ReGLU
GELU
Squared RelU
Swish Squared ReLU
Negative PPLX
Y
-2-1 0
Ux
22 ReLu
GELU
swish
ReGLU
SWiGLL
Sar Retu
-1
0
1
Figure 5: Left: Squared ReLU has starkly different asymptotics compared to other common activation functions. Center: Squared ReLU has significant overlap with GLU variants [39] that use activations with ReLU-like asymptotics, such as ReGLU and SwiGLU. Our experiments indicate that squared ReLU is better than these GLU variants in Transformer language models. Right: Comparison of different nonlinearities in Transformers trained on C4 auto nugo-regressive LM for 525K steps.
5

---

Multi-DConv-FEn Attention (MDHA): Another effective modification is adding 3x1 depthwise convolutions after each of the multi-head projections for query Q. key K and value V self-attention. These depthwise convolutions are performed over the spatial dimension of each dense projection's s output. Interestingly, this ordering of pointwise followed by depthwise convolution is the reverse of typical separable convolution, which we find to be less effective in Appendix A.6. We also find that wider depthwise convolution and standard convolution not only do not improve performance, but in several cases hurt it. Although depthwise convolutions have been used for Transformers before [41, 42], using them after each dense head projection has not been done to the best of our knowledge. MDHA is similar Convolutional Attention [43], which uses separable convolution instead of depthwise convolution and does not apply convolution operations per attention head as we do.
The other Primer modifications are less effective. Graphs for each modifica- tion can be found in Appendix A.5 and an ablation study can be found in Appendix A.7. We briefly describe the modifications and their usefulnesses here:
Other Modifications:
Shared Q and K Depthwise Representation: Primer shares weight matrices for Q and K. K is created using the previously described MDHA projection and Q = KW for learnable weight matrix W E Rdddd We find that this generally hurts performance.
Pre and Post Normalization: The standard practice for Transformers has become putting normalization before both the self-attention and feed forward transformations [44, 45]. Primer normalization before self-attention but applies the second normalization after the feed forward transformation. We find this is helpful in some but not all cases.
Custom Normalization: Primer uses modified version of layer normalization [46] that uses z(x μ4) instead of (x 2)2 but we find this is not always effective.
12X Bottleneck Projection: The discovered model uses a smaller dmmodet size of 384 (compared to the baseline's 512) and a larger d size of 4608 (compared to the baseline's 2048). We find this larger projection improves results dramatically at smaller sizes (~35MM parameters), but is less effective for larger models, as has been previously noted [9]. For this reason we do not include this modification when referencing Primer or Primer-EZ.
Post-Softmax Spatial Gating: The discovered model has set of per-channel learnable scalars after the attention softmax, which improves perplexity for fixed length sequences. However, these scalars cannot be applied to variable sequence lengths and so we do not include this modification in Primer for our experiments.
Extraneous Modifications: There are handful of additional modifications that produce no meaningful difference in the discovered architecture. For example, hidden states being multiplied by 12. Verifying that these modifications neither help nor hurt quality, we exclude them from discussion in the main text and do not include them when experimenting with Primer. These extraneous modifications can still be found in Appendix A.5.
4 Results
In our experiments, compare Primer against three Transformer variants:
Vanilla Transformer: The original Transformer [1] with ReLU activations and layer normaliza- tion [46] outside of the residual path.
Transforme++GELU: A commonly used variant of the vanilla Transformer that uses a GELU [40] approximation activation function [2, 7].
Transformer++: : A Transformer with the following enhancements: RMS normalization [47], Swish activation [48] and GLU multiplicative branch [38] in the feed forward inverted bottleneck (SwiGLU) [39]. These modifications were benchmarked and shown to be effective in T5 [49].
We conduct our comparisons across three different codebases: Tensor2Tensor (T2T) [25]. T5 [5]. and Lingvo [50]. Tensor2Tensor the codebase we use for searching and so a majority of our side-by-sides are done in T5 and Lingvo to prove transferabilityy In all cases, we use the default Transformer hyperparameters for each codebase, with regularization disabled. See Appendix A.8 for more hyperparameter details.

---

In the following sections, we will present our results in four main experiments on auto-regressive language modeling. First, we will show that Primer outperforms the baseline models on the search task. Next, we will show that the relationship between Primer's s compute savings over Transformers and model quality follow power law at optimal model sizes. These savings also transfer across datasets and codebases. Then, we will study Primer's gains in an established training regime and show that enables 4.2X compute savings at a 500M parameter size using full compute T5 training. Finally, we will demonstrate that these gains transfer to the pretraining and one-shot downstream task setup established by GPT-3 [7].
4.1 Search Task Comparison
We first analyze Primer's performance on the search task: LMIB language modeling with sequence length 64, ~35M model parameters, batches of 4096 tokens and 24 hours of training. We compare against the baseline models in both Tensor2Tensor (T2T) [25] and T5 [5] and on TPUv2s and V100 GPUs. We grade each model's performance according how much faster it reaches the vanilla Transformer's s final quality, which we will refer to as its speedup factor. Figure 6 shows that Primer provides a speedup factor of 1.7X or more over Transformer in all cases. Figure 6 also shows that both Primer and Primer-EZ generalize to other hardware platforms and codebases.
T2T LM1B TPUv2
T2T LM1B V100 GPU
T5 LM1B TPUV2
2.4
222
2.0
1.6
2.0
18
1.6
Spoedup
1.6
1.4
1.2
1.2
Transt
Tranot aad
Trans ++
Primer
Primer-
Transt
Trand kard
Trans ++
Primer
Primes 12
0.8
Trant
Tranef+ -BeL
Transr me
Ayinet
Primer EZ
Figure 6: Comparison on the LMIB search task using 35M parameter models. "Speedup Factor" refers to the fraction of compute each model needs to reach quality parity with the vanilla Transformer trained for 24 hours. Primer and Primer-EZ both achieve over 1.7X speedup in all cases. Note that results transfer across codebases and different hardware.
2.00
21.7
Validation Loss
Validation Loss
2.85
21.71
Tranff+ GZLL
Tranc
2170
2
24.0
23.5
27. 7amm
2
25.
25.0
27.
TPUV2 Hours
TPUv2 Hours
Figure 7: Left: When sweeping over optimal model sizes, the relationship between Transformer language model quality and training compute roughly obeys a power law [9]. Right: Comparison of these power law lines for varying Transformer modifications, fit with smoothed MSE loss. That the model lines are parallel to one another implies that compute savings by using superior modeling also scales as a power law with quality. Spacings between vertical dotted lines represent 2X differences in compute. Note that the and y-axes in both plots are log.
Next we study the scaling laws of Primer. Here we compare Primer to our baselines over many sizes by training each model using every permutation of L E (6, 9, 12) layers, dynoeel E (384, 512, 1024) initial embedding size, and p E {4, 8, 12) feed forward upwards projection ratio, creating parameter range from 23M to 385M. The results, shown in Figure 7, corroborate previous claims that, at optimal parameters sizes, the relationship between compute and language model quality roughly follows a power law [9]. That is, the relationship between validation loss, 1, and training compute, c, follows the relationship / = ac k for empirical constants and k. This is represented as line in double log space (Figure 7): log l -k log e + log a. However, these lines are not the same for each architecture. The lines are roughly parallel but shifted up and down. In Appendix A.9 we show that,
7

---

given a vertical spacing of log parallel lines such as these indicate compute savings, s, for superior modeling also follow a power law of the form a (1 - 1/b)* -k, The intuition behind this is that b is constant compute reduction factor for all and thus power law investment of training compute with relation to results in a power law savings with relation to I as well (see Appendix A.9).
Primer also has the capacity to improve inference, despite our search focusing on training compute. Figure shows Pareto front compar- ison of quality vs. inference, when using feed forward pass timing
Negative perplexity
as a proxy for inference. We use forward pass timing as proxy for
inference because there are multiple ways to decode language model,
each with varying compute costs. A more in depth study could be conducted analyzing Primer's 's inference performance across different decoding methods, serving platforms, datasets, etc., but that is beyond the scope of this work.
Primer
Transt.
Transt.+GELU
Transf. ++
1.25
2330
Inference Time (s)
4.2 Primer
Figure 8: Pareto optimal inference comparison on LMIB. Primer demonstrates improved inference at ma- jority of target qualities. We observe these models have a 0.97 correlation between their train step and inference times.
Transferability to Other Codebases, Datasets, and Model Types
We now study Primer's ability to transfer to larger datasets, PG19 and C4, in another codebase, T5. We additionally scale up to higher compute regime that has been used as proxy for large scale training by previous studies [49, 5]: the batches are increased to 65K tokens, the sequence lengths are longer 512, each decoder is 1 10M parameters (dmodel = 768, d = 3072, L = 12) and each model is trained to ~525K steps on 4 TPUv3 chips. We also continue training each model to IM steps to study the effect of larger compute budgets on Primer savings. The results, shown in Figure 9, indicate that the Primer models are as strong in larger data, higher compute regimes, as they are in the smaller LMIB regime. Compared to the vanilla baseline, Primer and Primer-EZ are at least 1.8X more efficient at the end of training on both PG19 and C4.
C4
PG19
C4
C4
2.4
Speedup Factor
2.2
127 Hours
127 Hours
142 Hours
124 Hours
2.0
242 Hours
242 Hours
270 Hours
236 Hours
1.8
1.6
1.4
1.2
1.0
0.8
Transf Transf- transf++GEU +GELL Transf.+ + Primer
Primer-Er EZ
Transf Transf +GELL Transt + Primer Primer-Fe EZ
Switch Transf. Switch Primer
syn. Syn +Sar ReLu
Figure 9: Comparison transferring Primer to larger datasets (C4 and PG19) and different model families (Switch Transformer and Synthesizer) in a different codebase (T5) with an order of magnitude more compute than the search task. Compared at least 1.8X more efficient at the end of training on both PG19 and C4. In all cases, the fraction of Primer compute savings increases as more compute is invested. Primer-EZ. modifications also improve Switch Transformer (550M params) and Synthesizer (145M params), showing that it is compatible with other efficient methods. Compute budgets are selected according to how long takes each baseline to train for 525K and IM steps. See Appendix 10 for exact numbers.
the vanilla baseline, Primer and Primer-FZ.
Figure 9 also shows that the Primer modifications are compatible with other efficient model families, such as large sparse mixture-of -f-eeeerts like Switch Transformer [8] and efficient Transformer approx- imations like Synthesizer [23]. For these experiments, we use the T5 implementations provided by Narang et al. [49]. The Primer-EZ techniques of added depthwise convolutions and squared ReLUs reduce Switch Transformer's 's compute cost by a factor of 1.5X; this translates to 0.6 perplexity im- provement when controlling for compute (see Appendix A.1.). Adding squared ReI. to Synthesizer reduces training costs by a factor of 2.0X and improves perplexity by 0.7 when fully trained.
8

---

4.3 Large Scale TS Auto-Regressive Language Model Training
In large scale compute configurations, the Primer compute savings ratios are even higher. To demon- strate Primer's s savings in an established high compute training setup, we scale up to the full T5 compute regime, copying Raffel al. exactly [5]. This is the same as the C4 configuration in the previous section, but uses batches of ~M tokens, 64 TPUv3 chips and 537M parameters (dmodel 1024, y 8192, = 24). Primer 4.2X more compute efficient than the original T5 model and 2X more efficient than strengthened Transformer++ baseline (Table 1).
Model
Steps
TPUv3 Hours
PPLX
Original T5 T5++
1M
15.7K
13.25
251K
4.6K
13.25
Primer
207K
3.8K
13.25
T5++
1M
16.5K
12.69
Primer
480K
8.3K
12.69
Primer
IM
17.3K
12.35
The reason why savings are even better here is be- cause, at fixed sizes, more compute invested yields higher Primer compute savings. Figure 10 shows how the fraction of compute Primer needs to achieve parity with the original TS architecture shrinks as the models are trained for longer; this due to the asymptotic nature of both the control and variable per- plexity training curves. This differs from the power law savings described in Section A.6. There, the optimal number of parameters for each compute budget, and so the compute saving factor, remains constant. For fixed model sizes, the compute saving factor grows invested, meaning that compute savings can exceed the power law estimation. Note, this means that comparisons such as the ones given here can be "gamed" by investing more compute than is necessary for baseline models. is for this reason that we use an exact replica Raffel et al.'s [5] training regime: demonstrate Primer's savings in already published training configuration.
Comparison in compute usage to reach target qualities on C4 LM 537M pa- rameters using the full T5 compute scale. Tar- get qualities are selected according to the final performances of the baseline models. Primer achieves the same quality as the original T5 architecture using 4.2X less compute.
Table I:
compute is
LM 500M Params
Primer Savings vs. T5
-12.0
4.30
Negative PPLX
Saving Factor (F)
-13.5
T5
Primer
-15.0g
8000
16000
0.00
8000
16000
TPUv3 Hours
TPUV3 Hours
Figure 10: Compute savings of Primer vs. the original T5 architecture on C4 LM over time. The more compute invested in training, the higher the savings due to the asymptotic nature of both perplexity curves. Primer achieves the same performance as the original T5 architecture with 4.2X less compute.
4.4 Primer Transferability to Downstream One-Shot Tasks
In our final comparison, we demonstrate Primer's improvements also hold in the pretraining and one-shot downstream task transfer regime. Recent trends in language modeling have moved towards training large models on large datasets, which is referred to as "pretraining." These models are then transferred unseen datasets and tasks, and, without much any additional training, demonstrate the capacity to perform well on those "downstream" tasks [2, 51]. In the decoder-only auto-regressive language modeling configuration we study here, the most impressive results have been achieved by GPT-3 [7], which showed that large language models can exhibit strong performance on unseen tasks given only example referred to "one-shot" learning. In this section, demonstrate that Primer's training compute savings stretch beyond reaching a target pretraining perplexity and indeed transfer downstream one-shot task performance.
To do this, we replicate the GPT-3 pretraining and one-shot evaluation setup. This replication is not exactly the same the one used for GPT-3 because GPT-3 was not open sourced. Thus, these

---

experiments are not meant to compare directly to GPT-3, as there are configuration differences. Instead, these experiments are used as a controlled comparison of the Transformer and Primer architectures. We conduct these experiments in the Lingvo codebase using proprietary pretraining dataset. The downstream tasks are configured in the same one-shot way descrihed by Brown al. [7], with single prefix examples fed into each model with each task's inputs. We compare (1) baseline 1.9B parameter Transformer (dmodet = 2048, dyf = 12288, L = 24) with GELU activations, meant to approximate the GPT-3 XL architecture, and (2) a full Primer without shared QK representations, which only hurt performance according to Appendix A.7. Each model is trained using batches of 2M tokens using 512 TPUv4 chips for 140 hours (~71.8K total accelerator hours ~ IM train steps). We once again use the T5 training hyperparemeters without any additional tuning.
One Shot Pretraining
-15
One-Shot
One-Shot
Pretraining
QA Tasks
Multi-Choice Tasks
Negative PPLX
-13
Average Score
37
Average Score
57
-16
Neg PPLX
Transf. 24K Hrs
Transformer Primer
Primer 24K Hrs
Primer 72K His
-19
-16
29
52
26000
72000
TPUv4 Hours
Figure 11: Comparison between TTansformer+GELU and Primer at 1.9B parameters and varying training compute budgets on downstream one-shot tasks, similar to GPT-3. Primer achieves slightly better performance than Transformer when given 3X less pretraining compute and substantially better performance when given the same pretraining compute. Here we stop at 72K TPUv4 hours to roughly match the quality of GPT-3 XL, but the compute savings of Primer would be larger if we let the two models run longer (see Figure 10), Note, this is crude comparison that uses averaged scores from the 27 one-shot tasks we evaluate. See Appendix A. (Table 6 and Figure 27) for exact task scores.
Figure 11 shows that Primer achieves the same pretraining perplexity and one-shot downstream performance as Transformer+GEEUU while using 3X less compute. Table 6 in the Appendix gives the exact performance numbers for each of the 27 evaluated downstream tasks. Primer, despite using 3X less compute, outperforms Transfomer+GELU on tasks, does worse on task, and performs equivalently on the remaining 21 tasks. The same table shows that when given equivalent compute, Primer outperforms Transforme++GELU on 15 tasks, does worse on 2 tasks, and performs equivalently the remaining 10 tasks. This result shows that not only can Primer improve language modeling perplexity, but the improvements also transfer to downstream NLP tasks.
Conclusion
There are limitations this study. First, although our comparisons are at substantial sizes, they are still orders of magnitude smaller than sate-of--thearr models such as the full-scale GPT-3 [7]. Second, we focus primarily on decoder-only models, while encoder-based sequence models are still widely used [1, 2, 3, 4, 6, 52]. In Appendix A.12, we perform encoder-decoder masked language modeling comparisons in T5, but do not study the results in significant depth. The main finding there is that, although Primer modifications improve upon vanilla Transformer, they perform only as well as Transformer++. This result suggests that architectural modifications that work well for decoder-only auto-regressive language models may not necessarily be as effective for encoder-bate based masked language models. Developing better encoders a topic of our future research.
Limitations:
We recommend the adoption of Primer and Primer-EZ for auto-regressive language modeling because their strong performance, simplicity, and robustness to hyperparameter and codebase changes. To prove their potential, simply dropped them into established codebases and, without any changes, showed that they can give significant performance boosts. Furthermore, in practice, additional tuning could further improve their performance.
Recommendations and Future Directions:
We also hope our work encourages more research into the development of efficient Transformers. For example, important finding of this study is that small changes to activation functions can yield more efficient training. In the effort to reduce the cost of Transformers, more investment the development of such simple changes could be a promising area for future exploration.
10

---

Acknowledgements
We thank Zhen Xu for his help with infrastructure. We also thank Gabriel Bender, Hallie Cramer, Andrew Dai, Nan Du, Yanping Huang. Daphne Ippolito, Norm Jouppi, Lluis-Miquel Munguia, Sharan Narang, Ruoming Pang, David Patterson, Yanqi Zhou, and the Google Brain Team for their help and feedback.
References
[1] Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, 2017.
Ming-Wei [2] Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In NAACL, 2018.
[3] Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Ruslan Salakhutdinov, and Quoc
Le. Xlnet: Generalized autoregressive pretraining for language understanding. In Advances in Neural Information Processing Systems, 2019.
[4] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: 11692, 2019.
[5] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research, 21(140):1-67, 2020.
[6] Daniel Adiwardana, Than-TThag&aa Luong, David R. So, J. Hall, Noah Fiedel, R. Thoppilan,
Z. Yang. Apoory Kulshreshtha, Gaurav Nemade, Yifeng Lu, and Quoc V. Le. Towards human-like open-domain chatbot. arXiv preprint arXiv:2001.09977 2020.
[7] Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. In Advances in Neural Information Processing Systems, 2020.
[8] William Fedus, Barret Zoph, and Noam M. Shazeer. Switch transformers: Sealing to trillion parameter models with simple and efficient sparsity. arXiv preprint vrXiv:2101.0366 2021.
[9]
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.088 083611 2020.
[10] Martín Abadi, P. Barham, Chen, Z. Chen, Andy Davis, J. Dean, M. Devin, Sanjay Ghemawat, Geoffrey Irving, M. Isard, M. Kudlur, Josh Levenberg, Rajat Monga, Sherry Moore, D. Murray, Benoit Steiner, P. Tucker, Vijay Vasudevan, Pete Warden, Martin Wicke, Y. Yu, and Xiaoqiang Zhang. Tensorflow: A system for large-scale machine learning. In OSDI, 2016.
[11] Esteban Real, Sherry Moore, Andrew Selle, Saurabh Saxena, Yutaka Leon Suematsu, Quoe V.
Le, and Alex Kurakin. Large-scale evolution of image classifiers. In ICML, 2017.
[12] Hanxiao Liu, Karen Simonyan, Oriol Vinyals, Chrisantha Fernando, and Koray Kavukcuoglu. Hierarchical representations for efficient architecture search. In ICLR, 2018.
[13] David R. So, Chen Liang, and Quoc V. Le. The evolved transformer. In ICML, 2019.
[14] Hanxiao Liu, Andrew Brock, Karen Simonyan, and Quoc V Le. Evolving normalization- activation layers. In NewrIPS, 2020.
[15] Xin Yao. Evolving artificial neural networks. Proceedings of the IEEE, 870)) 1999.
[16] Jurgen Schmidhuber. Evolutionary principles in self-referential learning. (on learning how to learn: The meta-meta-.. hook.). Diploma thesis, Technische Universitat Munchen, Germany, 1987.
[17] Kenneth Stanley, Jeff Clune, Joel Lehman, and Risto Miikkulainen. Designing neural networks through neuroevolution. Nature Machine Intelligence, 1:24-35, 2019.
[18] Alec Radford, Jeffrey Wu, R. Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models unsupervised multitask learners. In Technical report, OpenAl, 2019.

---

[19] Timo Schick and Hinrich Schütze. It's not just size that matters: Small language models are also few-shot learners. arXiv preprint rrXiv:2009.07118. 2021.
[20] Sinong Wang, Han Fang, Madian Khabsa, Hanzi Mao, and Hao Ma. Entailment as few-shot learner. arXiv preprint arXiv:2104. 14690, 2021.
[21]
Tianyu Gao, Adam Fisch, and Danqi Chen. Making pre-trained language models better few-shot learners. arXiv preprint arXiv:2012.15723, 2020.
[22] Jack W. Rae, Anna Potapenko, Siddhant M. Jayakumar, and T. Lillicrap. Compressive trans- formers for long-range sequence modelling. ArXiv, abs/1991.05500 2020.
[23] Yi Tay, Dara Bahri, Donald Metzler, Da-Cheng Juan, Zhe Zhao, and Che Zheng. Synthesizer: Rethinking self-attention transformer models, arXiv preprint arXiv:2005.007: 2020.
[24] Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, Phillipp Koehn, and Tony Robinson. One billion word benchmark for measuring progress in statistical language modeling. In Interspeech, 2014.
[25] Ashish Vaswani, Samy Bengio, Eugene Brevdo, Francois Chollet, Aidan N. Gomez, Stephan Gouws, Llion Jones, Lukasz Kaiser, Nal Kalchbrenner, Niki Parmar, Ryan Sepassi, Noam Shazeer, and Jakob Uszkoreit. Tensor2tensor for neural machine translation. arXiv preprint arXiv:1803.074: 07441, 2018.
[26] Mingxing Tan and Quoc Le. Efficientnet: Rethinking model scaling for convolutional neural networks. In /CML, 2019.
[27] Mingxing Tan, Bo Chen, Ruoming Pang, Vijay Vasudevan, and Quoc V. Le. Mnasnet: Platform- aware neural architecture search for mobile. 2019 IEEE/CVF Conference Computer Vision and Puttern Recognition (CVPR), pages 2815-2823, 2019.
[28] Han Cai, Ligeng Zhu, and Song Han. Proxylessnas: Direct neural architecture search on target task and hardware. arXiv preprint: 1 082200333, 2019.
[29] Thomas Elsken, Jan Hendrik Metzen, and Frank Hutter. Efficient multi-objective neural architecture search via lamarckian evolution. In ICLR, 2019,
[30] Esteban Real, Alok Aggarwal, Yanping Huang, and Quoe V. Le. Regularized evolution for image classifier architecture search. In AAAI, 2019.
[31] Thomas Elsken, Jan Hendrik Metzen, and Frank Hutter. Neural architecture search: A survey. Journal of Machine Learning Research, 20(55):1-21, 2019.
[32] Liam Li and Ameet S. Talwalkar. Random search and reproducibility for neural architecture search. In UAI, 2019.
[33] Kaicheng Yu, Christian Sciuto, Martin Jaggi, Claudiu Musat, and Mathieu Salzmann. Evaluating the search phase of neural architecture search. In ICLR, 2020.
[34] Gabriel Bender, Hanxiao Liu, Bo Chen, Grace Chu, Shuyang Cheng, Pieter-Jan Kindermans, and Quoe V. Le. Can weight sharing outperform random architecture search? An investigation with tunas. In CVPR, 2020.
[35] Esteban Real, Chen Liang, David R. So, and Quoc V. Le. Automl-zero: Evolving machine learning algorithms from scratch. In ICML, 2020.
[36] Dmitry Krotov and John J. Hopfield. Dense associative memory for pattern recognition. In Advances in Neural Information Processing Systems, 2016.
[37] Siddhant M. Jayakumar, Jacob Menick, Wojeiech M. Czarnecki, Jonathan Schwarz, Jack W. Rac, Simon Osindero, Y. Teh, Tim Harley, and Razvan Pascanu. Multiplicative interactions and where to find them. In ICLR, 2020.
[38] Yann N Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In ICML, 2017.
[39] Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002..5202,, 2020.
[40] Dan Hendrycks and Kevin Gimpel. Bridging nonlinearities and stochastic regularizers with gaussian error linear units. arXiv preprint arXiv: 1o06iv:600..00 1606.0844444 08415, 2016.
[41] Adams Wei Yu, David Dohan, Minh-Thane Luong, Rui Zhao, Kai Chen, Mohammad Norouzi, and Quoe V. Le. QANet: Combining local convolution with global self-attention for reading comprehension. In ICLR, 2018
12

---

[42] Anmol Gulati, James Qin, Chung-Cheng Chiu, Niki Parmar, Yu Zhang, Jiahui Yu, Wei Han, Shibo Wang, Zhengdong Zhang. Yonghui Wu, and Ruoming Pang. Conformer: Convolution- augmented transformer for speech recognition. In Interspeech, 2020.
[43] Haiping Wu, Bin Xiao, Noel Codella, Mengehen Liu, Xiyang Dai, Lu Yuan, and Lei Zhang. CvT: Introducing convolutions to vision transformers. arXiv preprini arXiv:2103, 15808, 2021.
[44]
Alexei Baevski and Michael Auli, Adaptive input representations for neural language modeling. In arXiv preprini arXiv: /809:/889.000 10099.00853 10853, 2019.
[45] Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, S. Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, L. Wang, and T. Liu. On layer normalization in the transformer architecture. arXiv 2020.
preprint arXiv:2002.04774
[46] Jimmy Ba, Jamie Kiros, and Geoffrey E. Hinton. Layer normalization. arXiv preprint 1607:0660..644 2016.
[47] Biao Zhang and Rico Sennrich. Root mean square layer normalization. In NewrlPS, 2019.
[48] Prajit Ramachandran, Barret Zoph, and Quoc V. Le. Searching for activation functions. arXiv preprint arXiv: 1710.05944, 2018.
[49] Sharan Narang. Hyung Won Chung, Yi Tay, William Fedus, Thibault Févry, Michael Matena, Karishma Malkan, Noah Fiedel, Noam Shazeer, Zhenzhong Lan, Yanqi Zhou, Wei Li, Nan Ding, Jake Marcus, Adam Roberts, and Colin Raffel. Do transformer modifications transfer across implementations and applications? arXiv preprint arXiv:2102.1 2021.
[50] Jonathan Shen, P. Nguyen, Yonghui Wu, Z. Chen, M. Chen, Ye Jia, Anjuli Kannan, T. Sainath, Yuan Cao. C. Chiu, Yanzhang He, Chorowski, Smit Hinsu, S. Laurenzo, James Qin, Orhan Firat, Wolfgang Macherey, Suyog Gupta, Ankur Bapna, Shuyuan Zhang, Ruoming Pang, Ron Weiss, Rohit Prabhavalkar, Qiao Liang, Benoit Jacob, Bowen Liang, HyoukJoong Lee, Ciprian Chelba, Sébastien Jean, Bo Li, M. Johnson, Rohan Anil, Rajat Tibrewal, Xiaobing Liu, Akiko Eriguchi, Navdeep Jaitly, Naveen Ari, Colin Cherry, Parisa Haghani, Otavio Good, Youlong Cheng. R. Álvarez, Isaac Caswell, Wei-Ning Hsu, Zongheng Yang. Kuan-Chich Wang. Ekaterina Gonina, Katrin Tomanek, Ben Vanik, Zelin Wu, Llion Jones, M. Schuster, Y. Huang, Dehao Chen, Kazuki Irie, George Foster, John Richardson, Uri Alon, and E. al. Lingvo: a modular and scalable framework for sequence-fo-sequence modeling. ArXiv, abs/1902.08295. 2019.
[51] Andrew M. Dai and Quoc V. Le. Semi-supervised sequence learning. In Advances in Neural Information Processing Systems, 2015.
[52] Iyy Sutskever, Oriol Vinyals, and Quoc V Le. Sequence to sequence learning with neural networks. In Advances in Neural Information Processing Systems, 2014.
[53] Zohar Karnin, Tomer Koren, and Oren Somekh. Almost optimal exploration in multi-armed bandits. In Proceedings of the 30th International Conference on Machine Learning, 2013.
[54] Lisha Li, Kevin Jamieson, Giulia DeSalvo, Afshin Rostamizadeh, and Ameet Talwalkar. Hy- perband: A novel bandit-based approach to hyperparameter optimization. Journal of Machine Learning Research, 18(18)):1-52, 2018.
[55] Thomas Helmuth, N. McPhee, and L.. Spector. Program synthesis using uniform mutation by addition and deletion. Proceedings of the Genetic and Evolutionary Computation Conference, 2018.
[56] Noam M. Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. arXiv preprint arXiv:1804.04233 2018.
[57] Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position repre- sentations. In NAACL-HLT, 2018
[58] Taku Kudo and Richardson. Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. In EMNLP, 2018.
[59] David Patterson, Joseph Gonzalez, Quoe V. Le, Chen Liang, Lluís-Miquel Munguía, D. Rothchild, David R. So, Maud Texier, and J. Dean. Carbon emissions and large neural network training. arXiv preprint arXiv:2104.1 10350, 2021.
[60] 24/7 carbon-free energy: Methodologies and metrics. httpp::/wwir http:///tt/tttttii.....mmmmm. com/ gumdrop/ mustainability /1lltyyyyy////i/ii -ttp://22xaacaoooooaaa frenn--reeeeeeee methodollghoeeeesssseeeesseseeee meweooiiiiieeeeiiii..ioeo.ee .df. Ac-
cessed: 2021-09-14.
13


Reviews:
[{"review_id": "SzGrBWuyovX", "reviewer": "Reviewer_LGXf", "summary": "This work performs a neural architecture search to find an alternative to Transformer language models for efficient training and inference. Primer is a resulting architecture and its two main modifications are (1) squaring ReLU activations and (2) adding depthwise convolution to attention multi-head projections. Primer and Primer-EZ (applied these two modifications only) achieve better efficiency-accuracy trade-off while also following a scaling law like Transformer. ", "questions": "", "limitations": "", "rating": 9, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is well-organized and clearly written. Extensive experimental results and analysis on them support the effectiveness of the proposed method. Although the scope of this work is for autoregressive language modeling, the impact of this work could be quite significant in that Primer can be used for the alternative to Transformer which is currently a defacto standard architecture for NLP and many other ML problems. Primer is found by the automatic architecture search rather than designed by human intuition \n\nStrong initialization with Transformer what the authors called conceptual initialization is a neat idea for the efficient search. However, it reduces the scope of possible architecture significantly, limiting a chance to find a largely different novel architecture. It would be nice if there is a way to alleviate this hard limit with some kind of relaxation.\n\nIt would be interesting to check whether the search space contains previous works on modifying a Transformer architecture in near search steps from the initialization after an exhaustive literature survey on those works or fix the search strategy inspired by those methods.\n\nIt is questionable that the most training efficient architecture (having the best training accuracy with a fixed amount of budget) is also the most accurate model having the best test accuracy because generalization ability might be sacrificed.\n\nIs there any way to include efficiency metrics directly to a fitness score during the search?\n\nIn Section 3, other modifications of Primer except two major modifications are removed by future post-processing. Is there a way to prevent these marginal or not generalizable modifications as a result of the search in advance without requiring human intervention?\n\nMinor: the term of DNA is used without mentioning the unabbreviated version (maybe Differential Neural Architecture) first.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The paper honesty reveals its limitations and suggests meaningful future directions. This work focuses on autoregressive language modeling (decoder-only). Nevertheless, Appendix A.11 includes results on encoder-decoder masked language modeling. Results on encoder-only language models like BERT would be also interesting. Since Transformer is universally used for encoder-only, decoder-only, and encoder-decoder, I am curious whether Primer can take a similar role.\n\nI am wondering Primer can perform well even on a GPT-3 scale as the authors mentioned in the limitation section and other modalities.\n\nDoes the proposed method generalize to other deep learning libraries other than TensorFlow like PyTorch?", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "12 hours"}, "extra_text": {}}, {"review_id": "Druy5qy215_", "reviewer": "Reviewer_rJ28", "summary": "The paper defines a new neural architecture search framework based on generating TensorFlow-based neural network programs. The paper discusses the search space and the evolutionary search-based searching method. The authors perform the search algorithm on a language modeling benchmark and show two specific modifications to the original Transformer model in the resulting model generalizes well to other tasks: squaring ReLU activation and adding a spatial depth-wise convolution after the QKV transformation of the multi-head attention layer.\n\n", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I lean towards accepting the paper. The paper performs neural architecture search on a much larger searching space compared to previous works and provides insights from search results that would be useful for a much wider set of language modeling tasks. Detailed comments:\n\nStrengths:\n- This paper works on a much larger search space compared to previous works.\n- The distilled insights (i.e. two techniques in Primer-EZ) would be greatly beneficial for the community. I believe this actually will ease the reproduction of the paper because it can now be verified by the evaluation of the effectiveness of these two techniques.\n- The evaluation in the paper is comprehensive and the paper is clear and easy to understand.\n\nWeaknesses:\n- It would be better to include deeper discussions on why the resulting modification performs well.\n- The search result is only a slight modification of the original transformer structure. This seems to contradict the starting point of the paper, which is to introduce a larger search space compared to previous works.\n\nMinor issues:\n- Line 59 “a evolutionary search DNA” -> “an evolutionary”\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Please refer to the main review section.\n\n", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "D0NyHBi0TQ", "reviewer": "Reviewer_WYyB", "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n* I really like the idea of Primer-EZ, where they strip down all the complexity induced by the architecture search and focus on the smallest number of modifications that they can make while still getting good results. One of the downsides of typical NAS is that the resulting architecture is a mess and impossible to implement from scratch---this is fortunately not true for Primer-EZ.\n* I like that they tested across numerous codebases and datasets. My takeaway from this part was that Primer-EZ actually works, as opposed to other transformer modifications which arent reproducible across different codebases https://arxiv.org/abs/2102.11972.\n\nWeaknesses:\n* The only thing that would make this a “home run” paper for me would be experiments that show gains on masked language model pretraining + finetuning. Right now it seems that Transformer++ is a strong baseline in this setting (see the Appendix). However, I agree with the authors that this is out of scope and that focusing on autoregressive language modeling itself is a fine goal. \n* I would like to see more discussion and comparison with the Evolved Transformer architecture https://arxiv.org/abs/1901.11117. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes, this work can improve training efficiency and reduce costs across the field.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "6-ZJ1JYMZZ3", "reviewer": "Reviewer_dssT", "summary": "This paper mainly investigates how to perform architecture search to obtain the efficient transformer for auto-regressive language modeling. Compared with previous works, this searched model introduces two simple and effective modifications: squaring ReLU and depthwise convolution. Experimental results also demonstrate that the proposed model can achieve comparable results with fewer computations.", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\n1. This paper applies neural architecture search, which adopts regularized evolution and conceptual initialization, to obtain the efficient transformer model for language modeling tasks.\n2. Experimental results have validated the effectiveness and efficiency of the searched model.\n\nWeaknesses:\n\n1. The contribution of this paper is incremental. First, the searching algorithm is based on the existing works (Evolutionary search). The main difference is this paper aims to obtain a trade-off between the training time and sample efficiency, but what is the optimal trade-off has not been discussed. Second, as a NAS method, the searching algorithm is more important than the searched model. However, the authors cost many columns for introducing the searched model. \n2. The part of the introduction section is like a program instruction. I am not sure whether it is appropriate for a research paper.  \n3. Authors highlight that the multi-DConv-Head attention is important, but no ablation study has been conducted for analyzing the impact of this module.\n\nQuestions:\n\n1. Is Tensorflow Program necessary? If the core contributions of this paper are the searching algorithm and the searched models, I think it should not be sensitive to the code platform (like PyTorch or MXNet). In other words, I think it is unnecessary to highlight Tensorflow usage in a research paper.\n2. This paper argues the improvements are mainly from Squared ReLU and Multi-DConv-Head Attention. So these modules are determined by neural architecture search, or just hand-crafted design?\n\nComments:\n\n1. Line 115, \"search search\" => \"search\"\n\n=======================\n\nAfter read authors' response, I have raised my score to 5. But I still lean negative opinion on this paper. Here are my responses:\n1.\tJust as aforementioned, your paper focuses more attention on the model (named Primer) rather than the searching algorithm. Therefore, I think your title needs to be changed, like \"Primer: xxx\" rather than \"searching for efficient transformer\". Otherwise, it is unfair for other NAS papers, let alone you have not compared with any other NAS-based methods. And your paper is mainly focused on the autoregressive language model, called \"for language modeling\", which is also a little over-claim.\n2.\tSearching efficient transformer has been surveyed in many previous works [a][b][c]. The difference between these works is that they only focus on lightweight compressed transformers due to the limitation of resources.  Considering no one has searched such a large-scale model previously, it can be viewed as a contribution and I have raised my score to 5 based on this merit. However, a question is if under the same computation, is it possible to obtain a better model than other NAS methods? \n3.\tAbout Squared ReLU and Separable Conv plus multi-head attention, I have privately conduct translation experiments on the original Transformer via replacing ReLU with Squared ReLU, and do not observe any improvement in bleu (maybe it is more suitable for the searched architecture or requires different specific hyper-parameters). And for Separable Conv plus multi-head attention has been used in many previous works, including NAS-based methods [a][b][c] and some computer vision works [d][e] also point out that Separable Conv plus multi-head attention can bring additional benefits. Therefore, it is unnecessary to emphasize the improvement is from Squared ReLU and Separable Conv plus multi-head attention in the search space. And more importantly, a convincing approach to validate the effectiveness of Squared ReLU and Separable Conv plus multi-head attention is using a search space with the constraints that we can not use Squared ReLU and Separable Conv plus multi-head attention. In other words,  under these constraints, maybe we can get other advanced architectures, rather than Squared ReLU and Separable Conv plus multi-head attention. However, the authors only tried to replace the corresponding module on the original Transformer and Primer in the ablation study, rather than replace the search space and re-search architecture. \n4.\tWill the searching code be released? \n \n \n[a] AdaBERT: Task-Adaptive BERT Compression with Differentiable Neural Architecture Search.\n\n[b] NAS-BERT: Task-Agnostic and Adaptive-Size BERT Compression with Neural Architecture Search\n\n[c] AutoTinyBERT: Automatic Hyper-parameter Optimization for Efficient Pre-trained Language Models.\n\n[d] CvT: Introducing Convolutions to Vision Transformers\n\n[e] Co-Scale Conv-Attentional Image Transformers\n \n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "12"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The submission makes two main contributions: defining neural architecture search using lower level primitives than previous work, and using this search to introduce a Transformer variant they call Primer. Results show that Primer can outperform vanilla Transformers on various language modeling tasks.

While three of the reviewers support acceptance, Reviewer dssT raises several concerns, arguing that a new search space is not enough of a contribution. The reviewer also would like to see a search space that does not include the squared ReLU and Separable conv plus multi-head attention that are used in the final solution. Other reviewers appreciated the larger, low level search space. I think that the authors convincingly rebut the need to exclude modules from the search space, as for example, excluding squared ReLU would require excluding multiplication. There are also concerns that it is only evaluated on language modeling, but the authors are up front about this limitation, and I agree with most of the reviewers that this is an important and general enough task to focus on.

Overall, the consensus is that this is a good paper that should be accepted.
Decision: Accept (Poster)