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 Simple framework to enhance the long-content processing of off-the-shelf pre-trained transformers via three steps: Chunk, Align, and Select (SimCAS). More specifically, we first divide each longsequence 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 longsequence processing baselines. The code is at https://github.com/xjw-nlp/SimCAS.

heading: 1 Introduction
text: Transformers (Vaswani et al., 2017) have become a fundamental model architecture for sequential data modeling (Sun et al., 2019a; Dosovitskiy et al., 2021; Wysocki et al., 2023), especially in Natural Language Processing (NLP) (Devlin et al., 2018; Brown et al., 2020), 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 (Yang et al., 2019), controllable generation (Byrne et al., 2021; Cheng and Li, 2022), summarization (Rush et al., 2015; Nallapati et al., 2016) and logic reasoning (Wei et al., 2022; Cheng et al., 2024a).
*Work done during an internship at Tencent AI Lab. †Corresponding author.
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 (Kwiatkowski et al., 2019; Gong et al., 2020; Pang et al., 2022) and longtext summarization (Huang et al., 2021; Ma et al., 2022).
To enhance transformers with more efficient long-sequence processing, prior works focus on two perspectives, efficient attention operations (Beltagy et al., 2020; Zaheer et al., 2020; Choromanski et al., 2020) and sub-sequence processing (Liu et al., 2022). 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 (Phang et al., 2022).
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 (Moro et al., 2022). 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 (Ivgi et al., 2023), which however significantly increases the computational cost of each chunk and runs counter
13500
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 (Ouyang et al., 2022), we adopt the Proximal Policy Optimization (PPO) (Schulman et al., 2017) 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:
• We propose a simple yet effective long-sequence processing framework, which noticeably extends the application range of existing full-attention PLMs. Unlike prior works compromising the performance on short sequences, our method maintains satisfying performances processing either short or long contexts.
• We discover that transformers can be conceptualized as simulation environments for policy learning. We leverage transformers’ attention scores and output likelihoods to optimize the selector policy to compress the long-sequence information. The optimized selector policy meanwhile facilitates the transformer to concentrate more on hidden states with high semantic importance.
• We conduct comprehensive experiments illustrating that SimCAS consistently surpasses previous baselines. Furthermore, we provide model scal-
ing, efficiency comparisons, and ablation studies to substantiate the superior performance of our proposed method.

heading: 2 Background
text: 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 X = (x1,x2, · · · ,xN ) and its corresponding output Y = (y1,y2, · · · ,yM ), we model the following conditional probability:
pϕ(Y |X) = M∏
m=1
pϕ(ym|Y<m,X), (1)
where Y<m = (y1,y2, . . . ,ym−1), and ϕ represents the model parameters.
Proximal Policy Optimization In the domain of reinforcement learning (RL), Proximal Policy Optimization (PPO) (Schulman et al., 2017) is a widely used policy gradient method (Kakade, 2001) for its remarkable performance and efficiency in solving complex control and decision-making tasks (Vinyals et al., 2019; Akkaya et al., 2019; Cheng et al., 2024b). The vanilla policy gradient estimator has the form: ∇θEπθ(at|st)[Aπt (at, st)] ≈ Êt[∇θ log πθ(at|st)Ât], where st ∈ S is the state at t-step, πθ(at|st) is a stochastic policy acting at at st, Ât is the estimated value of the advantage function Aπt (at, st), and Êt denotes the empirical average over a sample batch. The PPO algorithm improves the training stability of the policy gradient, by optimizing the following objective:
L = Êt[min(rt(θ)Ât, clip(rt(θ), 1−ε, 1+ε)Ât)], (2) where rt(θ) = πθ(at|st) πθold (at|st)
is the probability ratio between new and old policies, and ε > 0 is a hyperparameter for clipping.

heading: 3 Methodology
text: Given a long-input text X = (x1,x2, . . . ,xN ) with a fairly large input length N , we aim to design a model pϕ(Y |X) to predict a corresponding label sequence Y = (y1,y2, . . . ,yM ), 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(N2).
Encoder Input (N tokens) Decoder Input (M tokens)
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.

