Attention

Attention is the first operation performed inside a transformer layer. It enriches the hidden state of each token with information retrieved from the previous tokens in the sequence.

Why attention is needed

Let’s take this sequence as an example:

The cat sat on the mat.

When the model processes the token mat, it should understand that:

  • cat is the subject,
  • sat describes the action,
  • mat is where that action happened.

Those relationships are not stored in the embedding but are discovered by the attention mechanism. Attention allows each token to gather information from the rest of the context.

Query, Key and Value

To perform that retrieval efficiently, every token produces three vectors from its current hidden state:

  • Query (Q) describes the information the token is looking for
  • Key (K) determines how other tokens can find this token
  • Value (V) contains the information that will actually be retrieved

Query, Key, and Value vectors are computed using learned projection matrices inside every transformer layer.

The following diagram illustrates this process. The hidden state of mat produces its Query vector. Then, since each processed token has already produced its own Key and Value vectors, the Query is compared against each Key to determine which tokens are most relevant.

Those relevance scores are then used to combine the corresponding Value vectors into a single output. The result is a new vector containing contextual information gathered from the rest of the sequence.

This diagram also shows that Attention does not replace the hidden state, but instead it produces an attention update that is added to the current hidden state. The hidden state then continues through the transformer layer, where it is further refined by the MLP.

In practice, Transformers perform several attention computations in parallel, called attention heads. This allows the model to capture different kinds of relationships simultaneously.

What comes next?

Attention computes new Query, Key, and Value vectore for every token. During inference, the Key and Value vectors are stored and reused rather than recomputed. This optimization is explained in the KV Cache article. But, first we’ll have a closer look into the MLP step.