Sparse is Enough in Scaling Transformers
Authors: Sebastian Jaszczur
Aakanksha Chowdhery
Afroz Mohiuddin
Lukasz Kaiser
Wojciech Gajewski
Henryk Michalewski
Jonni Kanerva
Venue: NeurIPS 2021 Poster (2021)
Keywords: machine learning
transformers
sparsity
scaling

Abstract: Large Transformer models yield impressive results on many tasks, but are expensive to train, or even fine-tune, and so slow at decoding that their use and study becomes out of reach. We address this problem by leveraging sparsity. We study sparse variants for all layers in the Transformer and propose Scaling Transformers, a family of next generation Transformer models that use sparse layers to scale efficiently and perform unbatched decoding much faster than the standard Transformer as we scale up the model size. Surprisingly, the sparse layers are enough to obtain the same perplexity as the standard Transformer with the same number of parameters. We also integrate with prior sparsity approaches to attention and enable fast inference on long sequences even with limited memory. This results in performance competitive to the state-of-the-art on long text summarization.

Sparse is Enough in Scaling Transformers
Sebastian Jaszezur' University Warsaw
Aakanksha Chowdhery
Afroz Mohiuddin
Lukasz Kaiser" OpenAl
Google Research
Google Research
Wojciech Gajewski
Henryk Michalewski
Jonni Kanerva Google Research
Google Research
Google Research
Abstract
Large Transformer models yield impressive results many tasks, but are expen- sive to train, or even fine-tune, and so slow at decoding that their use and study becomes out of reach. We address this problem by leveraging sparsity. We study sparse variants for layers in the Transformer and propose Scaling Transformers, family of next generation Transformer models that use sparse layers to scale efficiently and perform unbatched decoding much faster than the standard Trans- former as scale up the model size. Surprisingly, the sparse layers enough to obtain the same perplexity as the standard Transformer with the same number of parameters. We also integrate with prior sparsity approaches attention and enable fast inference on long sequences even with limited memory. This results performance competitive to the ofett---rr-ttttttt long text summarization.
Introduction
The field of natural language processing has seen dramatic improvements recent years due to large neural networks based on the Transformer architecture. The original Transformer [42] significantly advanced state-of-the-as in machine translation. BERT [7] surpassed all previous methods on question answering, language inference and other NLP tasks and was followed by line of models like T5 [30] that further improved these results. The GPT line of models [29, 3] elevated language generation to the point that GPT-2 was invited to write short passages for the Economist and GPT-3 created whole articles almost indistinguishable from human-writte ones.
The benefits of this progress are undercut by the huge costs such models incur. Strubell et al. [36] estimate that training single base BERT model costs $4k-$112 and emits as much CO2 as one passenger's s share of 4-hour flight and later Patterson et al. [27] estimate that training GPT-3 has three times as much (CO2e (metric tons of CO2 equivalent) emissions as a SF-NY round trip flight. Data and serving costs are also forbidding: single training run of BERT, for example, processes 128B tokens, and Google Translate reportedly' serves over 143B words per day.
With the growing popularity and size of these models, it is increasingly valuable to make them scale efficiently. In this work we propose Scaling Transformers with separate sparse mechanism for the query, key, value and output layers (QKV layers for short) and combine it with sparse feedforward blocks get fully sparse Transformer architecture.
To quantify the computational complexity of inference Transformer models, recall the architecture of Transformer decoder block. It consists of three parts: masked self-attention layer, encoder- decoder attention layer and a feedforward block. The sizes of these layers are parameterized by dmndel and dØ. The base BERT model sets dmodel 768, the large BERT has dmodel 1024, the largest
Work done while at Google Research. btttps://cuutt.l//akk///
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Params
Dec. time
Dec. time
per block
1,700
Avg baselines
baseline Transl.
Sparse FF
800M
0.160s
0.093s
3.Ims
5.9ms
1.675
Spensel FT, N-b4
Sparse PT. N=128
+ Sparse QKV
0.152s
6.2ms
1/650
Sparse
Sparse FF+OKV
0,061s
1.9ms
Sparse FF+OKV
Sparse OKV, mult
Speedup
2.62x
3.05x
baseline Transf.
17B
3.690s
0.581s
1,575
+Sparse FF
1.595s
0.259s
1.550
+Sparse QKV
3.154s
0,554s
1,225
+Sparse FF+QKV
0.183s
0.014s
400036
420000
440000
460000
430000
500000
Speedup
20.0x
42.5x
Kumber of steps
Tuble 1: Decoding speed (in seconds) of a single token. For Transformer model lequivalent T5 large with ap- proximately 800M parameters), Scaling Transformers with proposed sparsity mechanisms (FF+QKV) achieve up to 2x speedup in decoding compared to baseline dense model and 20x speedup for 17B param model.
Figure 1: Log-perplexity of Sealing Transformers (equivalent to T5 large with approximately 800M pa- rameters) on C4 dataset with proposed sparsity mecha- nisms (FF, QKV, FF+QKV) is similar to baseline dense model. Other models used in this paper are shown in grey lines; raw data is available in the appendix.
GPT-2 has dmodet 1600 and GPT-3 reaches dmodel 12288. For both BERT and GPT models the authors use dæ = 4 dmodel. While decoding a token, the self-attention layer needs to activate four matrices of size duandel x drwodet: one each for the queries, keys and values input to the attention and one for merging the output. In the encoder-decoder attention, the keys and values may already be cached, only two matrices of size dmodee dmodel activated. The feedforward block consists of two matrices of size dmnddel X de, omitting small additional contribution of biases. The total adds up to: 4 model 2d d clodd + dmmdde dg. This describes both the number of trainable weights of a single block and approximates well the number of floating-point operations needed for decoding single token, except for the attention operations (discussed later). The complexity is quadratic in dmodet; for example, as dmmddl increases 16-fold from base BERT to GPT-3, the complexity of a single block grows 256-fold.
Scaling Transformers use only 2dmodet = 2d parameters in QKV layers and yield results good as the baseline (fully dense) Transformer with the same number of parameters and complexity: 8 cl 1.5 + 4 cl 1-5 model + 4 1.5 medel We were surprised that the fully sparse Scaling Transformers are indeed enough to match the results of the baseline Transformer on the large C4 dataset [30] (Figure 1). The improvement in complexity holds not just asymptotically but yields over 2.6x speedup in wall-clock hed decoding time already for model with 800M parameters and 20x improvement for a model with 17B parameters, as shown in Table
In comparison
To verify that Scaling Transformers can be used with other Transformer improvements real tasks, we create Terraformer - a Transformer model that uses reversible layers for memory efficiency and sparse attention to handle long sequences. We pre-train Terraformer on the C4 dataset and fine-tune it on the challenging task of summarizing arxiv articles. Terraformer yields results competitive to the state-of the-art BigBird-Pegasus without using the Pegasus loss in pre-training (Table 5).
2 Related Work
As discussed in the previous section, large Transformer models brings significant improvements in performance, as seen in models such as GPT-3 [3, 17] or T5 [44, 30]. Training and inference incur a high computational cost at the scale of hundreds of billions of parameters. Numerous techniques improve the efficiency of Transformer models, and Gupta and Agrawal [11] divide them into several classes, including pruning, knowledge distillation, quantization, parameter sharing, efficient attention, and efficient feedforward.
Model compression. Model pruning [24, 2] makes matrices smaller by removing unneeded weights after during training, however, the gains in computational complexity on sparse matrices often do
Thhe 800M model 16, ff ifasspasity- 64. We scale up this model to approximately 17B parameters with daodel = 9216 and get up to 20x speedup in decoding compared to baseline dense model. This 17B param model has six layers of Encoder & Decoder, 96 attn heads, iicention-sparsiii 64, ff ff-sparsity 256.
layers of Encoder Decoder, dmakk 1024, 16 attn heads, attention attention-sparsity

