Module 5 — Attention Mechanism

Introduction

The Attention Mechanism is the core innovation behind modern Transformer models. Before attention, Recurrent Neural Networks (RNNs) and LSTMs compressed an entire input sequence into a single fixed-size context vector, making it difficult to retain information from long sequences.

Attention solves this problem by allowing the model to focus on the most relevant parts of the input when generating each output token.

Instead of remembering everything equally, attention learns where to look.

This idea became the foundation for Transformer-based models such as:

• BERT
• GPT
• T5
• RoBERTa
• LLaMA
• Qwen

In this module, you'll learn:

• Why Attention?
• Query (Q)
• Key (K)
• Value (V)
• Attention Scores
• Dot Product Attention
• Additive Attention
• Scaled Dot Product Attention
• Self-Attention
• Cross-Attention
• Multi-Head Attention
• Implement Self-Attention from Scratch

1. Why Attention?

Consider the sentence:

textCopy
1The animal didn't cross the street because it was too tired.

When predicting the meaning of "it", the model should focus on:

textCopy
1animal ✓
2
3street ✗

RNNs struggle to remember distant words because they compress the sentence into a single hidden state.

Attention allows the model to examine all previous words before making a prediction.

Architecture

textCopy
1Input Tokens
2
3↓
4
5Attention
6
7↓
8
9Relevant Words
10
11↓
12
13Prediction

Benefits

• Better long-range understanding
• Parallel computation
• Improved translation
• Better language modeling

2. Query (Q)

A Query represents what the current token is looking for.

Example

textCopy
1"The cat sat on the mat."

Suppose we are processing:

textCopy
1sat

The Query asks:

textCopy
1Which words are important for "sat"?

Python

pythonCopy
1import torch
2
3query = torch.randn(1, 64)
4
5print(query.shape)

Output

textCopy
1torch.Size([1, 64])

3. Key (K)

Every token generates a Key vector.

Keys describe what information each token contains.

Example

textCopy
1Token
2
3↓
4
5Key Vector

Python

pythonCopy
1import torch
2
3key = torch.randn(6, 64)
4
5print(key.shape)

Output

textCopy
1torch.Size([6, 64])

4. Value (V)

Each token also has a Value vector.

Once attention determines which tokens are important, their Value vectors are combined to produce the output.

Example

textCopy
1Token
2
3↓
4
5Value Vector

Python

pythonCopy
1import torch
2
3value = torch.randn(6, 64)
4
5print(value.shape)

5. Attention Score

The attention score measures the similarity between a Query and each Key.

Formula

textCopy
1Score = Q × Kᵀ

Higher score

↓

Higher importance

Python

pythonCopy
1import torch
2
3Q = torch.randn(1, 64)
4
5K = torch.randn(6, 64)
6
7scores = torch.matmul(Q, K.T)
8
9print(scores.shape)

Output

textCopy
1torch.Size([1, 6])

Each score corresponds to one input token.

6. Dot Product Attention

Dot Product Attention computes similarity using the dot product.

Formula

textCopy
1Attention(Q,K,V)
2
3=
4
5Softmax(QKᵀ)V

Implementation

pythonCopy
1import torch
2
3Q = torch.randn(1,64)
4
5K = torch.randn(5,64)
6
7V = torch.randn(5,64)
8
9scores = torch.matmul(Q, K.T)
10
11weights = torch.softmax(scores, dim=-1)
12
13output = torch.matmul(weights, V)
14
15print(output.shape)

Output

textCopy
1torch.Size([1,64])

7. Additive Attention (Bahdanau Attention)

Instead of a dot product, Additive Attention uses a small neural network to compute similarity.

Architecture

textCopy
1Query
2
3↓
4
5Neural Network
6
7↓
8
9Score

Advantages

• Better for small hidden dimensions
• Used in early Seq2Seq models

Example

pythonCopy
1import torch
2import torch.nn as nn
3
4hidden = 128
5
6linear = nn.Linear(hidden * 2, hidden)
7
8query = torch.randn(1, hidden)
9
10key = torch.randn(1, hidden)
11
12combined = torch.cat([query, key], dim=-1)
13
14score = linear(combined)
15
16print(score.shape)

8. Scaled Dot Product Attention

The Transformer uses Scaled Dot Product Attention.

Formula

[ \text{Attention}(Q,K,V)= \text{Softmax} \left( \frac{QK^T} {\sqrt{d_k}} \right)V ]

