One Wide Feedforward Is All You Need (2023)

Abstract: The Transformer architecture has two main non-embedding components: Attention and the Feed Forward Network (FFN). Attention captures interdependencies between words regardless of their position, while the FFN non-linearly transforms each input token independently. In this work we explore the role of the FFN, and find that despite taking up a significant fraction of the model’s parameters, it is highly redundant. Concretely, we are able to substantially reduce the number of parameters with only a modest drop in accuracy by removing the FFN on the decoder layers and sharing a single FFN across the encoder. Finally we scale this architecture back to its original size by increasing the hidden dimension of the shared FFN, achieving substantial gains in both accuracy and latency with respect to the original Transformer Big.

## introduction the transformer architecture @xcite has become the de facto paradigm in many natural language processing (nlp) tasks, including machine translation (mt). several studies have shown that transformers exhibit impressive scaling-law properties @xcite @xcite , wherein increasing the number of model parameters leads to further accuracy gains. in parallel with this architecture's impressive scaling of the numbers of parameters @xcite , there is a growing trend towards reducing model footprints for real-world deployment, to satisfy practical constraints like latency requirements as well as memory and disk space limitations. in turn, researchers are actively exploring parameter sharing @xcite @xcite , reducing the dimensionality of transformer compo-nents, and pruning components like attention heads @xcite . although the role of attention in learning pairwise dependencies between tokens is relatively well understood @xcite @xcite , the role of the feed forward network (ffn) remains under-explored. recently, @xcite established a connection between the ffn and attention by positing that the ffn corresponds to learnable key-value pairs where the weights of the first layer of the ffn corresponds to the keys and those of the second to the values. they find that the keys are able to capture salient textual patterns at each layer, and they notice that the classes of patterns tend to overlap between neighboring layers, indicating redundancy in the representation. this observation motivates our work, where we revisit the conventional practice of allocating an individual ffn per layer. we investigate the effect of sharing and dropping the ffn across different layers on mt models. we conduct thorough experiments with different configurations of the transformer, across different language pairs, including a low resource language pair and multilingual. in addition, we investigate the effect of the ffn in a decoder-only transformer-based model. we find that a considerable level of redundancy exists between the encoder and decoder ffns. as a result, we are able to eliminate the decoder ffn and share a single ffn across the encoder without significantly compromising the model's accuracy. this step leads not only to significant parameter savings but also opens up opportunities for further improvements. we also suggest using wider ffns in the encoder while dropping the decoder's ffn, which results in a model with a similar size, but improved accuracy and reduced latency. finally we conduct a fine-grained analysis of the representational similarity between the original model, using one independent ffn per layer, and various models with shared ffns. our results reveal that both model accuracy and the internal representation of transformer blocks remain stable when sharing the ffn. ## experimental setup data in our experiments, we show results on wmt22 english (en) → german (de) (296m pairs), which we obtained using the provided mt-data scripts foot_3 , wmt16 @xmath0, and for the multilingual setup of @xcite , consisting of 10 languages: german, english, spanish, french, italian, japanese, korean, portuguese, swahili, and chinese. in our analysis, we mostly focus on wmt22 @xmath1. following @xcite , we use wmt'16 provided scripts to normalize the ro side. @xmath2. for more details refer to @xcite . for the multilingual experiments, we replicated the setup of @xcite , which includes all details, including data preprocessing and dataset sizes. latency we report inference time in tokens/second (the higher, the better), averaged over 5 runs. for the multilingual models, we use the @xmath3. our measurements were collected using a single nvidia v100 gpu on a single-threaded intel(r) xeon(r) gold 6148 cpu @ 2.40ghz with batch size of 1 and beam size of 5, in order to realistically mimic the inference of a deployed model. for experiments with larger batch sizes, see appendix d. ## related work weight pruning and parameter sharing are wellknown techniques to reduce a model's footprint. given the scale of the latest models @xcite , there have been multiple efforts to prune neurons based on different automatic methods @xcite @xcite , sharing parameters efficiently @xcite , and factorizing certain components @xcite . neuron pruning methods often focus on finding and pruning redundant neurons through correlation methods @xcite , but also on how transformer components like the multi-head attention can be pruned significantly due to model redundancy in the encoder or decoder either by checking the gradients salience @xcite or a differentiable relaxation of the l 0 regularization at training time @xcite . for parameter sharing, the universal transformer @xcite proposed a model where all layers are shared (i.e., in effect it reduced the model to a single shared layer). @xcite proposes finding an optimal configuration of shared layers in the encoder or decoder through different methods of sharing (in sequence, in cycle, or in reversed cycle) always keeping a specified number of final layers foot_7 . similarly, @xcite proposes an approach where just the middle layers are shared, while the bottom and top layers are independent, and using a lower dimensionality for the embedding layer. analogously, @xcite focus on minimizing the number of parameters and the number of calls to each parameters' group in order to optimise on-device models. they achieve this by sharing the encoder and decoder in a similar way to both previous methods, particularly by sharing all layer parameters in cycle like @xcite . previous works also focus on reducing the dimensionality of certain parameters, mostly through low rank factorization. ## conclusion in this work, we studied the importance of the ffn in transformer models. we analyzed the impact of removing and/or sharing the ffn across layers and found that, due to this component's redundancy, the model sizes can be substantially reduced with little impact on accuracy for machine translation. in particular, we found that sharing the ffn across all encoder layers while making it larger and removing it from the decoder layers leads to models that are more accurate and faster at inference. our findings are applicable across multiple settings, including decoder-only and multilingual models. in a low-resource setting the results are modest but our approach can still recover the base-line's performance with a faster inference. finally, we conducted a thorough similarity analysis between the vanilla transformer and our proposed architectures, and found that the latter's internal representations do not differ significantly from the former's, except in that they are less redundant.