Chunk, Align, Select: A Simple Long-sequence Processing Method for Transformers (2024)

Abstract: Although dominant in natural language processing, transformer-based models still struggle with long-sequence processing, due to the computational costs of their self-attention operations, which increase exponentially as the length of the input sequence grows. To address this challenge, we propose a **Sim**ple framework to enhance the long-content processing of off-the-shelf pre-trained transformers via three steps: **C**hunk, **A**lign, and **S**elect (SimCAS). More specifically, we first divide each long-sequence input into a batch of chunks, then align the inter-chunk information during the encoding steps, and finally, select the most representative hidden states from the encoder for the decoding process. With our SimCAS, the computation and memory costs can be reduced to linear complexity. In experiments, we demonstrate the effectiveness of the proposed method on various real-world long-text summarization and reading comprehension tasks, in which SimCAS significantly outperforms prior long-sequence processing baselines. The code is at [https://github.com/xjw-nlp/SimCAS](https://github.com/xjw-nlp/SimCAS).

## introduction transformers @xcite have become a fundamental model architecture for sequential data modeling @xcite @xcite , especially in natural language processing (nlp) @xcite , where texts are regarded as sequences of tokens. built with transformer blocks, pre-trained language models (plms) have recently shown astonishing empirical performance in various nlp tasks such as question answering @xcite , controllable generation @xcite , summarization @xcite and logic reasoning @xcite . however, one fatal weakness, that has hindered transformer-based models from being applied in broader application scenarios, is the quadratically raised computational consumption of self-attention operations when the input length increases. hence, vanilla transformers have continuously been challenged by long-context tasks, such as machine reading comprehension @xcite @xcite and longtext summarization @xcite . to enhance transformers with more efficient long-sequence processing, prior works focus on two perspectives, efficient attention operations @xcite @xcite and sub-sequence processing @xcite . efficient attention targets on reducing the memory and calculation cost of self-attention operations while preserving transformers' empirical performance on downstream tasks. unfortunately, most efficient attention methods require customized self-attention implementations, which demand from-scratch training instead of being directly plugged into existing pre-trained models. moreover, some empirical studies have demonstrated that efficient-attention methods inevitably sacrifice the short-sequence processing performance compared with full-attention models @xcite . the alternative solution to long-sequence processing is to decompose the long-sequence input into multiple sub-sequences and then feedforward each separately, known as sub-sequence processing @xcite . although fullattention blocks are sufficiently utilized in each sub-sequence, these methods have been challenged to capture the semantic information across different sub-contexts. to solve this, some works assign the same fragment to different chunks @xcite , which however significantly increases the computational cost of each chunk and runs counter to the original efficiency goal. to gather the advantages of the methods above, we introduce a simple yet effective learning framework with three typical operations: chunk, align, and select (simcas). in detail, simcas first chunks the input sequence into a batch of subsequences then feeds each subsequence forward the elaborately designed encoding blocks with an interchunk alignment mechanism. finally, the most semantically representative hidden representations are selected via a tailored selection policy to compress the overall sequence length. to align the semantic information across sub-sequence chunks, we introduce a sequential batch alignment operation to calibrate the start and end token embeddings of each sub-sequence in the encoder layers. to learn the selector policy, inspired by the recent success of reinforcement learning in nlp @xcite , we adopt the proximal policy optimization (ppo) @xcite algorithm with the decoder treated as the environment. moreover, we use the attention scores and output likelihood to calculate effective rewards for the selector's policy optimization. to verify the effectiveness of simca, we conducted comprehensive experiments on seven long-context datasets from the domains of document-level summarization, multidocument summarization, and reading comprehension. the empirical results show that simcas outperforms other long-sequence processing baselines with high-level scalability. the main contributions of this paper are: language modeling the training objective for sequence generation consists of a sequence of token decisions made in an auto-regressive manner. this is formulated as a product of decision probabilities corresponding to specific tokens. given an input sequence @xmath0, x 2 , • • • , x n ) and its corresponding output @xmath1, y 2 , • • • , y m ), we model the following conditional probability: ## methodology given a long-input text @xmath2, x 2 , . . . , x n ) with a fairly large input length n , we aim to design a model p ϕ (y |x) to predict a corresponding label sequence @xmath3, y 2 , . . . , y m ), where y can either be classification labels or output sequence tokens. the major difficulty of the task is that the input length n is so large that the original selfattention operations become infeasible with the quadratic complexity o(n 2 ). to address the challenge, we propose a novel method that intuitively splits the long inputs into chunks with feasible lengths, then selects the most representative tokens for decoding steps. to guarantee inter-chunk semantic information extracted during encoding, we design an aligning scheme in the encoder blocks. in the following, we will describe the chunking scheme, the aligning strategy, and the selector design in detail. ## related work efficient transformers the attention mechanism in transformers requires quadratically increased computational complexity with respect to the input sequence length, which limits the application scenarios, especially for long-text processing. to address this issue, various previous works have been proposed for designing more efficient attention operations @xcite . longformer @xcite , big-bird @xcite , gmat @xcite , and etc @xcite reduce the memory consumption of dense attentions by elaborate combinations of global attention and local attention mechanisms. longt5 @xcite is based on the original t5 with global-local attention sparsity patterns to handle long input and has an additional pre-training phase on c4 dataset. colt5 @xcite is an improved version of longt5, with the same parameter size and pretraining dataset as longt5, and it optimizes the computation of the attention module and the feedforward module and proposes a new pre-training objective. additionally, reformer @xcite leverages a locality-sensitive hashing to the attention mechanism, changing its complexity from o(n 2 ) to o(n log n), where n is the input text sequence length. routing transformer @xcite applies a sparse routing module based on online k-means to self-attention while reducing the overall complexity of attention. approximationbased methods, such as performers @xcite and rfa @xcite , use linear space and time complexity to estimate the attention matrix based on random features. luna @xcite attends only to a fixed number of hidden vectors. linformer @xcite calculates selfattention by a low-rank matrix. however, the vast majority of these methods are difficult to apply to existing plms. moreover, @xcite proposes that many efficient-attention transformers do not even perform as well as simple local-attention models on downstream language tasks. ## experiments to evaluate the performance of simcas, we conduct experiments on the long-text abstractive summarization and machine reading comprehension tasks. in the following, we introduce the information about the datasets, baselines, model implementations, and evaluation results of our experiments. ## discussion main result anaylses ablation study simcas is composed of three components: chunk, align, and select. to investigate the contributions of each component, we independently removed each one. notably, after removing chunk, the maximum input length was restricted to 1024, causing align to fail. as demonstrated in ## conclusion and future work in this paper, we introduced a simple method for long-text processing via chunking, aligning, and se-lecting, called simcas. we divided the input long sequence into chunks, and encoded them with sequential batch alignment to capture the inter-chunk semantics. to select the important token representations in the encoded output of the encoder, we introduced a reinforcement learning-based token selector with the ppo method. we leverage the transformer as an environment and design a reward scheme for the corresponding actions based on the output logits and decoder cross-attention feedback to optimize the hidden token selector. substantial experiment results and analyses demonstrate the satisfying effectiveness of simcas. besides, our method does not depend on any particular tasks or models, which have good generalization ability for various application scenarios. for future work, our method can be naturally adopted into non-language long-sequence processing tasks, such as molecular structure analysis. also, simcas has the potential to enhance the long text pre-training to transformers.