Simple linear attention language models balance the recall-throughput tradeoff

Simran AroraSabri EyubogluMichael ZhangAman TimalsinaSilas AlbertiJames ZouAtri RudraChristopher Ré

article2024ICML151 citations

statistical language models

Arora et al (2024).pdf

Topics

statistical language models

1. Introduction
2. Preliminaries and Related Work
3. No Free Lunch: Memory-Recall Tradeoff
3.1. Empirical study of memory-recall tradeoff
3.2. Theoretical Analysis
4. The BASED Architecture
4.1. Taylor Linear Attention
4.2. Local Exact Attention with Small Sliding Windows
5. Efficient Implementation
5.1. Preliminaries
5.2. Taylor Exponential Linear Attention
6. Results
6.1. Language Modeling Evaluations
6.2. Efficiency Benchmarks
7. Conclusion
Acknowledgments
Impact Statement
References
A. Extended Related Work
A.1. Efficient Language Modeling Architectures
A.1.1. EFFICIENT ATTENTIONS
A.1.2. ATTENTION ALTERNATIVES
A.2. Efficient Implementations
A.2.1. EFFICIENT ATTENTION IMPLEMENTATIONS
A.2.2. EFFICIENT ATTENTION-ALTERNATIVE IMPLEMENTATIONS
B. IO Aware Implementations
B.1. Forward / Generation Prefill
B.2. Next Token Prediction
B.2.1. TAYLOR LINEAR ATTENTION RECURRENT UPDATE
B.2.2. SLIDING WINDOW ATTENTION
C. Extended Architecture Details
D. Extended Results
D.1. Extended empirical study of memory-recall tradeoff
D.2. Downstream Language Results
D.3. DNA Modeling
D.4. BASED Quality Ablations
E. Experimental Details
E.1. Language Model Pretraining
E.2. Computing Recurrent State Size
E.3. Language Model Evaluation
F. Theoretical Results
F.1. Introduction
F.2. The Models
F.2.1. BASED
F.2.2. MAMBA
F.3. Equivalency to BaseConv
F.4. The Lower Bounds
F.4.1. THE SPACE COMPLEXITY OF AR
F.4.2. LOWER BOUND FOR RECURRENT MODELS
F.4.3. LOWER BOUND ON THE NUMBER OF LAYERS FOR AR
F.5. Lower Bound on the Number of Layers for d ≥ log 2 c with Specific Encodings
F.5.1. THE EQUALITY PROBLEM
F.5.2. THE p-HOT ENCODING FOR p ≥ 1
F.6. Upperbound on MQAR with sub-logarithmically many BaseConv layers
F.6.1. BASECONV PRIMITIVES
F.6.2. PROOF OF THEOREM F.30

Knowls

