Maths behind Attention Mechanism

In this article I'm going to explore the mathematics behind attention mechanisms, focusing primarily on text generation models (LLMs). Attention is a crucial component that allows these models to weigh the importance of different words in a sequence when predicting the next word.

Word Embeddings: Converting Words to Vectors

Suppose we have a vocabulary of 7 words. First, let's convert each word into a vector. (these representations are called embeddings.)

\begin{align*} \text{cat} &\mapsto \begin{bmatrix} 2 & 5 & 1 \end{bmatrix} \\ \text{the} &\mapsto \begin{bmatrix} 1 & 0 & 3 \end{bmatrix} \\ \text{sat} &\mapsto \begin{bmatrix} 3 & 4 & 0 \end{bmatrix} \\ \text{on} &\mapsto \begin{bmatrix} 0 & 2 & 4 \end{bmatrix} \\ \text{mat} &\mapsto \begin{bmatrix} 2 & 3 & 2 \end{bmatrix} \\ \text{dog} &\mapsto \begin{bmatrix} 1 & 4 & 2 \end{bmatrix} \\ \text{hat} &\mapsto \begin{bmatrix} 3 & 1 & 3 \end{bmatrix} \end{align*}

Now let's consider the sentence: <kbd>the cat sat on the</kbd>. We want to understand how attention mechanisms help predict the next word.

The Core Attention Equation

The fundamental equation for attention is:

\begin{equation} \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V \end{equation}

Where:

Q (Query): Represents what we're looking for
K (Key): Helps us find relevant information
V (Value): The information we extract once we know where to look
$d_k$: The dimensionality of the key vectors

Attention Mechanisms: Step by Step

In a transformer model, Q, K and V are created by multiplying input embeddings by different learned weight matrics. (X is the original vector.)

\begin{align*} \text{Q}=XW_Q \\ \text{K}=XW_K \\ \text{V}=XW_V \\ \end{align*}

For the simplicity let's assume the weight matrices are identiy matrices. In real transformer models, these matrices are learned during the training to capture different aspects of the input embeddings.

\begin{align*} \text{Q}=X \times 1 \\ \text{K}=X \times 1 \\ \text{V}=X \times 1 \\ \end{align*}

Step 1: Calculate Query Vector

Now we need to calculate Q, K and V for the last word (<b>the</b>) in our sentence. ( <kbd>the cat sat on the</kbd>) According to above assumption Q, K and V vectors are same to each other.

Therefore for the last word "the" query vector is [1,0,3]

Step 2: Compute Dot Products with Key Vectors

With the dot product we can calculate the similarity between "the" word with other words. For that we will calculate the dot product between the query vector and each key vector in our sequence.

\begin{align*} [1,0,3] \cdot [1,0,3] &= 1 \times 1 + 0 \times 0 + 3 \times 3 = 1 + 0 + 9 = 10 \\ [1,0,3] \cdot [2,5,1] &= 1 \times 2 + 0 \times 5 + 3 \times 1 = 2 + 0 + 3 = 5 \\ [1,0,3] \cdot [3,4,0] &= 1 \times 3 + 0 \times 4 + 3 \times 0 = 3 + 0 + 0 = 3 \\ [1,0,3] \cdot [0,2,4] &= 1 \times 0 + 0 \times 2 + 3 \times 4 = 0 + 0 + 12 = 12 \\ [1,0,3] \cdot [1,0,3] &= 1 \times 1 + 0 \times 0 + 3 \times 3 = 1 + 0 + 9 = 10 \end{align*}

Raw attention scores = [10, 5, 3, 12, 10]

Step 3: Scale by $\sqrt{d_k}$

We divide by the square root of the dimension of the key vectors to prevent the dot products from growing too large which could lead to unstable gradients during training.

\begin{equation} \text{Scaled scores} = \frac{Raw scores}{\sqrt{d_k}} = \frac{[10, 5, 3, 12, 10]}{\sqrt{3}} = [5.77, 2.89, 1.73, 6.93, 5.77] \end{equation}

Step 4: Apply softmax to get probability weights

After that apply softmax to get probability weights. Softmax converts these scores into a probability distribution.

\begin{equation} \text{Softmax}([5.77, 2.89, 1.73, 6.93, 5.77]) = [0.1908, 0.0106, 0.0033, 0.6044, 0.1908] \end{equation}

The probability weights show the model pays the most attention to "on" (0.6044) and less to the other words. This makes sense that after "the cat sat on the" words like "mat" are likely to follow. This shows how attention captures the contexual relationships in the sequence.

Step 5: Calculate the Context Vector

After that multiply each value vector by its corresponding attention weight and sum them. With that operation we can capture the weighted importance of all previous words in the sequence.

\begin{equation} \begin{aligned} \text{Context vector} &= 0.1908 \times [1,0,3] + 0.0106 \times [2,5,1] + 0.0033 \times [3,4,0] \\ & + 0.6044 \times [0,2,4] + 0.1908 \times [1,0,3] \\ &= [0.4127, 1.2750, 3.5730] \end{aligned} \end{equation}

Modern transformers use Multi-Head Attention to capture different types of relationships by running attention multiple times in parallel, Positional Encoding to incorporate word order since attention is permutation-invariant, and Masking in decoders to prevent attending to future tokens during training.

References

The Illustrated Transformer