Sparsifying Transformer Models with Trainable Representation Pooling

TL;DR

This paper introduces a trainable pooling method to sparsify attention in Transformer models, significantly reducing computational complexity while maintaining high performance on long document summarization tasks.

Contribution

It presents a novel trainable top-k operator for attention sparsification, enabling faster training and inference without sacrificing model quality.

Findings

Achieved sublinear complexity in attention computation.

Maintained state-of-the-art performance on long document summarization.

Reduced training time by 1.8x and inference time by 4.5x.

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.