I Found Traces of Transformer-VQ in Performer

我在Performer中发现了Transformer-VQ的踪迹 苏剑林 · 2023-11-29 · Machine Learning quantization language model attention

A few days ago, in Quantizing the Key: Transformer's Complexity Becomes Linear, we introduced "Transformer-VQ," a scheme that linearizes attention complexity by applying VQ (Vector Quantize) to the Key sequence. Admittedly, Transformer-VQ offers a beautiful transition from standard attention to linear attention, giving a sense of "great simplicity underlying great truth." But readers familiar with VQ can probably sense that as the codebook size or model parameter count grows further, VQ is likely to become a bottleneck for performance, since the gradients it estimates via STE (Straight-Through Estimator) are probably suboptimal (the experimental results of FSQ offer some corroborating evidence). Moreover, the gradient truncation that Transformer-VQ employs to keep training efficiency linear may also turn out to be a future performance bottleneck.

With this in mind, I spent some time thinking about alternative approaches to linearization that could replace VQ. From the $\exp\left(QC^{\top}\right)$ form in Transformer-VQ, I was reminded of Performer, and by "following the vine to find the melon," I discovered that Performer can actually be viewed as a soft version of Transformer-VQ. Going further, I tried to imitate Performer's derivation method to re-derive Transformer-VQ, in the hope of providing some reference for future optimization. more

Recap

First, let's take a moment to review Transformer-VQ. Let $Q,K\in\mathbb{R}^{n\times d_k},V\in\mathbb{R}^{n\times d_v}$; the key idea of Transformer-VQ is to apply the following VQ approximation to $K$:

\begin{equation}K\approx\hat{K}\triangleq\Delta C\end{equation}

Here $\Delta\in\{0,1\}^{n\times c},C\in\mathbb{R}^{c\times d_k}$ are all matrices, where $C$ is a trainable parameter, and $\Delta$ is defined as:

\begin{equation}\Delta_{i,j} = \left\{\begin{aligned}& 1, \quad j=\mathop{\text{argmin}}_{k=1,2,\cdots,c} \Vert K_i - C_k\Vert \\ & 0, \quad\text{other}\end{aligned}\right.\end{equation}

In plain terms, VQ approximates $K_i$ by the $C_j$ that is closest to $K_i$. Under this approximation, we have (for simplicity, taking the Encoder as an example):

\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)\label{eq:transformer-vq}\end{equation}

Readers familiar with linear attention will readily recognize that the computation in this last expression has linear complexity — it is precisely one of the protagonists of this post, Transformer-VQ (this gives the numerator; the denominator follows analogously).

With no particularly complicated derivation, linear attention just falls out. This gives the impression that, almost "by accident" while approximating the Key, the complexity of attention has been reduced to linear — a truly elegant result. Hence we return once more to the assessment we've made several times already: Transformer-VQ offers a very beautiful transition from standard attention to linear attention.

A Sense of Déjà Vu

The $\exp\left(QC^{\top}\right)$ in Transformer-VQ reminded me of an earlier post, Transformer Upgrade Road: 3. From Performer to Linear Attention. In that post, I simplified Performer's results and argued that the optimal activation function for $Q,K$ in linear attention is $\exp$. Since $\exp$ also shows up in Transformer-VQ, perhaps there's some connection between the two.

To dig into this connection, let's bring in Performer, which is based on an elegant approximation:

\begin{equation} e^{\boldsymbol{q}\cdot \boldsymbol{k}}=\mathbb{E}_{\boldsymbol{\omega}\sim \mathcal{N}(\boldsymbol{\omega};0,\boldsymbol{1}_d)}\left[e^{\boldsymbol{\omega}\cdot \boldsymbol{q}-\Vert \boldsymbol{q}\Vert^2 / 2} \,e^{\boldsymbol{\omega}\cdot \boldsymbol{k}-\Vert \boldsymbol{k}\Vert^2 / 2}\right]\approx\underbrace{\frac{1}{\sqrt{m}}\begin{pmatrix}e^{\boldsymbol{\omega}_1\cdot \boldsymbol{q}-\Vert \boldsymbol{q}\Vert^2 / 2} \\ e^{\boldsymbol{\omega}_2\cdot \boldsymbol{q}-\Vert \boldsymbol{q}\Vert^2 / 2}\\ \vdots\\ e^{\boldsymbol{\omega}_m\cdot \boldsymbol{q}-\Vert \boldsymbol{q}\Vert^2 / 2} \end{pmatrix}}_{\tilde{\boldsymbol{q}}} \cdot \underbrace{\frac{1}{\sqrt{m}}\begin{pmatrix}e^{\boldsymbol{\omega}_1\cdot \boldsymbol{k}-\Vert \boldsymbol{k}\Vert^2 / 2} \\ e^{\boldsymbol{\omega}_2\cdot \boldsymbol{k}-\Vert \boldsymbol{k}\Vert^2 / 2}\\ \vdots\\ e^{\boldsymbol{\omega}_m\cdot \boldsymbol{k}-\Vert \boldsymbol{k}\Vert^2 / 2} \end{pmatrix}}_{\tilde{\boldsymbol{k}}} \label{eq:performer}\end{equation}