heading: 3.1 Chunking Scheme
text: Assume the maximum input sequence length of a pre-trained transformer model is S. We first split the long input sequence into B = ⌈NS ⌉ chunks with each chunk length equal to S. To simplify the expression, we assume that a sentence is in only one chunk. In the practical experiment, we use the sentence-level segmentation (Moro and Ragazzi, 2022) to acquire a sequence of sentences, and then employ the greedy method to allocate sequentially these sentences to the chunks. If there is not enough space in the chunk for the current sentence, we will use the padding tokens to fill the chunk. Mathematically, these divided segments can be represented as: {
(ck1, c k 2, . . . , c k S) }B k=1 , (3)
where cki = x(k−1)S+i, and ⌈·⌉ is the ceiling function. Since the chunk might not have S tokens from
X , we append special padding tokens at the end (see Figure 1). After chunking, we add the start token ([S]) and end token ([E]) to each chunk, and treat the chunks as a normal transformer input batch C with batch size B:
C = { ([S], ck1, c k 2 . . . , c k S , [E]) }B k=1 . (4)

heading: 3.2 Sequential Batch Alignment
text: After chunking the input text into a standard token batch, we can encode it with the transformer encoder layers. As in Figure 1, we assume the encoder has N1 layers. Denote the hidden representations of k-th chunk in C (in equation 4) at l-th encoder layer as Hk,l = (hk,l0 ,h k,l 1 , . . . ,h k,l S ,h k,l S+1), where hk,l0 and h k,l S+1 are hidden representation of [S] and [E] tokens respectively, and hk,li is the embedding for cki with 1 ≤ i ≤ S.
As mentioned in Section 1, chunking methods make it difficult to capture the inter-chunk semantic information. Inspired by recent work using special tokens for full sequence information (Mohtashami and Jaggi, 2023), we align the information of [S] and [E] tokens at each encoding block. More specifically, at l-th layer, our batch alignment average the hidden states of [S] and [E] of all chunks:
h̄lBOS = 1
B
B∑
k=1
hk,l0 , h̄ l EOS =
1
B
B∑
k=1
hk,lS+1. (5)
Then we replace hk,l0 and h k,l S+1 with the aligned h̄lBOS and h̄ l EOS into the hidden states for the nextlayer encoding block, as shown in Figure 1.

