Visual Attention Mechanism in PyTorch

Build scaled dot-product and multi-head attention from scratch in PyTorch.

Introduction

30 Second Summary

Every time you ask ChatGPT a question or get a suggestion from Claude, a single mechanism does the heavy lifting behind the scenes. It decides which words in your input matter most for generating each word of the response.

In this project, you will build that mechanism from scratch. You will implement Scaled Dot-Product Attention and Multi-Head Attention using only PyTorch primitives, then visualize exactly how your model pays attention to different parts of a sentence.

What You'll Build

You will type a sentence, run it through your attention module in a Google Colab notebook, and see heatmaps light up showing exactly which words your model considers important when processing each token.

By the end of this project, you'll have:

  • A working Scaled Dot-Product Attention function that you can feed any sentence into and get back contextualized representations of every token.
  • Annotated heatmaps that visualize exactly which tokens attend to which other tokens, revealing how your model interprets relationships between words.
  • A complete Multi-Head Attention module where you can see how different heads specialize in different linguistic patterns, validated against PyTorch's official implementation.
  • Secret Mission: Add causal masking to make your attention autoregressive, the mechanism that lets GPT generate one token at a time.

Are there any prerequisites?

You should be comfortable with Python and have experience working with PyTorch tensors. Familiarity with basic linear algebra concepts like matrix multiplication, dot products, and softmax will help you follow the math behind each operation.

Before We Start

Before We Start...

Before we dive in, let's lock in what you're about to build and why it matters for your understanding of generative AI.

Set Up Your Colab Notebook

Every transformer model you have heard of (GPT, Claude, Gemini) has the same core engine: an attention mechanism that decides which tokens matter for understanding each word. Before you can build that engine from scratch, you need a running environment with all the right libraries ready to go.

In this step, you will open Google Colab, confirm that PyTorch and your visualization libraries are available, and create the input data that flows through every attention computation in this project.

In this step, get ready to:
  • Create a new Colab notebook named "Attention_From_Scratch".
  • Import all required libraries and confirm PyTorch and NumPy versions.
  • Define a sample sentence and create simulated token embeddings as your input tensor.
Open Google Colab and create a new notebook
  • Open your browser and navigate to colab.research.google.com.
  • Sign in with your Google account if prompted.
  • Click New Notebook (or go to File then New notebook).
  • Click the notebook title at the top left (it defaults to something like "Untitled0.ipynb") and rename it to Attention_From_Scratch.

What is Google Colab?

Google Colab is a free cloud-based notebook environment that runs Python code in your browser. It comes with PyTorch, NumPy, matplotlib, and seaborn pre-installed, so there is nothing to download or configure locally.

Import libraries and confirm versions

Your notebook needs access to PyTorch for tensor math, and matplotlib plus seaborn for the heatmap visualizations you will create later.

  • In the first code cell of your notebook, import all required libraries and print version numbers by pasting this code:
# Import PyTorch and visualization libraries
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import math

# Confirm library versions
print(f"PyTorch: {torch.__version__}")
print(f"NumPy: {np.__version__}")
  • Run the cell by pressing Shift+Enter.

What does this code do?

Each import brings in a library you will use throughout the project:

  • torch is the core PyTorch library for tensor operations and matrix math.
  • torch.nn provides neural network building blocks like nn.Linear (the weight matrices you will create in Step 2).
  • torch.nn.functional gives you the official scaled_dot_product_attention function you will validate against in Step 5.
  • matplotlib and seaborn handle all your heatmap visualizations.

You should see output like:

PyTorch: 2.10.0
NumPy: 2.0.2

Your version numbers may differ slightly

Colab updates its runtime regularly. As long as both lines print a version number without errors, you are good to go.

Seeing an ImportError?

All these libraries are pre-installed in Colab's default runtime. If you see an import error, make sure you are using a hosted runtime (not a local one). Check by clicking Runtime in the top menu and selecting Change runtime type. Confirm it is set to the default Python 3 environment.

help me debug an ImportError in Google Colab

Define your input sentence and create token embeddings

Attention operates on sequences of vectors. You need a sample sentence split into tokens, and a tensor of embeddings representing each token as a 64-dimensional vector.

  • Create a new code cell below the first one by clicking the + Code button, then paste this code:
