LOCOST : State-Space Models for Long Document Abstractive Summarization (2024)

Abstract: State-space models are a low-complexity alternative to transformers for encoding long sequences and capturing long-term dependencies. We propose LOCOST: an encoder-decoder architecture based on state-space models for conditional text generation with long context inputs. With a computational complexity of \mathcal{O}(L \log L) , this architecture can handle significantly longer sequences than state-of-the-art models that are based on sparse attention patterns. We evaluate our model on a series of long document abstractive summarization tasks. The model reaches a performance level that is 93-96% comparable to the top-performing sparse transformers of the same size while saving up to 50% memory during training and up to 87% during inference. Additionally, LOCOST effectively handles input texts exceeding 600K tokens at inference time, setting new state-of-the-art results on full-book summarization and opening new perspectives for long input processing.

## introduction nowadays the design of efficient models for long texts remains an open challenge despite the recent progress achieved in natural language processing (nlp). the introduction of transformer architectures @xcite indeed came as a major bump in performance and scalability for text generation. however the quadratic complexity in the input length still restricts the application of large pretrained models to long texts. for instance, bert @xcite and bart @xcite are limited to a context size of 512 and 1024 tokens respectively, which amounts to 2-3 paragraphs of standard text. to mitigate this issue, a straightforward approach is to leverage sparse-attention patterns @xcite to better cope with long texts. as key examples, @xcite and @xcite extended the context capacity of encoderdecoder models @xcite and showed drastic increases in the performance on long text summarization, motivating the quest to incorporate longer contexts. however, in practice, even the best sparse-transformers need heavy computational resources to handle sequences of length larger than 8k tokens (see figure 4 ). deep state-space models (ssms) @xcite have been proposed for sequence processing, with complexity o(l log l), initially for computer vision and audio and more recently for text. their recurrent architectures are designed for capturing long-range dependencies @xcite . up to now, their applications have been restrained to either unconditional autoregressive generation, i.e., with a decoder-only @xcite ; or sequence classification, i.e., with an encoder-only @xcite . tackling conditional text generation with ssms as required e.g. for summarization remains yet unexplored. in this paper, we propose locost an encoder-decoder architecture to explore the performance of ssms for conditional text generation tasks, through the lens of abstractive summarization. we demonstrate that ssms can be competitive with transformer-based models while drastically reducing their memory requirements. we opt for a lightweight architecture design, comparable to the average base transformers (roughly 250m parameters) in order to process extremely long sequences on standard compute resources. our experimentations with extremely long sequences yield stateof-the-art results on the challenging booksum-book. with an increase of up to 2 points in average rouge score compared to sparse attention baselines, our model is able to process entire books, without truncation, and on a single gpu. our contributions are threefold: to the best of our knowledge, this is the first encoder-decoder that performs competitively with sparse transformers with no attention in the encoder. furthermore, this work represents the first successful attempt at processing extremely long texts e.g. entire books without any truncation, all in a single pass. the proposed model opens new perspectives for addressing long texts with lesser resources. * ## related work in this section, we first review memory-efficient transformers and existing alternatives to the attention mechanism. then, we discuss recent literature on state-space models. memory efficiency for transformers. reducing the memory consumption of transformers is an active research field. optimization at the hardware level @xcite helped to improve the scaling of the attention computation on recent gpus. a line of work considers retrieving-augmented transformers, like @xcite , that use additional modules to enhance the language modeling backbone. while crucial in developing memory-efficient architectures, we consider these last two topics as being orthogonal to our work that focuses on the models' architecture. profuse literature focuses on tailoring the models' architecture for long inputs. since the computational complexity of attention comes from the computation of the self-attention matrix, a straightforward way to reduce its cost is to approximate it using sparse-attention patterns. these patterns typically incorporate a combination of local attention and a set of carefully selected tokens. for instance, in addition to global tokens, bigbird @xcite considers random tokens, while lsg (condevaux and harispe, 2023) considers sparse tokens through various strategy of sparsification. longt5 @xcite chunks the sequence into blocks and averages their representations, which gives a number of global tokens equal to the number of blocks. an overview of the complexity of various sparse-transformers can be found in in contrast, we propose an alternative, computationally efficient architecture, without the need of costly self-attention blocks nor sparse-attention patterns. attention-free transformers. some variants of transformers already avoid the standard attention mechanism. for example @xcite ; @xcite approximate the softmax similarity in the attention by a more efficient computation. more recently, mixing architectures were introduced in @xcite . they are the main component of the fnet (lee-thorp et al., 2022) model, an encoder that replaces self-attention with a discrete fourier transform (dft). fnet has a complexity of o(l log l) and is an encoder-only model, thus restricted to classification and regression tasks. our proposed model also bypasses attention in the encoder, reaching the same computational complexity as encoders such as fnet, while being a much more versatile model, specifically designed for conditional text generation. encoder architecture complexity per layer transformer (full) o(l 2 ) led o(lw) bigbird o(lw + l(g + r)) lsg o(lw + l(g + s)) longt5 (tglobal) o(lw + l ⌊l/c⌋) locost o(l log(l)) ## background for contextualization, we leverage state-space models instead of self-attention. throughout the paper, l denotes the sequence length, h the embedding dimension and n the dimension of the state-space hidden state (to be introduced in section 3). before delving into our model in section 4, we describe below the main components of the state-space architecture and elaborate on their potential for long sequence processing. state-space models. for unidimensional inputs @xmath0, ..., u l-1 ) ∈ r l , deep ssms @xcite are based on the recurrent equation: state-space convolution. by unrolling the recurrence above, the output sequence @xmath1: y @xmath2, ∀@xmath3, ..., l}. let * denote the causal convolution operator (details about this operator are in appendix a). then, we can define a convolution kernel κ ∈ r l that depends on a, b, c. a ssm layer is therefore parametrized by a, b, c, d through κ and its output is defined by y as in the following equation: for multidimensional @xmath4, we simply compute h convolutions with one kernel κ h for each dimension. ## model in this section, we present the locost model. we first introduce the bidirectional deep state-space model, then show how to use it to enable global contextualization of the tokens. then, we present bisssm dropout ssm forward convolution ssm backward convolution layer norm gated feedforward projection projection (a) the locost layer. dense (2048, 768) gelu dense (768, 2048) dropout dense (768, 2048) (b) gated feedforward net. the architecture of the locost layer with an efficient contextualization that can be used as a dropin replacement for the self-attention mechanism in transformers. ## experiments to validate our experiments, we focus on the long document abstractive summarization task as it represents a typical conditional generation problem with long input requirements. ## limitations though we investigated lightweight models for computational reasons, scaling the architecture to a larger size could be studied. we focused on long document abstractive summarization, we leave for future work the study of ssms on other long inputs abstractive tasks. although replacing self-attention with state-space encoders drastically reduces the computational complexity, the use of dense crossattention in the decoder still limits the output sequence length in terms of computation during training. ## ethics statement we performed pre-training on a subset of the c4 dataset, which has been identified to include inappropriate content like hate speech and explicit material, as noted in the studies conducted by @xcite and also exhibits negative biases towards certain ethnicities @xcite . it is important to investigate potential solutions for mitigating these problems through more meticulous preprocessing in order to prevent the emergence of such undesirable attributes in future research. nevertheless, it is worth mentioning that despite these concerns, the c4 dataset serves as a benchmark within the community, and the reported results solely focus on the quality of the summaries, thereby avoiding any unethical implications. in this paper, we consider a relatively small size for locost. we believe our work could be reproducible with limited resources. we tracked the gpu power consumption during pre-training. the average power usage was 190w per gpu. we trained for 140 hours on 16 gpus. given the local co 2 intensity of 58 gco 2 /kwh * , we can estimate that approximately 25kg of co 2 have been emitted during the pre-training, to be compared with the average emissions of 4.6t of co 2 par capita in 2019 * .