Grouped self-attention mechanism for a memory-efficient Transformer

Transformer (Vaswani et al. (2017)) is a popular deep learning architecture utilized in many different applications. Their powerful performance on sequence data has been demonstrated in numerous domains. They are thus utilized across many domains such as natural language processing (NLP) and computer vision (CV). Many state-of-the-art methods for various tasks in such domains use a Transformer architecture, which is an encoder-decoder architecture based on self-attention and cross-attention modules.

The core mechanism of Transformer is the self-attention module. The self-attention module utilizes scaled dot-product attention with an input comprising queries and keys, and values of dimension $d$ generated from each input node of the Transformer model. Scaled dot-product attention is calculated by dividing dot product results of queries and keys by $\sqrt{d}$ and applying a softmax function to obtain the weights of the values. To compute the attention function on a set of queries simultaneously, queries, keys, and values are concatenated into a matrix ${\bm{Q}},{\bm{K}},{\bm{V}}$ where ${\bm{Q}}\in\mathbb{R}^{l_{Q}\times d}$, ${\bm{K}}\in\mathbb{R}^{l_{K}\times d}$, and ${\bm{V}}\in\mathbb{R}^{l_{K}\times d}$ as given below.

However, the self-attention mechanism is also the bottleneck of the Transformer architecture owing to the quadratic growth of its computational complexity for the sequence length. This is caused by the calculation of the attention matrix where the dot-product operation is performed $l_{Q}\times l_{K}$ times and occupies memory in proportion to $l_{Q}\times l_{K}$ as well.

Capturing long-range dependency is a crucial factor in time-series data analysis. Compressive Transformer (Rae et al. (2019)) and Transformer-XL (Dai et al. (2019)) reinforced the ability of capturing long-range dependency by utilizing auxiliary hidden states. However, these approaches could amplify the computational order issue mentioned above, especially with longer input sequences.

Recent studies have addressed this issue to apply the Transformer architecture to tasks involving a long sequence length by creating an efficient way to calculate attention values.

In this section, we provide a detailed explanation of the proposed Grouped Self-Attention (GSA) mechanism. Overall, we utilize the concept that self-attention reflects the interaction among the input nodes into outputs. The input nodes are divided into several groups to calculate the local attention within the groups, and summarized nodes generated from each groups are concatenated to compute another self-attention to reflect the information interaction among other groups.

For detailed explanation, first, for each input node $n_{i}\in\mathbb{R}^{d}$ ($i\in\{1,...,l\}$), query ${\bm{q}}_{i}\in\mathbb{R}^{d}$, key ${\bm{k}}_{i}\in\mathbb{R}^{d}$, and value ${\bm{v}}_{i}\in\mathbb{R}^{d}$ are generated by linear projection as in canonical self-attention methods. The input nodes are divided into several groups within which the local attention is computed to capture the locality.

To compute the attention function efficiently, all queries, keys, and values are concatenated to form matrices ${\bm{Q}}\in\mathbb{R}^{l\times d}$, ${\bm{K}}\in\mathbb{R}^{l\times d}$ and ${\bm{V}}\in\mathbb{R}^{l\times d}$. Before computing the attention matrix, as shown in Fig.1 and Fig.2, the input nodes are divided into several groups with a length of $l_{g}$. We define these groups of nodes as G = $\{g_{j}\in\mathbb{R}^{l_{g}\times d}\}_{j=1}^{m}$, whereas $m$ is the number of divided groups. If $l$ cannot be divided by $l_{g}$, the zero-padded nodes are added to match the length $ml_{g}$. Queries, keys, and values are generated by linear projection for each group. For example, queries, keys, and values of group $g_{j}$ are defined as ${\bm{Q}}_{j}^{g}\in\mathbb{R}^{l_{g}\times d}$, ${\bm{K}}_{j}^{g}\in\mathbb{R}^{l_{g}\times d}$, and ${\bm{V}}_{j}^{g}\in\mathbb{R}^{l_{g}\times d}$ respectively. The attention output of group $g_{j}$, ${\bm{O}}_{j}^{g}\in\mathbb{R}^{l_{g}\times d}$, is computed as given below in equation 3

After computing the local attention among these groups, summarized nodes defined as S = $\{S_{j}\in\mathbb{R}^{l_{s}\times d}\}_{j=1}^{m}$ are computed. The queries, keys, and values of summarized nodes $S_{i}$ for group $g_{j}$ are calculated by applying linear projection matrices ${\bm{E}}_{q},{\bm{E}}_{k},{\bm{E}}_{v}\in\mathbb{R}^{l_{s}\times l_{g}}$ as given below in equation 4.

where ${\bm{Q}}_{j}^{s},{\bm{K}}_{j}^{s},{\bm{V}}_{j}^{s}\in\mathbb{R}^{l_{s}\times d}$.

Self-attention is again calculated only with these queries, keys, and values of summarized nodes to reflect the information of other groups. To compute the self-attention among summarized nodes, all ${\bm{Q}}_{j}^{s},{\bm{K}}_{j}^{s}$, and ${\bm{V}}_{j}^{s}$ are concatenated to form ${\bm{Q}}^{s}_{cat},{\bm{K}}^{s}_{cat},{\bm{V}}^{s}_{cat}\in\mathbb{R}^{ml_{s}\times d}$.

These concatenated queries, keys, and values are used to compute the global self-attention output ${\bm{O}}^{s}$ as illustrated in Fig.1 which represents the projected global features of all input nodes. ${\bm{O}}^{s}\in\mathbb{R}^{ml_{s}\times d}$ of equation 5 reflects the global information by calculating attention with summarized queries and keys across the groups.

