CAST: Clustering Self-Attention using Surrogate Tokens for Efficient Transformers (2402.04239v1)
Abstract: 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(\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.
- Etc: Encoding long and structured inputs in transformers, 2020.
- Longformer: The long-document transformer, 2020.
- Lukas Biewald. Experiment tracking with weights and biases, 2020. URL https://www.wandb.com/. Software available from wandb.com.
- Language models are few-shot learners, 2020.
- Distilling knowledge learned in bert for text generation, 2019. URL https://arxiv.org/abs/1911.03829.
- Classification of long sequential data using circular dilated convolutional neural networks, 2022.
- Generating long sequences with sparse transformers, 2019. URL https://arxiv.org/abs/1904.10509.
- Rethinking attention with performers, 2020. URL https://arxiv.org/abs/2009.14794.
- Flashattention: Fast and memory-efficient exact attention with io-awareness, 2022.
- Hungry hungry hippos: Towards language modeling with state space models, 2023.
- Smyrf: Efficient attention using asymmetric clustering, 2020.
- An image is worth 16x16 words: Transformers for image recognition at scale, 2020. URL https://arxiv.org/abs/2010.11929.
- Efficiently modeling long sequences with structured state spaces, 2022.
- Liquid structural state-space models, 2022.
- Long short-term memory. Neural Computation, 9(8):1735–1780, 1997.
- Highly accurate protein structure prediction with alphafold. Nature, 596:583 – 589, 2021.
- Sparse factorization of large square matrices, 2021.
- Chordmixer: A scalable neural attention model for sequences with different lengths, 2023.
- Reformer: The efficient transformer, 2020. URL https://arxiv.org/abs/2001.04451.
- Alex Krizhevsky. Learning multiple layers of features from tiny images. pp. 32–33, 2009. URL https://www.cs.toronto.edu/~kriz/learning-features-2009-TR.pdf.
- Fnet: Mixing tokens with fourier transforms. arXiv preprint arXiv:2105.03824, 2021.
- What makes convolutional models great on long sequence modeling?, 2022.
- Learning long-range spatial dependencies with horizontal gated-recurrent units. CoRR, abs/1805.08315, 2018. URL http://arxiv.org/abs/1805.08315.
- Transformer acceleration with dynamic sparse attention. CoRR, abs/2110.11299, 2021a. URL https://arxiv.org/abs/2110.11299.
- Swin transformer: Hierarchical vision transformer using shifted windows, 2021b. URL https://arxiv.org/abs/2103.14030.
- Effective approaches to attention-based neural machine translation, 2015. URL https://arxiv.org/abs/1508.04025.
- Luna: Linear unified nested attention, 2021. URL https://arxiv.org/abs/2106.01540.
- Mega: Moving average equipped gated attention, 2023.
- Learning word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies - Volume 1, HLT ’11, pp. 142–150, Stroudsburg, PA, USA, 2011. Association for Computational Linguistics. ISBN 978-1-932432-87-9. URL http://dl.acm.org/citation.cfm?id=2002472.2002491.
- Derek Miller. Leveraging bert for extractive text summarization on lectures, 2019. URL https://arxiv.org/abs/1906.04165.
- OpenAI. Gpt-4 technical report, 2023.
- The acl anthology network corpus. Lang. Resour. Eval., 47(4):919–944, dec 2013. ISSN 1574-020X. doi: 10.1007/s10579-012-9211-2. URL https://doi.org/10.1007/s10579-012-9211-2.
- Simplified state space layers for sequence modeling, 2023.
- Utilizing bert for aspect-based sentiment analysis via constructing auxiliary sentence, 2019. URL https://arxiv.org/abs/1903.09588.
- Synthesizer: Rethinking self-attention in transformer models, 2020a. URL https://arxiv.org/abs/2005.00743.
- Sparse sinkhorn attention, 2020b. URL https://arxiv.org/abs/2002.11296.
- Long range arena: A benchmark for efficient transformers, 2020c.
- Mlp-mixer: An all-mlp architecture for vision, 2021.
- Llama: Open and efficient foundation language models, 2023.
- Attention is all you need, 2017.
- Linformer: Self-attention with linear complexity, 2020.
- Paramixer: Parameterizing mixing links in sparse factors works better than dot-product self-attention, 2022.
- Big bird: Transformers for longer sequences, 2021.
- Improving deep neural networks using softplus units. In 2015 International Joint Conference on Neural Networks (IJCNN), pp. 1–4, 2015. doi: 10.1109/IJCNN.2015.7280459.
- Adjorn van Engelenhoven (1 paper)
- Nicola Strisciuglio (30 papers)
- Estefanía Talavera (20 papers)