Long Sequence Modeling with Attention Tensorization: From Sequence to Tensor Learning (2024)

Abstract: As the demand for processing extended textual data grows, the ability to handle long-range dependencies and maintain computational efficiency is more critical than ever. One of the key issues for long-sequence modeling using attention-based model is the mismatch between the limited-range modeling power of full attention and the long-range token dependency in the input sequence. In this work, we propose to scale up the attention receptive field by tensorizing long input sequences into compact tensor representations followed by attention on each transformed dimension. The resulting Tensorized Attention can be adopted as efficient transformer backbones to extend input context length with improved memory and time efficiency. We show that the proposed attention tensorization encodes token dependencies as a multi-hop attention process, and is equivalent to Kronecker decomposition of full attention. Extensive experiments show that tensorized attention can be used to adapt pretrained LLMs with improved efficiency. Notably, using customized Triton kernels, tensorization enables Llama-8B training under 32,768 context length and can steadily extrapolate to 128k length during inference with 11 times speedup (compared to full attention with FlashAttention-2).

## introduction in the swiftly advancing field of natural language processing (nlp), large language models (llms) like gpt-4 and llama have emerged as vital tools, demonstrating expertise in comprehending and producing human language. these complex scenarios often involve context lengths longer than those llms were trained on, posing significant challenges to transformer-based architectures in handling long sequences. for example, llms are expected to read long paragraphs or books to answer questions while they are usually trained with a much smaller context length (like 8k for llama-3 @xcite ). scaling up llms for long sequences presents significant challenges related to limited attention windows during training and length extrapolation during inference. first, transformers usually adopt a bounded context window in the training phase due to quadratic computation costs @xcite . to mitigate these costs and expand the training context window, existing works usually adopt attention approximation methods like sparse, low-rank, and softmax-free attention. despite their efficiency and scalability with relatively long sequences, these methods gain less popularity due to inefficient implementations and incompatibility with pre-trained llms widely used in community. given such open-sourced pretrained llms, researchers start to focus more on the length extrapolation challenge: leveraging the short-range modeling power acquired during pre-training to handle unseen long-range dependencies. addressing this challenge requires methods that take into account both efficiency and effectiveness perspectives. one line of thought is to employ full attention with hardware-efficient implementations such as flashattention and quantization for efficiency, paired with positional extrapolation @xcite or interpolation @xcite for enhanced performance. while full attention captures all potential correlations, it significantly increases running time, and many correlations in long sequences are in fact unnecessary and distracting. the second line of approaches involves using segmented or sliding windows for efficiency, supplemented by additional recurrent memory modules @xcite or global sinks @xcite to integrate segments. however, it remains difficult for these additional modules to retain past or global information as windows move and update, posing a persistent challenge in processing long sequences effectively. proposed work. to address these challenges in modeling long-range dependencies, we propose folding the one-dimensional long-range token interactions into a higher-order tensor, where each dimension can be effectively modeled with shortrange interactions. as shown in figure 1 (a), given a limited context window, tensorization reduces the interaction distance between two tokens, converting previously out-of-window interactions into within-window interactions. inspired by the compact representation of the input tensor, we generalize input at the attention layer from the conventional sequence to higher-order tensor format and propose a supporting tensorized attention mechanism that replaces traditional vector operations with corresponding tensor operations. unlike traditional attention approximations, tensorized attention is implemented with an efficient custom triton kernel and can be adapted from pretrained llms through continued pretraining. besides, the multi-dimensional encoding of token interactions provides tensorized attention with natural advantages for length extrapolation. given finite attention window size or memory budget, tensorized attention captures correlations along different dimensions and incorporates hierarchical multi-hop interactions in the original sequence at multiple scales as shown in figure 1 (b), which exponentially extends the effective context length and avoids the time consumption of using full attention. compared to segmented and sliding windows, tensorized attention hierarchically merges interactions from low to high levels and does not need additional global or recurrent modules to retrieve past information under exponential extrapolation speed. we then integrates tensorized attentions into llms to improve the model scalability to longersequence inputs. experiments with diverse nlp datasets show that tensorized attention improves the performance and efficiency of existing pretrained models on long-sequence tasks after continued pretraining. notably, tensorized attention enables llmama-3-8b training under 32,768 context length and can steadily extrapolate to 128k length during inference with 11× speedup (compared to full attention with flashattention-2). ## attention tensorization in this section, we first compare the representation of input sequences in both vector and tensor space. to accommodate the tensorized input, we then introduce attention tensorization by changing conventional attention from the vector space to the corresponding tensor space. finally, we show that low-rank approximation of attention can be more efficient in tensor space compared to vector space, both empirically and theoretically. ## experiments in this section, we conduct various experiments to demonstrate that tensorized attention can serve as an efficient backbone for llms in long-sequence applications. we first discuss the continued pretraining process, followed by an evaluation of its performance on downstream tasks and its extrapolation capabilities. additionally, we showcase the strong performance of tensorized attention in conventional nlp and time series models, as detailed in the appendix d. ## conclusion in this paper, we explore attention tensorization by replacing vector/matrix operations with corresponding tensor operations. from an efficiency perspective, decomposing long sequences into individual dimensions allows the limited context length 0 1 6 k 3 2 k 4 8 k 6 4 k 8 0 k 9 6 k 1 1 2 k 1 2 8 k length budget to be utilized more effectively for attention calculation along each dimension. from a performance perspective, attention tensorization naturally encodes multi-hop hierarchical token interactions, making attention more efficiently approximated compared to its representation in vector space. extensive results on datasets from various domains demonstrate that tensorization can extend the training context length, enhancing both efficiency and length extrapolation performance.