A Primal-Dual Framework for Transformers and Neural Networks

Introduction

Transformer models (Vaswani et al., 2017) have achieved impressive success with state-of-the-art performance in a myriad of sequence processing tasks, including those in computer vision (Dosovitskiy et al., 2021; Liu et al., 2021; Touvron et al., 2020; Ramesh et al., 2021; Radford et al., 2021; Arnab et al., 2021; Liu et al., 2022; Zhao et al., 2021; Guo et al., 2021), natural language processing (Devlin et al., 2018; Al-Rfou et al., 2019; Dai et al., 2019; Child et al., 2019; Raffel et al., 2020; Baevski & Auli, 2019; Brown et al., 2020; Dehghani et al., 2018), reinforcement learning (Chen et al., 2021; Janner et al., 2021), and other important applications (Rives et al., 2021; Jumper et al., 2021; Zhang et al., 2019; Gulati et al., 2020; Wang & Sun, 2022). Transformers can also effectively transfer knowledge from pre-trained models to new tasks with limited supervision (Radford et al., 2018; 2019; Devlin et al., 2018; Yang et al., 2019; Liu et al., 2019). The driving force behind the success of transformers is the self-attention mechanism (Cho et al., 2014; Parikh et al., 2016; Lin et al., 2017), which computes a weighted average of feature representations of the tokens in the sequence with the weights proportional to similarity scores between pairs of representations. The weights calculated by the self-attention determine the relative importance between tokens and thus capture the contextual representations of the sequence (Bahdanau et al., 2014; Vaswani et al., 2017; Kim et al., 2017). It has been argued that the flexibility in capturing diverse syntactic and semantic relationships is critical for the success of transformers (Tenney et al., 2019; Vig & Belinkov, 2019; Clark et al., 2019).

Step 1. The input sequence ${\mathbf{X}}$ is projected into the query matrix ${\mathbf{Q}}$, the key matrix ${\mathbf{K}}$, and the value matrix ${\mathbf{V}}$ via three linear transformations

Step 2. The output sequence ${\mathbf{H}}:=[{\bm{h}}_{1},\cdots,{\bm{h}}_{N}]^{\top}$ is then computed as follows

The self-attention computed by Eqn. (1) and (2) is called the scaled dot-product or softmax attention. In our paper, we call a transformer that uses this attention the softmax transformer. The structure that the attention matrix ${\mathbf{A}}$ learns from training determines the ability of the self-attention to capture contextual representation for each token. Additionally, a residual connection can be added to the output of the self-attention layer, {\bm{h}}_{i}={\bm{x}}_{i}+\sum_{j=1}^{N}{\rm softmax}\Big{(}{{\bm{q}}}_{i}^{\top}{{\bm{k}}}_{j}/\sqrt{D}\Big{)}{{\bm{v}}}_{j}.

Despite their remarkable success, most attention layers are developed based on heuristic approaches, and a coherent principled framework for synthesizing attention layers has remained elusive.

2 Contribution

We derive the self-attention as the support vector expansion of a given support vector regression (SVR) problem. The primal representation of the regression function has the form of a neural network layer. Thus, we establish a primal-dual connection between an attention layer in transformers and a neural network layer in deep neural networks. Our framework suggests a principled approach to developing an attention mechanism: Starting from a neural network layer and a support vector regression problem, we derive the dual as a support vector expansion to attain the corresponding attention layer. We then employ this principled approach to invent two novel classes of attentions: the Batch Normalized Attention (Attention-BN) derived from the batch normalization layer in deep neural networks and the Attention with Scaled Heads (Attention-SH) resulting from solving the support vector regression model with less amount of training data. Our contribution is three-fold.

We derive self-attention as a support vector expansion that solves a SVR problem, thus providing a principled primal-dual framework to study and develop self-attentions.

We re-derive popular attentions, such as the linear attention (Katharopoulos et al., 2020), the sparse attention (Child et al., 2019), and the multi-head attention (Vaswani et al., 2017), from our proposed framework.

We develop two new attention mechanism: the Batch Normalized Attention (Attention-BN) and the Attention with Scaled Heads (Attention-SH) using our proposed framework.

We empirically demonstrate that 1) the Attention-BN significantly outperforms the baseline softmax and linear attention and 2) the Attention-SH performs better while being more efficient than the same baselines on a variety of practical tasks including image and time-series classification.

Primal-Dual Interpretation of Self-Attention