heading: 3.3 Token Selector
text: After being encoded into the last hidden space with our sequential batch aligning scheme, the chunks should be reformed back to a sequence for the next decoding steps. Directly tiling the chunks’ representations back into a sequence is still infeasible because the overlong sequence makes it difficult to fuse information at the decoding stage. Therefore, we propose a token selector to select the most representative hidden representations for decoding steps. Inspired by Ramamurthy et al. (2023), we design the selector from the perspective of reinforcement learning (RL).
Selection Module Design Formally, the token selector takes the last hidden states HL = {hk,L0 ,hk,L1 , . . . ,hk,LS ,h k,L S+1}Bk=1 and selection actions At = (a1,a2, . . . ,at−1) as inputs and predicts the next selection action at, where each action at has two values "select" and "skip" for operating the t-th token xt in X . We set the state of RL as st = (HL,At), then the selector is a policy πθ(at|st) to predict next action at.
We implement the selector with the actor-critic style (Konda and Tsitsiklis, 1999). Both actor and critic are simple feed-forward networks, but the actor outputs a probability distribution over the action space and the critic outputs a single scalar value. At state st, to consider the sequential effect of previous action At, we first take the average of all selected hidden states as
h̄t =
∑ k,i I{aj="select",j<t}h
k,L i∑
k,i I{aj="select",j<t} , (6)
where j = (k − 1)S + i maps chunk indices back to the input sequence, and I{·} is the indicator function. Then we concatenate the current selector state h̄t and token information h k,L i , t = (k − 1)S + i, to predict next action at via the actor:
πθ(at|st) = actor(h̄t,hk,Li ). (7) Reward Design To train the selector within an RL scheme, we treat the transformer as an environment and design action rewards. Inspired by Yuan et al. (2021), we can directly utilize the language modeling likelihood as the generation quality reward for selection actions:
RLM = ξ exp{ 1
M log pϕ(Y |X)}, (8)
where ξ is a coefficient that magnifies the value of the reward for easier optimization of the selector. However, RLM is only a scalar value, which cannot provide fine-grained guidance to the selector. Therefore, we use the input-output cross-attention scores to calibrate RLM. More specifically, we denote the cross-attention matrix of the q-th attention head in the l-th layer of the decoder as Alq ∈ RM×N , and the overall cross-attention
Ā = 1
N2 ·Q
N2∑
l=1
Q∑
q=1
Alq, (9)
where N2 is the number of decoder layers, and Q is the number of cross-attention heads. With the overall feedback RLM and cross-attention Ā = (āij)M×N , we adjust the reward R+j to each selected tokens. For "skip" action, we intend to limit the selected sequence length. Assume the number of overall input tokens and selected token representations is Lall and Lselect respectively. We set a hyper-parameter Lhyper to control the size of Lselect with the skipping reward R−:
R+j = āj
1− ā0 RLM, āj =
1
M
M∑
i=1
āij ,
R− =
{ RLM Lall
if Lselect < Lhyper RLM Lselect otherwise.
(10)
With the selector and rewards designed above, we can optimize the selector with the PPO algorithm described in Section 2. Note that in our setups, the environment (the transformer) changes during the training steps. Therefore, we alternatively update the selector and transformer: in each interaction, we first fix the transformer and use the reward RLM and cross-attention scores to update the selector, then fix the selector and update the transformer with language modeling loss. In addition, the selector is similar to the chunk-wise RNN, so the time overhead of the selection process is low.

heading: 4 Related Work
text: 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 (Tay et al., 2023; Fournier et al.,
2023). Longformer (Beltagy et al., 2020), BIGBIRD (Zaheer et al., 2020), GMAT (Gupta and Berant, 2020), and ETC (Ainslie et al., 2020) reduce the memory consumption of dense attentions by elaborate combinations of global attention and local attention mechanisms. LongT5 (Guo et al., 2022) 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 (Ainslie et al., 2023) 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 (Kitaev et al., 2020) leverages a locality-sensitive hashing to the attention mechanism, changing its complexity from O(n2) to O(n log n), where n is the input text sequence length. Routing Transformer (Roy et al., 2021) 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 (Choromanski et al., 2020) and RFA (Peng et al., 2021), use linear space and time complexity to estimate the attention matrix based on random features. Luna (Ma et al., 2021) attends only to a fixed number of hidden vectors. Linformer (Wang et al., 2020) calculates selfattention by a low-rank matrix. However, the vast majority of these methods are difficult to apply to existing PLMs. Moreover, Xiong et al. (2022) proposes that many efficient-attention transformers do not even perform as well as simple local-attention models on downstream language tasks.
Chunking Methods for Long Sequence Another solution for long-sequence processing is to perform sequence chunking and then process them respectively (Zhong et al., 2022; Wang et al., 2023). Among chunking methods, SLED (Ivgi et al., 2023) splits the long sequence into overlapping chunks and processes each chunk with the encoder, then fuses cross-chunk information with the decoder. PageSum (Liu et al., 2022) separates the long sequence into non-overlapping chunks and effectively tackles them by the principle of locality (Denning, 2005). Unlimiformer (Bertsch et al., 2023) encodes long inputs in chunks and utilizes only the top-k input tokens for every attention head.
Length Extrapolation Length extrapolation in transformers refers to their ability to handle input sequences of varying lengths during both train-
ing and inference (Press et al., 2022; Sun et al., 2023). Transformers use a self-attention mechanism to capture dependencies across positions in a sequence, allowing them to generalize well to sequences of different lengths. This flexibility is essential for tasks like NLP and time series analysis, where input lengths can vary.
Sequence Length Reduction Reducing the length of hidden states (Guan et al., 2022; Kim et al., 2022) is the method of model compression from the width perspective, which is promising since some studies showed that there is redundant encoded information in token representations (Ethayarajh, 2019; Klafka and Ettinger, 2020). Among the redundancy, some tokens carry more task-specific information than others, suggesting that these tokens are more salient and imperative to be selected to feed into subsequent operations. Compared with model compression via layer-wise pruning, token-level pruning does not come at the expense of model performance in complex reasoning (Sanh et al., 2019; Sun et al., 2019b).

heading: 5 Experiments
text: 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.

heading: 5.1 Datasets
text: We conduct experiments on two types of NLP tasks: long-text summarization and machine reading comprehension. For long-text summarization, we use four single-document summarization datasets: arXiv, PubMed (Cohan et al., 2018), GovReport (Huang et al., 2021), SummScreen (Chen et al., 2022), and two multi-document summarization datasets: Multi-News (Fabbri et al., 2019b), WCEP (Gholipour Ghalandari et al., 2020). For the reading comprehension task, we test on the NarrativeQA (Kočiský et al., 2018) dataset. We list the introduction about these datasets below. More dataset details can be found in Appendix § A.
arXiv & PubMed1 are two long-document summarization datasets in the scientific research domain. Each document is a scientific paper whose summary is the corresponding abstract.
1https://github.com/armancohan/long-summarization
GovReport2 is a long-document summarization dataset based on reports published by the U.S. Government Accountability Office and Congressional Research Service.
SummScreen3, which includes TV series transcripts, often indirectly presents plot details through character dialogues scattered throughout the transcript. These details need to be consolidated to create concise plot descriptions.
Multi-News4 is a large-scale multi-document summarization dataset. It consists of news articles and human-written summaries of these articles. Each summary is professionally written by editors and with links to the original articles cited.
WCEP5 is a dataset for multi-document summarization (MDS). It contains short, human-written summaries about news events, obtained from the Wikipedia Current Events Portal (WCEP).
NarrativeQA6 is a reading comprehension dataset over entire books from Project Gutenberg and movie scripts from different websites.

heading: 5.2 Baselines
text: There are several baselines for comparison: among them, those based on the full-attention mechanism are HiMAP (Fabbri et al., 2019a), BERTREG (Gholipour Ghalandari et al., 2020), Submodular+Abs (Gholipour Ghalandari et al., 2020), BART (Lewis et al., 2020), PEGASUS (Zhang et al., 2020), DynE (Hokamp et al., 2020), GraphSum (Li et al., 2020), BART-Long-Graph (Pasunuru et al., 2021), SLED (Ivgi et al., 2023), Memorizing Transformers (Wu et al., 2022), Unlimi-
2https://github.com/luyang-huang96/LongDocSum 3https://github.com/mingdachen/SummScreen 4https://github.com/Alex-Fabbri/Multi-News 5https://github.com/allenai/PRIMER 6https://github.com/deepmind/narrativeqa
former (Bertsch et al., 2023). Moreover, several baselines introduce the tailored sparse attention to address longer sequences such as LED (Longformer Encoder-Decoder) (Beltagy et al., 2020), BIGBIRD (Zaheer et al., 2020), PRIMERA (Xiao et al., 2022), HEPOS (Huang et al., 2021) and LED+RELAX (Parnell et al., 2022). More details of baselines can be found in Appendix §C.

heading: 5.3 Implementation Details
text: Our implementation is based on PyTorch (Paszke et al., 2019) and Transformers (Wolf et al., 2020) libraries. We train our model by using 8 NVIDIA V100 32G GPUs.
During the training phase, the maximum input lengths for BARTlarge and BARTbase are set to 8192 and 16384, respectively, unless otherwise specified. To ensure efficient training, we update the parameters of the original backbone and the selector alternately. The reward estimation for each action is computed based on decoder cross-attention and the output feedback of the generative model. This estimation process is detached from the computation graph and does not participate in backpropagation.
At the inference stage, compared to the original generation process, our framework only adds a chunk-wise selection procedure between the en-
coder and the decoder, which takes very little time. At the decoding stage, the target sequence is generated with beam search in an auto-regressive manner (Wiseman and Rush, 2016).

heading: 5.4 Evaluations
text: Like most previous works, for abstractive summarization tasks, we measure the quality of generated summaries using the popular metric ROUGE (Lin, 2004). On the test set of arXiv, PubMed, GovReport, SummScreen, Multi-News, and WCEP, we report full-length F1-based ROUGE-1, ROUGE-2, and ROUGE-L scores computed with the standard ROUGE Perl package. Furthermore, we also use a popular model-based semantic metric BERTScore7 (Zhang* et al., 2020) to demonstrate the superiority of our approaches comprehensively. As shown in Table 2, ROUGE-1 and ROUGE-L are strongly correlated (Pearson correlation score of 0.913). In particular, the relatively low Pearson correlation coefficient (< 0.8) between ROUGE-2 and BERTScore indicates a significant difference in preferences.
As for the reading comprehension task, we use the F1 and exact match (EM) metrics defined in SCROLLS (Shaham et al., 2022) to evaluate the model performance on the NarrativeQA dataset. Both metrics normalize the reference and system output strings via lowercasing, removing punctuation and stopwords, and normalizing whitespace.

heading: 6 Discussion
text: Main Result Anaylses Table 1 and 3 report results over four long-document test sets. We note that casting the backbone BARTbase into our SimCAS can significantly improve the model performance, and increasing the size of the model can further improve performance. Our approach outperforms baseline on several metrics and achieves new state-of-the-art performance on the PubMed test set. In Table 4, the results of multi-document summarization tasks have a similar trend. Apart from PRIMERA, which customizes a pre-training objective for multi-document summarization, our method significantly outperforms previous results. See Figure 3, we observe a substantial improvement in the performance of our model on the NarrativeQA test set compared to previous works. Given the fact that this dataset has an extremely long average input length (≥100K) and short average output
7In order to make a fair comparison with the previous work, we also use checkpoint “facebook/bart-large-mnli” for BERTScore.
length (≤10), we attribute this significant performance enhancement to our method’s proficiency in filtering out an immense amount of task-irrelevant information. This, in turn, enables the decoding blocks to fuse information more efficiently.
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 Table 5, the performance experienced a significant decline when either Chunk or Select was removed, underscoring their effectiveness. The Align component improved performance on all datasets, with the exception of the Multi-News dataset. We hypothesize that more sophisticated alignment methods could potentially enhance performance further.
Effectiveness of Conditional Skipping Reward To prevent our selector from choosing an excessive number of token embeddings, we implement a conditional reward mechanism for skipping actions, which imposes a soft constraint on the number of selected token embeddings. Figure 2 (left) demonstrates the effectiveness of this reward mechanism. As training progresses, the number of selected embeddings gradually converges to a threshold of 2048. This indicates that our selection process is effective in identifying essential token embeddings. Additionally, as the input sequence length increases, the ratio of selected token embeddings to the total length decreases. This is due to the length penalty in equation 10, which discourages the selection of an excessive number of token embeddings.
Increasing Maximum Input Length We explore the performance changes of our model on
a single 32G V100 GPU as the input sequence length increases during inference. Initially, we assess the efficiency of our framework in processing long sequences using BARTbase as the backbone. We depict the fluctuations in the ratio of the average time required by SimCAS to the time taken by BART (1024) in Figure 2 (right). Specifically, a line graph represents the relative time per sample, while a bar graph displays the number of token embeddings selected by SimCAS. When the input length is short (< 214), the time cost exhibits near-linear growth. We attribute this phenomenon to our method’s significant reduction in computational complexity and the redundant computation resource for short sequences. However, with a dramatic increase in input length (> 217), the time cost rises sharply. As indicated in the blue area on the far right of the graph, our method can handle input sequences of up to 350k in length on a V100. Moreover, even as the number of input tokens increases significantly, the selected hidden states remain at a reasonable size, demonstrating our selector’s capability to retrieve high-contribution information.
Low-Resource Evaluation In real-world applications, the quantity of training samples available for downstream tasks is frequently quite constrained. Consequently, the performance of the model under low-resource conditions is of paramount importance. In our setup, we randomly sample a few (10 and 100) training examples from specific datasets to adapt these PLMs such as BARTbase (1024), LEDbase (4096), and SimCASbase (4096), for corresponding data distributions. The results in Figure 4 indicate that our model demonstrates superior sample efficiency in low-resource scenarios compared to previous robust baselines.

heading: 7 Conclusion and Future Work
text: 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.

heading: 8 Limitations
text: Even though our framework significantly reduces the overhead of processing long sequences, due to memory constraints, the maximum input length on a single V100 32G GPU during training is 16k, and there is an extremely large consumption of GPU memory even if the batch size is small, which might suppress the potential of our method. In addition, we use small-scale PLMs with good performance as the backbone for downstream tasks. The effectiveness of our framework on large language models needs to be investigated further.

heading: 9 Ethical Considerations
text: Similar to existing methods, there is no guarantee that the generated content is factually consistent and free from hallucination (Maynez et al., 2020; Kang and Hashimoto, 2020). Therefore caution is imperative when applying our method to scenarios with high demand of generation accuracy.

heading: A Dataset Statistics

heading: B Hyper-parameters & Packages
text: Table 7 and Table 8 delineate the hyper-parameter settings at the training and inference phases of the experiment, respectively. For evaluation metrics, we used the following packages:
• For the uniform data processing, before the evaluation, all the reference texts and model-generated texts are converted to the lowercase and tokenized using the PTB tokenizer: https://nlp.stanford.edu/ nlp/javadoc/javanlp/edu/stanford/nlp/ process/PTBTokenizer.html.
• For ROUGE metrics, we used the public rouge-score Perl package provided by the authors: https://github.com/ summanlp/evaluation/tree/master/ ROUGE-RELEASE-1.5.5.
• For BERTScore (Zhang* et al., 2020), we used the public bert-score package shared by the authors: https://github.com/Tiiiger/bert_ score.
C Introduction for Baselines
We use the competitive baselines that demonstrate the downstream task results for comparison. Among them, BART (Lewis et al., 2020) is a standard full-attention PLM for sequence generation. Compared with BART, PEGASUS (Zhang et al., 2020) has a tailored pre-training objective for abstractive text summarization. LED (Longformer Encoder-Decoder) (Beltagy et al., 2020) uses the sparse attention-based encoder and full-attention decoder. Before pre-training, its parameters are initialized from BART. BIGBIRD (Zaheer et al., 2020), for an encoder-decoder setup, also introduces their specified sparse attention mechanism only at the encoder side. PRIMERA based on
LED introduces a task-specific per-training objective for multi-document summarization. SLED (Ivgi et al., 2023) processes long sequences via short-context PLMs. The origin long sequence is partitioned into overlapping chunks. HEPOS (Huang et al., 2021) proposes head-wise positional strides to effectively pinpoint salient information from the source. Memorizing Transformers (Wu et al., 2022) employs a trainable attention gate to moderate between the standard crossattention and attention over retrieved keys from a datastore. Unlimiformer (Bertsch et al., 2023) uses the k-nearest-neighbor KNN index to choose full-length input to reduce computation overhead. HiMAP (Fabbri et al., 2019a) expands the existing pointer-generator network model into a hierarchical network. GraphSum (Li et al., 2020) employs well-known graph representations of documents to effectively process multiple input documents for abstractive summarization. BART-LongGraph (Pasunuru et al., 2021) is fine-tuned based on LED and additionally injects discourse graph. LED+RELAX (Parnell et al., 2022) introduces a RELAX gradient estimator with multi-document coverage reward. BERTREG (Gholipour Ghalandari et al., 2020) is a regression-based sentence ranking system with BERT embedding, which is used as an extractive summarization method. Submodular+Abs (Gholipour Ghalandari et al., 2020) consists of a submodular-based extractive summarizer and a bottom-up abstractive summarizer. DynE (Hokamp et al., 2020) ensembles multipledocument for abstractive summarization by single document summarization models.

heading: D Perplexity on Test Sets
text: To more comprehensively demonstrate the effectiveness of our approach, we additionally use perplexity as an evaluation metric for comparative experiments. The experimental results in Table 9 indicate that our system (SimCAS) significantly enhances model performance compared to the baselines (BART).

heading: E Time Latency at Inference Stage
text: At the inference stage, our framework additionally introduces a selection process compared with the standard transformer. Therefore, We study the proportion of our selector in the average inference time cost per sample. The results in Table 10 indicate that our selector only adds a small amount
of time cost while significantly improving model performance, demonstrating the effectiveness of our approach.

heading: F Potential to select representation
text: Because in the decoder cross-attention module, the encoded output from the encoder is used as “Key” and “Value” to participate in the calculation, its attention score can lead to the contribution of the corresponding token representation from encoded output to the current token decision at the decoding step. Figure6 demonstrates how the cross-attention scores change during the decoding of the reference output with one example in GovReport. We can observe that 1) most of the token decisions in the decoding phase focus only on a small set of encoded representations; 2) For each token decision, the contributions of different encoded token representations vary greatly. These phenomena suggest that there is still the feasibility of further filtering out low-contribution encoded representations.

