# Chunk, Align, Select:
A Simple Long-sequence Processing Method for Transformers

Author Name
thanksCorresponding author  Jiawen Xie12∗, Pengyu Cheng1 , Xiao Liang13, Yong
Dai1, Nan Du1
1Tencent AI Lab
2Shanghai Jiao Tong University
3Tsinghua University ∗Work done during internship at Tencent AI
Lab.†Corresponding author.

###### Abstract

Although dominant in natural language processing, transformer-based models
remain challenged by the task of long-sequence processing, because the
computational cost of self-attention operations in transformers swells
quadratically with the input sequence length. To alleviate the complexity of
long-sequence processing, we propose a simple framework to enable the off-the-
shelf pre-trained transformers to process much longer sequences, while the
computation and memory costs remain growing linearly with the input sequence
lengths. More specifically, our method divides each long-sequence input into a
batch of chunks, then aligns the inter-chunk information during the encoding
steps, and finally selects the most representative hidden states from the
encoder for the decoding process. To extract inter-chunk semantic information,
we align the start and end token embeddings among chunks in each encoding
transformer block. To learn an effective hidden selection policy, we design a
dual updating scheme inspired by reinforcement learning, which regards the
decoders of transformers as environments, and the downstream performance
metrics as the rewards to evaluate the hidden selection actions. Our empirical
results on real-world long-text summarization and reading comprehension tasks
demonstrate effective improvements compared to prior long-sequence processing
baselines.

##  1 Introduction

Transformers Vaswani et al. (2017) have become a fundamental model structure
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 architectures, pre-trained language models
(PLMs) have recently shown astonishing empirical performance in various NLP
tasks such as question answering Yang et al. (2019a), dialogue generation
Byrne et al. (2021), summarization Rush et al. (2015); Nallapati et al. (2016)
and machine translation Yang et al. (2019b); Pan et al. (2021). However, one
fatal weakness, which has hindered transformer-based models from being applied
in broader application scenarios, is that the computational consumption of
self-attention operations increases quadratically with the input sequence
length. Therefore, standard transformers have always been challenged by long
text processing tasks, such as machine reading comprehension Kwiatkowski et
al. (2019); Gong et al. (2020); Pang et al. (2022) and long-text summarization
Huang et al. (2021); Ma et al. (2022).

To enhance transformers’ capability in long-sequence processing, prior works
pay their research attention to two major perspectives, the efficient
attention operation and the sub-sequence processing. The efficient attention
Beltagy et al. (2020); Zaheer et al. (2020); Choromanski et al. (2020) aims to
reduce the memory and calculation cost of self-attentions while still
preserving transformers’ empirical performance. Unfortunately, most efficient
attention methods require customized self-attention implementations, which are
not compatible with existing pre-trained language models. Hence, introducing
efficient attention leads to re-training the whole transformer model, which is
incompatible with the fine-tuning mechanism. Moreover, some empirical results
show that sparse-attention methods inevitably sacrifice the performance on
short-sequence processing tasks compared with full-attention models Phang et
al. (2022).

Another perspective to long-sequence processing is to decompose the long-input
texts into multiple sub-sequences, then process them individually Moro et al.
(2022); Zhang et al. (2022); Mao et al. (2022); Liu et al. (2022). Although
utilizing the effectiveness of full-attention blocks, sub-sequence processing
methods are not good at capturing the global semantic information among
different chunks. Moreover, for better fusing information across chunks in the
decoder, Ivgi et al. (2023) put the same fragment in different chunks, which
significantly increases the computational cost.

To gather the advantages of the long-sequence processing methods above, in
this paper, we introduce a Simple learning framework with three typical
operations: Chunk, Align, and Select, which we call as SimCAS. In more detail,
SimCAS first chunks the input long text into a sub-sequence batch then feeds
the batch through specially designed encoding blocks with an inter-chunk
aligning mechanism, and finally selects semantically representative hidden
representation via a contextual selection module. To align the global semantic
information among chunks, we introduce a sequential batch alignment (SBA)
operation to calibrate the start and end token embeddings with each batch in
the encoder layers. For learning the selector, 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 to train the
selector, where the transformer decoder is treated as the environment and the
task evaluation metric is regarded as the reward for actions. To evaluate the
effectiveness of our SimCAS, we conduct experiments on three long-document
datasets (arXiv, GovReport, and PubMed), two multi-document datasets (Multi-
News and WCEP), and one reading comprehension dataset (NarrativeQA). Empirical
results show that SimCAS can outperform other long-sequence processing
baselines and is highly scalable.

##  2 Preliminaries

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
${\bm{X}}=({\bm{x}}_{1},{\bm{x}}_{2},\cdots,{\bm{x}}_{N})$ and its
corresponding output
${\bm{Y}}=({\bm{y}}_{1},{\bm{y}}_{2},\cdots,{\bm{y}}_{M})$, we model the
following conditional probability:

$\textstyle
p_{\phi}({\bm{Y}}|{\bm{X}})=\prod_{m=1}^{M}p_{\phi}({\bm{y}}_{m}|{\bm{Y}}_{<m},{\bm{X}}),$
(1)

