When Attention Meets Fast Recurrence: Training Language Models with Reduced Compute (2021)

Abstract: Large language models have become increasingly difficult to train because of the growing computation time and cost. In this work, we present SRU++, a highly-efficient architecture that combines fast recurrence and attention for sequence modeling. SRU++ exhibits strong modeling capacity and training efficiency. On standard language modeling tasks such as Enwik8, Wiki-103 and Billion Word datasets, our model obtains better bits-per-character and perplexity while using 3x-10x less training cost compared to top-performing Transformer models. For instance, our model achieves a state-of-the-art result on the Enwik8 dataset using 1.6 days of training on an 8-GPU machine. We further demonstrate that SRU++ requires minimal attention for near state-of-the-art performance. Our results suggest jointly leveraging fast recurrence with little attention as a promising direction for accelerating model training and inference.

## introduction many recent advances in language modeling have come from leveraging ever larger datasets and model architectures. as a result, the associated computation cost for developing such models have grown enormously, requiring hundreds of gpu hours or days per experiment, and raising concerns about the environmental sustainability of current research @xcite . as a consequence, it has become imperative to build computationally efficient models that retain top modeling power while reducing computational costs. the transformer architecture @xcite was proposed to accelerate model training and has become the predominant architecture in nlp. specifically, it is built entirely upon selfattention and avoids the use of recurrence to enable strong parallelization. while this change has 1 our code, experimental setup and models are available at https://github.com/asappresearch/sru . led to many empirical success and improved computational efficiency, we are interested in revisiting the architectural question: is attention all we need for modeling? the attention mechanism permits learning dependencies between any parts of the input, making it an extremely powerful neural component in many machine learning applications @xcite . we hypothesize that this advantage can still be complemented with other computation that is directly designed for sequential modeling. indeed, several recent works have studied and confirmed the same hypothesis by leveraging recurrence in conjunction with attention. for example, @xcite demonstrates that single-headed attention lstms can produce results competitive to transformer models in language modeling. other work have incorporated rnns into transformer, and obtain better results in machine translation @xcite and language understanding benchmarks @xcite . these results highlight one possibility -we could build more efficient models by combining attention and fast recurrent networks @xcite . in this work, we validate this idea and present a self-attentive recurrent unit that achieves strong computational efficiency. our work builds upon the sru @xcite , a highly parallelizable rnn implementation that has been shown effective in language and speech applications @xcite @xcite . we incorporate attention into the sru by simply replacing the linear transformation of input with a self-attention component. the proposed architecture, called sru++, enjoys enhanced modeling capacity and remains equally parallelizable. figure 1 compares its performance with the transformer-xl model @xcite on the enwik8 dataset. sru++ achieves better results while using a fraction of the training resources needed by the baseline. we evaluate sru++ on standard language modeling benchmarks including the enwik8, wiki-103 and billion word datasets. sru++ consistently outperforms various transformer models on these datasets, delivering better or on par results while using 3x-10x less computation. our model do not use positional encoding, multi-head attention and other techniques useful to transformer models. furthermore, we demonstrate that a couple of attention layers are sufficient for sru++ to obtain near state-of-the-art performance. these changes not only highlight the effectiveness of recurrence but also enable strong computation reduction in training and inference. finally, we also showcase the effectiveness of sru++ on the iwslt'14 @xmath0, and open source our implementation in pytorch to facilitate future research. ## background: sru we first describe the simple recurrent unit (sru) in this section. a single layer of sru involves the following computation: two important code-level optimizations are performed to enhance the parallelism and speed of sru. first, given the input sequence @xmath1, • • • , x[l]} where each x[t] ∈ r d is a ddimensional vector, sru combines the three matrix multiplications across all time steps as a single multiplication. this significantly improves the computation intensity (e.g. gpu utilization). specifically, the batched multiplication is a linear projection of the input tensor @xmath2: the second optimization performs all elementwise operations in an efficient way. this involves similar to other built-in operations such as attention and cudnn lstm @xcite , sru implements all these operations as a single cuda kernel to accelerate computation. note that each dimension of the hidden vectors is independent once u is computed. the computation can run in parallel across each hidden dimension (and each input sequence given a mini-batch of multiple sequences). ## sru++ the key modification of sru++ is to incorporate more expressive non-linear operations into the recurrent network. note that the computation of u (equation 1 ) is a linear transformation of the input sequence x. we can replace this linear transformation with self-attention operation to enhance modeling capacity. specifically, given the input sequence represented as a matrix @xmath3, the attention component computes the query, key and value representations using the following multiplications, next, we compute a weighted average output @xmath4, the final output u required by the elementwise recurrence is obtained by another linear projection, this implementation is post-layer normalization in which the normalization is added after the residual connection. alternatively, pre-layer normalization @xcite only applies to the nonlinear transformation. while pre-normalization tends to be less sensitive to different learning rates, we use post-normalization for better results following the observations in @xcite . we analyze the effectiveness of layer normalization in appendix a.2. model batch size b × m bpc ↓ trans-xl 24×512 1.06 sru++ 24×512 1.03 sru++ 16×768 1.02 ## experimental setup datasets we evaluate our model on four standard nlp benchmarks. we follow the same setup of @xcite and other previous work. the dataset uses a shared vocabulary of 14k bpe tokens. models all our language models are constructed with a word embedding layer, multiple layers of sru++ and an output linear layer followed by softmax operation. we use single-head attention in each layer and 10 sru++ layers for all our models. we use the same dropout probability for all layers and tune this value according to the model size and the results on the dev set. by default, we set the hidden dimension d : @xmath5: 1. we report additional analysis and tune this ratio for best results in section 5 and appendix a. for simplicity, sru++ does not use recent techniques that are shown useful to transformer such as multi-head attention, compressed memory @xcite , relative position @xcite , nearest-neighbor interpolation @xcite and attention variants to handle very long context @xcite . we compare with previous transformer models that incorporate one or several these techniques. however, we do not compare with results that use additional data or dynamic evaluation @xcite , for a fair comparison between all models. optimization we use radam @xcite with the default β values as our optimizer. radam is a variant of adam optimizer @xcite that is reported less sensitive to the choice of learning rate and warmup steps while achieving similar results at the end. we use a fixed weight decay of 0.1 and an initial learning rate of 0.0003 in our experiments. these values are selected based on enwik8 dev set and used for other tasks. see appendix a.3 for more details. we use a cosine learning rate schedule following @xcite . we do not change the initial learning rate unless otherwise specified. see appendix b for the detailed training configuration of each model. each training batch contains b sequences (i.e. the batch size) and m consecutive tokens for each sequence (i.e. the unroll size), which gives an effective size of b × m tokens per batch. following standard practice, the previous training batch is provided as additional context for attention, which results in a maximum attention length of 2 × m . for enwik8 and wiki-103 datasets, the training data is partitioned into b chunks by concatenating articles and ignoring the boundaries between articles. we train a 10-layer sru++ model with 42m parameters. for a fair comparison, we use the same hyperparameter setting including the effective batch size, attention context length, learning rate and the number of training iterations as the transformer-xl base model. notably, our base model can be trained using 2 gpus due to less gpu memory usage. after training, we set the at- how much attention is needed? @xcite demonstrated that using a single attention layer with lstm retains most of the modeling capacity compared to using multiple attention layers. we conduct a similar analysis to understand how much attention is needed in sru++. to do so, we only enable attention every k layers. the layers without attention become the variant with dimension projection illustrated in figure 2 (b). note that @xmath6, and @xmath7. ## conclusion we present a highly-efficient architecture combining fast recurrence and attention, and evaluate its effectiveness on various language modeling datasets. we demonstrate fast rnns with little attention not only achieve top results but also reduce training cost significantly. our work shares a different idea to accelerating attention, therefore providing an orthogonal direction to advancing stateof-the-art model architecture. as future work, we believe the model can be improved using stronger attention or recurrent implementations, better normalization or optimization techniques.