heading: G Compatible with short-input tasks
text: While the sparse-attention transformers have been proven effective on a wide range of long-sequence datasets, as shown in Figure 7, these methods tend to underperform traditional full-attention transformers on the more common short-sequence tasks. However, in real scenarios, short-sequence inputs and long-sequence inputs are often mixed together, and the former occupies the vast majority. This limits the application scope of sparse-attention transformer architecture. In contrast, applying our framework SimCAS to existing full-attention transformers has strong flexibility. Specifically, given the full-attention model under SimCAS, if the input sequence exceeds the maximum length of a single chunk, we will perform chunking and selecting,
otherwise, we can naturally switch to a standard short-text encoding form by skipping chunking procedures.

heading: H Efficacy of SBA
text: As our framework introduces SBA to align interchunk information during the encoding steps, we explore the change of hidden state in the forward propagation process under the influence of this component. Figure 8 demonstrates the visualized similarity matrix between each chunk’s start hidden state after the encoding block. It can be seen from an array of cosine similarity matrices in the top half that, through the alternation of SBA and encoding blocks, the cosine similarities between the start hidden states from different chunks are all close to 1, which indicates that their directions in high-dimensional space are almost the same.
Considering that cosine similarity can only reflect the closeness between vector directions, we design a similarity calculation method to add the measure of Euclidean distance. Given a pair of vectors v1,v2 ∈ Rd, the custom similarity Sim between them is formulated as follows:
Sim = (v1,v2)
∥v1∥ ∗ ∥v2∥ ∗ (1 + ∥v1 − v2∥) ,
by which, we scale the cosine similarity to observe its change. As can be seen from the lower part of Figure 8, after each encoder layer, the start hidden states will vary in scale due to different contexts.