We first provide a primal-dual interpretation of self-attention as a support vector regression problem in Section 2.1. Based on that primal-dual framework, we derive popular attention mechanisms as the support vector expansion in Section 2.2. Finally, we introduce two new attention mechanisms in Section 2.3, the Attention-BN and Attention-SH.

In order to derive the self-attention from the support vector regression defined by the optimization problem 5, the key idea to construct the Lagrangian from Eqn. 5 and find the representation of the ${\bm{w}}_{d}$, $d=1,\ldots,D_{v}$, in terms of the dual variables. We define the Lagrangian function as follows:

Deriving Softmax Attention. Choosing the appropriate $h({\bm{x}})$ and $\Phi({\bm{x}})$ allows us to derive the popular softmax attention given in Eqn. 1 and 2. In particular, if we choose $h({\bm{x}}):=\sum_{j}^{N}\Phi({\bm{x}})^{T}\Phi({\bm{k}}_{j})$, Eqn. 2.1 becomes

We then select $\Phi({\bm{x}})=\left(a^{(0)}_{l_{0}},a^{(1)}_{1},\dots,a^{(1)}_{l_{1}},\dots,a^{(t)}_{1},\dots,a^{(t)}_{l_{t}},\dots\right)$ where $l_{t}={D+t-1\choose t}$ and

Let ${\bm{x}}={{\bm{q}}}_{i}$, ${\bm{b}}=0$ and relax the boundness constraint of ${\bm{v}}_{j}$ in Remark 1. Eqn. D becomes Eqn. 2 of the softmax attention (Vaswani et al., 2017). We summarize our results in the following theorem.

Theorem 1 and its derivation can be easily extended to capture the full form of the softmax attention with the residual connection, the query matrix projection ${\mathbf{W}}_{Q}$, the key matrix projection ${\mathbf{W}}_{K}$, and the value matrix projection ${\mathbf{W}}_{V}$. We include this result in Appendix F.

The primal representation of the function $f$ as in Eqn. 4 has the form of a neural network layer where ${\mathbf{W}}$ is the weight, ${\bm{b}}$ is the bias term, $\Phi({\bm{x}})$ is the input, and $h({\bm{x}})$ is the normalization term. Thus, an attention layer and a neural network layer are primal-dual of each other.

A principled approach to developing an attention mechanism. The observation in Remark 4 suggests a principled way to construct an attention layer: Starting from a neural network layer and a support vector regression problem, we derive the dual as a support vector expansion to attain the corresponding attention layer. Using this approach, we derive popular attention mechanisms in Section 2.2 and propose our new attention mechanisms in Section 2.3.

2 Deriving Popular Attention Mechanisms as the Support Vector Expansion

In this section, we derive popular attentions such as the linear attention (Katharopoulos et al., 2020), the sparse attention (Child et al., 2019), and the multi-head attention (Vaswani et al., 2017).

The Eqn. 27, which is obtained when choosing $h({\bm{x}}):=\sum_{j}^{N}\Phi({\bm{x}})^{T}\Phi({\bm{k}}_{j})$, already matches the formula of the linear attention. Here, we can let ${\bm{b}}=0$ as above and select the function $\Phi$ that results in a positive similarity function, e.g. $\Phi({\bm{x}})=\text{elu}({\bm{x}})+1$, as in (Katharopoulos et al., 2020).

2.2 Sparse Attention

The sparse attention (Child et al., 2019) can be derived by fitting the function $f$ in Eqn. 4 using a different subset $\{({\bm{k}}_{m_{x}(1)},{\bm{y}}_{m_{x}(1)}),\dots,({\bm{k}}_{m_{x}(M)},{\bm{y}}_{m_{x}(M)})\}$ of training data $\{({\bm{k}}_{1},{\bm{y}}_{1}),\dots,({\bm{k}}_{N},{\bm{y}}_{N})\}$ for each input data ${\bm{x}}$, where $\mathcal{M}_{{\bm{x}}}=\{m_{{\bm{x}}}(1),\dots,m_{{\bm{x}}}(M)\}\subset\{1,\dots,N\}$. The support vector expansion of $f$ is then given by

where $\textbf{1}_{\mathcal{M}_{{\bm{x}}}}(j)=\left[j\in\mathcal{M}_{{\bm{x}}}\right]:=\begin{cases}1\,\,\text{if }j\in\mathcal{M}_{{\bm{x}}}\\ 0\,\,\text{otherwise}\end{cases}$. Note that the subsets $\mathcal{M}_{{\bm{x}}}$ are different for different ${\bm{x}}$. When letting ${\bm{x}}={\bm{q}}_{i}$ where ${\bm{q}}_{i}$, $i=1,\dots,N$, are the query vectors and choosing $\Phi,h,{\bm{b}}$ as in Section 2.1, we can obtain the sparse attention in (Child et al., 2019) where the binary matrix ${\mathbf{M}}=\left(\textbf{1}_{\mathcal{M}_{{\bm{q}}_{i}}}(j)\right)_{i,j=1}^{N}$ becomes the sparse masking matrix.

