Transformer Upgrade Path: 5. Standard Attention as an Infinite-Dimensional Linear Attention

Transformer升级之路：5、作为无限维的线性Attention 苏剑林 · 2021-08-06 · Machine Learning language model attention kernel method

In Performer: Linearizing Attention Complexity via Random Projections we looked at Google's Performer model, which proposed a random projection scheme that converts standard attention into linear attention while preserving a certain degree of approximation. In theory, as long as the projection dimension is large enough, this approximation can get arbitrarily close to standard attention. In other words, standard attention can be regarded as an infinite-dimensional linear attention.

This post introduces two other approaches I've devised for converting standard attention into an infinite-dimensional linear attention. Unlike Performer's random projection, both of the schemes I propose here are deterministic, and it is relatively easy to gauge the degree of approximation they achieve.

Brief Introduction

I won't go into much detail about standard attention and linear attention here; readers who are not yet familiar with them can refer to my earlier posts Exploring Linear Attention: Does Attention Need a Softmax? and Transformer Upgrade Path: 3. From Performer to Linear Attention. Briefly, standard attention is computed as

\begin{equation}a_{i,j}=\frac{e^{\boldsymbol{q}_i\cdot \boldsymbol{k}_j}}{\sum\limits_j e^{\boldsymbol{q}_i\cdot \boldsymbol{k}_j}}\end{equation}more

while linear attention is computed as

\begin{equation}a_{i,j}=\frac{\phi(\boldsymbol{q}_i)\cdot \varphi(\boldsymbol{k}_j)}{\sum\limits_j \phi(\boldsymbol{q}_i)\cdot \varphi(\boldsymbol{k}_j)}\end{equation}

So, to (approximately) transform standard attention into linear attention, we generally need to find a transformation $\phi,\varphi$ such that the following approximation holds:

\begin{equation}\phi(\boldsymbol{q})\cdot \varphi(\boldsymbol{k})\approx e^{\boldsymbol{q}\cdot \boldsymbol{k}}\end{equation}

Here $e^{\boldsymbol{q}\cdot \boldsymbol{k}}$ is precisely the "kernel function" from kernel methods.

Random Projection

Performer found the first fairly practical random-projection scheme. In essence, it is based on the following integral:

\begin{equation}\begin{aligned} e^{\boldsymbol{q}\cdot \boldsymbol{k}} =&\,\frac{1}{(2\pi)^{d/2}}\int e^{-\Vert\boldsymbol{\omega}-\boldsymbol{q}-\boldsymbol{k}\Vert^2 / 2 + \boldsymbol{q}\cdot \boldsymbol{k}}d\boldsymbol{\omega}\\ =&\,\frac{1}{(2\pi)^{d/2}}\int e^{-\Vert\boldsymbol{\omega}\Vert^2 / 2}\times e^{\boldsymbol{\omega}\cdot \boldsymbol{q}-\Vert \boldsymbol{q}\Vert^2 / 2} \times e^{\boldsymbol{\omega}\cdot \boldsymbol{k}-\Vert \boldsymbol{k}\Vert^2 / 2}d\boldsymbol{\omega} \\ \end{aligned}\end{equation}

which gives

\begin{equation}\begin{aligned} 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} \times e^{\boldsymbol{\omega}\cdot \boldsymbol{k}-\Vert \boldsymbol{k}\Vert^2 / 2}\right]\\[6pt] &\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}}_{\phi(\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}}_{\varphi(\boldsymbol{k})} \end{aligned}\end{equation}

where $\boldsymbol{\omega}_1,\boldsymbol{\omega}_2,\cdots,\boldsymbol{\omega}_m\sim \mathcal{N}(\boldsymbol{\omega};0,\boldsymbol{1}_d)$. In this way, using the idea of random projection, we approximately convert the exponential of the inner product of two $d$-dimensional vectors into the inner product of two $m$-dimensional vectors, and when $m\to\infty$, the two are theoretically equal.

This random-projection scheme is quite clever and not easy to come up with. Below I introduce two schemes of my own devising, which are relatively easier to understand—especially for readers already familiar with kernel functions, who may grasp them at a glance.

Taylor Expansion

My first idea is based on the Taylor expansion:

\begin{equation}e^{\boldsymbol{q}\cdot \boldsymbol{k}} = \sum_{m=0}^{\infty} \frac{(\boldsymbol{q}\cdot \boldsymbol{k})^m}{m!}\end{equation}

Truncating to the first $n+1$ terms gives a degree-$n$ polynomial in $\boldsymbol{q}\cdot \boldsymbol{k}$:

\begin{equation}e^{\boldsymbol{q}\cdot \boldsymbol{k}} \approx 1 + \boldsymbol{q}\cdot \boldsymbol{k} + \frac{1}{2}(\boldsymbol{q}\cdot \boldsymbol{k})^2 + \cdots + \frac{1}{n!}(\boldsymbol{q}\cdot \boldsymbol{k})^n\end{equation}

This is in fact a "polynomial kernel function." Notice that we have:

\begin{equation}\begin{aligned} (\boldsymbol{q}\cdot \boldsymbol{k})^m =&\, \left(\sum_i q_i k_i\right)^m = \left(\sum_{i_1} q_{i_1} k_{i_1}\right)\cdots\left(\sum_{i_m} q_{i_m} k_{i_m}\right) \\ =&\, \sum_{i_1,\cdots,i_m} (q_{i_1}\cdots q_{i_m}) (k_{i_1}\cdots k_{i_m}) \end{aligned}\end{equation}