# Define a sample sentence and simulate token embeddings
sentence = "The cat sat on the mat because it was tired"
tokens = sentence.split()
seq_len = len(tokens)
d_model = 64  # embedding dimension

# Set seed for reproducibility and create random embeddings
torch.manual_seed(42)
X = torch.randn(1, seq_len, d_model)  # (batch=1, seq_len=10, d_model=64)
print(f"Input shape: {X.shape}")
print(f"Tokens: {tokens}")
  • Run the cell by pressing Shift+Enter.

What does this code do?

This cell creates the raw material your attention mechanism will process:

  • sentence.split() breaks the sentence into 10 individual tokens (words).
  • d_model = 64 sets the embedding dimension. In real transformers this is often 512 or 768, but 64 keeps things fast and readable.
  • torch.manual_seed(42) ensures everyone gets the same random numbers, so your heatmaps will match the examples.
  • X is a tensor with shape (1, 10, 64). That is 1 batch, 10 tokens, 64 features per token.

What do you expect the output shape to be? Before looking, predict the dimensions for a 10-word sentence with 64 features per word.

You should see:

Input shape: torch.Size([1, 10, 64])
Tokens: ['The', 'cat', 'sat', 'on', 'the', 'mat', 'because', 'it', 'was', 'tired']

The shape [1, 10, 64] confirms you have 1 batch containing 10 tokens, each represented by a 64-dimensional vector. This tensor X is the input to every attention computation you build from here.

Shape looks different?

If seq_len is not 10, check that you copied the sentence exactly (10 words separated by single spaces). If the shape shows different numbers for the last dimension, confirm d_model = 64 is in your cell.

help me fix my input tensor shape

✔️ Awesome, I've got everything!

Great. Double check you have saved your notebook (Ctrl+S or Cmd+S) and that both cells ran without errors.

ⓧ I'd like to double check the full code

Here is everything you should have across both cells in your notebook:

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import math

# --- Setup ---
sentence = "The cat sat on the mat because it was tired"
tokens = sentence.split()
seq_len = len(tokens)
d_model = 64

torch.manual_seed(42)
X = torch.randn(1, seq_len, d_model)
print(f"Input shape: {X.shape}")
print(f"Tokens: {tokens}")

Your environment is ready and your input data is defined. Next up, you will use this tensor to compute Query, Key, and Value vectors and see your first attention heatmap.

Compute QKV Vectors and Raw Attention Scores

Your notebook is ready and your input tensor is defined. Now it's time to build the first real piece of the attention mechanism: the Query, Key, and Value vectors.

Before you can compute attention, you need to understand what these vectors actually are. They are learned linear projections of the input. Queries ask "what am I looking for?", Keys answer "what do I contain?", and Values carry the actual information to pass along.

In this step, get ready to:
  • Create weight matrices and project your input into Q, K, V vectors.
  • Compute raw attention scores by multiplying Q and K together.
  • Observe why unscaled scores are a problem for attention.
Create weight matrices and compute Q, K, V projections

Each of the three vectors (Query, Key, Value) is produced by multiplying the input through its own learned weight matrix. These are standard PyTorch linear layers with no bias.

  • Create a new code cell in your notebook and add the QKV projection code by pasting the following:
# --- Step 2: QKV Vectors and Raw Scores ---
# Create learned projection matrices for Q, K, V
W_q = nn.Linear(d_model, d_model, bias=False)
W_k = nn.Linear(d_model, d_model, bias=False)
W_v = nn.Linear(d_model, d_model, bias=False)

# Project input X through each weight matrix
Q = W_q(X)  # (1, 10, 64)
K = W_k(X)  # (1, 10, 64)
V = W_v(X)  # (1, 10, 64)
print(f"Q shape: {Q.shape}, K shape: {K.shape}, V shape: {V.shape}")

What does this code do?

Each nn.Linear creates a 64x64 weight matrix. When you pass X through it, the matrix multiplication projects your input into a new 64-dimensional space.

  • Q (Query) represents "what is this token looking for?"
  • K (Key) represents "what does this token contain?"
  • V (Value) carries the actual content to pass along when attention is high.
  • Run the cell by pressing Shift+Enter.
  • You should see all three shapes printed as torch.Size([1, 10, 64]).

