Attention

Creator: Chung-En Johnny Yu

Content update: 2025/10/04

Source:

• Build a Large Language Model From Scratch by Sebastian Raschka - Ch3

Hands-on practice this notebook on your Google Colab:

1. Click here to open this notebook

2. Now, run the code and practice it!

Additional resources (not included here):

• Attention as Soft Dictionary Lookup by Yuan Meng

• Understanding and Coding Self-Attention, Multi-Head Attention, Causal-Attention, and Cross-Attention in LLMs by Sebastian Raschka

• Transformers Laid Out by Pramod

Modeling long sequences problem¶

Translating a text word by word isn’t feasible due to the differences in grammatical structures between the source and target languages.

Recurrent Neural Networks (RNNs)

• Prior to the introduction of transformer models, encoder-decoder RNNs were commonly used for machine translation tasks.

• In this setup, the encoder processes a sequence of tokens from the source language, using a hidden state—a kind of intermediate layer within the neural network—to generate a condensed representation of the entire input sequence.

• Problem: Can’t handle long context.

  • It can’t directly access earlier hidden states from the encoder during the decoding phase. Consequently, it relies solely on the current hidden state, which encapsulates all relevant information. This can lead to a loss of context, especially in complex sentences where dependencies might span long distances.

Capturing data dependencies with attention¶

Through an attention mechanism, the text-generating decoder segment of the network is capable of selectively accessing all input tokens, implying that certain input tokens hold more significance than others in the generation of a specific output token.

Self-attention (More in the next section)

Self-attention in transformers is a technique designed to enhance input representations by enabling each position in a sequence to engage with and determine the relevance of every other position within the same sequence.

Self-attention without trainable weights¶

Simple illustration¶

This section:

• Explain a very simplified variant of self-attention, which does not contain any trainable weights.

• Purely for illustration purposes and NOT the attention mechanism that is used in transformers.

Problem setup in math for the following figure:

• Suppose we are given an input sequence $x^{(1)}$ to $x^{(T)}$:

  • The input is a text (for example, a sentence like “Your journey starts with one step”) that has already been converted into token embeddings.

  • For instance, $x^{(1)}$ is a d-dimensional vector representing the word “Your”, and so forth.

Goal:

• Compute context vectors $z^{(i)}$ for each input sequence element $x^{(i)}$ in $x^{(1)}$ to $x^{(T)}$ (where $z$ and $x$ have the same dimension).

  • A context vector $z^{(i)}$ is a weighted sum over the inputs $x^{(1)}$ to $x^{(T)}$.

  • The context vector is “context”-specific to a certain input.

    • Instead of $x^{(i)}$ as a placeholder for an arbitrary input token, let’s consider the second input, $x^{(2)}$.

    • And to continue with a concrete example, instead of the placeholder $z^{(i)}$, we consider the second output context vector, $z^{(2)}$.

    • The second context vector, $z^{(2)}$, is a weighted sum over all inputs $x^{(1)}$ to $x^{(T)}$ weighted with respect to the second input element, $x^{(2)}$.

    • The attention weights are the weights that determine how much each of the input elements contributes to the weighted sum when computing $z^{(2)}$.

    • In short, think of $z^{(2)}$ as a modified version of $x^{(2)}$ that also incorporates information about all other input elements that are relevant to a given task at hand.

(Please note that the numbers in this figure are truncated to one digit after the decimal point to reduce visual clutter; similarly, other figures may also contain truncated values)

Step 1: Compute unnormalized attention scores ($\omega$)

• By convention, the unnormalized attention weights are referred to as “attention scores” ($\omega$) whereas the normalized attention scores, which sum to 1, are referred to as “attention weights” ($\alpha$).

• Suppose we use the second input token as the query, that is, $q^{(2)} = x^{(2)}$, we compute the unnormalized attention scores via dot products:

  • $\omega_{21} = x^{(1)} q^{(2)\top}$

  • $\omega_{22} = x^{(2)} q^{(2)\top}$

  • $\omega_{23} = x^{(3)} q^{(2)\top}$

  • ...

  • $\omega_{2T} = x^{(T)} q^{(2)\top}$

  • The subscript “21” in $\omega_{21}$ means that input sequence element 2 was used as a query against input sequence element 1.

The figure depicts the initial step in this process, which involves calculating the attention scores ω between $x^{(2)}$ and all other input elements through a dot product operation.

Step 2: Normalize the unnormalized attention scores (“omegas”, $\omega$) to get the attention weights ($\alpha$).

• Here is a simple way to normalize the unnormalized attention scores to sum up to 1 (a convention, useful for interpretation, and important for training stability).

• However, in practice, using the softmax function for normalization, which is better at handling extreme values and has more desirable gradient properties during training, is common and recommended.

Step 3: Compute the context vector $z^{(2)}$ by multiplying the embedded input tokens, $x^{(i)}$ with the attention weights and sum the resulting vectors.

Attention weights for ALL input tokens¶

Now, we are generalizing this computation to compute all attention weights and context vectors.

Apply previous step 1 to all pairwise elements to compute the unnormalized attention score matrix.

Similar to step 2 previously, we normalize each row so that the values in each row sum to 1.

Quick verification that the values in each row indeed sum to 1.

Apply previous step 3 to compute all context vectors.

As a sanity check, the previously computed context vector $z^{(2)} = [0.4419, 0.6515, 0.5683]$ can be found in the 2nd row in above.

Self-attention with trainable weights¶

Step by step¶

This self-attention mechanism is also called “scaled dot-product attention”.

• The most notable difference is the introduction of weight matrices that are updated during model training.

• These trainable weight matrices are crucial so that the model (specifically, the attention module inside the model) can learn to produce “good” context vectors.