Knowl 1 — BASED hybrid sequence mixer
model/method
BASED is a hybrid language-model sequence mixer designed to combine long-range memory with precise local token interactions. Its core combines globally applied softmax-approximating Taylor linear attention with causal exact softmax attention restricted to a small sliding window. The models used in the main experiments additionally replace some mixer layers with gated short convolutions: approximately 20% Taylor linear-attention layers, 20% sliding-window-attention layers, and 60% gated-convolution layers. The gated convolutions use filter width 3, SiLU activation, and a projection expansion factor of 4. The main BASED models use no input-dependent decay when trained on 30B or 50B tokens. Varying the Taylor feature dimension and sliding-window width changes the recurrent-state size, allowing BASED to interpolate between low-memory models and high-recall models.
Knowl 2 — Fundamental memory–recall tradeoff
theoretical result
The paper establishes that recall quality is fundamentally constrained by the amount of recurrent state retained during generation. Empirically, on multi-query associative recall, increasing the recurrent state almost always improves accuracy within a given architecture class, while different architectures make substantially different use of the same state budget. Exact softmax attention attains perfect recall but stores a key–value cache whose size grows linearly with sequence length; fixed-state alternatives such as Mamba and H3 use much less memory but have limited recall capacity. In addition, any causal recurrent model with bounded-precision state entries that solves the associative-recall problem on an input of length $N$ requires a recurrent state containing at least $\Omega(N)$ bits. Thus, the observed memory–recall tradeoff is not merely an artifact of a particular architecture.
Knowl 3 — Taylor linear attention and recurrent state
equation
For causal inputs, Taylor linear attention replaces the softmax kernel with a feature-map inner product. Let $q_i,k_j\in\mathbb{R}^{d'}$ be projected query and key vectors, $v_j\in\mathbb{R}^{d}$ be a value vector, and $\phi:\mathbb{R}^{d'}\to\mathbb{R}^{\tilde d}$ be a feature map. The output at position $i$ is
$y_i=\frac{\phi(q_i)\left(\sum_{j=1}^{i}\phi(k_j)^{\top}v_j\right)}{\phi(q_i)\left(\sum_{j=1}^{i}\phi(k_j)^{\top}\right)}.$
The numerator and denominator can be maintained recurrently using a value-weighted state $s_i\in\mathbb{R}^{\tilde d\times d}$ and a key state $z_i\in\mathbb{R}^{\tilde d}$ :
$s_i=s_{i-1}+\phi(k_i)^{\top}v_i,\qquad z_i=z_{i-1}+\phi(k_i)^{\top},\qquad y_i=\frac{\phi(q_i)s_i}{\phi(q_i)z_i}.$
For BASED, the feature map is the second-order Taylor approximation to the exponential kernel. With scaling absorbed into $q$ and $k$ , it satisfies
$\phi(q)^{\top}\phi(k)=1+q^{\top}k+\frac{(q^{\top}k)^2}{2},$
which can be implemented as the concatenation of a constant term, first-order coordinates, and second-order pairwise products. If the projected query and key dimension is $d'$ , the recurrent state size is proportional to $(d+1)\tilde d$ , rather than growing with sequence length. The second-order feature map has naive $O(Nd^3)$ time and space cost when the head dimension and feature dimension are comparable, so BASED uses smaller projected dimensions to control the state size.
Knowl 4 — Tensor-core-aware local exact attention
model/method
BASED uses causal sliding-window attention to provide exact local comparisons and token shifts. For a window width $w$ , query $q_i$ attends only to keys $k_{\max(1,i-w+1)},\ldots,k_i$ , giving $O(Nwd)$ time and space for sequence length $N$ and head dimension $d$ , and a $w$ -token key–value cache during generation. Unlike large-window models using widths such as 4096, BASED chooses small widths that are multiples of 16 and generally no larger than 128. GPU tensor cores operate on $16\times16$ tiles, but the measured latency of $16\times16$ and $64\times64$ matrix multiplications is similar because the smaller operation underutilizes the tensor cores; latency begins increasing substantially beyond width 128. The resulting tensor-core-aware window provides fast exact local attention, while the global Taylor linear-attention component supplies long-range interactions. In the reported recall experiment, small-window exact attention recovered 90.8% of full softmax attention's recall accuracy at approximately $10^{-5}$ of its latency.
Knowl 5 — IO-aware Taylor attention implementation
algorithm
BASED introduces fused GPU kernels for Taylor linear attention during both prompt processing and recurrent generation. The prefill kernel takes projected $q,k,v$ tensors for batch size $B$ , $H$ heads, sequence length $N$ , head dimension $d$ , and expanded feature dimension $D$ . It uses $16\times16$ tiles, eight warps per batch–head block, double-buffered query/key/value tiles in SRAM, and running zeroth-, first-, and second-order Taylor states. For each sequence tile, the kernel computes the causal within-tile term through a masked quadratic attention calculation and the preceding-prefix term through multiplication by the accumulated linear-attention state. The first-order state is accumulated across warps in SRAM; the larger second-order state is partitioned across warp registers and combined before writing the output. The kernel optionally writes the final Taylor states for later generation.

