Sparsifying Transformer Models with Trainable Representation Pooling (2022)

Abstract: 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 top- k 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\times faster during training, 4.5\times faster during inference, and up to 13\times more computationally efficient in the decoder.

## introduction the introduction of transformer architecture led to an immense improvement in the performance of natural language processing systems @xcite @xcite . 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 @xcite @xcite @xcite @xcite @xcite . 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. ## novel approach of representation pooling it is suspected that when humans engage in information search, they use various cognitive processes depending on the relevance level of constituent text fragments @xcite . 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 @xcite . 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 constantlength, leveraging the already provided capacity. ## suitable top-k operator the choice of the selection operator is challenging, as it has to be trainable to instantiate a pooler. in case of the hard top-k 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 top-k 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 and v are computed as convex combinations of these pairs with weights based on their respective scores. each new vector is calculated as: to improve convergence towards selecting the real top-k, 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. ## evaluation the main focus of the experiments was to understand how to employ the successive halving top-k 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. ## limitations and social impact 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 4k. as discussed earlier, 4k 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 extractiveabstractive, 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 sustainability, 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.5x the price of respective pooling methods. the most prolonged training of the 8k 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 onefourth of the total computation spent, including false runs, unpublished work, and initial experiments. the dataset preparation took less than 10 hours on 1 cpu. ## summary we propose representation pooling as a method to reduce the complexity of transformer encoder-decoder models. specifically, we optimize self-attention com-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 @xcite . 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); kočiský et al. (2018)). we are looking forward to seeing these opportunities exploited. yen-chun chen and mohit bansal. 2018. fast abstractive summarization with reinforce-selected sentence rewriting. in proceedings of the 56th annual meeting of the association for computational linguistics (volume 1: long papers), pages 675-686, melbourne, australia. association for computational linguistics. rewon child, scott gray, alec radford, and ilya sutskever. 2019. generating long sequences with sparse transformers. eunsol choi, daniel hewlett, jakob uszkoreit, illia polosukhin, alexandre lacoste, and jonathan berant. 2017. coarse-to-fine question answering for long documents. in proceedings of the 55th annual meeting of the association for computational linguistics (volume 1: long papers), pages 209-220, vancouver, canada. association for computational linguistics. krzysztof marcin choromanski, valerii likhosherstov, david dohan, xingyou song, andreea gane, tamás sarlós, peter hawkins, jared quincy davis, afroz mohiuddin, lukasz kaiser, david benjamin belanger, lucy j. colwell, and adrian weller. 2021. rethinking attention with performers. in 9th international conference on learning representations, iclr 2021, virtual event, austria, may 3-7, 2021. openreview.net. 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. zihang dai, guokun lai, yiming yang, and quoc le. 2020. funnel-transformer: filtering out sequential redundancy for efficient language processing. in advances in neural information processing systems 33: annual conference on neural information processing systems 2020, neurips 2020, december 6-12, 2020, virtual. zihang dai, zhilin yang, yiming yang, jaime carbonell, quoc le, and ruslan salakhutdinov. 2019. transformer-xl: attentive language models beyond a fixed-length context. in proceedings of the 57th annual meeting of the association for computational linguistics, pages 2978-2988, florence, italy. association for computational linguistics.