Improving Transformer Models by Reordering their Sublayers (2020)

Abstract: Multilayer transformer networks consist of interleaved self-attention and feedforward sublayers. Could ordering the sublayers in a different pattern lead to better performance? We generate randomly ordered transformers and train them with the language modeling objective. We observe that some of these models are able to achieve better performance than the interleaved baseline, and that those successful variants tend to have more self-attention at the bottom and more feedforward sublayers at the top. We propose a new transformer pattern that adheres to this property, the sandwich transformer, and show that it improves perplexity on multiple word-level and character-level language modeling benchmarks, at no cost in parameters, memory, or training time. However, the sandwich reordering pattern does not guarantee performance gains across every task, as we demonstrate on machine translation models. Instead, we suggest that further exploration of task-specific sublayer reorderings is needed in order to unlock additional gains.

## introduction the transformer layer @xcite is currently the primary modeling component in natural language processing, playing a lead role in recent innovations such as bert @xcite and gpt-2 @xcite . each transformer layer consists of a self-attention sublayer (s) followed by a feedforward sublayer (f), creating an interleaving pattern of self-attention and feedforward sublayers (sfsfsf • • • ) throughout a multilayer transformer model. to the best of our knowledge, there is no reason to expect this particular pattern to be optimal. we conduct a series of explorations to obtain insights about the nature of transformer orderings that work well, and based on this, we design a new transformer ordering pattern that improves upon the baseline. first, we generate random transformer models, varying the number of each type of sublayer, and their ordering, while keeping the number of parameters constant. we train these models on the standard wikitext-103 word-level language modeling benchmark @xcite , and observe that some of these random models outperform the original interleaved transformer model, even when the number of self-attention and feedforward layers is not equal. our analysis shows that models with more self-attention toward the bottom and more feedforward sublayers toward the top tend to perform better in general. based on this insight, we design a new family of transformer models that follow a distinct sublayer ordering pattern: sandwich transformers (figure 1 ). our experiments demonstrate that a sandwich transformer outperforms the baseline of @xcite . this result is made more interesting by the fact that our sandwich transformer is simply a reordering of the sublayers in the baseline model, and does not require more parameters, memory, or training time. model ppl f s f s f f f s f f s f s s s f f s f s s f s s s s f f s f f s 20.74 s f s s f f s f f f f s s s s f s f f f s f s f f s f s s s s f 20.64 f s f f s s f f s s s s f f s s s s s f f s f s s f s f f f f f 20.33 f s f f f f f f s s s f s s f f s f s s f f s f s s s f f s s s 20.27 f s s f f f f f f s f s s s f f f s s s s f f f s s s s f f s s 19.98 s s s f s s f s f f f f s s f s f s f s s s f f s f s f f f s f 19.92 f f f s f s s s f s f f s f s f f s f f s s s s s f f s s f f s 19.69 f f f s f f s s f f s s s f s s f s s s f f f f f s f s s s f s 19.54 s f s f s f s f s f s f s f s f s f s f s f s f s f s f s f s f 19.13 f s f f s s f s s f f f s s s s f f f s s s f f f f s f s s f s 19.08 s f s f f s s s s f f s s f f f f s s s f f s s s f s f f s f f 18.90 s f s f s f s f s f s f s f s f s f s f s f s f s f s f s f s f 18.83 s s s s s s s f f s f f s f s f s f f f f s f f f s f s s f f s 18.83 s f f s f s f f s f s s s f f s s f s s s s s s f f f f f f f s 18.77 s s s f s s f f s f s s f s f f s f f f s s f f s f s f f s s f 18.68 f f f s s s s s f f f s f s s s s f f s f s f s f s s f f s f f 18.64 s f f f s s s f s f s s f s s s s s f s s f f f f f s f f f s f 18.61 s s f f s s f s s s s f f f f f f s s f f s s s f s f f s s f f 18.60 f s f s s s s s f s f s f f f f f s f f f s f f s s f f s s s s 18.55 s f s f s f s f s f s f s f s f s f s f s f s f s f s f s f s f 18.54 s f s f s f s f s f s f s f s f s f s f s f s f s f s f s f s f 18.49 f s f s s s s s f s f f f s s f s f f s f s f s f s f f f f s s 18.38 s f s s f f s f s f s f f s s s s s f f f s s s f f f s f f s f 18.28 s f s f s f s f s f s f s f s f s f s f s f s f s f s f s f s f 18.25 s f s f s s f s s s f f s f s f s f s f f f f s s f f s f s s f 18.19 ## notation each transformer layer consists of a self-attention sublayer followed by a feedforward sublayer, modifying a sequence of vectors x 0 as follows: 2 stacking multiple transformer layers creates an interleaved network of sublayers. we denote these 2 we omit dropout @xcite and layer normalization @xcite to simplify the notation. models as strings, with s and f representing selfattention and feedforward sublayers, respectively. a three-layer transformer network, for example, would be denoted s fs f s f, with the flow of computation moving from input on the left to output on the right. thus, any string in the regular language (s|f) * defines a valid network that uses the same building blocks as the original transformer. for simplicity, we refer to these alternatives as transformers as well. ## random search we conduct a series of experiments to understand which transformer networks work well and whether particular architectural patterns can improve performance. first, we generate random transformer models while keeping the number of parameters constant. we then train these random models to determine whether the interleaving pattern (sfsfsf • • • ) is optimal (section 3.1), and whether balancing the number of self-attention and feedforward sublayers is desirable (section 3.2). finally, we analyze additional properties of these random models, and find that those with more selfattention at the beginning and more feedforward sublayers near the end tend to outperform the standard interleaved model (section 3.3). experimental setup our baseline is the strong transformer language model of @xcite , trained on wikitext-103 @xcite . wikitext-103 contains roughly 103 million tokens from english wikipedia, split into train, development, and test sets by article. the baevski model ppl s f f f s s f s f s f s s f f f f s f s f f s f f f f f f 22.80 s f f s s f s s s s s s s s s s s s s f s f s s s f s f f s s s f s s s f s 21.02 s s s s s s f f s f f f f s s f f f f f s s s f s f s s s s s s s s s 20.98 f f f f f f f f f s f f s s f f s f f s s s s f s f s s s f 20.75 f s s f s s s f f f f f f s s f s s s f s f f f s s s s f s f s s 20.43 s f f s f f f f f f s f s f s s f s s s f s f s f s s f s s f s 20.28 s f f s s f f s f f f s f s f s s s s f f f f f f s s s s f f 20.02 f s f f s f s s f f f f s f s f f f s f f f s s f f f s s s 19.93 s f f s f f s s f f s f s f f s s s f s s s s s f s s s f f f s s s 19.85 s s f f f f f f f s s f f f s s f s s f f s f s f s f f s f 19.82 s f s f s f f f s f f f s s f s f f f s f f s s f s f s f s s 19.77 s f s f f s s s f f s f f s s s f s s f f f f f s s s s f s s s f 19.55 s f f s f s s f f f s f f s f s s s s f s f s f f f f s f s s s 19.49 s f f f f s f f s s s s f s s s f s s f f f s s s f s s s s f s f s 19.47 f s s s f f s s s s s s f s f s f s f f s f f f f s s f s f s s s s 19.25 s f s f s f s f s f s f s f s f s f s f s f s f s f s f s f s f 19.13 f s s s s s s f s f s f s f f f s f s s s f s s f f s s s s f s f f 18.86 s f s f s f s f s f s f s f s f s f s f s f s f s f s f s f s f 18.83 s s f s f s s s f s s s s s f f s f s f s s s f s s f s f s s s s s s s f 18.62 s f s f s f s f s f s f s f s f s f s f s f s f s f s f s f s f 18.54 s f s f s f s f s f s f s f s f s f s f s f s f s f s f s f s f 18.49 s s s f s f f s f s s f s s s f f s f f f f f f s s f s f f f 18.34 s s s f s f s f f s s s f s f f f f f s f s f f f f s s s f f 18.31 s f s f s f s f s f s f s f s f s f s f s f s f s f s f s f s f 18.25 s s s s s s f s s s f f f f s f s f f f f f f f f f f f s f 18.12 and auli model contains 16 transformer layers of @xmath0, with 16 heads in each self-attention sublayer, and feedforward sublayers with an inner dimension of 4096. in this setting, each self-attention sublayer contains 4d 2 parameters, while each feedforward sublayer contains 8d 2 parameters (excluding bias terms, which have a marginal contribution). thus, each f sublayer contains twice the parameters of a s sublayer, following the parameter ratio between self-attention and feedforward sublayers described in @xcite . all of our experiments use the same hyperparameters as baevski and auli's original model. to set an accurate baseline, we train the baseline model (the standard interleaved transformer) with five different random seeds, achieving 18.65 ± 0.24 perplexity on the development set. ## designing a better transformer our analysis in the previous section motivates designing a transformer model that is heavy on selfattention at the bottom and feedforward sublayers at the top, while at the same time containing a more-or-less balanced amount of both sublayer types. as a first attempt to manually design a better transformer, we take this hypothesis to the extreme, and train a transformer model of 16 self-attention sublayers followed by 16 feedforward sublayers (s 16 f 16 ). this model achieves 18.82 perplexity, which is comparable to the performance of the baseline with the same number of parameters. we next generalize this model and the original interleaved transformer, creating the family of sandwich transformers. a sandwich n k transformer consists of 2n sublayers in total (n of each type), conforming to the regular expression s k (sf) n-k f k . the first k sublayers are purely self-attention (s), while the last k are feedforward sublayers (f). in between, we use the original interleaving pattern (sf) to fill the remaining 2(n-k) sublayers. when @xmath1, we get the original transformer model, and when @xmath2. we refer to k as the transformer's sandwich coefficient. we train sandwich transformers for @xmath3, . . . , 15}. figure 5 shows the transformer's performance as a function of the sandwich coefficient k. with the exception of @xmath4, 15, all sandwich transformers achieve lower perplexities ## one reordering to rule them all? results show that as we drift away from our original setting, sandwich transformers provide diminishing gains, but always perform at least as well as the baseline transformers (provided that the sandwich coefficient is properly tuned). this finding suggests that different settings may benefit from different sublayer reordering patterns. ## analysis at the time of writing, we do not have an explanation for why sublayer reordering improves performance on language modeling. however, we are able to determine that sandwich transformers spread their attention in a different fashion than interleaved models. we analyze two baseline models and two sandwich 16 6 models trained with different seeds on the wikitext-103 dataset, by first recording the attention values that each token's heads assign to all other tokens during inference on the validation set. given the attention outputs of two models, we then compute the models' attention distance for each token, and for each self-attention sublayer. this metric compares the attention distribution in the ith self-attention sublayer of the first model to that of the ith self-attention sublayer of the second model, for a specific token. ## conclusion we train random transformer models with reordered sublayers, and find that some perform better than the baseline interleaved transformer in language modeling. we observe that, on average, better models contain more self-attention sublayers at the bottom and more feedforward sublayer at the top. this leads us to design a new transformer stack, the sandwich transformer, which significantly improves performance over the baseline at no cost in parameters, memory, or runtime. we then show that the sandwich ordering also improves language modeling performance on a different word-level language modeling benchmark, and that the sandwich pattern can be used to achieve state of the art results on character-level language modeling. although sandwich ordering does not improve translation models, we show that they are robust to layer order changes, and that even extreme reorderings (all attention sublayers at the bottom, and all the feedforward sublayers at the top) perform as well as the baseline. sublayer reordering can improve the performance of transformer models, but an ordering that improves models on one group of tasks (word/character-level language modeling) might not improve the performance on another task. by showing that sublayer ordering can improve models at no extra cost, we hope that future research continues this line of work by looking into optimal sublayer ordering for other tasks, such as translation, question answering, and classification.