---

not result in inference speedups on actual hardware [9]. Structured pruning based approaches [47, 22, 43] account for this challenge by leveraging sparsity hardware CPU and GPU architectures [1]. Our paper is different from pruning approaches in that it relies on dynamic sparsity wherein the feedforward layer loads only a subset of weights in the layer for each token. Our approach is complementary to model quantization studies [35, 38, 28] that use fewer bits for the weights.
Model distillation. Several natural language models used for mobile inference [13, 39] rely on distillation [32] to speed up inference from the pretrained large models. For example, [18] pretrains a large model and uses knowledge distillation along with pruning to get more than 10x faster inference. Instead of distilling a large model, our approach speeds up inference by reducing the number of weights loaded in memory from the model.
Sparse attention. Sparse caanttion-baeed approaches have made the attention layer more efficient, especially for long sequences, by incorporating additional combinatorial mechanisms, in [40], or selecting a subset of tokens this layer attends to [31, 5, 19, 37, 15, 4] or other approaches [12]. Our work is complementary to these approaches for sparse attention and reuses the advances on SOTA therein. Inference speedups in the attention layers also use bottleneck layers [39] or grouped convolutions [13]. Our work extends beyond the idea of grouped convolutions approach because each attention head is limited to using only a fixed part of the embedding while our work is able to permute the embeddings to improve model quality; see Section 3.2 for details.
Tensor Decomposition. and handling of long sequences, however none of them addresses the fundamental scaling issue: even we distill into a smaller model, quantize it and prune a percentage of the weights, the complexity still grows quadratically with dmmdeel The final approach, which does attack this scaling issue, is called tensor decompositions in [11]. Unluckily, as the authors there note, the approach is most effective in dealing with large input and output embedding matrices and tends to produce lower performance than unstructured models if used inside the decoder block.
The approaches discussed above significantly improve Transformer speed
Sparse feedforward. Mixture of experts approaches have been shown to achieve computational efficiency in training [33, 21, 34], scaling up to trillion parameters [8]. The key idea is to partition the dizsi dimension into parts (called experts) and retrieve only one part per token, which reduces the complexity of the feedforward block from 2d to 2dmeddeddy Dexperts- These speedups are mostly measured in training speed, and the method focuses on feedforward blocks. In contrast to prior methods, we train a full weight matrix and then only activate specific parts of it for each input token during decoding; see Section 3.1.
3 Sparse is Enough
We study how to sparsify every part of the Transformer model-otherwis the non-sparse parts dominate decoding time and become bottleneck. This means we need sparse equivalents for the feedforward blocks, for the dense Q K, and output layers attention, and for the final dense layer before the softmax and loss.
3.1 Sparse Feedforward Layer
In a baseline Transformer, decoding speed is dominated by the execution cost of the feedforward block. Recall that this block consists of two fully-connected (dense) layers with a ReL.U nonlinearity in between. The dimensionality of activation vectors between these layers is usually denoted by de and is often 4 or 8 times larger than the dimensionality of the activations in other places (dmodel).
We make use of the structure of the feedforward block to sparsify it. One main observation is that the ReLU in the middle creates lot of zeros?. We impose fixed structure on this middle activation vector: only one float in every block of N will be allowed to be non-zero. Prior techniques prune weights or blocks from weight matrices and can be referred to as static sparsity. Our proposed technique will train a full weight matrix but only activate specific parts of it for each input token during decoding. We call this dynamic sparsity, because the model dynamically selects only a fraction of its parameters, and the selection is independent for each token.
2GeLU is another non- non-linearity often used in the Transformer feedforward block. Table 1 in [26] shows the same final loss for ReLU and GeLU the C4 dataset, though, in this work for simplicity, we focus ReLU.
3

---