If we regard $q_{i_1}\cdots q_{i_m},k_{i_1}\cdots k_{i_m}$ as a large $d^m$-dimensional vector, then $(\boldsymbol{q}\cdot \boldsymbol{k})^m$ is just the inner product of these two large vectors. In fact, the operation of forming such a "big vector" out of several vectors is called the "outer product" of vectors, also known as the "tensor product," which is generally denoted as $\otimes$. In this notation,

\begin{equation} \frac{1}{m!}(\boldsymbol{q}\cdot \boldsymbol{k})^m = \frac{1}{m!}\underbrace{(\boldsymbol{q}\otimes\cdots\otimes\boldsymbol{q})}_{m\text{count}\boldsymbol{q}}\cdot\underbrace{(\boldsymbol{k}\otimes\cdots\otimes\boldsymbol{k})}_{m\text{count}\boldsymbol{k}} = \left(\frac{\otimes^m\boldsymbol{q}}{\sqrt{m!}}\right)\cdot\left(\frac{\otimes^m\boldsymbol{k}}{\sqrt{m!}}\right) \end{equation}

Here $\otimes^m\boldsymbol{q},\otimes^m\boldsymbol{k}$ is shorthand for the outer product of $m$ copies of $\boldsymbol{q},\boldsymbol{k}$ (the $m$-th tensor power). Using this result, we get

\begin{equation} e^{\boldsymbol{q}\cdot \boldsymbol{k}}\approx \sum_{m=0}^n \left(\frac{\otimes^m\boldsymbol{q}}{\sqrt{m!}}\right)\cdot\left(\frac{\otimes^m\boldsymbol{k}}{\sqrt{m!}}\right) =\underbrace{\begin{pmatrix} 1 \\ \boldsymbol{q}\\ \frac{\otimes^2\boldsymbol{q}}{\sqrt{2}} \\ \vdots\\ \frac{\otimes^n\boldsymbol{q}}{\sqrt{n!}}\end{pmatrix}}_{\phi(\boldsymbol{q})} \cdot \underbrace{\begin{pmatrix} 1 \\ \boldsymbol{k}\\ \frac{\otimes^2\boldsymbol{k}}{\sqrt{2}} \\ \vdots\\ \frac{\otimes^n\boldsymbol{k}}{\sqrt{n!}}\end{pmatrix}}_{\varphi(\boldsymbol{k})} \end{equation}

This completes the conversion from standard attention to linear attention.

Definition via the Exponential Limit

Compared with Performer's random projection, the Taylor-expansion approach above should be easier to understand. But there is an even more simple and direct approach, based on the defining limit of the natural exponential:

\begin{equation}e^x = \lim_{n\to\infty} \left(1+\frac{x}{n}\right)^n\end{equation}

Therefore, by choosing an appropriate $n$, we obtain

\begin{equation}e^{\boldsymbol{q}\cdot \boldsymbol{k}} \approx \left(1+\frac{{\boldsymbol{q}\cdot \boldsymbol{k}}}{n}\right)^n = \left(\begin{pmatrix} 1 \\ \frac{\boldsymbol{q}}{\sqrt{n}}\end{pmatrix} \cdot \begin{pmatrix}1 \\ \frac{\boldsymbol{k}}{\sqrt{n}}\end{pmatrix}\right)^n \end{equation}

Combining this with the polynomial-kernel conversion result from the previous section, we get

\begin{equation}e^{\boldsymbol{q}\cdot \boldsymbol{k}} \approx \underbrace{\left(\otimes^n\begin{pmatrix} 1 \\ \frac{\boldsymbol{q}}{\sqrt{n}}\end{pmatrix}\right)}_{\phi(\boldsymbol{q})} \cdot \underbrace{\left(\otimes^n\begin{pmatrix}1 \\ \frac{\boldsymbol{k}}{\sqrt{n}}\end{pmatrix}\right)}_{\varphi(\boldsymbol{k})}\end{equation}

This is perhaps the simplest and most direct scheme for converting standard attention into linear attention.

Analysis of the Results

In terms of practical value, these two deterministic schemes are far inferior to Performer's random-projection scheme, because the output dimension of random projection can be controlled fairly flexibly, whereas the output dimension of the two deterministic schemes scales as $d^n$, which is usually far larger than the sequence length itself. So using them for linear attention is basically even less efficient than standard attention.

However, theoretically speaking, these latter two schemes provide a more concise and convenient way of thinking, letting us equate standard attention with an infinite-dimensional linear attention. This equivalence often helps us better understand the attention mechanism, the most direct application being an understanding of the rank of attention.

Readers who have worked on linear attention will know that if linear attention is used for bidirectional attention tasks (such as MLM), performance drops noticeably. This is because for linear attention, $\phi(\boldsymbol{Q}),\varphi(\boldsymbol{K})\in\mathbb{R}^{n\times d}$ (where $d$ is the head size of each head) generally satisfies $n \gg d$, so the attention matrix of shape $n\times n$ obtained from $\phi(\boldsymbol{Q})\varphi(\boldsymbol{K})^{\top}$ has rank at most $d$. This is the low-rank problem of linear attention, and this low rank limits its expressive power.

By contrast, all three transformations discussed above tell us that standard attention can be viewed as an infinite-dimensional linear attention, so the rank of standard attention is, in theory, not bounded by $d$. This is why standard attention, given the same parameter count, tends to perform better than linear attention. As we also mentioned in Transformer Upgrade Path: 3. From Performer to Linear Attention, if we want to switch from standard attention to linear attention, $d$ also needs to be scaled up correspondingly in order to maintain the approximation to a reasonable degree.

Summary

This post introduced three ways of understanding standard attention as an infinite-dimensional linear attention. These different perspectives let us connect standard attention with linear attention, giving us a more comprehensive understanding of the attention mechanism from multiple angles.