[ICLR 2022] Part 1: Efficient Transformer

This year marks the 10th anniversary of the International Conference on Learning Representations (ICLR), the most prominent conference in the world in the field of deep learning that was first held in 2013. It is also a venue for the dissemination of cutting-edge research and publications in all fields that employ deep learning, including artificial intelligence, data science, vision, speech recognition, and robotics. We introduce some of the papers presented by LG AI Research at the online ICLR 2022 conference held in April. - Part 1: Efficient Transformer — Junghee Kim, Applied AI Research Lab - Part 2: Time Series Anomaly Detection — Yeseul Shim, Applied AI Research Lab - Part 3: Reinforcement Learning as a sequence modeling problem — Deunsol Yoon, Applied AI Research Lab

Introduction

In 2017, Google introduced Transformer, a new sequence-to-sequence model. Since then, Transformer has been utilized in several fields, including Natural Language Processing (NLP), computer vision, and audio. This blog post introduces “Three Ways to Implement Transformer Efficiently,” which was presented at ICLR 2022. We will first review the seq-to-seq model and Transformer, and then introduce the “Three approaches to Efficiently Implementing Self-Attention of Transformer,” which is the most computationally-intensive component in Transformer.

Seq-to-Seq model Overview

Figure 1 Seq-to-Seq Model

The seq-to-seq model refers to a model with an encoder-decoder structure that outputs a sequence in a domain from the input sequence in another domain. The Recurrent Neural Network (RNN) model is the most commonly used model for the seq-to-seq modeling. The RNN can save past information of the time series as a hidden state, which enables it to model time series effectively. However, since this model recursively updates the current hidden state by using the previous hidden state, the previous information of the hidden state is not effectively transmitted due to the vanishing gradient if the distance between the sequences is large, resulting in the problem of long-term dependency.

Figure 2 Long Short Term Memory (LSTM) Structure

Long Short Term Memory (LSTM) has been proposed to resolve the aforementioned problem. In addition to the hidden state, LSTM adds the cell state, which allows it to effectively propagate the gradient like a conveyor belt. However, it does not completely resolve the problem of long-term dependency, and its learning speed is slow due to its sequential nature. Transformer model has been developed to address these RNN issues.

Transformer Overview

Figure 3 Transformer Structure

Transformer^[1] was proposed by Google in 2017, which was titled Attention is All You Need. As shown in Figure 1, the conventional seq-to-seq models consist of an encoder-decoder structure. The encoder condenses the input sequence into a single context vector representation, and the decoder generates an output sequence from this representation. When compressing the input sequence into a single representation, it is possible to lose the information of the input sequence. To resolve this issue, Transformer with the encoder-decoder structure employs self-attention only, and it demonstrates superior performance compared to the existing RNNs.

Figure 4 Bert Structure

Transformer has been also used as a backbone for the language models in the NLP fields, e.g., Bert and RoBerta.

Figure 5 Self-Attention Mechanism

The reason for the superiority of Transformer in the seq-to-seq modeling is due to self-attention. For instance, we have the following sentence: “The animal didn’t cross the street because it was too tired.”. When interpreting the sentence, it is important to figure out the meaning of it. With self-attention, all words corresponding to it are considered within the sentence, and it is learned to mean the “animal,” which is the most likely. Thus, self-attention enables the model to more effectively comprehend the meaning of words within sentences.

Figure 6 Computational Volume of Self-Attention

In self-attention, the representation of the input sequence is learned by using Query, Key, and Value, and the attention weight is computed by matrix multiplication of Query, Key, and Value. This attention weight determines to focus which part of the entire input sequence. As shown in Figure 6, the self-attention requires the computational complexity of $O\left(N^2\right)$, making it challenging to apply to various applications requiring long-sequence input, such as NLP and demand forecasting. Here, N is the length of the input sequence.

Figure 7 How to Replace Self-Attention Effectively^[4]

Accordingly, a number of methods have been proposed for effectively implementing self-attention, which can be divided into two major categories. The first approach is to approximate self-attention. Because most attention weights have close-to-zero values, it is more efficient to calculate only the important and necessary attention weights rather than all of the values. Typical examples include sparse attention^[2] and low-rank approximation^[3].

The second approach is to substitute self-attention with an effective structure. For example, the matrix multiplication of the self-attention is replaced with addition, or the self-attention uses a single query, such as Attention-free ^[4].