2.3 Multi-head Attention (MHA)

where each function $f^{s}({\bm{x}}^{s})={\mathbf{W}}^{s}\frac{\Phi^{s}({\bm{x}}^{s})}{h^{s}({\bm{x}}^{s})}+{\bm{b}}^{s}$ is fitted to the training dataset $\{({\bm{k}}^{s}_{1},{\bm{y}}^{s}_{1}),\dots,({\bm{k}}^{s}_{N},{\bm{y}}^{s}_{N})\}$. Following the same derivation and choosing $\{\Phi^{s},h^{s},{\bm{b}}^{s}\}_{s=1}^{H}$ as in Section 2.1, we can rewrite $f({\bm{x}})$ in terms of the support vector expansions of the individual functions $f^{s}({\bm{x}}^{s})$, which are the individual softmax attentions

Comparing Eqn. 18 and Eqn. 3, we see that Eqn. 18 computes the MHA when choosing ${\bm{x}}^{s}={\bm{q}}^{s}_{i}$ where ${\bm{q}}^{s}_{i}$, $i=1,\dots,N$, are the query vectors at the $s^{\text{th}}$ head.

3 Deriving New Attention Mechanisms: Batch Normalized Attention and Multiresolution Head Attention

In this section, we employ our primal-dual framework to develop new attention mechanisms. In particular, we derive: 1) the Batch Normalized Attention from employing the batch normalization (Ioffe & Szegedy, 2015); and 2) the Attention with Scaled Heads from using different amounts of training data. By 1) and 2), we demonstrate that new attentions can be invented by modifying the primal neural network layer and the support vector regression problem in our framework, respectively.

Here, $d=1,\dots,D$, and the mean subtraction and division by the standard deviation is performed element-wise along the feature dimension of ${\bm{x}}$. Following the same derivation as in Section 2.1, we derive the following support vector expansion of $f$

Given a set of the key and value vectors $\{{\bm{k}}_{j},{\bm{v}}_{j}\}_{j=1}^{N}$, for each query vector ${\bm{q}}_{i}$, $i=1,\dots,N$, the Batch Normalized Attention (Attention-BN) computes the corresponding output vector ${\bm{h}}_{i}$ of the query ${\bm{q}}_{i}$ by the following attention formula:

The Effect of Normalization. Expanding the dot product in the Attention-BN (see Appendix E), Eqn. 22 becomes

Eqn. 24 implies that in the Attention-BN, the similarity between the query ${\bm{q}}_{i}$ and the key ${\bm{k}}_{j}$ is adjusted by the similarity between the key ${\bm{k}}_{j}$ and all the keys ${\bm{k}}_{j^{\prime}}$, $j^{\prime}=1,\dots,N$. In particular, if the key ${\bm{k}}_{j}$ is too similar to other keys, the query ${\bm{q}}_{i}$ will attend to it less and vice versa.

3.2 Attention with Scaled Heads

Given $H$ sets of the key and value vectors $\{{\bm{k}}^{1}_{j},{\bm{v}}^{1}_{j}\}_{j=1}^{N_{1}},\dots,\{{\bm{k}}^{H}_{j},{\bm{v}}^{H}_{j}\}_{j=1}^{N_{H}}$, for each set of H query vectors ${\bm{q}}^{1}_{i},\dots,{\bm{q}}^{H}_{i}$, $i=1,\dots,N$, the Attention with Scaled Heads (Attention-SH) computes the corresponding output vector ${\bm{h}}_{i}$ of the queries ${\bm{q}}^{1}_{i},\dots,{\bm{q}}^{H}_{i}$ by the following attention formula:

Linear Attention with Batch Normalization and Scaled Heads. The Attention-BN/SH can be extended to use with the linear attention. In particular, in the Linear Attention-BN/SH, we replace the softmax kernel in Eqn. 22 and Eqn. 25 by the linear kernel, respectively.

Experimental Results