heading: I More Experimental Details
text: Since we aim to minimize dataset-specific modeling, we unify the processing of the used long-
Algorithm 1 Token embedding selection with PPO 1: Initialize agent and hyper-parameters 2: Initialize advantages, rewards, and returns 3: Compute Rewards and Advantages 4: Initialize next_value← 0, last_gae_lam← 0 5: valid_reward_num← size of select_rewards 6: for t = 0 to num_steps-1 do 7: Compute action sum for each step and up-
date rewards 8: end for 9: cur_num← 0
10: for t = 0 to num_steps-1 do 11: Update rewards and actions based on select_rewards 12: end for 13: for t = num_steps-1 to 0 do 14: Compute delta, advantages, and returns 15: end for 16: Reverse advantages and returns 17: Policy and Value Function Update 18: for epoch = 0 to update_epochs-1 do 19: Shuffle batch indices 20: for batch_st = 0 to seq_len by mini_batch do 21: Compute policy loss, value loss, and entropy loss 22: Perform gradient ascent 23: end for 24: if approx_kl > threshold then 25: Break 26: end if 27: end for
document summarization, multi-document summarization, and machine reading comprehension datasets. In the NarrativeQA dataset for machine reading comprehension, we concatenate together the question and reference content, each of which prepends a prefix to indicate the type of text. For a fair comparison, we uniformly use the full-attention PLM BART8 on the seven public datasets above. Built on the BART, our framework introduces an additional parameterized selector to focus on more task-specific token representations. The selector follows the actor-critic style (Konda and Tsitsiklis, 1999) and contains around 8M parameters. There are two Adam optimizers with β1 = 0.9, β2 = 0.999 for BART and selector respectively.
Additionally, to handle longer sequences during the training phase, we set the chunk size to 512 instead of 1024 in all experiments (considering the start token and end token). To maintain chunk-wise alignment, we pad the chunks to uniform the size of each chunk. During the forward propagation, the embedding layer embeds position representations for each chunk independently.9
8The checkpoints of BART are “facebook/bart-base” and “facebook/bart-large” containing around 139 M and 406 M parameters respectively, whose maximum encoding length is 1024.
9Considering the context of each chunk is different, although our design may lead to repeated position representation, we argue that after the encoding stage, the combination of the same token and position information would still produce various representations (Kazemnejad et al., 2023)
Maximum Input Length R-1 R-2 R-L BS
Training Details In this paper, for the training of BARTbase on various datasets we uniformly employ the Adam optimizer with the following dynamic learning rate:
lr = γmin(step−0.5, step × warmup−1.5),
where γ affects the maximum learning rate and warmup indicates the warmup steps, step is the number of updating steps, and lr is the learning rate. In addition, there is another separate Adam optimizer for our reinforcement learning-based parameterized selector, in which the learning rate is fixed to 1× 10−4. The optimization of the selector and transformer is performed alternately, with one model trained while the other model is fixed.
Inference Details In practice, We also investigate how the model performs when the maximum effective input length is increased during inference. In Table 11, we can observe that although the maximum input length is set to 16384 during the training phase due to the memory limitation, increasing this maximum length during inference still improves the model performance.
The Details of Selector Our token embedding selection process based on PPO is shown in Algorithm 1. In our setting, the selector consists of an actor component and a critic component with PPO optimization, both of which are the feed-forward network, except for the final layers. In order to enable the selector to choose more diverse token representations instead of becoming homogeneous during the chunk-wise selection process, the state space consists of the current token representation and the selector’s hidden state that is affected by previous actions. At the beginning of the selection procedure, the initial selector hidden state is the average of the start hidden state of all chunk representations. At each time step, we input the
current selector hidden state and a chunk of token representations. Then the actor of the selector outputs a probability distribution over action space, and the critic of the selector outputs a single estimated scalar value for each token based on the hidden state selector and the corresponding token representation. For the state transition after executing current actions, we update the hidden state of the selector using the selected token representation from the current chunk. Note that in order to avoid the extreme case where the selector skips all tokens, in each chunk decision, if all actions are “skipping”, they are all switched to “selecting”, and the corresponding action probability and value are re-obtained. Additionally, to increase the training stability, the advantage in the PPO algorithm is approximated using the Generalized Advantage Estimation (GAE) (Schulman et al., 2016).

heading: J Case Study on GovReport Dataset
text: In addition to using regular automatic evaluation metrics to measure the effect of model generation, we also present some actual output to support the results. Figure 9 displays several examples of summaries generated by the fine-tuned base model BARTlarge and our BARTbase-SimCAS. We can observe that the system output of our model has fewer grammatical errors and higher coherence compared with the base model. Furthermore, since our model is able to perceive longer sequences, our output is more informative and better aligned with the reference text.

heading: Future Work
text: 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