LongVQ: Long Sequence Modeling with Vector Quantization on Structured Memory

1 Introduction

Researchers and practitioners have been focusing on longrange dependency learning in sequence modeling. To tackle this challenge, popular sequence models like LSTM [Hochreiter and Schmidhuber, 1997] and GRU [Chung et al., 2014] have been introduced. These models can capture the relationships among input tokens (such as words [Ainslie et al., 2020], pixels [Dosovitskiy et al., 2021], and speech signals [Pascual et al., 2019]) during the learning process. By utilizing contextual information, these models can learn intricate patterns of token dependencies in training samples. However, traditional sequential models are inefficient due to their non-parallelizability and limited memory. Consequently, scaling up the model parameters fails to enhance performance and significantly slows down training and inference speed. In response to the limitations of traditional sequence models, the Transformer [Vaswani et al., 2017] architecture was introduced to achieve faster training and SotA performance. A key innovation enabling their parallelizability and superiority is the self-attention mechanism with positional encoding, which explicitly learns dependencies between all combinations of input tokens with order information encoded, aiming to create flexible kernels for sequences.

Corrsponding author.

Transformers have revolutionized natural language processing (NLP) and achieved significant breakthroughs. However, they have a major drawback when dealing with very long input sequences that may contain millions of tokens. That is a quadratic O(L2 ) complexity relative to the input sequence length in the self-attention mechanism, resulting in rapid scaling of compute and memory needs with sequence length L. Thus, many researchers in recent years focused on improving Transformer efficiency. Proposed techniques include approximating the costly self-attention matrix through various means like sparsity [Tay et al., 2020a; Liu et al., 2022] or conditional probability assumptions [Ren et al., 2021]. Other innovations like matrix factorization, optimized memory access patterns, and architectural constraints have also been explored [Wang et al., 2020; Beltagy et al., 2020; Lingle, 2023]. While all of these operations reduce the complexity of the model, at the level of expressiveness, transformers can easily get lost in long sequence data. The middle of Fig. 1 demonstrates attention will fail in the long sequence scenarios, e.g., resulting high-entropy attention matrix. In parallel, wholly novel models have been introduced to better handle long sequences. Models like applying linear recurrent units, which integrate state space models (SSMs) and polynomial projections, to reconstruct latent histories and demonstrate promising performance, especially on long-sequence tasks [Gu et al., 2020b; Gu et al., 2021; Gu et al., 2022; Fu et al., 2023a; Fu et al., 2023b; Poli et al., 2023]. As shown in Fig. 1(left), it has to be admitted, however, that the SSMmodel is far less capable of screening and selecting information than models with attention mechanisms [Fu et al.,

2023a].

Therefore, in order to combine the respective advantages of the SSMs and Transformers, some hybrid models have been proposed recently and also achieved good results on long sequence tasks [Ma et al., 2022; Zuo et al., 2022]. To accommodate the computational complexity of the SSMmodel, linear attention is generally used as a collocation. Nevertheless, we can clearly see in the experimental results that the overall performance has a noticeable degradation after switching self-attention to linear attention (see right of Fig. 1). The question that needs to be addressed is: is it possible to combine the advantages of both and match the performance of a hybrid model with self-attention to linear complexity?

To answer this question, we present Long Sequence Modeling with Vector Quantization (LongVQ), an attention-SSM hybrid model in linear time with respect to the sequence length. Our method combines vector quantized keys, localized positional biases, and a dynamic fixed-size memory mechanism to store and retrieve complex data structures in long sequences efficiently. It has transformed how we approach data analysis and processing and unlocked new possibilities for various fields, including natural language processing, speech recognition, and image recognition. We benchmarked LongVQ on five popular datasets, including Long Range Arena (LRA) [Tay et al., 2020b], which has diverse and challenging data modalities, the image dataset sCIFAR [Krizhevsky et al., 2009], the natural language datasets WikiText-103 [Merity et al., 2016] and enwik8, and the speech data SpeechCommand [Warden, 2018]. In these datasets our method demonstrates promising performance.

2 RelatedWork

2.1 Efficient Transformers

Kernel method and clustering Kernelizable attention, a technique introduced in studies by [Katharopoulos et al., 2020], involves computing query and key features and then applying the same nonlinearity to both of them separately. It omits additional nonlinearities when computing attention weights. Clustering attention, as described by [Vyas et al., 2020], involves using VQ queries and can be kernelized.