In this section, we empirically demonstrate the advantages of our Attention-BN, Attention-SH, and their combination (Attention-BN+SH) over the baseline softmax attention on the UEA time-series classification benchmark (Bagnall et al., 2018), the Long Range Arena benchmark (Tay et al., 2021), and the image classification task on the Imagenet dataset (Deng et al., 2009; Russakovsky et al., 2015). We aim to show that: (i) Attention-BN significantly outperforms the softmax baseline across tasks; (ii) Attention-SH achieves better or comparable accuracy while saving computation and memory compared to the baseline; (iii) Attention-BN+SH, which combines both Attention-BN and Attention-SH, results in the best model performance in term of accuracy and efficiency; (iv) all our proposed models help reduce the redundancy in multi-head attention and benefit learning of the long-term dependency in long input sequences; (v) Attention-BN and Attention-SH can be applied on other attention mechanisms beyond the softmax attention. When combined with the linear attention (Katharopoulos et al., 2020), the resultant Linear Attention-BN and Linear Attention-SH yield similar advantages mentioned in (i), (ii), (iii) and (iv) over the baseline linear attention.

In our experiments, we compare the proposed models with the baseline softmax and linear attentions of the same configuration. For the Attention-BN and Attention-BN+SH, we observe that recentering queries and keys alone is sufficient for improving the model performance. In addition, weighting ${\bm{\mu}}$ with a constant $\beta$, as in Eqn. 26 in the Appendix, enables the Attention-BN/BN+SH to adjust the effect of normalization to the attention score and help increase the accuracy. Our results are averaged over 5 runs. Details on datasets, models, and training are provided in Appendix A.

UEA Time Series Classification. Table 1 compares the accuracy of the Attention-BN and Attention-SH with the baseline softmax attention on 10 tasks in the UEA Time Series Classification benchmark (Bagnall et al., 2018). Both Attention-BN and Attention-SH significantly outperform the softmax baseline on most tasks and on average among all tasks. When combining two models, the resulting Attention-BN+SH yields the best accuracy with more than $1\%$ overall improvement over the softmax baseline. Notably, the Attention-SH and Attention-BN+SH are much more efficient than the baseline since they need much fewer keys and values in computing the attention output. The efficiency advantage of the Attention-SH/BN+SH is discussed and analyzed in detail in Section 4.

Long Range Arena (LRA) benchmark. In this experiment, we verify the advantage of our methods over the softmax baseline on tasks that involve very long sequences (e.g., the sequence length can be up to 4K) in the LRA benchmark (Tay et al., 2021). Those tasks require the model to capture long-range dependency in the input sequence. The summarized results in Table 2 indicate significant improvements of Attention-BN/SH/BN+SH over the baseline softmax attention. Same as in the UEA Time Series experiment, on this LRA benchmark, Attention-BN and Attention-SH both outperform the softmax attention on most five tasks. Moreover, Attention-BN+SH, which combines these two attention mechanisms, results in the most accurate models on average across tasks. Specifically, for the retrieval task, the most challenging task with the largest sequence length in the LRA benchmark, Attention-BN+SH achieve a remarkable improvement of more than $1.5\%$ over the baseline.

Image Classification on Imagenet. We corroborate the advantage of our proposed attention over the baseline softmax attention when scaled up for the large-scale ImageNet image classification task. We summarize the results in Table 3. The Deit model (Touvron et al., 2021) equiped with the Attention-BN yields better performance than the softmax baseline. Meanwhile, Attention-SH/BN+SH Deit perform on par with the baseline while being more efficient. These results, together with other results above justify the benefits of our proposed methods across various tasks and data modalities, proving the effectiveness of our primal-dual approach to develop new attentions.

Empirical Analysis

Efficiency Analysis. The Attention-BN+SH not only improves the accuracy of the model remarkably but also help reduce the computational and memory cost significantly. Fig.1 presents the efficiency benefits of our Attention-BN+SH trained on the retrieval task when the model dimension $D$ and sequence lengths $N$ grow. The efficiency advantage of our model increase as $N$ increase. In addition, the scaled-up models (with large $D$) remains significantly more efficient than the baseline. When the model dimension is 64 and sequence length is 4096, which is the standard configuration of the task, the model’s FLOPS, in both training and inference, reduce almost $25\%$, whereas the reductions for memory usage in training and testing are $31.9\%$ and $47.3\%$, respectively. Notably, this efficient model also outperforms the baseline with more than $1.5\%$ improvement in accuracy. These results prove the benefits of applying the Attention-BN+SH for long-sequence tasks and large-scale models.

