$\dagger$$\dagger$affiliationtext: Stanford
University$\ddagger$$\ddagger$affiliationtext: University at
Buffalo$\triangle$$\triangle$affiliationtext: Purdue
University$\dagger$$\dagger$affiliationtext:
<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>
#  Simple linear attention language models balance the recall-throughput
tradeoff

Simran Arora Corresponding authors; equal contribution and random ordering for
SA, SE, MZ (SSM). Sabri Eyuboglu††footnotemark:  Michael Zhang††footnotemark:
Aman Timalsina  Silas Alberti  Dylan Zinsley  James Zou  Atri Rudra
Christopher Ré

###### Abstract

Recent work has shown that attention-based language models excel at recall,
the ability to ground generations in tokens previously seen in context.
However, the efficiency of attention-based models is bottle-necked during
inference by the KV-cache’s aggressive memory consumption. In this work, we
explore whether we can improve language model efficiency (e.g. by reducing
memory consumption) without compromising on recall. By applying experiments
and theory to a broad set of architectures, we identify a key tradeoff between
a model’s state size and recall ability. We show that efficient alternatives
to attention (e.g. H3, Mamba, RWKV) maintain a fixed-size recurrent state, but
struggle at recall. We propose Based a simple architecture combining linear
and sliding window attention. By varying Based window size and linear
attention feature dimension, we can dial the state size and traverse the
pareto frontier of the recall-memory tradeoff curve, recovering the full
quality of attention on one end and the small state size of attention-
alternatives on the other. We train language models up to $1.3$b parameters
and show that Based matches the strongest sub-quadratic models (e.g. Mamba) in
perplexity and outperforms them on real-world recall-intensive tasks by 6.22
accuracy points. Implementations of linear attention are often less efficient
than optimized standard attention implementations. To make Based competitive,
we develop IO-aware algorithms that enable $24\times$ higher throughput on
language generation than FlashAttention-2, when generating 1024 tokens using
$1.3$b parameter models.

##  1 Introduction

The choice of sequence mixer (e.g. attention, convolution) in a language model
affects both its quality and efficiency [1, 2]. Prior work shows that
attention excels at recall, the ability to ground generations in previously
seen tokens [3, 1]. On the other hand, the throughput of attention-based
models is bottle-necked during training by quadratic compute complexity and
during inference by aggressive memory consumption. The natural question is:
can we improve the real-world speed and memory-use of language models without
comprising on quality?

Figure 1: Based overview. Combining linear attention with tiny sliding window
softmax attention (e.g., 64 or 128 tokens in width) enables improved recall
accuracy with limited efficiency overhead vs. smaller tile sizes. (Left) Time
to execute Cutlass GEMMs ($y$) vs. sliding window attention size ($x$), with
batch size $512$ on tensor cores. (Center) Model recall accuracy ($y$) vs.
sliding window attention size ($x$). We compare linear attention alone (dark
blue), sliding window attention alone (light blue), and their combination
(Based, orange). (Right) Schematic diagram of Based illustrating how the two
components complement each other.

Recently, a number of architectures have been proposed that enable
substantially higher throughput while matching attention in perplexity [4, 5,
6, 7, 8]. However, coarse metrics like overall perplexity can obscure
important differences in model quality. For example, recent work shows that a
specific class of architectures, gated-convolutions, despite complexity
scaling sub-quadratically in sequence length, are less efficient than
attention at performing recall [1]. Building on this analysis, we evaluate a
broader class of architectures across real-world recall-intensive tasks and
show attention improves over the best attention-free alternative, Mamba, by
46.7 accuracy points (Table 1). 111Examples of recall-intensive tasks include
information extraction, reading comprehension, and code generation using
custom variable and function names. These require using in context
(contrasting memorized) information during generation.

Motivated by these observations, we explore the pareto frontier of the
tradeoff between high-recall and high-throughput models. We evaluate a range
of architectures (e.g. attention, SSMs, and convolutions) on a popular
synthetic associative recall task [1, 9, 3]. Since generation throughput is
bottle-necked by memory consumption, we vary hyperparameters (e.g. model
dimension) that affect the size of the recurrent state during generation and
demonstrate a fundamental recall-memory tradeoff that holds across
architecture classes (Figure 2). Attention performs associative recall
perfectly, but the recurrent state (i.e. the KV-cache) grows linearly with the
sequence length. Sliding window attention can cap the size of the recurrent
state at the cost of worse long-range recall [10]. However, Mamba, a recently
proposed SSM architecture expands the Pareto frontier beyond sliding window.
This begs the question: are there other, perhaps simpler, models that can also
expand the pareto frontier?

To reduce the memory consumption, we consider using two simple techniques:
sliding window attention and softmax-approximating linear attention. Our
results on language modeling (Table 1) and synthetic recall experiments
(Figure 1, center) suggest neither primitive alone suffices to navigate the
Pareto frontier.

  1. 1.

We find that linear attention alone struggles to solve associative recall
(Figure 1, center). We hypothesize that this is because linear attention lacks
the precision to perform local token shifts and comparisons [9, 1].

  2. 2.

In sliding window attention, associative recall range is limited by the width
of the windows (Figure 1, center). As we increase the window size, the
recurrent state grows linearly and has a non-linear affect on speed during
parallel training and inference (Figure 1, left).

We combine these two techniques into a single architecture, which we call
Based (Figure 1, right). We find that sliding window attention and linear
attention complement each other, enabling Based to expand the pareto frontier
of the recall-memory tradeoff (Figure 2). We suspect that (1) the large
recurrent memory of linear attention could help model long-range token
interactions in the sequence and (2) sliding window attention handles the
precise local shifts needed to perform associative recall.

To make Based competitive with SoTA attention [11] and recurrent [5] models
under wall-clock and throughput metrics, we introduce several IO-aware
optimizations.

  1. 1.

Despite the theoretically improved complexity, linear attention
implementations are often slower than well-optimized attention implementations
[12]. In Based, we use the 2nd-order Taylor approximation of softmax as the
linear attention feature map With sequence length $N$ and head dimension $d$,
this naïvely requires $\mathcal{O}(Nd^{3})$ time and space complexity [13,
14]. To make our attention competitive in real-world wall-clock time and
memory usage, we provide hardware-efficient algorithms and custom CUDA
implementations. Relative to the baseline, our algorithm reduces data movement
from HBM (slower-to-acccess memory) to SRAM (faster-to-access memory) by
$\mathcal{O}(Nd^{2})$ bytes and from SRAM to register (fastest memory) by
$O(Nd^{3})$ bytes (Section 5).

  2. 2.

Sliding window attention exploits tensor cores, specialized units on modern
GPUs for performing matrix multiplications (GEMMs). While popular
architectures use long window sizes (e.g. 4096 for Mistral-7B [10]), we choose
fixed size $64$ windows, guided by hardware properties. In particular, we use
just enough occupancy to hide the tensor core kernel launch latency. Although
tensor cores operate on $16\times 16$ tiles, in Figure 1 (left), we see that
the latencies for performing $16\times 16$ vs. $64\times 64$ (or even
$128\times 128$) dimension matrix multiplications on NVIDIA H100 tensor cores
are similar, informing our window size.

In experiments, we show that Based competes in quality with strong
Transformer++ [15] and SoTA sub-quadratic baselines in models up to the 1.3Bn
parameters across language modeling on the Pile language, DNA modeling, and
the LM Eval Harness [16]. Beyond this, Based outperforms prior sub-quadratic
architectures on the associative recall slice of the Pile and in downstream
recall-intensive tasks by $0.14$ perplexity points and $6.22$ accuracy points,
respectively. In efficiency, Based enables up to $24\times$ higher throughput
than the strong FlashAttention-2 implementation on generation. Code for this
work is provided at: https://github.com/HazyResearch/based.

##  2 Preliminaries and Related Work

We discuss the key relevant work in this section and provide an extended
discussion in Appendix A.

##### Attention

Popularized by Transformers [2] as the de facto language modeling primitive,
softmax attention takes inputs $\bm{x}\in\mathbb{R}^{N\times d}$ of length $N$
and head dimension $d$, and computes outputs $\bm{y}\in\mathbb{R}^{N\times d}$
via the softmax over projections
$\bm{q},\bm{k},\bm{v}=\bm{x}\bm{W}_{q},\bm{x}\bm{W}_{k},\bm{x}\bm{W}_{v}$,
i.e.,

$\bm{y}_{i}=\sum_{j=1}^{i}\frac{\exp(\bm{q}_{i}^{\top}\bm{k}_{j}/\sqrt{d})\bm{v}_{j}}{\sum_{m=1}^{i}\exp(\bm{q}_{i}^{\top}\bm{k}_{m}/\sqrt{d})}$
(1)

in the causal case where
$\bm{W}_{q},\bm{W}_{k},\bm{W}_{v}\in\mathbb{R}^{d\times d}$ are learnable
matrices . While effective at recall [1] and efficient to train (Eq 1 is
parallelizable on GPUs and $\mathcal{O}(N)$ in memory with recent advances
[12]), attention remains expensive for generation. For every new output
$\bm{y}_{n}$, we require $nd$ operations over a growing KV-cache of prior
$\\{\bm{k}_{i},\bm{v}_{i}\\}_{i=1}^{n-1}$. This results in larger memory
consumption and lower-throughput for longer sequences.

##### Efficient attentions

Various works thus try to improve on attention’s efficiency without
sacrificing quality. _Sparse attentions_ reduce attention’s time and memory
requirements by only attending over specific strided patterns or local
_sliding windows_ [17, 18, 19]. While further popularized in large language
models (Mistral, Jiang et al. [10]), prior works either underperform full
attention with sparse patterns that fail to capture dense interactions, or use
large window sizes that still permit large KV-caches and subsequent
inefficiency.

Meanwhile, _linear attentions_ replace the softmax in standard attention with
alternative kernel functions [20, 21, 22, 23, 14]. By removing the
$\exp(\bm{q}^{\top}\bm{k})$ in favor of feature map dot-products
$\phi(\bm{q})^{\top}\phi(\bm{k})$, these methods use matrix product
associativity to compute attention in $\mathcal{O}(Nd^{2})$ time and space
[24]. Furthermore, they permit a _recurrent view_ for constant memory and
$\mathcal{O}(1)$ time per-token generation [25, 26]. However, present linear
attention feature maps either fail to match standard attention on recall or
remain expensive to compute [13]. Linear attentions also do not achieve faster
wall-clock time or lower memory versus modern standard attention in practice
[12].

##### Attention alternatives

Finally, various models use attention-free sequence mixers such as state-space
models (SSMs) [27, 28], gated convolutions [9, 7] and input-dependent
recurrences [8, 5] to rival attention performance while improving its
efficiency. However, while recent such models can match attention in overall
perplexity, further study suggests they may underperform Transformers on tasks
such as recall and in-context learning [1, 29].

##  3 No Free Lunch: The Memory-Recall Tradeoff

In this section, we demonstrate a fundamental tradeoff between a model’s
memory consumption during inference (i.e., the size of its recurrent state)
and its capacity to perform recall. We use a combination of experiments on
synthetic data and theoretical analysis.

  * •

Empirical study of memory-recall tradeoff : In Section 3.1, we evaluate a
number of popular architecture classes (e.g. Mamba, Hyena) on a synthetic
associative recall task, varying hyperparameters that affect the model’s
recurrent state size (Figure 2). Within each architecture class, we observe a
clear tradeoff: the larger the recurrent state size, the better recall.
However, for a fixed recurrent state size, performance is not consistent
across architectures. We observe that some sequence mixers fall well-below the
pareto-frontier. This motivates the design of sequence mixers that can expand
the pareto frontier.

  * •

Lower bounds on memory required for recall: In Section 3.2, we lower bound the
recurrent state size required to perform exact recall with any recurrent model
Theorem F.1. This analysis reinforces our empirical observations on the
throughput-recall tradeoff.

Figure 2:  Throughput (memory) - recall tradeoff. $x$-axis shows state size
(bytes) during generation; $y$-axis shows accuracy on the MQAR recall task
[1]. For each architecture, we train several models varying hyperparameters
that affect the recurrent state size (e.g. model dimension). The plot shows a
fundamental tradeoff between the recurrent state size and recall capacity that
applies to broad class of models [1, 5, 9].

###  3.1 Empirical study of memory-recall tradeoff

Setup. We use a synthetic AR task called Multi-Query Associative Recall
(${\mathrm{MQAR}}$) [1] to demonstrate the trade-off. In this task, input
sequences consist of a number of key-value pairs followed by queries. For a
given query, the model must recall the corresponding key-value pair from
earlier in the sequence in order to predict the next token. For example, the
correct output for input below would be 4, 6, 1, 2, 3:

$\text{A 4 B 3 C 6}\underbrace{\text{F 1}}_{\mathclap{\textbf{Key-
Value}}}\text{E 2}\rightarrow\text{A ? C ?}\underbrace{\text{F
?}}_{\mathclap{\textbf{Query}}}\text{E ? B ?}$

We train on sequences of length 256 tokens containing between 4 and 64 key-
value pairs. During evaluation, we measure accuracy on sequences of length
1,024 tokens containing between 4 and 256 key-value pairs.

We train and evaluate six sequence mixers: attention [2], sliding window
attention [19], Mamba [5], H3 [9], Hyena [7], and Based. For each, we vary
hyperparameters that affect the memory consumption during inference. For
example, in sliding window attention we vary the window width. We compare how
MQAR accuracy varies with the size of the recurrent state. Section E.1
contains details on how the state size is calculated for each architecture.

Figures 2 and 3 can be reproduced or extended to new architectures using the
scripts provided at https://github.com/HazyResearch/zoology.

##### Results

In Figure 2, we demonstrate a fundamental tradeoff between recurrent state
size and accuracy on ${\mathrm{MQAR}}$ that holds within and across
architecture classes. Within each architecture class (e.g. H3 models),
increasing the recurrent state size almost always leads to an improvement in
accuracy. Across architecture classes, we see a tradeoff as well. Attention
achieves perfect recall accuracy, but its recurrent state size grows with the
length of the sequence. Other architecture classes like Mamba and H3 admit
models with much smaller recurrent states, but these models have limited
recall capacity.

Given a fixed recurrent state, not all architectures have the same recall
capacity. Among architectures proposed in prior work, Mamba makes the best use
of a limited memory budget. Notably, architectures with a convolutional view
(e.g. Hyena and H3) fall well below the pareto frontier. Our proposed
architecture, Based (introduced in Section 4), expands the pareto-frontier
beyond Mamba. By varying hyper-parameters that determine its state size (e.g.
feature dimension and model dimension), we can smoothly navigate the tradeoff
between efficient models and memory-hungry models with high recall capacity.

###  3.2 Theoretical Analysis

Our theoretical analysis provides further insight into the empirical
observations described above. First, using results from communication
complexity theory, we show that the recall capacity of any causal model (e.g.
Mamba, Attention) is bounded by the size of its recurrent state (Theorem F.3
in Appendix F).

######  Theorem 3.1.

Any recurrent model222In particular, for Mamba [5], see Corollary F.1.
depending causally on input ${\bm{u}}\in\\{0,1\\}^{N\times d}$ requires
$\Omega(N)$-bits333Here, we need the entries of the state to be bounded. in
state size to solve $\mathrm{MQAR}$.

This result suggests that the tradeoff observed in Figure 2 is fundamental,
not an artifact of architectural quirks.

