We propose a novel method to sparsify attention in the Transformer model by learning to select the most-informative token representations during the training process, thus focusing on the task-specific parts of an input. A reduction of quadratic time and memory complexity to sublinear was achieved due to a robust trainable topk operator. Our experiments on a challenging long document summarization task show that even our simple baseline performs comparably to the current SOTA, and with trainable pooling we can retain its top quality, while being 1 . 8  faster during training, 4 . 5  faster during inference and up to 13  more computationally efficient in the decoder. 1 1 Introduction The introduction of Transformer architecture led to an immense improvement in the performance of Natural Language Processing systems (Vaswani et al., 2017; Radford et al., 2018; Devlin et al., 2019). Nevertheless, the underlying attention mechanism is marked by the original sin of quadratic memory complexity w.r.t. the input sequence length. It results from the attention matrix reflecting inter-connections between every two representations in the input sequence. Previous approaches either reduce the full connectivity of its elements to its non-empty subset or approximate the self-attention matrix (Dai et al., 2019; Beltagy et al., 2020; Kitaev et al., 2020; Tay et al., 2020; Zaheer et al., 2020a; Wang et al., 2020; Shen et al., 2021; Choromanski et al., 2021; Roy et al., 2021). In particular, in these models, each word at every layer attends to at least one other word. In contrast, we disregard attention for a given representation completely in the case of non-informative ones (Figure 1 and 2). In particular, we optimize the attention complexity by learning to select encoded representations for the given task and promoting only the chosen ones to the next layer of the model. This mechanism will be referred to as representation pooling . Consequently, a significantly 1 Code publicly available at https://github.com/ applicaai/pyramidions along with trained models. lower memory consumption and an improved processing time are achieved. As the selection operation has to be trainable, we provide a suitable high-performance continuous relaxation of topk , robust for every k value and input sequence length. We demonstrate this idea's applicability by performing on par to state-of-the-art on the challenging problem of long document summarization. Simultaneously, the proposed end-to-end model is a significant theoretical improvement over the previous systems, which are based on independently trained extractive and abstractive models. Contribution. The specific contributions of this paper are the following: (1) We propose a method to sparsify Transformer architecture in a novel, previously unrecognized way, achieving sublinear time and memory complexity. Our model learns to select the subset of best representations depending on the advantage they give on a downstream task. (2) Additionally, we demonstrate an improvement of the decoder's cross-attention complexity. It is beneficial for both train/inference time and memory consumption. (3) We demonstrate an elegant way to train extractive-abstractive models in an end-to-end manner with only a cross-entropy loss function. (4) We present a Successive Halving Topk operator that outperforms previous approaches in terms of approximation quality and speed. We provide a detailed analysis of its differential properties and prove that it is trainable in an end-to-end manner, making it applicable within our neural networks. (5) We achieve state-of-the-art performance level in long document's summarization and show that previous models can be outperformed by a straightforward baseline. Word-vector elimination. It has been previously shown that the progressive elimination of word vectors occurring layer after layer can improve inference time of transformer-based language models used in a text classification scenario (Goyal et al., 2020). We extend this notion to tasks demanding text generation in a way that, contrary to previous work, is trainable and optimized concerning a downstream task. A similar approach has been taken in the Funnel Transformer proposed concurrently to our work (Dai et al., 2020). We directly compare to both methods' adaptations (see Section 5), and consider our work to surpass it in two aspects: 1) results were improved due to a better pooling mechanism than mean/max; 2) training was accelerated, which we attribute to the significant reduction of the decoder's complexity. Sparse attention. Several authors proposed to limit attention connectivity, e.g., by dividing input into smaller 'blocks' (Child et al., 2019; Beltagy et al., 2020; Rae and Razavi, 2020). Blockwise attention is an optional element of our architectures, used in addition to trainable pooling. Summarization. In terms of the type of summarization task we target, our representation pooling mechanism can be considered an end-to-end extractive-abstractive model. This is a conceptual breakthrough compared to recently proposed two-stage hybrids that extract and paraphrase in two independent steps, using separately trained modules (Pilault et al., 2020; Hsu et al., 2018; Gehrmann et al., 2018; Chen and Bansal, 2018). It is suspected that when humans engage in information search, they use various cognitive processes depending on the relevance level of constituent text fragments (Gwizdka et al., 2017). The method we propose is inspired by this search for relevant fragments, which is an important aspect of human cognition when engaged in reading to do actions (Mosenthal, 1996; Mosenthal and Kirsch, 1992). We intend to mimic relevance judgments and hypothesize that it is possible to answer problems involving natural language with only selected passages of the input text. These passages may be of substantially shorter length than the original text. One may compare this to a person reading the paper and highlighting in such a way that it is possible to provide a summary using only the highlighted parts. The end-to-end mechanism we introduce performs such highlighting by scoring the representations and passes only the selected ones to the next layer of the neural network (Figure 3). The role of the selection is to reduce data resolution in a roughly similar way to how pooling works in CNNs, where the feature map is downsampled and only the most informative activations are retained. When pooling in a trainable manner at the bottleneck of the encoder-decoder, it impacts the encoding process because the additional, orthogonal, informational bottleneck forces the model to compress more context into one representation vector of constant-length, leveraging the already provided capacity. Let n denote the number of input tokens that are projected onto d dimensions, resulting in a matrix of embedding representations E  R n  d . We want to assign scores v i to embedding vectors E i , in such a way that v i measures the usefulness of E i for further layers and the training objective. Typically, this can be achieved by defining a scoring function S : R d  R (which we allow to depend on additional parameters, thus making it trainable) that assigns a usefulness score to every embedding vector, and putting v i = S ( E i ) . (1) Next, we use our soft topk operator : R n  d  R n  R k  d to reduce the number of embeddings from 8617 encoding Tok 1 Tok N Tok N+1 Tok M Tok M+1 Tok L ... ... ... ... ... ... n to k , based on their usefulness scores. The k vectors produced by  form the input for the next network layer. The path of residual connections starts on a reduced number of tokens. Flavors. We consider two architectures in this work: with single or multiple pooling layers (Figure 1). Specifically, the latter is a generalization of the former to any given number of pooling layers. We use the term Transpooler when a single pooling layer is placed after the encoder. This setup directly limits the amount of information passed to the decoder through the network's bottleneck. However, pooling can be applied between any subsequent layers, such that multiple operations of this type will be used in the network and gradually introduce the bottleneck along the encoding process. As a result, the same model bottleneck size can be achieved as when using Transpooler. Moreover, the decision to pool earlier has the advantage of attaining more substantial memory complexity reduction. This model will be referred to as the Pyramidion. Blockwise attention. When propagating through layers, we use blockwise attention and split input into nonoverlapping chunks in such a way that the full quadratic attention is computed for each chunk. The score is then determined for each representation vector, and after selecting with the topk operator, chosen representations are passed to the next layer. We assure our topk operator selects representations without permuting their order, keeping them in line with their original position. Scoring functions. Multiple scoring methods can be proposed. The most straightforward is to use a linear scoring function as used in conventional token classification, S ( e ) = e T w + b , where w  R d and b  R are trainable parameters. We found it to work best with our pooling method. In the Appendix A we perform ablations on different scoring functions. Table 1 presents the complexity of attention in our models, and compares it to different architectures. The vanilla encoder depends on the number of layers l , the number of tokens in the input n and the number of tokens each attends to n . Likewise, the decoder's cross-attention depends on l , n and the target length t . The m denotes the effective number of tokens one can attend to, resulting from the attention's block size, allowed window size or the clustering of key-values. The number of parallel LSH hashes is denoted by h . The rank of the factorization matrix is r , which can be a constant that is independent of n . Similarly, the number of best task-specific representations k , selected after encoding, is independent of n . c is an effective number of layers in a hierarchically decreasing encoder of the Pyramidion. The Pyramidion's c can be as low as  . Blockwise sparse attention improved 8618 the vanilla Transformer's complexity by limiting the number of tokens each attends to from n (input length) to m (block size) as seen in Table 1. As we keep the encoding of blockwise attention, the m improvement also applies to our self-attention. For the Pyramidion model, we narrow down the size of the representation on the output of each chosen layer, leading to the exponential reduction of memory consumption as the encoding proceeds. For example, when pooling after every layer is considered, the total memory complexity across l layers would be (cid:80) pi =0 2  i mnd = (2  k/n ) mnd where p denotes the number of passes p = log 2 ( n/k ) , assuming k  n and n, k  { 2 i | i  Z + } . Hence, the effective complexity of all layers is lower than mnd , which means it is lower than  times the complexity of the full-size first layer. For the decoder cross-attention, the number of input representations that t target tokens can attend to is limited by k , thus decreasing the memory complexity of cross attention from O ( tn ) to O ( tk ) . Optimization over quadratic sentence-length complexity is even more powerful and needed on the decoder side, as O ( tn ) complexity hurts performance of real-world applications based on auto-regressive decoding. The blockwise attention itself reduces encoder complexity proportionally to the number of chunks. We further reduce the decoder layer's complexity in Transpooler models by a factor of n/k , thanks to representation pooling. The Pyramidion we propose offers an additional improvement on the encoder side, where time and memory consumption are reduced in each of the consecutive layers compared to the Transformer featuring blockwise attention. In other words, when b denotes the number of blocks, l stands for the number of layers, and the sequence length is halved in each layer, we reduce memory from b + b + ... + b = lb to b + b/ 2+ b/ 4+ ... + b/ (2 l )  2 b . Because the beneficial impact of pooling accumulates, we are able to improve complexity from one that is linearly dependent on l to one that is constant, independent of l . In the further DeepPyramidion's experiments, we will proceed with a higher reduction factor, where the length of a sequence is cut in four. As a result, the Pyramidion achieves an effective self-attention time and space complexity linear of n and logarithmic of l . For comparison, other sparse models such as, e.g., Linformer depend linearly on n and linearly on l . The analysis of Figure 4 found evidence that our method scales well with an increasing number of layers. In the evaluation (see Section 5), we demonstrate that our model achieves a 2 . 5  computation reduction in the encoder's self-attention and a 16  reduction in the decoder's cross-attention comparing to blockwise baseline, while both models are close to SOTA results on the task of long-document summarization. All things considered, we introduce Pyramidion with sublinear complexity that achieves remarkable results. The advantage of our approach is that it complements all other proposed sparsification techniques, thus paving a new interesting avenue of potential research. It can be effortlessly applied in-between layers and simultaneously with other improvements since representation pooling addresses a different aspect of the attention's complexity problem. The choice of the selection operator is challenging, as it has to be trainable to instantiate a pooler. In case of the hard topk operator, back-propagation through the scores is impossible and prevents training the scoring function. It could be seen as an extreme case of the vanishing gradient problem. In this section we introduce a mechanism not prone to this issue, while the Appendix B is dedicated to a theoretical analysis of its differential properties, from a geometrical point of view. The crux of our approach is the Successive Halving Topk selection mechanism that finds k convex combinations of vector representations E i , dominated by those achieving the highest scores v i (pseudocode available in the Appendix B.1). 2 The general idea is to perform a tournament soft selection, where candidate vectors are compared in pairs ( i, j ) , until only k remained. After each tournament's round new E (cid:48) and v (cid:48) are computed as convex combinations of these pairs with weights based on their respective scores. Each new vector is calculated as: E (cid:48) i = w i E i + w j E j , where the w i , w j are the result of a peaked softmax over the scores v i , v j . Analogously, we use v (cid:48) i = w i v i + w j v j as the new-round's scores. 2 Preliminary work regarding this method was previously presented in the form of a Student Abstract, see Pietruszka et al. (2020). Weights are calculated using a PeakedSoftmax function (Goyal et al., 2017), increasing the pairwise difference in scores between v i and v j . One round halves the number of elements in E and v . We perform it iteratively unless the size of E and v matches the chosen value of k . To improve convergence towards selecting the real topk , it is desired to permute v and E first. In our algorithm, we sort the vectors E i in descending order of their scores v i and then put them into the tournament in pairs of the form ( i, n + 1  i ) . This method of pairing guarantees that the weights w i depend monotonically on the scores v i , which is the main motivation for using it. Extended benchmarks for time and accuracy are covered in details in Appendix B.5. The main focus of the experiments was to understand how to employ the Successive Halving Topk operator within neural networks to build models that have better training and inference time and are expressive enough to achieve results comparable to state-of-the-art models. The first experiment was specifically designed to compare to other sparse Transformers and Vanilla baselines. Choice of tasks. We demonstrate the benefit of pooling on the arXiv and PubMed summarization datasets (Cohan et al., 2018) available under Apache License 2.0 license. Both tasks demand text generation and have the highest average input sequence length ( 6 k and 3 k words on average for arXiv and PubMed respectively). Assuming an embedding of dimensionality 768 , it is important to note that for inputs shorter than approx. 4 k tokens, more multiplications happen in the Transformer's FFN layers and projection layers than in the attention layers. Hence, the validation of the sparsification mechanism should be proved by showing that it works for longer inputs. Time benchmarks. The average time of processing a batch of documents is reported to evaluate the computational improvements experimentally. Decoding experiments were synthetic with a forced fixed length of 512 output tokens to discount for the lower processing time of models predicting an earlier sequence end. We recorded time in seconds on batches of size 64 and 8 for training and generation, respectively. Details regarding the hyperparameters and test environment are reported in Appendix C. 8620 Ablations on input and decoder lengths. Table 2 presents evaluation metrics and time benchmarks depending on encoder and decoder lengths, as well as used sparsification mechanisms. At this stage, we use shallow 4 -layer models to perform ablation studies and estimate each approach's strengths and weaknesses. We observe that all sparse models deliver on the promise of accelerating training time over Vanilla Transformers for longer sequences in this setup. Methods requiring the elimination of word vectors scale well with the sequence length but incur additional pooling costs, which may be notable for shorter sequences. Nevertheless, inference time was significantly reduced only when methods eliminating word vectors were employed. The introduction of blockwise attention and pooling does not decrease scores while lowering the computational cost. The detailed training procedure for all models is provided in Appendix C. Scaling deeper. In preliminary experiments it was estimated that the fastest-to-train model that performs comparably to the Vanilla Transformer is the Blockwise Transformer. Here, we scale it to 6 -layers in each encoder and decoder and provide an interesting baseline for our model, since Transpooler's backbone is blockwise attention. We undertook the empirical analysis of scaling Transpooler to many layers in Appendix C.2 and found that in order to balance performance and speed, it is crucial to delay the first pooling and not to perform it directly on the first layer's output. It was also revealed that appending more layers at the end of the encoder (after pooling) results in a negligible increase in time while considerably improving scores. Both changes to the block size and reduction of the bottleneck harmed the performance. Thus, the data supports the premise that the 6 -layers encoder should consume 8 k tokens on the input and output representations of lengths 8 k, 8 k, 2 k, 512 , 512 , 512 after each successive layer. We refer to this model as DeepPyramidion (note that pooling happens twice in the encoder). The decoder also has six layers, making our model directly comparable to the deeper Blockwise Transformer. We confront DeepPyramidion with the Blockwise baseline by training models from scratch on arXiv and PubMed datasets separately and report results in comparison to the state-of-the-art summarization models (Table 3). Results. The evaluation of the data presented in Table 3 leads to the unexpected conclusion that our Blockwise Transformer baseline, despite its simplicity, is sufficient to outperform deeper, denser, and additionally pretrained models that were recently reported as state-of-the-art. We demonstrate that DeepPyramidion retains or improves the performance of the competitive baseline we produced. The training time speedup by 1 . 8  supports the notion that our model scales better to long sequences, assuming deeper models. This result stands in line with evidence in Figure 4. While our baseline Blockwise model reduces the computational demand of self-attention in encoder by a factor of 16  when comparing to Vanilla Transformer, it does not improve the decoder's computational complexity. It is interesting to highlight that DeepPyramidion further lowers the cost of self-attention by 2 . 5  and improves 16  over Blockwise's cross-attention in the decoder, and leads to overall 13  improvement in the number of multiplication operations in the decoder. Time benchmarks show a 4 . 5  improvement in the generation times for our method, proving how vital the improvement in the decoder's cross-attention complexity is for inference time. DeepPyramidion achieves a ROUGE-2 score indistinguishable from SOTA on arXiv and performs competitively on PubMeb. At the same time, an entire DeepPyramidion costs five times less than a single Transformer layer consuming 8 k tokens. However, when comparing our results to those of older studies, it must be pointed out that our models were trained from scratch only on the targeted dataset, whereas prior works often base on already pretrained models such as BART or RoBERTa and leverage unsupervised training on additional datasets. On the contrary, a longer input sequence was consumed by both Blockwise and DeepPyramidion, which we speculate, is the reason for their strong performance. 3 Impact of longer inputs. The results achieved in our paper are comparable to other, much heavier, and more costly models due to two main reasons, that will be briefly discussed below. Firstly, to perform well on a long document summarization task, there is a need to strike the right balance not only between the depth and width of the network but also it is required for design optimization to take into account the length of the input. All previous work seem to underperform when considering all three factors, as they were designed and optimized for shorter tasks and generally have more parameters, denser computations, or even a hard limit on the range of positional encoding. The authors were thus bounded by the maximal sequence length of 512 or 1024 tokens. One can argue that within this prefix (corresponding to the first 2  3 pages), any data point from the arXiv/PubMed datasets (a scientific paper) usually provides enough information to write a meaningful summary, but also, important details will be missing to some degree. Hence, increasing the length of the input that can be consumed on GPUs, at the price of using a shallower network, with sparser computation, may be considered a better fit for the task. 3 This view is supported by results of PoolingFormer that are concurrent to our work (Zhang et al., 2021). Despite that, at first sight, the methods seem similar and the authors present an interesting use of pooling in the attention, we argue that the mentioned model suffers from several weaknesses that are not present in our work. First of all, in the PoolingFormer model vectors are not removed from computations in further layers. Hence logarithmic complexity of the number of layers does not apply. PoolingFormer's approach suffers from having three orders of magnitude more calculations than when a global pooling based on scores of individual tokens is considered. Secondly, we think that pretraining in the Pyramidion's case may be disregarded due to an interesting length exploiting hypothesis. That is, while we consume longer sequences on the input, the network learns more efficiently, as more information is available, and thus, the training signal is stronger. This can be convincingly portrayed in the case of embedding layers, as during training they see many more words and sentences from the chosen dataset, and hence, can provide more meaningful representations to the further layers. One can think that making the most of already available domain texts and consuming longer inputs is an advantageous approach to masked pretraining on out-of-domain datasets. While the latter approach may aid general' language understanding, it has insufficient transferability potential to domain-specific document understanding (e.g., scientific or medical texts). To sum up, the Pyramidion has improvements that allow consuming longer inputs cheaply, which turns out to be a more cost-effective strategy compared to other models. This aspect is crucial for achieving strong results on the presented datasets. At this stage of understanding, we believe that sparsification based on trainable pooling is unlikely to improve processing time for short sequences specific to some NLP tasks, e.g., sentence-level Neural Machine Translation. In addition, the score improvement may be attainable for tasks characterized by at least an order of magnitude shorter outputs than inputs, as it was previously shown on classification, or, as in the case of this work, on summarization. However, the extent to which it is possible to replace full-attention in Transformer with the sparse attention we propose is unknown. However, we argue that the benefits are visible starting from the inputs of length 4 k . As discussed earlier, 4 k is a break-even point where more calculations are needed for attention than for FFNs and projecting layers. As such, we recommend applying sparsification methods on datasets featuring sequences of length over that value. While we focus on the long end of the possible inputs, one can continue our analysis, to find improvements that work for shortest sequences, such as, e.g., concentrating on employing lighter projection layers and FFNs or stacking more attention blocks. Although our method is a hybrid extractive-abstractive, it does not provide interpretable explanations to which specific representations were selected as the pooling operates in the latent space. How to match the selected vectors to the vocabulary tokens remains an open question. Moreover, framing the trainable pooling for language modeling remains a challenge to address in future works, especially as in this task the Markov assumption may serve as a basis for competitive pooling heuristics. We did not consider Relative Positional Encoding in our work as pooling mechanism is not trivially applicable with it and some generalization of our method may be needed. In that case, as it demands more experiments and proofs, we will leave the generalization of the pooling method for future work. Regarding the social impact and environmental sus-tainability, we actively considered the Earth's wellbeing by contributing a technique for reducing the computational demand of recent Deep Learning models. Our near-state-of-the-art DeepPyramidion model costs us 3 days of training on 8 NVIDIA A100 GPUs. Shallow models featuring trainable pooling were finished in about 2 days each, given the same hardware. Blockwise baselines cost us about 3 . 5 x the price of respective pooling methods. The most prolonged training of the 8 k Vanilla Transformer lasted for about 2 weeks. The total cost of training the models covered in this paper is about 2 months on the mentioned hardware, plus an additional month for models and ablations described in the appendices. We roughly estimate that it is between half and one-fourth of the total computation spent, including false runs, unpublished work, and initial experiments. The dataset preparation took less than 10 hours on 1 CPU. We propose representation pooling as a method to reduce the complexity of Transformer encoder-decoder models. Specifically, we optimize self-attention com-8622 plexity and address the decoder's cross-attention complexity optimization, which has so far not been widely acknowledged by the research community. Moreover, the DeepPyramidion we introduced establishes results comparable to state-of-the-art, outperforming not only other systems relying on progressive word-vector elimination but also deeper, denser, and additionally pretrained models. We tackle the problem by introducing a novel method of applying successive halving to a model's input in a tournament style. It is a theoretical improvement over existing approaches in terms of both computational complexity and approximation quality. Trainable Top-k selection allows to train scorer for a task and outperforms other pooling methods. From the summarization task's point of view, the proposed end-to-end model is a significant theoretical improvement over the previous systems, where the extractive model was trained independently of the abstractive one. In contrast, our mechanism does not require the introduction of an additional training objective or training stage. Our approach can be easily applied to other problems from Natural Language Processing and Computer Vision. E.g., in a recent work later than ours, Multiscale Vision Transformers were proposed. These, similarly to our Pyramidion model, introduce the bottleneck gradually along the encoding process of videos and images, leading to better results, and complexity (Fan et al., 2021). As it comes to Natural Language Processing, possible applications include Key Information Extraction, Machine Reading Comprehension, and Question Answering in scenarios where encoder-decoder models struggle or would struggle with input sequence length (see, e.g., Choi et al. (2017); Townsend et al. (2021); Kocisk et al. (2018)). We are looking forward to seeing these opportunities exploited. For easy reproduction of the results, we release our utilities, code and pretrained models on the MIT license for all researchers not affiliated or working for Russian state-controlled institutions and public companies. The reason to ostracize scientists under those affiliations is the violent invasion of their armed forces on Ukraine, recklessly intended to inflict pain, threaten world peace and civilians life with nonhuman aggression against a sovereign nation. The authors would like to thank Zofia Prochoroff and Pawe Morawiecki for the helpful discussions on the draft of the paper. Moreover, we thank the reviewers for their comments and suggestions that helped improve the paper. The Smart Growth Operational Programme supported this research under project no. POIR.01.01.01-00-0877 /19-00 ( A universal platform for robotic automation of processes requiring text comprehension, with a unique level of implementation and service automation ).
objective abstain result abstain abstain abstain abstain abstain other abstain objective method abstain other abstain method objective abstain abstain objective abstain objective abstain objective result method result other other objective abstain abstain other other method other method other other objective objective other other method other other abstain abstain abstain method abstain abstain abstain abstain method abstain method abstain abstain abstain abstain abstain abstain method abstain abstain abstain abstain abstain abstain method method abstain abstain abstain abstain abstain abstain abstain abstain abstain method method abstain abstain abstain abstain abstain method abstain method abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain method abstain abstain abstain abstain abstain abstain abstain other other abstain abstain abstain abstain abstain abstain abstain abstain abstain method abstain abstain abstain abstain abstain abstain method abstain abstain abstain abstain abstain abstain other abstain abstain abstain abstain method abstain abstain abstain abstain abstain abstain other abstain abstain abstain abstain abstain abstain abstain abstain other abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain method abstain method abstain abstain method abstain abstain abstain abstain abstain abstain abstain method abstain abstain abstain method abstain abstain abstain abstain abstain abstain method abstain abstain abstain abstain other other other other other other