We will start by introducing the three training weight matrices $W_q$, $W_k$, and $W_v$.

• These three matrices are used to project the embedded input tokens, x^{(i)} , into query, key, and value vectors via matrix multiplication:

  • Query vector: $q^{(i)} = x^{(i)}\,W_q$

  • Key vector: $k^{(i)} = x^{(i)}\,W_k$

  • Value vector: $v^{(i)} = x^{(i)}\,W_v$

• The embedding dimensions of the input $x$ and the query vector $q$ can be the same or different, depending on the model’s design and specific implementation

• In GPT models, the input and output dimensions are usually the same, but for illustration purposes, to better follow the computation, we choose different input and output dimensions.

Step 1, we initialize the three weight matrices.

(Note that we are setting requires_grad=False to reduce clutter in the outputs for illustration purposes, but if we were to use the weight matrices for model training, we would set requires_grad=True to update these matrices during model training.)

Next we compute the query, key, and value vectors.

As we can see below, we successfully projected the 6 input tokens from a 3D onto a 2D embedding space.

Step 2, we compute the unnormalized attention scores by computing the dot product between the query and each key vector.

Since we have 6 inputs, we have 6 attention scores for the given query vector.

Step 3, we compute the attention weights (normalized attention scores that sum up to 1) using the softmax function we used earlier.

• The difference to earlier is that we now scale the attention scores by dividing them by the square root of the embedding dimension, \sqrt{d_k} (i.e., d_k**0.5 ).

• The reason for the normalization by the embedding dimension size is to improve the training performance by avoiding small gradients that lead to gradient vanishing.

Step 4, we now compute the context vector for input query vector 2.

Putting all in SelfAttention class¶

Putting it all together, we can implement the self-attention mechanism as follows:

We can streamline the implementation above using PyTorch’s Linear layers, which are equivalent to a matrix multiplication if we disable the bias units.

Another big advantage of using nn.Linear over our manual nn.Parameter(torch.rand(...) approach is that nn.Linear has a preferred weight initialization scheme, which leads to more stable model training.

Hiding future words with causal attention¶

In causal attention, the attention weights above the diagonal are masked, ensuring that for any given input, the LLM is unable to utilize future tokens while calculating the context vectors with the attention weight.

Causal attention mask¶

Causal self-attention ensures that the model’s prediction for a certain position in a sequence is only dependent on the known outputs at previous positions, not on future positions.

To illustrate and implement causal self-attention, let’s work with the attention scores and weights from the previous section:

Create a mask via PyTorch’s tril function is the simplest way.

Then, we can multiply the attention weights with this mask to zero out the attention scores above the diagonal.

However, if the mask were applied after softmax, like above, it would disrupt the probability distribution created by softmax. (Softmax ensures that all output values sum to 1)

• In simpler terms, after masking and renormalization, the distribution of attention weights is as if it was calculated only among the unmasked positions to begin with. This ensures there’s no information leakage from future (or otherwise masked) tokens as we intended.

To make sure that the rows sum to 1 (prevent information leak), we can normalize the attention weights as follows:

Be more efficent, instead of zeroing out attention weights above the diagonal and renormalizing the results, we can mask the unnormalized attention scores above the diagonal with negative infinity before they enter the softmax function:.

• Softmax function treats $-∞$ as 0 during calculation.

As we can see below, now the attention weights in each row correctly sum to 1 again:

Masking weights with dropout¶

In addition, we also apply dropout to reduce overfitting during training.

• Dropout can be applied in several places:

  1. for example, after computing the attention weights (more common);

  2. or after multiplying the attention weights with the value vectors.

If we apply a dropout rate of 0.5 (50%), the non-dropped values will be scaled accordingly by a factor of 1/0.5 = 2.

• The scaling is calculated by the formula 1 / (1 - dropout_rate).

Compact causal self-attention class¶

One more thing is to implement the code to handle batches consisting of more than one input so that our CausalAttention class supports the batch outputs produced by the data loader.

• Note that dropout is only applied during training, not during inference

Multi-head attention¶

Stacking single-head attention¶

Single-head attention (What we just went through.)

Multi-head attention: Stack multiple single-head attention.

• The main idea is to run the attention mechanism multiple times (in parallel) with different, learned linear projections. This allows the model to jointly attend to information from different representation subspaces at different positions.

The following code uses the previous CausalAttention class for illustration. (In the next section will reveal more.)

In the implementation above, the embedding dimension is 4, because we d_out=2 as the embedding dimension for the key, query, and value vectors as well as the context vector. And since we have 2 attention heads, we have the output embedding dimension 2*2=4.

More efficient with weight splits¶

Now, we don’t concatenate single attention heads for this stand-alone MultiHeadAttention class. Instead, we create single W_query, W_key, and W_value weight matrices and then split those into individual matrices for each attention head:

Note that in addition, we added a linear projection layer (self.out_proj ) to the MultiHeadAttention class above. This is simply a linear transformation that doesn’t change the dimensions. It’s a standard convention to use such a projection layer in LLM implementation, but it’s not strictly necessary. (Recent research has shown that it can be removed without affecting the modeling performance.)

• Check out torch.nn.MultiheadAttention class in PyTorch.

Since the above implementation may look a bit complex at first glance, let’s look at what happens when executing attn_scores = queries @ keys.transpose(2, 3):

In this case, the matrix multiplication implementation in PyTorch will handle the 4-dimensional input tensor so that the matrix multiplication is carried out between the 2 last dimensions (num_tokens, head_dim) and then repeated for the individual heads.

• For instance, the following becomes a more compact way to compute the matrix multiplication for each head separately:

Prepare for interview¶

• Code from scratch: SelfAttention_v2, CausalAttention, MultiHeadAttention.