Symmetric Dot-Product Attention for Efficient Training of BERT Language Models (2024)

Abstract: Initially introduced as a machine translation model, the Transformer architecture has now become the foundation for modern deep learning architecture, with applications in a wide range of fields, from computer vision to natural language processing. Nowadays, to tackle increasingly more complex tasks, Transformer-based models are stretched to enormous sizes, requiring increasingly larger training datasets, and unsustainable amount of compute resources. The ubiquitous nature of the Transformer and its core component, the attention mechanism, are thus prime targets for efficiency research.In this work, we propose an alternative compatibility function for the self-attention mechanism introduced by the Transformer architecture. This compatibility function exploits an overlap in the learned representation of the traditional scaled dot-product attention, leading to a symmetric with pairwise coefficient dot-product attention. When applied to the pre-training of BERT-like models, this new symmetric attention mechanism reaches a score of 79.36 on the GLUE benchmark against 78.74 for the traditional implementation, leads to a reduction of 6% in the number of trainable parameters, and reduces the number of training steps required before convergence by half.

## introduction since its introduction in 2017, the transformer architecture powered by its scaled dot-product attention mechanism @xcite has become the core component of modern deep-learning architectures and has enabled researchers to achieve breakthroughs in both natural language processing (nlp) and computer vision tasks such as language modelling @xcite , machine translation @xcite , speech processing @xcite , and image recognition @xcite . one of the many successes of the transformer lies in its ability to operate and learn in an unsupervised setting from unstructured textual data, as well as its ability to handle complex and varied structures such as graphs, images, and sentences by increasing the model's number of layers. however, this trend has led to the emergence of machine learning models so enormous that the gap between the amount of compute resources available to many research groups and the amount needed to stay competitive is increasing year after year @xcite , and by training larger and larger models, brought deep-learning's energy consumption to unsustainable amounts @xcite . efficient transformer implementations are a popular area of research with many recent contributions on encoding and dense representation of tokens @xcite , hardware-optimized implementation of attention @xcite , or implementations for long document processing @xcite . while the attention mechanism itself has been studied extensively @xcite , and several improvements to its computational complexity have been achieved @xcite , it is still primarily computed via the dot-product between a query and a key (see figure 1 ). @xcite highlight the difficulty of determining a proper compatibility function, and suggest that a more sophisticated compatibility function than dot product may be beneficial. in this work, we propose alternative compatibility functions for the attention mechanism, i. e., the scaled dot-product attention mechanism. with this approach, we aim to improve the training efficiency of transformer-based models and to reduce their resource consumption. we especially focus on the self-attention mechanism of bert @xcite , a transformer-based encoder model. our contributions can be summarized as follows: we demonstrate that our new attention formula reduces the number of parameters of the model by 6%, and achieves a reduction of the number of training steps required for model convergence by 50% without sacrificing accuracy (section 4). finally, we discuss the effects of our proposed compatibility function on training efficiency, and situate our approach within the context of research on efficient transformer-based models (section 5). ## improving the attention mechanism modern transformer-based models are neural networks that rely on the scaled dot-product attention mechanism introduced by @xcite . we propose two variations of this mechanism: a symmetric dot-product and a symmetric with pairwise factors dot-product, that lead to a reduction in the number of parameters of the self-attention layer. ## experiments to evaluate the symmetric and pairwise operators against the original operator, we train and evaluate several transformer-based encoder models, each using a different compatibility operator as part of the self-attention mechanism. the models are trained under the same conditions. first, we pre-train the models, because we want to measure the evaluation loss during training to see if our modifications have an impact on the training efficiency and the accuracy of the model. then, we evaluate each model on the glue benchmark @xcite to evaluate the model's accuracy on relevant downstream tasks, such as, sentence acceptability @xcite , sentiment analysis @xcite , sentence similarity @xcite , and natural language inference @xcite . finally, we select model checkpoints during training and evaluate those checkpoints on glue to measure the models' accuracy on downstream tasks during training. ## results in this section, we present the results of our experiments, the pre-training of our three variants (figure 2 ), the scores they reach on the glue benchmark ( ## related work while the transformer architecture @xcite popularized the use of the attention mechanism, and contributed to its adoption in the field of nlp, the attention mechanism was first introduced to nlp with recurrent neural networks applied to machine translation @xcite . in this setting, the compatibility operator is a simple multi-layer perceptron with non-linear activation operating on the concatenation of inputs encoded by the recurrent neural network. this definition of the attention mechanism was then extended to other compatibility operator: @xcite mention the use of the dot-product between the recurrent neural network's hidden state, propose to explicitly integrate token positions into the compatibility operator, and even suggest the use of a general dot-product operator score(h t , h s ) = h t t w h s . those initial influences have also been documented by @xcite and @xcite , where the general dot-product appears as a weighted dot-product between query and keys f (q, k) = q t w k. thus the pairwise compatibility operator we introduce is an evolution of the general dot-product, where we constrain it to a single and shared linear operator q before applying the bilinear form of matrix s, resulting in the following operator a(x, y) = q(x)sq(y) t . to the best of our knowledge, our work is the first application of the general dot-product with enforced symmetry to the self-attention mechanism of the transformer architecture. while we focused on the compatibility operator, recent improvements have been made on other parts of the attention mechanism. namely, @xcite proposed to simplify the entire transformer block by carefully removing components and achieved an impressive 15% weight reduction, while still relying on the traditional scaled dot-product. ## conclusions in this work, we revisited the traditional scaled dot-product used in the transformer self-attention mechanism. we challenged the use of two dis-tinct operators to compute the dot-product between queries and keys, in favor of single shared operator and a weighted dot-product with pairwise factors. by doing so, we enforced a symmetric structure to the compatibility operator of the attention mechanism, reducing the number of parameters used in the transformer layer by a third. as a result, when applied to bert models, our pairwise compatibility operator reduces the overall number of parameters of the model by 6%, reduces the number of pre-training steps required by half and improves accuracy on the glue benchmark, making transformer-based encoders more efficient, faster to train and lowering their resource requirements. we believe our work can be applied to other transformer architectures like decoder and encoder-decoder models, as well as to other nlp tasks like machine translation and language modeling. and, more generally, to the concept of attention as a whole, where it would bring improvement in other fields such as computer vision. for future work, we plan to evaluate the pairwise dot-product attention mechanism on larger models reaching into the billion parameters, and to evaluate our attention mechanism on other benchmarks, like superglue @xcite and squad2.0 @xcite . we plan on implementing the pairwise compatibility operator for the crossattention mechanism, and evaluating it on decoder and encoder-decoder tasks like language modeling and machine translation. finally, we want to evaluate our pairwise dot-product attention not only on natural language processing tasks, but also on tasks from other fields, computer vision, time series forecasting and reinforcement learning.