Next, we focus on gated-convolutions, a broad class of architectures built
from gating and convolutions (e.g. H3, Hyena, RWKV v4). To make progress in
theoretically analyzing the broad set of gated convolution proposals, prior
work develops a canonical gated-convolution, referred to as BaseConv which can
provably simulate any architecture built from gating and convolution
primitives.

Building on this work, we show that BaseConv cannot solve $\mathrm{MQAR}$ in
constant-many layers (Theorem F.5 and Theorem F.6 in Appendix F).

######  Theorem 3.2.

Given an input sequence ${\bm{u}}\in\\{0,1\\}^{3N\times d}$, where $N$ and $d$
denote the sequence length and head dimension, respectively, a data-
independent BaseConv model needs $\log(2d)$-layers to solve $\mathrm{MQAR}$
for $d=\log_{2}(c)$, where $c$ denotes the vocabulary size444That is, each
token from the vocabulary has the natural binary encoding in
$\\{0,1\\}^{\log_{2}(c)}$.

In contrast, Arora et al. [1] show that attention solves ${\mathrm{MQAR}}$ in
constant-many layers. This result helps to explain why the gated-convolution
architectures (H3 and Hyena) in Figure 2 lie below the pareto frontier
established by newer architectures.

Architecture | Params | Efficiency | Language Modeling (Pile) | Info. Extraction | Question Answering
---|---|---|---|---|---
Prefill | Generate | All | AR | Other | SWDE | FDA | SQUAD | Common
Tok./ms $\uparrow$ |  Tok./ms $\uparrow$ |  Ppl. $\downarrow$ |  Ppl. $\downarrow$ |  Ppl. $\downarrow$ |  Acc $\uparrow$ |  Acc $\uparrow$ |  F1 $\uparrow$ |  Avg. Acc. $\uparrow$
Transformer++ | 1.33b | 103.50 | 0.99 | 7.26 | 1.74 | 8.10 | 41.97 | 73.23 | 36.19 | 47.64
Based | 1.35b | 161.71 | 24.28 | 7.43 | 1.87 | 8.26 | 30.83 | 24.41 | 30.46 | 46.68
Mamba | 1.32b | 112.22 | 25.69 | 7.48 | 1.96 | 8.29 | 25.93 | 12.89 | 28.20 | 46.84
Transformer++ | 360m | 207.77 | 23.82 | 8.39 | 1.87 | 9.42 | 37.62 | 58.00 | 27.18 | 44.08
Based | 363m | 514.57 | 47.23 | 8.65 | 2.07 | 9.64 | 22.81 | 11.71 | 25.07 | 43.03
Mamba | 358m | 267.09 | 39.95 | 8.64 | 2.21 | 9.59 | 25.61 | 6.53 | 24.06 | 43.51
GLA | 362m | — | — | 9.12 | 2.36 | 10.68 | — | — | — | —
RWKV v5 | 362m | — | — | 9.79 | 2.40 | 10.90 | — | — | — | —
H3 | 362m | — | — | 10.60 | 4.88 | 11.23 | 17.59 | 0.64 | 7.87 | 39.35
Table 1: Evaluation of pre-trained language models. All models were trained on
the same set of 10 billion tokens drawn from the Pile [30]. We report
inference throughput on $4,096$ tokens ($16,384$ for 360m param.) of pre-fill
and $2,048$ tokens of recurrent generation for a subset of architectures. We
report language model perplexity on the overall Pile test set as well as
perplexity on two slices of the test set: associative recall tokens and other
tokens (see Section 6.1, [1]). We report zero-shot performance on three
recall-intensive tasks: information retrieval on SWDE and FDA as well as
question answering on SQUAD. Finally, we report average performance on the set
of LM Eval Harness [16] common sense reasoning tasks used in Gu and Dao [5],
details in Table 5. These tasks do not require significant recall capacity
because the input text is typically very short. See Section 6.1. Some proposed
architectures that do not implement recurrent views for generation are marked
with a —.

######  Remark 3.1.

For a class of input encodings that generalizes one-hot encodings termed as
$p$-hot encodings (Definition F.7), input-dependent BaseConv needs at least
$\lfloor\log(2p)\rfloor$-layers to solve ${\mathrm{MQAR}}$ where
$d=p\cdot\sqrt[p]{c}$.

Finally, we show that we can simulate linear attention [20], the foundation of
Based, using BaseConv [1] with a poly-log blowup in the number of layers
(Proposition F.1 in Appendix F), pointing to the relative efficiency of linear
attention over gated-convolution architectures.

##  4 The Based Architecture

In this section, we introduce Based. Our objective in designing this
architecture is to demonstrate how we can navigate the pareto-frontier of the
memory-recall tradeoff using well-known architectural building blocks.

Softmax attention excels at recall, but since its recurrent state, the KV-
cache, grows unconstrained with the length of sequence, it is stuck in the
upper right quadrant of Figure 2. We study two simple approaches for
constraining the size of attention’s recurrent state: linear attention and
sliding window attention. The recurrent state size of linear attention (i.e.
attention without softmax) does not grow with the sequence length and can be
modulated by changing simple hyperparameters [20]. With sliding window
attention, we cap the recurrent state size to be the width of the window.

However, our experiments on real-world language modeling (Table 4) and
synthetic associative recall (Figure 1 middle) suggest that neither primitive
alone suffices to navigate the pareto frontier. Linear attention lacks the
precision to perform local token shifts and comparisons [31, 1]. In sliding
window attention, associative recall range is limited by the width of the
windows (Figure 2, center). As we increase the window size, the recurrent
state grows linearly and has a non-linear effect on speed during parallel
training and inference (Figure 2, left).

Based models simply combine (1) softmax-approximating linear attention applied
globally and (2) exact softmax attention applied locally in small sliding
windows (Figure 1, right). This combination allows us to use softmax attention
in surprisingly small sliding windows (e.g., 64 tokens) that recover $90.8\%$
of full softmax attention’s recall accuracy at 1e-5$\times$ its latency.
Sliding window and linear attention alone fail (Figure 1, left).

###  4.1 Taylor Linear Attention

By approximating softmax attention using linear feature maps, we can constrain
the size of the recurrent state while maintaining global token interactions
(i.e. each token depends on every token before it in the sequence).

Katharopoulos et al. [20], Tsai et al. [32], Choromanski et al. [21] show that
we can select a feature map
$\phi:\mathbb{R}^{d}\rightarrow\mathbb{R}^{\tilde{d}}$ such that
$\phi(\bm{q}_{i})^{\top}\phi(\bm{k}_{j})\approx\exp(\bm{q}_{i}^{\top}\bm{k}_{j}/\sqrt{d})$.
We can then rewrite the formula for softmax attention in Equation 1 as

$\sum_{j=1}^{i}\frac{\phi(\bm{q}_{i})^{\top}\phi(\bm{k}_{j})\bm{v}_{j}}{\phi(\bm{q}_{i})\sum_{j=1}^{i}\phi(\bm{k}_{j})}=\frac{\phi(\bm{q}_{i})\sum_{j=1}^{i}\big{(}\phi(\bm{k}_{j})^{\top}\bm{v}_{j}\big{)}}{\phi(\bm{q}_{i})\sum_{j=1}^{i}\phi(\bm{k}_{j})}$
(2)

where every query attends to every past key in $\mathcal{O}(Nd^{2})$ time and
space complexity. Furthermore, Katharopoulos et al. [24] show that linear
attention has a fixed size recurrent state during generation. Letting
$\bm{s}_{i}=\sum_{j=1}^{i}\phi(\bm{k}_{j})^{\top}\bm{v}_{j}$ and
$\bm{z}_{i}=\sum_{j=1}^{i}\phi(\bm{k}_{j})^{\top}$ be a “KV-state” and
“K-state” respectively, we can compute Equation 2 as

$\bm{s}_{i}=\bm{s}_{i-1}+\phi(\bm{k}_{i})^{\top}\bm{v}_{i},\;\;\;\bm{z}_{i}=\bm{z}_{i-1}+\phi(\bm{k}_{i})^{\top},$
$\bm{y}_{i}=\frac{\phi(\bm{q}_{i})\bm{s}_{i}}{\phi(\bm{q}_{i})\bm{z}_{i}}$ (3)

where $\bm{s}_{i}\in\mathbb{R}^{d\times\tilde{d}}$ and
$\bm{z}_{i}\in\mathbb{R}^{\tilde{d}}$.

##### Feature map.

To approximate $\exp(\bm{q}_{i}^{\top}\bm{k}_{j}/\sqrt{d})$, we use the
$2^{\text{nd}}$-order Taylor series feature map, picking
$\phi:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d^{2}}$ such that

$\phi(\bm{q}_{i})^{\top}\phi(\bm{k}_{j})=1+\bm{q}_{i}^{\top}\bm{k}_{j}+\frac{(\bm{q}_{i}^{\top}\bm{k}_{j})^{2}}{2}$
(4)

While Zhang et al. [13] note that picking a feature map with $\tilde{d}=d^{2}$
results in linear attention with $\mathcal{O}(Nd^{3})$ time and space
complexity and large recurrent state of size $O(d^{3})$, we can tradeoff
efficiency for recall capacity by projecting queries and keys to smaller
dimensions i.e., $\bm{W}_{q},\bm{W}_{k}\in\mathbb{R}^{d\times d^{\prime}}$
with $d^{\prime}=16$. By changing $d^{\prime}$ we modulate the size of the
recurrent state.

Figure 3:  Linear attention feature maps on AR. $x$: state size (bytes) during
generation or param. count; $y$: MQAR accuracy. This setting is harder than
fig. 2 (256 key-value pairs).

How does the choice of feature map affect the memory-recall tradeoff? Prior
work demonstrates the strong performance of the Taylor feature map on
associative recall [13]. Building on this analysis, we evaluate a broad set of
feature maps ($\phi_{\text{ReLU}}(x)=\max(x,0)$,
$\phi_{\text{PosELU}}(x)=\text{ELU}(x)+1$, $\phi_{\text{Square}}(x)=x^{2}$,
$\phi_{\text{Identity}}(x)=x$, $\phi_{\text{CosFormer}}$ as defined in [23],
and $\phi_{\text{Performer}}$ as defined in [21]) using the experimental setup
described in Section 3.1. In Figure 3 (top), we plot the memory-recall
tradeoff curves for these feature maps. The Taylor series feature map, along
with the simple $\phi_{\text{PosELU}}$ and $\phi_{\text{ReLU}}$ feature maps,
sits at the pareto frontier. One advantage of the Taylor feature map over
these alternatives is that it expands the recurrent state size (improving
recall capacity) without changing the number of parameters. As shown in Figure
3 (bottom), the Taylor series feature map requires fewer parameters than
alternatives to achieve high recall capacity. This analysis and the ablations
in Table 4 informed our decision to use the Taylor approximation, though other
simple feature maps may be effective as well.

###  4.2 Local Exact Attention with Tensor Core Sliding Windows

To efficiently model fine-grained local interactions, Based uses sliding
window attention with window sizes set at small multiples of $16$ (up to 64
tokens). Similar to past (causal) implementations [18, 19], for window size
$w$ each query $\bm{q}_{i}$ only attends to past keys
$\\{\bm{k}_{i-w+1},\ldots,\bm{k}_{i}\\}$. This enables $\mathcal{O}(Nw)$ time
and space complexity for linear scaling in sequence length $N$, with a
$w$-sized KV-cache for constant-memory generation.

However, unlike past sliding window attentions that keep $w$ at sizes 256 [17]
to 4096 [10], Based uses only $w=$ $16$, $32$, or $64$ to best exploit modern
GPUs. In Section 5, we discuss how this “Tensor core-aware” window (tcWindow)
achieves 1e-5$\times$ the latency than the $w=4096$ windows in modern LLMs
(e.g., Mistral 7B [10]).

While the small $w$ in tcWindow enable fast local and exact attention, it
presents a challenge for long range modeling. With just $w=64$, for every
layer of $w=4096$ Mistral sliding window attention we would require $64$
layers of Based to achieve the same receptive field. Controlling for model
depth and sequence length, Figure 2 indeed shows smaller $w$ linearly
decreasing in associative recall accuracy. Based’s global _linear attention_
described above overcomes the lack of long-range modeling presented with low
$w$.

Additional architectural details for Based are discussed in Appendix C and the
hybridization of layers used in experiments are provided in Table 7. We
include ablations of architectural choices in Table 4 and evaluate the overall
quality and efficiency of Based in Section 6.

##  5 Efficient Implementation

In this section we focus on the efficiency of Based. A naïve implementation is
slower than the most efficient standard attention implementations (shown in
Figure 4) as it requires large amounts of high latency memory movement. We
first describe preliminaries of the GPU execution model and memory hierarchy.
We next present the baseline and our hardware-aware algorithms for linear
attention in Section 5.1 and for sliding window attention in Section 5.2.

##### Preliminaries

GPU operations, or kernels, are executed by thousands of parallel threads. In
NVIDIA terms, GPU streaming multiprocessors launch thread blocks at the
software level. These blocks are divided into warps (e.g. 32 threads) that are
assigned to cores at the hardware level. Threads need to read inputs into
their registers to perform computations and write the outputs. The time taken
to read and write is referred to as the IO cost.

Operations could either be memory or compute bound, depending on the time to
load data vs. perform computations on loaded data. In designing our IO-aware
algorithms, we would like to exploit two key properties of modern GPUs. First,
tensor core units (fast matrix multiply units) achieve 312 TFLOP/s speeds
relative to 19 TFLOP/s for the non-matrix multiply cores. Second, GPUs face a
memory hierarchy with large amounts of slow-to-access memory and smaller
amounts of fast-to-access memory. The access speed is governed by the
proximitiy of the processor to the data address. For instance, the hierarchy
on a modern NVIDIA 80GB A100 GPU is: 80GB of HBM with 2 TB/s bandwidth, 80MB
of L2 cache, 192KB of L1 cache / shared memory (implemented via SRAM) with 19
TB/s bandwidth per SM, and 256 KB of register file per SM [33]. Note that
register memory is private to an executing thread, so threads need to write to
shared memory to communicate data to other threads in the block. To reduce the
time required for reads and writes, a key principle is to fuse multiple
operations on the same data slice while it’s in fast memory before writing it
back to slow memory.

###  5.1 Taylor Exponential Linear Attention

Despite the theoretically improved complexity, the linear attention methods
demonstrated in prior work are often less efficient than highly-optimized
softmax attention implementations (Flash Attention [12]) when measured in
real-world wall-clock time and memory usage. We next present hardware-aware
algorithms to make Taylor linear attention efficient. We focus on two
operations: (1) prefill, corresponding to processing the prompt during
generation or the forward pass during training, and (2) next token prediction
during generation, which also requires updating the recurrent hidden state
state.

In this section, we refer to the batch size as $B$, number of heads as $H$,
head dimension as $d$, sequence length as $N$ and feature dimension as
$d^{\prime}$, following Section 4. For ease of notation, let
$D=1+d^{\prime}+d^{\prime 2}$ in this section. Additional details for the IO-
Aware algorithms are included in Appendix B.

####  5.1.1 Forward Pass / Generation Prefill

##### Baseline Implementation

The naïve implementation detailed in Appendix B (1) produces the feature maps
$Q,K$, (2) computes and materializes the large hidden state
$KV\in\mathbb{R}^{H\times d\times D}$, then (3) computes the causal dot
product between $Q$ and $KV$. Prior work has released popular CUDA kernels for
linear attention to efficiently perform the causal dot product / step (3)
[34]. The kernel parallelizes computation across the heads and batches, loads
tiles of $V$ and $K$ to SRAM, updates the running $KV$ state in SRAM, loads
tiles of $Q$ to SRAM, produces the final output in SRAM, and writes the result
to HBM.