Hierarchical Attention Combiner, proposed by [Ren et al., 2021], approximates softmax using a graphical model and enables decoder-only self-attention. H-Transformer-1D [Zhu and Soricut, 2021] reduces the complexity of encoder-only self-attention. Transformer-LS [Zhu et al., 2021] downsamples long-range features in Transformers.

Compressed Attention Compressive Transformers [Rae et al., 2019; Ainslie et al., 2020; Zhu et al., 2021], learn a compression function for long-range features directly. On the other hand, LUNA [Ma et al., 2021] uses cross-attention to compress long-range features into a recurrent state. Our LongVQ differs from compressive/recurrent Transformers in that it equates to quadratic-time attention over vectorquantized keys. This means that if the keys are already vectorquantized, the LongVQ cache reduces the cost to linear time without any loss. In contrast, Perceivers [Jaegle et al., 2021] use cross-attention to attend to long sequences and compute self-attention over only a narrow stack of "latents." However, LongVQ computes dense self-attention in linear time instead of just cross-attention. Therefore, while Perceivers' long-range layers incur a quadratic time complexity during sampling, LongVQ generates sequences in linear.

2.2 State Space Models

Recurrent neural networks and their linear counterparts, such as state-space models, can retain memory of the past. Among them, S4 [Gu et al., 2021] is notable because it can be implemented through convolutions thanks to its linear recurrence. However, the long convolution kernel for this model is as long as the input sequence, and its efficient computation requires sophisticated parameterization and approximation techniques. Although recent advances have found solutions to this issue, initializing these models still requires special effort [Gupta et al., 2022; Gu et al., 2020b]. Many of these models use the HiPPO [Gu et al., 2020a] initialization mechanism, which aims to memorize historical data through projection to orthogonal polynomials. Based on a structure similar to SSM, an increasing number of works focusing on either linear recurrence or global convolution have been developed recently [Fu et al., 2023b; Fu et al., 2023a; Poli et al., 2023].

2.3 Vector Quantization

Various techniques such as k-means, vector quantization, and codebooks have been used in Transformers for various applications aiming to learn semantic discrete tokens [Roy et al., 2020]. These techniques typically involve codebooks or similar structures within a Transformer architecture. Some models have proposed using a codebook outside the Transformer, such as when Transformers are priors for VQ-VAEs [van den Oord et al., 2018]. LongVQ utilizes one codebook within each layer and performs dense self-attention with SSMs in linear time, which distinguishes it from all of the aforementioned works.

2.4 Data-dependent Gating

Gated attention is a fusion of attention sublayers and GLUbased MLP sublayers, which was introduced in FLASH [Hua et al., 2022b]. Various gating mechanisms have been previously used to stabilize the training of Transformers and other sequence models. Examples of such models include GSS [Mehta et al., 2022], MEGA by [Ma et al., 2022], and RWKV by [Peng et al., 2023].

3 Preliminary

3.1 Self-attentionMechanism

Given the input to the layer is X 2 RL⇥d , where L is the sequence length and d is the embedding dimension, then the self-attention mechanism outputs pair-wise score M:

Here Wq↪Wk↪Wv 2 Rd⇥d are learnable matrices, and σ is the self-attention function e.g., softmax. The self-attention mechanism facilitates the computation of pair-wise relationships of all input tokens as cache, thereby enabling the modeling of long-range dependencies that surpasses the capabilities of recurrent neural networks. For example, denote the attention score M = softmax(QK>◁pd) 2 RL⇥L. Then, Mij computes the similarity between the i-th and the j-th tokens.

3.2 Linear Recurrent Models

Continuous linear state-space model. The center of the LSSL layer is a linear continuous-time SSM which maps a 1-dimensional function x(t) to y(t) through the hidden state h(t) 2 RN⇥1 . Concretely, LSSL can be formally defined as:

where A 2 RN⇥N ↪B↪C 2 RN⇥1 ↪D 2 R1⇥1 . The SSM can also be viewed as a convolution over the input signal.

Many researchers use Eq. 2 to model long sequences. For instance, [Gu et al., 2020a] argue that randomly initialized parameters A, B, and C aren't effective in modeling longrange dependencies. To address this issue, they propose a set of matrices called HiPPO (high-order polynomial projection operators) for initializing A. The HiPPO matrices are designed in such a way that the state x(t) at time t can remember the history of the input u(t) up to time t.

Discrete-time form. The continuous-time LSSL in 2 needs to be discretized for real-world data. Using the bilinear method and a step size of ∆:

We unroll the above recurrent representation, then we have:

Convolution form. This can be written as a convolutional representation y = K ⇤ u, where the convolution kernel

Here, the SSM transformation can be viewed as embedding mapping from input to output by using a long convolution kernel. "⇤" is the discrete convolution operator, u represents the input, and y represents the corresponding output. In Eq. 3, the output y can be efficiently computed if the convolution kernel K is known. However, computing the kernel can be challenging, with most existing algorithms having O(L2 ) time and space complexity.

Structured State Space SequenceModel [Gu et al., 2021] develop the Structured State Space Sequence (S4) model to efficiently compute Eq. 3. Specifically, D is the residual connection, C in Eq. 2 is randomly initialized, and A and B are initialized as structured matrices as following:

The convolution can be computed efficiently with Fast Fourier Transforms (FFTs) when the kernel K is known. However, for sequence length L, computing the kernel K takes O(LN 2 ) operations for unstructured ¯A. The paper proposed an efficient albeit complex algorithm for a class of structured ¯ Awith ˜O(N +L).

4 Sequence Modeling with Vector Quantization on Structured Memory

This section describes our proposed model, long sequence modeling with vector quantization (LongVQ). We first introduce the linear-time VQ attention equipped with the SSM model. Then, we present the whole block of LongVQ, including feed-forward and normalization layers. At last, we provide the network blocks and training algorithm of LongVQ.

4.1 Vector Quantization

Vector quantization (VQ) is a widely used information compression technique [Lingle, 2023]. In our work, VQ will be employed in self-attention layers to ensure linear complexity and make SSM operators more expressive.

What can VQ bring to SSM-attention hybrid? As discussed in Sec. 1, SSM and self-attention mechanisms each have their own limitations despite their wide applications and impressive successes in sequence modeling. The hybrid models enjoy both side benefits by leveraging their properties to complement each other. However, the linear attention, e.g., chunk-wise attention, degrades the performance dramatically [Ma et al., 2022]. The linear mixing methods designed artificially may lose the global view. VQ is a two-way solution that retains linear complexity and dynamically compresses global information to address this issue. The resulting model benefits from strong inductive bias while also being able to efficiently learn complex dependency patterns.

Vector Quantizers and Training A vector quantizer is a function VQ(·; C) compresses a continuous input vector x to a discrete embedding space ˆx 2 RS⇥D with fixed size:

where C 2 RS⇥D is known as the codebook. The row indices ¶0↪ ▷ ▷ ▷ ↪ S - 1♦ of C are called shortcodes, and the rows themselves are called codewords. Because of the non-differentiable nature of the quantization operation, the straight-through estimator is typically used to update the network and codebook in an end-to-end manner:

where SG(·) denotes for stop-gradient operator.Since the computation of attention varies with the length of the sequence, linear complexity can be achieved as long as we can compress the K to a constant length.

4.2 LongVQLayer

The gated attention method in LongVQ is inspired by MEGA [Ma et al., 2022] and Gated Attention Unit [Hua et al., 2022a] as the basic block in architectures. With an SSM layer embedded to capture long-term patterns, the input of our Gated VQ attention G↪ Q↪ K↪ V can be computed as follows:

where Z can be regarded as the structured memory compresses the whole sequence. Ga is the gating variable that empowers the expressiveness of attention. Note that φsilu is the self-gated activation function (SiLU) [Ramachandran et al., 2017]. The output of Gated Attention is:

where ⊙ denotes the element-wise production.

Linear-time VQ Attention Note that the above computation of attention is still in quadratic complexity. Recall the vector quantization process in Sec.3. By following [van den Oord et al., 2018],we could compress the K into a dynamic codebook with a fixed length:

where STVQ(·;C) represents the row-wise application of vector-quantization with a straight-through gradient estimator. Thus, the attention score can be refactored as:

where ∆ could be regarded as Kronecker delta function that builds a mapping between each token in K and code in C. Here, we take the element-wise activation function as an example, i.e., φelement = relu2 (·). The resulting attention computation achieves linear-time complexity for efficient sequence modeling.

Local window If C ⌧ L, the code will appear uniformly throughout the sequence and will also cause the model to lose the ability to distinguish between near and far. To address this, a local bias B is added to the VQ attention for enhancing local information:

Moreover, we add an output gate Go after the attention to control the final output O dependent on the input X:

Fig 2 middle plots the block design of LongVQ architecture.

4.3 LongVQBlocks