where ${\bm{Y}}_{<m}=({\bm{y}}_{1},{\bm{y}}_{2},\dots,{\bm{y}}_{m-1})$, and
$\phi$ 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). The vanilla policy gradient estimator has the
form:
$\nabla_{\theta}{\mathbb{E}}_{\pi_{\theta}({\bm{a}}_{t}|{\bm{s}}_{t})}[A^{\pi}_{t}({\bm{a}}_{t},{\bm{s}}_{t})]\approx\hat{{\mathbb{E}}}_{t}[\nabla_{\theta}\log\pi_{\theta}({\bm{a}}_{t}|{\bm{s}}_{t})\hat{A}_{t}]$,
where ${\bm{s}}_{t}\in{\mathcal{S}}$ is the state at $t$-step,
$\pi_{\theta}({\bm{a}}_{t}|{\bm{s}}_{t})$ is a stochastic policy acting
${\bm{a}}_{t}$ at ${\bm{s}}_{t}$, $\hat{A}_{t}$ is the estimated value of the
advantage function $A_{t}^{\pi}({\bm{a}}_{t},{\bm{s}}_{t})$, and
$\hat{{\mathbb{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:

$\hat{{\mathbb{E}}}_{t}[\min(r_{t}(\theta)\hat{A}_{t},\text{clip}(r_{t}(\theta),1-\varepsilon,1+\varepsilon)\hat{A}_{t}],$
(2)

where
$r_{t}(\theta)=\frac{\pi_{\theta}({\bm{a}}_{t}|{\bm{s}}_{t})}{\pi_{\theta_{\text{old}}}({\bm{a}}_{t}|{\bm{s}}_{t})}$
is the probability ratio between new and old policies, and $\varepsilon>0$ is
a hyper-parameter for clipping.

Figure 1: The learning framework of SimCAS: The long-text inputs are first
divided into a batch of chunks and padded with ([P]), with [BOS] ([S]) and
[EOS] ([E]) tokens added. Then the chunk batch is encoded with [BOS] and [EOS]
representations aligned in each encoder layer. Next, representative hidden
states are selected as a subsequence for decoding steps. The decoder output
logits are utilized as rewards for updating the token selector.

##  3 Method

Given a long-input text
${\bm{X}}=({\bm{x}}_{1},{\bm{x}}_{2},\dots,{\bm{x}}_{N})$ with a fairly large
input length $N$, we aim to design a model $p_{\phi}({\bm{Y}}|{\bm{X}})$ to
predict a corresponding label sequence
${\bm{Y}}=({\bm{y}}_{1},{\bm{y}}_{2},\dots,{\bm{y}}_{M})$, where ${\bm{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 self-attention operations become infeasible with the quadratic
complexity ${\mathcal{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.

###  3.1 Chunking Scheme

Assume the maximum input sequence length of a pre-trained transformer model is
$S$. We first split the long input sequence into $B=\lceil\frac{N}{S}\rceil$
chunks with each chunk length equal to $S$:

$\big{\\{}({\bm{c}}^{k}_{1},{\bm{c}}^{k}_{2},\dots,{\bm{c}}^{k}_{S})\big{\\}}_{k=1}^{B},$
(3)

where ${\bm{c}}^{k}_{i}={\bm{x}}_{(k-1)S+i}$, and $\lceil\cdot\rceil$ is the
ceiling function. Since the last chunk might not have $S$ tokens from
${\bm{X}}$, we append special padding tokens at the end, i.e.,
${\bm{c}}^{k}_{i}=\texttt{[PAD]}$ if $(k-1)S+i>N$ (as the [P] token in Fig.
1). After chunking, we add additional beginning of sentence ([BOS]) and end of
sentence ([EOS]) tokens to each chunk, and treat the chunks as a normal
transformer input batch ${\mathcal{C}}$ with batch size $B$:

${\mathcal{C}}=\big{\\{}(\texttt{[BOS]},{\bm{c}}^{k}_{1},{\bm{c}}^{k}_{2}\dots,{\bm{c}}^{k}_{S},\texttt{[EOS]})\big{\\}}_{k=1}^{B}.$
(4)

###  3.2 Sequential Batch Alignment

After chunking the input text into a standard token batch, we can encode it
with the transformer encoder layers. As in Fig. 1, we assume the encoder has
$N_{1}$ layers. Denote the hidden representations of $k$-th chunk in
${\mathcal{C}}$ (in equation 4) at $l$-th encoder layer as
${\bm{H}}^{k,l}=({\bm{h}}^{k,l}_{0},{\bm{h}}^{k,l}_{1},\dots,{\bm{h}}^{k,l}_{S},{\bm{h}}^{k,l}_{S+1})$,
where ${\bm{h}}^{k,l}_{0}$ and ${\bm{h}}^{k,l}_{S+1}$ are hidden
representation of [BOS] and [EOS] tokens respectively, and
${\bm{h}}^{k,l}_{i}$ is the embedding for ${\bm{c}}^{k}_{i}$ with $1\leq i\leq
S$.

As mentioned in Section 1, chunk-level methods are difficult to capture the
inter-chunk global semantic information. Hence, we design a sequential batch
alignment, which aligns the information of [BOS] and [EOS] tokens at each
encoding block, for [BOS] and [EOS] have been well-recognized representative
of sentence semantics Devlin et al. (2018). More specifically, at $l$-th
layer, our batch alignment average the hidden states of [BOS] and [EOS] of all
chunks:

$\bar{{\bm{h}}}^{l}_{\text{BOS}}=\frac{1}{B}\sum_{k=1}^{B}{\bm{h}}^{k,l}_{0},\
\
\bar{{\bm{h}}}^{l}_{\text{EOS}}=\frac{1}{B}\sum_{k=1}^{B}{\bm{h}}^{k,l}_{S+1}.$
(5)

Then we replace ${\bm{h}}^{k,l}_{0}$ and ${\bm{h}}^{k,l}_{S+1}$ with the
aligned $\bar{{\bm{h}}}^{l}_{\text{BOS}}$ and
$\bar{{\bm{h}}}^{l}_{\text{EOS}}$ into the hidden states for the next-layer
encoding block, as shown in Fig. 1.

###  3.3 Token Selector

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 of the overlong sequence length.
Therefore, we propose a token selector to filter 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
${\bm{H}}^{L}=\\{{\bm{h}}^{k,L}_{0},{\bm{h}}^{k,L}_{1},\dots,{\bm{h}}^{k,L}_{S},{\bm{h}}^{k,L}_{S+1}\\}_{k=1}^{B}$
and selection actions
${\bm{A}}_{t}=({\bm{a}}_{1},{\bm{a}}_{2},\dots,{\bm{a}}_{t-1})$ as inputs and
predicts the next selection action ${\bm{a}}_{t}$, where each action
${\bm{a}}_{t}$ has two values "select" and "skip" for operating the $t$-th
token ${\bm{x}}_{t}$ in ${\bm{X}}$. We set the state of RL as
${\bm{s}}_{t}=({\bm{H}}^{L},{\bm{A}}_{t})$, then the selector is a policy
$\pi_{\theta}({\bm{a}}_{t}|{\bm{s}}_{t})$ to predict next action
${\bm{a}}_{t}$.

We implement the selector with the actor-critic style Konda and Tsitsiklis
(1999). Both actor and critic are simple three-layer feed-forward networks,
but the actor outputs a probability distribution over the action space and the
critic outputs a single scalar value. At state ${\bm{s}}_{t}$, to consider the
sequential effect of previous action ${\bm{A}}_{t}$, we first take the average
of all selected hidden states as

$\bar{{\bm{h}}}_{t}=\frac{\sum_{k,i}\mathbb{I}_{\\{{\bm{a}}_{j}=\textit{"select"},j<t\\}}{\bm{h}}_{i}^{k,L}}{\sum_{k,i}\mathbb{I}_{\\{{\bm{a}}_{j}=\textit{"select"},j<t\\}}},$
(6)

where $j=(k-1)S+i$ maps chunk indices back to the input sequence, and
$\mathbb{I}_{\\{\cdot\\}}$ is the indicator function. Then we concatenate the
current selector state $\bar{{\bm{h}}}_{t}$ and token information
${\bm{h}}_{i}^{k,L}$, $t=(k-1)S+i$, to predict next action ${\bm{a}}_{t}$ via
the actor:

$\textstyle\pi_{\theta}({\bm{a}}_{t}|{\bm{s}}_{t})=\text{actor}(\bar{{\bm{h}}}_{t},{\bm{h}}_{i}^{k,L}).$
(7)

Reward Design To train the selector within an RL scheme, we treat the
transformer as an environment and design action rewards based on the decoder
outputs. Inspired by Yuan et al. (2021), we can directly utilize the language
modeling likelihood as the generation quality reward for selection actions:

$\textstyle R_{\text{LM}}=\xi\exp\\{\frac{1}{M}\log
p_{\phi}({\bm{Y}}|{\bm{X}})\\},$ (8)

where $\xi$ is a coefficient that magnifies the value of the reward for easier
optimization of the selector.

However, $R_{\text{LM}}$ 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 $R_{\text{LM}}$. More specifically, we denote
the cross-attention matrix of the $q$-th attention head in the $l$-th layer of
the decoder is denoted as ${\bm{A}}^{l}_{q}\in\mathbb{R}^{M\times N}$, and the
overall cross-attention

$\textstyle\bar{{\bm{A}}}=\frac{1}{N_{2}\cdot
Q}\sum_{l=1}^{N_{2}}\sum_{q=1}^{Q}{\bm{A}}^{l}_{q},$ (9)

where $N_{2}$ is the decoder layer number, and $Q$ is the cross-attention head
number. With the overall cross-attention
$\bar{{\bm{A}}}=(\bar{a}_{ij})_{M\times N}$, we adjust the reward to each
selected tokens:

$\textstyle R_{j}^{+}=\frac{\bar{a}_{j}}{1-\bar{a}_{0}}R_{\text{LM}},\ \
\bar{a}_{j}=\frac{1}{M}\sum_{i=1}^{M}\bar{a}_{ij}.$ (10)

For "skip" action, we intend to limit the selected sequence length. Assume the
number of overall input tokens and selected token representations is
${L}_{\text{all}}$ and ${L}_{\text{select}}$ respectively. We set a hyper-
parameter ${L}_{\text{hyper}}$ to control the size of ${L}_{\text{select}}$
with the following skipping reward:

$R^{-}=\begin{cases}\frac{R_{\text{LM}}}{L_{\text{all}}}&\text{if
${L}_{\text{select}}<{L}_{\text{hyper}}$}\\\
\frac{R_{\text{LM}}}{L_{\text{select}}}&\text{otherwise.}\end{cases}\\\ $ (11)

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 decoder) is changing during the training steps.
Therefore, we alternatively update the selector and the decoder: in each
interaction, we first fix the transformer and use the reward $R_{\text{LM}}$
to update the selector, then fix the selector and update the transformer with
language modeling loss.

We call our method a Simple long-sequence processing method via Chunking,
Aligning, and Selecting (SimCAS). The whole framework is demonstrated in Fig.
1.

##  4 Related Work

Efficient Transformers The attention mechanism in transformers requires
quadratically increased computational complexity with respect to the input
sequence length, which limited 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. Reformer Kitaev et al. (2020)
leverages a locality-sensitive hashing to the attention mechanism, changing
its complexity from $\mathcal{O}(n^{2})$ to $\mathcal{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. Approximation-based
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 self-
attention 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 effective solution for long-
sequence processing is to chunk the long text into sub-sequence and then
process them respectively. Among chunking methods, SLED (Ivgi et al., 2023)
splits the long sequence into overlapping chunks and processes each chunk with
a full-attention encoder, then fuse cross-chunk information with a transformer
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.
Sequence Length Reduction Reducing the length of hidden states is the method
of model compression from the width perspective, which is promising since some
analytical 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 layers. 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).

##  5 Experiments

To evaluate the performance of our proposed SimCAS, we primely conduct
experiments on two NLP tasks: long-text summarization and machine reading
comprehension. In the following, we introduce detailed information about the
datasets, baselines, model implementations, and evaluation results of our
experiments.

###  5.1 Datasets

We conduct experiments on two NLP tasks: long-text summarization and machine
reading comprehension. For long-text summarization, we use three single-
document summarization datasets: arXiv, PubMed (Cohan et al., 2018), GovReport
(Huang et al., 2021), 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 (Kočiský et al., 2018)
dataset.
arXiv & PubMed111https://github.com/armancohan/long-summarization are two
long-document summarization datasets in the scientific research domain. Each
document is a scientific paper whose summary is the corresponding abstract.
GovReport222https://github.com/luyang-huang96/LongDocSum is a long-document
summarization dataset based on reports published by the U.S. Government
Accountability Office and Congressional Research Service.
Multi-News333https://github.com/Alex-Fabbri/Multi-News 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.
WCEP444https://github.com/allenai/PRIMER 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).
NarrativeQA555https://github.com/deepmind/narrativeqa is a reading
comprehension dataset over entire books from Project Gutenberg and movie
scripts from different websites.

More detailed information about the datasets is provided in Appx. §A.

###  5.2 Baselines

There are several competitive baselines for comparison: 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), HEPOS Huang et al. (2021), SLED Ivgi et al. (2023),
Unlimiformer Bertsch et al. (2023), LED (Longformer Encoder-Decoder) Beltagy
et al. (2020), BIGBIRD Zaheer et al. (2020), PRIMERA Xiao et al. (2022), and
LED+RELAX Parnell et al. (2022). More details of baselines can be found in
Appx. §C.

###  5.3 Implementation Details

Base Model | Method | arXiv | GovReport | PubMed
---|---|---|---|---
R-1 | R-2 | R-L | BS | R-1 | R-2 | R-L | BS | R-1 | R-2 | R-L | BS
LED${}_{\text{large}}$ | Standard | 46.63 | 19.62 | 41.83 | - | 56.20 | 26.60 | 28.80 | - | - | - | - | -
LED${}_{\text{large}}$ | PRIMERA | 47.60 | 20.80 | 42.60 | - | - | - | - | - | - | - | - | -
PEGASUS${}_{\text{large}}$ | Standard | 44.21 | 16.95 | 38.83 | - | - | - | - | - | 45.97 | 20.15 | 41.34 | -
PEGASUS${}_{\text{large}}$ | BIGBIRD | 46.63 | 19.02 | 41.77 | - | - | - | - | - | 46.32 | 20.65 | 42.33 | -
BART${}_{\text{base}}$ | SLED | - | - | - | - | 54.70 | 24.40 | 25.40 | - | - | - | - | -
BART${}_{\text{large}}$ | SLED | - | - | - | - | 57.50 | 26.30 | 27.40 | - | - | - | - | -
BART${}_{\text{large}}$ | HEPOS | 47.87 | 20.00 | 41.50 | - | 56.86 | 22.62 | 53.82 | - | 47.93 | 20.74 | 42.58 | -
BART${}_{\text{base}}$ | Unlimiformer | - | - | - | - | 56.60 | 26.30 | 27.60 | 68.20 | - | - | - | -
PRIMERA | Unlimiformer | - | - | - | - | 57.40 | 26.20 | 28.00 | 68.10 | - | - | - | -
BART${}_{\text{large}}$ | Standard | 42.97 | 15.54 | 37.02 | 61.18 | 52.94 | 19.78 | 50.43 | 64.44 | 42.87 | 15.44 | 36.93 | 63.08
BART${}_{\text{base}}$ | SimCAS | 47.08† | 19.28† | 42.17† | 63.44† | 59.22† | 25.97† | 56.50† | 68.15† | 48.01† | 21.03† | 43.72† | 66.18†
Table 1: Average results on three long-document test sets. $\dagger$:
significantly better than BART ($p<0.01$). R-1/2/L is the ROUGE-1/2/L F1
score. BS refers to the neural model-based metrics BERTScore. The best results
are bolded.

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.

For a fair comparison, we uniformly use the full-attention PLM
BART${}_{\text{base}}$666The checkpoint of BART is “facebook/bart-base”
containing around 139 M parameters, whose maximum encoding length is 1024.
with 6 layers each for the encoder and decoder as our backbone on the six
public datasets above. In detail, the hidden size of each layer is 768, which
is converted into 12 attention heads with a hidden unit size of 64 for multi-
head attention. Built on the BART${}_{\text{base}}$, our framework introduces
an additional parameterized selector to filter out 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 $\beta_{1}=0.9$, $\beta_{2}=0.999$ for BART${}_{\text{base}}$
and selector respectively.

Additionally, we use a chunk size of 512 (considering the start token and end
token) in all experiments. To maintain chunk-wise alignment, we pad the origin
long sequence to ensure that all chunks have the same size. During the forward
propagation, the embedding layer embeds position representations for each
chunk independently.777Considering 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)

During training, we set the maximum input length to 16384 if not stated
otherwise. For efficient training, we update the parameters of the original
backbone and selector alternatively. The reward estimation of each action is
computed based on decoder cross-attention and the training objective 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 encoder 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).

###  5.4 Evaluations

Like most previous works, for summarization tasks, we measure the quality of
generated summaries using the popular metric ROUGE Lin (2004). On the test set
of arXiv, PubMed, GovReport, 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
BERTScore888In order to make a fair comparison with the previous work, we also
use checkpoint “facebook/bart-large-mnli” for BERTScore. Zhang* et al. (2020)
to demonstrate the superiority of our approaches comprehensively.

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.

##  6 Discussion

###  6.1 Results on Long-document Dataset

Tab. 1 reports results over three long-document test sets. We note that
casting the backbone BART${}_{\text{base}}$ into our SimCAS can significantly
improve the model performance on abstractive summarization tasks. On the arXiv
and GovReport test set, the base-size BART enhanced by our framework SimCAS
achieves comparable performance to original large-size models. Our approach
achieves new state-of-the-art performance using the base-size BART on the
PubMed test set.

###  6.2 Results on Multi-document Dataset

System | R-1 | R-2 | R-L | BS
---|---|---|---|---
BART${}_{\text{large}}$ | 42.04 | 14.88 | 23.34 | -
PEGASUS${}_{\text{large}}$ | 47.52 | 18.72 | 24.91 | -
HiMAP | 44.17 | 16.05 | 21.38 | -
GraphSum | 45.87 | 17.56 | 23.39 | -
BART-Long-Graph | 49.24 | 18.99 | 23.97 | -
LED+RELAX | 47.23 | 18.86 | 25.03 | -
PRIMERA | 49.94 | 21.05 | 25.85 | -
BART${}_{\text{base}}$+SimCAS | 48.91† | 19.92† | 45.02† | 65.01†
Table 2: Average results on Multi-News test set. $\dagger$: significantly better than BART ($p<0.01$). R-1/2/L is the ROUGE-1/2/L F1 score. BS refers to the model-based metrics BERTScore. The best results are bolded. System | R-1 | R-2 | R-L | BS
---|---|---|---|---
BERTREG | 35.00 | 13.50 | 25.50 | -
SUBMODULAR+ABS | 34.40 | 13.10 | 25.00 | -
DynE | 35.40 | 15.10 | 25.60 | -
LED | 39.79 | 18.94 | 32.10 | -
LED+RELAX | 41.11 | 19.46 | 33.13 | -
PRIMERA | 46.08 | 25.21 | 37.86 | -
BART${}_{\text{base}}$+SimCAS | 45.20† | 22.48† | 37.30† | 70.46†
Table 3: Average results on WCEP test set. $\dagger$: significantly better
than BART ($p<0.01$). R-1/2/L is the ROUGE-1/2/L F1 score. BS refers to the
neural model-based metrics BERTScore. The best results are bolded.

The results of multi-document summarization tasks in Tab. 2 and 3 have a
similar trend. With the exception of PRIMERA, which tailors a pre-training
objective for multi-document summarization, our method substantially surpasses
previous results.

###  6.3 Results on reading comprehension Dataset

Figure 2: System performance comparison on NarrativeQA test set. Besides the
commonly used metric F1, we additionally use metric EM to measure the
performance of our model.

See Fig. 2, we are surprised to find that the performance of our model on the
NarrativeQA test set is greatly improved compared with prior works. Given the
fact that this dataset has an extremely long average input length (>100K) and
short average output length (<10), we believe that the large improvement in
performance is due to our method’s ability to filter out a tremendous of task-
irrelevant information, which allows the decoding blocks to fuse information
more efficiently.

###  6.4 Analysis

We further analyze the properties of our framework to gain more insights.

Model | R-1 | R-2 | R-L | BS
---|---|---|---|---
SimCAS | 59.87 | 26.36 | 56.95 | 68.52
w/o SBA | -0.89 | -0.45 | -0.81 | -0.28
w/o selector | -4.93 | -2.02 | -2.82 | -1.73
Table 4: Ablation study results on GovReport development set. Performance
changes compared with the full model are reported. Larger decreases in metrics
are shaded with a darker orange.

Ablation Study We verify the contributions of various components in SimCAS on
the GovReport development set and show the ablation study results in Tab. 4.
Specifically, we consider taking out SBA and the reinforcement learning-based
selector respectively. On the grounds of the results, we can come to the
conclusion that 1) removing SBA and the selector substantially hurt model
performance; 2) although using a selector reduces the number of the token
representations, it can still improve the performance of the model because it
can acquire and fuse the information that is more task-specific.
Effectiveness of Conditional Skipping Reward In practice, to circumvent the
selector selecting exceedingly many token representations, we design a
conditional reward mechanism for skipping actions to soft constrain the number
of selected representations. Fig. 3 exhibits the effectiveness of our reward
mechanism. We can see that as the training progresses, the number of selected
tokens gradually converges to the threshold (2048).
Results on Larger Model

Figure 3: As the training progresses, the actual number of tokens entered into the model and the number of tokens selected by our selector. The red dash line represents the boundary for the conditional skipping reward. Model | R-1 | R-2 | R-L | BS
---|---|---|---|---
PRIMERA (SOTA) | 46.08 | 25.21 | 37.86 | -
BART${}_{\text{base}}$+SimCAS | 45.20 | 22.48 | 37.30 | 70.46
BART${}_{\text{large}}$+SimCAS | 46.29 | 24.45 | 38.61 | 71.38
Table 5: Performance on WCEP test set. R-1/2/L are the ROUGE-1/2/L F1 score.
BS refers to the model-based metric BERTScore. The best results are bolded.

We apply SimCAS to larger-scale backbone BART${}_{\text{large}}$ on the WCEP
dataset. Due to memory limitations, we use a maximum input length of 10240
instead of 16384 in the training phase. The experimental results in Tab. 5
show that with the increase in the number of parameters, the performance of
our method can be further improved by a large margin, and even outperform the
previous task-specific strong baseline.
Increasing Maximum Input Length We explore the performance change of our model
on a single 32G V100 GPU with the increase of input sequence length during
inference. First, we test the efficiency of our framework in processing long
sequences. In our setup, we use BARTbase as the backbone. Each text is encoded
and filled to the maximum input length using padding tokens. Fig. 4 shows that
in the case of small magnitude, with the exponential increase of input length,
the time cost still increases approximately linearly. We believe that this
phenomenon occurs because our method significantly reduces the computation
complexity and short sequences make insufficient use of the computation cores.
We also note that processing the sequence containing 350K tokens in the
inference phase approaches the limit of memory. Additionally, we notice that
even as the number of input tokens increases dramatically, the selected tokens
are kept at a reasonable size, indicating the capability of our selector to
filter out low-contribution information.

We also investigate how the model performs when the maximum effective input
length is increased during inference. In Tab. 6, we can observe that although
the maximum input length is set to 16384 during the training phase due to the
memory limitation, exceeding this maximum length during inference still
improves the model performance.

##  7 Conclusion and Future Work

Figure 4: Effect of increasing the number of input tokens in the inference phase on the time latency and the number of selected tokens. The area marked in blue on the right represents the limit of the number of tokens that the V100 can handle (350K tokens). Maximum Input Length | R-1 | R-2 | R-L | BS
---|---|---|---|---
4096 | 44.95 | 22.58 | 37.15 | 69.74
8192 | 45.38 | 22.78 | 37.34 | 70.46
16384⋆ | 45.20 | 22.48 | 37.30 | 70.46
32768 | 45.30 | 22.56 | 37.37 | 70.52
$+\infty$ | 45.56 | 22.76 | 37.61 | 70.55
Table 6: Performance of BART${}_{\text{base}}$-SimCAS on WCEP with different
maximum input lengths during inference. R-1/2/L is the ROUGE-1/2/L F1 score.
BS refers to the model-based metric BERTScore. $\star$: the maximum input
length in the training phase. $+\infty$: the input sequence to the model is
complete. The best results are bolded.

In this paper, we introduced a simple method for long-text processing via
chunking, aligning, and selecting, called SimCAS. We divided the input long
sequence into chunks, and encode 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 optimization method.
We leverage the transformer decoder 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 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
have the potential to enhance the long text pre-training to transformers.

## Limitations

Although our method significantly reduces the computational complexity of
full-attention Transformers, training on long-text datasets still costs a
large amount of GPU memory, even if the batch size per device is set to $1$.
Although we have achieved competitive results with base-size pre-trained
language models, how to improve efficiency for extremely large language models
remains unexplored. Besides, similar to most of the language generation
models, our method has not paid attention to the generation controllability Hu
et al. (2017); Prabhumoye et al. (2020); He et al. (2022). Therefore, we
cannot guarantee that the generated content never contains repetitive,
useless, or harmful information.

## Ethics Statement

Similar to existing generative models, there is no guarantee that the
generated content of our proposed method is factually consistent and free from
hallucination Maynez et al. (2020); Kang and Hashimoto (2020). Therefore,
caution is required when applying the model to real-world scenarios.

## References

  * Ainslie et al. (2020) Joshua Ainslie, Santiago Ontanon, Chris Alberti, Vaclav Cvicek, Zachary Fisher, Philip Pham, Anirudh Ravula, Sumit Sanghai, Qifan Wang, and Li Yang. 2020.  ETC: Encoding long and structured inputs in transformers.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)_ , pages 268–284, Online. Association for Computational Linguistics.
  * Akkaya et al. (2019) Ilge Akkaya, Marcin Andrychowicz, Maciek Chociej, Mateusz Litwin, Bob McGrew, Arthur Petron, Alex Paino, Matthias Plappert, Glenn Powell, Raphael Ribas, et al. 2019.  Solving rubik’s cube with a robot hand.  _arXiv preprint arXiv:1910.07113_.
  * Beltagy et al. (2020) Iz Beltagy, Matthew E Peters, and Arman Cohan. 2020.  Longformer: The long-document transformer.  _arXiv preprint arXiv:2004.05150_.
  * Bertsch et al. (2023) Amanda Bertsch, Uri Alon, Graham Neubig, and Matthew R. Gormley. 2023.  Unlimiformer: Long-range transformers with unlimited length input.
  * Brown et al. (2020) Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. 2020.  Language models are few-shot learners.  In _Advances in Neural Information Processing Systems_ , volume 33, pages 1877–1901. Curran Associates, Inc.
  * Byrne et al. (2021) Bill Byrne, Karthik Krishnamoorthi, Saravanan Ganesh, and Mihir Kale. 2021.  TicketTalk: Toward human-level performance with end-to-end, transaction-based dialog systems.  In _Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers)_ , pages 671–680, Online. Association for Computational Linguistics.
  * Choromanski et al. (2020) Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. 2020.  Rethinking attention with performers.  _arXiv preprint arXiv:2009.14794_.
  * Cohan et al. (2018) Arman Cohan, Franck Dernoncourt, Doo Soon Kim, Trung Bui, Seokhwan Kim, Walter Chang, and Nazli Goharian. 2018.  A discourse-aware attention model for abstractive summarization of long documents.  In _Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 2 (Short Papers)_ , pages 615–621, New Orleans, Louisiana. Association for Computational Linguistics.
  * Denning (2005) Peter J. Denning. 2005.  The locality principle.  _Commun. ACM_ , 48(7):19–24.
  * Devlin et al. (2018) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2018.  Bert: Pre-training of deep bidirectional transformers for language understanding.  _arXiv preprint arXiv:1810.04805_.
  * Dosovitskiy et al. (2021) Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. 2021.  An image is worth 16x16 words: Transformers for image recognition at scale.  In _International Conference on Learning Representations_.
  * Ethayarajh (2019) Kawin Ethayarajh. 2019.  How contextual are contextualized word representations? comparing the geometry of bert, elmo, and gpt-2 embeddings.  _arXiv preprint arXiv:1909.00512_.
  * Fabbri et al. (2019a) Alexander Fabbri, Irene Li, Tianwei She, Suyi Li, and Dragomir Radev. 2019a.  Multi-news: A large-scale multi-document summarization dataset and abstractive hierarchical model.  In _Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics_ , pages 1074–1084, Florence, Italy. Association for Computational Linguistics.
  * Fabbri et al. (2019b) Alexander R Fabbri, Irene Li, Tianwei She, Suyi Li, and Dragomir R Radev. 2019b.  Multi-news: A large-scale multi-document summarization dataset and abstractive hierarchical model.  _arXiv preprint arXiv:1906.01749_.
  * Fournier et al. (2023) Quentin Fournier, Gaétan Marceau Caron, and Daniel Aloise. 2023.  A practical survey on faster and lighter transformers.  _ACM Comput. Surv._ Just Accepted.
  * Gholipour Ghalandari et al. (2020) Demian Gholipour Ghalandari, Chris Hokamp, Nghia The Pham, John Glover, and Georgiana Ifrim. 2020.  A large-scale multi-document summarization dataset from the Wikipedia current events portal.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_ , pages 1302–1308, Online. Association for Computational Linguistics.
  * Gong et al. (2020) Hongyu Gong, Yelong Shen, Dian Yu, Jianshu Chen, and Dong Yu. 2020.  Recurrent chunking mechanisms for long-text machine reading comprehension.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_ , pages 6751–6761, Online. Association for Computational Linguistics.
  * Gupta and Berant (2020) Ankit Gupta and Jonathan Berant. 2020.  Gmat: Global memory augmentation for transformers.  _arXiv preprint arXiv:2006.03274_.
  * He et al. (2022) Junxian He, Wojciech Kryscinski, Bryan McCann, Nazneen Rajani, and Caiming Xiong. 2022.  CTRLsum: Towards generic controllable text summarization.  In _Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing_ , pages 5879–5915, Abu Dhabi, United Arab Emirates. Association for Computational Linguistics.
  * Hokamp et al. (2020) Chris Hokamp, Demian Gholipour Ghalandari, Nghia The Pham, and John Glover. 2020\.  Dyne: Dynamic ensemble decoding for multi-document summarization.
  * Hu et al. (2017) Zhiting Hu, Zichao Yang, Xiaodan Liang, Ruslan Salakhutdinov, and Eric P Xing. 2017\.  Toward controlled generation of text.  In _International conference on machine learning_ , pages 1587–1596. PMLR.
  * Huang et al. (2021) Luyang Huang, Shuyang Cao, Nikolaus Parulian, Heng Ji, and Lu Wang. 2021.  Efficient attentions for long document summarization.  In _Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies_ , pages 1419–1436, Online. Association for Computational Linguistics.
  * Ivgi et al. (2023) Maor Ivgi, Uri Shaham, and Jonathan Berant. 2023.  Efficient Long-Text Understanding with Short-Text Models.  _Transactions of the Association for Computational Linguistics_ , 11:284–299.
  * Kakade (2001) Sham M Kakade. 2001.  A natural policy gradient.  _Advances in neural information processing systems_ , 14.
  * Kang and Hashimoto (2020) Daniel Kang and Tatsunori B. Hashimoto. 2020.  Improved natural language generation via loss truncation.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_ , pages 718–731, Online. Association for Computational Linguistics.
  * Kazemnejad et al. (2023) Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy. 2023.  The impact of positional encoding on length generalization in transformers.
  * Kitaev et al. (2020) Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. 2020.  Reformer: The efficient transformer.  In _International Conference on Learning Representations_.
  * Klafka and Ettinger (2020) Josef Klafka and Allyson Ettinger. 2020.  Spying on your neighbors: Fine-grained probing of contextual embeddings for information about surrounding words.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_ , pages 4801–4811, Online. Association for Computational Linguistics.
  * Konda and Tsitsiklis (1999) Vijay Konda and John Tsitsiklis. 1999.  Actor-critic algorithms.  In _Advances in Neural Information Processing Systems_ , volume 12. MIT Press.
  * Kočiský et al. (2018) Tomáš Kočiský, Jonathan Schwarz, Phil Blunsom, Chris Dyer, Karl Moritz Hermann, Gábor Melis, and Edward Grefenstette. 2018.  The NarrativeQA Reading Comprehension Challenge.  _Transactions of the Association for Computational Linguistics_ , 6:317–328.
  * Kwiatkowski et al. (2019) Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, Kristina Toutanova, Llion Jones, Matthew Kelcey, Ming-Wei Chang, Andrew M. Dai, Jakob Uszkoreit, Quoc Le, and Slav Petrov. 2019.  Natural Questions: A Benchmark for Question Answering Research.  _Transactions of the Association for Computational Linguistics_ , 7:453–466.
  * Lewis et al. (2020) Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Veselin Stoyanov, and Luke Zettlemoyer. 2020.  BART: Denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehension.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_ , pages 7871–7880, Online. Association for Computational Linguistics.
  * Li et al. (2020) Wei Li, Xinyan Xiao, Jiachen Liu, Hua Wu, Haifeng Wang, and Junping Du. 2020.  Leveraging graph to improve abstractive multi-document summarization.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_ , pages 6232–6243, Online. Association for Computational Linguistics.
  * Lin (2004) Chin-Yew Lin. 2004.  ROUGE: A package for automatic evaluation of summaries.  In _Text Summarization Branches Out_ , pages 74–81, Barcelona, Spain. Association for Computational Linguistics.
  * Liu et al. (2022) Yixin Liu, Ansong Ni, Linyong Nan, Budhaditya Deb, Chenguang Zhu, Ahmed Hassan Awadallah, and Dragomir Radev. 2022.  Leveraging locality in abstractive text summarization.  In _Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing_ , pages 6081–6093, Abu Dhabi, United Arab Emirates. Association for Computational Linguistics.
  * Ma et al. (2022) Congbo Ma, Wei Emma Zhang, Mingyu Guo, Hu Wang, and Quan Z. Sheng. 2022.  Multi-document summarization via deep learning techniques: A survey.  _ACM Comput. Surv._ , 55(5).
  * Ma et al. (2021) Xuezhe Ma, Xiang Kong, Sinong Wang, Chunting Zhou, Jonathan May, Hao Ma, and Luke Zettlemoyer. 2021.  Luna: Linear unified nested attention.  In _Advances in Neural Information Processing Systems_.
  * Mao et al. (2022) Ziming Mao, Chen Henry Wu, Ansong Ni, Yusen Zhang, Rui Zhang, Tao Yu, Budhaditya Deb, Chenguang Zhu, Ahmed Awadallah, and Dragomir Radev. 2022.  DYLE: Dynamic latent extraction for abstractive long-input summarization.  In _Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 1687–1698, Dublin, Ireland. Association for Computational Linguistics.
  * Maynez et al. (2020) Joshua Maynez, Shashi Narayan, Bernd Bohnet, and Ryan McDonald. 2020.  On faithfulness and factuality in abstractive summarization.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_ , pages 1906–1919, Online. Association for Computational Linguistics.
  * Moro et al. (2022) Gianluca Moro, Luca Ragazzi, Lorenzo Valgimigli, and Davide Freddi. 2022.  Discriminative marginalized probabilistic neural method for multi-document summarization of medical literature.  In _Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 180–189, Dublin, Ireland. Association for Computational Linguistics.
  * Nallapati et al. (2016) Ramesh Nallapati, Bowen Zhou, Cicero dos Santos, Çağlar Gulçehre, and Bing Xiang. 2016.  Abstractive text summarization using sequence-to-sequence RNNs and beyond.  In _Proceedings of the 20th SIGNLL Conference on Computational Natural Language Learning_ , pages 280–290, Berlin, Germany. Association for Computational Linguistics.
  * Narayan et al. (2018) Shashi Narayan, Shay B. Cohen, and Mirella Lapata. 2018.  Don’t give me the details, just the summary! topic-aware convolutional neural networks for extreme summarization.  In _Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing_ , pages 1797–1807, Brussels, Belgium. Association for Computational Linguistics.
  * Ouyang et al. (2022) Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. 2022\.  Training language models to follow instructions with human feedback.  _Advances in Neural Information Processing Systems_ , 35:27730–27744.
  * Pan et al. (2021) Xiao Pan, Mingxuan Wang, Liwei Wu, and Lei Li. 2021.  Contrastive learning for many-to-many multilingual neural machine translation.  In _Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers)_ , pages 244–258, Online. Association for Computational Linguistics.
  * Pang et al. (2022) Richard Yuanzhe Pang, Alicia Parrish, Nitish Joshi, Nikita Nangia, Jason Phang, Angelica Chen, Vishakh Padmakumar, Johnny Ma, Jana Thompson, He He, and Samuel Bowman. 2022.  QuALITY: Question answering with long input texts, yes! In _Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies_ , pages 5336–5358, Seattle, United States. Association for Computational Linguistics.
  * Parnell et al. (2022) Jacob Parnell, Inigo Jauregi Unanue, and Massimo Piccardi. 2022.  A multi-document coverage reward for relaxed multi-document summarization.
  * Pasunuru et al. (2021) Ramakanth Pasunuru, Mengwen Liu, Mohit Bansal, Sujith Ravi, and Markus Dreyer. 2021\.  Efficiently summarizing text and graph encodings of multi-document clusters.  In _Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies_ , pages 4768–4779, Online. Association for Computational Linguistics.
  * Paszke et al. (2019) Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zach DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. 2019.  _PyTorch: An Imperative Style, High-Performance Deep Learning Library_. Curran Associates Inc., Red Hook, NY, USA.
  * Peng et al. (2021) Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah Smith, and Lingpeng Kong. 2021.  Random feature attention.  In _International Conference on Learning Representations_.
  * Phang et al. (2022) Jason Phang, Yao Zhao, and Peter J Liu. 2022.  Investigating efficiently extending transformers for long input summarization.  _arXiv preprint arXiv:2208.04347_.
  * Prabhumoye et al. (2020) Shrimai Prabhumoye, Alan W Black, and Ruslan Salakhutdinov. 2020.  Exploring controllable text generation techniques.  In _Proceedings of the 28th International Conference on Computational Linguistics_ , pages 1–14, Barcelona, Spain (Online). International Committee on Computational Linguistics.
  * Ramamurthy et al. (2023) Rajkumar Ramamurthy, Prithviraj Ammanabrolu, Kianté Brantley, Jack Hessel, Rafet Sifa, Christian Bauckhage, Hannaneh Hajishirzi, and Yejin Choi. 2023.  Is reinforcement learning (not) for natural language processing: Benchmarks, baselines, and building blocks for natural language policy optimization.  In _The Eleventh International Conference on Learning Representations_.
  * Roy et al. (2021) Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. 2021.  Efficient Content-Based Sparse Attention with Routing Transformers.  _Transactions of the Association for Computational Linguistics_ , 9:53–68.
  * Rush et al. (2015) Alexander M. Rush, Sumit Chopra, and Jason Weston. 2015.  A neural attention model for abstractive sentence summarization.  In _Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing_ , pages 379–389, Lisbon, Portugal. Association for Computational Linguistics.
  * Sanh et al. (2019) Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. 2019.  Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter.  _arXiv preprint arXiv:1910.01108_.
  * Schulman et al. (2016) John Schulman, Philipp Moritz, Sergey Levine, Michael I. Jordan, and Pieter Abbeel. 2016.  High-dimensional continuous control using generalized advantage estimation.  In _4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings_.
  * Schulman et al. (2017) John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. 2017\.  Proximal policy optimization algorithms.  _CoRR_ , abs/1707.06347.
  * Shaham et al. (2022) Uri Shaham, Elad Segal, Maor Ivgi, Avia Efrat, Ori Yoran, Adi Haviv, Ankit Gupta, Wenhan Xiong, Mor Geva, Jonathan Berant, and Omer Levy. 2022.  SCROLLS: Standardized CompaRison over long language sequences.  In _Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing_ , pages 12007–12021, Abu Dhabi, United Arab Emirates. Association for Computational Linguistics.
  * Sun et al. (2019a) Chen Sun, Austin Myers, Carl Vondrick, Kevin Murphy, and Cordelia Schmid. 2019a.  Videobert: A joint model for video and language representation learning.  In _Proceedings of the IEEE/CVF international conference on computer vision_ , pages 7464–7473.
  * Sun et al. (2019b) Siqi Sun, Yu Cheng, Zhe Gan, and Jingjing Liu. 2019b.  Patient knowledge distillation for bert model compression.  _arXiv preprint arXiv:1908.09355_.
  * Tay et al. (2023) Yi Tay, Mostafa Dehghani, Samira Abnar, Hyung Won Chung, William Fedus, Jinfeng Rao, Sharan Narang, Vinh Q. Tran, Dani Yogatama, and Donald Metzler. 2023.  Scaling laws vs model architectures: How does inductive bias influence scaling?
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017.  Attention is all you need.  In _Advances in Neural Information Processing Systems_ , volume 30. Curran Associates, Inc.
  * Vinyals et al. (2019) Oriol Vinyals, Igor Babuschkin, Wojciech M. Czarnecki, Michaël Mathieu, Andrew Dudzik, Junyoung Chung, David H. Choi, Richard Powell, Timo Ewalds, Petko Georgiev, Junhyuk Oh, Dan Horgan, Manuel Kroiss, Ivo Danihelka, Aja Huang, L. Sifre, Trevor Cai, John P. Agapiou, Max Jaderberg, Alexander Sasha Vezhnevets, Rémi Leblond, Tobias Pohlen, Valentin Dalibard, David Budden, Yury Sulsky, James Molloy, Tom Le Paine, Caglar Gulcehre, Ziyun Wang, Tobias Pfaff, Yuhuai Wu, Roman Ring, Dani Yogatama, Dario Wünsch, Katrina McKinney, Oliver Smith, Tom Schaul, Timothy P. Lillicrap, Koray Kavukcuoglu, Demis Hassabis, Chris Apps, and David Silver. 2019.  Grandmaster level in starcraft ii using multi-agent reinforcement learning.  _Nature_ , pages 1–5.
  * Wang et al. (2020) Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. 2020.  Linformer: Self-attention with linear complexity.
  * Wiseman and Rush (2016) Sam Wiseman and Alexander M. Rush. 2016.  Sequence-to-sequence learning as beam-search optimization.  In _Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing_ , pages 1296–1306, Austin, Texas. Association for Computational Linguistics.
  * Wolf et al. (2020) Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Remi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander Rush. 2020.  Transformers: State-of-the-art natural language processing.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations_ , pages 38–45, Online. Association for Computational Linguistics.
  * Wysocki et al. (2023) Oskar Wysocki, Zili Zhou, Paul O’Regan, Deborah Ferreira, Magdalena Wysocka, Dónal Landers, and André Freitas. 2023.  Transformers and the Representation of Biomedical Background Knowledge.  _Computational Linguistics_ , 49(1):73–115.
  * Xiao et al. (2022) Wen Xiao, Iz Beltagy, Giuseppe Carenini, and Arman Cohan. 2022.  PRIMERA: Pyramid-based masked sentence pre-training for multi-document summarization.  In _Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 5245–5263, Dublin, Ireland. Association for Computational Linguistics.
  * Xiong et al. (2022) Wenhan Xiong, Barlas Oguz, Anchit Gupta, Xilun Chen, Diana Liskovich, Omer Levy, Scott Yih, and Yashar Mehdad. 2022.  Simple local attentions remain competitive for long-context tasks.  In _Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies_ , pages 1975–1986, Seattle, United States. Association for Computational Linguistics.
  * Yang et al. (2019a) Wei Yang, Yuqing Xie, Aileen Lin, Xingyu Li, Luchen Tan, Kun Xiong, Ming Li, and Jimmy Lin. 2019a.  End-to-end open-domain question answering with BERTserini.  In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics (Demonstrations)_ , pages 72–77, Minneapolis, Minnesota. Association for Computational Linguistics.
  * Yang et al. (2019b) Zonghan Yang, Yong Cheng, Yang Liu, and Maosong Sun. 2019b.  Reducing word omission errors in neural machine translation: A contrastive learning approach.  In _Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics_ , pages 6191–6196, Florence, Italy. Association for Computational Linguistics.
  * Yuan et al. (2021) Weizhe Yuan, Graham Neubig, and Pengfei Liu. 2021.  Bartscore: Evaluating generated text as text generation.  In _Advances in Neural Information Processing Systems_ , volume 34, pages 27263–27277. Curran Associates, Inc.
  * Zaheer et al. (2020) Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. 2020.  Big bird: Transformers for longer sequences.  _Advances in neural information processing systems_ , 33:17283–17297.
  * Zhang et al. (2020) Jingqing Zhang, Yao Zhao, Mohammad Saleh, and Peter Liu. 2020.  PEGASUS: Pre-training with extracted gap-sentences for abstractive summarization.  In _Proceedings of the 37th International Conference on Machine Learning_ , volume 119 of _Proceedings of Machine Learning Research_ , pages 11328–11339. PMLR.
  * Zhang* et al. (2020) Tianyi Zhang*, Varsha Kishore*, Felix Wu*, Kilian Q. Weinberger, and Yoav Artzi. 2020.  Bertscore: Evaluating text generation with bert.  In _International Conference on Learning Representations_.
  * Zhang et al. (2022) Yusen Zhang, Ansong Ni, Ziming Mao, Chen Henry Wu, Chenguang Zhu, Budhaditya Deb, Ahmed Awadallah, Dragomir Radev, and Rui Zhang. 2022.  Summn: A multi-stage summarization framework for long input dialogues and documents.  In _Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 1592–1604, Dublin, Ireland. Association for Computational Linguistics.

