LLM chevron_right The Transformer Architecture: Attention, Multi-Head Attention, and Positional Encoding

The Transformer Architecture: Attention, Multi-Head Attention, and Positional Encoding

calendar_today schedule 6 min read

Learning Objectives: By the end of this lesson, you will be able to:

  • Explain the motivation for the Transformer and the self-attention mechanism
  • Implement a simplified self-attention layer and understand multi-head attention
  • Describe the role and implementation of positional encoding

1. Welcome and Introduction (5 minutes)

Icebreaker Activity

Have you ever wondered how modern AI models like ChatGPT can “pay attention” to the right words in a sentence, no matter how long it is? Why do you think older models struggled with long texts?

Lesson Overview

  • Duration: Approximately 2.5 hours including all exercises (or 60–90 minutes when selected activities are assigned)
  • Prerequisites: Python, basic understanding of neural networks, previous lessons on sequence models
  • Materials Needed: Python (Jupyter Notebook or IDE), NumPy or PyTorch, sample text data

2. Learning Objectives and Session Outline (5 minutes)

What We’ll Cover Today:

  1. The limitations of RNNs and the need for attention
  2. The self-attention mechanism (Query, Key, Value, scaled dot-product)
  3. Multi-head attention and its benefits
  4. Positional encoding and why it’s necessary
  5. Hands-on practice and exercises
  6. Common pitfalls and how to avoid them

Success Criteria:

By the end of this session, you should be able to explain and implement the core components of the Transformer encoder.

3. Foundation Knowledge (15-20 minutes)

Core Concepts

Why Move Beyond RNNs?

RNNs process sequences step-by-step, making them slow and unable to capture long-range dependencies efficiently. Attention mechanisms allow models to focus on relevant parts of the input, regardless of their position.

Real-World Analogy: Imagine reading a book and being able to instantly reference any page, not just the previous one. Attention lets models “look anywhere” in the input.

Why This Matters: Attention enables parallel processing and better handling of long sequences, which is crucial for modern NLP tasks.

Self-Attention Mechanism

Self-attention computes a weighted sum of all input tokens for each token, allowing the model to capture dependencies between any pair of words.

Multi-Head Attention

Instead of a single attention calculation, multi-head attention runs several in parallel, allowing the model to capture different types of relationships.

Positional Encoding

Since Transformers have no recurrence, they need a way to encode the order of tokens. Positional encoding adds this information to the input embeddings.

Key Terminology

  • Query, Key, Value: Vectors derived from input embeddings used in attention calculation
  • Scaled Dot-Product Attention: The core operation in self-attention
  • Head: An independent attention calculation in multi-head attention
  • Positional Encoding: A method to inject sequence order into the model

4. Detailed Explanation and Examples (25-30 minutes)

Understanding the Transformer Encoder

The Transformer encoder is built from layers of multi-head self-attention and feed-forward networks, with residual connections and normalization.

Practical Applications

  • Machine translation
  • Text summarization
  • Question answering

Common Misconceptions

  1. Misconception: Attention is only useful for long sequences. Reality: Attention helps with both short and long sequences by focusing on relevant information.
  2. Misconception: Transformers don’t need any information about word order. Reality: Positional encoding is essential for sequence understanding.

5. Live Coding Demonstration (20-25 minutes)

Example 1: Simplified Self-Attention in NumPy

import numpy as np

def self_attention(Q, K, V):
    d_k = Q.shape[-1]
    scores = np.dot(Q, K.T) / np.sqrt(d_k)
    # Subtracting the row maximum makes softmax stable for large scores.
    scores = scores - np.max(scores, axis=-1, keepdims=True)
    exp_scores = np.exp(scores)
    weights = exp_scores / np.sum(exp_scores, axis=-1, keepdims=True)
    return np.dot(weights, V)

# Example usage:
Q = np.random.rand(3, 4)
K = np.random.rand(3, 4)
V = np.random.rand(3, 4)
output = self_attention(Q, K, V)
print(output)

Explanation:

  • Computes attention scores between all pairs of tokens
  • Applies softmax to get attention weights
  • Produces a weighted sum of values

Example 2: Multi-Head Attention Skeleton (PyTorch)

import torch
import torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        self.mha = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
    def forward(self, x):
        attn_output, _ = self.mha(x, x, x)
        return attn_output

# Example usage:
mha = MultiHeadAttention(embed_dim=32, num_heads=4)
x = torch.randn(8, 10, 32)  # (batch, seq_len, embed_dim)
output = mha(x)
print(output.shape)  # (8, 10, 32)

