Note each matrix consists of row vectors, not column vectors.
The interpretation is that each query $\mathbf{q}_j$ corresponds to a different feature which aggregates a convex combination of the values $\mathbf{v}_i$, where the weights of the convex combination are given by a scaled softmax over the dot products between query $\mathbf{q}_j$ and the keys $\mathbf{k}_i$. To see this, we can expand
where the softmax is applied row-by-row, corresponding to a softmax per query. Here, the weights corresponding to the query $\mathbf{q}_j$ comprise an $n$-dimensional vector
We normalize by $\frac{1}{\sqrt d_k}$ (the dimensionality of the queries and keys) to prevent the dot products $q_i^\top k_j$ from blowing up in value and thereby vanishing gradients in the softmax.
(Recall that if $\mathbf{q} \sim N(\mathbf{0}, \mathbf{I})$ and $\mathbf{k}\sim N(\mathbf{0},\mathbf{I})$ then $\mathbf{q}^\top \mathbf{k}\sim N(0, d)$.)
Runtime is $O(mnd)$ i.e. "quadratic" in the self-attention case where $m=n$.
Multi-Head Attention
Multi-head attention projects $\mathbf{Q},\mathbf{K},\mathbf{V}$ using $h$ heads through linear embeddings. Then it applies attention to each of the embeddings, concatenates the resulting features at the end, then sends the result through another linear transform to create outputs $\mathbf{O}$.
Note that this operation is permutation-equivariant with respect to the inputs $\mathbf{Q}$, and $\mathbf{K}$, $\mathbf{V}$.
Linear Attention
Consider the output of attention corresponding to the query $\mathbf{q}_j$. Earlier we saw that it is equal to a convex combination of the values $\mathbf{v}_i$, with weights given using keys $\mathbf{k}_i$. We can generalize the computation of these weights by expressing it as a kernel function $\mathrm{sim}(\cdot, \cdot)$ between queries and keys. Of course, we'll have to normalize so that the weights still sum to one. Scaled dot-product attention is equivalent to the following.
Recall that we can express any kernel function as an inner product in a feature space, say $p$-dimensional. Furthermore recall the associative property of matrix multiplication. We can write the following.
Observe that both $\phi(\mathbf{K})^\top \mathbf{V}$ and $\phi(\mathbf{K})^\top \mathbf{1}$ only need to be computed once. The runtime is $O((m+n)pd)$. This is extremely favorable because typically $m \approx n \gg p \approx d$, and this brings runtime down to linear.
The authors propose the following feature representation in which $p=d$.
$$\phi(\mathbf{x}) = \mathrm{ELU}(\mathbf{x})+1$$
Concurrent work Shen et al. 2021 proposes the feature representation
$$\phi(\mathbf{x}) = \exp(\mathbf{x}),$$
which approximates the original kernel function via
We can interpret linear attention as an RNN in the case of a causal attention mask (typically used in a Transformer decoder). Then each query $\mathbf{q}_j$ can only attend to keys $\mathbf{k}_i$ where $i \leq j$. That is, $\mathrm{sim}(\mathbf{q}_j, \mathbf{k}_i) = 0$ for $i >j$. We can interpret $\sum_i \phi(\mathbf{k}_i)\mathbf{v}_i$ and $\phi(\mathbf{k}_i)$ as a "hidden state" that is updated in a recurrence relation for each subsequent query $\mathbf{q}_j$.
Sinusoidal Positional Embeddings
The traditional sinusoidal embeddings from Vaswani et al. 2017 are features which interpolate between dimensions of very high frequency to dimensions of very low frequency.
In practice initializing them to random features and learning positional embeddings seems to work about as well empirically.
Relative Position Biases
Instead of explicitly adding positional embeddings to tokens, Press et al. 2022 propose the use of relative positional biases in the softmax attention. This allows the model to extrapolate i.e. plausibly generalize to sequences of lengths longer than the context window on which it was trained ("train short, test long").
Specifically, they add a positional bias to the softmax input weighted by a learnable per-head parameter $m$.
using a set of Transformer encoder blocks followed by an attention-based pooling layer.
Note that the English vocabulary consists of roughly 30k tokens in size.
References
[1] Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A.N., Kaiser, Ł., and Polosukhin, I. (2017). Attention is All you Need. In Advances in Neural Information Processing Systems 30, I. Guyon, U.V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, eds. (Curran Associates, Inc.), pp. 5998–6008.
[2] Lee, J., Lee, Y., Kim, J., Kosiorek, A., Choi, S., and Teh, Y.W. (2019). Set Transformer: A Framework for Attention-based Permutation-Invariant Neural Networks. In International Conference on Machine Learning, pp. 3744–3753.
[3] Nguyen, T. Q. & Salazar, J. Transformers without Tears: Improving the Normalization of Self-Attention. in Proceedings of the 16th International Conference on Spoken Language Translation (Association for Computational Linguistics, 2019).
[4] Katharopoulos, A., Vyas, A., Pappas, N. & Fleuret, F. Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention. in International Conference on Machine Learning 5156–5165 (PMLR, 2020).
[5] Shen, Z., Zhang, M., Zhao, H., Yi, S. & Li, H. Efficient Attention: Attention With Linear Complexities. in 2021 Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision 3531–3539 (2021).
[6] Press, O., Smith, N. & Lewis, M. Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation. in International Conference on Learning Representations (2022).
[7] Zhang, B. & Sennrich, R. Root Mean Square Layer Normalization. in Advances in Neural Information Processing Systems vol. 32 (Curran Associates, Inc., 2019).
[8] Xiong, R. et al. On Layer Normalization in the Transformer Architecture. in Proceedings of the 37th International Conference on Machine Learning 10524–10533 (PMLR, 2020).