Analysis In overall IO cost, ignoring the input and output projections in the
linear attention layer, this procedure requires $2BHND$ bytes for writing
featurized $Q,K$ to HBM. During the causal dot product, this requires
$2BHND+BHNd$ bytes to read $Q,K,V$ tiles and $BHNd$ bytes to write the result.
Throughout the computation, $\mathcal{O}(BHNDd)$ bytes (note this is the shape
$KV$ state during the forward pass) are read in and out of thread registers to
SRAM to update the running output and $KV$ state at 19TB/s bandwidth.

##### Algorithm

To improve efficiency, our algorithm computes both the feature map and the
causal dot product in fast memory. Our overall algorithm is detailed in
Algorithm 1 and we discuss the overall process here. The algorithm first
parallelizes over the batch $B$ and head $H$ dimensions, since computation is
independent for each head in each batch. We consider the three terms
$T_{0},T_{1},T_{2}\in\mathbb{R}^{N\times d}$ in the linear attention output,
corresponding to the three terms in the $2^{nd}$-order Taylor polynomial for
the exponential. We also consider the corresponding three terms, in the
cumulative $KV$ state as we process a sequence. For a given tile (i.e. sub-
matrix) of $Q,K,V$, we load the tiles (e.g., tile of
$Q,K\in\mathbb{R}^{16\times 16}$ and $V\in\mathbb{R}^{16\times 64}$ given $16$
tokens, $16$ feature dimension, and $64$ head dimension) into fast memory and
update the running $KV$ state and output for each of the three terms. We load
the tiles into the register file and fuse-operations to compute the causal dot
product in register, only writing to SRAM to synchronize the cumulative $KV$
state across parallel warps. We use explicit writes to SRAM for this
synchronization because register memory is warp-specific (threads in one warp
do not access the registers for a thread in a different warp).

Analysis In IO cost, again ignoring the input and output projections in the
linear attention layer, our procedure requires $2BHNd^{\prime}$ bytes for
reading $q,k$ and $2BHNd$ bytes for reading $v$ and writing output $y$ between
HBM and SRAM. Overall, our algorithm avoids in HBM $\mathcal{O}(2BHND)$ bytes
in HBM to SRAM data movement. We additionally improve upon the baseline by
performing computation in-register to avoid the $\mathcal{O}(BHNDd)$ bytes in
SRAM to register data movement.

####  5.1.2 Next token prediction

During next token prediction, an important consideration is how to efficiently
update the recurrent state $KV_{t}\in\mathbb{R}^{BHDd}$ at timestep $t$. The
expensive operation during next token prediction is computing the outer
product between projected hidden states $k_{t+1}\in\mathbb{R}^{BHD}$ and
$v_{t+1}\in\mathbb{R}^{BHd}$. The outer product requires $\mathcal{O}(BHDd)$
computation and space, and the result is summed with $KV_{t}$ to produce
$KV_{t+1}$. We provide an IO-aware algorithm for the state updates in
Algorithm 2. This algorithm incurs $\mathcal{O}(BHD+BHd)$ bytes of HBM to SRAM
data movement (to load the $q,k,v$ projections).

###  5.2 tcWindow

Next we motivate the choice of window size for tcWindow. In contrast to
sliding-window style models such as the popular Mistral models, which use
large window sizes $w=4096$ [10], Based chooses a window size based on
hardware specifications. GPU tensor cores operate on $16\times 16$ tiles.
Large GEMMs are compute bound (for e.g. in long-context attention). But, we
need sufficient occupancy to hide the latency of the tensor core units. Figure
1 (Right) shows $64\times 64$ dimension matrix multiplications are
approximately the same latency as $16\times 16$. Based sets $w$ to use
$64\times 64$ tiles (Figure 1). To distinguish from prior sliding windows, we
refer to this approach as tcWindow. We use the Flash Attention sliding window
implementation during training [11] and in Appendix B Algorithm 3, we provide
an IO-aware algorithm of tcWindow for next token prediction.

End-to-end benchmarks for Based implemented with these IO-aware algorithms are
provided in Section 6. Micro-benchmarks for each kernel against the baseline
implementations are provided in Appendix B.

##  6 Results

Figure 4: (Left) Throughput numbers for the varied prefill sequence lengths at
a fixed batch size of $2$. Right generation throughput at varied batch sizes
at a fixed generation length of $1024$ tokens. The $y$-axis shows the in
latency (ms). Lines are cutoff when the model runs out of memory. We show
results for both 360M and 1.3Bn, and all numbers are computed on a single
NVIDIA H100 GPU.

In this section, we present results for the following claims:

  1. 1.

Language modeling overall. We evaluate architectures in pretraining from
scratch on the Pile [30] and on standard benchmarks from the LM Eval Harness
[16]. We find Based matches or outperforms the strongest sub-quadratic
architectures (e.g. Mamba [5]) across these settings.

  2. 2.

Language modeling recall. Based closes the gap to attention on the challenging
associative recall slice of the real-world Pile language modeling corpus (see
Table 1), outperforming prior sub-quadratic architectures. We apply these
pretrained models zero-shot to a suite of recall-intensive tasks (e.g.
information extraction, QA), showing that Based systematically outperforms
other efficient architectures with comparable throughput.

  3. 3.

Generation throughput. Our IO-aware implementation of recurrent generation in
Based enables $40-60\%$ speedups relative to FlashAttention-2 and Mamba for
prefill at $4k$ sequence length and up to $24\times$ higher throughput over
FlashAttention-2 in generating $1024$ tokens at batch size $128$ (see Figure
4).

  4. 4.

Ablations of key design choices. We ablate four key design choices and
hyperparameters: (1) the feature map, (2) the feature dimension, (3) use of
local sliding-window softmax, and (4) the use of short gated-convolutions. The
results of these ablations are shown in Table 4.

##### Baselines

We compare to key baselines in the literature. We compare to Transformer (GPT
architecture) and Transformer++ (Llama architecture [15]), which adds rotary
encodings [35] and gated linear units. We compare an early class of efficient
architectures built from gating and long-convolution primitives including
Hyena [7], RWKV [8], and H3 [9]. We finally compare to recent state-of-the-art
architectures, that use input-dependent sequence aggregation to improve in
quality upon the long-convolution models, including Mamba [5] and Gated Linear
Attention [6].

###  6.1 Language Modeling Evaluations

##### Language Modeling Benchmarks

We pretrain language models from scratch at two parameter scales (355M and
1.3Bn parameters) on the Pile [30]. Each model sees the same 10 billion tokens
of pretraining data in the same order. The Pile data is tokenized using the
GPT-2 BPE tokenizer [36]. We measure perplexity on the Pile and report results
in Table 1 and further experimental details are provided in Section E.1.

We additionally evaluate the pretrained models on the standard LM Eval Harness
[16], as is standard protocol for our baselines [5, 6]. A detailed breakdown
of tasks and metrics can be found in Table 5. In both pretraining and on the
LM Eval Harness, Based consistently competes with the strongest Transformer++
and Mamba baselines. While these overall metrics are helpful, we next turn to
a fine-grained analysis of recall and in-context learning ability on real-
world data.

##### Recall and In-Context Learning Evaluations

We evaluate our pretrained models on a suite of in-context learning tasks
selected to test the downstream recall capacity in Table 1. These tasks fall
into three categories: (1) Real-world AR Beyond perplexity scores, we slice
the next token predictions on the Pile to understand each architecture’s AR
quality ( Section E.1). (2) Information extraction (IE) SWDE and FDA are
popular semi-structured and unstructured document IE benchmarks respectively
[37, 38, 39]. SWDE has raw HTML for $8$ Movie and $5$ University websites
(e.g. IMDB, US News) and annotations for 8-274 attributes per website (e.g.,
Movie `runtime`), and (3) Question answering from in-context passages. We find
Based outperforms the baseline sub-quadratic architectures across these
evaluations, closing the gap to Transformer++. These trends tracks the
${\mathrm{MQAR}}$ synthetic results from Section 3.1.

Model | Shots | BoolQ | CB | COPA | MultiRC | ReCoRD | RTE | WiC | WSC | Avg
---|---|---|---|---|---|---|---|---|---|---
|  |  Acc. $\uparrow$ |  Acc. $\uparrow$ |  F1 $\uparrow$ |  Acc. $\uparrow$ |  Acc. $\uparrow$ |  F1 $\uparrow$ |  EM $\uparrow$ |  Acc. $\uparrow$ |  Acc. $\uparrow$ |  Acc. $\uparrow$ |
Based (363m) | 0 | 59.0 | 41.1 | 19.4 | 69.0 | 54.9 | 14.5 | 14.0 | 52.0 | 50.0 | 36.5 | 45.7
1 | 57.5 | 37.5 | 26.8 | 68.0 | 52.5 | 19.9 | 19.2 | 47.7 | 50.9 | 49.0 | 47.2
5 | 56.6 | 44.6 | 28.9 | 73.0 | 53.6 | 24.9 | 24.1 | 48.7 | 51.1 | 39.4 | 48.0
Transformer++ (360m) | 0 | 57.3 | 41.1 | 21.3 | 67.0 | 57.0 | 16.6 | 16.1 | 53.8 | 50.0 | 37.5 | 46.3
1 | 54.2 | 39.3 | 25.3 | 69.0 | 51.5 | 22.2 | 21.6 | 50.9 | 47.0 | 55.8 | 47.8
5 | 50.7 | 58.9 | 49.9 | 64.0 | 46.9 | 24.2 | 23.6 | 47.3 | 52.2 | 51.9 | 48.9
Mamba (358m) | 0 | 57.5 | 35.7 | 24.4 | 71.0 | 57.2 | 18.8 | 18.3 | 52.4 | 50.0 | 36.5 | 46.6
1 | 51.1 | 39.3 | 27.4 | 71.0 | 52.9 | 21.6 | 21.0 | 46.6 | 46.2 | 52.9 | 46.9
5 | 41.1 | 37.5 | 23.6 | 69.0 | 49.2 | 20.4 | 19.9 | 48.4 | 51.7 | 51.9 | 45.2
Table 2: Few-shot downstream evaluation on SuperGLUE of pre-trained language
models. The same set of models as in table 1, all were trained on the same 10
billion tokens drawn from the Pile [30], evaluated on the SuperGLUE benchmark
[40] using the LM eval harness by EleutherAI [16]. When computing the average,
we first average the metrics by task and then average across tasks.

In order to understand in-context-learning performance, we next perform few-
shot evaluations on the SuperGLUE benchmark [40] for Based, Mamba and
Transformer++ in Table 2. Each model was evaluated on all tasks using under 0
shot (i.e., number of in-context examples), 1 shot and 5 shot prompting,
respectively. Transformer++ and Based both see monotonic improvement from
increasing the number of shots. For Mamba, however, albeit getting a slight
improvement from 0-shot to 1-shot, it performs worse on 5-shot than even on
0-shot. This result suggests that the limited recall ability observed in Mamba
could also impact few-shot abilities.

##### DNA modeling

Towards understanding the capability of Based beyond natural English language,
we next evaluate each architecture on its ability to perform DNA next token
prediction (Table 3). We evaluate architectures on the HG38 (human genome)
benchmark at $1k$, $4k$, and $8k$ sequence lengths used in prior architecture
evaluations [41, 5]. The DNA tasks uses a byte-level tokenizer wherein the
vocabulary consists of characters corresponding to the nucleotide bases. We
find Based is competitive with state-of-the-art architectures across evaluated
sequence lengths. In Table 5, we evaluate these checkpoints on downstream DNA
classification tasks.

Model | Params | HG38 PPL $\downarrow$
---|---|---
$N=$1024 |  $N=$4096 |  $N=$8192
Transformer++ | 46.2 | 2.52 | 2.50 | 2.51
Mamba | 46.1 | 2.51 | 2.49 | 2.49
Based | 48.8 | 2.51 | 2.50 | 2.49
Table 3: DNA modeling performance on the HG38 dataset. All models are
pretrained from scratch for 10Bn tokens at $N=$ 1k, 4k, and 8k sequence
lengths respectively. We report results after hyperparameter sweeping the
learning rate for each architecture.

###  6.2 Efficiency Benchmarks

We benchmark the throughput of Based, with and without our proposed IO-Aware
algorithms (Section 5). We consider both the forward pass / generation prefill
and next token prediction stages. Experiments were run using an H100 NVIDIA
GPU and averaged over $20$ repetitions. Results are shown in Figure 4.

##### End-to-end benchmarks

Using our efficient implementation (Section 5), Based achieves 56% faster
prefill than FlashAttention-2 [11] and 44% faster than Mamba at $4k$ sequence
length and 1.3Bn parameters (28% faster than FlashAttention-2 and 76% faster
than Mamba at 360M parameters). We find that next token generation, with no
prefill, provides $24\times$ higher throughput (tokens/second) over the highly
optimized FlashAttention-2 implementation and achieves 95% and the throughput
of the recurrent Mamba architecture at batch size $128$ and 1.3Bn parameters
(98% higher throughput vs. FlashAttention-2 and 118% higher throughput vs.
Mamba at 360M parameters). All benchmarks is on a single NVIDIA H100 GPU,
using CUDA cache graphs during next token prediction [42].

In Figure 4, we also include results for the baseline implementation of Based
that uses the popular Fast Transformers CUDA kernel in the linear attention
literature to compute the causal dot product [34] (discussed in Section 5).
The custom kernel introduced in our work unlocks the efficiency of Based.

##### Micro benchmarks

As the end-to-end Based architecture is a hybrid architecture, we provide
micro benchmarks of the individual IO-aware kernels against the key baseline
implementations in Appendix B.

###  6.3 Quality Ablations

Our objective with Based is to measure the throughput and recall of the
simplest possible linear attention model that achieves strong performance.
Therefore, we ablate the key design decisions — choice of feature map, feature
dimension for the Taylor map, use of sliding window and convolutions — to
understand their contributions to the quality of Based. We perform these
ablations on the Pile corpus [30] using the same number of tokens and data
ordering as the prior experiments.

In feature map ablations, we consider the CosFormer [23] and Performers [21]
feature maps, which have been demonstrated as strong linear attentions in
prior work [13]. We also include a baseline that expands the state size using
learned projections and applies CosFormer towards comparing to the larger
state size of the Taylor map. For these baselines, we keep the rest of the
Based architecture the same (i.e. in the number of linear attention layers and
hybridization with sliding window and gated convolution layers). We observe
that with the larger state size, CosFormer quality is increasingly competitive
with the Taylor map. We note that expanding the state size requires increasing
the model’s overall parameter count (due to the learned projections) in the
CosFormer baseline, in contrast to the Taylor map.

Next, we ablate the feature dimension, holding the feature map fixed to the
Taylor map. We find larger feature dimension improves quality, with
diminishing returns going from $24$ to $32$ dimension. Note that feature
dimension $\sqrt{1024}=32$, where $1024$ is the attention model dimension at
the $360$ parameter scale in our experiments.

Finally, the ablations show that eliminating the convolutions and/or the
sliding window attention degrades quality. We observe that adding either
convolutions or sliding window helps on the associative recall slice relative
to neither (e.g. $2.29$ AR Ppl. on the Pile with neither vs. $2.09$ or $2.11$
with sliding window or convolutions.). Further, increasing the window size
from $0$ to $64$ vs. $64$ to $128$ (also an efficient design point in Figure
1, left) continues to help quality, but with marginal improvements.