${\bm{O}}^{s}$ is then divided into $m$ segments where each segment ${\bm{O}}_{j}^{s}\in\mathbb{R}^{l_{s}\times d}$ represents the output that reflects global features with local information in group $g_{j}$. ${\bm{O}}_{j}^{s}$ is pooled by average pooling to form $\overline{{\bm{O}}_{j}^{s}}\in\mathbb{R}^{1\times d}$, where $\overline{{\bm{O}}_{j}^{s}}=\sum_{k=1}^{l_{s}}({\bm{O}}_{j}^{s})_{k}$, and it is added to the local attention output ${\bm{O}}_{j}^{g}$ to form the final output of group $g_{j}$, $\widetilde{{\bm{O}}_{j}^{g}}\in\mathbb{R}^{l_{g}\times d}$ according to equation 6.

$\alpha$ and $\beta$ are learnable parameters for each group that determine the reflection ratio of global output into local attention. All $\widetilde{{\bm{O}}_{j}^{g}}$ are concatenated to form final output of the GSA layer ${\bm{O}}\in\mathbb{R}^{ml_{g}\times d}$. This process allows the local attention to consider the global features while capturing locality.

The computational complexity of our proposed approach increases with the number of dot-production calculations in the attention, which is given as $m\times l_{g}^{2}+(m\times l_{s})^{2}$. Because $m$ is the number of groups and $l_{g}$ is the number of nodes included in one group, $m\times l_{g}$ equals the sequence length $l$, and because $l_{s}$ is the length of summarized nodes that qualifies $l_{s}<<l_{g}<l$, the algorithm has a complexity of order $O(l_{g}\times l)$. $l_{g}$ is a fixed length representing the number of nodes included in a single group, and thus the final complexity order becomes linear with respect to sequence length $l$.

To summarize, the GSA module is a linear-order self-attention mechanism based on a local attention mechanism that considers global features by receiving information from interactions of summarized nodes.

The cross-attention module in the canonical Transformer reflects the encoded input features from encoder outputs to the decoder layer. The keys and values for cross-attention are generated from the encoder output, and the queries are generated from the decoder input. The attention is computed using those queries, keys, and values similar to equation 2. The length of queries $l_{Q}$, generated from encoder output is equal to that of the input sequence, and the lengths of keys and values $l_{K}$ generated from decoder input are equal to that of the prediction sequence. This module thus exhibits a quadratic computational and space complexity of order $O(l_{Q}\times l_{K})$ when the values of $l_{Q}$ and $l_{K}$ are similar.

We propose a Compressed Cross-Attention (CCA) module that performs linear projection on encoder output across all decoder layers to reduce complexity while minimizing information loss during projection. The output of the encoder is compressed by linear projection to a fixed length $l_{comp}$ in each cross-attention module in the decoder with separate weights. Information loss is minimized by compressing the encoder outputs multiple times in each decoder layer because using separate weights across all layers can extract various different features from the encoder output. Moreover, the computational complexity is reduced to order $O(l_{Q}\times l_{comp})$, which is linear because $l_{comp}$ is a changeable hyperparameter of fixed length.

The hyperparameters that need to be set for the proposed model are listed as the group node length $l_{g}$, the number of layers in the encoder and decoder $e_{l},andd_{l}$, the summarized node length $l_{s}$, and a fixed length $l_{comp}$ in CCA module. Other parameters common to the canonical Transformer model follow Vaswani et al. (2017), such as the dimension of features in the multi-head attention module. A grid search was conducted for hyperparameter settings, and values of 64, 90 for $l_{g}$, 1, 3 for $e_{l}$ and $d_{l}$, and 4, 8 for $l_{s}$ were obtained. The parameters selected according to model performance were 64 for $l_{g}$, 3 for $e_{l}$ and $d_{l}$, and 4 for $l_{s}$. The parameters $e_{l}$ and $d_{l}$ were shared with the baseline models, and thus same values were used for other networks. Finally, a fixed compressed length $l_{comp}$ in the CCA module is set to 256.

Informer, Linformer, Image Transformer, and Reformer models were used as baselines in the experiments. As we utilized frameworks and code from Informer, the default parameters from Zhou et al. (2021) were used with the Informer architecture. The projected dimension in the low-rank approximation of Linformer was set to 256, as in Wang et al. (2020). For local1d attention in Image Transformer, the block length for each local block was set to 64 as in $l_{g}$ in grouped attention for a fair comparison with the proposed method. The bucket sizes required in Reformer were set to 48, 42, 72, 42, 42, 60, and 60 for sequence lengths 96, 168, 288, 336, 672, 720, and 1440 respectively because half of the sequence length should be divisible by the bucket size.

As mentioned above, ETTh1, ETTh2, ETTm1, ECL, and Weather datasets were used to evaluate the performance of models. For each model, global time stamp embedding and generative style decoder structure were applied as in Zhou et al. (2021) to focus on comparing the self-attention modules. One difference among the networks other than the self-attention module was that Informer used encoder distillation to reduce the computational complexity of calculating cross-attention, and the CCA module was applied to reduce the computational complexity of other methods, including the proposed approach.

In the experiment with multivariate settings, we fed the multi-dimensional features as input to predict all the features, including not only the target feature but also other features as well. In contrast, in the experiments following univariate settings, we input the single-dimensional feature of the target value to predict target features using the output of the models. We focused on comparing the self-attention modules as much as possible, and we also compared the performance and computational complexity for all the variations of self-attention.