VQ the Key, and Transformer Complexity Becomes Linear

VQ一下Key，Transformer的复杂度就变成线性了 苏剑林 · 2023-11-09 · Machine Learning quantization encoding gradient attention

Efficient Transformer generically refers to any line of work aimed at reducing the quadratic complexity of Transformers. It started out specifically targeting improvements to attention, and later broader ideas—Fourier transforms, linear RNNs, and so on—were also folded into this category. It has to be said that, in the effort to bring down Transformer's quadratic complexity, all sorts of experts have shown their unique tricks, and a wonderful variety of ingenious ideas have blossomed; I've learned quite a bit of theory from following this literature. And yet, although Efficient Transformer work has been theoretically brilliant, in practice this field has always simmered rather than boiled—there's never really been a model with truly outstanding real-world performance. Now, in today's LLM-dominated landscape, it has gradually faded from view, and, frankly, from my own interest as well.

Recently, however, a paper called Transformer-VQ: Linear-Time Transformers via Vector Quantization made me sit up and applaud. The authors made a remarkably clever observation: simply apply VQ (Vector Quantization) to the Keys of standard attention, and the complexity automatically drops to linear! This linearization approach preserves the form of standard attention, making it a near-perfect bridge between standard attention and linear attention, while retaining as much of standard attention's capability as possible.

The Efficiency Puzzle

As it happens, this blog has been following Efficient Transformer work for quite a while—dating back to a 2019 post covering Sparse Transformer, Born to Save: From Standard Attention to Sparse Attention. Since then, I've written a number of other posts on Efficient Transformer topics, including:

Exploring Linear Attention: Must Attention Have a Softmax?
Performer: Linearizing Attention Complexity via Random Projections
Nyströmformer: A Linearized Attention Scheme Based on Matrix Factorization
The Road to Transformer Upgrades: 3. From Performer to Linear Attention
Linear Transformers Are Probably Not the Model You're Waiting For
FLASH: Perhaps the Most Interesting Efficient Transformer Design in Recent Memory
Google's New Attempt to "Resurrect" RNNs: Can RNNs Shine Again?

And yet, as noted at the outset, despite a considerable amount of Efficient Transformer work—and despite the high hopes placed on it—the field has never produced anything that truly broke into the mainstream. The reasons might be:

1. Many Efficient Transformer methods trade away performance for speed;
2. For many Efficient Transformer methods, the reduced complexity is only theoretical, with little practical speedup;
3. Some Efficient Transformer methods are hard to adapt for training causal LMs, leaving them with no place in today's LLM-dominated world;
4. The advent of Flash Attention showed that even standard Transformers still have plenty of room for speedup.

VQ It

So why does Transformer-VQ have the potential to break out?

Simply put, Transformer-VQ "clusters" the sequence of Key vectors in attention and approximates each original vector by the center of the cluster it belongs to. This alone brings attention's complexity down to linear. In other words, Transformer-VQ only changes the form of the Key—everything else remains (in theory) completely unchanged. This makes it a linearization scheme with an extremely small footprint of change to attention, and it also makes very clear exactly where the precision loss from linearization comes from (namely, the gap between the cluster center and the original vector).

Enough preamble—let's get into the details of Transformer-VQ. First, suppose $Q,K\in\mathbb{R}^{n\times d_k},V\in\mathbb{R}^{n\times d_v}$; standard attention is then

\begin{equation}softmax\left(QK^{\top}\right)V\end{equation}

For simplicity, the scale factor is omitted here. Transformer-VQ changes this to

\begin{equation}softmax\left(Q\hat{K}^{\top}\right)V,\quad \hat{K} = \color{skyblue}{\mathcal{VQ}}(K, C)\label{eq:vq-att}\end{equation}

