self-attention

Self-attention

The transformer accepts a set of observations $\mathbb{X}=\{{\mathbf{x}}_1,{\mathbf{x}}_2,\dots ,{\mathbf{x}}_T\}\in {\mathbb{R}}^{T\times d}$, (${\mathbf{x}}_t$ could be a word embedding, image patch, …). $n$ is the number of tokens and $d$ is the embedding dimension of each token.
To attend to another, each token ${\mathbf{x}}_t\in {\mathbb{R}}^{1\times d}$ (a row vector) emits a query ${\mathbf{q}}_t={\mathbf{x}}_t{\mathbf{W}}_q$ ($q_t\in {\mathbb{R}}^{1\times d_q},{\mathbf{W}}_q\in {\mathbb{R}}^{d\times d_q}$) and compares it to the keys of other tokens ${\mathbf{k}}_{\tau }={\mathbf{x}}_{\tau }{\mathbf{W}}_k$ ($k_{\tau }\in {\mathbb{R}}^{1\times d_q},{\mathbf{W}}_k\in {\mathbb{R}}^{d\times d_q}$), in order to decide what weight to assign to their values ${\mathbf{v}}_{\tau }={\mathbf{x}}_{\tau }{\mathbf{W}}_v$ ($v_{\tau }\in {\mathbb{R}}^{1\times d_v},{\mathbf{W}}_v\in {\mathbb{R}}^{d\times d_v}$):

$$z_t=\sum _{\tau }^T\kappa ({\mathbf{q}}_t,{\mathbf{k}}_{\tau }){\mathbf{v}}_{\tau }$$

Here, $\kappa (q_t,k_{\tau })\to \mathbb{R}$ is the similarity measure of ${\mathbf{q}}_t$ to ${\mathbf{k}}_{\tau }$.
The resulting output values $z_t\in {\mathbb{R}}^{1\times d_v}$ are a weighted sum of the values that each token encodes. Tokens with large $\kappa (q_t,k_{\tau })$ contribute more to the output $z_t$.
Typically, the similarity measure is the cosine similarity. Combined with a scaling factor and softmax, we get the standard scaled dot product attention $\frac{q_t\cdot \mathbf{K}}{\sqrt{d_k}}\in {\mathbb{R}}^{1\times T}$.
There are also other variants which remove softmax, e.g. forms of linear attention.
These equations can be succinctly expressed in matrix form, with all elements updated simultaneously:

$${\mathbf{z}}_t=\text{softmax}\left(\frac{{\mathbf{q}}_t{\mathbf{K}}^T}{\sqrt{d_k}}\right)\mathbf{V}\to \mathbf{Z}=\text{softmax}\left(\frac{\mathbf{Q}{\mathbf{K}}^T}{\sqrt{d_k}}\right)\mathbf{V}$$

Here, $\mathbf{Q}\in {\mathbb{R}}^{T\times d_q}$, $\mathbf{K}\in {\mathbb{R}}^{T\times d_q}$, and $\mathbf{V}\in {\mathbb{R}}^{T\times d_v}$ are matrices with rows filled by ${\mathbf{q}}_t$, ${\mathbf{k}}_t$, and ${\mathbf{v}}_t$, respectively.
The softmax is taken independently for each row (normalizing over the keys dimension). This creates a probability distribution over which nodes to attend to - ensuring attention weights sum to 1 for each query. Without this normalization, attention weights could explode/vanish and lose their “relative importance” interpretation.
The result of the ${\mathbf{QK}}^T$ operation is also called the attention matrix.

QKV?

We project the input nodes to Q,K,V, in order to give the scoring mechanism more expressivity.
In principle we could do without them, see message passing.

Q: Here is what I’m interested in (like text in a search bar).
K: Here is what I have (like titles, description, etc.).
V: Here is what I will communicate to you (actual content, filtered / refined by value head).
For every query, all values are added in a weighted sum, weighted by similarity between that query and the keys.

Attention matrix: queries = rows, keys = columns.

Self-attention?

What makes it Self- attention is, that the K, Q and V are all coming from the same source $x$, the input.
Keys and Values can also come from an entirely seperate source / nodes: cross-attention.

See also

causal attention
KV-cache

The qk - and v if it is the same size - dim is sometimes referred to as “head size”, esp. in the context of MHA.

Attention is a communication mechanism / message passing scheme between nodes.

Attention in principle be applied to any arbitrary directed graph.
Attention can be thought of as a message passing scheme.
Every node (token) has some vector of information (${\mathbb{R}}^{1\times d}$) at each point in time (data-dependent) and it gets to aggregate information via a weighted sum from all the nodes that point to it.
Imagine looping through each node:

• The query - representing what this node is looking for - gathers the keys of all edges that are input to this node and then calculates the the unnormalized “interestingness” of information / compatibilities with other nodes by taking the dot-product between the keys and the queries.
• We then normalize these scores and multiply them with the value of the input nodes.
• This happens in every head in parallel, and every layer in sequence, with different weights (in both cases).

The attention graph of encoder-transformers is fully connected.
The attention graph of the decoder is fully connected to the encoder values, and tokens are fully connected to every token that came before them (triangular attention matrix structure).
Graph of autoregressive attention (self-loops are missing from the illustration):

Attention acts over a set of vectors in a graph: There is no notion of space.

Nodes have no notion of space, e.g. where they are relative to another.
This is why we need to encode them positionally:

The transformer has a very minimal inductive bias.

In the core transformer inductive biases are mostly factored out. Self-attention is a very minimal (and verry useful) inductive bias, with a most general connectivity, where everything attends to everything..
Without positional encodings, there is no notion of space.
If you want to have a notion of space, or other constraints, you need to specifically add them.
Positional encodings, for example, are a type of inductive bias, same as for example the Swin Transformer, where you limit the attention node-connectivity to local windows, somewhat like the biologically inspired inductive bias of CNNs.
Causal attention is another example of an inductive bias, where tokens can only attend to previous tokens in the sequence.

Link to original

Desirable properties of attention

• It handles variable sequence lengths with ease
• It allows for a global receptive field without needing to scale parameters (as opposed to CNNs).
• It enables dynamic routing of information between tokens
• It is parallelizable and incredibly GPU-friendly

Transformers and The Bitter Lesson (from The Bittersweet Lesson)

→ If it’s possible for a human domain expert to discover from the data the basis of a useful inductive bias, it should be obvious for your model to learn about it too, so no need for an inductive bias.
→ Instead focus on building biases that improve either scale, learning or search .
→ In the case of sequence models, any bias towards short-term dependency is needless, and may inhibit learning (about long-term dependency).
→ Skip connections are good because they promote learning.
→ MHSA is good because it enables the Transformer to (learn to) perform an online, feed-forward, parallel ‘search’ over possible interpretations of a sequence.

Link to original

self-attention is an indirect encoding

Computing an attention matrix as $A=X{X}^T$ from the raw token features produces a gram matrix. This already captures associations/relational structure between tokens, but since there are no free (learnable) parameters,, it isn’t suitable for tasks beyond simple association.
Learnable $W_q,W_k\in {\mathbb{R}}^{d_{\text{emb}},d_{qk}}$ matrices are introduced, projecting the raw features into a new space, where $A$ is more meaningful for the task at hand.
These weight matrices can be seen as the small genotype encoding the much larger phenotype attention matrix:
When we scale number of learnable parameters while the input dimension is held constant, the free parameters scale with $O(d_{\text{qk}})$, while $A$ grows with $O(n^2)$, and typically, $n^2\gg d_{\text{qk}}$, where $n$ are the number of input tokens.

E.g. for gpt 3, there’s a 4 OOM difference!

>>> ctx, d_qk = 2048, 2*128; ctx**2/d_qk
16384.0

Link to original

The hamiltonian of the Hopfield Network is identical to the ising model, except that the interaction strength is not constant, and very similar to attention!

For a hopfield network:

$$H=-\sum _{i\ne j}{W}_{ij}{s}_i{s}_j$$

$W=X_i^T{X}_j$ is the weight / interaction matrix, which encodes how patterns are stored, rather than being a simple interaction constant.
In a simplified form, the update rule of the HFN is:

$$\Delta s_i\approx \sum _j(X_i^T{X}_j)s_j$$

And for attention:

$$\Delta s_i\approx \sum _j\sigma (K_i^T{Q}_j)V_j$$

Tokens in self-attention are just like spins in the ising model, but instead of spins, they are vectors in higher dimensional space, with all to all communication, self-organizing to compute the final representation.

Link to original

Attention matrix $\times$ values “visualization”.

http://matrixmultiplication.xyz/

Self-Attention is permutation invariant and permutation equivariant.

NOTE: The full transformer is not - set transformer changes this.
(from set transformer)

Transclude of Relating-transformers-to-models-and-neural-representations-of-the-hippocampal-formation#^97a897

References

Relating transformers to models and neural representations of the hippocampal formation
Stanford CS25: V2 I Introduction to Transformers w/ Andrej Karpath

Karpathy notebook and YT
Karpathy Attention timstamped)

AssemblyAI YT intro.
Deep-learning bible reference
Implementation: Pytorch Lightning Notebook with explanations.