Spacetime Chapter: Viewing Attention as an RNN with Quadratic Complexity

时空之章：将Attention视为平方复杂度的RNN 苏剑林 · 2024-03-18 · Mathematics language model RNN attention complexity

In recent years, RNNs have regained the interest of many researchers and practitioners thanks to their linear training and inference efficiency—there's even something of a "renaissance" underway, with representative works including RWKV, RetNet, and Mamba. When RNNs are used for language modeling, their defining characteristic is that each generation step has constant space complexity and time complexity, so over the whole sequence this amounts to constant space complexity and linear time complexity. Of course, everything has two sides: compared to Attention's dynamically growing KV cache, an RNN's constant space complexity often raises doubts about its memory capacity, and it's hard for it to match Attention's performance on long context.

In this article, we show that Causal Attention can be rewritten in the form of an RNN, and that in theory each of its generation steps can also be carried out with $\mathcal{O}(1)$ space complexity (at the cost of an extremely high time complexity, far exceeding quadratic). This suggests that Attention's advantages (if any) are achieved by piling up computation, not by intuitively piling up memory—it is, at its core, just as constant-capacity in memory (a memory bottleneck) as an RNN.

RNNs Beyond Linear Complexity

Supporters of RNNs often make an argument that seems hard to refute: think about whether your brain works like an RNN or like Attention?

Intuitively, RNN inference has constant space complexity, while Attention's KV cache grows dynamically, and given that human brain capacity is finite, one has to admit that from this angle RNNs are indeed closer to the human brain. However, even if it's reasonable to say that brain capacity limits each inference step to constant space complexity, this doesn't mean each step also has constant time complexity. Or put another way: even if a person's per-step time complexity is constant, when processing a sequence of length $L$ a person doesn't necessarily scan the sequence only once (think of "flipping back through a book"), so the total number of inference steps may well exceed $L$, leading to non-linear time complexity.

With this in mind, the author had a sudden idea: could we generally consider RNN models with constant space complexity but non-linear time complexity, in order to make up for a capability that mainstream RNNs lack (such as the "book-flipping" mentioned above)? For a language modeling task, suppose the sample is a b c d e; then the training task is to input a b c d and predict b c d e. A typical RNN looks like the diagram below:

Figure 1: A conventional RNN

The problem with this kind of RNN is that it has no book-flipping ability—once an input is read, it's discarded. Attention's defining trait, by contrast, is that every time it reads a new token, it does a full "flip-through" of the history. Although this approach may raise efficiency concerns, it is undeniably the simplest and most brute-force way to introduce book-flipping capability. And to give RNNs this book-flipping ability, we can simply mimic what Attention does:

Figure 2: An RNN that keeps "flipping back through the book"

Just like Attention, every time a new token is read, the entire history is flipped through again. Of course, one could argue that this isn't really designing a new kind of RNN at all, but merely a new way of using an RNN—simply modifying the input—and this trick can be bolted onto RWKV, Mamba, or anything else. Under this usage, decoding can still be done within constant space complexity, but the time complexity of each inference step grows linearly, so the total time cost is $\mathcal{O}(L^2)$.

Attention Is Also an RNN

In fact, the model represented by Figure 2 is extremely general—so general that Attention itself is merely a special case of it, as shown below:

Figure 3: The RNN corresponding to Causal Attention

Compared with Figure 2, several arrows in Figure 3 are faded, indicating that those connections are actually cut. So Attention is just a special case of Figure 2. Specifically, Attention's computation formula is:

\begin{equation}o_i = \sum_{j=1}^i a_{i,j}v_j = \frac{\sum_{j=1}^i e^{q_i\cdot k_j} v_j}{\sum_{j=1}^i e^{q_i\cdot k_j}}\end{equation}

Clearly, both the numerator and denominator sums can be written in recursive form:

\begin{equation} \begin{pmatrix} y_i^{(t)} \\ z_i^{(t)} \end{pmatrix} = \begin{pmatrix} y_i^{(t-1)} \\ z_i^{(t-1)} \end{pmatrix} + e^{q_i\cdot k_{i-t+1}}\begin{pmatrix} v_{i-t+1} \\ 1 \end{pmatrix}\quad,\quad o_i = \frac{y_i^{(i)}}{z_i^{(i)}} \end{equation}