Key Points:

  • Multi-head attention allows the model to focus on different relationships in parallel
  • PyTorch provides a built-in module for this

Example 3: Positional Encoding Implementation

import numpy as np

def positional_encoding(seq_len, d_model):
    PE = np.zeros((seq_len, d_model))
    for pos in range(seq_len):
        for i in range(0, d_model, 2):
            # The sine and cosine in a pair use exactly the same frequency.
            angle = pos / (10000 ** (i / d_model))
            PE[pos, i] = np.sin(angle)
            if i+1 < d_model:
                PE[pos, i+1] = np.cos(angle)
    return PE

print(positional_encoding(5, 8))

Debugging Demonstration

Show common errors and how to fix them:

# Common mistake: forgetting to scale dot-product by sqrt(d_k)
def bad_attention(Q, K, V):
    scores = np.dot(Q, K.T)  # Missing scaling
    weights = np.exp(scores) / np.sum(np.exp(scores), axis=-1, keepdims=True)
    return np.dot(weights, V)

# Corrected version
def good_attention(Q, K, V):
    d_k = Q.shape[-1]
    scores = np.dot(Q, K.T) / np.sqrt(d_k)
    scores -= np.max(scores, axis=-1, keepdims=True)
    exp_scores = np.exp(scores)
    weights = exp_scores / np.sum(exp_scores, axis=-1, keepdims=True)
    return np.dot(weights, V)

6. Hands-On Practice (30-40 minutes)

Exercise 1: Basic Practice (10 minutes)

Objective: Implement a simple self-attention function in NumPy or PyTorch.

Instructions:

  1. Write a function that takes Q, K, V and returns the attention output
  2. Test it on random matrices

Starter Code:

import numpy as np
# Your code here

Solution:

def self_attention(Q, K, V):
    d_k = Q.shape[-1]
    scores = np.dot(Q, K.T) / np.sqrt(d_k)
    weights = np.exp(scores) / np.sum(np.exp(scores), axis=-1, keepdims=True)
    return np.dot(weights, V)

Exercise 2: Intermediate Challenge (15 minutes)

Objective: Build a multi-head attention module using PyTorch.

Requirements:

  • Use nn.MultiheadAttention
  • Pass a batch of sequences through it
  • Print the output shape

Hints:

  • Remember to use batch_first=True
  • Input shape should be (batch, seq_len, embed_dim)

Exercise 3: Creative Application (15 minutes)

Objective: Visualize attention weights for a sample sentence.

Challenge:

  • Use a small model to compute attention weights
  • Plot the attention matrix as a heatmap

Extension Ideas:

  • Try with real sentences
  • Compare attention patterns for different words

7. Common Pitfalls and Troubleshooting (10 minutes)

Frequent Mistakes

  1. Problem: Forgetting to scale dot-product by sqrt(d_k) Cause: Leads to unstable gradients Solution: Always divide by sqrt(d_k)

  2. Problem: Ignoring positional encoding Cause: Model loses sense of word order Solution: Always add positional encoding to input embeddings

Debugging Tips

  • Print attention weights to check for uniformity or dominance
  • Visualize positional encodings
  • Use small test cases to verify implementation

8. Best Practices and Professional Tips (5 minutes)

Industry Standards

  • Use multi-head attention for all Transformer models
  • Always include positional encoding
  • Document model architecture and hyperparameters

Next Steps in Your Learning Journey

  • Explore full Transformer encoder and decoder stacks
  • Learn about training objectives (e.g., translation)
  • Try implementing a mini-Transformer

9. Summary and Review (5 minutes)

Key Takeaways

  1. Attention allows models to focus on relevant parts of the input
  2. Multi-head attention captures diverse relationships
  3. Positional encoding is essential for sequence understanding

Quick Review Questions

  1. What problem does attention solve compared to RNNs?
  2. Why is positional encoding necessary in Transformers?
  3. How does multi-head attention improve model performance?

For Instructors

Teaching Tips

  • Use diagrams to illustrate attention and positional encoding
  • Encourage students to experiment with different sequence lengths
  • Discuss real-world applications (translation, summarization)

Assessment Ideas

  • Quiz on attention and positional encoding concepts
  • Coding assignment: implement self-attention
  • Peer review: students explain their attention visualizations

Adaptation Notes

  • For beginners: focus on single-head attention and basic positional encoding
  • For advanced students: implement multi-head attention from scratch
  • For online delivery: use Colab or Jupyter for interactive coding