Not all parameters are born equal: Attention is mostly what you need (2021)

Abstract: Transformers are widely used in state-of-the-art machine translation, but the key to their success is still unknown. To gain insight into this, we consider three groups of parameters: embeddings, attention, and Feed-Forward Neural network (FFN) layers. We examine the relative importance of each by performing an ablation study where we initialise them at random and freeze them, so that their weights do not change over the course of the training. Through this, we show that the attention and FFN are equally important and fulfil the same functionality in a model. We show that the decision about whether a component is frozen or allowed to train is at least as important for the final model performance as its number of parameters. At the same time, the number of parameters alone is not indicative of a component’s importance. Finally, while the embedding layer is the least essential for machine translation tasks, it is the most important component for language modelling tasks.

## introduction bizarrely, we find that a transformer model @xcite for machine translation loses just 1.1-2% bleu when any one of the three main components (embeddings, attention, and feedforward) is frozen with randomly initialised parameters. in more extreme cases we see that with just 6% of trainable parameters, we maintain 72% of the baseline bleu. this suggests that redundancy in the model architecture is able to compensate for poor components. we perform a case study on transformer style models, exploring the built-in redundancy and the limits of what can be achieved when the majority of parameters are frozen after random initialisation. we do this in several different transformer architecture presets for neural machine translation: transformer-big, transformerbase and a transformer-tiny student preset, and one preset for a language modelling task. our key findings are: our study sheds light on the inner workings of the transformer neural network architecture and examines in detail the relative contribution of each component towards the final performance. ## methodology we perform a parameter-freezing case study on transformer style models. a transformer consists of three basic components: an embeddings layer, attention layers, and ffn, as shown in figure 1 . the embeddings encompass the source, target embeddings and the output layer as they are tied @xcite . we initialise the model with the glorot distribution @xcite and pick a component (or components) to mark as untrainable or "frozen". 1 this effectively means that during training the model will benefit from the extra transformations that those parameters provide, but they are random. it is up to the remaining components to accommodate them. freezing parameters is frequently performed in transfer learning @xcite @xcite in order to prevent the child model from deviating in an undesired manner from the parent. through these experiments we identify which components are essential for the model, which are interchangeable and which are complementary. an essential trainable component by itself can achieve high performance with the other components being random, while a non-essential component is not able to achieve good performance by itself. finally, we exchange the frozen random initialisation of components with frozen diagonal (identity) initialisation to verify how important the transformations of the frozen components are. in the next few sections, we first train a baseline model with no frozen components. then we train 3 models with one frozen component: either the embedding layer, the ffn layer or the attention. finally we perform experiments where two out of the three components are frozen. ## transformer-big experiments for our main experiments we used 1m parallel turkish-english sentences from the wmt 18 news translation task @xcite ). data was preprocessed and tokenised using moses @xcite and a shared vocabulary of 36k byte-pair encoding (bpe) tokens @xcite . for training the model, we used the marian @xcite implementation of the transformer-big architecture @xcite with tied embeddings and 8 attention heads for a total of 213m parameters. this model, and all subsequent were trained until 10 consecutive stalls on the perplexity of the dev set or 50 component parameter ratio emb att ffn bleu epochs trainable/all (0) 24.3 19 1 one frozen component (1) 22.6 26 .82 (2) 22.3 23 .64 (3) 23.2 26 .52 multiple frozen components (4) 21.5 36 .35 (5) 20.8 37 .47 (6) 4.4 25 .17 freezing a subset of the model leads to a tiny increase in the training speed in terms of words processed per second for all models, due to the reduced computational cost when performing the backward pass. however, this speed benefit is negated by the extra epochs required to achieve convergence, to the point that every single model was slower to converge in terms of real time. the increased training time is likely due to the optimiser running into more saddle points, as there are less parameters that can be pushed into a better model direction @xcite . freezing components also always leads to reduced performance in terms of bleu compared to the baseline, which is expected and is also observed when transferring pretrained components @xcite @xcite . our results show that when there is one frozen component, all models perform very similarly, with the frozen ffn one (3) being slightly better than the other two. however, when we look at the experiments where we have multiple frozen components, it is evident that the attention and the ffn layer are more essential for the overall model performance: if both of them are initialised randomly and left frozen, the model fails to learn completely (6). based on this we can conclude that the embedding layer is the least essential to be learned in the context of the machine translation task and the remaining components can easily learn to work with random embeddings. this result confirms the findings of @xcite , who show that transferring the embedding layer during transfer learning is the least important for the overall model performance. in the rest of the transformer-big experiments, we omit the system with just the embedding layer trainable, as the model quality in this setting is very poor. we also note that a model with trainable embeddings and attention (3) outperforms a model with frozen embeddings (1), despite the ffn and attention being more essential to be learned. this suggests that the embeddings perform a different role from the attention and the ffn in the translation process. based on the similar performance that the trainable attention and the ffn offer, we hypothesise that they are able to encode the same information and in case that one of them is not trainable, the other one can mostly fulfil its role. our hypothesis is supported by the work of @xcite , who show that feed forward transformer layers are key-value memories. on the other hand, trainable embeddings, while least essential for the model performance likely learn complementary, albeit non-essential, information to the attention and the ffn. we find a weak correlation between the number of parameters of a component and its importance to the model: the embeddings layer has the least amount of parameters and is the least essential to be trained, however we can't rank the other two components: the ffn has considerably more parameters than the attention, but they both seem to be equally important for the model's performance. we note that the model with just attention trainable (4) is able to reach competitive performance, despite having just 35% of the model parameters. ## experiments with smaller models the transformer-big preset used in the previous experiments is an overly expressive neural network with enough parameters, so that even when preventing a big part of them from being trained, the model likely has enough built-in redundancy to make up for it. in order to avoid this assumption, we performed analogous experiments with a transformerbase model, as well as a knowledge-distilled student transformer model. these experiments aim to reduce the model size in a systematic way in order to see whether our findings in the previous sections are dependent on the overall model size. ## language model experiments in order to verify whether our findings generalise to other tasks, we performed a language modelling experiment. we trained a transformer language model, using 78m sentences taken from newscrawl 2017, 2018 and 2019. our model is based on the transformer-base preset @xcite with 32k bpe tokens vocabulary and 38m parameters. we show our results in a change in the task can lead to a big change in the overall behaviour of components. for this task, the embedding layer loses the majority of its effectiveness if it is not trainable, unlike in the translation experiments. ## randomness and transformations in order to measure how much the transformations provided by the random parameters help the trainable parameters, we performed an experiment where we initialise either the attention or the ffn layer of a transformer-big model to a diagonal matrix, with zeroes everywhere except at the main diagonal. for the attention this means that all three attention matrices in self and context attention are diagonal and just copy over their input to the output, however since attention is interpreted as different heads, it does mean that the individual heads are very sparse and attend only to a single token. in the case of the ffn layer, since it is not square, the upper part of the matrix is diagonal and the lower part is just zeroes. 2 we present our results on transformer-big model from section 3 in our results are conflicting. when the attention is diagonal (3), the bleu score drops significantly from when it is initialised to random (3.1). this suggests that the model indeed makes use of and learns to work with random parameters and the transformations they provide, however it can not adapt to work with diagonal parameter matrices, as they don't provide useful input. in the case of the ffn layer, we see the opposite results. making the ffn layer matrices have 1 only across their main diagonal (2.1) seems to marginally improve performance compared to the baseline case (2). since the ffn layer is bypassed by residual connections which connect the attention to the output layer, we don't lose any information from the attention, and the ones across the main diagonal of the ffn allows for some minimal mixing of the output of the attention. in this case we can clearly see that the random transformations provided by the untrainable ffn are not as important, as the attention can do its job without them. we conclude that in certain cases the model does benefit from the transformation provided by the random parameters, however this does not generalise to every component of a transformer. we did not perform this experiment with the embedding layer, as it is not square (the shape is (32000,512) or similar) and it is used to encode different vocabulary items. rendering it diagonal across the main diagonal would leave the majority of the vocabulary items encoded as a 512 zero vector, which is not something that the model can learn from. ## discussion we find that depending on the task, different components of a transformer are essential. language modelling tasks require the embeddings to be trained, while translation can mostly cope without them. a large amount of research has focused on using bert @xcite or elmo @xcite pre-trained embeddings in downstream tasks, but as shown here, embedding representation is not as important as embedding width in a translation task. furthermore, there is evidence that in some downstream tasks, randomised language models perform competitively @xcite . we urge researchers to always evaluate how essential a component is by setting it to random and frozen, before performing transfer learning from pretrained components. we find that the "width" of a component is at least as important for the final model performance as whether it is trainable or not, but the raw parameter size of a component is not indicative of whether it is essential to have it trained or not. this clearly suggests that when there is a frozen and random component of a model, the rest of the components are able to work around just as if it were an extra activation function. when the attention is rendered completely linear by making it diagonal, we show a more significant drop in the performance compared to when this is done to the ffn layer. this suggests that the transformation that happens in the attention is more important than the transformation that happens in the ffn layer. we find that the smaller the model is, the more sensitive it is to having some of its components nontrainable. we suspect that transformer-big models have more built-in redundancy into them and adapt better to this situation. curiously, the less parameters a model has, or the less trainable parameters the model has, the more epochs it takes to train. this is especially visible in the experiments of transformer-base ( finally we note how similarly the attention and ffn layers behave: when the attention layer is small, the ffn layer is more essential and vice versa. in a setting where they are trainable by themselves, they achieve similar performance, with the attention having slight edge in larger models. if we decrease the size of the embeddings layer, or use a transformer-base preset, the performance gap between the ffn and the attention widens in favour of the attention. this echoes the findings of @xcite and suggests that the attention and the ffn layer are interchangeable and more essential for the model's performance, however there is a degree of redundancy between them. ## related work concurrently and independently from us, @xcite perform an ablation study to show how gradually setting different components of an encoderdecoder mt system from random and frozen to trainable increases the bleu scores from 4 to 39 for an english-french model. unlike garg et al. we use a transformer model and we show that wider and deeper models behave differently from small and narrow models. furthermore we show that depending on the task, the model benefits the most from different components. our work follows a long line of research that explores the performance of random (echo state) neural networks in various tasks @xcite @xcite . @xcite ; @xcite explore random encoders for sentences and evaluate their performance in downstream tasks. ## conclusion and future work we performed a study of transformer neural machine translation component freezing that identified the attention and the ffn as more essential to be learned, but also interchangeable. the embedding layer can not perform the translation task by itself, but the information it learns is complementary to either the attention or the ffn layer. our findings do not necessarily generalise over other tasks: in the language modelling setting, the components behave differently. our findings could serve as the basis for research into new transformer machine translation architectures for production @xcite . we have shown that models can retain competitive performance even when the ffn layer, which accounts for up to 48% of the parameters, is just a randomly initialised glorot distribution. this opens the door for models where part of the parameters are generated on the fly during decoding, as opposed to stored statically on disk, making it particularly suitable to embedded devices with limited memory.