Hyperparameters | Language Modeling (Pile) | Info. Extraction | QA
---|---|---|---
Feat. Map | Feat. Dim. | Sliding | Convs. | Decay | All | AR | Other | FDA | SQUAD
Ppl. $\downarrow$ |  Ppl. $\downarrow$ |  Ppl. $\downarrow$ |  Acc. $\uparrow$ |  Acc. $\uparrow$
Taylor Exp. (2) | 16 (153) | ✓(64) | ✓ | ✓ | 8.65 | 2.07 | 9.64 | 11.71 | 25.07
Performer | 16 (16) | ✓(64) | ✓ | ✓ | 9.08 | 8.53 | 11.62 | 0.36 | 7.47
CosFormer | 16 (32) | ✓(64) | ✓ | ✓ | 9.03 | 2.42 | 9.98 | 7.71 | 24.63
CosFormer | 64 (128) | ✓(64) | ✓ | ✓ | 8.82 | 2.18 | 9.80 | 9.07 | 27.85
Taylor Exp. (2) | 32 (561) | ✓(64) | ✓ | ✓ | 8.56 | 2.00 | 9.57 | 12.89 | 26.74
Taylor Exp. (2) | 24 (325) | ✓(64) | ✓ | ✓ | 8.58 | 2.02 | 9.58 | 20.87 | 24.77
Taylor Exp. (2) | 16 (153) | ✓(64) | ✓ | ✓ | 8.65 | 2.07 | 9.64 | 11.71 | 25.07
Taylor Exp. (2) | 8 (45) | ✓(64) | ✓ | ✓ | 8.77 | 2.18 | 9.75 | 12.79 | 22.35
Taylor Exp. (2) | 16 (153) | ✓(64) | ✓ | ✓ | 8.65 | 2.07 | 9.64 | 11.71 | 25.07
Taylor Exp. (2) | 16 (153) | ✓(64) | ✓ | ✗ | 8.65 | 2.04 | 9.66 | 1.72 | 10.80
Taylor Exp. (2) | 16 (153) | ✗ | ✓ | ✓ | 8.91 | 2.11 | 9.94 | 10.16 | 24.5
Taylor Exp. (2) | 16 (153) | ✓(64) | ✗ | ✓ | 8.74 | 2.09 | 9.74 | 2.36 | 18.87
Taylor Exp. (2) | 24 (325) | ✗ | ✗ | ✓ | 9.49 | 2.29 | 10.58 | 8.71 | 11.33
Taylor Exp. (2) | 16 (153) | ✓(128) | ✓ | ✓ | 8.61 | 2.06 | 9.60 | 14.39 | 31.84
Taylor Exp. (2) | 16 (153) | ✓(64) | ✓ | ✓ | 8.65 | 2.07 | 9.64 | 11.71 | 25.07
Table 4: Ablations of design choices and hyperparameters in Based. All models
are 362M param variants of the Based architecture described in Section 4,
trained to 10 billion tokens on the Pile. We ablate the hyperparameters
central to the design of Based: (1) the choice of feature map $\phi$ (see
Section 4.1), (2) the size of the feature dim $d^{\prime}$ (we show the
effective size of the feature after applying the feature map in parantheses,
see Section 4.1), (3) the use of local sequence mixers (sliding window
attention and short convolutions), and (4) the data-dependent decay defined in
Section 4.

##  7 Conclusion

This work studies the properties of high-quality and efficient sequence
mixers. We identify a fundamental tradeoff between recall, key to in-context
learning, and throughput through theory and experiments. Attention performs
recall perfectly, but requires retaining a KV cache that grows with the
sequence length. As an alternative, Based combines two simple techniques —
local fine-grained attention and long-range linear attention via a Taylor
approximation of the softmax exponential function – that are sub-quadratic
during training and permit an efficient recurrent inference view. To enable
wall clock efficiency, we introduce IO-aware algorithms of the Taylor linear
attention computation that leads Based to perform generation up to $24\times$
faster than FlashAttention-2 at the 1.3Bn parameter scale, generating 1024
tokens. Beyond competing in overall perplexity, Based outperforms prior sub-
quadratic architectures in recall quality by up to 6.2 accuracy points,
pointing to Based’s promise as a language modeling architecture.

#### Acknowledgments

We thank Tri Dao, Daniel Fu, Songlin Yang, Jessica Grogan, Albert Gu, Eric
Nguyen, Michael Wornow, Alyssa Unell, and Gautam Machiraju for their helpful
feedback and discussion during this work. We thank the Hazy Research lab and
Together AI for supporting this work. We gratefully acknowledge the support of
NIH under No. U54EB020405 (Mobilize), NSF under Nos. CCF2247015 (Hardware-
Aware), CCF1763315 (Beyond Sparsity), CCF1563078 (Volume to Velocity), and
1937301 (RTML); US DEVCOM ARL under Nos. W911NF-23-2-0184 (Long-context) and
W911NF-21-2-0251 (Interactive Human-AI Teaming); ONR under Nos. N000142312633
(Deep Signal Processing), N000141712266 (Unifying Weak Supervision),
N000142012480 (Non-Euclidean Geometry), and N000142012275 (NEPTUNE); Stanford
HAI under No. 247183; NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC,
Toshiba, TSMC, ARM, Hitachi, BASF, Accenture, Ericsson, Qualcomm, Analog
Devices, Google Cloud, Salesforce, Total, the HAI-GCP Cloud Credits for
Research program, the Stanford Data Science Initiative (SDSI), and members of
the Stanford DAWN project: Facebook, Google, and VMWare. The U.S. Government
is authorized to reproduce and distribute reprints for Governmental purposes
notwithstanding any copyright notation thereon. Any opinions, findings, and
conclusions or recommendations expressed in this material are those of the
authors and do not necessarily reflect the views, policies, or endorsements,
either expressed or implied, of NIH, ONR, or the U.S. Government. AR’s
research is supported by NSF grant CCF#2247014.

## References

  * Arora et al. [2023a] Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Ré.  Zoology: Measuring and improving recall in efficient language models.  _International Conference on Learning Representations_ , 2023a.
  * Vaswani et al. [2017] 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.
  * Olsson et al. [2022] 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.
  * Wang et al. [2022] Junxiong Wang, Jing Nathan Yan, Albert Gu, and Alexander M Rush.  Pretraining without attention.  _arXiv preprint arXiv:2212.10544_ , 2022.
  * Gu and Dao [2023] Albert Gu and Tri Dao.  Mamba: Linear-time sequence modeling with selective state spaces.  _arXiv preprint arXiv:2312.00752_ , 2023.
  * Yang et al. [2023] 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.
  * Poli et al. [2023a] Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré.  Hyena hierarchy: Towards larger convolutional language models.  _arXiv preprint arXiv:2302.10866_ , 2023a.
  * Peng et al. [2023] 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.
  * Fu et al. [2023a] Daniel Y. Fu, Tri Dao, Khaled K. Saab, Armin W. Thomas, Atri Rudra, and Christopher Ré.  Hungry Hungry Hippos: Towards language modeling with state space models.  In _International Conference on Learning Representations_ , 2023a.
  * Jiang et al. [2023] 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.
  * Dao [2023] Tri Dao.  FlashAttention-2: Faster attention with better parallelism and work partitioning.  2023\.
  * Dao et al. [2022] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré.  FlashAttention: Fast and memory-efficient exact attention with IO-awareness.  In _Advances in Neural Information Processing Systems_ , 2022.
  * Zhang et al. [2024] 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.  URL https://openreview.net/forum?id=4g02l2N2Nx.
  * Keles et al. [2023] 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.
  * Touvron et al. [2023] 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.
  * Gao et al. [2023] 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.  URL https://zenodo.org/records/10256836.
  * Parmar et al. [2018] 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.
  * Child et al. [2019] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever.  Generating long sequences with sparse transformers.  _arXiv preprint arXiv:1904.10509_ , 2019.
  * Beltagy et al. [2020] Iz Beltagy, Matthew E Peters, and Arman Cohan.  Longformer: The long-document transformer.  _arXiv preprint arXiv:2004.05150_ , 2020.
  * Katharopoulos et al. [2020a] 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.
  * Choromanski et al. [2020] 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.
  * Choromanski et al. [2021] 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.
  * Qin et al. [2022a] 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.
  * Katharopoulos et al. [2020b] 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.  URL https://arxiv.org/abs/2006.16236.
  * Kasai et al. [2021] 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.  URL https://aclanthology.org/2021.emnlp-main.830.
  * Schlag et al. [2021] Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber.  Linear transformers are secretly fast weight programmers.  In _International Conference on Machine Learning_ , pages 9355–9366. PMLR, 2021.
  * Gu et al. [2021] Albert Gu, Karan Goel, and Christopher Ré.  Efficiently modeling long sequences with structured state spaces.  _arXiv preprint arXiv:2111.00396_ , 2021.
  * Sun et al. [2023] 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.
  * Akyürek et al. [2024] Ekin Akyürek, Bailin Wang, Yoon Kim, and Jacob Andreas.  In-context language learning: Architectures and algorithms.  2024\.  URL https://arxiv.org/abs/2401.12973.
  * Gao et al. [2020] 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.
  * Fu et al. [2023b] Daniel Y. Fu, Elliot L. Epstein, Eric Nguyen, Armin W. Thomas, Michael Zhang, Tri Dao, Atri Rudra, and Christopher Ré.  Simple hardware-efficient long convolutions for sequence modeling.  _arXiv preprint arXiv:2302.06646_ , 2023b.
  * Tsai et al. [2019] 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.
  * NVIDIA [2022] NVIDIA.  Nvidia H100 tensor core GPU architecture, 2022.
  * Vyas et al. [2020] A. Vyas, A. Katharopoulos, and F. Fleuret.  Fast transformers with clustered attention.  2020\.
  * Su et al. [2023] Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu.  Roformer: Enhanced transformer with rotary position embedding, 2023.
  * Radford et al. [2019] 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.
  * Wu et al. [2021] 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.
  * Deng et al. [2022] Xiang Deng, Prashant Shiralkar, Colin Lockard, Binxuan Huang, and Huan Sun.  Dom-lm: Learning generalizable representations for html documents.  2022\.
  * Arora et al. [2023b] Simran Arora, Brandon Yang, Sabri Eyuboglu, Avanika Narayan, Andrew Hojel, Immanuel Trummer, and Christopher Ré.  Language models enable simple systems for generating structured views of heterogeneous data lakes.  _arXiv:2304.09433_ , 2023b.
  * Wang et al. [2019] 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.
  * Nguyen et al. [2023] 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 Ré.  Hyenadna: Long-range genomic sequence modeling at single nucleotide resolution, 2023.
  * NVIDIA [2019] NVIDIA.  Getting started with cuda graphs, 2019.  URL https://developer.nvidia.com/blog/cuda-graphs/.
  * Kitaev et al. [2020] Nikita Kitaev, Łukasz Kaiser, and Anselm Levskaya.  Reformer: The efficient transformer.  _arXiv preprint arXiv:2001.04451_ , 2020.
  * Zaheer et al. [2020] 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.
  * Wang et al. [2020] Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma.  Linformer: Self-attention with linear complexity.  _arXiv preprint arXiv:2006.04768_ , 2020.
  * Zhu et al. [2021] 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.
  * Alberti et al. [2023] Silas Alberti, Niclas Dern, Laura Thesing, and Gitta Kutyniok.  Sumformer: Universal approximation for efficient transformers.  _arXiv preprint arXiv:2307.02301_ , 2023.
  * Tay et al. [2022] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler.  Efficient transformers: A survey.  _ACM Computing Surveys_ , 55(6):1–28, 2022.
  * Xiong et al. [2021] Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh.  Nyströmformer: A nyström-based algorithm for approximating self-attention.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 35, pages 14138–14148, 2021.
  * Chen et al. [2021a] Yifan Chen, Qi Zeng, Heng Ji, and Yun Yang.  Skyformer: Remodel self-attention with gaussian kernel and nystr\”om method.  In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, _Advances in Neural Information Processing Systems_ , 2021a.  URL https://openreview.net/forum?id=pZCYG7gjkKz.
  * De Brebisson and Vincent [2015] Alexandre De Brebisson and Pascal Vincent.  An exploration of softmax alternatives belonging to the spherical loss family.  _arXiv preprint arXiv:1511.05042_ , 2015.
  * Chen et al. [2021b] Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré.  Scatterbrain: Unifying sparse and low-rank attention approximation.  _arXiv preprint arXiv:2110.15343_ , 2021b.
  * Qin et al. [2022b] 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.  URL https://aclanthology.org/2022.emnlp-main.473.
  * Ba et al. [2016] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton.  Layer normalization.  _arXiv preprint arXiv:1607.06450_ , 2016.
  * Zhang and Sennrich [2019] Biao Zhang and Rico Sennrich.  Root mean square layer normalization.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Cooley and Tukey [1965] 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.
  * Romero et al. [2022] David W. Romero, Anna Kuzina, Erik J. Bekkers, Jakub M. Tomczak, and Mark Hoogendoorn.  Ckconv: Continuous kernel convolution for sequential data.  2022\.
  * Gupta et al. [2022] Ankit Gupta, Albert Gu, and Jonathan Berant.  Diagonal state spaces are as effective as structured state spaces, 2022.
  * Gu et al. [2022] Albert Gu, Ankit Gupta, Karan Goel, and Christopher Ré.  On the parameterization and initialization of diagonal state space models, 2022.
  * Mehta et al. [2022] Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur.  Long range language modeling via gated state spaces, 2022.
  * Ma et al. [2022] 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.
  * Elhage et al. [2021] 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.
  * Dauphin et al. [2017] 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.
  * Ren et al. [2023] 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.
  * Liu et al. [2023] Hao Liu, Matei Zaharia, and Pieter Abbeel.  Ring attention with blockwise transformers for near-infinite context.  _arXiv preprint arXiv:2310.01889_ , 2023.
  * Kwon et al. [2023] 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.
  * Fu et al. [2023c] Daniel Y Fu, Hermann Kumbong, Eric Nguyen, and Christopher Ré.  Flashfftconv: Efficient convolutions for long sequences with tensor cores.  _arXiv preprint arXiv:2311.05908_ , 2023c.
  * Rabe and Staats [2021] Markus N Rabe and Charles Staats.  Self-attention does not need $o(n^{2})$ memory.  _arXiv preprint arXiv:2112.05682_ , 2021.
  * Jang et al. [2019] 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.
  * Liu and Abbeel [2023] Hao Liu and Pieter Abbeel.  Blockwise parallel transformer for long context large models.  _arXiv preprint arXiv:2305.19370_ , 2023.
  * Hua et al. [2022] 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.
  * Poli et al. [2023b] Michael Poli, Jue Wang, Stefano Massaroli, Jeffrey Quesnelle, Ryan Carlow, Eric Nguyen, and Armin Thomas.  StripedHyena: Moving Beyond Transformers with Hybrid Signal Processing Models.  12 2023b.  doi: 10.57967/hf/1595.  URL https://github.com/togethercomputer/stripedhyena.
  * Hendrycks and Gimpel [2023] Dan Hendrycks and Kevin Gimpel.  Gaussian error linear units (gelus), 2023.
  * Gresova et al. [2022] 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.  URL https://www.biorxiv.org/content/early/2022/06/10/2022.06.08.495248.
  * Paperno et al. [2016] Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández.  The lambada dataset: Word prediction requiring a broad discourse context, 2016.
  * Zellers et al. [2019] Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi.  Hellaswag: Can a machine really finish your sentence?, 2019.
  * Bisk et al. [2019] Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi.  Piqa: Reasoning about physical commonsense in natural language, 2019.
  * Clark et al. [2018] 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.
  * Sakaguchi et al. [2019] Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi.  Winogrande: An adversarial winograd schema challenge at scale, 2019.
  * Lockard et al. [2019] 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.  URL https://aclanthology.org/N19-1309.
  * Bürgisser et al. [2013] Peter Bürgisser, Michael Clausen, and Mohammad A Shokrollahi.  _Algebraic complexity theory_ , volume 315.  Springer Science & Business Media, 2013.
  * Jayram et al. [2008] Thathachar S Jayram, Ravi Kumar, and Dandapani Sivakumar.  The one-way communication complexity of hamming distance.  _Theory of Computing_ , 4(1):129–135, 2008.
  * Kopparty [2020] Swastik Kopparty.  Topics in algorithms and complexity theory: Spring 2020.  2020\.

