Rethinking Softmax: Self-Attention with Polynomial Activations
Abstract
This paper challenges the conventional belief that softmax attention in transformers is effective primarily because it generates a probability distribution for attention allocation. Instead, we theoretically show that its success lies in its ability to implicitly regularize the Frobenius norm of the attention matrix during training. We then explore alternative activations that regularize the Frobenius norm of the attention matrix, demonstrating that certain polynomial activations can achieve this effect, making them suitable for attention-based architectures. Empirical results indicate these activations perform comparably or better than softmax across various computer vision and language tasks, suggesting new possibilities for attention mechanisms beyond softmax.
1 Introduction
Transformer architectures (Vaswani et al., 2017) have become the state-of-the-art model architecture for diverse areas such as natural language processing (NLP) (Vaswani et al., 2017; Devlin et al., 2018; Zhuang et al., 2021; Zhen et al., 2022), computer vision (Dosovitskiy et al., 2020; Carion et al., 2020; Liu et al., 2021; Touvron et al., 2021), and robotics (Fu et al., 2024; Maiti et al., 2023; Salzmann et al., 2020). A key component in the transformer architecture is the softmax attention block, enabling transformers to evaluate the importance of individual input elements during output generation. This feature facilitates an efficient method to attend to diverse input elements throughout training, allowing transformers to effectively capture spatial dependencies within sequential data. Unlike traditional recurrent neural networks (RNNs) and convolutional neural networks (CNNs), transformers exhibit the ability to scale to large datasets without a significant degradation in performance. This characteristic has made them an ideal architecture for handling large-scale machine learning tasks.
Softmax is widely recognized for its effectiveness in attention mechanisms due to its ability to produce an attention matrix that meets three key conditions (Vaswani et al., 2017; Dosovitskiy et al., 2020; Zhen et al., 2022; AUEB et al., 2016): (i) non-negativity, (ii) rows that are normalized to sum to 1 and (iii) sparsity. The general consensus is that non-negativity guarantees that attention weights remain positive, facilitating the model in assigning significance to various input elements. The normalization constraint ensures that the attention weights for all input elements collectively sum to , rendering the weights interpretable as probabilities. Additionally, sparsity aids the model in focusing on a select few key elements, thereby enhancing efficiency. It has been argued that these properties are crucial for enabling the attention mechanism to attend to pertinent segments of the input sequence while efficiently filtering out irrelevant details. However, this approach to attention has become somewhat axiomatic as it is mostly motivated by empirical results with little theoretical foundation. Despite the exploration of alternative activations in several studies (Shen et al., 2023; Fang et al., 2022; Correia et al., 2019), softmax attention continues to dominate, largely due to its interpretability.
In this paper, we question this view by proposing that the effectiveness of softmax stems from its implicit regularization of the Frobenius norm of the attention matrix during training, preventing attention weights from becoming excessively large or small. We then derive a theoretical framework that produces polynomial activations that deliberately violate one or more of the three conditions mentioned earlier, yet are able to regularize the Frobenius norm of the attention weights during training. Our findings demonstrate that such activations can achieve comparable or even superior performance to softmax across various vision and natural language processing (NLP) tasks, even though they seem to violate our understanding of attention.
We advise the reader that this paper diverges from the usual pursuit of creating cutting-edge transformer architectures for achieving state-of-the-art results on benchmark datasets. Instead, our focus is on critically examining softmax attention to determine whether its effectiveness is a result of true interpretability or a more nuanced regularization mechanism. By questioning established views, we aim to uncover deeper insights into transformer architectures that could lead to broader applications and improved understanding. Nonetheless, we validate our theory on multiple transformer-based tasks, including image classification, segmentation, object detection, and NLP, often achieving results that match or surpass those of softmax attention. Our main contributions are:
-
1.
We question the widely accepted notion that softmax’s effectiveness in attention mechanisms is solely due to its ability to produce normalized sparse attention weights. Instead, we theoretically show that softmax has a regularization effect on attention and argue this plays a crucial role in its success.
-
2.
We explore activations that deliberately deviate from traditional softmax attention conditions. These activations are found to regularize the Frobenius norm of the attention matrix during training, akin to softmax, and demonstrate comparable or superior performance across various vision and NLP tasks.
2 Related Work
Several studies have explored alternative activations for attention mechanisms in transformers. Shen et al. (2023) investigated ReLU activations, finding them to outperform softmax in tasks with long sequences, such as document translation. Banerjee et al. (2020) examined Taylor series approximations to softmax, which showed superior performance to softmax in image classification. Wang et al. (2021) proposed periodic alternatives to softmax, designed to provide better gradients for attention mechanisms and achieved better results than softmax on simple networks for image classification. Koohpayegani & Pirsiavash (2024) demonstrated that applying normalization to linear attention mechanisms can achieve performance comparable to softmax. Our work differs from all these in that we identify a clear theoretical relationship between the scale of the Frobenius norm of the self-attention matrix and the input sequence length. Using this insight to derive potential acitvations that can perform on par with softmax.
3 Preliminaries and Notation
In this section we outline the definition of a transformer via the transformer block and set the notation of various mathematical quantities we will be using in future sections. For more details on transformers the reader can consult Vaswani et al. (2017); Dosovitskiy et al. (2020).
Transformer architectures comprise of transformer blocks, defined as follows. A transformer block is a mapping defined as
(3.1) |
where is a feedforward MLP with a residual connection and is an attention head.
The attention head is defined as follows: It comprises of three learnable matrices, a query (), key () and value () defined by: , , for an input sequence with , and . The attention head is then defined by
(3.2) |
where is a similarity transformation and is an activation function. The most common used is the dot-product: , known as self-attention, and will be the one we focus on in this paper. The most common activation function that is used by authors is softmax. This leads to the most common form of the attention head given by
(3.3) |
The function is the matrix softmax map that applies the usual softmax function row-wise:
(3.4) |
The factor , as explained in Vaswani et al. (2017), is a scaling to prevent the gradients of softmax from being too small. For the theoretical analysis in this paper we will only use the dot-product similarity and call the matrix the softmax self-attention matrix. In the experiments section, Sec. 5, we will empirically validate our theoretical framework on more general softmax attention blocks.
For general transformer architectures, multiple heads for are used. Each attention head is defined by equation 3.3 and then all outputs of each attention head are concatenated together before going into the feedforward layer.
We will need notation for the derivative of the matrix softmax map defined by equation 3.4. Given a matrix we can differentiate the matrix map at and obtain the gradient linear map that is defined by the formula
(3.5) |
where is the Jacobian of at .
Given a matrix , we denote its Frobenius norm by . Additionally, we use the notation to represent the expectation of a random variable, where the specific random variable being considered will be clear from the context.
4 Theoretical Analysis
4.1 Implicit regulatization of Softmax
This section presents a theoretical result showing that the softmax activation imposes control over the Frobenius norm of the self-attention matrix in a way that grows sub-linearly with the input sequence’s token length. Additionally, we demonstrate that the gradient of the softmax with respect to the self-attention matrix also exhibits a similar degree of regularity. While previous work has analyzed the regularity of softmax self-attention through the lens of the Lipschitz constant (Kim et al., 2021; Castin et al., 2023), our theorem offers a novel perspective by directly linking the Frobenius norm regularity to the token length. This provides insights into how self-attention activations should scale with token length to maintain stability during training, especially with gradient descent-based algorithms.
Theorem 4.1.
Let be the matrix softmax map defined by equation 3.4 and let denote the gradient of at . We then have the following bounds on the Frobenius norms
(4.1) | ||||
(4.2) |
The key implication of theorem 4.1 is that during the training of a transformer with softmax self-attention, the Frobenius norm of each softmax self-attention matrix remains bounded by a value that grows as . This ensures that backpropagation through the weights of the self-attention matrix does not lead to excessively large gradients. The proof hinges on the fact that the row normalization inherent in softmax effectively controls the Frobenius norm. For a detailed proof see appendix A.1.1.
4.2 Polynomial activations for self-attention
In section 4.1, we demonstrated that softmax implicitly regularizes the Frobenius norm of the self-attention matrix. Building on this, we now show that by scaling specific polynomial activations, a similar regularization effect on the Frobenius norm can be achieved in expectation, closely replicating the impact of softmax.
Theorem 4.2.
Let and , be i.i.d random variables distributed according to and , . We have the following expectations of the Frobenius norms of powers of the matrix for
(4.3) |
By scaling such an activation by we can obtain a bound.
Corollary 4.3.
Assume the same conditions as in theorem 4.2. Then
(4.4) |
Corollary 4.3 establishes that activations of the form provide a level of regularization, in expectation, similar to that of softmax when applied to the self-attention matrix. The proof of theorem 4.2 can be found in appendix A.1.2. The next property we want to prove is one similar to the gradient bound obtained in theorem 4.1. Since the self-attention matrix has parameters given by the queries and keys (Vaswani et al., 2017), this implies that during the training of a transformer the and matrices are the only aspects of the self-attention matrix that get updated. Therefore, we compute a regularity result with respect to the and derivatives.
Theorem 4.4.
Let and , be i.i.d random variables distributed according to and , . Then the expectation of the of the derivative of the matrix w.r.t the parameter matrix for is given by
(4.5) |
The above theorem then suggests that if we scale the polynomial by the derivative will have growth.
Corollary 4.5.
Assume the same condition as in theorem 4.4. Then
(4.6) |
An analogous estimate holds for derivatives with respect to the matrix. The proof of theorem 4.4 can be found in appendix A.1.2.
Corollaries 4.3 and 4.5 suggest that polynomial activations of the form , with , can achieve performance comparable to softmax when applied to self-attention matrices. In section 5, we empirically compare these activations to softmax and observe that they outperform softmax on a variety of transformer tasks. We focus on and as these polynomials clearly violate key aspects of softmax based attention, such as normalized rows, positivity, and sparsity. For larger values of , performance declines due to the functions having smaller gradients around when is large, causing difficulties in training.
5 Experiments
In this section, we validate the theory from section 4 on a variety of transformer tasks. We perform the empirical validation on two primary activations from section 4, namely a cubic polynomial activation and a linear polynomial . The goal will be to show that by suitably scaling these activations using the theory in section 4, we can achieve competitive performance when compared to softmax. For the rest of this section we will simply denote these activations by and .
5.1 Image classification
5.1.1 ViT-Tiny on Tiny-Imagenet:
In this section we test the theory from section 4 on the ViT-Tiny architecture (Steiner et al., 2021) trained from scratch on the Tiny-Imagenet dataset (Le & Yang, 2015).
Our first experiment was to test how the Top- accuracy changes for a ViT-Tiny trained on Tiny-Imagenet as we change the sequence length of the input and the scale predicted in corollaries 4.3 and 4.5 when using the activation . According to the theory developed in section 4, the Frobenius norm scales according to when we scale by . Thus as the sequence length decreases we should see the amount of scaling in a Log scale decrease.
Figure 1 shows the results of this experiment. We considered four different input sequence lengths of sizes , , and . We ran several ViT-Tiny architectures with a variety of scalings of the form where ranged below to above the sequence length. As can be seen from figure 1 as the sequence length got smaller the amount of scaling, shown in Log scale on the x-axis, needed for good accuracy got smaller verifying the theory in section 4.2.
The second experiment compared activations and , along with scaled versions and , against softmax using the Tiny-ViT architecture on Tiny-Imagenet. With a sequence length of (), we decided to take as the scale of the polynomial activations. The experiment used a patch size of 4, 3 attention heads, and 12 layers as described in Steiner et al. (2021). Results in table 1 show outperforming softmax, while the unscaled version performed poorly. Similarly, performed competitively with a significant drop in performance without scaling.
Figure 2 displays the Frobenius norm of the self-attention matrix during training for five activations in layers 2 and 12 of ViT-Tiny, averaged across all heads. Norms for and are higher than softmax, but scaling by reduces them to more stable levels, improving training stability. Similarly, figure 3 shows the Jacobian’s Frobenius norm, where scaling also brings the norms closer to softmax, ensuring more stable gradients. Further plots for other layers are in appendix A.2.2.
softmax | |||||
Top-1% accuracy | 50.26 | 50.5 | 45.3 | 47.9 | 31.78 |
5.1.2 Larger vision transformers on ImageNet-1k
For this experiment we carried out an image classification task using a variety of different vision transformers from the literature on the ImageNet-1k dataset. We found that the scales worked best for both and .
We train all models on the ImageNet-1k dataset from scratch and report Top-1 accuracy on the validation set. We use PyTorch Paszke et al. (2019) and Timm Wightman (2019) libraries to train our models with similar setups to He et al. (2022) Liu et al. (2021). We examined our approach along with the following three transformer architectures to show its generalization:
- ViT: Dosovitskiy et al. (2020) is the pioneering work that interprets an image as a sequence of patches and processes it by a standard Transformer encoder as used in NLP. This simple, yet scalable, strategy works surprisingly well when coupled with pre-training on large datasets. We use ViT-Small which has the following settings: patch size = 16, embedding dimensions = 384, number of heads = 6, and layers = 12. Also, We use ViT-Base which has the following settings: patch size = 16, embedding dimensions = 768, number of heads = 12, and layers = 12.
- DeiT: Touvron et al. (2021) is a well-known transformer based on ViT. It is very similar to ViT except that it converged faster due to better training recipe since we do not use the distillation token imposed in DeiT. We use DeiT-Small which has the following settings: patch size = 16, embedding dimensions = 384, number of heads = 6, and layers = 12. Also, We use Deit-Base which has the following settings: patch size = 16, embedding dimensions = 768, number of heads = 12, and layers = 12.
- Swin Transformer: Liu et al. (2021) produces a hierarchical feature representation and proposes the shifted window based self-attention which is shown to be effective and efficient on vision problems. We use Swin-Small with 96 channels and Swin-Base with 128 channels in the hidden layers of the first stage. The window size is set to = 7 by default, the query dimension of each head is = 32, and layer numbers are for all experiments.
- XciT: Xiong et al. (2021) is a vision transformer architecture consisting of two different components compared to the standard ViT. Firstly, it has Local Patch Interaction in each block, which includes one depth-wise 3×3 convolution followed by Batch Normalization, GELU, and another depth-wise 3×3 convolution. Secondly, it uses Cross-Covariance attention, where the attention map is derived from the cross-covariance matrix computed over the key and query projections of the token features. We use XCiT-S12 with a patch size of 16 and XCiT-M24 with a patch size of 24.
The results are shown in table 2. The activation performed best on the ViT’s and the Swin transformers while softmax performed best on the DeiT architectures. Further ablations with different scales and activations can be fond in appendix A.2.1.
Activation | Models | |||||||
ViT-Base | ViT-Small | DeiT-Base | DeiT-Small | Swin-Base | Swin-Small | XciT-Medium | XciT-Small | |
softmax | 79.6 | 80.2 | 78.9 | 79.6 | 83.0 | 83.3 | 81.1 | 81.2 |
79.6 | 80.5 | 77.4 | 78.3 | 83.2 | 83.4 | 81.2 | 82.1 | |
77.8 | 78.6 | 76.1 | 76.3 | 79.7 | 79.9 | 78.1 | 78.3 | |
76.9 | 77.8 | 79.6 | 79.8 | 79.4 | 79.5 | 79.2 | 79.3 | |
73.2 | 73.9 | 77.7 | 77.9 | 77.8 | 77.9 | 76.4 | 76.6 |
Figure 4 plots the Frobenius norm of the self-attention matrix with the five different activations during training in layers 2 and 12, averaged over all heads within each layer, of the ViT-Small architecture. By scaling the activations and by we were able to control the scale of the Frobenius norm of the self-attention matrix and obtain scales comparable to softmax’s scale. Simiarly, figure 5 plots the Frobenius norm of the Jacobian of the self-attention matrix during training for layers 2 and 12, averaged over all heads. By scaling the activations and by we were able to control the scale of the Jacobian norm and obtain scales comparable to softmax’s scale during training. Plots for other layers of the architecture during training are given in appendix A.2.2.
5.1.3 Visualizing self-attention with ViT-Base
We plotted heat maps of self-attention matrices using the and softmax activations across two layers and heads after convergence, averaging over a fixed training batch of size 128. Figure 6 shows layer 2, head 8, highlighting the differences in attention patterns, with containing both positive and negative values. Similarly, Figure 7 for layer 12, head 6 shows distinct patterns for each activation. Overall, the ViT-Base architecture with exhibits notably different self-attention patterns compared to softmax.
We visualized how the self-attention matrix targets different regions of an image. Using an image from the ImageNet-1k validation set, we extracted the class token, reshaped it into a grid representing 196 patches, and then mapped it back to the original image size with nearest neighbor interpolation. Figure 8 shows the input image, while figure 9 illustrates the self-attention matrices for different activations in layer 12, head 6 of the ViT-Base architecture after convergence, highlighting their distinct focus areas.
5.2 Object Detection and Instance Segmentation
In this section, in order to examine the transfer learning ability of our models, we demonstrate our approach to object detection and segmentation tasks by fine-tuning our ImageNet-pretrained XCiT model on them. Our experiments are conducted on COCO 2017Lin et al. (2014), which has 118K training images and 5K validation images with 80 categories. We integrate the XCiT architecture as the backbone in the Mask R-CNN (He et al., 2017) detector with a Feature Pyramid Network (FPN). Due to XCiT’s inherently columnar design, we adapt it for FPN compatibility by extracting features from various layers for XCiT-S12. These features have a consistent stride of either 16. The feature resolutions are then adjusted to strides of [4, 8, 16, 32] This downsampling is accomplished through max pooling, while upsampling is achieved using a single transposed convolution layer. The model is trained for 36 epochs using the AdamW optimizer with learning rate of , 0.05 weight decay and 16 batch size. In table 3,we condunct experiments on XCiT-S12 models using 16×16 patches with the activations , and softmax. We found we couldn’t train with the activations and on this task well so only report the others.
Activation | ||||||
softmax | 44.9 | 66.1 | 48.9 | 40.1 | 63.1 | 42.8 |
44.8 | 66.3 | 49 | 40.1 | 63.1 | 42.9 | |
44.8 | 66.2 | 49 | 40.1 | 63.1 | 42.8 |
5.3 Natural language processing(NLP)
To assess the effectiveness of our approach on NLP tasks, we trained models on five benchmarks from the Long Range Arena (LRA) suite Tay et al. (2020): ListOps, Text Classification, Retrieval, Image Classification, and Pathfinder. We evaluated the activations and against softmax, finding that and did not train effectively on their own, so only results for these scaled activations and softmax are presented. Our implementation followed the guidelines from Xiong et al. (2021). The results are summarized in table 4.
Activation | ListOps | Text | Retrieval | Image | Pathfinder |
softmax | 37.1 | 63.8 | 79.8 | 39.9 | 72.9 |
37.5 | 63.3 | 80.9 | 37.2 | 68.7 | |
34.3 | 62.9 | 81.5 | 39.0 | 69.1 |
6 Limitations
While our work introduces novel activations that challenge the conventional softmax approach, there are some limitations to address. Our theoretical framework is primarily designed for dot-product self-attention and may not immediately extend to other attention mechanisms, although our empirical results showed competitive performance against softmax across different architectures. Additionally, we observed that while our activations performed well on vision tasks, their performance was less consistent on NLP tasks, suggesting that a more refined theoretical approach may be needed for these applications.
7 Conclusion
This work challenges the traditional view that transformer activations for attention must produce sparse probability distributions. We introduced a theoretical framework analyzing the Frobenius norm of the self-attention matrix, which suggests key scaling properties for activations in attention mechanisms. We proved that specific polynomial activations, which behave very differently from softmax, satisfy these properties. Through extensive experiments across vision and NLP tasks, we demonstrated that these alternative activations not only compete with but sometimes outperform softmax, offering a fresh perspective on attention mechanisms in transformers.
References
- AUEB et al. (2016) Titsias RC AUEB et al. One-vs-each approximation to softmax for scalable estimation of probabilities. Advances in Neural Information Processing Systems, 29, 2016.
- Banerjee et al. (2020) Kunal Banerjee, Rishi Raj Gupta, Karthik Vyas, Biswajit Mishra, et al. Exploring alternatives to softmax function. arXiv preprint arXiv:2011.11538, 2020.
- Carion et al. (2020) Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In European conference on computer vision, pp. 213–229. Springer, 2020.
- Castin et al. (2023) Valérie Castin, Pierre Ablin, and Gabriel Peyré. Understanding the regularity of self-attention with optimal transport. arXiv preprint arXiv:2312.14820, 2023.
- Correia et al. (2019) Goncalo M Correia, Vlad Niculae, and André FT Martins. Adaptively sparse transformers. arXiv preprint arXiv:1909.00015, 2019.
- Devlin et al. (2018) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.
- Dosovitskiy et al. (2020) Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929, 2020.
- Fang et al. (2022) Haishuo Fang, Ji-Ung Lee, Nafise Sadat Moosavi, and Iryna Gurevych. Transformers with learnable activation functions. arXiv preprint arXiv:2208.14111, 2022.
- Fu et al. (2024) Daocheng Fu, Xin Li, Licheng Wen, Min Dou, Pinlong Cai, Botian Shi, and Yu Qiao. Drive like a human: Rethinking autonomous driving with large language models. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pp. 910–919, 2024.
- He et al. (2017) Kaiming He, Georgia Gkioxari, Piotr Dollár, and Ross Girshick. Mask r-cnn. In Proceedings of the IEEE international conference on computer vision, pp. 2961–2969, 2017.
- He et al. (2022) Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 16000–16009, 2022.
- Kim et al. (2021) Hyunjik Kim, George Papamakarios, and Andriy Mnih. The lipschitz constant of self-attention. In International Conference on Machine Learning, pp. 5562–5571. PMLR, 2021.
- Koohpayegani & Pirsiavash (2024) Soroush Abbasi Koohpayegani and Hamed Pirsiavash. Sima: Simple softmax-free attention for vision transformers. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pp. 2607–2617, 2024.
- Le & Yang (2015) Ya Le and Xuan Yang. Tiny imagenet visual recognition challenge. CS 231N, 7(7):3, 2015.
- Lin et al. (2014) Tsung-Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects in context. In Computer Vision–ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13, pp. 740–755. Springer, 2014.
- Liu et al. (2021) Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 10012–10022, 2021.
- Maiti et al. (2023) Abhisek Maiti, Sander Oude Elberink, and George Vosselman. Transfusion: Multi-modal fusion network for semantic segmentation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 6536–6546, 2023.
- Paszke et al. (2019) Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
- Salzmann et al. (2020) Tim Salzmann, Boris Ivanovic, Punarjay Chakravarty, and Marco Pavone. Trajectron++: Multi-agent generative trajectory forecasting with heterogeneous data for control. arXiv preprint arXiv:2001.03093, 2, 2020.
- Shen et al. (2023) Kai Shen, Junliang Guo, Xu Tan, Siliang Tang, Rui Wang, and Jiang Bian. A study on relu and softmax in transformer. arXiv preprint arXiv:2302.06461, 2023.
- Steiner et al. (2021) Andreas Steiner, Alexander Kolesnikov, Xiaohua Zhai, Ross Wightman, Jakob Uszkoreit, and Lucas Beyer. How to train your vit? data, augmentation, and regularization in vision transformers. arXiv preprint arXiv:2106.10270, 2021.
- Tay et al. (2020) Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers. arXiv preprint arXiv:2011.04006, 2020.
- Touvron et al. (2021) Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International conference on machine learning, pp. 10347–10357. PMLR, 2021.
- Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
- Wang et al. (2021) Shulun Wang, Feng Liu, and Bin Liu. Escaping the gradient vanishing: Periodic alternatives of softmax in attention mechanism. IEEE Access, 9:168749–168759, 2021.
- Wightman (2019) Ross Wightman. Pytorch image models. https://github.com/rwightman/pytorch-image-models, 2019.
- Xiong et al. (2021) Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströmformer: A nyström-based algorithm for approximating self-attention. Proceedings of the AAAI Conference on Artificial Intelligence, 2021.
- Zhen et al. (2022) Q Zhen, W Sun, H Deng, D Li, Y Wei, B Lv, J Yan, L Kong, and Y Zhong. cosformer: rethinking softmax in attention. In International Conference on Learning Representations, 2022.
- Zhuang et al. (2021) Liu Zhuang, Lin Wayne, Shi Ya, and Zhao Jun. A robustly optimized bert pre-training approach with post-training. In Proceedings of the 20th chinese national conference on computational linguistics, pp. 1218–1227, 2021.
Appendix A Appendix
A.1 Theoretical analysis
A.1.1 Proofs for theorems in section 4.1
In this section we give the proof of theorem 4.1.
Proof of theorem 4.1.
We will start by proving the first inequality in theorem 4.1. Given a matrix we have that
(A.1) |
By definition of the Frobenius norm we then see that
(A.2) | ||||
(A.3) | ||||
(A.4) | ||||
(A.5) |
where the second inequality uses the fact that for non-negative numbers and we always have that .
It then immediately follows that and this proves the first inequality in the statement of theorem 4.1.
We move on to prove the second inequality in the statement of theorem 4.1. For this, let us write each entry of the matrix on the right of equation A.1 as follows:
(A.6) |
By applying the chain rule we then have the following derivative formulas
(A.7) | ||||
(A.8) | ||||
(A.9) |
We can then express the gradient as
(A.10) |
where
(A.11) |
From these computations we see that
(A.12) |
We will proceed by bounding each collection separately then add up all the bounds. We have
(A.13) | ||||
(A.14) | ||||
(A.15) | ||||
(A.16) | ||||
(A.17) | ||||
(A.18) | ||||
(A.19) | ||||
(A.20) |
We then observe that since we have that where means we don’t include in the sum. This means we get the bound
(A.21) | ||||
(A.22) | ||||
(A.23) | ||||
(A.24) | ||||
(A.25) |
Putting all the bounds together for each of the terms terms we get
(A.26) |
and this implies
(A.27) |
This finishes the proof of theorem 4.1.
∎
A.1.2 Proofs for theorems section 4.2
Proof of theorem 4.2.
We will split the matrix product and think of it as the product of two matrices. Suppose , and . Each element in the matrix can be written as a product of a row of with a column of . Since expectation is linear, we need to compute the expectation of each of these elements. We do the case of the entry which is the entry in in the first row and first column. For the case we can then compute
(A.28) | ||||
The Frobenius norm of the matrix is just the sum of these values for all elements and this proves the case.
For the case that we proceed in a similar way. The key observation is that odd powers, in the matrix expansion, will have expectaion , so we need only consider the even powers. Therefore, suppose . We will compute the expectation of the first entry :
(A.29) | ||||
Note that the first term only has a count of and the second term has a count of . Thus, we only need to consider the term where all the components have a power of 2. The count is similar to choosing items from ,
(A.30) | ||||
can always be bounded above by , so the expectation can be upper bounded by and thus we get a quantity of the form .
∎
Proof of theorem 4.4.
We will do the case first. We proceed similar to the proof of Theorem 4.2.
(A.31) | ||||
(A.32) | ||||
(A.33) | ||||
(A.34) | ||||
(A.35) | ||||
(A.36) | ||||
(A.37) | ||||
(A.38) | ||||
(A.39) | ||||
(A.40) | ||||
(A.41) |
When we can proceed in a similar way.
(A.42) | ||||
(A.43) | ||||
(A.44) | ||||
(A.45) |
We know that
(A.46) | ||||
(A.47) | ||||
(A.48) | ||||
(A.49) |
where . Let Thus we have
(A.50) | |||||
(A.51) | |||||
(A.52) | |||||
(A.53) | |||||
(A.54) | |||||
(A.55) | |||||
(A.56) |
showing that we can bound the gradient by a quantity of the form and the proof is complete.
∎
A.2 Experiments
A.2.1 Ablations
In this section we carry out ablations on the experiments we did in the main paper.
Scale ablations:
The theory we developed in section 4.2 suggested that we needed to scale our polynomial activations by to obtain a complexity bound of . In general, we could also scale the polynomial activations by to see if we can get a better accuracr. Figure 10 carries out an ablation on both the ViT-Base architecture and the ViT-Small architecture to see how different scales affect the accuracy for the activation. We found that in general a scale from to seemed to perform very well.
Activation ablations:
Our theory primarily compared the activations and with softmax where was the sequence length of the input. In this section we carry out further experimental comparisons of our activations with other activations. We compare with a variety of activations used in the literature as well as the exponential function which can be thought of as softmax without the row normalization scaling. Table 5 shows the results of the experiments.
Activation | ViT-Base | ViT-small |
softmax | 79.6 | 80.2 |
79.6 | 80.5 | |
76.9 | 77.8 | |
78.7 | 79.9 | |
77.3 | 77.5 | |
78.2 | 78.5 | |
78.2 | 78.4 | |
77.1 | 77.3 | |
77.9 | 78.0 |
A.2.2 Frobenius norm computations
In section 5.1 we showed plots of the Frobenius norm of the self-attention matrix and for the Jacobian of the self-attention matrix for different scalings of the and activations along with softmax. This was done for a ViT-Tiny architecture on the Tiny-ImageNet dataset. Figure 11 shows the plots of the Frobenius norm of the self-attention matrix for the Tiny-ViT architecture, during training, for all layers averaged over the heads within each layer. Figure 12 shows the Frobenius norm of the Jacobian of the self-attention matrix during training for each layer, averaged over the total number of heads within each layer.