Since in the end we still need to normalize the attention over all $\boldsymbol{k}$, dropping $\frac{1}{\sqrt{m}}$ and $-\Vert \boldsymbol{q}\Vert^2/2$ from the above expression won't affect the final result. Likewise, if we assume that the norms of $\boldsymbol{\omega}_1,\boldsymbol{\omega}_2,\cdots,\boldsymbol{\omega}_m$ are all equal (see the JL Lemma), then subtracting $\Vert\boldsymbol{\omega}_i\Vert^2/2$ from the exponents of $\boldsymbol{k}$ also won't change the result. Thus, Performer is equivalent to computing $\tilde{\boldsymbol{q}},\tilde{\boldsymbol{k}}$ in the following form:

\begin{equation}\underbrace{\begin{pmatrix}e^{\boldsymbol{\omega}_1\cdot \boldsymbol{q}} \\ e^{\boldsymbol{\omega}_2\cdot \boldsymbol{q}}\\ \vdots\\ e^{\boldsymbol{\omega}_m\cdot \boldsymbol{q}} \end{pmatrix}}_{\tilde{\boldsymbol{q}}} \cdot \underbrace{\begin{pmatrix}e^{\boldsymbol{\omega}_1\cdot \boldsymbol{k}-\Vert \boldsymbol{k}\Vert^2 / 2-\Vert \boldsymbol{\omega}_1\Vert^2 / 2} \\ e^{\boldsymbol{\omega}_2\cdot \boldsymbol{k}-\Vert \boldsymbol{k}\Vert^2 / 2-\Vert \boldsymbol{\omega}_2\Vert^2 / 2}\\ \vdots\\ e^{\boldsymbol{\omega}_m\cdot \boldsymbol{k}-\Vert \boldsymbol{k}\Vert^2 / 2-\Vert \boldsymbol{\omega}_m\Vert^2 / 2} \end{pmatrix}}_{\tilde{\boldsymbol{k}}} = \underbrace{\begin{pmatrix}e^{\boldsymbol{\omega}_1\cdot \boldsymbol{q}} \\ e^{\boldsymbol{\omega}_2\cdot \boldsymbol{q}}\\ \vdots\\ e^{\boldsymbol{\omega}_m\cdot \boldsymbol{q}} \end{pmatrix}}_{\tilde{\boldsymbol{q}}} \cdot \underbrace{\begin{pmatrix}e^{-\Vert \boldsymbol{k}-\boldsymbol{\omega}_1\Vert^2 / 2} \\ e^{-\Vert \boldsymbol{k} - \boldsymbol{\omega}_2\Vert^2 / 2}\\ \vdots\\ e^{-\Vert \boldsymbol{k} - \boldsymbol{\omega}_m\Vert^2 / 2} \end{pmatrix}}_{\tilde{\boldsymbol{k}}} \propto \underbrace{\begin{pmatrix}e^{\boldsymbol{\omega}_1\cdot \boldsymbol{q}} \\ e^{\boldsymbol{\omega}_2\cdot \boldsymbol{q}}\\ \vdots\\ e^{\boldsymbol{\omega}_m\cdot \boldsymbol{q}} \end{pmatrix}}_{\tilde{\boldsymbol{q}}} \cdot \underbrace{softmax\begin{pmatrix}e^{-\Vert \boldsymbol{k}-\boldsymbol{\omega}_1\Vert^2 / 2} \\ e^{-\Vert \boldsymbol{k} - \boldsymbol{\omega}_2\Vert^2 / 2}\\ \vdots\\ e^{-\Vert \boldsymbol{k} - \boldsymbol{\omega}_m\Vert^2 / 2} \end{pmatrix}}_{\tilde{\boldsymbol{k}}} \end{equation}

Comparing this last expression with $\eqref{eq:transformer-vq}$, we notice many similarities: doesn't $\boldsymbol{\omega}_1,\boldsymbol{\omega}_2,\cdots,\boldsymbol{\omega}_m$ play the same role as the codebook $C$? Doesn't $\tilde{\boldsymbol{q}}$ correspond to $\exp\left(QC^{\top}\right)$? As for the final $\tilde{\boldsymbol{k}}$, it takes softmax over logits $-\Vert \boldsymbol{k} - \boldsymbol{\omega}_i\Vert^2 / 2$, which highlights precisely the $\boldsymbol{\omega}_i$ closest to $\boldsymbol{k}$ — and isn't the limiting case of softmax exactly a one-hot vector? Doesn't that correspond precisely to the $\Delta$ matrix in Transformer-VQ? So while it's not an exact match, the resemblance is striking — maybe 60-70% identical.

