Softmax-Free Linear Transformers

Abstract

Vision transformers (ViTs) have pushed the state-of-the-art for visual perception tasks. The self-attention mechanism underpinning the strength of ViTs has a quadratic complexity in both computation and memory usage. This motivates the development of approximating the self-attention at linear complexity. However, an in-depth analysis in this work reveals that existing methods are either theoretically flawed or empirically ineffective for visual recognition. We identify that their limitations are rooted in the inheritance of softmax based self-attention during approximations, that is, normalizing the scaled dot-product between token feature vectors using the softmax function. As preserving the softmax operation challenges any subsequent linearization efforts. By this insight, a family of SOftmax-Free Transformers (SOFT) are proposed. Specifically, a Gaussian kernel function is adopted to replace the dot-product similarity, enabling a full self-attention matrix to be approximated under low-rank matrix decomposition. For computational robustness, we estimate the Moore–Penrose inverse using an iterative Newton–Raphson method in the forward process only, while calculating its theoretical gradients only once in the backward process. To further expand applicability (e.g., dense prediction tasks), an efficient symmetric normalization technique is introduced. Extensive experiments on ImageNet, COCO and ADE20K show that our SOFT significantly improves the computational efficiency of existing ViT variants. With linear complexity, much longer token sequences are permitted by SOFT, resulting in superior trade-off between accuracy and complexity. Code and models are available at https://github.com/fudan-zvg/SOFT.

Appendices

Appendix A: Nyström Method

Nyström method (Williams & Seeger, 2000) aims to calculate a low-rank approximation for a Gram matrix. For Transformers, the self-attention matrix can be viewed as a Gram matrix S with a Gaussian kernel k applied to the query Q, with each element $S_{ij}$ expressed as:

$$\begin{aligned} S_{ij} =k\big (Q_{i,:},Q_{j,:}\big ) = \text {exp} (-\frac{\Vert Q_{i,:}-Q_{j,:}\Vert _2^2}{2 \sqrt{d}}), \end{aligned}$$

(18)

k(x, y) means operating Gaussian kernel k to (x, y), which can be written in the feature space as:

$$\begin{aligned} k(x,y)=\sum \limits _{i=1}^n \lambda _i \phi _i(x)\phi _i(y), \end{aligned}$$

(19)

n is the dimension of a feature space, $\lambda _i$ denotes the eigenvalue and $\phi _i$ denotes the eigenfunction of kernel k. According to the eigenfunction’s definition, we can get:

$$\begin{aligned} \int k(y,x)\phi _i(x)p(x)dx=\lambda _i \phi _i(y), \end{aligned}$$

(20)

where p(x) is the probability distribution of x. And {$\phi _i(x)$} are p-orthogonal:

$$\begin{aligned} \int \phi _i(x)\phi _j(x)p(x)dx=\delta _{ij}. \end{aligned}$$

(21)

$\delta _{ij}$ is 0 when $i \ne j$, 1 when $i = j$. To get an approximation of the eigenfunctions, we sample $\{x_1,x_2,\cdots ,x_q\}$ from p(x), then:

$$\begin{aligned}{} & {} \frac{1}{q}\sum \limits _{t=1}^q k(y,x_t)\phi _i(x_t) \approx \lambda _i \phi _i(y), \end{aligned}$$

(22)

$$\begin{aligned}{} & {} \frac{1}{q}\sum \limits _{t=1}^q \phi _i(x_t)\phi _j(x_t) \approx \delta _{ij}. \end{aligned}$$

(23)

This inspires us to approximate the Gram matrix S. Let $S^{(m)}$ be a submatrix of S, consisting of $m \times m$ elements from S. Gram matrix is a symmetric positive semi-definite matrix, so it has a spectral decomposition:

$$\begin{aligned} S^{(m)}U^{(m)}=U^{(m)} \Lambda ^{(m)}, \end{aligned}$$

(24)

where $U^{(m)}$ is column orthogonal and $\Lambda ^{(m)}$ is a diagonal matrix with the diagonal elements as the eigenvalues of $S^{(m)}$. Substituting the y to $x_j$ and applying the approximation above to S, we can get:

$$\begin{aligned}{} & {} \phi _i(x_j) \approx \sqrt{m} U_{j,i}^{(m)}, \quad \lambda _i \approx \frac{\lambda _i^{(m)}}{m}, \end{aligned}$$

(25)

$$\begin{aligned}{} & {} \phi _i(y) \approx \frac{\sqrt{m}}{\lambda _i^{(m)}}\sum \limits _{t=1}^{m}k(y,x_t)\phi _i(x_t), \end{aligned}$$

(26)

$\lambda _i$ is eigenvalue of S and $\lambda _i^{(m)}$ is the eigenvalue of $S^{(m)}$. Denote $\tilde{S}$ as the rank-m approximation of S and $\tilde{U}, \tilde{\Lambda }$ as the approximation for spectral decomposition of S. Now we can get an approximation of S with rank m:

$$\begin{aligned} \tilde{S}=\tilde{U}\tilde{\Lambda }\tilde{U}^T=\sum \limits _{t=1}^m \tilde{\lambda _t}^{(n)} \tilde{u}_t^{(n)}(\tilde{u}_t^{(n)})^T. \end{aligned}$$

(27)

Similarly, we have:

$$\begin{aligned} \phi _i(x_j) \approx \sqrt{n} U_{j,i}(n),\quad \lambda _i \approx \frac{\tilde{\lambda _i}^{(n)}}{n}. \end{aligned}$$

(28)

Thus

$$\begin{aligned}{} & {} \tilde{\lambda _i}^{(n)} \approx \frac{n \lambda _i^{(m)}}{m}, \end{aligned}$$

(29)

$$\begin{aligned}{} & {} \tilde{u}_t^{(n)} \approx \sqrt{\frac{m}{n}}\frac{1}{\lambda _t^{(m)}}S_{n,m} u_t^{(m)}. \end{aligned}$$

(30)

Then we get an approximation of S: $\tilde{S} \approx S_{n,m} S_{m,m}^{\dagger } S_{m,n}$. S has a block representation below:

$$\begin{aligned} S= \begin{bmatrix} S_{m,m} &{} S_{m,n-m} \\ S_{n-m,m} &{} S_{n-m,n-m} \end{bmatrix}. \end{aligned}$$

(31)

Appendix B: Newton Method

Proof of Proposition 1

When $\alpha$ is sufficiently small, $A_{k+1}=2A_k-A_k A A_k$, $A_k$ converges to $A^{\dagger }$. $\square$

Proof

A is a symmetric positive semi-definite matrix and $A_{ij}\le 1$, $\forall 1\le i,j \le n$, $A_{ii}=1,\quad 1\le i\le n$ in our case. $A_0$ is chosen to be $\alpha A$, so the $A_k$ can be written as $A_k=C_k A =AD_k$ for some matrix $C_k,D_k$, leading to the fact that

$$\begin{aligned} A^{\dagger }AA_k=A_k,\quad A_k A A^{\dagger }=A_k. \end{aligned}$$

(B15)

This is because $A_{k+1}=A_k(2I_n - AA_k)=(2I_n - A_kA)A_k$ and $A_0=\alpha A$. We make a difference between $A^{\dagger }$ and $A_{k+1}$:

$$\begin{aligned} A^{\dagger }-A_{k+1}&= A^{\dagger } - 2A_k + A_k A A_k \nonumber \\&= A^{\dagger } -A_k A A^{\dagger } -A^{\dagger } A A_k +A_k A A_k \nonumber \\&= (A^{\dagger }-A_k) A (A^{\dagger }-A_k). \end{aligned}$$

(B16)

We norm both sides of the equation above:

$$\begin{aligned} \Vert A^{\dagger }-A_{k+1} \Vert&= \Vert ( A^{\dagger }-A_{k})A( A^{\dagger }-A_{k})\Vert \nonumber \\&\le \Vert A^{\dagger }-A_{k} \Vert \Vert A( A^{\dagger }-A_{k})\Vert . \end{aligned}$$