The appendix is organized as follows:

  1. 1.

Appendix A includes an extended related works discussion.

  2. 2.

Appendix B includes details on the IO-aware implementation and benchmarking
for Based.

  3. 3.

Appendix C includes additional discussion of Based architectural details.

  4. 4.

Table 5 provides additional experimental results.

  5. 5.

Appendix E provides experimental details.

  6. 6.

Appendix F includes theoretical results and proofs.

##  Appendix A Extended Related Work

Our work relates broadly to various developments in efficient sequence
modeling. In this section, we organize these related works into (1) model-
based or algorithmic contributions (section A.1) and (2) implementation or
systems-based contributions (section A.2).

###  A.1 Efficient Language Modeling Architectures

While Transformers often achieve state-of-the-art language modeling quality,
their design motivates various efficiency improvements when both processing
input sequences and generating outputs. In particular, various works try to
retain their modeling quality, while improving on their quadratic scaling
($\mathcal{O}(N^{2})$ in input sequence length $N$) when processing inputs and
$\mathcal{O}(NM)$ time and space when decoding outputs for outputs of length
$M$ (when caching prior keys and values in the attention mechanism).

We note that most related lines of work build on one of two primitives:
_attention approximations_ (e.g., linear attentions, sparse attentions, sparse
and low-rank attentions), or _state-space models_ (SSMs) (which have
alternative parameterizations as either “long” convolutional models or
recurrent neueral networks). Both model classes achieve subquadratic time and
space complexity when processing inputs, while linear attentions and SSMs also
enable better than $\mathcal{O}(NM)$ decoding via their ability to process
inputs recurrently like a recurrent neural network (RNN).

We describe each of these model classes next.

####  A.1.1 Efficient Attentions

We focus on two of the most related paradigms for efficiently computing
attention here, _structured sparse attentions_ and _linear attentions_. We
acknowledge a great deal of prior work to compute attention more efficiently,
such as via locality-sensitive hashing [43], random sparse attentions [44],
and sequence compression [45, 46, 47]. Please see [48] for a comprehensive
survey.

##### Structured sparse attentions

Structured parse attentions reduce attention’s time and memory requirements by
only attending over specific strided patterns or local _sliding windows_ [17,
18, 19]. For example, [17] propose computing attention only over a local
window of the past $w$ tokens, such that processing sequences $N$ tokens long
only takes $\mathcal{O}{Nw}$ time and space. [18] note that this window alone
may not all capture all desired dependencies (such as long-term interactions),
and propose two strided patterns to compute dot products between queries and
keys further away. [19] further propose allowing specific tokens to attend to
all other tokens in a dense manner.

While further popularized in recent large language models (Mistral, Jiang et
al. [10]), we note that these implementations use large window sizes that
still leave room for improving efficiency. In Based, we introduce a hardware-
guided design (using small windows) and sliding window implementation that
allows us to capitalize on sparse attention’s efficiency.

##### Linear attentions

Linear attentions preserve the same “sequence-mixing” operations as standard
attention, computing dot products between queries and keys to weight
corresponding values. However, their key insight is to replace the softmax in
standard attention with alternative kernel functions [20]. Mechanically, by
removing the $\exp(\bm{q}^{\top}\bm{k})$ in favor of feature map dot-products
$\phi(\bm{q})^{\top}\phi(\bm{k})$, these methods use matrix product
associativity to compute attention in $\mathcal{O}(Nd^{2})$ time and space
[24] (Equation 2). Furthermore, they permit a _recurrent view_ for constant
memory and $\mathcal{O}(1)$ time per-token generation [25, 26] (Equation 3).

Prior works propose different feature maps $\phi$ to improve linear attention
modeling quality. [20] originally use the _positive elu_ function
$1+\text{elu}$ such that $\phi(\bm{q})^{\top}\phi(\bm{k})$ remains positive
and attention weights remain affine. [23] instead use the ReLU function
combined with a cosine-based reweighting function to add a locality bias.
Other approaches propose feature maps that aim to approximate the Softmax,
such as Random Fourier Features [21, 22] the Nystrom method [49, 50], or
deterministic low-degree polynomial approximations [13, 51, 14]. Finally,
recent works treat the feature map as a learnable function [25], and
optionally train the feature map explicitly to recover the softmax kernel
[13].

##### Combining sparse and linear attentions

Finally, another line of works similar in spirit to Based combine sparse and
linear attentions. Scatterbrain [52] use a sparse and low-rank decomposition
to unify past sparse and low-rank approximations. They theoretically and
empirically show that any low rank approximation of attention’s $\exp(QK^{T})$
will have a much larger approximation error than a sparse plus low rank
approximation, echoed in our observations. TransNormer [53] apply
normalizations such as LayerNorm [54] or RMSNorm [55] to linear attention
outputs in certain layers, and apply softmax attention in local chunks in
other layers.

####  A.1.2 Attention Alternatives

We now review other attention alternatives, which focus on improving upon the
quadratic scaling of attention. Initial work in this vein uses linear time
invariant state space models (SSMs) or long convolutions, which can
efficiently process sequences of length $N$ in $O(N\log N)$ time invoking the
FFT-convolution theorem [56], as the sequence mixer [27, 57, 58, 59, 60, 61,
4, 31]. SSMs can also be rewritten as recurrences to permit fast $O(1)$
inference.

Subsequent work identified that the long convolution alone is not expressive
enough to perform particular sub-tasks in language modeling. Prior work shows
pure linear SSMs cannot perform associative recall, a skill that is correlated
with a model’s in-context learning capability [62, 3], and introduces
multiplicative interactions (via gating or Hadamard product [63]) between
tokens to allow the model to compare tokens in the sequence [9, 7, 8].
However, Arora et al. [1] show empirically and theoretically the class of
gated convolution architectures, any architectures built from the two gating
and convolution primitives, struggles to learn associative recall (on
synthetic and real language data) as efficiently as attention. They show that
while attention solves AR in constant many layers / with model dimension that
is independent of sequence length, any gated convolution architecture uses
dimensionality that scales with the sequence length — we build upon their
upper bound theoretical results with a lower bound argument in Section 3.2. We
also study a broader set of architectures in this work beyond gated
convolutions.

Gu and Dao [5], Arora et al. [1], Yang et al. [6] identify that the use of
input-dependent sequence mixers is important for an architecture to perform AR
as efficiently as attention. AR requires shifting information that appears
prior in a sequence to interact with the current (last) tokens in the
sequence, in order to predict the next token [9]. While gating is one way to
introduce data-dependence [7], allowing comparing tokens in two (e.g. a
shifted and unshifted) sequences, it is difficult to select which information
from the prefix of the sequence to shift forwards in the first place, using
gating alone. Intuitively, the information to shift depends on the input’s
properties. Thus, several subquadratic architectures consider alternate
strategies to introduce input-dependence [24, 5, 64, 61, 6]. We present
another strategy for efficient input-dependent sequence mixing in our work.

###  A.2 Efficient Implementations

Beyond designing new model architectures, various works introduce systems-
level innovations to improve training and inference efficiency. These include
alternative implementations of architecture primitives such as attention [11,
65, 66], long convolutions [67, 31], and linear attention [20, 6]. They
frequently achieve both reduced memory and increased computational speed on
modern GPUs by “fusing” operations such as matrix multiplications into a
single CUDA kernel, and designing “IO-aware” ways to distribute and compute
the results of various read and write operations between different levels of
GPU memory.

####  A.2.1 Efficient Attention Implementations

[12] introduce FlashAttention, an alternative yet exact implementation of
softmax attention that improves memory and speed by both fusing attention
operations into a single CUDA kernel and distributing the attention operations
to better exploit High Bandwidth Memory (HBM) and Static Random Access Memory
(SRAM). They first compute attention’s query-key-value dot-products, masking,
and softmax, together as a single kernel. By doing so after a single load to
SRAM before moving the output back to HRAM, they exploit SRAM’s fast compute
and reduce the total number of read-write operations. To get around SRAM’s
small memory size and avoid attention’s quadratic memory size over input
sequence length, they use _tiling_ to split up the query, key, and value
inputs into smaller “blocks”, compute the attention operations for each block,
and adjust the outputs after computing all blocks to properly normalize the
softmax [68, 69]. To perform backpropagation fast on SRAM, they get around
SRAM’s limited storage by _recomputing_ the gradients rather than storing
them. Despite the extra operations, this IO-aware implementation still
significantly improves wall-clock time during training.

Similarly making use of block-wise computation, [65] instead compute attention
blocks across different _devices_ in RingAttention, enabling training and
inference over much larger context lengths that scale with device count. They
distribute and compute the attention operations in each block across multiple
hosts in parallel, likewise keeping track of summary statistics to gather
results correctly into exact attention. However, they introduce an
“overlapping” mechanism to coordinate communication of blocks to reduce
overhead. They further make use of Blockwise Parallel Transformers [70] to
reduce memory, which similar to FlashAttention removes the quadratic in memory
scaling of attention by dividing the attention operation into separate blocks
before gathering back the adjusted softmax output with block-wise
normalization statistics.

As a complement to attention training and inference, [66] improve attention
generation with PagedAttention. PagedAttention similarly uses block-wise
computation to address memory utilization issues during generation, where the
KV cache can grow an undetermined amount. Existing systems may naïvely handle
this by pre-allocating large amounts of contiguous memory. However, this can
result in low utilization and computational bottlenecks. Accordingly,
PagedAttention divides attention’s growing KV cache into _KV blocks_ that can
be stored separately on physical memory. This enables more flexible memory
management, where smaller chunks can be allocated in different locations when
needed to reduce memory-based bottlenecks.

In Based, we use similar blocking strategies to more efficiently compute both
the second-order Taylor series linear attention and the sliding window softmax
attention, and for both training and inference.

####  A.2.2 Efficient Attention-Alternative Implementations

Beyond optimizations for attention, various works also introduce similar “IO-
aware” implementations to improve memory usage and speed for convolutional and
recurrent operations. We overview the most relevant works to Based, which make
use of similar techniques such as fusing operations and blocking (tiling) to
compute results in SRAM.

##### Long convolutions

[67] improve the efficiency of long convolutions on modern GPUs. They build on
using the Fast Fourier Transform (FFT), which enables computing convolutions
with filter sizes equal to input sequence length from $\mathcal{O}(N^{2})$ (if
$N$ is filter size and sequence length) to $\mathcal{O}(N\log N)$. However, to
compute this algorithm efficiently on GPUs, they break down the convolution
into separate matrix multiply operations via a _Monarch_ decomposition of the
FFT, which allows both (1) fusing multiple steps of the FFT together (for
reduced read-write operations) and (2) scheduling these operations for fast
computation in SRAM while remaining under the smaller SRAM memory constraints.

##### Recurrence

[5] improve the efficiency of recent neural state-space models (SSMs) [27]
using several similar techniques to FlashAttention, specifically with regard
the recurrent view. They load the SSM parameters into SRAM for computation
before saving results back in HBM, and also use _recomputation_ where during
backpropagation the intermediate states are not saved but rather recomputed
when inputs are loaded from HBM to SRAM. They finally improve wall-clock time
by parallelizing the recurrent view of the SSM as a parallel scan.

##### Linear Attention

Finally, several works propose techniques to improve the real-world wall-clock
time and memory-usage of linear attention. [20] fuse several operations in the
causal dot product of linear attention. [6] use blocking to divide the linear
attention matrices into SRAM-computable chunks in FlashLinearAttention. As a
trade-off between the slow yet memory-efficient RNN view of linear attention
and faster but memory-intensive parallel “standard attention” view, they
further optimize a “chunk-wise” implementation of linear attention [71]. When
processing input sequences, the input is first divided into several non-
overlapping chunks, where we save memory by computing “kv states” at the end
of each chunk, and save time by computing the tokens in a given chunk in
parallel.

##  Appendix B IO Aware Implementations

In this section, we provide additional details pertaining to the benchmarking
experiments and we provide micro-benchmarking results for the individual
kernels we contribute, to complement the end-to-end benchmarking results in
the Section 6.

###  B.1 Forward / Generation Prefill

##### Baselines

In Figure 4, we implement Based using our IO-aware Taylor linear attention
Algorithm 1. The baseline approach presented in [13], prior to our kernel,
uses the popular linear attention CUDA kernel from Fast Transformers for
computing the causal dot product [20, 34]. 555https://github.com/idiap/fast-
transformers/blob/master/fast_transformers/attention/causal_linear_attention.py.
The listing below shows the baseline implementation for reference (where line
76-77 can be computed using pure PyTorch or the Fast Transformers kernel).

##### Micro Benchmark

To complement the end-to-end architecture benchmarks in Section 6, we provide
micro benchmark results for only the linear attention forward pass in Figure
5.

Figure 5: Time (ms) for different ways of computing the Taylor linear
attention forward pass — using Pure PyTorch (shown in the Listing and
introduced in [13]), Fast Transformers kernel (as indicated in the listing)
[34, 24], or our Based kernel (Algorithm 1). (Left) Varying the batch size at
fixed sequence length $1024$. (Right) Varying the sequence length at fixed
batch size $4$. (All) Benchmarking uses $16$ feature dimension, $16$ heads,
$64$ head dimension, and focuses on the numerator of the linear attention.
Each point represents the median across $10$ iterations is measured on a
single NVIDIA H100 GPU. Lines terminate on out-of-memory errors.

⬇

1from einops import rearrange

2import torch

3from torch import nn

4

5class TaylorExp(nn.Module):

6 """

7 Feature map to compute 2nd-order Taylor approx. of exp(q^T k / sqrt(d))

8 """

9

10 def __init__(self, input_dim, head_dim_idx, temp=None, eps=1e-12):

11 super().__init__()

12

13 self.input_dim = input_dim

14 self.head_dim_idx = head_dim_idx

15 self.temp = 1.0 if temp is None else temp

16 self.eps = eps

17

18 self.r2 = math.sqrt(2)

19 self.rd = math.sqrt(self.input_dim)

20 self.rrd = math.sqrt(self.rd)

21

22 def forward(self, x: torch.Tensor):

23 # Get 2nd-order terms (rearrange(x * x), ’... m n -> ... (m n)’)

24 x2 = (x.unsqueeze(-1) * x.unsqueeze(-2)).flatten(start_dim=-2) / self.r2

25 term1 = torch.ones(x[..., :1].shape).to(x.device)

26 term2 = x / self.rrd

27 term3 = x2 / self.rd

28 terms = [term1, term2, term3]

29 return torch.cat(t for t in terms), dim=self.head_dim_idx)

30

31

32class TaylorLinAttn(nn.Module):

33 def __init__(self):

34 super().__init__()