Imitating the Pattern

Of course, the above is more of a suggestive analogy than a genuine equivalence, because Performer is fundamentally based on a completely different approximation approach — for instance, the $\boldsymbol{\omega}_1,\boldsymbol{\omega}_2,\cdots,\boldsymbol{\omega}_m$ in it are randomly sampled and then fixed, meaning that their approximation quality as "center vectors" is actually quite poor. But this analogy prompted a thought: could we imitate Performer's approach to re-derive Transformer-VQ from scratch? That is, following the pattern of equation $\eqref{eq:performer}$, first construct an exactly equal expression, and then convert it into a sampling-based approximation to obtain a linear version.

After several days of thinking, I found a scheme that produces exactly the derivation I was hoping for. First, using the Dirac delta function, we can write

\begin{equation}e^{\boldsymbol{q}\cdot \boldsymbol{k}} = \int e^{\boldsymbol{q}\cdot \boldsymbol{\omega}}\delta(\boldsymbol{\omega} - \boldsymbol{k})d\boldsymbol{\omega}\end{equation}

This is purely an identity given by the definition of the Dirac delta function, and doesn't yet involve any clever manipulation or approximation. However, when we substitute it into the (numerator of the) attention computation, something interesting happens:

\begin{equation}\sum_j e^{\boldsymbol{q}\cdot \boldsymbol{k}_j} \boldsymbol{v}_j = \sum_j \boldsymbol{v}_j\int e^{\boldsymbol{q}\cdot \boldsymbol{\omega}}\delta(\boldsymbol{\omega} - \boldsymbol{k}_j)d\boldsymbol{\omega} = \int e^{\boldsymbol{q}\cdot \boldsymbol{\omega}} \left[\sum_j \delta(\boldsymbol{\omega} - \boldsymbol{k}_j) \boldsymbol{v}_j\right]d\boldsymbol{\omega}\label{eq:inf-vq}\end{equation}

Isn't the final equality exactly the form of linear attention?! Of course, since it requires integrating over $\boldsymbol{\omega}$, this — like in Transformer Upgrade Road: 5. Linear Attention as Infinite-Dimensional — is an "infinite-dimensional" linear attention, which for now has only formal value.

Typically, we'd interpret $\delta(\boldsymbol{\omega} - \boldsymbol{k}_j)$ as the limit of a normal distribution $\mathcal{N}(\boldsymbol{\omega};\boldsymbol{k}_j,\sigma^2\boldsymbol{I})$ as $\sigma\to 0$, which also means that $\delta(\boldsymbol{\omega} - \boldsymbol{k}_j)$ carries the meaning of a conditional distribution $p(\boldsymbol{\omega}|\boldsymbol{k}_j)$. However, from the perspective of generative models, the Dirac delta function is just a single-point distribution — in plain terms, it simply memorizes the training set, and so it has no capacity for abstraction or generalization. To alleviate this, we approximate $p(\boldsymbol{\omega}|\boldsymbol{k}_j)$ using a GMM (Gaussian Mixture Model):

\begin{equation}p(\boldsymbol{\omega}|\boldsymbol{k}_j) \approx \sum_{y=1}^m \mathcal{N}(\boldsymbol{\omega};\boldsymbol{c}_y,\sigma^2\boldsymbol{I}) \,p(y|\boldsymbol{k}_j) \end{equation}

Substituting this into equation $\eqref{eq:inf-vq}$ and then taking the limit as $\sigma\to 0$, we obtain

\begin{equation}\sum_j e^{\boldsymbol{q}\cdot \boldsymbol{k}_j} \boldsymbol{v}_j \approx \sum_{y=1}^m e^{\boldsymbol{q}\cdot \boldsymbol{c}_y} \left[\sum_j p(y|\boldsymbol{k}_j) \boldsymbol{v}_j\right]\end{equation}

This gives us a finite-dimensional linear attention. If we align $p(y|\boldsymbol{k}_j)$ with the one-hot distribution $\Delta$ defined in Transformer-VQ, the result we get is precisely equation $\eqref{eq:transformer-vq}$ of Transformer-VQ.

Summary

This post presented a discovery of mine: the earlier linear attention work "Performer" can be viewed as a "soft" version of Transformer-VQ. Building on this observation, I further obtained a new derivation of Transformer-VQ: using the Dirac delta function to convert standard attention into infinite-dimensional linear attention, and then applying a GMM approximation to arrive at Transformer-VQ.