##  Appendix A Dataset Statistics

Dataset | Train | Valid | Test | Avg. Input Tokens
---|---|---|---|---
arXiv | 203.0K | 6.4K | 6.4K | 6021
PubMed | 119.9K | 6.6K | 6.7K | 3049
GovReport | 17.5K | 1.0K | 1.0K | 9616
Multi-News | 44.9K | 5.6K | 5.6K | 1793
WCEP | 8.1K | 1.0K | 1.0K | 3866
NarrativeQA | 32.7K | 3.5K | 10.6K | 121736
Table 7: Statistics of used datasets.

##  Appendix B Potential to select representation

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. Fig.5 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.

##  Appendix 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. 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-Long-Graph 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 multiple-
document for abstractive summarization by single document summarization
models.

##  Appendix D Compatible with short-input tasks

Figure 5: Visualization of the average cross-attention of all attention heads
at all layers in the decoder (excluding start position). This example is
generated by the BART${}_{\text{base}}$+SimCAS trained on GovReport. Figure 6:
Comparison of the performance of full-attention PLM BART and sparse-attention
PLM LED on short-text summarization datasets CNNDM Nallapati et al. (2016) and
XSum Narayan et al. (2018). AVG ROUGE denotes the average of ROUGE-1/2/L F1
scores.