The LongVQ layer can be used instead of regular selfattention in the Transformer model. Once added, it is followed by position-wise feed-forward networks (FFNs) and normalization layers to create a LongVQ block. As the gated residual connection is already included in equation 20, the original residual connection is omitted, and a normalization layer is directly applied to Y :

where Y 0 is the output of a single LongVQ block. The overall LongVQ block is shown in Fig 2 left.

Training loss Let ✓ denote the set of non-codebook parameters, and C denote the set of layers' codebooks. For modeling sequences X for classification, our loss can be defined:

where γ is a hyperparameter known as the commit loss coefficient. The LCE is the classification loss, and the LVQ is employ mean squared error (MSE) between K and C:

The main cross-entropy loss is augmented with commitment losses to update the codebook across all layers. As [van den Oord et al., 2018], codebooks are parameterized through smoothed quantizer statistics.

5 Experiments

To evaluate LongVQ, we conduct experiments on five benchmark sequence modeling tasks across various data types, comparing them with current state-of-the-art models on each task. In all cases, we use codebook size 512. All experiments were realized based on NVIDIA A100-80G and Pytorch. We used float32 parameters, with bfloat16 precision for most computations. We adopt AdamW as the optimizer with a gradient clip of 0.1. The codebook commit coefficient was always γ = 0▷0001, and the codebook EMA rate was always ⌘ = 0▷99. All models were trained with a global batch size of 128 sequences. In our table, bold indicates the bestperforming model and underlines the second-best.

5.1 Pixel-Level Sequential Image Classification

Let's begin with tasks related to image classification, where images are treated as a sequence of pixels in one dimension. In such tasks, models are not allowed to use any two-dimensional bias from the image. Therefore, the model should be capable of identifying patterns at various timescales, including pixels that are close to the original image but far from each other in its sequence representation. We assess our model's performance on the Sequential CIFAR-10 dataset, which is widely used as a standard benchmark for modeling long-term dependencies in RNNs. The CIFAR-10 dataset's standard train and test split is used, and 10% of the training set is withheld as the validation set. To accomplish the classification task, we compute the mean of all tokens in the output sequences and feed the result into a fully connected layer to generate class logits. The test accuracy is reported based on the model with the highest validation accuracy.

Results The Table 1 displays the results. Our model has achieved state-of-the-art performance and the best test accuracy on the sequence classification task, surpassing numerous strong competitors such as Transformers [Vaswani et al., 2017], RNNs, state space models, and other convolutional models. Specifically, the LongVQ model has outperformed the previous convolution-based models by more than ten percentage points. It is noteworthy to mention that our model has achieved impressive results by surpassing the previously established performance benchmark, despite utilizing a relatively simple architecture. The model primarily uses a hybrid approach that compresses long historical information based on the output of SSM.

5.2 Long-Context Sequence Modeling

We have started conducting experiments by evaluating sequence models on the Long Range Arena (LRA) benchmark, which was recently introduced by [Tay et al., 2020b]. The benchmark aims to evaluate the performance of sequence models in long-context scenarios. The benchmark consists of a total of six tasks, each of which is designed to test different aspects of long-sequence modeling. These tasks include ListOps, Text, Retrieval, Image, Pathfinder, and Path-X. To ensure a comprehensive evaluation of our model, the tasks involve input sequences of varying lengths, ranging from 1K to 16K tokens. Additionally, the tasks cover a diverse range of data types and modalities, meaning that models must be capable of handling text, images, and other forms of input in order to perform well on the benchmark.

Results In Table 2, LongVQ is compared to various baselines, such as Transformer and its efficient versions, as well as the top-performing S4 models. Specifically, LongVQ outperforms the second-ranked model by 1.84%, 0.79%, 0.20%, 0.45%, 0.28%, 3.63% on each of the six datasets. To ensure a just comparison, we equalize the number of layers and model dimensions on each task so that LongVQ has a similar number of parameters with S4. The results are an average of five runs with distinct random seeds, and the tuning information and model details are given in the Appendix. The LongVQ has shown remarkable performance over all six tasks, with an average performance of 88.36%, outperforming all the other baselines. Furthermore, we have evaluated the speed and memory efficiency of LongVQ using the byte-level classification task with an input length of 4K. Our VQ mechanism has proven to be highly efficient, being 5.3 times faster and consuming only 10% as much memory as the vanilla Transformer. It is worth noting that the LongVQ layer, with its VQ-equipped hybrid design, is even more efficient than both Transformers and State Space Models by a large margin.

5.3 Auto-Regressive Language Modeling