New Attentions Helps Reduce Head Redundancy. We compute the average $\mathcal{L}_{2}$ distances between heads to analyze the attention diversity. Given our trained models for the retrieval task, the layer-average mean and standard deviation of distances between heads are reported in Table 4. All our introduced attentions attain greater $\mathcal{L}_{2}$ distances compared to the baseline, reducing the risk of learning redundant heads. In particular, Attention-SH has the highest head difference, indicating the model’s attention patterns are most spread out between heads.

Combining Attention-BN and Attention-SH with Other Attentions. Our methods can be extended to combine with other attention mechanisms. We study the Linear Attention-BN/SH/BN+SH, that combine the Attention-BN/SH/BN+SH with the linear attention (Katharopoulos et al., 2020) as explained at the end of Section 2.3. We summarize our results in Table 5 in Appendix B.1.

Related Work

Interpretation of Attention Mechanism. Recent works have focused on understanding the attention mechanism in transformers from different perspectives. (Tsai et al., 2019) considers attention as a weighted moving average over the inputs via a smoothing kernel. (Nguyen et al., 2022) draws a connection between self-attention and nonparametric kernel regression. With this understanding, the work explores better regression estimators, e.g. the generalized Fourier nonparametric regression estimator, to improve transformers. In addition, (Cao, 2021) then shows that the linear transformer (Katharopoulos et al., 2020) corresponds to a Petrov-Galerkin projection (Reddy, 2004) and proves that the softmax normalization in the softmax attention is sufficient but not necessary. Other works that employ ordinary/partial differential equations to provide an interpretation for attention include (Lu et al., 2019; Sander et al., 2022). From a probabilistic perspective, (Tang & Matteson, 2021; Gabbur et al., 2021; Zhang & Feng, 2021) propose Gaussian mixture model frameworks to study the self-attention in transformers. Using graph-structured learning and message passing in graphical models is another attempt at understanding the attention mechanism Wang et al. (2018); Shaw et al. (2018); Kreuzer et al. (2021). Optimization perspectives of attention mechanisms are recently explored. (Sander et al., 2022) connects transformers with an optimization process across iterations by specifically constructing the core energy function. (Sahiner et al., 2022) derive finite-dimensional convex equivalence of attentions that can be solved for global optimality. Different from these approaches, our primal-dual framework focuses on deriving attention as the dual expansion of a primal neural network layer via solving a support vector regression problem. This framework allows us to not only explain many different types of attention mechanisms but also create new ones.

Efficient Transformers. Recently, efficient transformers have been studied (Roy et al., 2021). Among them are sparse transformers which incorporate sparse structures into the attention matrix (Parmar et al., 2018; Liu et al., 2018; Qiu et al., 2019; Child et al., 2019; Beltagy et al., 2020). Another class of efficient transformers are models that aim to have better coverage by integrating different access patterns (Child et al., 2019; Ho et al., 2019), which can also be learned from the data (Kitaev et al., 2020; Roy et al., 2021; Tay et al., 2020). An emerging body of work is proposed to distill and prune the model, including (Sanh et al., 2019; Sun et al., 2019; Voita et al., 2019; Sajjad et al., 2020). In other works, a side memory module is utilized in order to access multiple tokens simultaneously (Lee et al., 2019; Sukhbaatar et al., 2019; Asai & Choi, 2020; Beltagy et al., 2020). Low-rank and kernelization methods have been proposed to improve the efficiency of self-attention calculation (Tsai et al., 2019; Wang et al., 2020; Katharopoulos et al., 2020; Choromanski et al., 2021; Shen et al., 2021; Peng et al., 2021). Our Attention-SH/BN+SH is orthogonal to these methods.

Concluding Remarks

Acknowledgements

This material is based on research sponsored by the NSF under Grant# 2030859 to the Computing Research Association for the CIFellows Project (CIF2020-UCLA-38). SJO acknowledges support from the ONR N00014-20-1-2093 and N00014-20-1-2787 and the NSF DMS 2208272 and 1952339. RGB acknowledges support from the NSF grants CCF-1911094, IIS-1838177, and IIS-1730574; ONR grants N00014-18-12571, N00014-20-1-2534, and MURI N00014-20-1-2787; AFOSR grant FA9550-22-1-0060; and a Vannevar Bush Faculty Fellowship, ONR grant N00014-18-1-2047. ALB acknowledges support from the NSF grants DMS-2152717 and DMS-1952339. NH acknowledges support from the NSF IFML 2019844 and the NSF AI Institute for Foundations of Machine Learning.

References

Appendix A Additional details on the experiments