35 self.d_model = d_model

36 self.feature_dim = 16

37 self.num_heads = 16

38 self.num_key_value_heads = 16

39 self.head_dim = self.d_model // self.num_key_value_heads

40 self.eps = 1e-12

41

42 feature_map_kwargs = {

43 "input_dim": self.feature_dim,

44 "head_dim_idx": -1,

45 "eps": 1e-12,

46 }

47 self.feature_map = TaylorExp(**feature_map_kwargs)

48 self.proj_q = nn.Linear(

49 self.d_model, self.feature_dim * self.num_heads, bias=False

50 )

51 self.proj_k = nn.Linear(

52 self.d_model, self.feature_dim * self.num_heads, bias=False

53 )

54 self.proj_v = nn.Linear(

55 self.d_model, self.num_key_value_heads * self.head_dim, bias=False

56 )

57 self.proj_o = nn.Linear(

58 self.num_heads * self.head_dim, self.d_model, bias=False

59 )

60

61 def forward(self, hidden_states: torch.Tensor, *args, **kwargs):

62 b, l, _ = hidden_states.size()

63 q = self.proj_q(hidden_states)

64 k = self.proj_k(hidden_states)

65 v = self.proj_v(hidden_states)

66 q = q.view(b, l, self.num_heads, self.feature_dim).transpose(1, 2)

67 k = k.view(b, l, self.num_key_value_heads, self.feature_dim).transpose(1,
2)

68 v = v.view(b, l, self.num_key_value_heads, self.head_dim).transpose(1, 2)

69

70 # Linear attention

71 q, k = self.feature_map(q), self.feature_map(k)

72 q, k, v = q.unsqueeze(-2), k.unsqueeze(-2), v.unsqueeze(-1)

73

74 # Compute attention causal (alternatively use the Fast Transformers kernel)

75 num = (q * (k * v).cumsum(dim=2)).sum(dim=-1)

76 denom = (q * k.cumsum(dim=2)).sum(dim=-1) + self.eps

77 y = (num / denom)

78

79 y = rearrange(y, "b h l d -> b l (h d)")

80 y = self.proj_o(y)

81 return y

Listing 1: PyTorch implementation of Taylor linear attention. Algorithm 1
Computing the $0^{th}$ ($T0$), $1^{st}$ ($T1$), $2^{nd}$ ($T2$) Order Taylor
Linear Attention Terms

Input projected hidden states $q,k,v\in\mathbb{R}^{N\times d}$.

Output $y=T0+T1+T2\in\mathbb{R}^{N\times d}$

Parallelize into $\mathrm{batch}\times\mathrm{heads}$ parallel computations,
with $\mathrm{n_{warps}}=8$ warps per block.

Within a block:

Define tile size $T$ $\triangleright$ $T=16$ in Based

Define $\mathrm{n_{tiles}}=\frac{N}{T}$ $\triangleright$ Block along the
sequence dimension

Define $\mathrm{n_{blocks}}=\mathrm{n_{tiles}}/\mathrm{n_{warps}}$
$\triangleright$ Block along the number of warps

Define $\mathrm{tic}=0$, $\mathrm{toc}=1$ $\triangleright$ Flags for
asynchronous data loading

Create SRAM buffers $B_{q}$, $B_{k}$ (Size $2\times\mathrm{n_{warps}}\times
T\times T$) and $B_{v}$ (Size $2\times\mathrm{n_{warps}}\times T\times 4T$)

Create SRAM buffers $\mathrm{A0},\mathrm{A1},\mathrm{A2}$ (Size
$\mathrm{n_{warps}}\times T\times 4T$) for storing interim. results for
$T0,T1,T2$ as warps process the sequence

Create SRAM buffers $total_{A0}$ and $total_{A1}$ to hold cumulative (“KV”)
state corresponding to $T0,T1$

Create SRAM buffers $y$ of (Size $\mathrm{n_{warps}}\times T\times 4T$) for
storing the final output

Create register fragments
$\mathrm{q_{a}},\mathrm{q_{b}},\mathrm{k_{a}},\mathrm{k_{b}},\mathrm{q_{frag}},\mathrm{k_{frag}}$,
$\mathrm{qk_{accum}}$ of size $16\times 16$. We create register fragments
$\mathrm{v_{frag}},\mathrm{a0_{frag}}$, $\mathrm{a1_{accum}}$, $A2_{0}$,
$A2_{1}$, $\mathrm{qA2_{accum}}$, $\mathrm{o_{accum}}$ of size $16\times 64$.
These fragments are for holding data during in-register computation.
Initialize the fragments to $0$.

Each warp loads initial tiles $B_{q}[\mathrm{tic}][\mathrm{warpid}]\leftarrow
Q_{t},B_{k}[\mathrm{tic}][\mathrm{warpid}]\leftarrow K_{t}$ and
$B_{v}[\mathrm{tic}][\mathrm{warpid}]\leftarrow V_{t}$ $\triangleright$ HBM
into SRAM

for $\mathrm{cur_{block}}\in[0..\mathrm{n_{blocks}-1}]$;
$\mathrm{tic}=0\oplus=1$, $\mathrm{toc}\oplus 1$ do $\triangleright$ XORs
$\mathrm{tic}$ and $\mathrm{toc}$ to toggle.

Warp loads $B_{q}[\mathrm{toc}][\mathrm{warpid}]\leftarrow Q_{t}$ for
$\mathrm{cur_{block}}+1$ $\triangleright$ HBM to SRAM

Warp loads $B_{k}[\mathrm{toc}][\mathrm{warpid}]\leftarrow K_{t}$ for
$\mathrm{cur_{block}}+1$

Warp loads $B_{v}[\mathrm{toc}][\mathrm{warpid}]\leftarrow V_{t}$ for
$\mathrm{cur_{block}}+1$

Warp loads $\mathrm{q_{frag}}\leftarrow q[\mathrm{tic}][\mathrm{warpid}]$
$\triangleright$ SRAM into register

Warp loads $\mathrm{k_{frag}}\leftarrow k[\mathrm{tic}][\mathrm{warpid}]$

Warp loads $\mathrm{v_{frag}}\leftarrow v[\mathrm{tic}][\mathrm{warpid}]$

Compute the warp-local cumulative sum on
$\mathrm{v_{frag}}\rightarrow\mathrm{a0_{frag}}$. $\triangleright$ T0
computation

Add the running $A0$ to the current $\mathrm{a0_{frag}}$

Compute $\mathrm{q_{frag}}\mathrm{k_{frag}}^{T}$ (attention) and make it
causal and store in a $\mathrm{qk_{accum}}$ $\triangleright$ T1 computation

Compute $\mathrm{qk_{accum}}\mathrm{v_{frag}}\rightarrow\mathrm{o_{accum}}$
$\triangleright$ Store causal $qk^{T}v$

Warps store
$\mathrm{k_{frag}^{T}}\mathrm{v_{frag}}\rightarrow\mathrm{a1_{accum}}$ and
write $\mathrm{a1_{accum}}\rightarrow A1[\mathrm{warpid}$] $\triangleright$
Register to SRAM

Compute cumulative sum over $A1$ in SRAM, updating $A1$ entries

Warps read $A1$ tiles back to registers $\triangleright$ Each warp now
contains its preceeding $A1$

Warps multiply the values in register with $\mathrm{q_{frag}}$ to update
$\rightarrow\mathrm{o_{accum}}$ $\triangleright$ Add in T1 to the running
result

Update $\mathrm{a0_{frag}}\rightarrow\mathrm{o_{accum}}$ $\triangleright$ Add
in T0 to the running result

Square $\mathrm{qk_{accum}}$, multiply with $\mathrm{v_{frag}}$ and add
$\rightarrow\mathrm{o_{accum}}$ $\triangleright$ Add in diagonal T2 to the
running result

Sum the values of $\mathrm{o_{accum}}$ into $y[\mathrm{warpid}]$

for $\mathrm{block}$ in $\mathrm{n_{warps}}$ iterations do $\triangleright$
Remaining T2 computation; Assumes feature dimension $16$

Each of $8$ warps copies the same slice of $q[\mathrm{tic}][\mathrm{warpid}]$
to $2$ registers $q_{a}$, $q_{b}$

Each thread $j$ in the warp computes $q_{a}[:,2j]q_{a}$ for dimension $2j$,
and for $2j+1$ (and for $q_{b}$). Together the threads compute the $256$
elements resulting from the second order outer product in the feature map.

Each warp stores two slices of $A2$: $A2_{0}$ and $A2_{1}$ $\triangleright$
Partitioning the large $A2$ across warp registers

Accumulate both $q_{a}A2_{0}$ and $q_{b}A2_{1}$
$\rightarrow\mathrm{qA2_{accum}}$

Warp writes $\mathrm{qA2_{accum}}\rightarrow A2[\mathrm{warpid}]$
$\triangleright$ Register to SRAM

Sum results across all in $A2[\mathrm{warpid}]$ and store the sum in
$y[\mathrm{block}]$ $\triangleright$ Add in T2

Each of $8$ warps copies the same slice of $k[\mathrm{tic}][\mathrm{block}]$
to $2$ registers $k_{a}$, $k_{b}$ $\triangleright$ KV state update

Square $k_{a}$ and $k_{b}$

Each of the $8$ warps loads $v[\mathrm{tic}][\mathrm{block}]$ to
$\mathrm{v_{frag}}$ in register

Multiply $k_{a}$ and $\mathrm{v_{frag}}$, $k_{b}$ and $\mathrm{v_{frag}}$ and
accumulate the results into $A2_{0}$ and $A2_{1}$, the two in-register slices
of $A2$ for the warp, respectively

End. Store $y$. Optionally store $A0$, $A1$, $A2$ (comprising the “KV state”)
for generation. $\triangleright$ SRAM to HBM

##### Algorithm

Here we revisit the key equations we aim to compute and then describe
Algorithm 1 in detail.

Objective First recall from Section 4:

$\bm{o}_{i}=\sum_{j=1}^{i}\frac{\phi(\bm{q}_{i})^{\top}\phi(\bm{k}_{j})\bm{v}_{j}}{\phi(\bm{q}_{i})\sum_{j=1}^{i}\phi(\bm{k}_{j})}=\frac{\phi(\bm{q}_{i})\sum_{j=1}^{i}\big{(}\phi(\bm{k}_{j})^{\top}\bm{v}_{j}\big{)}}{\phi(\bm{q}_{i})\sum_{j=1}^{i}\phi(\bm{k}_{j})}$
(5)

where $q_{i}$ reflects the $i^{th}$ of $N$ total tokens in the sequence and
every query attends to every past key in $\mathcal{O}(Nd^{2})$ time and space
complexity for embedding dimension $d$.

To approximate $\exp(\bm{q}_{i}^{\top}\bm{k}_{j}/\sqrt{d})$, we use the
$2^{\text{nd}}$-order Taylor series feature map, picking
$\phi:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d^{2}}$ such that

$\phi(\bm{q}_{i})^{\top}\phi(\bm{k}_{j})=1+\bm{q}_{i}^{\top}\bm{k}_{j}+\frac{(\bm{q}_{i}^{\top}\bm{k}_{j})^{2}}{2}$
(6)

In this section, we will refer to $q_{i}$ as a tile of data (e.g. of $16$
tokens) instead of as a single token since the hardware operates on chunks of
data in parallel.

Algorithm description In Algorithm 1, we allow each thread block to compute
the result for a particular $(\mathrm{batch},\mathrm{head})$ input. Within the
thread block, we use $8$ warps / workers to produce the result. We initialize
data structures $B_{q},B_{k},B_{v}$ in SRAM and
$\mathrm{q_{a}},\mathrm{q_{b}},\mathrm{k_{a}},\mathrm{k_{b}},\mathrm{q_{frag}},\mathrm{k_{frag}},\mathrm{v_{frag}}$
in register to hold chunks or tiles of the $q,k,v$ inputs. We initialize data
structures $A0,A1,A2$ in SRAM and $\mathrm{a0_{frag}}$, $\mathrm{a1_{accum}}$,
$\mathrm{qA2_{accum}}$ in register to hold computation for the running $KV$
state for the $0^{th},1^{st},2^{nd}$ order Taylor polynomial terms.

We partition the computation along the sequence dimension into
$\mathrm{n_{blocks}}$, where in each loop from $1$ to $\mathrm{n_{blocks}}$,
the warps load the next $8$ chunks into fast memory. Note that for $2048$
sequence length and $8$ warps, $16$ tile size, we end up with
$\mathrm{n_{tiles}}=128$ and $\mathrm{n_{blocks}}=16$. In each iteration, each
warp loads in $16\times 16$ tiles of $q,k$ and $16\times 64$ tiles of $v$,
where $16$ indicates a chunk of $16$ tokens along the sequence dimension and
$16,64$ are the feature and head dimensions respectively. Once tiles are
streamed in, we do not need to reuse them, which is key to the efficiency of
linear attention.

Zeroeth order Taylor terms: During the computation, for the $0^{th}$ term in
the Taylor polynomial, $q,k$ are $1$ after we apply the feature map (Equation
6). Therefore, computing a cumulative sum over $q(k^{T}v)$ reduces to
maintaining a cumulative sum of $v$ as we iterate across the sequence.

First order Taylor terms: Next we consider the $1^{st}$ order terms. On-
diagonal: First consider the on-diagonal blocks, e.g. with respect to tiles
$q_{i},k_{i},v_{i}$. For these, we simply multiply $q^{T}k$, masking (making
it causal), and then multiplying with $v$, following the order of operations
in standard attention. This makes it easy to apply the masking ($0$ out non-
causal elements). Now each warp contains a local result for its set of on-
diagonal tiles of $q_{i},k_{i},v_{i}$. Off-diagonal: However, we need to
obtain a global cumulative sum where $(q_{i}^{T}k_{j})v_{j}$ depends on all
$j\in[1..i]$ (Equation 5). Each warp is therefore missing values for tiles
$j\in[1..i-1]$. To incorporate this computation, we will now compute the
cumulative $KV$ hidden state for the warp up until $i-1$ and multiply this
with the local tile of $q$ (i.e. $\mathrm{q_{frag}}$). To accomplish this,
note in Algorithm 1, we multiply $\mathrm{k_{frag}}^{T}$ and
$\mathrm{v_{frag}}$ to compute local tiles of the hidden state, local to each
warp, in thread register. To perform the global cumulative sum across the $8$
warps’ local results, we write from registers (thread specific) to $A1$ in
SRAM (shared across warp threads). After computing the global cumulative sum
in shared memory, each warp loads back the $KV$ state (in $A1$) into its
registers such that it contains all the preceeding $KV$ (history) for tiles
$[1..i-1]$. We then multiply the local $\mathrm{q_{frag}}$ in register with
this $KV$ state to update the final output for the $1^{st}$ up until the
current $\mathrm{n_{blocks}}$. Note that we maintain the running $KV$ state
corresponding to the $1^{st}$ order term in $A1$ shared memory for the next
iteration along $\mathrm{n_{blocks}}$.

Second order Taylor terms: We finally need to compute the $2^{nd}$ order term.
Similar to the $1^{st}$ order term, we’ll consider On-diagonal: We can
leverage the computation from above. We’ll square the causal $(qk^{T})^{2}$
from above and multiply with $\mathrm{v_{frag}}$ to obtain the portion of the
$2^{nd}$ order term corresponding to the on-diagonal tiles
$q_{i},k_{i},v_{i}$. Off-diagonal: Again, we also need to compute the result
with respect to tiles $[1..i-1]$.

  * •

