Self-Attention

To understand how self-attention works, let's use a simplified example.

Suppose we have a sequence of tokens, and that our embedding dimension is 2; in real LLMs the embedding dimension is many thousands of dimensions, but here we will use 2 so that we can see what is going on.

Our example sequence

Let's say our sequence of tokens is:

and that, after embedding, our sequence of tokens is:

Three projections

To implement self-attention, we project each token 3 times, using 3 linear projections, called W_query, W_key, and W_value.

These projections are linear neural network layers, and, again for simplification, let's suppose they take our 2 dimensional token embeddings, and project into 3 dimensional space:

import torch.nn as nn

W_query = nn.Linear(2, 3, bias=False)   # initialised with random weights
W_key   = nn.Linear(2, 3, bias=False)   # initialised with random weights
W_value = nn.Linear(2, 3, bias=False)   # initialised with random weights

Note that these are initialised with random weight parameters, which will be optimised during training. For the purposes of this toy example, we'll construct some specific projections that mimic where the weights might end up after training, so that we can illustrate the concepts involved:

The intuition behind the projections

To understand the intuition behind the function of these three projections, we can think of the attention process as each token querying for specific information from the other tokens. In our toy example, after training, our token she might make the following query:

Here every token makes a query, but we are just focussing on the one token of interest, she. The W_query projection is responsible for this; it projects to a representation which "advertises" the desire to find a she-matching referent.

The W_key is the counterpart to this; it projects to a representation which "advertises" features that may want to be matched again. In this example, it could take the form:

Computing attention scores

We use these two projections to compute attention scores as follows:

attention_scores = W_query(x) @ W_key(x).T

First we project every token with W_query and W_key to get its query and key vectors:

Taking she's query and its dot product with every token's key — for example against dog and ball:

score(she, dog)  = [2.64, 1.46, 0.94] · [2.88, 3.06, 1.89]
                 = 2.64×2.88 + 1.46×3.06 + 0.94×1.89  =  13.8

score(she, ball) = [2.64, 1.46, 0.94] · [2.73, 0.81, 1.52]
                 = 2.64×2.73 + 1.46×0.81 + 0.94×1.52  =   9.8

Doing that for all nine keys gives the score row for she:

From scores to attention weights

Next we scale these scores so that they all add up to 1, and call these the attention weights. Doing this is needed to stabilise training and increase likelihood of convergence to a good model:

attention_weights = torch.softmax(attention_scores / 3**0.5, dim=-1)

Here, softmax is a widely used scaling function in machine learning. Note that we also scale the scores before applying softmax, by dividing them by the square root of the dimension of the space. This is a measure that helps to avoid outlying weights being boosted too much by softmax. Applied to she's score row, we get:

Causal attention

In the example so far, every token attends to every other token — including ones that come later in the sentence. But an LLM generates text one token at a time, from left to right: when it is processing a given token, the tokens that come after it don't exist yet. So each token must only attend to itself and the tokens before it.

We enforce this with a causal mask. Laying the attention weights out as a grid — one row per query token, one column per key token — we block out every cell above the diagonal, so no token can attend to a later one:

In practice we apply the mask to the scores just before the softmax, setting every masked position to -inf so that softmax sends its weight to 0 (and the surviving weights still sum to 1):

# block attention to future tokens (everything above the diagonal)
mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()
attention_scores = attention_scores.masked_fill(mask, float('-inf'))

For our query she, this zeroes out the (already tiny) weights on the later tokens was and happy; for a token near the start, like the first the, it leaves nothing to attend to but itself.

The W_value projection

Finally, we come to the W_value projection. We use W_value to project the input, and sum over it using the attention weights we just calculated:

result = attention_weights @ W_value(x)

Here are the value vectors, and the output for she — a weighted sum that is 0.883 × dog + 0.087 × ball + …:

Summary

In plain English, we can think of the whole process as follows. Our input:

"The dog fetched the ball and she was happy"

• she emits a query: "I'm a pronoun, I need my referent."
• dog and ball each emit keys advertising what they are.
• The query·key dot products score dog much higher than ball.
• Softmax turns that into weights like ~0.9 on dog.
• The output for she becomes mostly dog's value — so she is enriched with the meaning of dog.

Below we can see what we've learned so far in a simple implementation that could be plugged into a neural network like an LLM:

import torch
import torch.nn as nn


class SelfAttention(nn.Module):

    def __init__(self, input_dimension, output_dimension):  # in our example the dimensions were 2 and 3
        super().__init__()
        self.W_query = nn.Linear(input_dimension, output_dimension, bias=False)
        self.W_key   = nn.Linear(input_dimension, output_dimension, bias=False)
        self.W_value = nn.Linear(input_dimension, output_dimension, bias=False)

    def forward(self, x):
        keys    = self.W_key(x)
        queries = self.W_query(x)
        values  = self.W_value(x)

        attention_scores = queries @ keys.T

        # causal mask: a token can only attend to itself and earlier tokens
        seq_len = x.shape[0]
        mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()
        attention_scores = attention_scores.masked_fill(mask, float('-inf'))

        attention_weights = torch.softmax(
            attention_scores / keys.shape[-1] ** 0.5, dim=-1
        )

        result = attention_weights @ values
        return result

What's next

The above example covers the main concept behind self-attention, but it is a slightly simplified version. In the next sections we will look at the final pieces which will give us fully-fledged self attention:

• Multi-head attention