Seeing a NameError?

  • Make sure you ran the first cell (from Step 1) before this one. Colab cells depend on previous cells having been executed.
  • Check that d_model and X are defined by running your setup cell first.

Why is my Q, K, V cell throwing a NameError?

Compute raw attention scores and visualize them

The attention score between two tokens is simply the dot product of one token's Query with another token's Key. A high dot product means "this token is very relevant to what I'm looking for."

  • Create a new cell and compute the raw scores matrix by pasting this code:
# Compute raw attention scores: dot product of Q and K^T
raw_scores = torch.matmul(Q, K.transpose(-2, -1))  # (1, 10, 10)
print(f"Raw scores shape: {raw_scores.shape}")
print(f"Score range: min={raw_scores.min().item():.2f}, max={raw_scores.max().item():.2f}")

# Visualize raw scores as a heatmap
plt.figure(figsize=(8, 6))
sns.heatmap(raw_scores[0].detach().numpy(), xticklabels=tokens, yticklabels=tokens,
            cmap="rocket", annot=True, fmt=".1f", square=True)
plt.title("Raw Attention Scores (UNSCALED)")
plt.xlabel("Key (attending to)")
plt.ylabel("Query (attending from)")
plt.tight_layout()
plt.show()

What does this code do?

  • torch.matmul(Q, K.transpose(-2, -1)) multiplies every Query vector with every Key vector. The result is a 10x10 matrix where each cell [i, j] tells you how much token i wants to attend to token j.
  • sns.heatmap renders this matrix as a colour grid with your token labels on both axes, so you can visually see which tokens are "looking at" which other tokens.
  • Run the cell with Shift+Enter.
  • You should see the shape printed as torch.Size([1, 10, 10]) and a heatmap with wildly different magnitudes (some scores above 10, some below -10).

Why is the shape 10x10?

You have 10 tokens. Each token's Query is compared against every token's Key. That gives you 10 x 10 = 100 attention scores. This is the "attention map" that decides how information flows between tokens.

Heatmap not displaying?

  • Make sure plt.show() is the last line in the cell. Colab renders plots inline when the cell finishes executing.
  • If you see a blank output, try clicking Runtime then Run all to re-execute all cells in order.

Why isn't my seaborn heatmap showing in Colab?

Apply softmax to unscaled scores and observe the problem

To turn raw scores into proper attention weights (probabilities that sum to 1), you pass them through softmax. But what happens when the raw scores have wildly large magnitudes?

  • Before you run the next cell, predict: if softmax receives scores ranging from -10 to +15, will the resulting attention be spread across multiple tokens or concentrated on just one?
  • Create a new cell and apply softmax to the unscaled scores by pasting this code:
# Apply softmax to unscaled scores and inspect the distribution
unscaled_weights = torch.softmax(raw_scores, dim=-1)
print(f"Unscaled weight distribution for token 'cat':")
print(f"  Max weight: {unscaled_weights[0, 1].max().item():.4f}")
print(f"  Min weight: {unscaled_weights[0, 1].min().item():.6f}")
print("  Notice: weights are extremely peaked (nearly one-hot). The model can barely spread attention.")

What does this code do?

  • torch.softmax(raw_scores, dim=-1) converts each row of raw scores into probabilities that sum to 1. Each row represents how much one token "pays attention" to all other tokens.
  • We look specifically at token index 1 (the word "cat") to see how its attention is distributed across the other 9 tokens.
  • Run the cell with Shift+Enter.

You should see that the max weight is very close to 1.0 and the min weight is nearly zero. The attention is almost entirely concentrated on a single token. This is the problem with unscaled attention.

Why is this a problem?

When attention weights are nearly one-hot, the model can only look at one other token. It cannot blend information from multiple relevant tokens. A good attention mechanism should be able to spread focus across several tokens when multiple are relevant.

The large raw scores push softmax into saturation. This is exactly the problem that "scaling" solves in the next step.