While the sparse-attention transformers have been proven effective on a wide
range of long-sequence datasets, as shown in Fig. 6, 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,
aligning, and selecting procedures.

##  Appendix E Efficacy of SBA

Figure 7: Evolution of similarities between start hidden states of each chunk
after each encoding layer during one forward propagation. Only cosine
similarity is used in the top half of the figure. The lower part takes into
account both cosine similarity and Euclidean distance.

As our framework introduces SBA to align inter-chunk information during the
encoding steps, we explore the change of hidden state in the forward
propagation process under the influence of this component. Fig. 7 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
${\bm{v}}_{1},{\bm{v}}_{2}\in\mathbb{R}^{d}$, the custom similarity Sim
between them is formulated as follows:

$\textbf{Sim}=\frac{({\bm{v}}_{1},{\bm{v}}_{2})}{\lVert{\bm{v}}_{1}\rVert*\lVert{\bm{v}}_{2}\rVert*(1+\lVert{\bm{v}}_{1}-{\bm{v}}_{2}\rVert)},$

by which, we scale the cosine similarity to observe its change. As can be seen
from the lower part of Fig. 7, after each encoder layer, the start hidden
states will vary in scale due to different contexts.

##  Appendix F More Experimental Details

Input Structure Since we aim to minimize dataset-specific modeling, we unify
the processing of the used long-document summarization, multi-document
summarization, and machine reading comprehension datasets. Specifically, as
for the long-document datasets arXiv, PubMed, and GovReport, we simply input
the long document and truncate it to the maximum input length, which means
that all sections are separated by newlines. Following the tokenization form
of BART, for the multi-document datasets Multi-News and WCEP, all the
documents are concatenated to form a long sequence like the long document
example and separated with </s>. In the NarrativeQA dataset for machine
reading comprehension, we concatenate together the question, summary, and
source document, each of which prepends a prefix to indicate the type of text.
Before the embedding operation, we conduct the chunking procedure in our
method.
Training Details In this paper, for the training of BART${}_{\text{base}}$ on
various datasets we uniformly employ the Adam optimizer with the following
dynamic learning rate:

$lr=5\times
10^{-3}\min(\textit{step}^{-0.5},\textit{step}\times\textit{warmup}^{-1.5}),$

where warmup indicates the warmup steps, which is set to 2500, 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\times
10^{-4}$. The optimization of the selector and transformer is performed
alternately, with one model trained while the other model is fixed.
The Details of Selector In our setting, the selector consists of an actor
component and a critic component with PPO optimization, both of which are a
simple three-layer 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. And 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 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).

##  Appendix G Case Study on GovReport Dataset

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. Fig. 8 displays several examples of summaries generated by the fine-
tuned base model BART${}_{\text{large}}$ and our
BART${}_{\text{base}}$-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.

Figure 8: Example summary generated by BART and SimCAS trained on GovReport
dataset. The maximum input length of standard BART and SimCAS is 1024 and
16384 respectively. The sentence in green is included in the SimCAS summary,
while the one in red is discarded.


Experimental Setup Our experiments were performed on a single p3.8xlarge
Amazon EC2 instance with 32 vCPUs and 4 Nvidia V100 GPUs (with 16 of memory
each).

Micro-Benchmarks

For this analysis, we use a simple chained matrix multiplication micro-
benchmark. Inputs are taken from the data layer and fed through a series of
multiplications on the GPU before writing the final matrix output back to the
data layer. Each multiplication uses a constant matrix that is read from the
data layer but will not change between invocations (enabling caching). We
configured the benchmark to use three layers with square matrices with a side
length of 1024 single precision floats. The baseline gtask implementation
required the numpy, pickle, and pycuda python modules. For cold starts, the
system has not seen any requests for the benchmark. For gtask, this means that
it must create a new worker and initialize any python dependencies before
running the gtask code. For ktask, the system already has kaas executors
initialized since they are independent of any particular request. However,
those executors must parse the request, link against the specified cuda
libraries, and load all data from the data layer. On warm starts, both task
types already have a worker ready to execute the request immediately.