Partitioning KV hidden state for $2^{nd}$ order Because the hidden state for
the second order term is large ($\mathcal{O}(d^{2}D)$ in feature dimension $d$
and head dimension $D$) and warps have a limited number of registers, we slice
its storage across the registers of the $8$ warps. Considering the the
$16^{2}\times 64$ ($d^{2}\times D$) hidden state (stored in $A2$ SRAM in
Algorithm 1), we divide this into $16$ slices along the sequence dimension and
let each of the $8$ warps handle $2$ of the $16\times 64$ slices (stored in
$A2_{0},A2_{1}$ fragments in thread registers in Algorithm 1). Warp $i$ will
maintain slices $2i$ and $2i+1$ in two registers per thread.

  * •

Computing output for $2^{nd}$ order Each warp $i$ loads in one tile of $q_{i}$
into $2$ registers. We will use the $32$ threads in the warp to compute the
$256$ outer product terms for each token computed by the Taylor $2^{nd}$ order
term (for feature dimension $16$).

Next, the threads multiply these $256$ terms with the running $A2_{0}$ and
$A2_{1}$ slices. The results for the two slices are summed in register and
then stored in SRAM ($A2[\mathrm{warpid}]$). Since $o_{i}$ is ultimately the
sum of $q_{i}$ terms multiplied with all slices of $A2$ (Equation 5), we then
sum the results from all the warps together (which hold the remaining slices
of $A2$) and store the result in $y[\mathrm{block}]$. We can think of
$y[\mathrm{block}]$ as holding the result up until the
$(8\times\mathrm{cur_{block}}+\mathrm{block})$ tile of tokens (note $8$ is
because in each increment of $\mathrm{cur_{block}}$, the $8$ warps handle $8$
different tiles of the sequence).

  * •

Updating the $KV$ state: For $\mathrm{block}=i$, we load in $k[i],v[i]$ tiles
of size $16\times 16$ and $16\times 64$ respectively to registers
$k_{a},k_{b},\mathrm{v_{frag}}$. We compute the $256$ outer product terms on
$k[i]$ using the $32$ threads, multiply with $\mathrm{v_{frag}}$, and store
the result in the $A2_{0},A2_{1}$ running state.

The final result in $y$ is summed into the output to complete the $2^{nd}$
order computation.

###  B.2 Next Token Prediction

During next token prediction in generation, we use IO-aware algorithms for the
expensive $KV$ state update in Taylor linear attention and for the sliding
window attention computation.

####  B.2.1 Taylor linear attention

The KV update in PyTorch is provided in the following listing. In Figure 6 we
benchmark the speed of the PyTorch implementation against our kernel.

Figure 6: Time (ms) for computing the Taylor linear attention recurrent update
using Pure PyTorch (shown in the Listing and introduced in [13]) vs. our Based
kernel (Algorithm 2). Benchmarking uses $16$ feature dimension, $16$ heads,
$64$ head dimension, and focuses on the numerator of the linear attention.
Each point represents the median across $10$ iterations is measured on a
single NVIDIA H100 GPU.

⬇

1from einops import rearrange

2import torch

3from torch import nn

4

5def step(self, kv_state: torch.Tensor, k_state: torch.Tensor, q:
torch.Tensor, k: torch.Tensor, v: torch.Tensor):

6 """

7 Compute linear attention with recurrent view

8 -> Assume q.shape is (b, h, 1, D); k and v.shape are (b, h, l, d), where D
is the dimension after applying the feature map and d is the head dimension.

9 """

10 b, h, l, d = q.shape

11 assert l == 1, f’q.shape is {q.shape} but should be ({b}, {h}, 1, {d})’

12 # Expand dims for broadcasting to compute linear attention

13 q, k, v = q.unsqueeze(-2), k.unsqueeze(-2), v.unsqueeze(-1)

14

15 kv_state += k[:, :, -1:] * v[:, :, -1:]

16 k_state += k[:, :, -1:]

17

18 # Compute linear attention

19 num = (q * kv_state).sum(dim=-1)

20 y = num / ((q * k_state).sum(dim=-1) + self.eps)

21

22 y = rearrange(y, ’b h l d -> b l (h d)’).to(q.dtype)

23 return self.dropout(self.out_proj(y))

Listing 2: PyTorch implementation of Taylor linear attention KV update
Algorithm 2 Computing $KV$ State Updates

$KV_{t-1}$ state $\in\mathbb{R}^{Hd^{\prime 2}d}$, at time $t$. Featurized
$q,k\in\mathbb{R}^{B\times H\times 1\times D}$ and $V\in\mathbb{R}^{B\times
H\times 1\times d}$, for $d$ as the head dimension (e.g. $64$) and $D$ as the
expanded feature map dimension (e.g. $273=1+16+16^{2}$ for feature dim $16$).
To be hardware-friendly, we let $D=320$ (s.t. $320\mod 64=0$) via padding.

Updated $KV_{t}$ state.

Parallelize into $\mathrm{batch}\times\mathrm{heads}$ parallel computations,
with $\mathrm{n_{warps}}=8$ warps per block.

Within a block:

Define $\mathrm{n_{threads}}=\mathrm{n_{warps}}\times 32$ $\triangleright$
Assuming $32$ threads per warp

Define $\mathrm{buffer_{size}}=\mathrm{n_{warps}}\times 8\times d$

Define $\mathrm{total_{batches}}=\frac{D}{\mathrm{n_{warps}\times 8}}$
$\triangleright$ E.g. $\mathrm{total_{batches}}=5$ if $D=320$; For $k$,
$\frac{320}{5}=64$ values per batch

Define $\mathrm{tic}=0,\mathrm{toc}=1$

Create SRAM buffer $B_{q}$ (Size $D$) for $q$

Create SRAM buffer $B_{k}$ (Size $D$) for $k$

Create SRAM buffer $B_{v}$ (Size $d$) for $V$

Create SRAM buffer $B_{kvs}$ (Size $2\times\mathrm{buffer_{size}}$) for
storing blocks of $\mathrm{kv_{state}}$

Create SRAM buffer $\mathrm{o}$ (Size $d$) for output.

Create SRAM buffer $\mathrm{A}$ (Size $\mathrm{n_{warps}}\times d$) for
intermediate computation

Create register buffer $\mathrm{v_{reg}}$ (Size $2$) to store $V$ data

Create register $\mathrm{A_{reg}}$ (Size $2$) for intermediate computation

Warps load $B_{q}\leftarrow q$ $\triangleright$ HBM to SRAM; Load all $D=320$
elements of $q$

Warps load $B_{k}\leftarrow k$

Warps load $B_{v}\leftarrow V$

Warps load chunk $B_{kvs}[\mathrm{tic}]\leftarrow\mathrm{kv_{state}}$
$\triangleright$ Load $(1\times 64)\times 64$ of the
$(\mathrm{total_{batches}}\times 64)\times 64$ elements in $KV_{t-1}$

Initialize $m=0$

for Threads $j\in[0..31]$; $j<d$; $j+=32,m+=1$ do $\triangleright$ Each thread
holds $2$ values ($d=64$; $32$ threads)

Load $\mathrm{v_{reg}}[m]\leftarrow{v[j]}$ $\triangleright$ SRAM to Register;
Now $v[j]$ is stored in thread $j\mod 32$

for $i\in[0..\mathrm{total_{batches}}]$; $i=i+1$, $\mathrm{tic}\oplus 1$,
$\mathrm{toc}\oplus 1$ do

Loads $B_{kvs}[\mathrm{toc}]\leftarrow$ next batch of $\mathrm{kv_{state}}$
$\triangleright$ Asynchronous loads of next batch

for $j=\mathrm{warpid}$; $j<d$; $j+=\mathrm{n_{warps}}$ do $\triangleright$
Each of the $8$ warps loads $8$ of the $64$ rows of $k$, $q$ in the batch

$\mathrm{k_{val}}\leftarrow B_{k}[i*d+j]$ $\triangleright$ Grab single rows
$q[i]$ and $k[i]$, Broadcast to all threads

$\mathrm{q_{val}}\leftarrow B_{q}[i*d+j]$

$p=B_{kvs}[\mathrm{tic}]+j*d$ $\triangleright$ Point to output rows of
$KV_{t}$; We write $d\times\frac{D}{\mathrm{total_{batches}}}$ sub-matrix for
this batch

Initialize $m=0$

for Thread $k\in[0..31]$; $k<d$; $k+=32,m+=1$ do

$p[k]+=\mathrm{k_{val}}*\mathrm{v_{reg}}[m]$ $\triangleright$ Update running
state by multiplying broadcasted $\mathrm{k_{val}}$ with the full
$\mathrm{v_{reg}}$

$\triangleright$ This updates a $1\times d$ strip of the $d\times D$ full
$KV_{t}$ outer product

$\mathrm{A_{reg}}[m]+=\mathrm{q_{val}}*p[k]$ $\triangleright$ Multiply
$\mathrm{q_{val}}$ with the running state, updating all values in the $1\times
d$ output

Write out new $KV_{t}$ state for this batch: $B_{kvs}[\mathrm{tic}][k]$
$\triangleright$ SRAM to HBM

Initialize $m=0$

for Threads $j\in[0..31]$; $j<d$; $j+=32,m+=1$ do $\triangleright$ Each thread
holds info for $2$ of the $64$ output values

Store $A[\mathrm{warpid}][j]\leftarrow\mathrm{A_{reg}}[m]$ $\triangleright$
Register to SRAM

for Thread $j$; $j<d$; $j+=\mathrm{n_{threads}}$ do $\triangleright$ $d=64$
threads put values from first warp in $n_{j}$

$n_{j}=A[0][j]$ $\triangleright$ Each warp had only computed output values for
a subset of (e.g. $8$) rows of $k$ and $q$

for $w\in[0..\mathrm{n_{warps}}]$ do

Sum the $n_{j}+=A[w][j]$ across $\triangleright$ Need to combine results
across warps

Store $o[j]\leftarrow n_{j}$

Write output $\mathrm{o}$ $\triangleright$ SRAM to HBM

Figure 7: Time (ms) for different ways of computing sliding window attention
next token prediction — using PyTorch, Flash Attention (which supports a
sliding window function), or our inference kernel. Each point represents the
median across query tokens at different token positions in the generation
$\in\\{100,250,500,750\\}$.

####  B.2.2 Sliding window attention

##### Baselines

During training / prefill, we use the Flash Attention sliding window
implementation [11].

Our IO-aware implementation focuses on next token prediction. In the listing
below, we include a Torch reference. Our IO-aware sliding window attention
algorithm is provided in 3. The key insight is to fuse operations in thread
registers to minimize slower SRAM to register data movement.

##### Micro Benchmark

