# CAST: Clustering Self-Attention using Surrogate Tokens for Efficient Transformers

arXiv: https://arxiv.org/abs/2402.04239 · HTML: https://mlsanity.com/paper/2402.04239/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 25 of its 40 references that are in the map.
- First posted on arXiv 6 Feb 2024.

## Builds on (most cited first, 25 of 25)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [MLP-Mixer: An all-MLP Architecture for Vision](https://mlsanity.com/paper/2105.01601/) (2021)
- [Effective Approaches to Attention-based Neural Machine Translation](https://mlsanity.com/paper/1508.04025/) (2015)
- [Long Range Arena: A Benchmark for Efficient Transformers](https://mlsanity.com/paper/2011.04006/) (2020)
- [FNet: Mixing Tokens with Fourier Transforms](https://mlsanity.com/paper/2105.03824/) (2021)
- [Utilizing BERT for Aspect-Based Sentiment Analysis via Constructing Auxiliary Sentence](https://mlsanity.com/paper/1903.09588/) (2019)
- [Learning long-range spatial dependencies with horizontal gated-recurrent units](https://mlsanity.com/paper/1805.08315/) (2018)
- [Luna: Linear Unified Nested Attention](https://mlsanity.com/paper/2106.01540/) (2021)
- [Leveraging BERT for Extractive Text Summarization on Lectures](https://mlsanity.com/paper/1906.04165/) (2019)
- [SMYRF: Efficient Attention using Asymmetric Clustering](https://mlsanity.com/paper/2010.05315/) (2020)
- [Transformer Acceleration with Dynamic Sparse Attention](https://mlsanity.com/paper/2110.11299/) (2021)
- [Classification of Long Sequential Data using Circular Dilated Convolutional Neural Networks](https://mlsanity.com/paper/2201.02143/) (2022)
- [ChordMixer: A Scalable Neural Attention Model for Sequences with Different Lengths](https://mlsanity.com/paper/2206.05852/) (2022)
- [Paramixer: Parameterizing Mixing Links in Sparse Factors Works Better than Dot-Product Self-Attention](https://mlsanity.com/paper/2204.10670/) (2022)
- [Sparse Factorization of Large Square Matrices](https://mlsanity.com/paper/2109.08184/) (2021)

## Abstract (arXiv)

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(N2) to O(α N) where N is the sequence length, and α 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.