where $C\in\mathbb{R}^{c\times d_k}$ is a trainable parameter—this is exactly the Codebook of VQ. (Here "VQ" refers to the same VQ used in VQ-VAE; readers unfamiliar with it can check out A Concise Introduction to VQ-VAE: The Quantized Autoencoder and The Embarrassingly Simple FSQ: "Rounding" Surpasses VQ-VAE—I won't repeat that background here.) In short, after applying $\color{skyblue}{\mathcal{VQ}}$, the most immediate effect is that each vector of $K$ gets replaced by whichever entry of $C$ is closest to it. This means every vector of $\hat{K}$ is now one of the vectors in $C$; mathematically speaking, $K\in\mathbb{R}^{n\times d_k}$ has become $\hat{K}\in C^n$.

Encoder

Of course, if we implement Transformer-VQ directly according to formula $\eqref{eq:vq-att}$, the complexity is still quadratic. But since every vector of $\hat{K}$ is one of the vectors of $C$, we can first compute $\exp\left(QC^{\top}\right)$ and then "pick out" the results corresponding to $\exp\left(Q\hat{K}{}^{\top}\right)$. Since the size of $C$ is fixed, the key operation $QC^{\top}$ has linear complexity—this is the principle behind Transformer-VQ's linearization (let's call it the "pick out" trick).

As a warm-up, let's first consider the bidirectional-attention Encoder case. Since

\begin{equation}softmax\left(QK^{\top}\right)V = \frac{\exp\left(QK^{\top}\right)V}{\exp\left(QK^{\top}\right)1_{n\times 1}}\label{eq:softmax-qkv}\end{equation}

—here $1_{n\times 1}$ denotes an all-ones matrix of size $n\times 1$, and the denominator can be viewed as a special case of the numerator—we only need to consider the numerator $\exp\left(QK^{\top}\right)V$. Since every vector of $\hat{K}$ is one of the vectors in $C$, we can construct a one-hot matrix $\Delta\in \{0,1\}^{n\times c}$, where $\Delta_i\in\{0,1\}^c$ is a one-hot vector: if the dimension where the 1 sits is $j$, then $\hat{K}_i = C_j$, and hence $\hat{K}=\Delta C$.

For Transformer-VQ, this gives us

\begin{equation}\exp\left(Q\hat{K}{}^{\top}\right)V = \exp\left(QC^{\top}\Delta^{\top}\right)V = \exp\left(QC^{\top}\right)\Delta^{\top}V = \exp\left(QC^{\top}\right)(\Delta^{\top}V)\end{equation}

The key point here is clearly the second equality! For the one-hot matrix $\Delta$, right-multiplying by its transpose lets us factor it out of $\exp$—this is the mathematical expression of the "pick out" trick mentioned above. Once factored out, by associativity of matrix multiplication, $\Delta^{\top}$ can first be multiplied by $V$, yielding a matrix of size $c\times d_v$; $\exp\left(QC^{\top}\right)$ is a matrix of size $n\times c$, and multiplying it by $\Delta^{\top}V$ gives a matrix of size $n\times d_v$. The overall theoretical complexity is $\mathcal{O}(ncd_k + ncd_v + ncd_v) = \mathcal{O}(n)$.

Finally, substituting the result of $\exp\left(Q\hat{K}{}^{\top}\right)V$ into formula $\eqref{eq:softmax-qkv}$ lets us compute the full attention output (possibly with some extra details to avoid overflow), and the entire process can be completed with linear complexity.

Decoder

Now let's consider unidirectional attention in the Decoder—this is the key case for training generative models and the basis of today's LLMs. With the Encoder as groundwork, understanding the Decoder isn't so hard. Suppose $Q_i, \hat{K}_j \in \mathbb{R}^{1\times d_k}, V_j\in\mathbb{R}^{1\times d_v}$ is one of the row vectors of the vector sequence $Q,\hat{K},V$; then for the Decoder's numerator we have

\begin{equation}\begin{aligned} O_i =&\, \sum_{j\leq i}\exp\left(Q_i\hat{K}{}_j^{\top}\right)V_j = \sum_{j\leq i}\exp\left(Q_i C^{\top}\Delta_j^{\top}\right)V_j \\ =&\, \sum_{j\leq i}\exp\left(Q_i C^{\top}\right)\Delta_j^{\top}V_j = \exp\left(Q_i C^{\top}\right)\sum_{j\leq i}\Delta_j^{\top}V_j \end{aligned}\end{equation}

If $c\times d_v$ is not too large, the final expression can be computed directly with the $\text{cumsum}$ operator. In general, however, especially in the multi-head case, to save GPU memory we typically convert this into an RNN and compute recursively, just as in the "autoregressive generation" section of Exploring Linear Attention: Must Attention Have a Softmax?. That is, letting $U_i = \sum_{j\leq i}\Delta_j^{\top}V_j\in\mathbb{R}^{c\times d_v}$, we get

