Scaled Dot-Product Attention in PyTorch

Hello! Welcome back.

In our last lesson, we explored tokenization, the crucial first step in preparing text for a Transformer model. We saw how subword tokenization algorithms like BPE and WordPiece create a vocabulary that balances richness and efficiency, turning raw text into sequences of numerical IDs.

Today, we'll build directly on that foundation. Once we have our token embeddings (the vector representations of our tokens), the model needs a way to process them. This is where the attention mechanism comes in. You were introduced to the theory of attention in Module 7; now it's time to get our hands dirty and see how it works in practice.

Our learning goal for this lesson is to compute scaled dot-product attention step-by-step in PyTorch on a small example. We will break down the famous formula from the "Attention Is All You Need" paper into its constituent parts and implement each one.

From Averages to Attention: The Core Idea

At its heart, attention is a mechanism that allows a model to perform a weighted aggregation of information. For any given token, it can "look" at all other tokens in the sequence and decide which ones are most relevant, assigning higher weights to them.

A simple form of aggregation is just taking an average. However, this is too simplistic as it treats all tokens as equally important. Self-attention makes this process data-dependent and far more powerful. To build the intuition from a simple average up to the matrix operations of attention, let's watch a segment from Andrej Karpathy's excellent "Let's build GPT" video.

As you saw, the key ideas are:

1. A weighted sum of token information can be performed efficiently using a matrix multiplication.
2. Self-attention makes these weights data-dependent. Instead of a fixed averaging matrix, the weights are dynamically calculated based on the relationships between the tokens themselves.
3. This is achieved by projecting each input token's embedding into three separate vectors: a Query, a Key, and a Value.

Let's look at a diagram that lays out the entire process.

Implementing Scaled Dot-Product Attention in PyTorch

We are now ready to implement the diagram above step-by-step. We'll use a simple example: a sequence of 5 tokens, each represented by a 3-dimensional embedding vector.

For this walkthrough, we will follow the structure of a very clear tutorial. You can refer to it alongside my explanations.

Setup: Input Embeddings

First, let's create our input data. We'll have a batch of 1 sentence, with 5 words (tokens), and each token will have an embedding of dimension 3. The shape of our input X will be (5, 3). In PyTorch:

import torch
import torch.nn.functional as F
import math




# For reproducibility
torch.manual_seed(42)




# Our input sequence of 5 words, each with a 3-dimensional embedding
X = torch.rand(5, 3) 
print("Input Embeddings (X):")
print(X)
print("Shape of X:", X.shape)

Step 1: Create Query, Key, and Value Vectors

The first step is to project our input embeddings X into Q, K, and V matrices. This is done using three separate, learnable linear layers (or weight matrices). For this example, we'll create the weight matrices W_q, W_k, and W_v manually.

The operation is a matrix multiplication: $Q=X{W}_q$, $K=X{W}_k$, and $V=X{W}_v$.

embedding_dim = 3




# Create random weight matrices (these would be learned during training)
W_q = torch.randn(embedding_dim, embedding_dim)
W_k = torch.randn(embedding_dim, embedding_dim)
W_v = torch.randn(embedding_dim, embedding_dim)




# Project embeddings into Q, K, and V
Q = X.matmul(W_q)
K = X.matmul(W_k)
V = X.matmul(W_v)

print("Q matrix shape:", Q.shape)
print("K matrix shape:", K.shape)
print("V matrix shape:", V.shape)

Notice that the shapes of Q, K, and V are the same as our input X, which is (5, 3).

Step 2: Calculate Attention Scores

Now, we calculate the "affinity" or "compatibility" scores between each token's query and every other token's key. This is done with a dot product between the Query matrix and the transpose of the Key matrix.

$$\text{scores}=Q{K}^T$$

# Calculate dot product of Q and K.T
scores = Q.matmul(K.T)

print("Scores matrix shape:", scores.shape)
print("Scores matrix:\n", scores)

The resulting scores matrix has a shape of (5, 5). The element at scores[i, j] represents the raw attention score from token i to token j.

Step 3: Scale the Scores

As discussed in the "Attention Is All You Need" paper and Karpathy's video, the dot products can grow large, pushing the softmax function into regions with tiny gradients. To counteract this, we scale the scores by dividing by the square root of the dimension of the key vectors, $d_k$. In our case, $d_k=3$.

$$\text{scaled\_scores}=\frac{Q{K}^T}{\sqrt{d_k}}$$

d_k = K.shape[-1] # Dimension of key vectors




# Scale the scores
scaled_scores = scores / math.sqrt(d_k)

print("Scaled scores matrix:\n", scaled_scores)

This scaling helps stabilize training by keeping the variance of the scores under control.

Step 4: Apply Softmax to Get Attention Weights

Next, we apply the softmax function along the last dimension of the scaled_scores matrix. This converts the raw scores into a probability distribution. Each row will now sum to 1.

$$\text{attention\_weights}=\text{softmax}(\frac{Q{K}^T}{\sqrt{d_k}})$$

# Apply softmax to get attention weights
attention_weights = F.softmax(scaled_scores, dim=-1)

print("Attention weights matrix shape:", attention_weights.shape)
print("Attention weights matrix:\n", attention_weights)



# Verify that each row sums to 1
print("\nSum of first row:", attention_weights[0].sum())

The attention_weights matrix, still of shape (5, 5), now tells us how much "attention" each token should pay to every other token when creating its new representation.

Step 5: Compute the Final Output

Finally, we compute the output of the attention layer by multiplying the attention_weights by the Value matrix V.

$$\text{Output}=\text{softmax}(\frac{Q{K}^T}{\sqrt{d_k}})V$$

This operation creates a weighted sum of the value vectors for each token, where the weights are the attention scores we just computed.

# Multiply attention weights by V
output = attention_weights.matmul(V)

print("Output matrix shape:", output.shape)
print("Output matrix:\n", output)

The output matrix has the same shape as our original input, (5, 3). Each row output[i] is the new, context-aware representation for the i-th token. It has aggregated information from all other tokens in the sequence, weighted by their relevance.

To see this entire process demonstrated live, with a focus on the tensor shapes and dimensions, please watch the following clip.

Conclusion

In this lesson, we demystified scaled dot-product attention by implementing it from scratch in PyTorch. We translated the mathematical formula into a series of concrete matrix operations on a small example, bridging the gap between theory and code.

Key Takeaways:

• Scaled dot-product attention computes a context-aware representation for each token in a sequence.
• It begins by projecting input embeddings into Query (Q), Key (K), and Value (V) matrices.
• The core computation follows five steps:
  1. Calculate scores: $Q{K}^T$.
  2. Scale scores: Divide by $\sqrt{d_k}$ to stabilize training.
  3. Apply softmax: Convert scores to attention weights (a probability distribution).
  4. Multiply by Value: Compute a weighted sum of the value vectors.
• The final output is a new set of embeddings of the same shape as the input, but with contextual information mixed in.

Preview of the Next Lesson

While building attention from scratch is essential for understanding, in practice, we rarely do this. High-level libraries like HuggingFace Transformers abstract these details away. In our next lesson, we will learn how to load a pre-trained model and its corresponding tokenizer from the HuggingFace Transformers library, moving from low-level mechanics to high-level application.