Is the multi-head attention in the transformer a weighted adjacency matrix?

Question

Are multi-head attention matrices weighted adjacency matrices?

The job of the multi-head-attention mechanism in transformer models is to determine how likely a word is to appear after another word. In a sense this makes the resulting matrix a big graph with nodes and edges, where a node represents a word and an edge the likelihood to appear after that. So basically it is an adjacency matrix that is created.

Answer 1

Short answer, yes I believe we can! One way feels more meaningful that the other. First, let's look at some nuance in the definition of attention. If $\text{score}(x_i, x_j) = \text{score}(x_j, x_i)$, then the attention matrix is symmetric and naturally has the form of a weighted adjacency matrix. For example, this happens when attention is given by a simple dot product $\text{score}(x_i, x_j) = \langle x_i, x_j \rangle = x_i^Tx_j$. This also happens if we have a learnable matrix $A$ and $\text{score}(x_i, x_j) = x_i^TAx_j$ if $A$ is symmetric. We can then view the attention matrix $\alpha_{i,j} = \text{Attn}_{i,j}(X)$ as a weighted adjacency matrix where the nodes represent input tokens, and edge weights correspond to similarity scores (as defined by the inner product, scaled inner product, or the symmetric matrix $A$). Now, for the following definition of attention, this makes a little less sense, as the edge connecting token nodes in graph corresponding to tokens $x_i$ and $x_j$ is not the same as the weight on the edge connecting $x_j$ to $x_i$, in general. Suppose $X \in \mathbb{R}^{d \times n}$ has as columns $X_i$ the $d$-dimensional embeddings of the $n$-tokens $x_1, x_2, ..., x_n$ from your input. Now, let

$$W_QX = Q$$ $$W_KX = K$$ $$W_VX = V$$

be the learned weight matrices giving the "query" $q_i = W_QX_i$, "key" $k_i = W_KX_i$, and "value" $v_i = W_VX_i$ vectors. Then we can define attention as

\begin{align} \text{Attn}(X) &= \text{softmax}\left( \frac{Q^TK}{\sqrt{d}} \right)V \\ &= \text{softmax}_j\left(\text{score}(x_i, x_j)\right) \end{align}

From this we can derive,

$$ \text{Attn}_{i,j}(X) = \frac{\exp \left( \frac{\langle q_i, k_j \rangle}{\sqrt{d}}\right)}{\sum_k \exp\left( \frac{\langle q_i, k_k \rangle}{\sqrt{d}} \right)}.$$

Now, note, we have turned each column into a probability distribution by applying the softmax so we have

$$ P(X_i) = \text{softmax}\begin{pmatrix} \frac{\langle q_i, k_1 \rangle}{\sqrt{d}}\\ \frac{\langle q_i, k_2 \rangle}{\sqrt{d}}\\ \vdots \\ \frac{\langle q_i, k_n \rangle}{\sqrt{d}} \end{pmatrix}. $$

Now, if we adjust for masked self attention, we can represent the attention mechanism as a directed graph (with weighted self-loops) as the entries in the attention matrix above the diagonal are zero, and so all edges are directed at token nodes that come "after" them.

There is another way to understand attention using graphs when we view attention through the lens of (complete) graph attention networks as explained here.

There is also a visualization you can find here that shows the attention mechanism as a graph. It is equivalent to the graph attention formulation, but is more interactive and pretty.