a linear transformer based on attention pooling

a linear transformer based on attention pooling

Abstract

Currently, transformer-based neural networks have dominated the field of artificial intelligence, achieving widespread applications in natural language processing, computer vision, image generation, and multimodal domains. Especially since OpenAI (reference) released ChatGPT in 2023, GPT has become the undisputed leader in the language domain. However, from the perspective of first principles, there are significant differences between GPT and human thinking. This is mainly reflected in the following aspects: First, GPT needs to store all historical information, while humans have the ability to forget; second, the complexity of GPT is $O(N^2)$, while human reading complexity is $O(N)$. Given these differences, applying GPT to lifelong learning domains still faces several challenges. This paper proposes a linear transformer model based on attention pooling, aiming to achieve linear complexity and forgetting capabilities while retaining the core attention mechanism of GPT. In summary, this paper introduces a linear transformer model based on attention pooling to achieve the aforementioned goals.

Introduction

Due to its ability to capture long-range dependencies between words, the GPT architecture has achieved great success in natural language processing in recent years. From the mainstream solutions adopted by major companies, it is rare to see models like RNNs and LSTMs being used. RNNs and LSTMs have significant drawbacks in training and deployment, mainly due to their inability to parallelize training and their catastrophic forgetting issues. Although the GPT architecture has achieved such great success, it still has several limitations. First, the principles of GPT are fundamentally different from how human intelligence works. Human intelligence includes a forgetting mechanism, while GPT does not. While the lack of a forgetting mechanism might make GPT more powerful, it poses a fatal problem in lifelong learning, mainly due to the increasing size of stored historical information and the growing computational complexity of retrieving this information. Both factors are positively correlated with inference time.

To address these issues, compressing historical information is an unavoidable step. This paper proposes a attention pooling method to compress historical information for querying by the current token. This method supports both parallel training and causal inference.

In summary, the contributions of this paper are as follows:

1. Proposes a linear transformer named WQKV, which uses attention pooling to compress historical information.
2. Presents parallel training and recursive causal inference methods for WQKV.
3. Conducts ablation experiments on WQKV.

Related Work

There has been significant work in reducing the complexity of transformers from $O(N^2)$ to $O(N)$.

RWKV is inspired by Apple’s Attention Free Transformer. The RWKV architecture is carefully simplified and optimized to be converted into an RNN. Additionally, many techniques, such as TokenShift and SmallInitEmb, are used to make RWKV perform comparably to GPT.

In “Were RNNs All We Needed?”, minLSTM and minGRU are proposed by removing nonlinear factors that hinder parallel computation. These RNNs can perform recursive inference and parallel training using the parallel scan algorithm.

However, these methods have some drawbacks. To support large contexts (state variables), they often require a very large number of parameters. Our method does not have this requirement. Moreover, these methods lack a clear connection to the attention mechanism of transformers, and their effectiveness has not been widely validated. In contrast, our method fully inherits the attention mechanism of transformers, with the addition of a attention pooling layer. Therefore, the scalability of our method is guaranteed.

Attention pooling.

Method

A straightforward approach to reducing the complexity of GPT from $O(N^2)$to $O(N)$ is to compress the queried historical information into a fixed-length vector, so that the current token queries a fixed-size context. Clearly, this process is similar to pooling. Therefore, we designed a softmax-weighted pooling method to compress historical information into a fixed-length vector, followed by the original GPT architecture for attention operations. Below, we describe how this is achieved.

For simplicity, this paper assumes a batch size of 1 and does not use multi-head attention.

The attention pooling used in this paper is described as follow pictures.

Consistent with GPT, for the current sequence $X$ with shape $[L,C]$, we obtain $Q,K,V$ with shape $[L,C]$ using matrices $W_q,W_k,W_v$. The specific calculations are as follows:

$$\begin{gathered} K=W_{k}X \\ V=W_{v}X \end{gathered}$$

In addition, we use an extra weight matrix $W_w$ to obtain $W=W_{w}X$, with shape $[L,D]$. This weight matrix is used to compress historical information. The attention pooling results are:

$$\begin{gathered} K^{\prime }=\text{softmax}(W^T)K \\ {V}^{\prime }=\text{softmax}(W^T)V \end{gathered}$$

Thus, the shape of $K,V$ is compressed from $[L,C]$ to $[D,C]$.

After the sequence information is compressed, the subsequent operations are similar to GPT, using $Q$ to perform attention on $[K^{\prime },V^{\prime }]$, i.e.,

$$X^{\prime }=\text{softmax}(Q^T{K}^{\prime })V^{\prime }$$

From the above formulas, the new WQKV model is highly similar to GPT. The difference is that the new model adds a pooling module to compress information. Clearly, the WQKV model includes the transformer model, as long as the historical information is consistent with the full sequence information. When the memory dimension $D$ is reduced, the model gains the ability to forget. Compared to the $O(N^2)$ complexity of GPT, the new model has a complexity of $O(ND)$.

The above derivation assumes that each token can see all other tokens, equivalent to the encoder layer of a transformer. For infinitely long sequences, it is necessary to derive the parallel training method and recursive inference method for the causal form. Fortunately, the WQKV model also has a corresponding causal inference form. Below, we present its derivation.

From the previous derivation, we note that the computation of $K^{\prime },V^{\prime }$ has a recursive form. For convenience, we denote $\text{exp}(W^T)$ as $\left(\begin{array}{cccc}{w}_1& w_2& ...& w_l\end{array}\right)$, and $K$ as $\left(\begin{array}{c}{k}_1\\ k_2\\ ...\\ k_l\end{array}\right)$. We rewrite $K_l^{\prime }$ as follows:

$$\begin{gathered} K_l^{\prime }=\frac{\left(\begin{array}{cccc}{w}_1& w_2& ...& w_l\end{array}\right)}{{\sum }_i^l{w}_i}\left(\begin{array}{c}{k}_1\\ k_2\\ ...\\ k_l\end{array}\right) \\ =\frac{\left(\begin{array}{cccc}{w}_1& w_2& ...& w_l\end{array}\right)\left(\begin{array}{c}{k}_1\\ k_2\\ ...\\ k_l\end{array}\right)}{{\sum }_i^l{w}_i} \\ =\frac{{\sum }_i^l{w}_i{k}_i}{{\sum }_i^l{w}_i} \\ =\frac{{\sum }_i^{l-1}{w}_i{k}_i+w_l{k}_l}{{\sum }_i^{l-1}{w}_i+w_l} \end{gathered}$$

Let $M_l={\sum }_i^l{w}_i{k}_i,N_l={\sum }_i^l{w}_i$, then the following recurrence relation exists:

$$\begin{gathered} K_l^{\prime }=\frac{M_{l-1}+w_l{k}_l}{N_{l-1}+w_l} \\ {K}_{l-1}^{\prime }=\frac{M_{l-1}}{N_{l-1}} \end{gathered}$$

From the above recursive formula, if parallel training is desired, cumulative sums must be used, which is similar to existing methods.

Experiments

There is a clear linear relationship between the number of tokens memorized in the context and performance.

Future Work

This linear transformer network based on attention pooling can be reinterpreted from the perspective of cross-attention.

For a text sequence $L1$ with shape $[L,C]$, we add a random weight sequence $W$ with shape $[D,C]$ to form a new sequence $L$ with shape $[L+D,C]$.

The basic principle of the cross-attention-based linear transformer is to perform cross-attention repeatedly on $L1$ and $L2$.

The formula for $L2$ performing cross-attention on $L1$ is:

$$W^{\prime }=\text{softmax}(WL{1}^T)L1$$

The formula for $L{1}^{\prime }$ performing cross-attention on $L2$ is:

$$L{1}^{\prime }=\text{softmax}(L1W^{\prime T})W^{\prime }$$

The overall formula is:

$$L{1}^{\prime }=\text{softmax}(L1(\text{softmax}(WL{1}^T)L1{)}^T)\ast \text{softmax}(WL{1}^T)L1$$

From the above simple transformation, it can be seen that this is essentially equivalent to the attention pooling method. The connection between these two methods is worth further exploration.

Limitations:

To achieve parallel training, GPU memory usage is excessively high When use our method.

REFERENCE