CAST: Clustering Self-Attention using Surrogate Tokens for Efficient
  Transformers
Adjorn van Engelenhoven Nicola Strisciuglio Estefanía Talavera
The Transformer architecture has shown to be a powerful tool for a wide range
of tasks. It is based on the self-attention mechanism, which is an inherently
computationally expensive operation with quadratic computational complexity:
memory usage and compute time increase quadratically with the length of the
input sequences, thus limiting the application of Transformers. In this work,
we propose a novel Clustering self-Attention mechanism using Surrogate Tokens
(CAST), to optimize the attention computation and achieve efficient
transformers. CAST utilizes learnable surrogate tokens to construct a cluster
affinity matrix, used to cluster the input sequence and generate novel cluster
summaries. The self-attention from within each cluster is then combined with
the cluster summaries of other clusters, enabling information flow across the
entire input sequence. CAST improves efficiency by reducing the complexity from
$O(N^2)$ to $O(\alpha N)$ where N is the sequence length, and {\alpha} is
constant according to the number of clusters and samples per cluster. We show
that CAST performs better than or comparable to the baseline Transformers on
long-range sequence modeling tasks, while also achieving higher results on time
and memory efficiency than other efficient transformers.