This blog will introduce three methods — Pyraformer^[5], cosFormer^[6], and PoNet^[7] — for implementing self-attention presented at ICLR 2022. Pyraformer and cosFormer are methods for approximating self-attention, whereas PoNet is for replacing the self-attention with an efficient structure via a pooling mechanism.

Pyramial Attention based Transformer (Pyraformer)

Figure 8 Signal Traversing Path of Various Seq-to-Seq Models

Let's first take a look at Pyraformer. In the seq-to-seq model, the sequence dependency can be better captured when the signal traversing path is shorter. For instance, in the CNN and RNN, when the length of the sequence is N, a signal traversing path of O(N) is required to capture all the input signal dependencies, as illustrated in Figure 8 (b) and (c). In Transformer (full attention), since the self-attention is performed between all input sequences for each layer, the maximum signal traversing path is O(1), as depicted in the figure. Pyraformer was proposed to decrease the computational complexity of Transformer by employing a compact multi-resolution structure while O(1) of the maximum signal traversal path is maintained.

Figure 9 Pyraformer Structure

The structure of Pyraformer is shown in Figure 9. The mainly different modules of Pyraformer from the conventional Transformer are Coarser scale construction module (CSCM) and Pyramidal Attention Module (PAM).

Figure 10 Coarser Scale Construction Module (CSCM) and Pyramidal Structure

As its name suggests, Pyraformer employs a pyramidal structure for compact multi-resolution representations. Figure 10 depicts how CSCM creates a C-ary tree structure (i.e., a pyramidal structure) by repeatedly applying convolution operations with stride C. In PAM, the self-attention is then performed within a limited range from the parent node  ${\mathrm{\mathbb{P}}}_{\mathcal{l}}^{\left(s\right)}$, the neighboring node ${\mathbb{A}}_{\mathcal{l}}^{\left(s\right)}$ on the same scale, and the child node ${\mathrm{ℂ}}_{\mathcal{l}}^{\left(s\right)}$ in the pyramid structure, as depicted in the figure. Specifically, the attention weight of node $n_{\mathcal{l}}^{\left(s\right)}$ is calculated as follows.

$y_i=\sum _{\mathcal{l}\in {\mathrm{\mathbb{N}}}_{\mathcal{l}}^{\left(s\right)}}\frac{\exp \left(\frac{{\mathit{q}}_{\mathit{i}}{\mathit{k}}_{\mathit{l}}^{\mathit{T}}}{\sqrt{{\mathit{d}}_{\mathit{K}}}}\right){\mathbf{v}}_{\mathcal{l}}}{\sum _{\mathcal{l}\in {\mathrm{\mathbb{N}}}_{\mathcal{l}}^{\left(s\right)}}\exp \left(\frac{{\mathit{q}}_i{\mathit{k}}_{\mathcal{l}}^T}{d_K}\right)}$ 

Here,  ${\mathrm{\mathbb{N}}}_{\mathcal{l}}^{\left(s\right)}={\mathrm{\mathbb{P}}}_{\mathcal{l}}^{\left(s\right)}\cup {\mathbb{A}}_{\mathcal{l}}^{\left(s\right)}\cup {\mathrm{ℂ}}_{\mathcal{l}}^{\left(s\right)}$. Therefore, Pyraformer computes a smaller number of Query-Key pairs than the conventional Transformer, and finally has computational burden ofO(N).

 

Based on the experimental results, Pyraformer calculates 5–6 times fewer Q-K pairs than the existing Transformer. Nevertheless, Pyraformer exhibited improved performance..

cosFormer

 

Figure 11 Linearized Self-Attention in cosFormer

The second method is cosFormer. As shown in Figure 11, the conventional self-attention calculates the attention weights by (a) multiplying matrices of Q (query) and K (key), (b) applying the softmax function, and (c) multiplying matrices of the result and V (value). Finally, the computational burden of the self-attention should be O(N²). On the other hand, cosFormer linearized self-attention by employing the decomposable similarity function $\varphi \left(·\right)$, as shown in fthe igure. Therefore, K (key) and V (value) are firstly calculated prior to Q (query), and the computational complexity can be reduced to O(N).

 