As far as the author is aware, the earliest work to propose the above formula and use it to optimize Attention computation is Self-attention Does Not Need O(n^2) Memory; the block-matrix version of the above formula is exactly the theoretical basis of Flash Attention, today's mainstream acceleration technique. Since in Self-Attention, Q, K, and V are all obtained from the same input via token-wise operations, the recursive form above can indeed be represented as Figure 3.

Of course, Figure 3 only depicts a single layer of Attention. Multiple layers can naturally be drawn as well, although the connections start to look rather complicated—for instance, the two-layer case is shown below:

Figure 4: The RNN corresponding to two-layer Attention

Constant Space Complexity

As stated at the beginning of this article, a common advantage of RNNs is that inference can be done with constant space complexity and linear time complexity. Since Attention can also be written as an RNN, a natural question arises: does it also enjoy these two advantages in this formulation?

Clearly, since the RNN corresponding to Attention is one whose sequence length has grown to $\mathcal{O}(L^2)$, linear time complexity is out of the question. The only thing worth pondering is whether constant space complexity can be achieved. Most people's first instinct is "no," since it's common knowledge that Attention decoding involves a dynamically, linearly growing KV cache. But that's only the case for the usual, relatively efficient implementation. If we trade time for space without regard for cost, how far can we push the space complexity down?

The answer might be surprising: if we really push the trade-off between time and space to the extreme, the space complexity can indeed be brought down to $\mathcal{O}(1)$!

Actually, this conclusion isn't hard to imagine. First, the single-layer Attention shown in Figure 3 has exactly the same form as an ordinary single-layer RNN, so it's clear that inference can be done with a fixed amount of storage. Next, consider the multi-layer Attention shown in Figure 4—the connections between layers are more complex, so normally one needs to cache the historical K's and V's to compute efficiently. But if we resolutely refuse to store any KV cache, then the K and V fed into every layer, at every step, can be recomputed from scratch from the original input (recomputation). This causes an enormous amount of redundant computation, so the total time complexity will far exceed quadratic complexity—very "environmentally unfriendly"—but the space complexity can indeed be kept at $\mathcal{O}(1)$.

Take two-layer Attention as an example. The second Attention layer takes the output of the first Attention layer as input, and every output of the first Attention layer can be computed within $\mathcal{O}(1)$ space. So as long as we're willing to sacrifice efficiency for recomputation, the second Attention layer also only needs $\mathcal{O}(1)$ space to complete its computation. By the same logic, the third Attention layer uses the output of the second layer as input, the $N$-th layer uses the output of the $N-1$-th layer as input, and since the previous layer can always be computed within $\mathcal{O}(1)$ space via recomputation, every layer—indeed the whole model—can complete its computation within $\mathcal{O}(1)$ space.

This brings us back to the point made at the start of the article: if Attention really does have some advantage over RNNs, it is achieved purely through more computation—intuitively "expanding memory" is just an illusion created by trading space for time. At its core, Attention has exactly the same constant-capacity memory bottleneck as an RNN.

Of course, some readers might think: isn't trading time for space a fairly common technique? Doesn't this seem like a rather unremarkable conclusion? Indeed, trading time for space is common, but it isn't always possible to do. In other words, not every problem can have its space complexity reduced to $\mathcal{O}(1)$ via a time-for-space trade-off—this is a common but non-trivial property.

Reflections on Model Capability

The reason for pointing out this property of Attention is not really to use it for practical inference, but to use it to help us think further about the capacity bottlenecks of Attention.

First, if we really want to be precise about the details, $\mathcal{O}(1)$ is not quite correct; it would be more rigorous to say $\mathcal{O}(L)$, because an RNN with quadratic complexity needs to repeatedly scan the historical sequence, which requires at minimum storing both the original input and the outputs generated so far—that is, at least $L$ integer token IDs need to be stored, and the space this requires is $\mathcal{O}(L)$. If $L$ is large enough, $\mathcal{O}(L)$ will end up larger than $\mathcal{O}(1)$. However, what $\mathcal{O}(1)$ mainly refers to here is the minimum space needed by the LLM's intermediate computation layers—equivalent to the hidden state when viewed as an RNN, with at least (hidden_size num_layers 2) components—whereas the $\mathcal{O}(L)$ space manifests in the input and output. An intuitive analogy is to think of Attention as a computer with unlimited disk space but fixed memory: it continuously reads data from disk, performs computation in memory, and writes the results back to disk.