Weights don't look peaked?

  • Make sure you are printing unscaled_weights[0, 1] (token index 1, which is "cat"). Other tokens may show different distributions.
  • Confirm that you applied softmax to raw_scores (not a scaled version). The unscaled scores should produce peaked weights.

Why are my softmax attention weights not showing a peaked distribution?

✔️ Awesome, I've got everything!

Great. Confirm you have three cells for this step, each producing the expected output.

ⓧ I'd like to double check the full code

Scroll through your notebook and confirm you have these three cells in order after your setup cell:

  • Cell 1: QKV projections - defines W_q, W_k, W_v and computes Q, K, V. Prints all three shapes as torch.Size([1, 10, 64]).
  • Cell 2: Raw scores and heatmap - computes raw_scores with torch.matmul(Q, K.transpose(-2, -1)) and displays a 10x10 heatmap titled "Raw Attention Scores (UNSCALED)".
  • Cell 3: Softmax on unscaled scores - computes unscaled_weights and prints the peaked distribution for token "cat".

You have just uncovered the core problem with unscaled attention: large dot products push softmax into saturation, making the model fixate on a single token. In the next step, you will fix this with a simple but powerful trick from the original "Attention Is All You Need" paper.

Apply Scaling and Complete Scaled Dot-Product Attention

In the previous step, you saw that unscaled attention scores produce nearly one-hot weights after softmax. Large dot products push softmax into saturation, so the model can barely spread attention across multiple tokens.

The fix is simple: divide by the square root of the key dimension. This is the "Scaled" in "Scaled Dot-Product Attention." In this step, you will apply that scaling factor, see the distribution spread out, and wrap everything into a reusable function.

In this step, get ready to:
  • Apply the sqrt(d_k) scaling factor and observe how it distributes attention weights.
  • Compute the final attention output and wrap the full mechanism into a reusable function.
  • Visualize the difference between unscaled and scaled attention weights side by side.
Apply the scaling factor and recompute attention weights

The scaling factor is sqrt(d_k) where d_k is the dimensionality of your key vectors (64 in this project). Dividing the raw scores by this value keeps the variance stable so softmax can produce a spread distribution.

Before you run this, think: the unscaled max weight for token "cat" was close to 1.0 (nearly one-hot). Do you think dividing scores by 8.0 will spread that weight more evenly?

  • Create a new cell below your previous code and apply the scaling factor by pasting this code:
d_k = Q.shape[-1]
scale_factor = math.sqrt(d_k)
scaled_scores = raw_scores / scale_factor

attention_weights = torch.softmax(scaled_scores, dim=-1)
print(f"\nScaled weight distribution for token 'cat':")
print(f"  Max weight: {attention_weights[0, 1].max().item():.4f}")
print(f"  Min weight: {attention_weights[0, 1].min().item():.6f}")
print("  Much more distributed! The model can now attend to multiple relevant tokens.")

What does this code do?

  • d_k = Q.shape[-1] grabs the last dimension of Q, which is 64.
  • math.sqrt(d_k) computes the square root (8.0). This is the scaling constant from the "Attention Is All You Need" paper.
  • Dividing raw_scores by 8.0 shrinks their magnitude, preventing softmax from saturating into a one-hot spike.
  • Run the cell.

The max weight for "cat" should be notably lower than the unscaled version (no longer near 1.0). The min weight should be notably higher (no longer near zero). The model can now attend to multiple relevant tokens at once.

Seeing the same peaked distribution?

Make sure you divided by scale_factor (8.0), not by d_k (64) directly. Dividing by 64 would over-compress the scores.

Also confirm you ran all the cells above first (Setup, then Step 2) so raw_scores exists.

Why are my scaled weights still peaked?

Compute the attention output and define a reusable function

Now that you have proper attention weights, the final step of the mechanism is multiplying them by V. This produces a weighted combination of value vectors for each token.

  • In a new cell, compute the attention output and define the complete function by pasting this code:
attention_output = torch.matmul(attention_weights, V)
print(f"Attention output shape: {attention_output.shape}")


def scaled_dot_product_attention(Q, K, V):
    """Scaled Dot-Product Attention from 'Attention Is All You Need'"""
    d_k = Q.shape[-1]
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    weights = torch.softmax(scores, dim=-1)
    output = torch.matmul(weights, V)
    return output, weights