During next-token generation, the recurrent kernel loads the current projected query, key, value, and previous state, updates the state with the outer product $k_t^{\top}v_t$ , and computes the output by multiplying the query by the updated numerator and denominator states. BASED uses eight warps per block and pads the expanded feature dimension to a hardware-friendly multiple; for projected feature dimension 16 and head dimension 64, it uses $D=320$ rather than the unpadded $273$ . The state update requires $O(BHDd)$ arithmetic and moves $O(BHD+BHd)$ bytes from HBM to SRAM per step. The corresponding sliding-window generation kernel uses four warps, fuses query–key multiplication, softmax normalization, and value multiplication in registers, and uses a 64-token window with head dimension 64 in the reported kernel benchmarks.

Relative to the baseline implementation, the fused prefill kernel avoids $O(2BHND)$ bytes of HBM-to-SRAM movement for materialized feature maps and avoids $O(BHNDd)$ bytes of SRAM-to-register movement for repeatedly updating the Taylor state. These reductions come from computing feature maps and causal products inside one kernel and retaining active state fragments in registers.
Knowl 6 — Lower bounds for gated convolutions
theoretical result
The paper analyzes BaseConv, a canonical data-independent gated-convolution layer of the form
$z=(uW_B+B_B)\odot(K*u)+B_K,$
where $u\in\mathbb{R}^{N\times d}$ is a length- $N$ sequence with model dimension $d$ , $W_B$ is a learned projection, $K*u$ is a convolution across sequence positions, $\odot$ is elementwise multiplication, and $B_B,B_K$ are biases. BaseConv can represent a broad class of architectures built from gating and convolution.

