$$
\newcommand{\R}{\mathbb{R}}
\newcommand{\Z}{\mathbb{Z}}
\newcommand{\N}{\mathbb{N}}
\newcommand{\sft}{\text{softmax}}
\newcommand{\List}{\text{List}}
\newcommand{\Seq}{\text{Seq}}
\newcommand{\SeqT}{\text{SeqT}}
\newcommand{\CSeqT}{\text{CSeqT}}
\newcommand{\Dist}{\text{Dist}}
\newcommand{\SM}{\text{SM}}
\newcommand{\Fn}{\text{Fn}}
\newcommand{\Tok}{\text{Tok}}
\newcommand{\Aij}{ A_{[i,j]}}
$$

# Linear Self Attention is Faster

# 0) Introduction
What would happen if you took an exact replica GPT2 using flash attention 2 and made one simple modification: remove all the exponentials that appear in the attention layers (due to the softmax). Essentially, just change all the terms $e^{K_j^T Q_i} \to K_j^T Q_i$. [^1] In the plot below you can see the results for the 4 Transformer scales of the GPT2 paper which range from $120$ million to $1.5$ billions parameters.

[^1]: To avoid dividing by $0$ in the normalization of the atteniton layer we also applied a positive nonlinearity when computing $K, Q$ to ensure all terms $K_j^T Q_i >0$. See section X for all the details.


Well, it looks like we found a way to make GPT2 slighly worse. But there is a catch. Note that the x axis is the number of training updates, not the time elapsed. See what happens if you click on the `time` x_axis? Now our modified transformer is way better at minimizing the loss. This is because the new transformer, running on equivalent hardware, is 20x faster. Look at what hapens for context sizes larger and smaller than 4096. As the context size grows, the speedup only gets more extreme, to the point that for the largest context size we testd, 100k, it's 200x faster.

But how can such a tiny modification have such radical impact on the performance of the architecture? Section X contains all the details and performance analysis, but we can explain the core of the idea quite simply. So let's get on with it.

## 0.1) Linear Self Attention
The inputs to the layer are sequences of $Q_i, K_i, V_i \in \R^d$ of query, key and values, where $i$ ranges from $1$ to $t$, the sequence length. The job of an attention layer is to compute another sequence of outputs $Y_i\in \R^d$. The core of the transformer attention layer is
$$
\hat{Y}_i = \sum_{j=0}^i e^{K_j^T Q_i} V_j
$$

Let's think about the cost of computing $Y_i$. We'll need to perform inner products $K_j^T Q_i$, each requiring $O(d)$ operations. And to compute $Y_i$ we'll need to perform $i$ such inner products, one for each $j\le i$. The inner products dominate, so the cost of computing $Y_i$ is $O(id)$. Since we need to compute the entire output sequence for $0\le i < t$, the total cost will be $O(d + 2d +\cdots + td) = O(t^2d)$.

TODO: explain that the normalization factor is actually important for performance. The experiments we show include it. But normalization doesn't have any impact on the deep ideas behind the reason RSA can be so much faster than classical transformers. So we omit them for now to not clutter the notation. To see details see the full description on section 2.X.

The story doesn't change much for the modified attention layer. It has the equation
$$
Y_i = \sum_{j=0}^i K_j^T Q_i V_j
$$

All the same inner products appear in this linear attention layer, so the cost should also be $O(t^2d)$. The optimized way to implement attention is with specialized kernels but the core of the idea is that you do

def linear_attention(Q, K, V):
    """ TODO: explain shapes of the inputs """
    M = causal_mask(K.shape[1])
    A = (K.T @ Q) * M
    return V @ A

The difference with the standard transformer layer is that we can perform the follwoing factorization:
$$
Y_i = \sum_{j=0}^i K_j^T Q_i V_j = \underbrace{ \left ( \sum_{j=0}^i V_j K_j^T\right )}_{S_i} \; \; Q_i
$$

Where $S_i \in \R^{d\times d}$ can be thought of as a state summarizing all the relevant information up to time $i$. We can further massage the formula into the following recurrent equations
$$
S_i = S_{i-1} + V_i K_i^T \;\;\;\;\;\;\;\; Y_{i} = S_i Q_i
$$

where we assume $S_{-1} = 0$. Let's think about the computational cost of this approach. Every step, when we are trying to compute $Y_i$ we already have access to $S_{i-1}$. We need to compute $V_i K_i^T$ which takes $O(d^2)$ operations, adding it to $S_{i-1}$ is also $O(d^2)$ and the matrix vector multiplication $S_i Q_i$ is also $O(d^2)$. Thus, to compute the entire output sequence it will require $O(td^2)$ operations. The efficient way to implement this is just

def linear_attention(Q, K, V):
    S = cumsum(V.T @ K)
    return Q @ S

We have a choice between two modes of computation! The first one mirrors the intuitions of self attention, each output $Y_i$ is computed by seeing how $Q_i$ interacts with all the previous keys $K_j$ for $0\ge j \ge i$. This **self attention mode** has cost $O(t^2 d)$. The other mode is sequential in nature, there is nothing in it that involves the interaction of every moment on time with all the privious ones. The **sequential mode** has cost $O(td^2)$.