What does this code do?

  • torch.matmul(attention_weights, V) multiplies the (10x10) weight matrix by the (10x64) value matrix. Each token's output is a weighted blend of all value vectors, where the weights came from the scaled softmax.
  • The scaled_dot_product_attention function wraps the entire 3-step process (score, scale+softmax, multiply by V) into a reusable unit you will call in later steps.
  • Run the cell.

You should see Attention output shape: torch.Size([1, 10, 64]). The shape matches your input X. Each of the 10 tokens now has a 64-dimensional representation that incorporates context from the other tokens.

Getting a NameError?

If you see a NameError for attention_weights or V, make sure you ran the scaling cell above first. Colab cells must be run in order.

Help me fix a NameError in my attention output cell.

Visualize scaled vs unscaled attention weights

Numbers tell part of the story, but a side-by-side heatmap makes the difference between unscaled and scaled attention impossible to miss.

  • In a new cell, create the comparison visualization by pasting this code:
fig, axes = plt.subplots(1, 2, figsize=(16, 6))

sns.heatmap(unscaled_weights[0].detach().numpy(), xticklabels=tokens, yticklabels=tokens,
            cmap="rocket", annot=True, fmt=".2f", square=True, ax=axes[0],
            vmin=0, vmax=1)
axes[0].set_title("UNSCALED Attention Weights\n(Nearly one-hot)")
axes[0].set_xlabel("Key")
axes[0].set_ylabel("Query")

sns.heatmap(attention_weights[0].detach().numpy(), xticklabels=tokens, yticklabels=tokens,
            cmap="rocket", annot=True, fmt=".2f", square=True, ax=axes[1],
            vmin=0, vmax=1)
axes[1].set_title("SCALED Attention Weights\n(Distributed attention)")
axes[1].set_xlabel("Key")
axes[1].set_ylabel("Query")

plt.tight_layout()
plt.show()

What does this visualization show?

The left heatmap shows the unscaled weights from Step 2. Most rows have one bright cell and everything else is near zero. The model is forced to pick a single token.

The right heatmap shows the scaled weights. Attention is spread across multiple tokens per row. This lets the model blend information from several relevant positions, which is essential for understanding context.

  • Run the cell.

You should see two heatmaps side by side. The left one has sharp, bright single-cell spikes in each row. The right one shows softer, more distributed color across each row.

Heatmaps look identical?

Make sure the left plot uses unscaled_weights (from Step 2) and the right uses attention_weights (the scaled version from this step). If you accidentally used the same variable for both, they will look the same.

Why do my two heatmaps look the same?

✔️ Awesome, I've got everything!

Great. Make sure all three cells from this step ran without errors before moving on.

ⓧ I'd like to double check the full code

Scroll through your notebook and confirm you have these cells in order:

  • Setup cell: imports, sentence, tokens, X tensor.
  • Step 2 cell 1: W_q, W_k, W_v projections and Q, K, V computation.
  • Step 2 cell 2: raw_scores computation and heatmap.
  • Step 2 cell 3: unscaled_weights with softmax and the peaked distribution print.
  • Step 3 cell 1: d_k, scale_factor, scaled_scores, and attention_weights with scaled distribution print.
  • Step 3 cell 2: attention_output computation and scaled_dot_product_attention function definition.
  • Step 3 cell 3: side-by-side heatmap comparing unscaled vs scaled weights.

You now have a complete, working Scaled Dot-Product Attention implementation. Next, you will extend this into Multi-Head Attention, where multiple attention computations run in parallel to capture different types of token relationships.

Build Multi-Head Attention

Your single-head attention from the previous step can only capture one type of relationship at a time. If the head learns to track adjacent tokens, it loses its ability to link pronouns back to their referents.

Multi-Head Attention solves this by running several attention computations in parallel on different learned subspaces. Each head specializes in a different pattern, then the results are combined into a single output.

In this step, get ready to:
  • Implement a MultiHeadAttention class that splits input across parallel heads.
  • Run the module and confirm output shapes.
  • Visualize each head's attention pattern as separate heatmaps.
Define the MultiHeadAttention class