We know that when memory itself is large and the data being processed is small, we programmers tend to get a bit more "carefree"—we might even load all the data into memory and let the entire intermediate computation proceed without touching the disk at all. Similarly, an LLM trained under the "large model, short sequence" paradigm will tend to rely on the fixed "memory" of scale $\mathcal{O}(1)$ brought about by model scale, rather than the dynamic "disk" brought about by sequence length, because at current LLM scales the former is large enough. SGD will "take the lazy way out" and train the model as if it were a machine with unlimited static memory (because for short sequences, memory is always sufficient). But in reality, the model's static memory is finite, and so for tasks that cannot possibly be completed within $\mathcal{O}(1)$ space, Attention-based models cannot generalize to inputs of arbitrary length either.

Here's an example. Suppose we want to compute the decimal representation $y$ of $2^x$, modeling it with Attention as a conditional model $p(y|x)$, with the training corpus consisting of $\{x,\color{red}{[sep]},y\}$ concatenated together, and computing the loss only on $y$. Note that $y$ here can be uniquely determined by the input $x$, so in principle it should be possible to learn 100% accuracy. But without a chain of thought (CoT) to dynamically extend the sequence length, the model can only implicitly stuff the entire computation process into "memory," which always works for short inputs. In reality, though, memory is finite, while the space required to compute $2^x$ grows as $x$ grows, so there must inevitably exist some sufficiently large $x$ for which the accuracy of $p(y|x)$ cannot reach 100% (even the training accuracy). This is different from the length-extrapolation problem discussed in Transformer Upgrade Path: 16, "Reviewing" Length Extrapolation Techniques—it's not caused by out-of-distribution positional encodings, but rather a capability defect arising from "large model, short sequence" training when there isn't enough CoT guidance.

So why does the current mainstream direction of scaling up LLMs still focus on increasing the model's memory—i.e., increasing hidden_size and num_layers—rather than researching approaches like CoT that increase seq_len? The latter is of course also a mainstream research direction, but the core issue is that if memory becomes the bottleneck, it lowers the model's learning efficiency and generality. It's just like when memory is small and the amount of data is huge: we need to promptly save results to disk and clear memory, which means the algorithm needs to be more elaborate and harder to write, and might even need to be customized for the specific task at hand. So under what circumstances does a memory bottleneck arise? Take LLAMA2-70B as an example: its num_layers is 80 and hidden_size is 8192; multiplying these gives 640K, and multiplying by 2 gives roughly 1M. In other words, once the input length reaches the 1M-token range, LLAMA2-70B's "memory" may well become the bottleneck. Although training LLMs at the 1M-token level is still not easy at present, it is no longer out of reach—for instance, Kimi has already launched a private beta of a model at the 1M level.

So, continuously increasing the model's context length (disk space), to accommodate more input and CoT, while at the same time increasing the model's own scale so that "memory" doesn't become the bottleneck—this has become the dominant theme of current LLM development.

At the same time, this also invalidates an idea the author previously entertained: could we shrink the model size and increase seq_len to achieve results comparable to a large model? The answer is probably no, because a small model has a memory bottleneck, and to make up for it using the "disk" provided by seq_len, every sample would need a sufficiently long CoT—which is even harder than directly training a large model. If seq_len is increased merely through simple tricks like repetition, there's no substantial benefit, since no additional information is introduced. However, if the increase in seq_len is achieved via prefix tuning, it may actually be possible to make up for the gap in space complexity, because the prefix parameters are not computed from the input sequence but are trained separately—this is effectively like plugging in an extra set of "memory sticks," thereby increasing the model's memory.

A Brief Summary to Close

In this article, we examined Attention from the perspective of a quadratic-complexity RNN, and found that it has a constant-space-complexity bottleneck. This shows that Attention, compared to RNNs, doesn't fundamentally add more "memory"—it merely adds a great deal more computation. The existence of this bottleneck suggests that Attention may face theoretical difficulties (insufficient memory) in generalizing to certain tasks over long lengths. How to better guide models to make use of the dynamic "disk" provided by the seq_len dimension may well be the key to overcoming this difficulty.