\begin{equation}O_i = \exp\left(Q_i C^{\top}\right)U_i,\quad U_i = U_{i-1} + \Delta_i^{\top}V_i \end{equation}

At inference time, this kind of step-by-step recursion is fine, but during training, step-by-step computation can be quite slow. We can instead go block by block to speed things up: without loss of generality, let $n=lm$, with $l$ denoting the block size and $m$ the number of blocks; write the block slice $[il:(i+1)l]$ as $[i]$. Then

\begin{equation}\begin{aligned} O_{[i]} =&\, \exp\left(Q_{[i]}\hat{K}{}_{[i]}^{\top} + M\right)V_{[i]} + \sum_{j\lt i}\exp\left(Q_{[i]}\hat{K}{}_{[j]}^{\top}\right)V_{[j]} \\ =&\, \exp\left(Q_{[i]}\hat{K}{}_{[i]}^{\top} + M\right)V_{[i]} + \sum_{j\lt i}\exp\left(Q_{[i]}C^{\top}\Delta_{[j]}^{\top}\right)V_{[j]} \\ =&\, \exp\left(Q_{[i]}\hat{K}{}_{[i]}^{\top} + M\right)V_{[i]} + \exp\left(Q_{[i]}C^{\top}\right)\sum_{j\lt i}\Delta_{[j]}^{\top}V_{[j]} \\ \end{aligned}\end{equation}

where $M\in\{-\infty,0\}^{l\times l}$ is a lower-triangular attention mask, i.e., $M_{i,j}=0$ when $i \geq j$, and $M_{i,j}=-\infty$ otherwise. Denoting $U_i = \sum_{j\lt i}\Delta_{[j]}^{\top}V_{[j]}$, we then have

\begin{equation}O_{[i]} = \exp\left(Q_{[i]}\hat{K}{}_{[i]}^{\top} + M\right)V_{[i]} + \exp\left(Q_{[i]}C^{\top}\right)U_{i-1},\quad U_i = U_{i-1} + \Delta_{[i]}^{\top}V_{[i]} \end{equation}

This reduces the number of recursive steps to $m$, letting us maintain linear efficiency while better exploiting hardware parallelism. The denominator can be computed the same way, and dividing the two gives the full attention output.

Local Enhancement

Is that the whole story? Not quite. If this were all there is to it, Transformer-VQ might not differ much from earlier matrix-factorization-based Kernelized Attention methods like Performer. When the sequence length $n$ is much larger than the codebook size $c$, the pigeonhole principle tells us that some codebook vectors must inevitably recur repeatedly—and we might even reasonably guess that all codebook vectors end up roughly uniformly distributed across the sequence. As a result, the attention weight for a nearby token could end up equal to the attention weight for some far-away token, meaning the model can't distinguish near from far. This is essentially the low-rank problem that afflicts all Kernelized Attention methods.

Existing experience tells us that for language models, nearby tokens tend to matter more than distant ones, so a good language model architecture should have the ability to distinguish near from far. To this end, Transformer-VQ chooses to add, after $Q\hat{K}$, a Sliding-Window-shaped attention bias (denoted $B$) to up-weight nearby tokens, as shown in the figure below:

Illustration of the window attention bias

From the last panel, we can see that if we set the window size directly equal to the block size $l$—i.e., $B_{i,j}=0$ when $i < j$ or $i - j \leq l$—then, when computing block by block, the matrix $B$ affects at most the two nearest blocks, and blocks farther away can still be linearized using the "pick out" trick. To simplify the derivation below, let's write $B_{[i,j]} = B_{[il:(i+1)l,jl:(j+1)l]}$; then

