Efficient Classification of Long Documents via State-Space Models (2023)

Abstract: Transformer-based models have achieved state-of-the-art performance on numerous NLP applications. However, long documents which are prevalent in real-world scenarios cannot be efficiently processed by transformers with the vanilla self-attention module due to their quadratic computation complexity and limited length extrapolation ability. Instead of tackling the computation difficulty for self-attention with sparse or hierarchical structures, in this paper, we investigate the use of State-Space Models (SSMs) for long document classification tasks. We conducted extensive experiments on six long document classification datasets, including binary, multi-class, and multi-label classification, comparing SSMs (with and without pre-training) to self-attention-based models. We also introduce the SSM-pooler model and demonstrate that it achieves comparable performance while being on average 36% more efficient. Additionally our method exhibits higher robustness to the input noise even in the extreme scenario of 40%.

## introduction since the emergence of large-scale pre-trained language models such as bert @xcite and @xcite , these transformerbased models have become popular solutions for many text classification and generation tasks. however, their benefit is constrained to short-length inputs when the computation resource is limited because attention module requires quadratic computation time and space. more specifically, each token in a sequence of length n requires pairwise computation with all n tokens, which results in o(n 2 ) complexity. such limitation makes transformerbased models hard to process long sequential data efficiently. there are many works aiming to improve the performance on long document classi-fication for transformers @xcite . one of the common approaches is to simply truncate long texts to a pre-defined length, e.g. 512, which makes pre-trained models to be applicable for them. some work demonstrated this technique is not sufficient for long documents due to the missing of important information @xcite . another sort of technique attempts to reduce the computation overhead of attention-based systems. this problem has several relevant solutions, e.g. sparse attention models @xcite and hierarchical attention models @xcite . one of the important sparse attention methods is longformer, which leverages local and global attention to reduce the computational complexity of the models and increases the input length up to 4096 tokens. another popular sparse attention method is bigbird @xcite : besides the global and local attention, it includes extra random attention modules to attend to a predefined number of random tokens. apart from designing sparse attention mechanisms, hierarchical transformers (han) like tobert @xcite propose to construct systems on top of the conventional transformer @xcite . basically, the long text is first split into several chunks less than a fixed number, e.g. 200. next, every chunk is encoded by a vanilla transformer one by one to form a collection of chunk representations and then another transformer processes the sequence of chuck representations. apart from the computation complexity issue, some works pointed out that the length extrapolation ability of self-attention-based models is quite limited @xcite , namely, the transformer models perform poorly during inference if the sequence length of test data is beyond the sequence length of training data. as an order-invariant encoding mechanism, the selfattention-based encoder heavily relies on the position embedding (pes) to model input orders, how-ever, these works demonstrate that the failure of transformers on long sequence is due to the limited length generalization ability of these position embedding methods. this also encourages exploring alternative architectures for the challenge of long document classification problems. recently, @xcite propose modeling sequential data with state-space models and show surprisingly good performance on a benchmark for comparing transformers over long sequential data @xcite . however, this benchmark only consists of one single text classification task. it is still unclear about the efficacy and efficiency of state-space models for long document tasks compared to transformer-based models. in this paper, we aim to fill this gap with a more comprehensive by conducting extensive experiments and analysis on six long document classification benchmarks. besides, we develop an efficient long document system and show that the state-space-based models outperform self-attentionbased models (including sparse attention and han) and yield consistent performance across datasets with much higher efficiency and more robustness to input noise. ## experiments in this section, we present our experimental setup and results for multiple datasets. ## discussion in this section, we investigate the training efficiency, the impact of different down-sampling structures and the robustness to input noise of ssmbased models. ## conclusion in this work, we conduct a comprehensive evaluation of the ssm-based models and show their superiority over self-attention-based models on long document classification tasks in terms of performance and training time. we further propose an efficient ssm-based system by imposing input length reduction for deeper layers and show our method performs on par with the original ssm models while greatly improving time efficiency.