(B17)

And we left multiply A on the both sides of (Eq B16), then norm the equation:

$$\begin{aligned} \Vert A A^{\dagger }-A A_{k+1}\Vert&=\Vert A (A^{\dagger }-A_k) A (A^{\dagger }-A_k)\Vert \nonumber \\&\le \Vert A A^{\dagger } -A A_{k}\Vert ^2 . \end{aligned}$$

(B18)

We choose $\alpha$ sufficiently small so that the initial value satisfy $\Vert AA^{\dagger }-A A_0\Vert < 1$. We set $\alpha =\frac{2}{\Vert A\Vert _1^2}$ to ensure it is small enough (Ben-Israel & Cohen, 1966). Then the $\Vert A A^\dagger - A A_{k}\Vert \rightarrow 0$, when $k \rightarrow \infty$. The inequality (B17) implies that $A_k \rightarrow A^{\dagger },\ k \rightarrow \infty$. $\square$

Proof of equation 12

$$\begin{aligned} \nabla _x {\mathcal {L}} = -Y^\top (\nabla _Y {\mathcal {L}}) Y^\top , \end{aligned}$$

Proof

Since $XY=I$, we take derivatives on the both sides

$$\begin{aligned} \nonumber YdX + XdY&= 0\\ dY&= -YdXY \end{aligned}$$

(B19)

For the loss function ${\mathcal {L}}$, we have

$$\begin{aligned} d{\mathcal {L}} = \langle \nabla _X {\mathcal {L}}, dX \rangle = \langle \nabla _Y {\mathcal {L}}, dY \rangle \end{aligned}$$

Recall that the inner product $\langle A,B\rangle = \text {Tr}(A^\top B)$ and $\text {Tr}(ABC)=\text {Tr}(CAB)=\text {Tr}(BCA)$, we have

$$\begin{aligned} \langle \nabla _X {\mathcal {L}}, dX \rangle&= \langle \nabla _Y {\mathcal {L}}, dY \rangle \\&= \langle \nabla _Y {\mathcal {L}}, -YdXY \rangle \\&= -\text {Tr}(\nabla _Y {\mathcal {L}}^\top YdXY )\\&= \langle -Y^\top \nabla _Y {\mathcal {L}} Y^\top , dX \rangle \end{aligned}$$

Therefore, Eq (12) is proved. $\square$

Appendix C: Attention Normalization

Proof of Porpostion 5

Assume the bottleneck matrix of softmax-free attention, $A\in {\mathbb {R}}^{m\times m}$ is k-connected. If $\lambda _1\ge \lambda _2\ge \cdots \lambda _m \ge 0$ are eigenvalues of $A^\dag$, then $\lambda _1 = {\mathcal {O}}(m^2)$ and $\Vert A^\dag \Vert _2 = {\mathcal {O}}(m^2)$. $\square$

To prove Proposition 5, we first introduce the following lemmas.

Lemma 1

If F is a non-singular and symmetric matrix, then $\Vert F\Vert _2\le \Vert F\Vert _1$.

Proof

According to the definition of 2-Norm, we have $F^TFz=\mu ^2z$ with $\mu = \Vert F\Vert _2$. Therefore, $\mu ^2\Vert z\Vert _1 = \Vert F^TFz\Vert _2\le \Vert F^T\Vert _1\Vert F\Vert _1\Vert z\Vert _1 = \Vert F\Vert _{\infty }\Vert \Vert F_1\Vert z\Vert _1$. Since A is non-singular and symmetric,

$$\begin{aligned} \Vert F\Vert _2 \le \sqrt{\Vert F\Vert _{\infty }\Vert F\Vert _1} = \Vert F\Vert _1 \end{aligned}$$

(C19)

$\square$

Lemma 2

If F is a real matrix and $\Vert F\Vert _p < 1$, then

$$\begin{aligned} (I-F)^{-1} = \sum _{k=0}^\infty F^k \end{aligned}$$

(C20)

with

