Combiner: Full Attention Transformer with Sparse Computation Cost

By Hongyu Ren, Hanjun Dai, Zihang Dai, Mengjiao Yang, Jure Leskovec, Dale Schuurmans, Dale Schuurmans, Bo Dai. 2021

99 citations recorded by the source corpus.

Transformers provide a class of expressive architectures that are extremely effective for sequence modeling. However, the key limitation of transformers is their quadratic memory and time complexity $\mathcal{O}(L^2)$ with respect to the sequence length in attention layers, which restricts application in extremely long sequences. Most existing approaches leverage sparsity or low-rank assumptions in the attention matrix to reduce cost, but sacrifice expressiveness. Instead, we propose Combiner, which provides full attention capability in each attention head while maintaining low computation and memory complexity. The key idea is to treat the self-attention mechanism as a conditional expectation over embeddings at each location, and approximate the conditional distribution with a structured factorization. Each location can attend to all other locations, either via direct attention, or through indirect attention to abstractions, which are again conditional expectations of embeddings from corresponding local regions. We show that most sparse attention patterns used in existing sparse transformers are able to inspire the design of such factorization for full attention, resulting in the same sub-quadratic cost ($\mathcal{O}(L\log(L))$ or $\mathcal{O}(L\sqrt{L})$). Combiner is a drop-in replacement for attention layers in existing transformers and can be easily implemented in common frameworks. An experimental evaluation on both autoregressive and bidirectional sequence tasks demonstrates the effectiveness of this approach, yielding state-of-the-art results on several image and text modeling tasks.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Model Architectures

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Model Architectures

Citations in the corpus

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022, background): “One can even combine sparse and low-rank approximation for better accuracy (e.g., Longformer [3], Big-Bird [96], Scatterbrain [9], Long-short transformer [98], Combiner [76]).”
Self-attention Does Not Need $O(n^2)$ Memory (2021, background): “…and architectures with more favorable complexity classes (Kitaev et al., 2020; Roy et al., 2021; Zaheer et al., 2020; Choromanski et al., 2020; Wang et al., 2020; Ren et al., 2021; Child et al., 2019; Tay et al., 2021; Wang et al., 2020; Ma et al., 2021; Shen et al., 2021; Qiu et al., 2020).”
Landmark Attention: Random-Access Infinite Context Length for Transformers (2023, methodology): “Combiner [29] utilizes a hierarchical attention mechanism and heuristic reduction techniques, such as max-pooling, to derive key and query vectors for input blocks.”
General-purpose, long-context autoregressive modeling with Perceiver AR (2022, methodology): “1024 latents 2048 latents 64 batch 128 batch (Roy et al., 2021), and Combiner (Ren et al., 2021), and the diffusion model Variational Diffusion Model (Kingma et al., 2021).”
MEGABYTE: Predicting Million-byte Sequences with Multiscale Transformers (2023, background): “Routing Transformer (Roy et al., 2020) 3.43 Combiner (Ren et al., 2021) 3.42 Perceiver AR (Hawthorne et al., 2022)”
Memorizing Transformers (2022, methodology): “Pooling strategies such as Hierarchical 1D attention (Zhu & Soricut, 2021), and Combiner (Ren et al., 2021) apply pooling or averaging over tokens at longer distances.”
Extending Context Window of Large Language Models via Positional Interpolation (2023, background): “…that focuses on decreasing the memory and computational complexity of the multi-head attention (MHA) mechanism through approximation or sparsiﬁcation (Child et al., 2019; Zaheer et al., 2020; Beltagy et al., 2020; Wang et al., 2020; Choromanski et al., 2021; Kitaev et al., 2020; Ren et al., 2021).”
Transformer Quality in Linear Time (2022, methodology): “Similar to our method, Com-biner (Ren et al., 2021) also splits the sequence into non-overlapping chunks and utilizes quadratic local attention within each chunk.”
FLatten Transformer: Vision Transformer using Focused Linear Attention (2023, background): “One of the possible reasons may give credit to the rank of the attention matrix [36, 53], where a significant difference can be seen.”