Why divide by √dk?

Without scaling:

• Dot products become very large
• Softmax becomes saturated
• Gradients become small

Scaling stabilizes training.

Implementation

pythonCopy
1import torch
2import math
3
4Q = torch.randn(4,64)
5
6K = torch.randn(4,64)
7
8V = torch.randn(4,64)
9
10scores = torch.matmul(Q,K.T)
11
12scores = scores / math.sqrt(64)
13
14weights = torch.softmax(scores, dim=-1)
15
16output = torch.matmul(weights,V)
17
18print(output.shape)

9. Self-Attention

Self-Attention allows each token to attend to every other token in the same sequence.

Sentence

textCopy
1The cat chased the mouse

For the token:

textCopy
1cat

Attention may focus on

textCopy
1cat
2
3↓
4
5chased
6
7↓
8
9mouse

Every token generates

• Query
• Key
• Value

from the same input sequence.

Pipeline

textCopy
1Input
2
3↓
4
5Q,K,V
6
7↓
8
9Attention Scores
10
11↓
12
13Softmax
14
15↓
16
17Weighted Values
18
19↓
20
21Output

10. Cross-Attention

Cross-Attention is used in Encoder-Decoder Transformers.

Here:

• Query comes from the Decoder.
• Keys and Values come from the Encoder.

Architecture

textCopy
1Encoder Output
2
3↓
4
5Keys
6Values
7
8↓
9
10Decoder Query
11
12↓
13
14Cross Attention
15
16↓
17
18Decoder Output

Used in

• Machine Translation
• Image Captioning
• Speech Recognition

11. Multi-Head Attention

Instead of one attention operation, Transformers perform multiple attention operations in parallel.

Example

textCopy
1Head 1
2
3↓
4
5Head 2
6
7↓
8
9Head 3
10
11↓
12
13Head 4
14
15↓
16
17Concatenate
18
19↓
20
21Linear Layer

Benefits

• Captures multiple relationships
• Learns syntax and semantics simultaneously
• Better contextual representations

PyTorch Example

pythonCopy
1import torch
2import torch.nn as nn
3
4attention = nn.MultiheadAttention(
5    embed_dim=128,
6    num_heads=8,
7    batch_first=True
8)
9
10x = torch.randn(2,10,128)
11
12output, weights = attention(x,x,x)
13
14print(output.shape)

Output

textCopy
1torch.Size([2,10,128])

Practice — Implement Self-Attention from Scratch

pythonCopy
1import torch
2import torch.nn as nn
3import math
4
5class SelfAttention(nn.Module):
6
7    def __init__(self, embed_dim):
8
9        super().__init__()
10
11        self.query = nn.Linear(embed_dim, embed_dim)
12
13        self.key = nn.Linear(embed_dim, embed_dim)
14
15        self.value = nn.Linear(embed_dim, embed_dim)
16
17    def forward(self, x):
18
19        Q = self.query(x)
20
21        K = self.key(x)
22
23        V = self.value(x)
24
25        scores = torch.matmul(Q, K.transpose(-2, -1))
26
27        scores = scores / math.sqrt(K.size(-1))
28
29        weights = torch.softmax(scores, dim=-1)
30
31        output = torch.matmul(weights, V)
32
33        return output
34
35
36model = SelfAttention(64)
37
38x = torch.randn(2,8,64)
39
40output = model(x)
41
42print(output.shape)

Output

textCopy
1torch.Size([2,8,64])

What You'll Learn

• Generate Query, Key, and Value vectors
• Compute attention scores
• Apply Softmax to obtain attention weights
• Produce context-aware output representations
• Understand the core computation inside every Transformer layer

Module Summary

After completing this module, you will be able to:

• Explain why the Attention Mechanism was introduced.
• Understand the roles of Query, Key, and Value vectors.
• Compute attention scores using dot products.
• Differentiate between Dot Product and Additive Attention.
• Implement Scaled Dot Product Attention.
• Explain how Self-Attention captures contextual relationships.
• Understand Cross-Attention in encoder-decoder architectures.
• Explain why Multi-Head Attention improves model performance.
• Build a complete Self-Attention layer from scratch using PyTorch.

Next Module: Module 6 – Transformer Architecture, where you'll build every Transformer component from scratch, including embeddings, positional encoding, multi-head attention, feed-forward networks, encoder blocks, decoder blocks, layer normalization, residual connections, and the complete Transformer model.