# SparQ Attention: Bandwidth-Efficient LLM Inference

Luka Ribar Ivan Chelombiev Luke Hudlass-Galley Charlie Blake Carlo Luschi
Douglas Orr

Graphcore Research

{lukar, ivanc, lukehg, charlieb, carlo<EMAIL_ADDRESS>
###### Abstract

Generative large language models (LLMs) have opened up numerous novel
possibilities, but due to their significant computational requirements their
ubiquitous use remains challenging. Some of the most useful applications
require processing large numbers of samples at a time and using long contexts,
both significantly increasing the memory communication load of the models. We
introduce SparQ Attention, a technique for increasing the inference throughput
of LLMs by reducing the memory bandwidth requirements within the attention
blocks through selective fetching of the cached history. Our proposed
technique can be applied directly to off-the-shelf LLMs during inference,
without requiring any modification to the pre-training setup or additional
fine-tuning. We show how SparQ Attention can decrease the attention memory
bandwidth requirements _up to eight times_ without any loss in accuracy by
evaluating Llama $2$ and Pythia models on a wide range of downstream tasks.

$\bm{q}_{[{\color[rgb]{0.6015625,0.1484375,0.1796875}\definecolor[named]{pgfstrokecolor}{rgb}{0.6015625,0.1484375,0.1796875}\bm{i}_{1}}]}$$\bm{0.8}$$-0.2$$\bm{-1.3}$$0.4$called$\bm{\otimes}$$\bm{K}_{[:,{\color[rgb]{0.6015625,0.1484375,0.1796875}\definecolor[named]{pgfstrokecolor}{rgb}{0.6015625,0.1484375,0.1796875}\bm{i}_{1}}]}$thedograntowardsthetreeastheownercalledapproximateattentionscores
$\bm{\hat{s}}$$\bm{q}$$0.8$$-0.2$$-1.3$$0.4$$\bm{\otimes}$$\bm{K}_{[{\color[rgb]{0.1484375,0.3515625,0.6015625}\definecolor[named]{pgfstrokecolor}{rgb}{0.1484375,0.3515625,0.6015625}\bm{i}_{2}},:]}$sparseattentionscores
$\bm{s}$$\bm{\otimes}$$\bm{V}_{[{\color[rgb]{0.1484375,0.3515625,0.6015625}\definecolor[named]{pgfstrokecolor}{rgb}{0.1484375,0.3515625,0.6015625}\bm{i}_{2}},:]}$$\alpha=\displaystyle\sum_{\mathclap{i\in{\color[rgb]{0.1484375,0.3515625,0.6015625}\definecolor[named]{pgfstrokecolor}{rgb}{0.1484375,0.3515625,0.6015625}\bm{i}_{2}}}}\hat{s}_{i}$$\bm{\bar{v}}=\text{mean}(\bm{V})$$\bm{\oplus}$$\bm{y}$$\times\,(1-\alpha)$$\times\,\alpha$

Figure 1: SparQ Attention for a single attention head. The algorithm consists
of three steps. First, we find the $r$ largest components of the incoming
query vector and gather the corresponding components along the hidden
dimension of the key cache $\bm{K}$. This allows us to approximate the full
attention scores ($\bm{\hat{s}}$). In the second step, we identify the top-$k$
largest scores in the approximation and proceed to gather the corresponding
full key and value vectors from the cache. As a final step, to compensate for
the missing value vectors, we additionally maintain and fetch the running mean
value vector and reassign it the leftover mass based on approximate score
weightings. The attention output is then calculated as usual using the top-$k$
fetched key and value pairs, together with the averaged vector. Figure 2:
Llama $2$ SQuAD $1$-shot performance versus memory transfers over a range of
compression ratios. SparQ Attention achieves matching performance, while
transferring between $1/8$ and $1/4$ as much data as the original dense model.
The curves show the mean $\pm$ one standard error evaluated over $4000$
examples. This pattern is representative of the performance across multiple
models and tasks, shown in Figure A1.

##  1 Introduction

Transformer models trained on large corpora of text have recently shown
remarkable capabilities in solving complex natural language processing tasks
[Kaplan et al., 2020, Brown et al., 2020]. With the dramatic increase in model
scale, LLMs became useful generalist tools that can be used for a multitude of
text-based tasks due to in-context learning capabilities that emerge in LLMs
at scale. These capabilities are unlocked by incorporating information and
instructions through textual prompts [Wei et al., 2022], which are absorbed
during generation as part of the attention cache without modifying the model
weights. Therefore, one of the key obstacles to efficient inference deployment
of LLMs remains their high memory bandwidth requirement when processing long
sequences [Pope et al., 2023], i.e. long instruction prompts, lengthy chat
histories or information retrieval from documents.

The main bottleneck appears due to the auto-regressive nature of transformer
inference: the output is generated token by token, and for each model call,
the full previous state (i.e. the _key-value cache_ , or KV cache) needs to be
fetched from memory. The size of the KV cache scales linearly with the
sequence length as well as the batch size, thus rendering inference using long
sequence lengths increasingly memory-bandwidth limited.111Note that this
problem is less severe during training, since teacher forcing allows whole
sequences to be processed in parallel.

However, tokens generally only attend to a small part of the sequence at a
time [Vig, 2019, Yun et al., 2020]; thus, if it were possible to efficiently
predict the tokens that will have high attention scores, memory traffic could
be significantly reduced by only transferring the key and value pairs of high-
scoring tokens. Building upon this idea, we present SparQ Attention, a
technique for significantly reducing the memory bandwidth requirements of
transformer inference. By approximating attention scores using a subset of
query/key vector components, we fetch only the most relevant tokens for each
inference step, greatly decreasing memory traffic without affecting task
performance.

In order to evaluate the effectiveness of our technique, we curate a selection
of downstream tasks that aim to test the model’s abilities to effectively
utilise the information within the provided textual context. This setup allows
us to evaluate the trade-off between task performance and data transfer
reduction, as well as compare different sparse attention techniques with
respect to memory transfer efficiency. Thus, we clearly show how SparQ
Attention performs favourably compared to state of the art and can lead up to
$8\times$ compression with no loss in accuracy.

##  2 Attention Memory Transfer

Consider a standard attention block and its associated memory requirements. A
single head within an attention layer has a head dimension $d_{h}$, and
processes an input token sequence of sequence length $S$. During token-by-
token inference the output of an attention head is calculated as:

$\bm{y}=\mathrm{softmax}\Bigl{(}\dfrac{\bm{q}\cdot\bm{K}^{\top}}{\sqrt{d_{h}}}\Bigr{)}\cdot\bm{V}$
(1)

where $\bm{q}$ is the query, and $\bm{K}\in\mathbb{R}^{S\times d_{h}}$ and
$\bm{V}\in\mathbb{R}^{S\times d_{h}}$ are the key and value caches
respectively.

For each forward pass, we need to fetch the key and value matrices from
memory, as well as write (append) $\bm{k}$ and $\bm{v}$ vectors for the
current token, giving a total number of elements transferred per attention
head:

$\mathcal{M}_{\mathrm{base}}=\underbrace{2\,S\,d_{h}}_{\text{Read $\bm{K}$ and
$\bm{V}$}}+\underbrace{2\,d_{h}}_{\text{Write current $\bm{k}$ and $\bm{v}$}}$
(2)

Total memory transfers for inference also include model parameters, however
these can be amortised over a large batch of sequences, unlike the $\bm{K}$
and $\bm{V}$ caches. Fetching these from memory dominates processing time for
large $S$, thus creating a bottleneck to efficient transformer inference.

##  3 SparQ Attention

Algorithm 1 SparQ Attention Algorithm

procedure
SparseQAttention($\bm{q}\in\mathbb{R}^{d_{h}},\bm{K}\in\mathbb{R}^{S\times
d_{h}},\bm{V}\in\mathbb{R}^{S\times
d_{h}},\bm{\overline{v}}\in\mathbb{R}^{d_{h}},r\in\mathbb{N},k\in\mathbb{N},l\in\mathbb{N}$)

$\bm{i}_{1}\leftarrow\mathrm{argtopk}\left(|\bm{q}|,r\right)$ $\triangleright$
Indices of top $r$ elements of $\bm{|q|}$

$\bm{\hat{s}}\leftarrow\mathrm{softmax}\left(\bm{q}_{[\bm{i}_{1}]}\cdot\bm{K}_{[:,\bm{i}_{1}]}^{\top}/\sqrt{d_{h}\tfrac{||\bm{q}_{[\bm{i}_{1}]}||_{1}}{||\bm{q}||_{1}}}\right)$
$\triangleright$ Approximate attention scores (all positions)

$\bm{m}\leftarrow[1$ if $i>S-l$ else $0]_{i=1}^{S}$ $\triangleright$ Local
mask of last $l$ positions

$\bm{i}_{2}\leftarrow\mathrm{argtopk}\left(\bm{\hat{s}}+\bm{m},k\right)$
$\triangleright$ Indices of top $k$ approximate scores or local

$\alpha\leftarrow\mathrm{sum}\left(\bm{\hat{s}}_{[\bm{i}_{2}]}\right)$
$\triangleright$ Total approximate score of top $k$

$\bm{s}\leftarrow\mathrm{softmax}\left(\bm{q}\cdot\bm{K}_{[\bm{i}_{2},:]}^{\top}/\sqrt{d_{h}}\right)$
$\triangleright$ Final attention scores (top $k$ positions)

$\bm{y}\leftarrow\alpha\,\bm{s}\cdot\bm{V}_{[\bm{i}_{2},:]}+\left(1-\alpha\right)\,\overline{\bm{v}}$
$\triangleright$ Mixed scores and values, interpolating with
$\overline{\bm{v}}$

return $\bm{y}$

end procedure

Figure 3: *

Algorithm 1: SparQ Attention is, for the most part, a stateless replacement of
the attention operation in a pre-trained transformer. The exception is
$\overline{\bm{v}}=\mathrm{mean}(\bm{V})$, which should be cached and updated
during inference using a running-mean algorithm in order to avoid downloading
the entire $\bm{V}$ matrix at every step. In our experiments, we set
$r\in\left\\{16,32,64\right\\}$ ($r$ is the main determinant of the
compression ratio), $k\in\left\\{64,128\right\\}$ and $l=k/4$. PyTorch code is
provided in Appendix B.

To facilitate an accurate approximation of attention, without transferring the
entire $\bm{K}$ and $\bm{V}$ matrices, we make the following observations:

  1. Observation 1:

The output of the $\mathrm{softmax}$ function in Equation 1 is dominated by a
small number of components, with most of the components being close to $0$.
For example, $32$ from $S\\!>\\!1000$ scores often cover over $90$% of the
total (Figure A2). This suggests that only the $k$ rows of $\bm{V}$
corresponding to the largest attention scores could be fetched from memory
with a minimal degradation in accuracy.

  2. Observation 2:

The indices of the largest attention scores can be efficiently predicted
without fetching the full $\bm{K}$ matrix by sparsifying the query vector
$\bm{q}$, keeping the $r$ largest-magnitude components (see Figure A3).

Using these observations, we propose SparQ Attention (see Figure 1 and
Algorithm 1) consisting of three steps:

  1. Step 1:

Find the indices of $r$ largest components of $|\bm{q}|$ and only fetch
$\bm{K}$ along the dimensions corresponding to these indices. Calculate
approximate attention scores $\bm{\hat{s}}$ using the sliced query and keys.

  2. Step 2:

Find the top-$k$ positions in the approximate attention scores and fetch the
corresponding _full_ key and value vectors. Calculate the output of the
attention block using the top-$k$ keys and values.

  3. Step 3:

Estimate the total score $\alpha$ assigned to the top-$k$ positions using the
approximate attention scores. Use this total score to interpolate between the
attention output from the top-$k$ positions, and a _mean value_ vector,
$\overline{\bm{v}}$.222Further justification and implementation details can be
found in Appendix C.

The memory transfer of the SparQ Attention algorithm for a single attention
head forward-pass:

$\mathcal{M}_{\mathrm{SparQ\
}}=\underbrace{S\,r\,\vphantom{d_{h}}}_{\text{Read $r$ columns of
$\bm{K}$}}+\underbrace{2\,k\,d_{h}}_{\text{Read top-$k$ rows of $\bm{K}$ and
$\bm{V}$}}+\underbrace{4\,d_{h}}_{\text{Write current $\bm{k}$ and $\bm{v}$,
read/write $\overline{\bm{v}}$}}$ (3)

By varying $r$ and $k$ we can tune the total memory bandwidth requirement of
the scheme, trading-off approximation accuracy for inference speed-up. Since
typically $S\gg d_{h}$, $r$ is the most important parameter controlling the
approximation and bandwidth compression. Typical compression ratios are given
in Table A2.

##  4 Experiments

###  4.1 Setup

##### Models

We evaluate our method on two widely-used open-source language model variants:
Llama $2$ [Touvron et al., 2023] and Pythia [Biderman et al., 2023],
evaluating model sizes up to $7$ billion parameters. Both of these model
families are decoder-only transformers [Radford et al., 2018], pre-trained on
causal language modelling. They share similar architectural components such as
Multi-Head Attention and Rotary positional encoding [Su et al., 2021], while
also having some notable differences such as different layer normalisation
implementations, activation functions and execution of modules in parallel.

##### Tasks

In order to evaluate the technique on a spectrum of relevant NLP tasks, our
evaluation setup tests the models on various tasks requiring information
retrieval and reasoning from the provided context: question answering,
summarisation, perplexity/bits-per-character (BPC), text repetition. For all
tasks, we aimed to generate examples of sequence lengths between $1000$ and
$2000$ tokens.333As we wanted to define the tasks independently of the
selected models, our examples were chosen to have context lengths between
$4000$ and $8000$ characters, roughly giving the desired lengths in tokens.

For question answering, we use the SQuAD [Rajpurkar et al., 2016] and TriviaQA
[Joshi et al., 2017] datasets in the open-book setting. In order to construct
the SQuAD examples, we augment the provided context (required to answer the
question) with seven additional confusion contexts from unrelated questions.
This ensures that the examples have the appropriate context length, as well as
making the task harder as the model needs to distinguish the relevant
information from the context from the unrelated paragraphs. We use SQuAD
v$1.1$, as it does not include unanswerable questions included in SQuAD
v$2.0$, since we aim to measure the model’s ability to extract useful
information from the KV cache. For both question answering tasks we use exact
string match accuracy as the evaluation metric.

Summarisation is evaluated on the CNN/DailyMail dataset [See et al., 2017]
using the ROUGE-L F-score [Lin, 2004] as the metric. We use the WikiText-$103$
dataset [Merity et al., 2016] with bits per character (BPC)444We quote
performance for sub-word language modelling in BPC, to account for any
differences in vocabulary across models. for evaluating language modelling
performance.

Finally, we construct an artificial “Text Repetition” task to evaluate the
capability of the model to repeat sentences from its context verbatim. Such a
task can commonly appear in a dialogue setting where the LLM agent is required
to retrieve a piece of text from a possibly long context provided, and can be
challenging for sparse attention techniques. We construct examples using the
Tiny-Shakespeare dataset [Karpathy, 2015] by chunking the text into contexts
of the appropriate size, appending them with the prompts containing a subset
of the context, and evaluating the output exact character length match with
the continuation from the context.

##### Baselines

We consider the cache eviction technique H2O [Zhang et al., 2023] and local
windowing with the initial-tokens enhancement proposed by Han et al. [2023] as
baselines. For consistency, we allocate a fixed KV cache transfer budget $k$
for both techniques. With H2O, we set the local window size $l=k/4$ (with
$3k/4$ heavy hitters), and for local windowing we always include the first
$16$ positions (with $k-16$ local positions).

###  4.2 Results

Dataset Name | SQuAD $\uparrow$ | TriviaQA $\uparrow$ | CNN/DM $\uparrow$ | WikiText $\downarrow$ | Repetition $\uparrow$
---|---|---|---|---|---
Compression |  1 | 1/2 | 1/8 |  1 | 1/2 | 1/8 |  1 | 1/2 | 1/8 |  1 | 1/2 | 1/8 |  1 | 1/2 | 1/8
Pythia $6.9$B |  Local |  57.8 | 38.5 | 14.5 |  52.6 | 41.6 | 23.5 |  20.2 | 14.9 | 13.6 |  0.68 | 0.71 | 0.83 |  150 | 64 | 18
H2O | 52.9 | 45.5 | 52.6 | 52.3 | 20.2 | 17.9 | 0.69 | 0.73 | 47 | 17
SparQ | 58.0 | 57.9 | 52.4 | 51.2 | 20.6 | 20.4 | 0.68 | 0.69 | 151 | 133
Llama$\,2$ $7$B |  Local |  76.9 | 46.1 | 25.5 |  77.6 | 70.1 | 65.2 |  21.4 | 16.5 | 14.1 |  0.77 | 0.81 | 0.91 |  207 | 66 | 10
H2O | 66.2 | 47.1 | 77.2 | 71.3 | 21.8 | 7.4 | 0.77 | 0.88 | 41 | 18
SparQ | 76.4 | 70.2 | 77.3 | 75.3 | 21.7 | 20.5 | 0.77 | 0.79 | 191 | 188

Table 1: Results from the largest models tested are presented above. SQuAD and
TriviaQA measure performance in accuracy. CNN/DailyMail uses ROUGE-L score.
Repetition counts the number of characters before the generated sequence
diverges and WikiText task measures perplexity in bits per character (BPC).
Values in bold represent the best score for a particular model, task and
sparsity setting. Standard errors are: SQuAD $0.8$, TriviaQA $0.9$,
CNN/DailyMail $0.4$, WikiText $0.006$, Repetition $3$.

Our experiments span four distinct models: Llama $2$ with $7$ billion
parameters and three Pythia models with $1.4$, $2.8$ and $6.9$ billion
parameters. Results from the largest models are presented in Table 1, with
further results in Figure A1. We observe:

  * •

SparQ Attention performance is robust across all tasks and model sizes tested.
Compression ratios of $2\times$ to $8\times$ are readily achievable with
little to no loss in task performance.

  * •

The simple recipe of using $k=64$ or $128$, choosing $r$ to set the
compression/performance trade-off seems generally robust. The exception is
Repetition, where $k=64$ performs much worse than $k=128$ for many models.

  * •

Certain tasks are more challenging for H2O (Repetition, SQuAD), while others
are more forgiving (TriviaQA, WikiText-$103$ LM).

  * •

Local attention degrades performance across all tasks, demonstrating that the
tasks do not permit the trivial solution of discarding the long context.

##  5 Related Work

Efficient attention methods have been a very active area of research [Tay et
al., 2020b]. Schemes such as _Sparse Transformers_ [Child et al., 2019],
_Combiner_ [Ren et al., 2021], _Longformer_ [Beltagy et al., 2020], _BigBird_
[Zaheer et al., 2020], _Reformer_ [Kitaev et al., 2020] and _Sparse Sinkhorn
Attention_ [Tay et al., 2020a] have been developed to increase efficiency of
the attention mechanism by extracting information from the most salient tokens
in the sequence or approximating the dense attention maps. These methods form
part of the architecture, they must be implemented at pre-training stage and
carry varying task performance trade-offs.

Another line of work focuses on reducing the memory footprint requirements of
the attention mechanism, while maintaining quadratic complexity. These methods
include popular schemes such as _Multi-Query Attention_ (MQA) [Shazeer, 2019]
and _Grouped-Query Attention_ (GQA) [Ainslie et al., 2023] that share a single
key and value head across multiple query heads. This leads to a reduction in
memory footprint and bandwidth. While these schemes boost efficiency, they
must be implemented at pre-training stage and may affect model quality and
stability.

Alternatively, an emerging area of research similar to SparQ Attention aims to
only adapt the inference procedure of a frozen pre-trained model. The simplest
method of this category is used as part of _FlexGen_ [Sheng et al., 2023]
called _top- $k$ sparse attention_ which calculates exact attention scores and
retrieves only the values associated with the top-$k$ scores. This process
uses the full key cache to produce attention scores, thus limiting the
asymptotic reduction of the memory transfers to only $50\%$.

_Eviction_ schemes describe a number of methods that only cache a subset of
keys and values, and continually delete tokens that are uninformative for
future outputs. By reducing the cache size itself, both memory and memory
bandwidth overheads are reduced. _H 2O_ [Zhang et al., 2023] and
_Scissorhands_ [Liu et al., 2023b] are examples of such eviction methods. H2O
uses a greedy eviction policy that maintains the most salient “Heavy Hitter”
tokens that contribute most to the attention scores on every new generation
step. Similarly, Scissorhands identifies and maintains “pivotal tokens” by
counting how many times a token’s attention score exceeds an importance
threshold over the generative process. While these methods reduce the memory
footprint of the attention cache as well as bandwidth, they also lead to
permanent loss of information from the context window, which can lead to
mistakes for queries seeking less attended parts of the context.

_IceFormer_ [Anonymous, 2023] evaluates a multitude of existing approximate
nearest neighbour algorithms for efficiently approximating attention scores of
pre-trained models, similarly to SparQ Attention. Yet, this work focuses
solely on serving models using CPU due to the implementation of approximate
nearest neighbour, while our scheme is hardware-agnostic and can be used with
any accelerated computing platform. _Scatterbrain_ [Chen et al., 2021] employs
similar techniques, but for computer vision applications.

An orthogonal line of work saves memory traffic by compressing the KV cache
with $4$-bit number formats [Liu et al., 2023a, Sheng et al., 2023]. Liu et
al. [2023b] demonstrate that $4$-bit compression can be complementary to
techniques that reduce the number of transferred elements.

##  6 Discussion

In this work we have presented SparQ Attention, a novel technique for
unlocking faster inference for pre-trained LLMs. Our proposed technique
modifies the multi-head attention mechanism to access only the relevant tokens
from the KV cache on every inference step, leading to drastic memory transfer
savings. This is a significant improvement as it allows for pre-trained models
to be executed more efficiently, without any fine-tuning or modifications to
the weights of the model.

LLM inference speed is often bottlenecked by memory transfers rather than
computation, especially in regimes of long sequence lengths and limited memory
bandwidth. Traffic of the past KV values can dominate the memory bandwidth
when using large batch sizes and sequence lengths longer than the hidden state
of the model, in which case the KV cache reaches a higher memory footprint
than model weights. This mode of execution is particularly suited to the SparQ
technique, with up to eightfold decrease in memory transfer of the KV cache
with little to no accuracy degradation. Our method significantly improves on
existing methods and does not lead to any loss of information from the context
window, as tokens are maintained in memory.

SparQ Attention has some limitations: while maintaining very strong
performance at high bandwidth compression ratios, this is achieved by keeping
all cached values in memory which is sparsely accessed during the generation
steps. It therefore does not save any memory capacity, only bandwidth. Another
possible limitation is the unclear efficiency saving of SparQ Attention when
used with transformer models using MQA and GQA, which were not evaluated in
this work. We leave it as future work to extend SparQ Attention to cover more
attention mechanisms.

##### Acknowledgements

We would like to thank Daniel Justus, Paul Balança and Andrew Fitzgibbon for
their helpful input and feedback on this work.

## References

  * Ainslie et al. [2023] Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai.  GQA: Training generalized multi-query transformer models from multi-head checkpoints.  _arXiv preprint arXiv:2305.13245_ , 2023.
  * Anonymous [2023] Anonymous.  Iceformer: Accelerated inference with long-sequence transformers on CPUs.  In _Submitted to The Twelfth International Conference on Learning Representations_ , 2023.  URL https://openreview.net/forum?id=6RR3wU4mSZ.  under review.
  * Beltagy et al. [2020] Iz Beltagy, Matthew E Peters, and Arman Cohan.  Longformer: The long-document transformer.  _arXiv preprint arXiv:2004.05150_ , 2020.
  * Biderman et al. [2023] Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O’Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al.  Pythia: A suite for analyzing large language models across training and scaling.  In _International Conference on Machine Learning_ , pages 2397–2430. PMLR, 2023.
  * Brown et al. [2020] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.  Language models are few-shot learners.  _Advances in Neural Information Processing Systems_ , 33:1877–1901, 2020.
  * Chen et al. [2021] Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré.  Scatterbrain: Unifying sparse and low-rank attention.  _Advances in Neural Information Processing Systems_ , 34:17413–17426, 2021.
  * Child et al. [2019] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever.  Generating long sequences with sparse transformers, 2019.
  * Han et al. [2023] Chi Han, Qifan Wang, Wenhan Xiong, Yu Chen, Heng Ji, and Sinong Wang.  LM-infinite: Simple on-the-fly length generalization for large language models.  _arXiv preprint arXiv:2308.16137_ , 2023.
  * Joshi et al. [2017] Mandar Joshi, Eunsol Choi, Daniel S Weld, and Luke Zettlemoyer.  TriviaQA: A large scale distantly supervised challenge dataset for reading comprehension.  _arXiv preprint arXiv:1705.03551_ , 2017.
  * Kaplan et al. [2020] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei.  Scaling laws for neural language models.  _arXiv preprint arXiv:2001.08361_ , 2020.
  * Karpathy [2015] Andrej Karpathy.  The unreasonable effectiveness of recurrent neural networks.  https://github.com/karpathy/char-rnn, 2015.
  * Kitaev et al. [2020] Nikita Kitaev, Łukasz Kaiser, and Anselm Levskaya.  Reformer: The efficient transformer.  _arXiv preprint arXiv:2001.04451_ , 2020.
  * Lin [2004] Chin-Yew Lin.  ROUGE: A package for automatic evaluation of summaries.  In _Text Summarization Branches Out_ , pages 74–81, Barcelona, Spain, July 2004. Association for Computational Linguistics.  URL https://aclanthology.org/W04-1013.
  * Liu et al. [2023a] Zechun Liu, Barlas Oguz, Changsheng Zhao, Ernie Chang, Pierre Stock, Yashar Mehdad, Yangyang Shi, Raghuraman Krishnamoorthi, and Vikas Chandra.  Llm-qat: Data-free quantization aware training for large language models, 2023a.
  * Liu et al. [2023b] Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anastasios Kyrillidis, and Anshumali Shrivastava.  Scissorhands: Exploiting the persistence of importance hypothesis for llm kv cache compression at test time.  _arXiv preprint arXiv:2305.17118_ , 2023b.
  * Merity et al. [2016] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher.  Pointer sentinel mixture models.  _arXiv preprint arXiv:1609.07843_ , 2016.
  * Pope et al. [2023] Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean.  Efficiently scaling transformer inference.  _Proceedings of Machine Learning and Systems_ , 5, 2023.
  * Radford et al. [2018] Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever.  Improving language understanding by generative pre-training, 2018.  URL https://openai.com/research/language-unsupervised.
  * Rajpurkar et al. [2016] Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang.  SQuAD: 100,000+ questions for machine comprehension of text.  _arXiv preprint arXiv:1606.05250_ , 2016.
  * Ren et al. [2021] Hongyu Ren, Hanjun Dai, Zihang Dai, Mengjiao Yang, Jure Leskovec, Dale Schuurmans, and Bo Dai.  Combiner: Full attention transformer with sparse computation cost.  _Advances in Neural Information Processing Systems_ , 34:22470–22482, 2021.
  * See et al. [2017] Abigail See, Peter J Liu, and Christopher D Manning.  Get to the point: Summarization with pointer-generator networks.  _arXiv preprint arXiv:1704.04368_ , 2017.
  * Shazeer [2019] Noam Shazeer.  Fast transformer decoding: One write-head is all you need.  _arXiv preprint arXiv:1911.02150_ , 2019.
  * Sheng et al. [2023] Ying Sheng, Lianmin Zheng, Binhang Yuan, Zhuohan Li, Max Ryabinin, Beidi Chen, Percy Liang, Christopher Ré, Ion Stoica, and Ce Zhang.  Flexgen: high-throughput generative inference of large language models with a single GPU.  In _International Conference on Machine Learning_ , pages 31094–31116. PMLR, 2023.
  * Su et al. [2021] Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu.  Roformer: Enhanced transformer with rotary position embedding.  _CoRR_ , abs/2104.09864, 2021.  URL https://arxiv.org/abs/2104.09864.
  * Tay et al. [2020a] Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, and Da-Cheng Juan.  Sparse sinkhorn attention.  In _International Conference on Machine Learning_ , pages 9438–9447. PMLR, 2020a.
  * Tay et al. [2020b] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler.  Efficient transformers: A survey.  _CoRR_ , abs/2009.06732, 2020b.  URL https://arxiv.org/abs/2009.06732.
  * Touvron et al. [2023] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al.  Llama 2: Open foundation and fine-tuned chat models.  _arXiv preprint arXiv:2307.09288_ , 2023.
  * Vig [2019] Jesse Vig.  A multiscale visualization of attention in the transformer model.  In _Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics: System Demonstrations_ , pages 37–42, 01 2019.
  * Wei et al. [2022] Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M. Dai, and Quoc V Le.  Finetuned language models are zero-shot learners.  In _International Conference on Learning Representations_ , 2022.  URL https://openreview.net/forum?id=gEZrGCozdqR.
  * Yun et al. [2020] Chulhee Yun, Yin-Wen Chang, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar.  $o(n)$ connections are expressive enough: Universal approximability of sparse transformers.  _Advances in Neural Information Processing Systems_ , 33:13783–13794, 2020.
  * Zaheer et al. [2020] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al.  Big bird: Transformers for longer sequences.  _Advances in Neural Information Processing Systems_ , 33:17283–17297, 2020.
  * Zhang et al. [2023] Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al.  H2O: Heavy-hitter oracle for efficient generative inference of large language models.  _arXiv preprint arXiv:2306.14048_ , 2023.

Figure A1: Compression versus performance trade-off curves over all models and
tasks evaluated. The y-axis minimum is set to ($0.5$, $0.5$, $0.5$, $1.25$,
$0.0$)$\times$ the dense baseline for the tasks, reading top-to-bottom, in
order to give a consistent view of the performance loss across models.
Vertical dotted lines show ($1/4$) and ($1/8$) compression versus dense.
Shaded lines show $\pm 1$ standard error of the mean.

##  Appendix A Detailed Results

Figure A1 reports the compression/performance trade-off curves for all models
and tasks that were evaluated.

##  Appendix B Code

##  Appendix C SparQ Attention Procedure

The SparQ Attention procedure consists of three steps. In the first step we
find the indices $\bm{i}_{1}\in\mathbb{N}^{r}$ corresponding to the top
$r\in\mathbb{N}$ components of the query vector $\bm{q}\in\mathbb{R}^{d_{h}}$
and proceed to calculate the approximate attention scores
$\bm{\hat{s}}\in\mathbb{R}^{S}$ across the sequence as follows, where
$S\in\mathbb{N}$ is the sequence length:

$\displaystyle\hat{a}_{i}=\sum_{j\in\bm{i}_{1}}q_{j}K_{i,j}$
$\displaystyle\bm{\hat{s}}=\sigma\Bigl{(}\dfrac{\bm{\hat{a}}}{\tau}\Bigr{)}$
(A1)

where $\sigma(\bm{x})_{i}=e^{x_{i}}/\sum_{j}(e^{x_{j}})$ is the
$\mathrm{softmax}$ function and $\bm{K}\in\mathbb{R}^{S\times d_{h}}$ is the
key cache matrix. In the case of standard attention, the scaling factor $\tau$
is chosen to be $\sqrt{d_{h}}$, corresponding to the length of the vector
$\bm{q}$ over which the dot product summation is performed. However, since the
dot product is approximated using only $r$ components here, the scaling factor
needs to be appropriately changed. Empirically we found the appropriate factor
to be:

$\tau=\sqrt{d_{h}\dfrac{\sum_{i\in\bm{i}_{1}}|q_{i}|}{\sum_{i}|q_{i}|}}$ (A2)

which takes into account the relative L1-norm of the selected top $r$
components of the vector $\bm{q}$.

In the second step, we proceed to find indices $\bm{i}_{2}\in\mathbb{N}^{k}$
corresponding to the top $k\in\mathbb{N}$ components of the approximate score
vector $\bm{\hat{s}}$. We then calculate the attention scores using only the
top-$k$ positions, which we can express by defining a corresponding boolean
mask $\bm{b}\in\\{0,1\\}^{S}$:

$\displaystyle b_{i}=\begin{cases}1&\text{if }i\in\bm{i}_{2}\\\
0&\text{else}\end{cases}$
$\displaystyle\bm{s}=\sigma\Bigl{(}\dfrac{(\bm{q}\cdot\bm{K}^{\top})\circ\bm{b}}{\sqrt{d_{h}}}\Bigr{)}$
(A3)

In the third step, we perform a weighted sum of the value vectors in
$V\in\mathbb{R}^{S\times d_{h}}$ using the calculated scores $\bm{s}$. In
order to improve the accuracy of the final approximation, we add an additional
weighted contribution of the mean value vector
$\bm{\bar{v}}=\frac{1}{S}\sum_{i}\bm{V}_{i,*}$:

$\bm{y}=\alpha\sum_{i\in\bm{i}_{2}}s_{i}\bm{V}_{i,*}+(1-\alpha)\bm{\bar{v}}$
(A4)

where $\alpha\in[0,1]$ is the relative weight of the top-$k$ terms. We choose
$\alpha$ using the approximate attention scores $\bm{\hat{s}}$ from the first
step as:

$\alpha=\sum_{i\in\bm{i}_{2}}\hat{s}_{i}$ (A5)

##  Appendix D Attention Sparsity Analysis

In order to understand how to approximate the attention calculation of
Equation 1, we analysed the intermediate scores vector (softmax output). We
took $20$ examples from our SQuAD 1-shot task, and generated the first
completion token using the dense Pythia $1.4$b model, capturing the $\bm{q}$
vector and $\bm{K}$ matrix from every layer and attention head.

To test if the scores are sufficiently sparse to facilitate dropping positions
within $\bm{V}$ entirely, we show the distribution of the sum of the
top-$k=32$ scores, termed coverage, in Figure A2. If the attention scores were
flat, we would expect coverage $\approx k/S\approx 0.03$. Instead, we see most
(example, layer, head) instances show coverage $>0.9$, implying a very sparse
distribution. We note from Figure A2b that coverage does vary per-head and
per-layer, for example earlier layers have “flatter” heads, but even these
heads show coverage $>0.3$.

(a) Coverage over (example, layer, head) (b) Coverage by (layer, head)

Figure A2: Coverage of the softmax output for the top $32$ scores, from a
total of $S\in[1039,1792]$ (varying across examples). Calculated for the first
token generated in answer to a SQuAD prompt, using Pythia-$1.4$B, for $20$
examples.

(a) Attention approximation top-$k$ overlap, over (batch, layer, head) (b)
Attention approximation top-$k$ overlap, by (layer, head)

Figure A3: Proportion of the top-$k$ positions that are correctly predicted by
an approximated softmax, using either the top $16$ absolute components of
$\bm{q}$, or the first $16$ components. Note $d_{h}=128$, so these
approximations have a bandwidth compression ratio of $8\times$. Calculated for
the first token generated to answer a SQuAD prompt, using Pythia-$1.4$B, for
$20$ examples.

It is useful to drop positions in $\bm{V}$ given attention scores, but this
can save at most half of the memory transfers, since the whole of $\bm{K}$ is
needed to calculate these scores. We propose approximating these scores using
a subset of the components of $\bm{K}$. To test such an approximation, we
measure the proportion of overlap between the top $32$ positions in the
approximated and true scores. If overlap is high, we can use the approximation
to avoid downloading the whole $\bm{K}$ matrix, instead only downloading some
components of $\bm{K}$ for all positions, then all components of $\bm{K}$ for
some positions.

Our hypothesis is that the $r$ largest-magnitude components of $\bm{q}$ are
most useful to predicting the score, $\bm{q}\bm{K}^{\top}$. The coverage of
this technique against an arbitrary-component baseline is shown in Figure A3.
These results show that it is possible to achieve reasonably high overlap even
using $r=d_{h}/8$, but that some later layers are harder to predict. Using the
top-$r$ components outperforms the first $r$ baseline considerably.

##  Appendix E Practical Implementation Considerations

One limitation of our theoretical model of memory transfers is that it does
not account for the granularity of memory access. Since the $\bm{K}$ matrix is
indexed on different axes in Step 1 and Step 2 of SparQ Attention, a naive
implementation would load non-contiguous elements in one of the two steps. To
mitigate this, we propose storing $\bm{K}$ twice, once in $S$-major format and
once in $d_{h}$-major format. This increases KV cache memory usage by $50\%$,
but uses only a small amount of extra bandwidth to write $\bm{k}$ twice. This
extra write is non-contiguous, but small, so should not form a bottleneck.

##  Appendix F Methodology

We provide a comprehensive set of hyperparameters for reference in Table A1.
Typical compression ratios for settings of $(r,k)$ are given in Table A2.

We use our own implementation of H2O [Zhang et al., 2023], which differs from
the author’s implementation in that it uses a fixed cache size $k$, rather
than a ratio of the current sequence length. To validate that these
implementations are sufficiently similar, we ran their implementation through
our harness on a small model and sample size. On SQuAD $1$-shot, with
Pythia-$1.4$B, using $k=256$, $l=64$, our implementation was correct for $60$
of $200$ examples, theirs for $57$ (the dense baseline achieved $74$). Perhaps
more importantly, we found that of the $79$ times that either output differed
from dense, $41$ occurrences showed a $20$-character prefix match between our
implementation and theirs. The fact that the two implementations often
generate the same errors (despite minor implementation differences) reassures
us that our results should be a fair representation of $\mathrm{H_{2}O}$.

Group | Hyperparameter |  Value or range
---|---|---
Dense model | Family |  Llama $2$ ($7$B), Pythia ($6.9$B, $2.8$B, $1.4$B)
$d_{h}$ |  $\\{80,128\\}$
Max $S$ |  $\\{2048,4096\\}$
Tasks | QA (#samples) |  SQuAD $1$-shot ($4000$)
TriviaQA $0$-shot ($2992$)
Summarisation (#samples) |  CNN/DailyMail $0$-shot ($500$)
Language Modelling (#samples) |  WikiText-$103$ LM ($500$)
Artificial (#samples) |  Repetition ($1000$)
Baselines | Eviction |  keep $(k-l)$ tokens with highest $\mathrm{score}(n)=\sum_{i}s_{in}$
and the most recent $l=k/4$
$k\in\\{128,192,256,384,512,768\\}$
Local |  take the first $16$ tokens, and most recent $k-16$
$k\in\\{128,192,256,384,512,768\\}$
SparQ | Rank $r$ |  $\\{16,32,64\\}$
Number of values $k$ |  $\\{64,128\\}$
Local window $l$ |  $k/4$
Table A1: Experiment hyperparameters Method | $r$ | $k$ | Compression
---|---|---|---
SparQ Attention | 16 | 64 | 0.10-0.12
128 | 0.13-0.17
32 | 64 | 0.16-0.18
128 | 0.19-0.23
64 | 64 | 0.28-0.30
128 | 0.31-0.36
H2O | – | 128 | 0.07-0.11
192 | 0.10-0.17
256 | 0.13-0.22
384 | 0.20-0.33
512 | 0.26-0.43
768 | 0.39-0.65
Table A2: Range of compression ratios for different settings of $(r,k)$, for
Llama 2 7B and Pythia 6.9B. The compression ratio achieved varies across
models and tasks, based on the sequence length and head size.

###  F.1 Examples

We illustrate the task setup with a single example per task, showing the
prompt formatting and a cherry-picked example. In each case, we show outputs
from a dense Llama 2 model, SparQ ($r=16,k=64$), Eviction and Local ($k=192$)
Attention. Where “…” appears, we have truncated the line of text for brevity.

####  F.1.1 Question Answering (SQuAD 1-shot)

    ### PROMPT (56de0ed14396321400ee257b)
    Title: Private school. Background: In many parts o...
    Title: University of Chicago. Background: The Univ...
    Title: Nikola Tesla. Background: Near the end of h...
    Title: Ctenophora. Background: Since all modern ct...
    Title: Chloroplast. Background: Lepidodinium virid...
    Title: Construction. Background: Construction proj...
    Title: Normans. Background: Institutionally, the N...
    Title: Genghis Khan. Background: At this time, non...
    From what you’ve just read about Normans, please answer the following questions.
    Question: What does Kitab Rudjdjar mean in English?
    Answer: The Book of Roger
    Question: What kind of bureaucracy did the Normans institute?
    Answer:

    ### OUTPUT
    DENSE: Meritocratic
    SPARQ: Meritocratic
      H2O: Mer Mer Mer Mer Mer Mer Mer Mer Mer
    LOCAL: Feudalism

####  F.1.2 Question Answering (TriviaQA 0-shot)

    ### PROMPT (dpql_5685)
    Apritifs and digestifs ( and) are drinks, typical...
    Apritifs
    An apritif is an alcoholic beverage usually serve...
    "Apritif" may also refer to a snack that precedes...
    "Apritif" is a French word derived from the Latin...
    ...
    ...
    * Distilled liquors (ouzo, tequila, whisky or akva...
    * Liquor cocktails (Black Russian, Rusty Nail, etc...
    In certain areas, it is not uncommon for a digesti...
    Bitter digestifs typically contain carminative her...
    In many countries, people drink alcoholic beverage...
    Question: Which aperitif is named for the Paris chemist who created it in 1846?
    Answer:

    ### OUTPUT
    DENSE: Dubonnet
    SPARQ: Dubonnet
      H2O: Dubonnet
    LOCAL: Vermouth.

Note that for Pythia, the prompt “Single-word answer:” was used in place of
“Answer:”, as this helped prevent the model from restating the question in the
answer (often qualitatively correct, but not a regex match).

####  F.1.3 Summarisation (CNN/DailyMail)

    ### PROMPT (417e9c57edd888a811e860429219b3a66f16d92c)
    Article: Being at Old Trafford on Saturday was a surreal experience  and Im no...
    Summary:

    ### OUTPUT
    DENSE: Manchester United fans have been left frustrated by the club’s lack of attacking...
    SPARQ: Manchester United fans have been left frustrated by the team’s performances this...
      H2O: Manchester United fans are not happy with the team’s performance this season. Th...
    LOCAL: The rebuild is underway. But its unfair to expect ’Fergie football’ so soon.

####  F.1.4 Repetition (Shakespeare)

    ### PROMPT (5049)
    First Citizen:
    Before we proceed any further, hear me speak.
    ...
    ...
    Our steed the leg, the tongue our trumpeter.
    With other muniments and petty helps
    In this our fabric, if that they--

    MENENIUS:
    What then?
    ’Fore me, this fellow speaks! What then? what then?

    First Citizen:
    Should by the cormorant belly be restrain’d,
    ...
    ...
     the tongue our trumpeter.
    With other muniments and petty helps
    In this our fabric, if that they--

    MENENIUS:
    What then?
    ’Fore me,

    ### OUTPUT
    DENSE: this fellow speaks! What then? what then?
    SPARQ: this fellow speaks! What then? what then?
      H2O: this is most
    LOCAL: I cannot tell.

####  F.1.5 Language Modelling (WikiText-103)

    ### QUERY (0)

     = County Road 595 ( Marquette County , Michigan ) =

     County Road 595 ( CR 595 , Co . Rd . 595 ) is a proposed primary county road in...
     In 2003 , a flood along the Dead River destroyed or forced the closure of sever...
     The Marquette County Road Commission ( MCRC ) applied for permits from the stat...

     = = Route description = =

     The county road would start at its southern end at an intersection with US 41 /...

     = = History = =

     = = = Background = = =

     On May 14 , 2003 , a section of the Silver Lake Dam failed in northwestern Marq...

Since the flood a new , roughly 100 @-@ foot @-@ tall ( 30 m ) bridge has been
constructed on CR 510 , ensuring that any future flooding on the Dead River
would not necessitate closure of the modern bridge . This structure was built
at that height in order to benefit commercial interests in the area , as well
as to provide the area with reliable public and emergency access . Rio Tinto
Group , the British @-@ based parent company of Kennecott Minerals received
permission

    ### BPC
    DENSE: 0.788
    SPARQ: 0.790
      H2O: 0.859
    LOCAL: 0.904