The idea is to split the 64-dimensional embedding space into 4 heads of 16 dimensions each. Each head runs its own attention on a different subspace of the input.

  • In a new cell in your Colab notebook, define the class structure and head-splitting logic by pasting this code:
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)

    def split_heads(self, x):
        batch_size, seq_len, _ = x.shape
        x = x.view(batch_size, seq_len, self.num_heads, self.d_k)
        return x.transpose(1, 2)

What does this code do?

The __init__ method sets up four weight matrices. W_q, W_k, and W_v project the input into query, key, and value spaces. W_o is the output projection that combines all heads back together.

The split_heads method reshapes a tensor from (batch, seq_len, d_model) into (batch, num_heads, seq_len, d_k). This lets each head operate on its own 16-dimensional slice independently.

  • Still in the same cell, add the forward method directly below the split_heads method:
    def forward(self, x):
        Q = self.split_heads(self.W_q(x))
        K = self.split_heads(self.W_k(x))
        V = self.split_heads(self.W_v(x))

        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        attn_weights = torch.softmax(scores, dim=-1)
        attn_output = torch.matmul(attn_weights, V)

        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(x.shape[0], -1, self.d_model)
        output = self.W_o(attn_output)

        return output, attn_weights

What does the forward method do?

  • Projects input through Q, K, V weight matrices and splits into separate heads.
  • Runs scaled dot-product attention on each head independently (the same operation you built in Step 3, but now per head).
  • Concatenates heads back together with .transpose(1, 2).contiguous().view() and projects through W_o.
  • Returns both the output tensor and the per-head attention weights for visualization.

Why is d_k now 16 instead of 64?

With 4 heads and d_model = 64, each head gets d_k = 64 // 4 = 16 dimensions. The total computation stays the same as single-head attention, but each head now operates in a smaller subspace where it can specialize.

Run the module and inspect outputs
  • In a new cell below the class definition, instantiate the module and run it on your input by pasting this code:
num_heads = 4
mha = MultiHeadAttention(d_model, num_heads)

output, attn_weights = mha(X)
print(f"MHA output shape: {output.shape}")
print(f"Attention weights shape: {attn_weights.shape}")

What does this code do?

You create a 4-head attention module and pass your same input tensor X through it. The module returns both the transformed output and the attention weights from every head.

  • Run the cell.

You should see output shape (1, 10, 64) and attention weights shape (1, 4, 10, 10). The output matches your original input dimensions. The weights now have 4 separate 10x10 attention matrices, one per head.

Seeing an assertion error?

If you see d_model must be divisible by num_heads, make sure your d_model variable is still set to 64 from an earlier cell. Re-run your setup cell if needed.

Help me debug my MultiHeadAttention assertion error

Visualize each head's attention pattern

The real payoff of multi-head attention is that each head learns to focus on different relationships. Before you visualize, think about what you expect to see.

Do you think all 4 heads will show the same attention pattern, or different ones? Why?

  • In a new cell, visualize all 4 heads side by side by pasting this code:
fig, axes = plt.subplots(1, num_heads, figsize=(20, 5))

for head_idx in range(num_heads):
    head_weights = attn_weights[0, head_idx].detach().numpy()
    sns.heatmap(head_weights, xticklabels=tokens, yticklabels=tokens,
                cmap="rocket", square=True, ax=axes[head_idx],
                vmin=0, vmax=head_weights.max(), cbar=False)
    axes[head_idx].set_title(f"Head {head_idx + 1}")
    axes[head_idx].set_xlabel("Key")
    if head_idx == 0:
        axes[head_idx].set_ylabel("Query")
    else:
        axes[head_idx].set_yticklabels([])

plt.suptitle("Multi-Head Attention: Each Head Learns Different Patterns", fontsize=14)
plt.tight_layout()
plt.show()

What does this visualization show?

Each subplot is one head's 10x10 attention weight matrix. Brighter cells mean stronger attention between that query-key pair. Because each head has its own learned W_q, W_k, W_v projections, each head highlights different token relationships.

  • Run the cell.

You should see a row of 4 heatmaps with visibly different patterns. One head might show a strong diagonal (adjacent token attention). Another might have broader, more distributed weights. This is the power of multi-head attention: different heads specialize automatically.