To achieve the linearization of the self-attention, cosFormer compared the performances of several decomposable linear functions $\varphi \left(·\right)$ and softmax function. Let's consider the identity, the leaky ReLU, and the ReLU functions as a decomposable linear function $\varphi \left(·\right)$. When evaluating the performance in terms of the loss and accuracy for three datasets (QQP, SST-2, and MNLI), the ReLU function demonstrated the best performance among the three considered functions. Interpreting this result, it can be seen that the ReLU function satisfies the non-negativity of the attention weight, whereas the identity function and the leaky ReLU function do not. This suggests that the non-negativity of the attention weight has a significant impact on performance. In addition, when comparing the performance of the ReLU function and the softmax function, it shows that non-linear re-weighting of the softmax function has a significant impact on performance as well.

Considering the above observations, cosFormer accomplished the linearization of the self-attention i) by exploiting the ReLU function for the non-negativity of the attention weight, and ii) by employing the cos function based on non-linear re-weighting. Specifically, we can first calculate the key and value prior to the query through the ReLU function as follows.

$\mathrm{S}\left(Q,K\right)=s\mathit{(}{\varphi }_{linear}\mathit{\left(}Q\mathit{\right)}\mathit{,}{\varphi }_{linear}\mathit{\left(}K\mathit{\right)}\mathit{)}\mathit{=}s\mathit{(}\mathrm{ReLU}\mathit{\left(}Q\mathit{\right)}\mathit{,}\mathrm{ReLU}\mathit{\left(}K\mathit{\right)}\mathit{)}$
$o_i=\frac{\sum _{j=1}^N\left(\mathrm{ReLU}\left(Q_i\right)\mathrm{ReLU}{\left(K_j\right)}^T\right)V_j}{\sum _{j=1}^N\left(\mathrm{ReLU}\left(Q_i\right)\mathrm{ReLU}{\left(K_j\right)}^T\right)}=\frac{\mathrm{ReLU}\left(Q_i\right)\sum _{j=1}^N\mathrm{ReLU}{\left(K_j\right)}^T{V}_j}{\mathrm{ReLU}\left(Q_i\right)\sum _{j=1}^N\mathrm{ReLU}{\left(K_j\right)}^T}$

In addition, the re-weighting of the softmax function can be replaced with the cos function-based non-linear re-weighting. The cos function is utilized because it can be divided into two summations using the sum formula for the cos function, making the linearization of the self-attention possible. The cos-based re-weighting is applied as follows.

$s\left(Q_i^{\text{'}},K_j^{\text{'}}\right)=Q_i^{\text{'}}{K_j^{\text{'}}}^{T}cos\left(\frac{\pi }{2}\times \frac{i-j}{M}\right)$    $Q^{\text{'}}=\mathrm{ReLU}\left(Q\right), K^{\mathit{\text{'}}}\mathit{=}\mathrm{ReLU}\left(K\right)$

In the equation above, the function of (i-j) is used to give more weight to nearby sequences than to distant ones. When further expanding the above equation using the cos sum formula, it can be as follows.

$$\begin{gathered} Q_i^{\text{'}}{K_j^{\text{'}}}^{T}cos\left(\frac{\pi }{2}\times \frac{i-j}{M}\right)=Q_i^{\text{'}}{K_j^{\text{'}}}^T\left(cos\left(\frac{\pi i}{2M}\right)cos\left(\frac{\pi j}{2M}\right)+sin\left(\frac{\pi i}{2M}\right)sin\left(\frac{\pi j}{2M}\right)\right) \\ =\left(Q_i^{\text{'}}cos\left(\frac{\pi i}{2M}\right)\right){\left(K_j^{\text{'}}cos\left(\frac{\pi j}{2M}\right)\right)}^T+\left(Q_i^{\text{'}}sin\left(\frac{\pi i}{2M}\right)\right){\left({K_j^{\text{'}}}^{T}sin\left(\frac{\pi j}{2M}\right)\right)}^T \end{gathered}$$ 

Therefore, we can finally obtain the following attention weight.

$$\begin{gathered} O_i=\frac{\sum _{j=1}^{ \\ N}{Q}_i^{\cos }\left({\left(K_j^{\cos }\right)}^T{V}_j\right)+\sum _{j=1}^N{Q}_i^{\sin }\left({\left(K_j^{\sin }\right)}^T{V}_j\right)}{\sum _{j=1}^N{Q}_i^{\cos }{\left(K_j^{\cos }\right)}^T+\sum _{j=1}^N{Q}_i^{\sin }{\left(K_j^{\sin }\right)}^T} \end{gathered}$$

