Sparsifying Transformer Models With Trainable Representation Pooling

Pietruszka Michał, Borchmann Łukasz, Garncarek Łukasz. Arxiv 2020

[Paper]    
Applications Attention Mechanism Model Architecture Pretraining Methods Training Techniques Transformer

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.

Similar Work