[b]0.49 [width=]figs/microWarm.png [b]0.49 [width=]figs/microCold.png [b]0.5
[width=]figs/microKaasOnly.png [b]0.35 [width=]figs/microLegend.png

Figure : Warm Start Figure : Cold Start Figure : kaas Cold vs Warm Start
Figure : Microbenchmark Results _Kernel Run_ is the time spent actually
executing the request’s kernels. _Kernel Init_ is the time to link any needed
precompiled cuda libraries. _GPU Malloc_ and _GPU Copy_ encompass all device
data management while _Data Layer_ measures host data management. _Overheads_
measure any additional tasks performed by the worker including request
parsing, worker initialization, and Ray framework overheads.

Figure  shows the results of this experiment. We begin by observing that ktask
and gtask have similar warm-start performance (Figure ). This is not
surprising as both implementations must perform the same steps. Any
differences are mostly down to implementation details. Next, we move our
attention to cold starts (Figure ). Here we see the drastic impact of starting
a new Python process for gtask. Even though the microbenchmark only imports a
minimal set of packages, it still takes an additional 400 to load (a 46x
increase). In contrast, the only cold-start overheads experienced by ktask
come from warming the data caches with the constant matrices. To see this in
more detail, Figure  compares warm and cold starts experienced by ktask. As
expected, the time to load an additional three matrices adds significant time
to both the data layer and GPU memory management phases, but there are no
other overheads. Multitenant Workload As we saw with the microbenchmarks,
ktask behave similarly to gtask for a single warm workload. However, when
there are multiple clients, the behavior changes dramatically. As discussed in
§LABEL:sec:pool, the gtask approach requires a policy that assigns GPUs to
workloads exclusively, falling back to terminating and restarting workers as
the number of workloads exceeds the available GPUs. These cold starts prevent
the system from fully utilizing its GPU resources. kaas, in contrast, is able
to share a limited pool of GPUs more effectively.

