linear attention math formulation

Topics

• linear attention
• linear algebra

Linear attention reduces computational complexity standard attention from $O(N^{2}d)$ to $O(Nd)$ using kernelization and associative matrix multiplication.

Standard scaled dot product attention output for query $Q_i$:

$${\text{Output}}_i^{1\times d_v}=\frac{{\sum }_{j=1}^N\exp \left(\frac{Q_i{K}_j^T}{\sqrt{d_k}}\right)V_j}{{\sum }_{j=1}^N\exp \left(\frac{Q_i{K}_j^T}{\sqrt{d_k}}\right)}$$

$Q_i,K_j$ row vectors ($1\times d_k$), $V_j$ row vector ($1\times d_v$), Output${}_i$ row vector ($1\times d_v$). $Q_i{K}_j^T$ is scalar dot product.

We can generalize (dot prod and softmax combo) using a similarity function $\text{sim}(Q_i,K_j)$:

$${\text{Output}}_i=\frac{{\sum }_{j=1}^N\text{sim}(Q_i,K_j)V_j}{{\sum }_{j=1}^N\text{sim}(Q_i,K_j)}$$

Linear attention approximates this similarity function using kernel decomposed by feature maps $\varphi :{\mathbb{R}}^{d_k}\to {\mathbb{R}}^c$. $\varphi$ maps $1\times d_k$ row vector to $1\times c$ row vector.

$$\text{sim}(Q_i,K_j)\approx \varphi (Q_i)\varphi (K_j{)}^T$$

Substitute into generalized attention:

$${\text{Output}}_i\approx \frac{{\sum }_{j=1}^N(\varphi (Q_i)\varphi (K_j{)}^T)V_j}{{\sum }_{j=1}^N\varphi (Q_i)\varphi (K_j{)}^T}$$

Using the linearity of dot product and the associative property of mat mul:

$$\boxed{{\text{Output}}_i\approx \frac{\varphi (Q_i)\left({\sum }_{j=1}^N\varphi (K_j{)}^T{V}_j\right)}{\varphi (Q_i)\left({\sum }_{j=1}^N\varphi (K_j{)}^T\right)}}$$

where:

• $\varphi (Q_i)$ is $1\times c$ row vector
• $\varphi (K_j{)}^T$ is $c\times 1$ column vector
• $V_j$ is $1\times d_v$ row vector
• Sum ${\sum }_{j=1}^N\varphi (K_j{)}^T{V}_j$ is $c\times d_v$ matrix. Also called as context summary matrix
• Sum ${\sum }_{j=1}^N\varphi (K_j{)}^T$ is $c\times 1$ column vector

Note

The terms ${\sum }_{j=1}^N\varphi (K_j{)}^T{V}_j$ (context-summary) and ${\sum }_{j=1}^N\varphi (K_j{)}^T$ (key-summary) are pre-computable.

Vectorized form (efficient) with $\varphi (Q)\in {\mathbb{R}}^{N\times c}$, $\varphi (K)\in {\mathbb{R}}^{N\times c}$, and $V\in {\mathbb{R}}^{N\times d_v}$:

• Global key-value (context) summary: $S=\varphi (K{)}^{T}V\in {\mathbb{R}}^{c\times d_v}$
• Global key summary: $Z={\sum }_{j=1}^N\varphi (K_j)=\varphi (K{)}^T{1}_N\in {\mathbb{R}}^c$, where $1_N$ is a column vector of ones
• Numerator: $\text{Numerator}=\varphi (Q)S\in {\mathbb{R}}^{N\times d_v}$
• Denominator (normalization term): $\text{Denominator}=\varphi (Q)Z\in {\mathbb{R}}^N$
• Output: $\text{Output}\approx \text{Numerator}/{\text{Denominator}}^{\prime }\in {\mathbb{R}}^{N\times d_v}$, where $/$ is element-wise division and ${\text{Denominator}}^{\prime }$ is broadcasted.

$$\boxed{\text{VecOutput}\approx \frac{\varphi (Q)\varphi (K{)}^{T}V}{\varphi (Q)\varphi (K{)}^T{1}_N}}$$

Complexity Analysis

• Computing $\varphi (Q)$ and $\varphi (K)$: $O(Nc{d}_k)$ or $O(Nc)$ depending on $\varphi$
• Computing $S=\varphi (K{)}^{T}V$: requires $O(Nc{d}_v)$ time
• Computing $Z=\varphi (K{)}^T{1}_N$: requires $O(Nc)$ time
• Computing Numerator: $O(Nc{d}_v)$
• Computing Denominator: $O(Nc)$
• Final element-wise division: $O(N{d}_v)$

Time complexity: $\approx O(N(c{d}_v+c))$. If $c$ and $d_v$ are constant with respect to $N$, the complexity is $O(N)$.

Space complexity: Determined by the need to store $\varphi (Q),\varphi (K),S,Z,\text{Output}$, which requires $O(Nc+c{d}_v+c+N{d}_v)\approx O(N)$ if $c$ and $d_v$ are constant.

Related