Scaled Dot-Product Attention in PyTorch

Introduction

In our previous lesson, we mapped out the high-level architecture of the original Transformer, identifying its encoder-decoder structure and the roles of its core components like self-attention, positional encodings, and feed-forward networks. You now have the blueprint.

Today, we drill down into the engine that powers the Transformer: scaled dot-product attention. Our goal is to move beyond the conceptual block diagram and implement this mechanism from scratch in PyTorch. This is a crucial step towards deep mastery, as it demystifies how the model learns to weigh the importance of different tokens in a sequence. By the end of this lesson, you will have translated the famous attention formula into working code.

The Attention Formula: A Retrieval System Analogy

At its core, the attention mechanism functions like a soft, differentiable database retrieval system. For every token in your sequence, you want to retrieve a piece of information from the entire sequence (including itself) to enrich its own representation.

This process involves three learned projections of each token's embedding:

1. Query (Q): A vector representing what a token is "looking for."
2. Key (K): A vector representing what a token "contains" or what it's about.
3. Value (V): A vector representing the information a token will "provide" if it's attended to.

The interaction is simple: for a given token's Query, we calculate a score by taking its dot product with every other token's Key. A high score means a strong match. These scores are then normalized to become weights, which are used to compute a weighted sum of all the Value vectors.

This entire process is captured in a single, elegant formula.

The mathematical expression for this, as defined in the original "Attention Is All You Need" paper, is:

$\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$

Let's break down how to implement each part of this formula in PyTorch.

Implementing Scaled Dot-Product Attention Step-by-Step

We'll assume we have our Q, K, and V matrices ready. In a real model, these are produced by passing the input embeddings through separate linear layers. For our implementation, we'll start with random tensors with the shape (batch_size, seq_len, d_k), where d_k is the dimension of the keys and queries.

Let's start by setting up our tensors:

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

# Let's define our dimensions
batch_size = 4
seq_len = 10  # Number of tokens in the sequence
d_k = 64      # Dimension of Query and Key vectors
d_v = 64      # Dimension of Value vector (often same as d_k)

# Initialize Q, K, V with random data
Q = torch.randn(batch_size, seq_len, d_k)
K = torch.randn(batch_size, seq_len, d_k)
V = torch.randn(batch_size, seq_len, d_v)

print(f"Shape of Q: {Q.shape}")
print(f"Shape of K: {K.shape}")
print(f"Shape of V: {V.shape}")

Now, let's implement the formula.

1. Calculate Raw Attention Scores: $Q{K}^T$

The first step is to compute the dot product between each query and all keys. This gives us a similarity score. For a batch of sequences, we can do this efficiently with a batched matrix multiplication.

• Q has shape (batch_size, seq_len, d_k)
• K has shape (batch_size, seq_len, d_k)

To multiply them, we need to transpose the last two dimensions of K to get a shape of (batch_size, d_k, seq_len).

# Transpose K to get shape (batch_size, d_k, seq_len)
K_transposed = K.transpose(-2, -1)

# Perform batched matrix multiplication
# (batch, seq_len, d_k) @ (batch, d_k, seq_len) -> (batch, seq_len, seq_len)
scores = torch.matmul(Q, K_transposed)

print(f"Shape of scores: {scores.shape}")

The resulting scores tensor has a shape of (batch_size, seq_len, seq_len). You can interpret scores[b, i, j] as the raw attention score from token i to token j in the b-th sequence of the batch.

2. Scale the Scores: $/\sqrt{d_k}$

The original authors found that for large values of d_k, the dot products could grow very large in magnitude. This pushes the softmax function into regions where its gradients are extremely small, hindering learning. To counteract this, we scale the scores by dividing by the square root of the key dimension, $\sqrt{d_k}$.

This is a critical, yet simple, step for stabilizing training.

# Scale the scores
d_k = K.size(-1)
scaled_scores = scores / math.sqrt(d_k)

3. Normalize with Softmax

Next, we apply the softmax function to the scaled scores. This converts the scores into a probability distribution, ensuring the weights for each query sum to 1. We apply it along the last dimension (dim=-1), which corresponds to the keys we are attending to.

# Apply softmax to get attention weights
# The weights for each query (each row) will sum to 1
attention_weights = F.softmax(scaled_scores, dim=-1)

print(f"Shape of attention_weights: {attention_weights.shape}")
# Verify that weights for one query sum to 1
print(f"Sum of weights for first query in first batch: {attention_weights[0, 0].sum()}")

4. Compute the Weighted Output: $...$V

Finally, we use the attention_weights to compute a weighted sum of the V vectors. This is another batched matrix multiplication.

# (batch, seq_len, seq_len) @ (batch, seq_len, d_v) -> (batch, seq_len, d_v)
output = torch.matmul(attention_weights, V)

print(f"Shape of final output: {output.shape}")

The output tensor now contains the new representations for each token in the sequence. Each vector output[b, i] is a mixture of all value vectors in the sequence V[b, :, :], weighted by how much token i decided to pay attention to them.

Bringing It All Together

For practical use, we encapsulate this logic into a reusable module. The following videos are excellent resources to see this process in action and gain a deeper intuition. Your background in production ML will make the class-based structure immediately familiar.

To achieve the "deep mastery" you're aiming for, understanding the why is as important as the how. The next video from Andrej Karpathy is a masterclass in building intuition for self-attention. It explains why matrix multiplication serves as an efficient mechanism for "weighted aggregation" and provides powerful analogies for Q, K, and V.

Finally, for a clean, text-based reference, you can consult this article. It provides a complete, modular PyTorch implementation that's easy to read and adapt.

Conclusion

Today, you have moved from theory to practice by implementing the scaled dot-product attention mechanism. This single operation is the fundamental building block that enables Transformers to dynamically route information and capture complex dependencies within data.

Key Takeaways:

• Q, K, V Analogy: Attention can be intuitively understood as a retrieval system where queries fetch information from a set of key-value pairs.
• Implementation Steps: The mechanism is implemented in four steps: calculating similarity scores (matmul(Q, K.T)), scaling (/sqrt(d_k)), normalizing with softmax, and calculating a weighted sum of values (matmul(weights, V)).
• Tensor Shapes: You can now trace the flow of tensors through the mechanism, understanding how shapes change from (batch, seq, dim) to the final output.
• The Role of Scaling: Scaling is not an arbitrary detail; it's a crucial technique for stabilizing training by preventing softmax from saturating.

What's Next?

The scaled dot-product attention we implemented is powerful, but it's just a single "head." What if a token needs to focus on different things simultaneously (e.g., the subject of the sentence, the tense of the verb, and a related adjective)? In our next lesson, we will build upon today's work to implement multi-head attention, which allows the model to learn different representation subspaces in parallel, dramatically increasing its expressive power.