[width=0.65]figs/expPipeline.png

Figure : Multitenant Environment: Clients submit requests over a ZeroMQ socket
to a server that submits per-request func pipelines to Ray. Each request
pipeline involves optional CPU-only pre and post-processing steps with a GPU-
based execution phase. The server proxies pipeline outputs back to the client.

We now evaluate four real-world applications in a multi-tenant environment
with two scenarios: online and offline. In the online scenario, we evaluate
how the system supports latency-sensitive workloads like model serving or
interactive data analytics where requests arrive according to a Poisson
process. To understand sustained throughput, we evaluate an offline scenario
where all workloads submit requests as fast as possible. Two of our workloads,
resnet50 and BERT, are deep-learning inference workloads generated by the TVM
interface. These workloads consist of many small kernels operating over many
small buffers. Deep learning models also have significant data re-use through
the model weights. The cGEMM workload is a chained complex number matrix
multiply using our Cutlass library interface. It multiplies a large constant
$10000\times 25000$ matrix by a narrow $100\times 10000$ matrix that changes
on each invocation. Both contain 32 complex floats. This results in small
inputs and outputs but a large amount of cacheable data. Finally, the Jacobi
workload uses the low-level kaas interface to implement an iterative solver
using the Jacobi method [jacobiSource]. This workload uses kaas’s iteration
control-flow mechanism to run a fast update kernel for 3000 iterations. The
input matrix contains $512\times 512$ float32s and is accumulated into a
$512\times 1$ float64 output. There is no data re-use between invocations.
Table  summarizes the properties of these workloads.