\begin{equation}\begin{aligned} O_{[i]} =&\, \exp\left(Q_{[i]}\hat{K}{}_{[i]}^{\top} + B_{[i,i]}\right)V_{[i]} + \exp\left(Q_{[i]}\hat{K}{}_{[i-1]}^{\top} + B_{[i,i-1]}\right)V_{[i-1]} + \sum_{j\lt i-1}\exp\left(Q_{[i]}\hat{K}{}_{[j]}^{\top}\right)V_{[j]} \\ =&\, \exp\left(Q_{[i]}\hat{K}{}_{[i]}^{\top} + B_{[i,i]}\right)V_{[i]} + \exp\left(Q_{[i]}\hat{K}{}_{[i-1]}^{\top} + B_{[i,i-1]}\right)V_{[i-1]} + \sum_{j\lt i-1}\exp\left(Q_{[i]}C^{\top}\Delta_{[j]}^{\top}\right)V_{[j]} \\ =&\, \exp\left(Q_{[i]}\hat{K}{}_{[i]}^{\top} + B_{[i,i]}\right)V_{[i]} + \exp\left(Q_{[i]}\hat{K}{}_{[i-1]}^{\top} + B_{[i,i-1]}\right)V_{[i-1]} + \exp\left(Q_{[i]}C^{\top}\right)\sum_{j\lt i-1}\Delta_{[j]}^{\top}V_{[j]} \\ \end{aligned}\end{equation}

so it's clear that (adopting the convention that $V_{[-1]},U_{[-1]},U_{[-2]}$ are all zero matrices)

\begin{equation}\begin{aligned} O_{[i]} =&\, \exp\left(Q_{[i]}\hat{K}{}_{[i]}^{\top} + B_{[i,i]}\right)V_{[i]} + \exp\left(Q_{[i]}\hat{K}{}_{[i-1]}^{\top} + B_{[i,i-1]}\right)V_{[i-1]} + \exp\left(Q_{[i]}C^{\top}\right)U_{i-2}\\[5pt] U_i =&\, U_{i-1} + \Delta_{[i]}^{\top}V_{[i]} \end{aligned}\label{eq:tvq}\end{equation}

I believe the introduction of $B$ is the key element that sets Transformer-VQ apart from other Kernelized Attention methods. To reduce the parameter count and support variable-length generation, we constrain the nonzero part of B to be a "Toeplitz matrix," i.e., $B_{i,j}$ is a function of $i-j$—in which case $B$ plays the role of an additive relative position encoding. Besides this approach, one could also consider replacing it with the ReRoPE scheme I proposed earlier, which is a windowed version of rotary position embedding and has the same relative-position-encoding shape as $B$.

Gradient Backpropagation

Wait—we seem to have forgotten something. Readers familiar with VQ-VAE know that "every vector of $\hat{K}$ is one of the vectors in $C$" only describes the forward pass; the backward pass uses the original $K$. This means that even though $\hat{K}_j$ at different positions may equal the same $C_k$, their gradients are not the same—this is the Straight-Through Estimator (STE). Because of STE, the "pick out" trick is in principle only usable at inference time; during training, it cannot be linearized.

Is there no other way? Indeed, if we insist on getting exact gradients, there is no linearly-efficient scheme available. However, given that the gradient of VQ is itself only an approximation, obtaining exact gradients for attention doesn't seem all that necessary either. So the authors devised a compromise: still perform the recursive computation according to formula $\eqref{eq:tvq}$, but only apply STE in the first two terms (so the Key sequence can obtain gradients), while the gradient for $U_{i-1}$ is simply stopped (via the $\text{stop_gradient}$ operator). This preserves the model's linearity while retaining the most important gradients (from the two nearest blocks)—a reasonably sound approximation scheme. In this respect, Transformer-VQ is quite similar to Transformer-XL, which also stops gradients through the history window during its recursion—i.e., the history window participates in the recursive computation but doesn't propagate gradients.

Once the gradient-flow issue is resolved, adding VQ's auxiliary codebook-update loss on top of the autoregressive cross-entropy loss gives us the full training objective. Since Transformer-VQ updates its codebook via a direct exponential moving average, only the auxiliary loss for the Key needs to be added; readers who are already familiar with VQ-VAE will understand these details quickly by glancing at the original paper.

Experimental Results

Let's take a look at the experimental results from the original paper. The authors have open-sourced their code here:

GitHub: https://github.com/transformer-vq/transformer_vq

It's worth pointing out that the base architecture the authors applied VQ to is not the conventional MHA (Multi-Head Attention), but rather GAU (Gated Attention Unit) + Softmax, an architecture I've long championed. A more accurate name for Transformer-VQ might really be "GAU-VQ." Readers unfamiliar with GAU can refer to FLASH: Perhaps the Most Interesting Efficient Transformer Design in Recent Memory and Attention and Softmax Are Apparently a Great Match. In short, GAU is already more efficient than MHA on its own, and combining it with the VQ trick makes it even more powerful.