This section provides datasets, models, and training details for experiments in Section 3. As mentioned in Section 3, for Attention-BN models, recentering queries and keys alone is sufficient for accuracy improvement, and we weight the mean ${\bm{\mu}}$ in Eqn 22 with a constant $\beta$. Hence Eqn 22 is simplified to:

In our experiments, we consider the constant $\beta$ in Attention-BN/BN+SH and the different downsampling scales in Attention-SH/SH+BN as hyper-parameters to finetune. All of our experiments are conducted on a server with 4 NVIDIA A100 GPUs.

Datasets and metrics The benchmark (Bagnall et al., 2018) consists of 30 datasets. Following (Wu et al., 2022), we choose 10 datasets, which vary in input sequence lengths, the number of classes, and dimensionality, to evaluate our models on temporal sequences. We report the test accuracy as evaluation for the benchmark. Models and baselines The experiment setups and configurations for the softmax/linear baseline and our models are the same as in (Wu et al., 2022) Implementation available at https://github.com/thuml/Flowformer. (for the PEMS-SF, SelfRegulationSCP2, UWaveGestureLibrary datasets) and (Zerveas et al., 2021) Implementation available at https://github.com/gzerveas/mvts_transformer. (for the other tasks). In all models, the number of heads is 8, whereas the model dimension and number of transformer layers are varied. For Attention-SH/SH+BN, we downsample keys and values by the factor of 2, after every two successive heads.

A.2 Long range arena benchmark

Datasets and metrics We adopt the tasks: Listops (Nangia & Bowman, 2018), byte-level IMDb reviews text classification (Maas et al., 2011), byte-level document retrieval (Radev et al., 2013), CIFAR-10 image classification (Krizhevsky et al., 2009) and the Pathfinder challenge (Linsley et al., 2018) in the LRA benchmark for our experiments. They consist of long sequences of length $2K$, $4K$, $4K$, $1K$, and $1K$ respectively. The evaluation protocol and metric are the same as in (Tay et al., 2021).

Models and baselines All our models and softmax/linear baselines follow the same architecture and configuration as in (Zhu et al., 2021)Implementation available at https://github.com/NVIDIA/transformer-ls.. Each model consists of two layers and 64 embedding dimensions. While one head at each layer remains intact, the keys and values of the other heads are halved in our Attention-SH/SH+BN experiments.

A.3 Image Classification on Imagenet

Datasets and metrics The ImageNet dataset (Deng et al., 2009; Russakovsky et al., 2015) consists of $1.28M$ training images and $50K$ validation images. The task is to classify 1000 categories. Top-1 and top-5 accuracies are reported. Models and baselines Our baseline is DeiT-tiny model (Touvron et al., 2021) with 12 transformer layers, 4 attention heads per layer, and the model dimension of 192. For model setting and setting and configuration, we follow (Touvron et al., 2021)Implementation available at https://github.com/facebookresearch/deit.. The downsampling scales in Attention-SH/BN+SH models are $$ for 4 heads at each layer, respectively.

Appendix B Additional Experimental Results

Table 5 summarizes the comparison between the Linear Attention-BN/SH/BN+SH and the baseline Linear Attention on the UEA Time Series Classification task. The Linear Attention-BN/SH/BN+SH achieve better accuracy than the Linear Attention baseline while being more efficient.

B.2 Convolution Attention

Table 6 demonstrates the advantage of Attention-Conv2D (Def. 3, Section G) over softmax Deit on the ImageNet image classification task. Furthemore, as shown in Table 7, the Attention-Conv1D (Def. 4, Section G) outperforms the baseline softmax attention on 5 tasks of the LRA benchmark (Tay et al., 2021).

B.3 Additional experiments on the UEA Timeseries Classification benchmark and the UCR Time Series Regression Archive

In this section, we further demonstrate the advantage of our Attention-BN/SH/BN+SH on additional 15 tasks in the UEA Time Series Classification benchmark and on 6 tasks in the UCR Time Series Regression benchmark. The results in Table 8 and 9 show that our Attention-BN and Attention-SH+BN outperform the baseline softmax transformers significantly on both of these benchmarks, while the attention-SH has comparable performance with the baseline but being more effiicient.

B.4 UEA Time Series Classification using the Sparse Attention-BN/SH/BN+SH

Table 10 summarizes the comparison between the Sparse Attention-BN/SH/BN+SH and the Sparse Attention baseline on a subset of the UEA Time Series Classification benchmark. Our models when combined with Sparse Attention achieve significantly better accuracy than the Sparse Attention baseline while the Sparse Attention-SH/BN+SH are more efficient (See Fig. 3 and Fig. 4 in Appendix C).