By following [Ma et al., 2022; Lingle, 2023], we evaluate LongVQ on two established language modeling benchmarks, i.e., WikiText-103 [Merity et al., 2016] and enwik8, which are next-token prediction tasks. WikiText-103 is a word-level language modeling dataset containing 103M training tokens from Wikipedia articles. Following previous work [Baevski and Auli, 2019], our approach involves the utilization of adaptive softmax and input embeddings, and we work with a vocabulary of 260K tokens. enwik8, on the other hand, is a challenging benchmark for character-level language modeling. It comprises 100M tokens of unprocessed Wikipedia articles and has a vocabulary size of around 200. We divide the test data into segments and sequentially handle each segment during testing.

Results In Table 3, we compare with previous topperforming models that are designed to take advantage of longer context, including Transformers [Baevski and Auli, 2019], Transformer-XL and S4 [Gu et al., 2021]. The LongVQ model has shown impressive performance on both WikiText-103 and enwik8 datasets, surpassing the baseline models by a considerable margin. Additionally, our model is able to achieve a faster (almost 10⇥) inference speed compared to the Pure Transformer model. One of the key benefits of the hybrid design of the SSM layer and VQ technique is that it enables LongVQ to naturally handle length extrapolation at inference time, allowing it to process longer sequences than those seen during training.

5.4 Raw Speech Classification

We aim to assess LongVQ's ability to model speech signals over a long range by utilizing it to classify unprocessed speech signals that have a length of 16000 instead of relying on conventional preprocessing methods, such as converting them to MFCC features. As per [Gu et al., 2021] approach, we conduct speech classification on the Speech Commands dataset's SC10 subset, which was introduced by [Warden, 2018]. As reported in [Ma et al., 2022], the Mega-chunk uses a chunk size of 1000 to enable processing the data.

Results As shown in Table 5, our LongVQ model with 542K parameters can achieve an accuracy of 97.67, which is the state-of-the-art method in this table. This is mainly because the large number of continuous and low-frequency signals in speech are well-suited to be processed using VQ. Coupled with SSM's ability to extract global information, it is only natural that such results are achieved.

5.5 Ablation Study

We focus our analysis on LongVQ's important components to answer two questions: (1) Is it possible to set parameters for data types in the codebook? (2) What is the best attention function for LongVQ?

Codebook size We take the subset of LRA to conduct the ablation study. As shown in Table 6, larger codebook sizes may allow more flexible attention patterns and improve the gradients' fidelity, and improve final performance. However, we can set the size of the codebook according to the resource requirements based on the data characteristics. For example, from images to logical reasoning, images with lower information density can use a lower size and vice versa.

Attention function We tested two attention functions, softmax and laplace, to evaluate performance. Based on LongVQ, our findings differ from Mega's [Ma et al., 2022]. Due to its compressed codebook on K, we found that LongVQ prefers softmax as the attention function for different data types; the results are shown in Table 6. This is mainly due to the existence of VQ technology, which dynamically compresses the data into codes, and softmax, which has the ability to filter information with strong sparsity, is more suitable for the data after compression. Hence, our LongVQ does not need parameter search in the choice of attention function.

6 Conclusion and Limitations

We presented LongVQ for robust and efficient memorization of long-term patterns in sequential data sources. LongVQ is an SSM-attention hybrid architecture that computes dense self-attention in linear time with respect to sequence length. Its superior performance is enabled by considering the global view of the SSM module and a gating mechanism. The vector quantization method to compress keys is the crucial point to making LongVQ more efficient when handling long sequences. By using the local window design, LongVQ can still capture long-range patterns even with small codebook sizes. Our large-scale and diverse experiments demonstrate that LongVQ is an efficient and flexible long sequence model with excellent performance on image, text, logical reasoning, and speech data. However, LongVQ also has some obvious limitations in empirical implementations. LongVQ has a relatively large training difficulty, and due to the online training codebook, we need to handle the gradient size of different modules carefully. Moreover, the exponential moving average (EMA) technique is essential for stabilizing codebook updates, but it also limits the update speed simultaneously. Whether a better-stabilized training model can be used on codebooks is still a question worth considering.

Acknowledgements

This work was supported by the National Key R&D Program of China (No. 2022ZD0115100), the National Natural Science Foundation of China Project (No. U21A20427), and Project (No. WU2022A009) from the Center of Synthetic Biology and Integrated Bioengineering of Westlake University. This work was done when Li Wang and Zedong Wang interned at Westlake University. We thank the AI Station of Westlake University for the support of GPUs.