  The transformer architecture has driven breakthroughs in recent years on
tasks which require modeling pairwise relationships between sequential
elements, as is the case in natural language understanding. However, long
seqeuences pose a problem due to the quadratic complexity of the attention
operation. Previous research has aimed to lower the complexity by sparsifying
or linearly approximating the attention matrix. Yet, these approaches cannot
straightforwardly distill knowledge from a teacher's attention matrix and often
require complete retraining from scratch. Furthermore, previous sparse and
linear approaches lose interpretability if they cannot produce full attention
matrices. To address these challenges, we propose SEA: Sparse linear attention
with an Estimated Attention mask. SEA estimates the attention matrix with
linear complexity via kernel-based linear attention, then subsequently creates
a sparse attention matrix with a top-k selection to perform a sparse attention
operation. For language modeling tasks (Wikitext2), previous linear and sparse
attention methods show roughly two-fold worse perplexity scores over the
quadratic OPT-1.3B baseline, while SEA achieves better perplexity than
OPT-1.3B, using roughly half the memory of OPT-1.3B, providing interpretable
attention matrix. We believe that our work will have a large practical impact,
as it opens the possibility of running large transformers on resource-limited
devices with less memory.