| resnet50 | BERT | cGEMM | Jacobi
---|---|---|---|---
[l]Constant |  |  |  |
Memory |  129 |  1.3 |  2.0 |  0
[l]Dynamic |  |  |  |
Memory |  6 |  6 |  8 |  1
GPU Time |  4 |  92 |  39 |  52
CPU Time |  10 |  132 |  0 |  0
Table : End to end workload properties.

In addition to the GPU-based func, the workloads also include pre and post
processing func that occur exclusively on the host with no GPU requirements.
These func form a graph that is lazily executed as inputs become available.
Finally, this experiment uses multiple external clients that communicate with
the Ray-based service through a network protocol using ZeroMQ [zmq]. Figure
shows the complete setup. Offline Workloads We begin with a throughput-
oriented scenario where workloads from many clients submit requests as quickly
as possible with no regard for latency. In this case, we measure the aggregate
throughput of the system. This indicates whether or not the GPU resources
remain highly utilized as the number of clients increases. Figures  and  show
the results from this experiment. We see that ktask and gtask both perform
well when there are sufficient GPU resources to support the workloads.
However, as we exceed the number of available GPUs, the gtask approach’s
throughput drops dramatically due to frequent worker cold starts. ktask do not
suffer significant slowdowns as we exceed available GPUs because the system is
able to handle requests from different clients without restarting the kaas
executor. The ktask implementation of cGEMM only begins to slow after about 20
replicas. This is the point at which we begin to exceed available GPU memory
for cached weights and kaas must begin evicting and re-loading constant
buffers. Unlike the sharp performance drop seen with gtask, there is a more
gradual performance impact with ktask since the cache can be managed at a fine
granularity. Even at 32 replicas, we still see significantly higher throughput
with kaas. This comes down to the difference in cold-start behaviors between
the two approaches. For gtask, the system must boot a fresh python process and
initialize any packages and state needed to run the model, while ktask need
only reload model weights from the host data cache. Unlike cGEMM, resnet50 is
a small model that fits easily within GPU memory. It experiences no
significant slowdowns, even at 32 replicas. However, single-client warm-start
performance is somewhat slower for ktask than the gtask implementation. This
is due to implementation differences between the two systems. The gtask
implementation uses TVM’s highly optimized C++ runtime while kaas uses a
simple Python implementation. While these differences are easily amortized by
the larger kernels used by BERT or cGEMM, resnet50 consists of many very small
kernels that accentuate system overheads. Fundamentally, the same work is done
in both implementations and a more optimized kaas executor implementation
should behave similarly to TVM.

[width=0.85]figs/throughputLowMem.png

Figure : Aggregate throughput of low-memory workloads. These workloads fit
easily within GPU memory and see no aggregate throughput loss in kaas as we
increase the number of replicas. gtask workloads require exclusive access to
GPUs and must begin cold-starting workers after 4 replicas.

[b]0.8 [width=]figs/cdfBert4.png [b]0.8 [width=]figs/cdfBert5.png [b]0.8
[width=]figs/cdfBert16.png

Figure : CDFs of BERT response latency for different numbers of replicas.

[width=0.85]figs/throughputHighMem.png

Figure : Aggregate throughput of high-memory workloads. These workloads
consume a large amount of GPU memory. When the aggregate memory requirements
exceed available GPU memory, kaas must begin evicting and re-loading objects
from the GPU memory cache. This leads to a gradual decline in performance as
cache pressure increases.

[b]0.335 [width=]figs/latDistResults_resnet50.png [b]0.335
[width=]figs/latDistResults_jacobi.png [b]0.335
[width=]figs/latDistResults_bert.png [b]0.335
[width=]figs/latDistResults_cGEMM.png

Figure : resnet50 Figure : Jacobi Figure : BERT Figure : cGEMM Figure : Online
benchmark results The median response latency is plotted along with the 90th
percentile tail.

Online Workloads In this experiment, we simulate a latency-sensitive online
environment using the MLPerf Inference load generator in server mode
[mlperfinference]. Clients submit requests to the framework according to a
Poisson process with their mean arrival rate set to 80% of peak throughput (to
ensure stability). We report the median and 90th percentile response latency
for each configuration. Figure  shows the results of this experiment. We see
very similar behavior to the offline scenario. When the number of clients does
not exceed available resources, both approaches behave reasonably well.
However, when we exceed available resources, the gtask approach is forced to
cold-start new actors for nearly every request, resulting in very poor tail.
In effect, the tail becomes a measure of cold start time. This latency depends
primarily on the python dependencies and initialization tasks for each
workload. ktask are able to effectively share GPUs, resulting in no
significant impact on tail, even as we exceed the number of available GPUs.
However, the cGEMM ktask begins to slow down after about 20 replicas as the
aggregate constant memory exceeds GPU memory. Still, even at 32 replicas, the
cGEMM ktask’s tail is far below the gtask version due to the difference in
cold start performance. Figure  shows the CDFs of response latency at key
points in the configuration space. Similar to Figure , we configure the
experiment to submit requests at 80% of peak throughput to ensure queue
stability. The maximum throughput of gtask is significantly lower than ktask
so we plot CDFs for ktask at both an equivalent submission rate to gtask, and
at peak ktask throughput. At 4 replicas, each replica has a dedicated GPU. In
this case, both ktask and gtask behave similarly with only minor differences
due to implementation details. As we exceed 4 GPUs, the system must begin
sharing GPUs between replicas. Even with only one extra replica, we see the
gtask tail suffer significantly since replicas may need to evict an existing
gtask before cold-starting. As we increase the concurrency to 16 replicas,
even the best 10th percentile latency suffers as replicas become nearly
guaranteed to require a cold start. In contrast, ktask experience very
consistent latencies as the number of replicas increase. Even at 16 replicas,
there is no noticeable reduction in latency.