All heads look the same?

Make sure you ran your setup cell with torch.manual_seed(42) earlier. The random initialization determines the initial weight differences between heads. If you re-ran cells in a different order, use Runtime > Restart and run all to reset.

Help me understand why my attention heads show the same pattern

✔️ Awesome, I've got everything!

Great. Double check you have three cells for this step: the class definition, the instantiation and shape check, and the visualization.

ⓧ I'd like to double check the full code

Scroll up through your Step 4 cells and confirm each one is present in this order:

  • Cell 1: The full MultiHeadAttention class definition (__init__, split_heads, and forward methods all inside one class).
  • Cell 2: num_heads = 4, instantiation of mha, and the two print statements confirming output shape (1, 10, 64) and weights shape (1, 4, 10, 10).
  • Cell 3: The plt.subplots(1, num_heads) visualization loop producing 4 side-by-side heatmaps.

You now have a complete Multi-Head Attention module that splits attention across parallel heads, each capturing different token relationships. Next, you will validate that your from-scratch implementation produces the exact same results as PyTorch's official optimized version.

Validate Against PyTorch's Official Implementation

You have built a complete Multi-Head Attention module from scratch. But how do you know the math is actually correct?

PyTorch provides F.scaled_dot_product_attention as an optimized, battle-tested implementation used in production models everywhere. By feeding the same Q, K, V tensors into both your code and PyTorch's, you can verify numerical equivalence and gain confidence your understanding matches reality.

In this step, get ready to:
  • Extract Q, K, V from your MHA module and run both implementations side by side.
  • Compare outputs numerically and confirm they match within floating-point tolerance.
Extract Q, K, V and run both implementations

To make a fair comparison, you need to feed the exact same Q, K, V tensors into both your manual attention math and PyTorch's official function. You will reuse the weight matrices from your mha instance so both paths start from identical inputs.

  • In a new cell in your Colab notebook, paste the following code to extract Q, K, V and compute attention both ways:
with torch.no_grad():
    Q_val = mha.split_heads(mha.W_q(X))
    K_val = mha.split_heads(mha.W_k(X))
    V_val = mha.split_heads(mha.W_v(X))

    # Your implementation
    my_scores = torch.matmul(Q_val, K_val.transpose(-2, -1)) / math.sqrt(mha.d_k)
    my_weights = torch.softmax(my_scores, dim=-1)
    my_output = torch.matmul(my_weights, V_val)

    # PyTorch's official implementation
    official_output = F.scaled_dot_product_attention(Q_val, K_val, V_val)

What does this code do?

torch.no_grad() disables gradient tracking since you are only comparing outputs, not training.

  • mha.split_heads(mha.W_q(X)) projects the input through the same learned weight matrix and reshapes into per-head format. This gives you the exact Q, K, V your MHA module uses internally.
  • The "Your implementation" block manually replays the three operations: score, softmax, weighted sum. This is the same math your forward() method performs.
  • F.scaled_dot_product_attention is PyTorch's optimized kernel that does the same math but dispatches to hardware-accelerated backends (FlashAttention-2 on GPU).
  • Run the cell.

This cell produces no printed output. Both tensors are computed silently. The next cell reveals whether they match.

Compare the results

Before you run the comparison, take a moment to predict: do you think your from-scratch output will match PyTorch's exactly, or will there be floating-point differences?

  • In the next cell, paste this code to print the numerical difference between the two outputs:
max_diff = (my_output - official_output).abs().max().item()
print(f"Maximum absolute difference: {max_diff:.2e}")
print(f"Outputs match: {torch.allclose(my_output, official_output, atol=1e-6)}")

if max_diff < 1e-6:
    print("\nYour from-scratch implementation is numerically equivalent to PyTorch's optimized version!")
    print("In production, you would use F.scaled_dot_product_attention for performance")
    print("(it dispatches to FlashAttention-2 or Memory-Efficient kernels on GPU).")

What does this code do?

  • (my_output - official_output).abs().max() computes the largest element-wise difference between the two tensors. If this is near zero, the implementations agree.
  • torch.allclose(..., atol=1e-6) returns True if every element is within the tolerance. This is PyTorch's standard way to check numerical equivalence despite floating-point rounding.
  • Run the cell.