We benchmark key baselines (Torch, Flash Attention-2 [11], and the Based
kernel on an NVIDIA H100 GPU in Figure 7. The benchmark uses window size $64$,
head dimension $64$, and number of heads $16$. We vary the batch size on the
$x$ axis and repeat the median timing across iterations on the $y$ axis. Note
that these timings include only the attention computation and not the time for
updating the $KV$-cache state. These timings also do not include any
processing for Rotary encodings (as shown below).

⬇

1import torch

2from torch import nn

3

4"""

5b: batch size

6h: number of heads

7n: sequence length

8d: head dimension

9

10w: window size

11

12qw: b x h x 1 x d

13kw: b x h x w x d

14vw: b x h x w x d

15"""

16

17w = torch.einsum("bhod, bhnd-> bhn",qw, kw)

18a = torch.nn.functional.softmax(w, dim=-1)

19result = torch.einsum("bhn,bhnd->bhd", a, vw)

Listing 3: PyTorch implementation of Sliding Window

Algorithm 3 Sliding window generation

$KV_{t-1}$ state $\in\mathbb{R}^{Hwd}$, at time $t$ and projected hidden
states $q,k,v\in\mathbb{R}^{B\times H\times 1\times d}$, for $H$ heads, head
dimension $d$, sliding window size $w$, and batch size $B$.

Updated $KV_{t}$ state.

Parallelize into $\mathrm{batch}\times\mathrm{heads}$ parallel computations,
with $\mathrm{n_{warps}}=4$ warps per block.

Within a block:

Define tile size $T$ $\triangleright$ $T=16$ in Based

Define $\mathrm{n_{threads}}=\mathrm{n_{warps}}\times 32$ $\triangleright$
Assuming $32$ threads per warp

Create SRAM buffers $B_{k}$ and $B_{v}$ (Each of size $4T\times 4T$) to hold
$k,v$. $\triangleright$ Assumes $4T=64$ is the $w$, $d$

Create SRAM vector $B_{q}$ (Size $1\times 4T$) to hold $q$ during the kernel
execution. $\triangleright$ Single query, assume $d=64$

Create SRAM vector $B_{w}$ (Size $1\times 4T$) of type float for intermediate
attention computation.

Create SRAM vector $B_{o}$ (Size $1\times 4T$) to hold the output.
$\triangleright$ Single output, assume $d=64$

Create SRAM buffers $\mathrm{max}$ and $\mathrm{sum}$ (Each of
$\mathrm{workers}$ by float size).

Create register fragments $\mathrm{q_{reg}}$, $\mathrm{k_{reg}}$,
$\mathrm{v_{reg}}$ to hold data during fused computation in-register.

Create register fragments $\mathrm{w_{reg}}$ (size $1\times 4T$) and
$\mathrm{{wv}_{reg}}$ (size $4T\times 1$) to store intermediate computation
in-register.

Create register fragment $\mathrm{o_{reg}}$ (size $4T\times 1$) to store
output in-register.

Loads $B_{k}\leftarrow k$ using $\mathrm{n_{threads}}$; $B_{v}\leftarrow v$
using $\mathrm{n_{threads}}$; $B_{q}\leftarrow q$ using one warp.
$\triangleright$ HBM to SRAM

Loads $\mathrm{q_{reg}}\leftarrow B_{q}$. $q$ gets broadcasted to all warps.
$\triangleright$ SRAM to Register

Loads $\mathrm{k_{reg}}\leftarrow B_{k}[\mathrm{warpid}]$. Each warp gets
$T\times 4T$ of the $4T\times 4T$ in $B_{k}$ (i.e. a column).

Loads $\mathrm{v_{reg}}\leftarrow B_{v}[\mathrm{warpid}]$. Each warp gets
$T\times 4T$ of the $4T\times 4T$ in $B_{v}$ (i.e. a column).

Initialize $\mathrm{w}_{reg}$ to zero

$\mathrm{w_{reg}}\leftarrow\mathrm{q_{reg}}\mathrm{k_{reg}}$ $\triangleright$
Matrix-vector (GEMV) multiplication

Initialize float $m=-\infty$ for the max $\triangleright$ Obtain the max
across tiles for Softmax

Update $m\leftarrow\max(\mathrm{w_{reg}})$ with the max from the local data

$\textrm{max[warpid]}\leftarrow m$ for all warps to access

Iterate over $\mathrm{n_{warps}}$ entries in max buffer to compute the global
max of $\mathrm{w_{reg}}$

Put global max back into each warp’s $m$ float

Initialize float $s=0$ for the sum $\triangleright$ Obtain the sum across
tiles for Softmax

Update $s\leftarrow\mathrm{sum}(\mathrm{w_{reg}})$ with the sum from the local
data

$\textrm{sum}[\mathrm{warpid}]\leftarrow s$ for all warps to access

Iterate over $\mathrm{n_{warps}}$ entries in sum buffer to compute the global
sum of $\mathrm{w_{reg}}$

Put global sum back into each warp’s $s$ float

$\mathrm{w_{reg}}\leftarrow\mathrm{w_{reg}}-m$ $\triangleright$ Start
attention computation in register

$\mathrm{w_{reg}}\leftarrow\exp(\mathrm{w_{reg}})$

$\mathrm{w_{reg}}\leftarrow\frac{\mathrm{w_{reg}}}{s}$

$B_{w}[\mathrm{warpid}]\leftarrow\mathrm{w_{reg}}$ $\triangleright$ Register
to SRAM; storing for the slice of $k$

$\mathrm{wv_{reg}}\leftarrow B_{w}$ $\triangleright$ SRAM to Register. Warp
loads entirety of $B_{w}$; all slices

Initialize $\mathrm{o_{reg}}$ to zero.

$\mathrm{o_{reg}}\leftarrow\mathrm{wv_{reg}}\mathrm{v_{reg}}$ $\triangleright$
Matrix-vector (GEMV) multiplication

Write $\mathrm{o_{reg}}$ to global memory $\triangleright$ Register to SRAM,
SRAM to HBM

##  Appendix C Extended Architecture Details

In this section, we describe two additional architectural details that can
enable small improvements in language model perplexity. We emphasize, however,
that the combination of Taylor linear attention and tcWindow layers alone is
sufficient to come within $0.1$ perplexity points of our best models using
these additional components (Table 4).

##### Convolution.

We find that replacing some of the linear attention and tcWindow layers with
gated convolutions [9, 1, 72] enables small improvements in language modeling
performance. We use BaseConv layers [1] with short convolutions and a SilU
non-linearity [73]. By keeping the convolutions short (e.g. width 3), we keep
the recurrent state size low and improve throughput. The projections expand
the dimensionality by a factor $c=4$. Note, this layer is nearly identical to
a Mamba block without the SSM [5].

$\displaystyle\bm{y}$
$\displaystyle:=(\underbrace{\left(\bm{u}\cdot\bm{W}_{1}+\bm{b}_{1}\right)}_{\mathclap{\textbf{Linear
Projection}}}\odot\sigma\underbrace{\left(\bm{h}\ast\bm{u}\cdot\bm{W}_{2}+\bm{b}_{2}\right)}_{\mathclap{\textbf{Convolution}}})\cdot\bm{W}_{3}+\bm{b}_{3}\quad$
(7)

where $\bm{u}\in\mathbb{R}^{N\times d}$ is a projected input,
$h\in\mathbb{R}^{N\times cd}$ is a learned filter, $\odot$ is the Hadamard
product, and $\bm{W}_{1},\bm{W}_{2}\in\mathbb{R}^{d\times cd}$,
$\bm{W}_{3}\in\mathbb{R}^{cd\times d}$,
$\bm{b}_{1},\bm{b}_{2}\in\mathbb{R}^{cd}$, and $\bm{b}_{3},\in\mathbb{R}^{d}$
define weights and biases of three linear projections.

##### Decay.

Recent recurrent architectures include the use of decay terms, implemented in
a variety of ways [27, 28, 5, 6]. As intuition, decay terms control how much a
token should attend to “recent” tokens vs. “early” tokens in the sequence.
Prior work falls in two categories: using input-independent [27, 28, inter
alia.] or input-dependent [5, 6] decay rates. The latter offers improved
quality, but requires the use of a parallel scan during sequence processing
[5]. Instead, we introduce a coarser input-dependent decay technique for the
linear attention layer, avoiding the parallel scan. We start with decay
matrices that use a fixed decay rate as in prior work and let each of the $h$
heads in the linear attention apply a unique decay schedule to the causal dot
product. We then introduce a learned projection $\bm{W}\in\mathbb{R}^{d\times
h}$ that effectively allows each input to weigh the combination of heads
(decay rates) when producing the output representation.

##  Appendix D Extended Results

Architecture | Params |  |  |  |  |  |  |  |
---|---|---|---|---|---|---|---|---|---
LAMBADA | HellaSwag | PIQA | Arc-E | Arc-C | WinoGrande | Average
Ppl. $\downarrow$ |  Acc. $\uparrow$ |  Acc. Norm. $\uparrow$ |  Acc $\uparrow$ |  Acc $\uparrow$ |  Acc. Norm. $\uparrow$ |  Acc. $\uparrow$ |  Acc. $\uparrow$
Transformer++ (LLaMa) | 360m | 18.39 | 42.52 | 33.48 | 63.98 | 46.04 | 24.49 | 53.99 | 44.08
Transformer (Pythia) | 356m | 25.17 | 37.16 | 31.32 | 63.76 | 44.82 | 23.8 | 51.54 | 42.08
Based | 363m | 21.80 | 38.66 | 33.43 | 64.42 | 45.79 | 24.66 | 51.22 | 43.03
Mamba | 358m | 20.23 | 39.65 | 33.63 | 65.02 | 47.01 | 25.00 | 50.75 | 43.51
H3 | 362m | 57.59 | 23.58 | 30.62 | 63.11 | 45.20 | 23.29 | 50.28 | 39.35
Transformer++ (LLaMa) | 1.33b | 11.12 | 49.10 | 39.29 | 66.16 | 51.68 | 26.19 | 53.43 | 47.64
Based | 1.35b | 12.35 | 46.96 | 39.11 | 66.32 | 50.72 | 26.54 | 50.43 | 46.68
Mamba | 1.32b | 13.11 | 46.13 | 39.41 | 66.38 | 52.36 | 25.94 | 50.83 | 46.84

Table 5: Downstream evaluation of pre-trained language models. The same set of
models as in table 1, all were trained on the same 10 billion tokens drawn
from the Pile [30], evaluated using the LM eval harness by EleutherAI [16]

.

Dataset | Enhancer Cohn | Enhancer Ens | Human Reg. | Non-TATA Promoters | Human OCR Ens.
---|---|---|---|---|---
CNN | 69.5 | 68.9 | 93.3 | 84.6 | 68.0
DNABERT | 74.0 | 85.7 | 88.1 | 85.6 | 75.1
GPT | 70.5 | 83.5 | 91.5 | 87.7 | 73.0
HyenaDNA | 74.2 | 89.2 | 93.8 | 96.6 | 80.9
Transformer++ | 73.4 | 89.5 | 89.9 | 94.4 | 79.5
Mamba | 73.0 | - | - | 96.6 | -
Based | 74.6 | 89.5 | 89.5 | 96.8 | 79.0
Table 6: Downstream evaluation of pre-trained DNA models on GenomicsBenchmarks
[74]. We report top-1 classification accuracy ($\%$) with pretrained models
(Transformer++, Mamba, Based) along with prior reported results in [41]. We
find the similar quality-matching in pretraining transfers to downstream
tasks. Modern architectures are also able to achieve state-of-the-art results
on the classification tasks.

##### Downstream Language Results

To further evaluate Based’s performance in language modeling, we evaluate the
PILE-pretrained models on several downstream tasks. We use the same protocol
as [5], utilizing the LM evaluation harness by EleutherAI [16]. In particular,
we use the following set of metrics and tasks:

  * •

LAMBADA (perplexity and accuracy) [75]

  * •

HellaSwag (normalized accuracy) [76]

  * •

PIQA (accuracy) [77]

  * •

ARC-challenge (normalized accuracy) and, separately, the easy subset ARC-easy
(accuracy) [78]

  * •

WinoGrande (accuracy) [79]

Normalized accuracy refers to accuracy normalized by sequence length and is
used to maintain the equivalent setting to [5]. We report results in Table 5.
For both 360 million and 1.3 billion parameter models, Based performs
competitively with recent and state-of-the art architectures, including Mamba
and modern Transformers (Transformer++ (LLaMa)).

##### Downstream DNA Classification

We further evaluate how different architectures compare for DNA modeling. We
use the pretrained models described in Section 6.1, and test how well they
transfer to DNA sequence classification on a popular benchmark
(GenomicBenchmarks) [74] used to evaluate recent DNA language models [41].

We compare our HG38-pretrained Transformer++ and Based models in Table 6. We
find similar performance across tasks, indicating that prior matching in
quality during pretraining transfers to downstream classification. For
reference, we also include results from [41]. Although not directly comparable
to due differences in tokenization, the evaluations suggest Based can perform
strongly on different modalities, and that recent sequence modeling
architectures are also able to outpeform or compete with prior state-of-the-
art on evaluated DNA tasks.

##  Appendix E Experimental Details

###  E.1 Language Model Pretraining

We use A100 80GB Nvidia GPUs to run all experiments. We use training
infrastructure closely adapted from the FlashAttention code base:
https://github.com/Dao-AILab/flash-attention/tree/main for all pretraining
runs [11]. The Pile data is tokenized using the GPT2BPETokenizer and all
models see the data in the same order. Here we provide details on the
hyperaparamters and configurations used for training each architecture. We
also provide details on the FLOPs computation.

  * •

Based We train using the specifications in Table 7. Our implementation is
provided here: https://github.com/HazyResearch/based. The initial models were
trained and evaluated using the Fast Transformer CUDA kernels discussed in
Appendix B [34, 24].

  * •

Transformer++ [15] We refer to the modern Llama architecture with Rotary
encodings, RMSNorm and SwiGLU as Transformer++, following prior work [5, 6].
We train using the the specifications in Table 8 using the Flash Attention
training code provided here: https://github.com/Dao-AILab/flash-
attention/tree/main [11].

  * •

Mamba [5] We train using the specifications in Table 9, where the parameters
are sourced from the Appendix of [5]. The implementation is sourced from the
provided reference at https://github.com/state-spaces/mamba.

  * •

Hyena [7] We train using the specifications in Table 10, where the parameters
are sourced from the Appendix of [7]. The implementation is sourced from the
provided reference at https://github.com/HazyResearch/safari.

  * •

H3 [9] We train using the specifications in Table 11. The implementation is
sourced from the provided reference at https://github.com/HazyResearch/safari.

  * •

RWKV [8] We train using the specifications in Table 12 and use the reference
implementation at https://github.com/BlinkDL/RWKV-LM. We specifically evaluate
RWKV-V5.

  * •

Gated Linear Attention (GLA) We train using the specifications in Table 13. We
train following the reference implementation at
https://github.com/berlino/gated_linear_attention.

We give all models the improved Transformer++ recipe (e.g., SwiGLU) as
relevant.

###  E.2 Computing Recurrent State Size

In this section, we provide details on how we compute the size of the
recurrent hidden state for the results described in Section 3.1. We train and
evaluate six sequence mixers on a synthetic associative recall task: attention
[2], sliding window attention [19], Mamba [5], H3 [9], Hyena [7], and Based.
For each, we vary hyperparameters that affect the memory consumption during
inference. We compare how MQAR accuracy varies with the size of the recurrent
hidden state.

#####  Based.

The recurrent state size in Based is determined by the model dimension $d$ and
the size of the hidden dimension after applying the feature map $\tilde{d}$.
The $+1$ accounts for the K-state required for computing the denominator. For
more details on the recurrent view of Based, see 4.

$\text{sizeof}(\bm{s}_{i})=(d+1)\times\tilde{d}$ (8)

In Based, we use the Taylor Exponential feature map after projecting $d$ down
to a smaller dimension $d^{\prime}$. With this approach, recurrent state size
is given by:

$\text{sizeof}(\bm{s}_{i})=(d+1)\times(1+\frac{3d^{\prime}}{2}+\frac{d^{\prime
2}}{2})$ (9)

In our synthetic experiments, we run Based with $d\in\\{48,64,128\\}$ and
$d^{\prime}\in\\{8,16,24\\}$.

##### Attention.

The recurrent state size (i.e. KV-cache size) in attention depends on two
parameters: the model dimension $d$ and the sequence length $N$. The $2$ in
the expression below accounts for the separate storage for keys and values in
the KV-cache.

$\text{sizeof}(\bm{s}_{i})=2\times d\times N$ (10)

In our synthetic experiments we run attention with $d\in\\{64,128\\}$. The
sequence length $N$ is determined by the task, not the model architecture.

##### Sliding window attention.

The recurrent state size in sliding window attention is given by the model
dimension $d$ and the width of the sliding window $k_{\text{sliding}}$. The
$2$ in the expression below accounts for the separate storage for keys and
values in the KV-cache.

$\text{sizeof}(\bm{s}_{i})=2\times d\times\min(N,k_{\text{sliding}})$ (11)

In our synthetic experiment we run sliding window attention with
$d\in\\{128\\}$ and $k_{\text{sliding}}\in\\{8,16,32,64,128,256,512,1024\\}$.

##### Mamba.

The recurrent state size in Mamba is determined by the model dimension $d$ and
the number of heads $h$. The $2$ in the expression below accounts for the
expansion in the Mamba block.

$\text{sizeof}(\bm{s}_{i})=2\times d\times d_{\text{state}}$ (12)

In our synthetic experiments, we run Mamba with $d\in\\{64,128,256\\}$ and
$d_{\text{state}}\in\\{8,16,24\\}$.

##### H3.

The recurrent state size in H3 is determined by the model dimension $d$ and
the number of heads $d_{\text{state}}$.

$\text{sizeof}(\bm{s}_{i})=d\times d_{\text{state}}$ (13)

In our synthetic experiments, we run H3 with $d\in\\{64,128,256\\}$ and
$d_{\text{state}}=\frac{d}{4}$.

##### Hyena.

The recurrent state size in Hyena is determined by the model dimension $d$ and
the number of heads $h$. The $2$ in the expression below accounts for the
separate storage for keys and values in the KV-cache.

$\text{sizeof}(\bm{s}_{i})=d\times N$ (14)

In our synthetic experiments, we run Hyena with $d\in\\{64,128,256\\}$.

###  E.3 Language Model Evaluation

In this section, we provide details on each of the evaluations (columns)
reported in tables 1 and 4.

##### Pile

(Language Modeling). First, we report overall perplexity on the Pile test set
[30]. Then , to understand how much of the perplexity gap is due to recall
capacity, we also evaluate perplexity on two slices (i.e. subsets) of the test
set:

  1. 1.

Associative recall(AR) tokens. Tokens in the final position of a bigram which
previously occured in context, but $\leq 1250$ times in the training data.

  2. 2.

Other tokens. All other tokens.

To construct these slices, we exactly follow the protocol in Arora et al. [1]
and refer the reader to that work for more details. We compute these slices on
the first 16 million tokens in the test set.

##### SWDE

(Information Extraction). The task in the SWDE benchmark is to extract semi-
structured relations from raw HTML websites. For example, given an IMBD page
for a movie (e.g. Harry Potter and the Sorcerer’s Stone) and a relation key
(e.g. release date), the model must extract the correct relation value (e.g.
2001). The SWDE benchmark was originally curated by Lockard et al. [80] for
the task of open information extraction from the semi-structured web. Because
we are evaluating the zero-shot capabilities of relatively small language
models, we adapt the task to make it slightly easier. Our task setup is
similar after to that used in Arora et al. [39].

##### FDA

(Information Extraction). The task is to extract key-value pairs from a set of
PDFs scraped from the FDA website. We use the dataset and labels collected in
[39]. We break apart the documents into chunks of 1,920 tokens. For every key-
value pair that appears in the chunk, we create a zero-shot prompt using the
simple prompt template:
{chunk} \n {key}:
We allow the model to generate a fixed number of tokens after the prompt and
check (with case insensitivity) if the value is contained within the
generation. We report accuracy, the fraction of prompts for which the
generation contains the value.

Below we include one example of a zero-shot prompt for the key-value pair
“Type of Test: Quantitative, colorometric, pyranose oxidase (PROD)”. The