B.5 Attention-BN/BN+SH with learnable β𝛽\beta

We experiment with our Attention-BN/BN+SH with learnable $\beta$ on the retrieval task. Table 11 shows that learning $\beta$ does not improve much over setting $\beta$ to be a hyperparameter.

Appendix C Additional Results on Efficiency Analysis

This section provides more efficiency analysis on our models.

Attention-SH. Fig.2 shows the efficiency benefits of our Attention-SH when trained on the retrieval task. Same as in the case of Attention-SH+BN, the efficiency benefits of our Attention-SH over the baseline Softmax attention grows when $N$ and $D$ increase.

Sparse Attention-SH/BN+SH. Fig.3 and Fig.4 show that the efficiency advantages of our Sparse Attention-BN+SH and Sparse Attention-SH, respectively, increase as the model dimension $D$ and sequence length $N$ grow. All models are trained on the LRA retrieval task. In addition to the efficiency advantage, the Sparse Attention-BN+SH also significantly outperforms the Sparse Attention baseline in terms of accuracy in this task (79.86% vs. 78.20%) while the Sparse Attention-SH achieves a comparable result to the baseline. More accuracy advantages of the Sparse Attention-BN/SH/BN+SH over the Sparse Attention baseline are given in Table 10.

Appendix D Deriving Softmax Attention.

Choosing the appropriate $h({\bm{x}})$ and $\Phi({\bm{x}})$ allows us to derive the popular softmax attention given in Eqn. 1 and 2. In particular, if we choose $h({\bm{x}}):=\sum_{j}^{N}\Phi({\bm{x}})^{T}\Phi({\bm{k}}_{j})$, Eqn. 2.1 becomes

We then select $\Phi({\bm{x}})=\left(a^{(0)}_{l_{0}},a^{(1)}_{1},\dots,a^{(1)}_{l_{1}},\dots,a^{(t)}_{1},\dots,a^{(t)}_{l_{t}},\dots\right)$ where $l_{t}={D+t-1\choose t}$ and

Let ${\bm{x}}={{\bm{q}}}_{i}$, ${\bm{b}}=0$ and relax the boundness constraint of ${\bm{v}}_{j}$ in Remark 1. Eqn. D becomes Eqn. 2 of the softmax attention (Vaswani et al., 2017).

Appendix E Batch Normalized Attention: Derivation of Eqn. 24

Appendix F Attention with the Residual Connection and Matrix Projections

In this supplement, we first discuss attention with the residual connection and matrix projections in Appendix F.

The Lagrangian of the optimization problem 33 is given by

Similar to the derivation in Section 2.1, the partial derivatives of $\mathcal{L}_{1}$ with respect to the primal variable ${\bm{w}}_{d}$, $d=1,\ldots,D_{v}$, have to vanish for optimality, which leads to

Note that here we only find the form of the optimal solution for ${\mathbf{W}}=[{\bm{w}}_{1},\dots,{\bm{w}}_{D_{v}}]^{\top}$. The optimal value of ${\mathbf{W}}^{\text{proj}}$ can then be found by optimization algorithm such as the (stochastic) gradient descent when training the transformer.

Here, the support vector expansion of $f$ already includes a residual connection. The softmax attention can then be derived by selecting $h({\bm{x}}):=\sum_{j}^{N}\Phi({\mathbf{W}}^{\text{proj}}{\bm{x}})^{T}\Phi({\mathbf{W}}^{\text{proj}}{\bm{x}}_{j})$ and choosing $\Phi$ as in Eqn. 28 in Section 2.1. Note that in Eqn. F, $\{{\bm{x}}_{j}\}_{j=1}^{N}$ and ${\bm{x}}$ are the training samples and test sample, respectively. In order to derive the key, query, and value matrix projections in attention, we can then relax Eqn. F by letting ${\mathbf{W}}^{\text{proj}}{\bm{x}}_{j}={\mathbf{W}}_{K}{\bm{x}}_{j}$, ${\mathbf{W}}^{\text{proj}}{\bm{x}}={\mathbf{W}}_{Q}{\bm{x}}$, ${\bm{v}}_{j}={\mathbf{W}}_{V}{\bm{x}}_{j}$ and choosing the test sample ${\bm{x}}$ among the training samples $\{{\bm{x}}_{j}\}_{j=1}^{N}$.