For experiments, the authors evaluated language modeling (ENWIK8, PG-19) and image generation (IMAGENET64), with a codebook size of $c=512$ used throughout. The largest model has 1.3B parameters—not on par with mainstream large models, but by no means small for a research setting. The overall experimental results are impressive:

PG-19 experimental resultsIMAGENET64 experimental results

Finally, remarkably, Transformer-VQ has a single author, and their affiliation is listed as "Independent Researcher."

Some Broader Thoughts

I find that starting from Transformer-VQ, one can connect to a great many research topics—this is one reason I admire it so much.

First, let me once again applaud the author's astonishing insight: the discovery that "simply applying VQ to the Key makes Transformer's complexity turn linear" is genuinely beautiful. It achieves a natural transition from standard attention to linear attention, and by adding an attention bias, it can be made more effective than many Kernelized Attention methods. Furthermore, the way clustering is achieved via VQ is more elegant than approaches like Linformer or Nyströmformer, because it prevents leakage of future information and can naturally be used for causal language modeling.

We know that VQ is, fundamentally, an operation that converts a sequence into discrete IDs—a role very similar to that of a tokenizer. Seen this way, Transformer-VQ is, like MegaByte, a model with a built-in tokenizer. And compared with MegaByte, the VQ operation is more similar to, and more intuitively aligned with, our traditional notion of a tokenizer. This makes Transformer-VQ especially well-suited for training "no-tokenizer" models that take raw bytes as input. Indeed, the ENWIK8 experiment mentioned above uses byte-level input, and Transformer-VQ clearly outperforms MegaByte there.

Compared with the recently released RetNet, Transformer-VQ has no explicit long-range decay, so its long-context ability may well be better. Also, since the Keys have gone through VQ and thus always belong to a finite set, there's no risk of encountering an "unseen" Key, so its length-generalization ability is likely to be better too. Although Transformer-VQ's underlying architecture, GAU, is single-head, the size of the model's memory state during recursion is $\Delta_i^{\top}V_i\in\mathbb{R}^{c\times d_v}$, which, under default settings, is actually larger than that of multi-head RetNet (RetNet's memory state size is $nd_k^2$, with $d_v = 2nd_k$ under default settings). So the memory capacity should theoretically be sufficient.

Since my previous post happened to cover The Embarrassingly Simple FSQ: "Rounding" Surpasses VQ-VAE, some readers might wonder whether the simpler FSQ could replace VQ here. I think that would be difficult, for reasons already given in that earlier post. First, $c=512$ still falls within the range of codebook sizes where VQ outperforms FSQ, so switching to FSQ would likely hurt performance. Second, since the Key of every attention layer needs to be VQ'd, on average neither the encoder nor decoder side of the VQ is particularly strong—and in such settings, VQ tends to give higher approximation precision, whereas FSQ is better suited to scenarios where both encoder and decoder are strong. Third, Transformer-VQ needs the center vector that a Key maps to under VQ, not just its ID, whereas FSQ directly produces an ID, which is not as easy to convert back into an approximate center vector.

Beyond that, using VQ rather than FSQ raises the hope that Transformer-VQ could be fine-tuned starting from an existing pretrained model such as LLaMA2, rather than only trained from scratch. Because VQ has clear geometric meaning and shares much in common with K-Means, one could start from an existing pretrained model, sample some data to compute Keys, run K-Means on those Keys to obtain center vectors as the codebook initialization, and then fine-tune the original model with VQ added on top. That said, Transformer-VQ doesn't play well with RoPE, so as mentioned earlier, a model using RoPE should probably be switched to ReRoPE before adding VQ—in which case the bias term wouldn't even be needed.

All told, in my view, among the many Efficient Transformer efforts out there, Transformer-VQ is one of the most distinctive, impressive, and deeply promising.

Summary

This post introduced an Efficient Transformer scheme called Transformer-VQ, built around the observation that "simply applying VQ to the Key makes Transformer's complexity turn linear." I personally think this is a remarkably distinctive and striking linearization idea, and the experimental results back it up well. It can be understood either as a smarter linear-attention/RNN model, or as an attention model equipped with a "trainable tokenizer."