You should see a maximum absolute difference of essentially zero (in scientific notation, something like 0.00e+00) and Outputs match: True. Your from-scratch attention is numerically equivalent to PyTorch's optimized version.

Why use the official version in production?

Your implementation is mathematically identical, but PyTorch's F.scaled_dot_product_attention automatically dispatches to hardware-optimized backends on GPU. Building from scratch teaches the mechanics. Using the official function gives you production-grade speed.

Seeing a difference larger than 1e-6?

  • Make sure you ran the MHA cell (Step 4) before this cell. The mha variable must exist in your notebook session.
  • Confirm you are dividing by math.sqrt(mha.d_k) (the per-head dimension, 16) not math.sqrt(d_model) (the full model dimension, 64).
  • If you restarted the Colab runtime, re-run all cells from the top. The random seed and weight matrices must be consistent.

Why doesn't my validation show matching outputs?

✔️ Awesome, I've got everything!

Great. Make sure both cells are saved in your notebook.

ⓧ I'd like to double check the full code

Scroll through your notebook and confirm the following cells are present in order:

  • Cell 1: imports and setup (torch.manual_seed(42), X created).
  • Cell 2: QKV projections and raw scores heatmap.
  • Cell 3: Unscaled softmax showing peaked distribution.
  • Cell 4: Scaled attention weights and scaled_dot_product_attention function.
  • Cell 5: Side-by-side comparison heatmap (unscaled vs scaled).
  • Cell 6: MultiHeadAttention class definition.
  • Cell 7: MHA instantiation and per-head heatmap grid.
  • Cell 8: torch.no_grad() block extracting Q_val, K_val, V_val and computing both outputs.
  • Cell 9: max_diff comparison printing Outputs match: True.

You have now proven that your hand-built attention mechanism produces exactly the same results as PyTorch's production implementation. You understand every matrix operation that powers the attention layers inside GPT, Claude, and every modern transformer. Up next: the Secret Mission extends your module with causal masking, the constraint that makes language generation possible.

Secret mission

Add Causal Masking (Making Attention Autoregressive)

Your attention mechanism can see all tokens at once. GPT can't do that. It generates one token at a time, left to right, never peeking at the future. Can you add a causal mask that blocks future tokens from being attended to?

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely in Google Colab with pre-installed libraries, so there are no ongoing costs. Your notebook will persist in your Google Drive until you choose to delete it.

Resources you used:

  • Google Colab notebook Attention_From_Scratch (stored in Google Drive).

Keep everything running

No action needed. Choose this if you plan to revisit and experiment with different sentences, embedding dimensions, or head counts.

  • Your notebook stays in Google Drive and you can reopen it anytime from colab.research.google.com.
  • Colab disconnects idle runtimes automatically, so no compute resources stay active.

Pause - I'll come back to this later

Disconnect the runtime to free up the session, but keep the notebook saved for later.

  • In your Colab notebook, click Runtime in the top menu bar.
  • Click Disconnect and delete runtime.
  • Your code and outputs remain saved in the notebook file on Google Drive.

Delete - I don't want to use this again

Remove the notebook entirely from your Google Drive.

  • In your Colab notebook, click File in the top menu bar.
  • Click Move to trash.
  • To permanently delete it, open Google Drive, navigate to Trash in the left sidebar, and click Empty trash.

That's a Wrap!

That's a Wrap!

Nice work! You just built the core mechanism that powers every modern large language model, entirely from scratch.

You've learned how to:

  • Build Query, Key, and Value projections from scratch and implement Scaled Dot-Product Attention with the sqrt(d_k) scaling factor that prevents softmax saturation.
  • Construct a complete Multi-Head Attention module that splits computation across parallel heads, each learning different attention patterns, and visualize those patterns as annotated heatmaps.
  • Validate your from-scratch implementation against PyTorch's official F.scaled_dot_product_attention to confirm numerical equivalence.
  • Secret Mission: Added causal masking to make attention autoregressive, the key mechanism enabling left-to-right text generation in GPT-style models.

Ready to quiz yourself?