Sparsifying Transformer Models with Trainable Representation Pooling

Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics Volume 1: Long Papers, pages 8616 - 8633 May 22-27, 2022 c⃝2022 Association for Computational Linguistics Sparsifying Transformer Models with Trainable Representation Pooling Michał Pietruszka1, 2 Łukasz Borchmann1, 3 Łukasz Garncarek1 1Applica.ai 2Jagiellonian University 3Poznan University of Technology {michal.pietruszka, lukasz.borchmann, lukasz.garncarek}@applica.ai Abstract We propose a novel method to sparsify atten- tion in the Transformer model by learning to select the most-informative token representa- tions during the training process, thus focusing on the task-specific parts of an input. A reduc- tion of quadratic time and memory complex- ity 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 train- able pooling we can retain its top quality, while being 1.8× faster during training, 4.5× faster during inference and up to 13× more computa- tionally 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). Neverthe- less, 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 con- nectivity 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 rep- resentation 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 1Code publicly available at https://github.com/ applicaai/pyramidions along with trained models. (A) Vanilla Transformer (B) Blockwise Encoder (D) Pyramidion Target length Input length Pooled length Block size Pool Pool Pool (C) Blockwise Encoder with Representation Pooling Figure 1: An illustration of sparse attention matrices as- suming a three-layer encoder and decoder (separated by the dashed line). The blue color reflects the memory con- sumption of self-attention (encoder) and cross-attention (decoder). (A) The complete input consumed at once. (B) Memory reduced with blockwise attention and (C) pooling applied after the encoder. (D) Gradual reduction of memory by pooling after every layer. Vanilla� Blockwise Pooling� Figure 2: Toy illustration of inter-connections consti- tuting the attention matrices in various approaches to attention. White dots denote disregarded representa- tions that are not attended to and removed from …

Content truncated to 500 words