Here, $Q_i^{\cos }=Q_i^{\text{'}}\cos \left(\frac{\pi i}{2M}\right)$, $Q_i^{\sin }=Q_i^{\text{'}}\sin \left(\frac{\pi i}{2M}\right)$, $K_j^{\cos }=K_j^{\text{'}}\cos \left(\frac{\pi j}{2M}\right)$, $K_j^{\sin }=K_j^{\text{'}}\sin \left(\frac{\pi j}{2M}\right)$.

 

Figure 12 Visualization of Attention Weights of Transformer and cosFormer

Figure 12 shows the visualization of the attention weight for Transformer and cosFormer. From the figure, it can be seen that (a) more weight is given near the diagonal through the cos function based re-weighting matrix, which confirms that more weight is given to sequences that are closer together than those that are further apart. (b) It can also be observed that the attention weight obtained by cosFormer is similar to that of the conventional Transformer. (c) Finally, it can be seen that the attention weight obtained by cosFormer with the re-weighting, is more similar to that by the conventional Transformert than that by cosFormer without re-weighting. It is shown from the experimental results that the ReLU function and the cos function based re-weighting can be replaced with the softmax function.

 

In addition, the effect of cos function-based re-weighting was validated through two experiments involving natural language — Bidirectional finetune and long-range-arena (LRA). As shown in the preceding table, the re-weighting of the cos function emphasizes locality, which has a significant impact on natural language tasks.

Pooling Network (PoNet)

Figure 13 Pooling Network (PoNet) Structure

Finally, we will examine PoNet, which replaces the self-attention structure with an efficient pooling network. PoNet uses three types of multi-granularity pooling mechanisms - global aggregation (GA), segment max-pooling (SMP), and local max-pooling (LMP) - to significantly reduce the computational burden. The GA, SMP, and LMP modules extract different levels of contextual information through the pooling network.

In GA, as depicted in Figure 13, the averaged representation of the entire input signal is first extracted as a single query via average pooling. Since the attention weight is computed using a single query, only the computational complexity of O(N) is required for the self-attention.

$\mathit{g}=\frac{1}{N}\sum _{n=1}^N{\mathit{h}}_Q$
${\mathit{g}}^{\text{'}}=\mathrm{Attention}\left\{\mathit{g},{\mathit{H}}_K,{\mathit{H}}_Q\right\}$

Next is SMP. GA has the disadvantage of losing significant information on the input sequence due to performing self-attention with just a single query. To compensate for this, SMP considers multiple sequences to be a single segment and performs max-pooling for each segment to extract important information at the segment level. In particular, as no matrix multiplication is required in SMP, no additional computation is required.

Similar to SMP, LMP extracts important information at the sequence level by applying max-pooling in sliding the windows between the sequence itself and adjacent ones. Like SMP, LMP has the advantage of no additional computation as max-pooling is employed.

PoNet combines the information extracted from GA, SMP, and LMP to obtain the final result P. Specifically, the representation g^' extracted from GA and the representation S_k(n) extracted from SMP are element-wisely multiplied with the original sequence H_on in order to extract various representations.

${\mathit{G}}_n={\mathit{g}}^{\text{'}}\circ {\mathit{H}}_{on}$
${\mathit{S}}_{\mathbf{n}}^{\mathbf{\text{'}}}={\mathit{S}}_{k\left(n\right)}\circ {\mathit{H}}_{on}$
$\mathit{P}={\mathit{G}}_n+{\mathit{S}}_n^{\text{'}}+\mathit{L}$

Due to element-wise multiplication, only O(N) of computation burden is required here, and PoNet needs O(N) of computation complexity only at each layer.

 

From the experimental results, we can see that in comparison to the conventional Transformer, PoNet with a max-pooling mechanism reduced learning speed and memory usage by up to 9 and 10 times, respectively.

Conclusion

This post provides a brief overview of Transformer and discusses three efficient ways to implement the self-attention of Transformer, which were introduced at ICLR 2022. Transformer has been widely utilized in various applications. In particular, how long the historical demand signal is used significantly influences the performance of the demand forecasting system. With such an efficient Transformer, we can expect to be able to more efficiently implement a demand forecasting system while maintaining the conventional Transformer's superior performance in the seq-to-seq modeling.