For a vocabulary of size $c$ encoded as $d=\log_2(c)$ -bit binary tokens, any data-independent BaseConv model that exactly solves multi-query associative recall requires at least $\log(2d)=\Omega(\log\log c)$ layers. For inputs with vocabulary size $c\le N$ and encoding dimension satisfying $\log c\le d\le 2(\log N)^{1-\epsilon}$ , where $\epsilon>0$ , a data-independent BaseConv model with parameters represented using $O(\log N)$ bits requires $\Omega(\epsilon\log\log N)$ layers even for ordinary associative recall. For $p$ -hot encodings, the lower bound becomes at least $\lfloor\log(2p)\rfloor$ layers when $d=p\sqrt[p]{c}$ . These bounds are complemented by an $O(\max(\log\log c,\log\log N))$ -layer BaseConv construction for certain one-hot associative-recall settings. The paper also shows that linear attention with feature dimension $d'$ can be simulated by BaseConv with $O(\log^2(Nd))$ layers, an inner dimension scaling as $O(N(d+d'^2))$ , and width scaling as $O(\max(d,d'^2))$ .
Knowl 7 — Associative-recall evaluation protocol
experimental setup
The synthetic evaluation uses Multi-Query Associative Recall (MQAR). An input contains key–value pairs followed by queries; for each query, the model must retrieve the value associated with an earlier matching key. For example, a sequence of pairs such as $A\,4$ , $B\,3$ , $C\,6$ followed by queries $A?$ , $C?$ requires the model to output the corresponding values.

Models are trained on sequences of length 256 containing 4–64 key–value pairs and evaluated on sequences of length 1,024 containing 4–256 pairs. The comparison includes softmax attention, sliding-window attention, Mamba, H3, Hyena, and BASED. For every architecture, hyperparameters controlling generation-state size are varied, including model dimension, window width, state dimension, and Taylor feature dimension. State size is measured in bytes during generation; for BASED with model dimension $d$ and Taylor feature dimension $d'$ , the reported state-size expression is
$\operatorname{size}(s_i)=(d+1)\left(1+\frac{3d'}{2}+\frac{d'^2}{2}\right),$
with the additional denominator state accounted for by the $+1$ term. The language-model experiments pretrain 360M- and 1.3B-parameter models from scratch on the Pile using the GPT-2 BPE tokenizer, identical token order, and 10B–50B training tokens. Evaluation covers overall Pile perplexity, associative-recall and non-recall Pile slices, SWDE and FDA information extraction, SQuAD question answering, common-sense LM-Eval tasks, few-shot SuperGLUE, and DNA modeling.

Knowl 8 — Language-model quality and recall results

data/table

The main Pile experiments show that BASED matches strong sub-quadratic baselines in perplexity while substantially improving recall-intensive tasks over Mamba. The following selected rows reproduce the reported 1.3B-parameter results; throughput is in tokens/ms, perplexities are lower-is-better, and accuracies or F1 scores are higher-is-better.

Model and training tokens	Prefill	Generate	Pile all PPL	AR PPL	Other PPL	SWDE Acc.	FDA Acc.	SQuAD F1	LM-Eval avg. Acc.
Transformer++ 1.33B/10B	103.50	0.99	7.26	1.74	8.10	71.92	73.23	36.19	47.64
BASED 1.35B/10B	161.71	24.28	7.43	1.87	8.26	48.06	24.41	30.46	46.68
Mamba 1.32B/10B	112.22	25.69	7.48	1.96	8.29	34.74	12.89	28.20	46.84
Transformer++ 1.33B/50B	103.50	0.99	6.28	1.65	6.82	76.50	80.47	43.47	53.33
BASED 1.35B/50B	161.71	24.28	6.30	1.71	6.82	64.45	30.40	41.62	53.81
Mamba 1.32B/50B	112.22	25.69	6.28	1.74	6.78	52.75	18.51	35.92	53.50

At 1.3B parameters and 50B training tokens, BASED exceeds Mamba by 11.70 SWDE accuracy points, 11.89 FDA accuracy points, and 5.70 SQuAD F1 points, while its LM-Eval average is similar. Averaged over the reported recall-intensive tasks, the paper summarizes the BASED advantage over Mamba as 10.36 accuracy points. BASED remains close to Transformer++ on the associative-recall Pile slice and on SQuAD, although Transformer++ is stronger on the information-extraction tasks.

Knowl 9 — Throughput gains from IO-aware kernels
empirical result
On a single NVIDIA H100 GPU, with timings averaged over 20 repetitions, the IO-aware BASED implementation improves both prompt processing and recurrent generation. At 1.3B parameters and a 4,096-token prefill, BASED is 56% faster than FlashAttention-2 and 44% faster than Mamba. At 360M parameters, it is 28% faster than FlashAttention-2 and 76% faster than Mamba for the same prefill benchmark.

For recurrent generation of 1,024 tokens at batch size 128, the 1.3B-parameter BASED implementation achieves up to 24 times the throughput of FlashAttention-2 and reaches 95% of Mamba's throughput. At 360M parameters, BASED provides 98% higher throughput than FlashAttention-2 and 118% higher throughput than Mamba. The generation benchmarks use CUDA graphs and exclude architectures that lack a recurrent generation implementation. These gains are attributed to fused feature-map computation, register-resident Taylor states, and fused local-window attention rather than to reducing the mathematical content of the attention computation.
Knowl 10 — Ablation and cross-domain findings
empirical result
The ablations show that the quality of BASED depends on both the Taylor feature-map capacity and the local mixer components. In 362M-parameter models trained for 10B Pile tokens, increasing the Taylor feature dimension from 16 to 32 changes overall Pile perplexity from 8.65 to 8.56 and associative-recall perplexity from 2.07 to 2.00; the corresponding SWDE, FDA, and SQuAD scores change from $(29.16,11.71,25.07)$ to $(37.62,12.89,26.74)$ . Replacing the Taylor map with a 16-dimensional Performer produces substantially worse associative-recall perplexity of 8.53, while a larger 128-effective-dimension CosFormer reaches 2.18 associative-recall perplexity. Removing sliding-window attention or short convolutions worsens associative-recall perplexity to 2.11 or 2.09, compared with 2.07 for the full configuration; removing both yields 2.29. Increasing the window from 64 to 128 improves some recall-intensive scores but gives only marginal overall gains. Input-dependent decay provides a small improvement at 10B tokens but is not required for the main trends.

BASED also transfers beyond natural-language modeling. On HG38 DNA modeling, BASED obtains perplexities of 2.51, 2.50, and 2.49 at sequence lengths 1,024, 4,096, and 8,192, respectively. On five GenomicBenchmarks classification datasets, its top-1 accuracies are 74.6, 89.5, 89.5, 96.8, and 79.0 percent. Its DNA results are competitive with Transformer++ and Mamba, indicating that the hybrid recall–throughput design is not restricted to English-language tokenization.

Coverage note — Detailed proof derivations, the full BaseConv construction implementing the one-hot upper bound, exhaustive baseline training configurations, and implementation listings for auxiliary kernels were omitted because they support the stated theorems or results rather than constituting separate load-bearing contributions.

References

1.Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Re. Zoology: Measuring and improving recall in efficient language models. International Conference on Learning Representations, 2023a.
2.Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. volume 30, 2017.
3.Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209.11895, 2022.
4.Junxiong Wang, Jing Nathan Yan, Albert Gu, and Alexander M Rush. Pretraining without attention. arXiv preprint arXiv:2212.10544, 2022.
5.Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023.
6.Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficient training. arXiv preprint arXiv:2312.06635, 2023.
7.Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Re. Hyena hierarchy: Towards larger convolutional language models. arXiv preprint arXiv:2302.10866, 2023.
8.Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, Xuzheng He, Haowen Hou, Przemyslaw Kazienko, Jan Kocon, and Jiaming et al. Kong. Rwkv: Reinventing rnns for the transformer era. arXiv:2305.13048, 2023.
9.Daniel Y. Fu, Tri Dao, Khaled K. Saab, Armin W. Thomas, Atri Rudra, and Christopher Re. Hungry Hungry Hippos: Towards language modeling with state space models. In International Conference on Learning Representations, 2023a.
10.Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint arXiv:2310.06825, 2023.
11.Tri Dao. FlashAttention-2: Faster attention with better parallelism and work partitioning. 2023.
12.Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Re. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In Advances in Neural Information Processing Systems, 2022.
13.Michael Zhang, Kush Bhatia, Hermann Kumbong, and Christopher Re. The hedgehog & the porcupine: Expressive linear attentions with softmax mimicry. In The Twelfth International Conference on Learning Representations, 2024.
14.Feyza Duman Keles, Pruthuvi Mahesakya Wijewardena, and Chinmay Hegde. On the computational complexity of self-attention. In 34th International Conference on Algorithmic Learning Theory, volume 201, page 1–23, 2023.
15.Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, and Shruti Bhosale. Llama 2: Open foundation and fine-tuned chat models. arXiv:2307.09288, 2023.
16.Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac’h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 12 2023.
17.Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In International conference on machine learning, pages 4055–4064. PMLR, 2018.
18.Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509, 2019.
19.Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150, 2020.
20.Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference on machine learning, pages 5156–5165. PMLR, 2020a.
21.Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009.14794, 2020.
22.Krzysztof Choromanski, Haoxian Chen, Han Lin, Yuanzhe Ma, Arijit Sehanobish, Deepali Jain, Michael S Ryoo, Jake Varley, Andy Zeng, Valerii Likhosherstov, et al. Hybrid random features. arXiv preprint arXiv:2110.04367, 2021.
23.Zhen Qin, Weixuan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Lingpeng Kong, and Yiran Zhong. cosformer: Rethinking softmax in attention. arXiv preprint arXiv:2202.08791, 2022a.
24.A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In Proceedings of the International Conference on Machine Learning (ICML), 2020b.
25.Jungo Kasai, Hao Peng, Yizhe Zhang, Dani Yogatama, Gabriel Ilharco, Nikolaos Pappas, Yi Mao, Weizhu Chen, and Noah A. Smith. Finetuning pretrained transformers into RNNs. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 10630–10643, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. doi: 10.18653/v1/2021.emnlp-main.830.
26.Imanol Schlag, Kazuki Irie, and Jurgen Schmidhuber. Linear transformers are secretly fast weight programmers. In International Conference on Machine Learning, pages 9355–9366. PMLR, 2021.
27.Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and et al. Big bird: Transformers for longer sequences. Proceedings of NeurIPS, 2020.
28.Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Re. Scatterbrain: Unifying sparse and low-rank attention approximation. arXiv preprint arXiv:2110.15343, 2021a.
29.Zhanpeng Zeng, Sourav Pak, Jeffrey Kline, Glenn Fung, and Vikas Sing. Multi resolution analysis (mra) for approximate self-attention. Proceedings of the 39 th International Conference on Machine Learning, 2022.
30.Albert Gu, Karan Goel, and Christopher Re. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00396, 2021.
31.Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models, 2023.
32.Ekin Akyurek, Bailin Wang, Yoon Kim, and Jacob Andreas. In-context language learning: Architectures and algorithms. 2024.
33.Daniel Y. Fu, Elliot L. Epstein, Eric Nguyen, Armin W. Thomas, Michael Zhang, Tri Dao, Atri Rudra, and Christopher Re. Simple hardware-efficient long convolutions for sequence modeling. arXiv preprint arXiv:2302.06646, 2023b.
34.Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis-Philippe Morency, and Ruslan Salakhutdinov. Transformer dissection: a unified understanding of transformer’s attention via the lens of kernel. arXiv preprint arXiv:1908.11775, 2019.
35.NVIDIA. Nvidia H100 tensor core GPU architecture, 2022.
36.A. Vyas, A. Katharopoulos, and F. Fleuret. Fast transformers with clustered attention. 2020.
37.Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027, 2020.
38.Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding, 2023.
39.Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019.
40.Eric Wu, Kevin Wu, Roxana Daneshjou, David Ouyang, Daniel Ho, and James Zou. How medical ai devices are evaluated: limitations and recommendations from an analysis of fda approvals. Nature Medicine, 27:1–3, 04 2021.
41.Xiang Deng, Prashant Shiralkar, Colin Lockard, Binxuan Huang, and Huan Sun. Dom-lm: Learning generalizable representations for html documents. 2022.
42.Simran Arora, Brandon Yang, Sabri Eyuboglu, Avanika Narayan, Andrew Hojel, Immanuel Trummer, and Christopher Re. Language models enable simple systems for generating structured views of heterogeneous data lakes. arXiv:2304.09433, 2023b.
43.NVIDIA. Getting started with cuda graphs, 2019. URL https://developer.nvidia.com/blog/cuda-graphs/.
44.Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint arXiv:2108.07258, 2021.
45.Laura Weidinger, John Mellor, Maribeth Rauh, Conor Griffin, Jonathan Uesato, Po-Sen Huang, Myra Cheng, Mia Glaese, Borja Balle, Atoosa Kasirzadeh, et al. Ethical and social risks of harm from language models. arXiv preprint arXiv:2112.04359, 2021.
46.Nikita Kitaev, Łukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001.04451, 2020.
47.Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768, 2020.
48.Chen Zhu, Wei Ping, Chaowei Xiao, Mohammad Shoeybi, Tom Goldstein, Anima Anandkumar, and Bryan Catanzaro. Long-short transformer: Efficient transformers for language and vision. Advances in neural information processing systems, 34:17723–17736, 2021.
49.Silas Alberti, Niclas Dern, Laura Thesing, and Gitta Kutyniok. Sumformer: Universal approximation for efficient transformers. arXiv preprint arXiv:2307.02301, 2023.
50.Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Computing Surveys, 55(6):1–28, 2022.
51.Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nystromformer: A nystrom-based algorithm for approximating self-attention. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 14138–14148, 2021.
52.Yifan Chen, Qi Zeng, Heng Ji, and Yun Yang. Skyformer: Remodel self-attention with gaussian kernel and nystrom method. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021b.
53.Alexandre De Brebisson and Pascal Vincent. An exploration of softmax alternatives belonging to the spherical loss family. arXiv preprint arXiv:1511.05042, 2015.
54.Emmanuel Candes, Xiaodong Li, Yi Ma, and John Wright. Robust principal component analysis? arXiv:0912.3599, 2009.
55.Zhenhai Zhu and Rau Soricut. H-transformer-1d: Fast onedimensional hierarchical attention for sequences. n Annual Meeting of the Association for Computational Linguistics, 2021.
56.Zhen Qin, Xiaodong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong, Nick Barnes, and Yiran Zhong. The devil in linear transformer. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 7025–7041, Abu Dhabi, United Arab Emirates, December 2022b. Association for Computational Linguistics. doi: 10.18653/v1/2022.emnlp-main.473.
57.Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450, 2016.
58.Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019.
59.James W Cooley and John W Tukey. An algorithm for the machine calculation of complex fourier series. Mathematics of computation, 19(90):297–301, 1965.
60.David W. Romero, Anna Kuzina, Erik J. Bekkers, Jakub M. Tomczak, and Mark Hoogendoorn. Ckconv: Continuous kernel convolution for sequential data. 2022.
61.Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are as effective as structured state spaces, 2022.
62.Albert Gu, Ankit Gupta, Karan Goel, and Christopher Re. On the parameterization and initialization of diagonal state space models, 2022.
63.Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces, 2022.
64.Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Zettlemoyer Luke. Mega: Moving average equipped gated attention. arXiv preprint arXiv:2209.10655, 2022.
65.Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, et al. A mathematical framework for transformer circuits. Transformer Circuits Thread, 1, 2021.
66.Yann N Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In International conference on machine learning, pages 933–941. PMLR, 2017.
67.Liliang Ren, Yang Liu, Shuohang Wang, Yichong Xu, Chenguang Zhu, and ChengXiang Zhai. Sparse modular activation for efficient sequence modeling. arXiv preprint arXiv:2306.11197, 2023.
68.Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring attention with blockwise transformers for near-infinite context. arXiv preprint arXiv:2310.01889, 2023.
69.Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Principles, pages 611–626, 2023.
70.Daniel Y Fu, Hermann Kumbong, Eric Nguyen, and Christopher Re. Flashfftconv: Efficient convolutions for long sequences with tensor cores. arXiv preprint arXiv:2311.05908, 2023c.
71.Markus N Rabe and Charles Staats. Self-attention does not need o(n2) memory. arXiv preprint arXiv:2112.05682, 2021.
72.Hanhwi Jang, Joonsung Kim, Jae-Eon Jo, Jaewon Lee, and Jangwoo Kim. Mnnfast: A fast and scalable system architecture for memory-augmented neural networks. In 2019 ACM/IEEE 46th Annual International Symposium on Computer Architecture (ISCA), pages 250–263, 2019.
73.Hao Liu and Pieter Abbeel. Blockwise parallel transformer for long context large models. arXiv preprint arXiv:2305.19370, 2023.
74.Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc Le. Transformer quality in linear time. In International Conference on Machine Learning, pages 9099–9117. PMLR, 2022.
75.Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus), 2023.
76.Alex Wang, Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. SuperGLUE: a stickier benchmark for general-purpose language understanding systems. Curran Associates Inc., Red Hook, NY, USA, 2019.
77.Denis Paperno, German Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernandez. The lambada dataset: Word prediction requiring a broad discourse context, 2016.
78.Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence?, 2019.
79.Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language, 2019.
80.Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge, 2018.
81.Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale, 2019.
82.Eric Nguyen, Michael Poli, Marjan Faizi, Armin Thomas, Callum Birch-Sykes, Michael Wornow, Aman Patel, Clayton Rabideau, Stefano Massaroli, Yoshua Bengio, Stefano Ermon, Stephen A. Baccus, and Chris Re. Hyenadna: Long-range genomic sequence modeling at single nucleotide resolution, 2023.
83.Katarina Gresova, Vlastimil Martinek, David Cechak, Petr Simecek, and Panagiotis Alexiou. Genomic benchmarks: A collection of datasets for genomic sequence classification. bioRxiv, 2022. doi: 10.1101/2022.06.08.495248.
84.Colin Lockard, Prashant Shiralkar, and Xin Luna Dong. OpenCeres: When open information extraction meets the semi-structured web. In Jill Burstein, Christy Doran, and Thamar Solorio, editors, Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pages 3047–3056, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10.18653/v1/N19-1309.
85.Peter Burgisser, Michael Clausen, and Mohammad A Shokrollahi. Algebraic complexity theory, volume 315. Springer Science & Business Media, 2013.
86.Thathachar S Jayram, Ravi Kumar, and Dandapani Sivakumar. The one-way communication complexity of hamming distance. Theory of Computing, 4(1):129–135, 2008.
87.Swastik Kopparty. Topics in algorithms and complexity theory: Spring 2020. 2020.

Citation

MLA

Arora, S., et al. “Simple Linear Attention Language Models Balance the Recall-throughput Tradeoff”. arXiv, 2024, https://doi.org/10.48550/arxiv.2402.18668.

APA

Arora, S., Eyuboglu, S., Zhang, M., Timalsina, A., Alberti, S., Zinsley, D., Zou, J., Rudra, A., & Ré, C. (2024). Simple linear attention language models balance the recall-throughput tradeoff. arXiv. https://doi.org/10.48550/arxiv.2402.18668

Chicago

Arora, S., S. Eyuboglu, M. Zhang, et al. 2024. “Simple Linear Attention Language Models Balance the Recall-throughput Tradeoff”. Preprint, ArXiv. https://doi.org/10.48550/arxiv.2402.18668.

Harvard

Arora, S. et al. (2024) “Simple linear attention language models balance the recall-throughput tradeoff”. arXiv. Available at: https://doi.org/10.48550/arxiv.2402.18668.

Vancouver

1. Arora S, Eyuboglu S, Zhang M, Timalsina A, Alberti S, Zinsley D, Zou J, Rudra A, Ré C (2024) Simple linear attention language models balance the recall-throughput tradeoff. https://doi.org/10.48550/arxiv.2402.18668

BibTeX

@misc{https://doi.org/10.48550/arxiv.2402.18668,
  doi = {10.48550/ARXIV.2402.18668},
  url = {https://arxiv.org/abs/2402.18668},
  author = {Arora, Simran and Eyuboglu, Sabri and Zhang, Michael and Timalsina, Aman and Alberti, Silas and Zinsley, Dylan and Zou, James and Rudra, Atri and Ré, Christopher},
  keywords = {Computation and Language (cs.CL), Machine Learning (cs.LG), FOS: Computer and information sciences, FOS: Computer and information sciences},
  title = {Simple linear attention language models balance the recall-throughput tradeoff},
  publisher = {arXiv},
  year = {2024},
  copyright = {Creative Commons Zero v1.0 Universal}
}

Metadata:DOI registry

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: https://creativecommons.org/licenses/by/4.0/

Simple linear attention language models balance the recall-throughput tradeoff

Topics

Table of Contents

Knowls

Knowl 1 — BASED hybrid sequence mixer

Knowl 2 — Fundamental memory–recall tradeoff

Knowl 3 — Taylor linear attention and recurrent state

Knowl 4 — Tensor-core-aware local exact attention

Knowl 5 — IO-aware Taylor attention implementation

Knowl 6 — Lower bounds for gated convolutions

Knowl 7 — Associative-recall evaluation protocol

Knowl 8 — Language-model quality and recall results

Knowl 9 — Throughput gains from IO-aware kernels

Knowl 10 — Ablation and cross-domain findings

References

Citation

BibTeX

Access the Paper