6.mose
d.model
low.rank
model
mat
6.11
activation vector
model
controller's output
ReLU
activation vector
Straight-Through Through
Gumbel-offmme
feed-forward output
(a)
(b)
Figure 2: (a) Sparse Feedforward Layer only activates I in N rows/columns of each block reduce the decoding time. Here only two rows/colums blocks of size 4 are loaded while the weights dark are not loaded from memory during inference. (b) Sparse Feedforward Controller with the ouput of blocks of size (1 in sparsity).
We train a controller to determine which activation in each block can be non-zero; the rest will be set zero. This can be represented
sparse max xW1 + 61) @ Controller(z))
SparseFFN (x) Y sparse W2 + 62
where @ element-wise multiplication. Note that each activation in Spsrss corresponds to a single column in W1 and single in W2. Therefore, if we compute Controller(z) output first, we don't 't have to use any columns in W1 or any rows in W2 that correspond to an activation set to zero by the controller. This allows for much faster decoding, and rows in W2 (see Figure 2(a)).
have process only in N columns in Wi
To design the controller to be computationally inexpensive, we project the input using a low-rank bottleneck dense layer. Figure 2(b) illustrates the controller which produces the output as follows
Controller (x) = arg maxx(Reshapp((((CC (-1,,))
where C1 E Relaced chewanic and C2 Rdhmma with diowrank usually set to (dmnntel/N..
During inference the controller uses a discrete argmax function, but during training the model softmax to calculate and sample from distribution. The model learns to select which row/column will be non-zero using the Gumbel-Softmax trick for discretization. To determine the active row/column in each block, we reparameterize sampling from a Bernoulli distribution by using the Gumbel-Soffmail trick [25]. Instead of using the logits in each block directly sample binary value, we add independent noise from the Gumbel distribution to each of the logits, and then select the binary value with the highest logit (i.e., argmax) as the sample z. The argmax operation not differentiable, but it can be approximated by softmax with annealing temperature. Therefore, on the forward pass, we use the argmax to obtain binary one-hot vector for each block, while on the backward pass, approximate it with softmax. This approach is known as the Straight-Through Humbel-Softmal estimator [14].
Ablations. We investigate the impact of sparse levels of sparsity. with dmodel 1024, du = 4096, and 16 attention heads. When we set the sparsity level to N (for e.g. inference. During training, the controller uses the bottleneck layer with diowrank = 64 and temperature of Gumbel softmax estimator set 0.1. To improve training stability, the controller in the forward pass will use the output of argmax that is binary one-hot vector for each block with probability of 30% and otherwise it uses the output softmax. Table 2 and Figure show the perplexity and the decoding time of this model with varying levels of sparsity in feedforward layer. As the level of sparsity increases from to 128, we observe a significant decrease the decoding time, while the neg-log-pent peolpppelyyxttyy of the model with N 64 sparsity is comparable to the baseline.
the model equivalent to T5-large with varying
64) then every block of size N has one non-zero value activated for

---

1.4
1.62
Avg of baselines SFF, N=
1.80
N-1
Dec, time
1.58
baseline Sparse FF 64
0.160s
1.56
0.093s
1.54
Sparse FF 128 0.089s
420000
42000G
440000
Numbe of steps
Table 2: Decoding time of singe ken decreases with increasing level of sparsity in FF layer.
Figure 3: Log- perplexiry of Scaling Transformers with Sparse Feedforward layer very similar dense buse- line for sparsity level but degrades slightly for N=128.
We also checked the performance of the feedforward block with Mixture-of- [33] style sparsity. As expected, this technique achieved decoding time comparable to sparse FF 0. is instead 0.09s but with its lack of granularity it achieved perpleetpy of 1.64, worse than both method and the dense baseline.
3.2 Sparse QKV Layer
The decoding speed for model with sparse feedforward blocks is dominated next by the query, key, value and output computation Each of these dense layers has dmodel parameters and computation cost. Unfortunately. QKV layers don't have ReI.Us, so the method used above to sparsify feedforward blocks not viable here.
dense layers in attention, which we jointly call a QKV layer.
To make QKV layers sparse, we subdivide the dimensionality of the layer, dmodel into S modules of size = dmodde/S SS, similar to splitting an activation vector into multiple heads. These modules can be processed with a convolutional layer with fewer weights and faster computation. However, with naïve design each module (and corresponding attention head) could access only small part of given token embedding. To alleviate that, develop multiplicative layer that represent arbitrary permutation and has fewer parameters and lower computation time than a dense layer. This multiplicative layer is inserted right before the convolutional layer, letting cach head access any part of the embedding (see Figure 4(a)). This solution yields well-performing models that also decode fast.
Multiplicative dense layer. Our new multiplicative dense layer can represent an arbitrary permuta- tion and has d /S S parameters, dependent the sparsity hyperparameter S. It processes an input vector E Bdolt by splitting it into S "modules" of size M dmeda// SS It produces output y E SSX as follows
y s.m
where the two weight matrices are X and E E X.M (see Figure 4(b)). This layer executes significantly faster during inference because of the decreased number of parameters which need to be loaded from memory. Unless stated otherwise, we use S 16.
The multiplicative layer is designed primarily to represent any permutation, so that each attention head can access information from any part of the embedding. We first verify that the multiplicative layer can indeed represent an arbitrary permutation (the proof is presented in the Appendix).
Theorem 1.
For any bijective function f (1.... -mmaeet 7 {1 .... 5} x {1 .1...M. M) there exists a pair of weights of multiplicative layer D, E such that x; Vs.m for {s, m) f()..
Convolutional layer. The output of the multiplicative layer is tensor of type/shape E lenghhi SMM We process this tensor with a two-dimensional convolutional layer, treating the length dimension and number of modules S like height and width of an image. This layer uses M filters and kernel size of F x F so that each filter looks modules ('S axis) of the last F tokens ('length' axis). Replacing the standard dense layer with such convolution reduces the parameter
5

---

weight
dense dense
weight
matrix
mooel m
baseline QKV
attention mechanism
dense
matrix D
model
dense
sparse QKV:
mult
mult
mult
attention mechanism
mult
d_model
element-wise
mat mul
mult
multiply
sparse QKV:
input:
mult-conv
cony
mult
conv
mechanism
attertion
output y
(a)
(b)
Figure 4: (a) Multiplicative layer can represent an arbitrary permutation, but has fewer parameters and reduced computation time compared to dense layer. (b) Sparse QKV layer replaces Q. K, and V dense layers by compoxing multiplicative and convolutional layers and reducing the number of parameters and decoding time.
count and computation time of the QKV layer. At the same time, by convolving over the 'length' axis, the model can incorporate more context into this computation [23].
The output of this layer has the same shape the input. The optimal value of S is less than y dmodel Empirically we set F' to 3, S equal to the number of heads in the attention mechanism and M to be the dimensionality of single attention head. In this case, we can feed the output of the convolution directly to the attention mechanism without reshaping the output. This convolutional layer has fewer parameters (9M22 + M F² (dmodat/ S)² + (dmodee / /S)), and lower computational complexity (oode S)). Unless stated otherwise, we use S = 16 and F 3.
Combining multiplicative and convolutional layers. There are four dense layers to replace in the original attention mechanism: Q K. V, and output. As shown in Figure 4(b), we replace Q. K, and V dense layers by composing multiplicative and convolutional layers, but with a multiplicative layer shared across all three: Q = conv (mult(z)), K = conv x (iiix<(muttt (mult(z)), V convy (mult(()) We remove the output dense layer. Note that the combined multiplicative- catiipiicaaieeecoooluutiaaall variant has the output dense layer removed, while the other variants have replaced with their respective sparse layers. Including this output layer negatively impacts decoding time. We can set the parameter S to around dmmodet, getting the number of layer parameters to scale proportionally to d compared to d doodet of standard QKV layer.
Interpretation of QKV layer. Note that when parameter S' in convolutional layer is equal to the number of heads the attention mechanism, which the case in our experiments, then each of the modules corresponds to a single attention head. Therefore. the model uses the convolution to process each head using the same linear projection. Without the multiplicative layer this projection would operate on a predetermined part of the embedding layer for each head. However, by adding it the model can perform arbitrary permutation of dimensions, so each head can have access to arbitrary subset of embedding dimensions, not a predetermined subset of them. This fact helps with keeping the expressibility of resulting QKV layer despite the reduced number of parameters.
Ablations. We investigate the impact of sparse QKV layers on the model equivalent to T5-laage in Figure 5. We increase the value of du from 4096 to 6144 to preserve the number of parameters (see the next subsection for details). The decoding time with sparse QKV layer variants is similar to the baseline because it is dominated by the dense feedforward layer (details in appendix).
Combined feedforward and QKV sparsity. Sparse QKV layers lower the total number of model parameters. To keep the model size matched to the baseline, we increase da to keep the number of parameters similar across all models we compare. For the T5-Large equivalent model, we increase du from 4096 to 6144. With increased da, decoding time in the feedforward layer increases and thus, Sparse QKV layers alone do not speed up the model. However, when we combine Sparse QKV layers with sparse FF layers, we get 3.05x speedup in decoding time of each decoding block with

---

1.60
log-perplexity
1.58
Avg of baselines
1.56
SOKV, mult-conv, S-8, Fw3
SOKV,
S-16, F-3
1.54
SOKV, mult-conv, 5-16, F-1
SOKV, conv, S-16
1.52
SOKV, mult, 5-16, F-3
400000
420000
440000
460000
480000
500000
Number steps
Figure Log-Lerplesn perplexity Scaling Transformers with Sparse QKV with different sparsity levels (S) and kernel sizes (F) is very similar to dense baseline within variance while multi layti-layer even improves perplexity.
RIE
MRPC
SST-2
QNLI
MNLI-m
QQP
Baseline Transformer (dense)
70.1
0.72
+
U.5
0..4
Scaling Transformer (Sparse FF+QKV)
68.4
81.2
91.6
90.1
82.9
89.9
Terraformer (Sparse FF+QKV)
66.1
84.6
92.3
88.3
79.1
85.5
Table 3: Accuracy of Scaling Transformer model and Terraformer model with sparse QKV+FF is comparable the baseline Transformer within variance. The results obtained by fine-tuning on selected downstream tasks from the GLUU dataset (validation split).
comparable perplexity (see Table and Figure 1). While the baseline these is vanilla Transformer, the decoding speed is almost the same for a Reformer model as well.
Table 3 shows the accuracy of fine-tuning the model for downstream tasks from the GLUE dataset. Note that the model with SparseFF+QKV achieves accuracy similar to the baseline.
Sparse loss layer.
A final dense layer maps the model embedding into vocabulary size to compute the loss. We can sparsify this part of the model by replacing the dense layer with multiplicative layer similar to previous sections; this speeds up decoding time but may degrade perplexity. The results presented in appendix.
Sparsity for Long Sequences
The above gains from sparsifying the dense layers are encouraging, but we omitted one fundamental issue. When applied to longer sequences, the gains would effectively be lost, as the decoding time will be dominated by attention operations. Luckily, a number of methods have been proposed to solve this problem for Transformers, see [41] for survey. We focus on the LSH Hashing) attention from Reformer [19] and show how integrate this sparse attention mechanism, as well recurrent blocks, into Scaling Transformer, yielding a Terraformer.
4.1 Architecture for Long Sequences
While integrating sparse attention layers into a Scaling Transformer, we notice that the architecture of the Transformer decoder block is suboptimal and can be redesigned to make better use of these layers. In particular, separating decoder self-attention and encoder-decoder attention not necessary any more from the perspective of efficiency. We therefore remove the encoder-decoder attention, but just concatenate the encoder representations before the decoder tokens. Doing this alone isn' enough though, since took away attention mechanism denoder- attention). We remedy this by having two attention mechanisms before the feedforward block. This simple architecture as fast the baseline Transformer while giving better results.
Putting this together, if Vene are the encoder activations and Udee are the decoder embeddings, the input the decoder block is their concatenation the length axis, LengthConcat( (Veney Vdec).

---

X
Y1
0
Attention,
Attention,
FFN
22
0
2
Figure 6: Reversible decoder block Terraformer.
Each decoder block can be represented as:
= x: - Deepuut(AteniionnLayyeNomm
1/2 y11 + DuapoutAAttntoonLLayynommmm]]]])
y y22 FFN(
where y becomes the input to the next decoder layer. See the appendix for full diagram of the resulting architecture.
4.2 Reversibility for Memory Efficiency
To enable training Terraformer with large batches, and to fine-tune even large models single machines, we apply ideas from the Reformer [19], in particular, reversible layers for the encoder and decoder blocks.
The original Reformer decoder block contained feedforward and attention layers in a 1-1 ratio. In the Terraformer architecture, as described above, there are two attention layers in the decoder block, so there are three swaps in the reversible layers in the decoder block (see Figure 6). In our experiments, this significantly improved performance.
Another issue with reversibility is that it is only formally correct for continuous functions. We find that this is not just a formal issue, but an important problem in practice. To make reversible layers train well with sparsity, we need to store the discrete decisions-i.e, the integers saying which rows to select- and use them for reversing. Recalculating these decisions on the backwards pass leads to worse results.
4.3 Recurrence for Generalization
In addition to incorporating sparse attention and reversibility, we also add recurrence to the feedfor- ward block of Terraformer. Recurrent layers allow information propagate in time, even in single decoder block. It is challenging though to use them without decreasing model speed, esp. in training. For that reason, use simple recurrent units [20] which parallelize well during training.
SRUs contain dense layers, so their use could negate the benefits of sparsity elsewhere. We tried few methods to alleviate that, but it turns out that simply reducing the dimensionality of the SRUs works. So we first project from dmmdde then project back to dmodel and add the result to the feedforward block. This low-rank recurrence is in our experiments sufficient transfer enough information through time for the network to generalize.
small dimension (32 our experiments), then apply the SRU, and
Since the effects of SRUs on C4 are minimal (as the training and evaluation data are very similar), use synthetic tasks to investigate nul-of-distribution generalization. We train the models on long addition and on the task of copying a decimal digit. We train on inputs with at most 128 digits and evaluate on inputs lengths from 256 to 300, so over 2x longer. As can be seen in the table below, the baseline Transformer does not generalize well, while Terraformer manages to get large portion correctly, even if not perfect like the Neural GPU [16].
4.4 Experiments
We designed Terraformer so that the benefits from sparsity would not be lost on long sequences, nor on downstream finetuning tasks. To test this, we chose the task of summarizing scientific papers

---

Model
copy
copy (seq)
add
add (seq)
Transformer
79.8%
0%
36.4%
0%
Terraformer
99.9%
93.9%
86.9%
32.4%
Table 4: Comparison of out-of distiiibbbtttinn generalization for Terraformer and Transformer on two toy tasks, long addition and copying decimal numbers. Under (seq) we report the number of fully correct sequences generated as answers.
Model
R-1
R-2
R-LSum
R-LSent
Terraformer
45.40
17.86
41.21
26.33
DANCER RUM
42.70
16.54
38.44
BIGBIRD-ROBERTA
41.22
16.43
36.96
Pegasus Large (C4)
44.21
16.95
38,83
25.67
DANCER PEGASUS
45.01
17.6
40.56
NIGBIRD-Pegasung
46.63
19.02
41.77
Table 5: Terraformer is competitive with strong baselines [46, 45, 10] on the ArXiv summarization task, without
using the Pegasus loss and without beam search. On R-1, R-2 and R-LSum, Terraformer outperforms all previous models except for BigBird-Pegasue
using the dataset of scientific papers from (6) In this task, the input a whole paper-a long sequence-and the model is asked output its abstract. Several recent papers studied this dataset and tasks and it has been shown [46, 45J that pretraining on C4 yields significant improvements on this task. We also pretrain Terraformer on C4 (like in all experiments in this paper) and fine-tuned on the ar summarization task. We find that Terraformer is competitive with the above baselines, even though we mask single words (we do not use the Pegasus sentence loss) and decode the answers in greedy way (no beam search). Note that ROUGE scores are computed using open-source scorer with the metries deseribed in its documentation", We also observe certain confusion between ROUGE-L metrics reported. As noted in the open-source scorer, there are two versions of ROUGEL-Sennence- Level (R-LSent) and HOUGGL-Suumary-LLen
(R-LSum).. For clarity, we report both of these metrics. Furthermore we only report the F1 measure of any ROUGE metric. We include a few examples of the generated abstracts in the appendix.
We pretrained Terraformer in the same way as all other baselines reported in this paper with the same number of parameters (800M), the same dimensions as mentioned before, and loss sparsity to get the fastest model. Compared to the sparse Transformer model from the previous section that achieves a decoding speed of 0.061s. Terraformer achieves decoding speed of 0.086s with a similar performance in terms of perplexity (see appendix for details). We also observe that the Terraformer model achieves accuracy similar to the Transformer model Table for selected downstream tasks on GLUE dataset.
Table shows the speedup in decoding with sparse layers when we scale up Terraformer 17B parameters. Note that sparsifying all the layers gives us 37x speedup in decoding.
Conclusion
When starting to investigate sparse variants of Transformers, we assumed that there would be price to pay for sparsity-that a sparse model would always underperform a dense one with the same number of parameters. To our surprise, this is not the case: sparse is enough!
In our experiments with large models on the C4 dataset, the sparse models match the performance of their dense counterparts while being many times faster at inference. And, when scaling the models up, the benefits of sparsity become even larger. This promises put Transformers back on a sustainable track and make large models more useful.
We provide full details of our datasets, hyperparamcters, and everything needed to reproduce the results in the appendix. The code is open-sourced part of Trax 1.4.0 https://github htt.:ithhhh....oeeeeeeeeeeeeeeee wwwggooogl//taa....ooooooooo haappp/orrgpoojee/rrrgeee phttps://githie http://gttuuu...cogooooooee ottp://ooo/ooo///////////////e//

---

Terraformer
Dec. time
Speedup
The current results have number of limitations. For one, the practical speedups we see are only for inference, not at training time. Moreover, we consider unbatched inference on CPUs, while often inference is ran in batched mode on GPUs. We believe with more work sparsity can bring improvements in these settings too, as our funda- mental result shows that the sparse models reach the same perplexity as their dense counterparts with the same number of parameters.
dense
3.651s
1x
Sparse FF 1.595s
2.29x
SparseFF+QKY 0.183s
19.98x
0.097s
37.64x
Table 6: Decoding speed of single token for Ter- roformer with 17B parameters 37x faster than dense baseline model, requiring less than 100ms/token for in- ference. Here attention-purrsity = 64, f-sparsity 256, and loxs- = 4.
So while we demonstrate that Scaling Trans- formers are possible, we consider this paper as a first step on the way to sustainable large models. There are numerous techniques for making models faster that could greatly benefit Terraformer and other Scaling Transformers. For example, we did not study quantization and we believe that it can make Scaling Transformers even faster. We also focused on inference speed and did not get improvements in training speed. The main reason is our use of Gumbel-Softtax when training the feedforward block (see Section 3.1). Fedus et al. [8] already provide a promising alternative, and we look forward to exploring it in future work.
Further, we hope that the community will take inspiration from Scaling Transformers and tune them for their needs. We ran experiments using layer sizes and hyperparameters borrowed from dense Transformers and they are most probably not optimal for Scaling Transformer. With proper tuning and further improvements we believe one could train a Scaling Transformer to match GPT-3 in accuracy but also run inference in reasonable time on a laptop. We put it as fascinating challenge to the community, since such Scaling Transformers will not only be more sustainable but will also make large models accessible to everyone.
References
[1] Nvidia Ampere Architecture.
[2] Christopher Brix, Parnia Bahar, and Hermann Ney. Successfully applying the stabilized lottery ticket hypothesis to the transformer architecture. arXiv preprint arXiv:2005.03445 2020.
[3] 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 are few-shot learners. urXiv preprint arXiv:2005. 14165, 2020.
[4] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 10509, 2019.
[5] Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009.14794 2020.
[6] Arman Cohan, Franck Dernoncourt, Doo Soon Kim, Trung Bui, Seokhwan Kim, Walter Chang, and Nazli Goharian. A discourse-aware attention model for abstractive summarization of long documents. Proceedings of the 2018 Conference of the North American Chapter of the Association for Compulational Linguistics: Human Language Technologies, Volume (Short Papers), 2018. doi: 10.18653/v htt////w//111//// 2097. URL attp://ds. doi. org/10. 18653/ 21/n18-2097.
[7] Jacob Devlin, Ming-Wai Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 18Xiv::8100 040005, 2018.
[8] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. arXiv preprint arXiv:2101.0300 2021.
[9] Trevor Gale, Erich Flsen, and Sara Hooker. The state of sparsity deep neural networks. arXiv 2019.
preprint
10

---

[10] A. Gidiotis and G. Tsoumakas. A divide-and and-conquer approach to the summarization of long documents. IEEE/ACM Transactions on Audio, Speech, and Language Processing, 28: 3029-3040, 2020. doi: 10.1 10.1109TAAIIP2200000374444
[1]] Manish Gupta and Puneet Agrawal. Compression of deep learning models for text: A survey. arXiv preprini urXiv:2008.0522. 2020.
[12] Tianyu He, Xu Tan, Yingce Xia, Di He. Tao Qin, Zhibo Chen, and Tie- Yan Liu. Layer-wise coordination between encoder and decoder for neural machine translation. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Basa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31, 2018.
[13] Forrest N landola, Albert E Shaw, Ravi Krishna, and Kurt W Keutzer. Squeezebert: What can computer vision teach nlp about efficient neural networks? arXiv preprint 13Xiv:2006.11316. 2020.
[14] Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmat arXiv preprini arXiv: 16/ 1016/11111444 2016.
[15] Lukasz Kaiser and Samy Bengio. Discrete autoencoders for sequence models. arXiv preprint arXiv:180..09997, 2018.
[16] Lukasz Kaiser and Ilya Sutskever. Neural gpus learn algorithms. arXiv preprint arXiv:1511.08228, 2015.
[17] 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.08361, 2020.
[18] Young Jin Kim and Hany Hassan Awadalla. Fastformers: Highly efficient transformer models for natural language understanding. arXiv preprint arXiv:2010. / 3382,, 2020.
[19] Nikita Kitaev, Lukasz. Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001.0445), 2020.
[20] Tao Lei, Zhang, and Yoav Artzi. Training rans as fast as cnns. CoRR, absl/aa9.0275.... 2017. URL httptt://aaxi..orrg///
02755.
[21] Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with condi- tional computation and automatic sharding. arXiv preprint arXiv:2006. 166666688.668 2020.
[22] Bingbing Li, Zhenglun Kong, Tianyun Zhang, Ji Li, Zhengang Li, Hang Liu, and Caiwen Ding. Efficient transformer-basee based large scale language representations using sardware-friende block structured pruning. arXiv preprint arXiv:2009.0s8:
2020.
[23] Shiyang Li, Xiaoyong Jin, Yao Xuan, Xiyou Zhou, Wenhu Chen, Yu-Xiang Wang, and Xifeng Yan. Enhancing the locality and breaking the memory bottleneck of transformer on time series forecasting, arXiv preprini arXiv: 90199000022555 2019.
[24] Zhuohan Li, Eric Wallace, Sheng Shen, Kevin Lin, Kurt Keutzer, Dan Klein, and Joey Gon- zalez. Train big, then compress: Rethinking model size for efficient training and inference of transformers. International Conference on Machine Learning. pages 5958-5968. PMLR. 2020.
[25] Krzysztof Maziarz, Ef Kokiopoulou. Andrea Gesmundo, Luciano Sbaiz, Gabor Bartok, and Jesse Berent. Gumbel-mathe routing for flexible multi-task learning. arXiv preprint 19X1v:1910004915. 2019.
[26] Sharan Narang, Hyung Won Chung, Yi Tay, William Fedus, Thibault Fevry, Michael Matena, Karishma Malkan, Noah Fiedel, Noam Shazeer, Zhenzhong Lan, Yanqi Zhou, Wei Li, Nan Ding, Jake Mareus, Adam Roberts, and Colin Raffel. Do transformer modifications transfer across implementations and applications? arXiv preprint 212. 1972, 2021.
11

---

[27] David Patterson. Joseph Gonzalez, Quoe Le. Chen Liang, Lluis Lluis-Miquel Munguia, Daniel Rothchild, David So, Maud Texier, and Jeff Dean. Carbon emissions and large neural network training. arXiv preprint arXiv:2 10350, 2021.
[28] Gabriele Prato, Ella Charlaix, and Mehdi Rezagholizadeh. Fully quantized transformer for machine translation. arXiv preprint arXiv: 111199911000.....1 2019.
[29] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
[30] 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 unified text-to-text transformer. Journal of Machine Learning Research, 21:1-67, 2020.
[31] Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. arXiv preprint arXiv:2003.05997, 2020.
[32] Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1910.01108, 2019.
[33] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outragcously large neural networks: The sparsely-gate gated mixture-of-experts 2017.
layer. arXiv preprint
[34] Noam Shazeer, Youlong Cheng, Niki Parmar, Dustin Tran, Ashish Vaswani, Penpom Koanan- takool, Peter Hawkins, HyoukJoong Lee, Mingsheng Hong, Cliff Young, et al. Mesh-tensorflow: Deep learning for supercomputers. In Advances in Neural Information Processing Systems, pages 14413-1423, 2018.
[35] Sheng Shen, Zhen Dong, Jiayu Ye, Linjian Ma, Zhewei Yao, Amir Gholami, Michael W Mahoney, and Kurt Keutzer. Q-bert: Hessian based ultra low precision quantization of bert. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 8815-8821, 2020.
[36] Emma Strubell, Ananya Ganesh, and Andrew McCallum. Energy and policy considerations for deep learning in nip. arXiv preprint arXiv:1906.000 2019.
[37] Sainbayar Sukhbaatar, Edouard Grave, Piotr Bojanowski, and Armand Joulin. Adaptive attention span in transformers. arXiv preprint 1rXiv:1955007999 2019.
[38] Xiao Sun, Jungwook Choi, Chia-Yo Yu Chen, Naigang Wang. Swagath Venkataramani, Vijay- alakshmi Viji Srinivasan, Xiaodong Cui, Wei Zhang, and Kailash Gopalakrishnan. Hybrid 8-bit floating point (hip8) training and inference for deep neural networks. Advances in neural information processing systems, 32:4900-4900 2019.
[39] Zhiqing Sun, Hongkun Yu. Xiaodan Song, Renjie Liu, Yiming Yang, and Denny Zhou. Mo- bilebert: a compact agnostti bert for sssource-limiter devices. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages
2020.
[40] Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, and Da-Cheng Juan. Sparse sinkhorn attention, 2020.
[41] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. arXiv preprini arXiv:2009.06732. 2020.
[42] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. arXiv preprint arXiv: 7706007762
2017.
[43] Hanrui Wang, Zhanhaao Wa, Zhijian Liu, Han Cai, Ligeng Zhu, Chuang Gan, and Song Han. Hat: Hardware-awar transformers for efficient natural language processing. arXiv preprint arXiv:2005.144 14187, 2020.
12

---

[44] Linting Xue, Noah Constant, Adam Roberts, Mihir Kale, Rami Al-Rfou, Aditya Siddhant, Aditya Barua, and Colin Raffel, mt5: massively multilingual pre-trained text-to-text trans- former, 2020.
[45] Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. arXiv preprint vrXf::2007.14062. 2020.
[46] Jingqing Zhang, Yao Zhao, Mohammad Saleh, and Peter J. Liu. Pegasus: Pre-training with extracted gap-sentences for abstractive summarization, 2020.
[47] Aojun Zhou, Yukun Ma, Junnan Zhu, Jianbo Liu, Zhijie Zhang, Kun Yuan, Wenxiu Sun, and Hongsheng Li. Learning n:m fine-grained structured sparse neural networks from scratch. In International Conference on Learning Representations, 2021. URL https:// //opgr/i///ee.iiee.. KtweooruuunddKKoyyyygpe
13


Reviews:
[{"review_id": "eMuEsx4sg6W", "reviewer": "Reviewer_BocG", "summary": "Update: I've moved my score from 7 to 8.\n\nThe paper first introduces a sparse feedforward layer, which uses the input at every timestep to decide how to zero out the hidden vector (the one right in the middle) of the MLP feedforward layer. It uses the gumbel-softmax trick and zeros out (N-1) out of every N activations. This zeroing out of each hidden vector allows us to not use some of the rows & column of the W_1 and W_2 matrices of the feedforward layer. \n\nThe next method is the sparse QKV layer, where each input is compressed into a few smaller vectors that each do attention independently. The authors show that a convolutional layer can make this faster, without loss of accuracy, and also toss away the final output projection from their self-attention module.\n\nThe authors then apply a few off-the-shelf tricks to their transformer, such as reversible layers and replacing the enc/dec attention by attending to the concatenation of the enc/dec words.\n\nThe empirical results are very strong and show good performance on challenging tasks such as language modeling, summarization, and GLUE, while massively improving inference speed (training speed is not improved). \n\n", "questions": "", "limitations": "", "rating": 8, "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:\nThe paper is well written, and the ideas presented are relatively simple to implement. \nThe paper presents strong empirical results on hard tasks, with big inference speedups. \n\nWeaknesses:\nI might be wrong about this, and I would be grateful if the authors corrected me in the author response period, but I believe that the methods shown in this paper mostly improve the speed of *unbatched* inference. If this is true than the authors should point it out in the paper, and if not then they should also show speed number of batched inference. \nIf my assumption is correct it might slightly weaken the results of the paper, since for most models I believe batched inference is used in production. It's still a strong paper that I believe should be accepted, but this limitation should be pointed out, and maybe future papers will be able to address it. \n\n\nMissing Citations:\nLine 222: \"We therefore remove the encoder-decoder attention, but just concatenate the encoder representations before the decoder tokens\". Authors must cite \"Layer-Wise Coordination between Encoder and Decoder for Neural Machine Translation\", He et al, which did exactly that. \n\nNitpicks:\nLine 27/100/and throughout the paper: \"output layers\"- consider referring to these as the 'output projection'. I feel like calling them \"output layers\" might confuse some readers. \nLine 103: \" decoding speed is dominated by the execution cost of the feedforward block\". Clarify that this is true only for shorter input sequences (the authors do mention this later in the paper but I feel it's important to also mention it in line 103). \nTable 2: Consider adding the final perplexity of each model to the table (I know you also have it in the figure but it would be nice to have it in the table).", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"ethical_concerns": "None.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "YSsvMkG1DTz", "reviewer": "Reviewer_ju7M", "summary": "This work focuses on improving Transformer inference efficiency by redesigning the aforementioned architecture such that the matrix multiplications can be sparse at inference time. This sparsity is introduced in the dense layers that compute the query, key, and value, and in the feed-forward layer as well. This results in a much faster decoding speed when compared to dense transformers.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "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": "_Originality_: The methods proposed are novel, and are clearly distinguished from related work.\n\n_Quality_: The changes to the Transformer architecture that are proposed in this paper are interesting and very useful. They allow for much faster decoding. However, it's not clear if there are differences in training time due to additional blocks, like the \"Controller\". Can the authors clarify this? Furthermore, I would like to see the proposed models compared to more baselines, for example, the Reformer architecture (in performance and decoding time).\n\n_Clarity_: The paper is mostly well-written, although I found it a bit confusing at times, particularly in §3. For example, are C_1 and C_2 additional parameter matrices? The method described in §3.2 also deserves a more thorough explanation.\n\n_Significance_: As mentioned, this work proposes a Transformer variant in which the decoding speed is greatly increased, at little cost in performance. This result is very useful for practioners and researchers with varying ammount of computing resources, making it a valuable contribution to the community. The authors mention in the conclusion that the proposed technique can help to fine-tune a large Transformer \"a single workstation\". Can the authors clarify? Isn't your method only faster at inference time?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The potential negative societal impact is adequately addressed.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6h"}, "extra_text": {}}, {"review_id": "MlxS64Fse8a", "reviewer": "Reviewer_vUEu", "summary": "The paper proposes a variant of the Transformer layer that includes sparsified dense and attention layers. For FFN layers, the author introduces a controller layer that selects only a subset of hidden dimensions to be used in the computation. For the attention layer, The author proposes a multiplicative transformation to simulate the permutation of the input hidden states and applies a convolution layer on top of it. The author shows no accuracy drop with the proposed method and speedup in the decoding phase.\n\n", "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": "Overall, I believe the proposed method in the paper is practically useful and can speed up the Transformers. However, I still have some doubts about the motivation and rationale of the proposed methods. Detailed comments:\n\nStrengths\n- The paper is clear and easy to follow. I can understand the proposed method very easily.\n- The author did thorough experiments to verify the effectiveness of the method, including end-to-end experiments on multiple datasets and ablation studies of each method.\n\nWeaknesses\n- The method used for the FFN layer is similar to the previous mixture-of-expert method. The argmax selection (i.e. controller) part in the proposed method can be analogized to the expert selection part in mixture-of-expert. The difference that the authors mentioned in the related work doesn’t convince me: In mixture-of-expert, all of the experts are also being trained and only a subset of experts will be activated during the decoding for each token.\n- The convolution layer in the proposed attention layer doesn’t have any semantic meaning. The convolution over the sequence length dimension can be understood because neighboring words in a sentence often have a more semantic correlation. But the convolution over the S dimension doesn’t have this kind of semantic meaning. Please elaborate on the rationale for the convolution design choice here. In addition, the sequence length and S dimension are very different, why is the convolution window size for both dimensions set to be F?\nI would like to see more results on the speedup. What’s the speedup on the training speed? What’s the speedup with inference on a sequence or a batch of sequences?\n\nOther questions:\n- I haven’t seen people use negative log perplexity before. Why don't use log perplexity or just perplexity?\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Please refer to the weaknesses in main review section.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "HDUX8IdSclc", "reviewer": "Reviewer_zmH4", "summary": "The paper presents an approach to training large-scale sparse transformers by sparsifying the activations of linear layers in the network (feedforward and Q, K, V projections). The approach is particularly appealing because it doesn't appear to require any hardware-specific or low-level optimizations to be efficient since it sparsifies activations. A controller network produces a block-wise one-hot mask on the activations and weights for the subsequent layer are selected on the fly based on this mask. Gradients are propagated through this hard-selection operation via the straight-through Gumbel softmax estimator.\n\nThe authors are able to achieve fairly substantial inference time speedups via sparsification. They also apply the same approaches to sparsifying models that are tailored towards handling long sequences such as the Reformer with LSH attention.", "questions": "", "limitations": "", "rating": 6, "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. The sparsification technique is simple and easy to implement and seems to work at scale.\n2. The approach shows empirical results on multiple tasks and large-scale benchmarks.\n3. Significant improvements in inference times.\n4. The paper is well written and easy to understand.\n\nWeaknesses:\n\n1. The paper seems to lack comparisons to more traditional weight pruning-like strategies which also promise minimal loss in task performance while being faster at inference.\n2. It would be interesting to see ablations that push for slightly stronger sparsity constraints and how it affects task performance.\n3. While the proposed method doesn't require hardware-specific or low-level library optimizations, it appears that block sparsity or other kinds of structured sparsity may be optimized on specialized hardware in the near future such as in https://developer.nvidia.com/blog/exploiting-ampere-structured-sparsity-with-cusparselt/. A discussion about the compatibility of the kinds of block-wise activation sparsity used in this work and structured weight sparsity would be useful.\n\nOverall comments:\n\n- I'd encourage the authors to rethink the title of the paper since it isn't clear what \"Enough\" is in the context of model scaling or even along what axis you are scaling models. The title is potentially misleading into having readers think that it enables efficient training, while this is only helps at inference.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors do not discuss the limitations of their work and its potential negative societal impact. There is plenty of prior literature to engage with here since there has been plenty of discussions recently on the societal impact of large language models.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper introduces a technique for improving the inference speed of transformers by sparsifying all the linear layers. During training, a controller module predicts a sparse mask for the activations (trained with a straight-through Gumbel softmax), and during decoding this mask is used to prune the weights for the linear layer. Reviewers appreciated that the method is simple, and dramatically improves the speed of decoding. Reviewer vUEu raises reasonable concerns about the lack of comparison with mixture-of-experts approaches (which improve both training and inference speed), however the authors argue that their method is more expressive, and include an empirical comparison in their response. I think the title and abstract should be clearer about that the paper is addressing purely unbatched decoding speed. Overall, I recommend acceptance.
Decision: Accept (Poster)