$$\begin{aligned} \Vert (I-F)^{-1}\Vert _p \le \frac{1}{1-\Vert F\Vert _p} \end{aligned}$$

(C21)

Proof

$$\begin{aligned} \left( \sum _{k=0}^N F^k \right) (I-F) = I- F^{N+1} \end{aligned}$$

Since $\Vert F\Vert _p<1$, $\lim _{k\rightarrow \infty } F^k=0$. Thus,

$$\begin{aligned} \left( \lim _{N\rightarrow \infty } \sum _{k=0}^N F^k \right) (I-F) = I \end{aligned}$$

Multiply the equation by $(I-F)^{-1}$ we obtain Eq (C20). From that, we can easily get

$$\begin{aligned} \Vert (I-F)^{-1}\Vert _p \le \sum _{k=0}^\infty \Vert F\Vert _p^k = \frac{1}{1-\Vert F\Vert _p} \end{aligned}$$

$\square$

Now we begin to prove Proposition 5.

Proof

Since the bottleneck matrix softmax-free attention $\Vert A\Vert _1$ is not less than 1, we normalize the matrix as $A_n = D^{-1}A$, where $D = \text {diag}(A\mathbb {1}_m)$. From Lemma 2, it follows that

$$\begin{aligned} \Vert A_n^{-1}\Vert _1 \le \frac{1}{1-\Vert 1-A_n\Vert _1} \end{aligned}$$

Since we assume the bottleneck matrix is k-connected, and $k<<m$,

$$\begin{aligned} \Vert I-A_n\Vert _1 \le \frac{(m-k)}{m} \end{aligned}$$

Then,

$$\begin{aligned} \Vert A_n^{-1}\Vert _1 \le \frac{1}{1 - \frac{(m-k)}{m}} = \frac{m}{k}= {\mathcal {O}}(m) \end{aligned}$$

Note that we $D={\mathcal {O}}(m)$ by assumption that the bottleneck matrix is k-connected, it follows that

$$\begin{aligned} \Vert A^{-1}\Vert _1 = \Vert A_n^{-1}D\Vert _1 \le \Vert A_n^{-1}\Vert _1\Vert D\Vert _1 = {\mathcal {O}}(m^2) \end{aligned}$$

$\square$

Fig. 7

Comparison of attention heatmaps for a selected query patch (indicated by a cross "+") against all patches in an image. Heatmaps are derived from the first head’s corresponding row in the attention maps, as calculated by Eq. 17. These heatmaps are normalized to a 0–1 scale, with warmer colors indicating higher relevance. The model variants compared are: a Transformer (Vaswani et al., 2017), b Performer (Choromanski et al., 2021), c Nystromformer (Xiong et al., 2021), and d Our SOFT approach

Appendix D: Non-linearized Gaussian Kernel Attention

Instead of directly calculating the Gaussian kernel weights, in our formulation they are approximated. More specifically, the relation between any two tokens is reconstructed via sampled bottleneck tokens. As the number m (e.g., 49) of bottleneck tokens is much smaller than the entire token sequence length, our attention matrix is of low-rank. This brings about two favorable consequences: (I) The model now focuses the attentive learning on latent salient information captured by the m bottleneck tokens. (II) The model becomes more robust against the underlying token noise due to the auto-encoder style reconstruction (Geng et al., 2021). This explains why a model with an approximated gram matrix performs better than one with a directly computed matrix. Further, we find that exact Gaussian kernel attention leads to training difficulties. As Proposition 4 reveals, this is due to lacking normalization that leads to explosion of the spectral norm especially in long token sequence cases. Big spectral norm could jeopardize the training and tend to collapse the model.

Appendix E: Attention Visualization

Figure 7 shows more visualization of the attention masks by various Transformers (Vaswani et al., 2017; Choromanski et al., 2021; Xiong et al., 2021) and our SOFT. For each model, we show the output from the first two attention heads (up and down row). It is noteworthy that SOFT exhibits better semantic diversity of the multi-head mechanism than other methods. Moreover, when we sample the patch at the boundary of multiple objects, SOFT is able to more precisely capture all these objects.