The 1.5B parameter model has $d=64$ and the context size $t$ would tipically be $1024$ or greater. Since $t$ is usually much larger than $d$, one would expect the sequential mode to be preferable over the attention approach which has a factor $t^2$ in the cost. Let's empirically measure this. The plot below shows the times to perform a training update for the 1.5B model with context sizes varying from 1024 to 132k

EXPERIMENTS COMPARING SPEEDS OF THE TWO MODES

It's perhaps surprizing how large the context size needs to be before the sequential approach becomes faster than the attention one. It turns out that attention is really well suited for GPUs, specially when using really well optimized kernels like flash attention. The attention mode is more parallelizable and requires less memory reads and writes. $t$ being much greater than $d$ does mean that the attention mode requires more floating point operations, but that turns out not to be the bottleneck for GPUs. At least it isn't until very large context sizes, $t=64k$ in the experiment above.

Two things are evident here.
* The attention formulation is curcial to fully utilize the power of GPUs
* For large $t$ the quadratic term in $O(t^2d)$ really starts to hurt us. So much so that we reach a point where even the inneficient sequential implementation ends up being faster

But why not get the best of both worlds? The great GPU utilization of attention together with the good growth in $t$ of the sequential mode. Let $c \in \Z^+$ be a positive integer that we call the chunk size. The key idea is that we are only only compute a subset of all states $S_0, S_c, S_{2c}, \cdots$ and we'll use attention of cost $O(c^2d)$ to compute all the in between outputs. For any $i$ find the $n\in \Z$ s.t. $cn < i \le c(n+1)$. Then we can easily see that the following equations are equivalent to the previous ones.
$$
Y_{i} = S_{cn}Q_i + \sum_{j=cn+1}^i K_j^T Q_i V_j
\;\;\;\;\;\;\;\;
S_{c(n+1)} = S_{cn} + \sum_{j=cn+1}^{c(n+1)} V_j K_j^T
$$

TODO: write pseudo code for the hybrid method

We can pick $c$ empirically. For the 1.5B transformer, the optimal value of the chunk size $c$ was $512$. So let's see how this hybrid approach compares with the other two:

EXPERIMENTS COMPARING SPEEDS OF THE THREE MODES

And now you can see why in the initial experiment SARNN was so much faster than a standard transofrmer.

TODO: add discusion of the hypotheis of why transformer is better when the x axis is updates:
* We haven't put a lot of work tunning the LSA architecute. Just used all standard hyperparameters that were tunned for GPT2. We'll push a bit further to see if we can close the gap
* Self attention with the softmax has more powerful representational capacity. In section X we present some theoretical evidence supporting this theory. If the exponential inside the softmax is appoximated with finitely many terms $e^{x} \simeq \sum_{k=0}^n \frac{x^k}{k!}$ then we can implement the classical transformer as an RNN just like we can for the linear self atteniton. The catch is that instead of having a state of size $d^2$, the state has size on the order of $O(d^n)$. So perhaps, when we are comparing a GPT2 vs RSA-GPT2 we are comparing two models with the same number of parameters with very different state sizes.

Regardless of what the gap reason is, the fact RSA is so much faster than attention with the softmax makes it somewhat of a mute point. For large context sizes, the speed up enables one to train a much larger model with the same amount of compute, and the gains from the extra

TODO: address BIG question about our work. Does the ability of RSA to utilize long context get worse as the context size grows? Since we are using datasets that don't contain long term information we aren't really testing that the models are capable of integrating large amounts of long term information. Essentially, it could be that the gap gets larger as the context size grows on tasks where there really is long context signal. Mention that this is ongoing work and we will update with more information once we have it. This is probably the biggest weakness of the work as is.

## 0.2) Sampling costs
But when working with language models, efficinet training is not the only performance consideration. Once the model is trained, what are the costs of using it?

The process of sampling looks something like this: we might have an initial sequence of input tokens $x_1 \cdots x_t$. We pass them to our model, which gives us a distirbution $p$ assigning probabilities to the next token . We sample from this distribution and append the new sample to the sequence and repeat the process with the sequence $x_1 \cdots x_{t+1}$.

During training time, transformers are great because they compute the next token distribuiton $p_1, \cdots, p_{t}$ all at the same time in parallel. This advantage is more than enough to overcome the quadratic costs $O(t^2)$, at least for short context sizes. But when we are sampling we are left with all the disadvantages but can't reap the parallelization benefits due to the sequential nature of the problem.

With the linear self attention, though, we have access to the recurrent formulation. It's has the advantage of lower cost

EXPERIMENT: Experiment comparing speeds of sampling between RSA-GPT2

NOTE: The comparasion is not completely fair. Our KV cache implementation could be improved quite a lot by using specialized kernels while it's unclear if much improvement is possible on the simple unrolling of the RNN. But the pattern will remain the same.

TODO: Make small point that GPT2 is not suited for sampling for arbitrarily long context due to the positional encodings.