Here, for self-attention, we choose the test sample ${\bm{x}}$ among the training samples $\{{\bm{x}}_{j}\}_{j=1}^{N}$ to compute the attention score of a token to other tokens in the same sequence. For cross-attention where a token in a sequence attends to tokens in another sequence, this constraint can be removed.

Appendix G 2D-Convolution Attention

The Lagrangian of the optimization problem 38 is given by

Similar to the derivation in Section 2.1 in the main text, the partial derivatives of $\mathcal{L}$ with respect to the primal variable ${\bm{w}}_{d}$, $d=1,\ldots,D_{v}$, have to vanish for optimality, which leads to

where ${\mathbf{A}}_{ij}:=\Phi(\text{Flatten}(\text{Conv2D}({\mathbf{X}},s))(i))^{\top}\Phi(\text{Flatten}(\text{Conv2D}({\mathbf{X}}^{train},s))(j))$.

where the $Conv2D(\cdot,s)$ is the depth-wise 2D-convolution (Howard et al., 2017) with the kernel size $s\times s$ and identical kernel channels.

The convolutional projections used in the Convolutional vision Transformer (CvT) (Wu et al., 2021) can be derived from Eqn. G by letting the training input tensor ${\mathbf{X}}^{train}$ to be the 2D input matrix of size $N\times D_{x}$ of the self-attention layer (see Section 1.1 in the main text) reshaped into a 3D tensor of size $N_{H}\times N_{W}\times D_{x}$ where $N=N_{H}\times N_{W}$. Here, to avoid confusion, we denote the input of the self-attention layer by ${\mathbf{X}}^{input}$ and its reshaped version by $\text{Reshape2D}({\mathbf{X}}^{input})$. We then replace the depth-wise 2D-convolution by the depth-wise separable 2D-convolution in (Wu et al., 2021) and remove the constraint that the kernels have identical channels. In order to derive the convolutional projections for the keys, queries, and values in CvT, for $i,j=1,\dots,N$, we let

Here, we specify the kernel/filter ${\mathbf{W}}_{K}$, ${\mathbf{W}}_{Q}$, and ${\mathbf{W}}_{V}$ to emphasize that the convolutional projections in CvT uses different kernels to compute keys, queries, and values in self-attention. Eqn. 7 matches the convolutional projects in CvT. By choosing $h$ and $\Phi$ similar to above, we can derive the convolutional attention in CvT.

Appendix H 1D-Convolution Attention

where the $Conv1D(\cdot,s)$ is the depth-wise 1D-convolution with the kernel size $s\times 1$ and identical kernel channels.

Appendix I Attention with Batch Normalization and scaled heads

Following the same derivation as in Section 2.1, we derive the following support vector expansion of $f^{s}$

Given $H$ sets of the key and value vectors $\{{\bm{k}}^{1}_{j},{\bm{v}}^{1}_{j}\}_{j=1}^{N_{1}},\dots,\{{\bm{k}}^{H}_{j},{\bm{v}}^{H}_{j}\}_{j=1}^{N_{H}}$, for each set of H query vectors ${\bm{q}}^{1}_{i},\dots,{\bm{q}}^{H}_{i}$, $i=1,\dots,N$, the Batch Normalized Attention with Scaled Heads (Attention-BN+SH) computes the corresponding output vector ${\bm{h}}_{i}$ of the queries ${\bm{q}}^{1}_{i},\dots,{\bm{q}}^{H}_{i}$ by the following attention formula:

Appendix J Hyperparameters

In this section, we provide the hyper-parameters for our best models.

For these two benchmarks, use the set of downsampling factors $\textbf{s}=$ for Attention-SH/BN+SH and Linear/Sparse Attention-SH/BN+SH models trained on the UEA benchmark. Table 13 and Table 12 provide the values of $\beta$ used for our best Attention-BN/BN+SH and Linear Attention-BN/BN+SH, Sparse Attention-BN/BN+SH models trained on subsets of the two benchmarks.

J.2 Long Range Arena benchmark

For all 5 tasks of the LRA benchmark, we set the downsampling factors s of Attention-SH/BN+SH, Linear/Sparse Attention-SH/BN+SH is $$and kernel size of Attention-Conv1D models is 5. In addition, Table 14 provides the values$\beta$ of Attention-BN/BN+SH models trained on the benchmark.

J.3 Imagenet Classification

This task’s $\beta$ of Attention-BN/BN+SH is 1. Attention-SH/BN+SH has the downsampling factor of $$, and the kernel size of Attention-Conv2D is$(2,2)$.