# Module: Attention — TinyTorch

Skip to main content
`Ctrl`+`K`

# Module: Attention

⭐⭐⭐ | ⏱️ 4-5 hours

## 📊 Module Info

- Difficulty: ⭐⭐⭐ Advanced

- Time Estimate: 4-5 hours

- Prerequisites: Tensor module

- Next Steps: Training, Transformers modules

Build the core attention mechanism that powers modern AI! This module implements the fundamental scaled dot-product attention that’s used in ChatGPT, BERT, GPT-4, and virtually all state-of-the-art AI systems.

## 🎯 Learning Objectives

By the end of this module, you will be able to:

- Master the attention formula: Understand and implement `Attention(Q,K,V) = softmax(QK^T/√d_k)V`

- Build self-attention: Create the core component that enables global context understanding

- Control information flow: Implement masking for causal, padding, and bidirectional attention

- Visualize attention patterns: See what the model “pays attention to”

- Understand modern AI: Grasp the mechanism that revolutionized natural language processing

## 🧠 Build → Use → Understand

This module follows TinyTorch’s Build → Use → Understand framework:

1. Build: Implement the core attention mechanism and masking utilities from mathematical foundations

2. Use: Apply attention to sequence tasks and visualize attention patterns

3. Understand: How attention enables dynamic, global context modeling that powers modern AI

## 📚 What You’ll Build

### Scaled Dot-Product Attention

    def scaled_dot_product_attention(Q, K, V, mask=None):
        """
        The fundamental attention operation:
        Attention(Q,K,V) = softmax(QK^T/√d_k)V

        This exact function powers ChatGPT, BERT, and all transformers.
        """
        d_k = Q.shape[-1]
        scores = Q @ K.transpose(-2, -1) / math.sqrt(d_k)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        attention_weights = softmax(scores)
        return attention_weights @ V, attention_weights

### Self-Attention Wrapper

    class SelfAttention:
        """
        Convenient wrapper for self-attention where Q=K=V.
        The most common use case in transformer models.
        """
        def __init__(self, d_model):
            self.d_model = d_model

        def forward(self, x, mask=None):
            # Self-attention: Q = K = V = x
            return scaled_dot_product_attention(x, x, x, mask)

### Attention Masking

    # Causal masking (GPT-style: can't see future tokens)
    causal_mask = create_causal_mask(seq_len)

    # Padding masking (ignore padding tokens)
    padding_mask = create_padding_mask(lengths, max_length)

    # Bidirectional masking (BERT-style: can see all tokens)
    bidirectional_mask = create_bidirectional_mask(seq_len)

## 🔬 Key Concepts

### Why Attention Revolutionized AI

- Global connectivity: Unlike CNNs, attention connects any two positions directly

- Dynamic weights: Attention adapts to input content, not fixed like convolution kernels

- Parallel processing: Unlike RNNs, all positions computed simultaneously

- Interpretability: You can visualize what the model pays attention to

### The Attention Formula Explained

    Attention(Q,K,V) = softmax(QK^T/√d_k)V

    Where:
    - Q (Query): "What am I looking for?"
    - K (Key): "What information is available?"
    - V (Value): "What is the actual content?"
    - √d_k scaling: Prevents extreme softmax values

### Attention vs Convolution

Aspect | Convolution | Attention
--- | --- | ---
Receptive field | Local, grows with depth | Global from layer 1
Computation | O(n) with kernel size | O(n²) with sequence length
Weights | Fixed learned kernels | Dynamic input-dependent
Best for | Spatial data (images) | Sequential data (text)

### Real-World Applications

- Language Models: GPT, BERT, ChatGPT use self-attention to understand context

- Machine Translation: Google Translate uses attention to align source and target words

- Image Understanding: Vision Transformers apply attention to image patches

- Multimodal AI: CLIP, DALL-E use attention to connect text and images

## 🚀 From Attention to Modern AI

This module teaches the core building block of modern AI:

What you’re building: The fundamental attention mechanism
What it enables: Multi-head attention, positional encoding, transformer blocks
What it powers: ChatGPT, BERT, GPT-4, and contemporary AI systems

Understanding this module gives you the foundation to understand:

- How ChatGPT generates coherent text

- How BERT understands language bidirectionally

- How Vision Transformers work without convolution

- How modern AI achieves human-like language understanding

## 📈 Module Progression

    Tensors → **ATTENTION** → Layers → Networks → CNNs → Training
      ↑              ↑
    Foundation   Modern AI Core

After completing this module, you’ll understand the mechanism that sparked the AI revolution, making you ready to work with state-of-the-art models and architectures.

## 🎯 Success Criteria

You’ll know you’ve mastered this module when you can:

- Implement scaled dot-product attention from scratch

- Explain why the √d_k scaling prevents gradient problems

- Create different types of attention masks for various use cases

- Visualize and interpret attention weights

- Understand why attention enabled the transformer revolution

- Connect this foundation to modern AI systems like ChatGPT

Choose your preferred way to engage with this module:

🚀 Launch Binder

Run this module interactively in your browser. No installation required!

https://mybinder.org/v2/gh/mlsysbook/TinyTorch/main?filepath=modules/source/07_attention/attention_dev.ipynb
⚡ Open in Colab

Use Google Colab for GPU access and cloud compute power.

https://colab.research.google.com/github/mlsysbook/TinyTorch/blob/main/modules/source/07_attention/attention_dev.ipynb
📖 View Source

Browse the Python source code and understand the implementation.

https://github.com/mlsysbook/TinyTorch/blob/main/modules/source/07_attention/attention_dev.py

💾 Save Your Progress

Binder sessions are temporary! Download your completed notebook when done, or switch to local development for persistent work.

Ready for serious development? → 🏗️ Local Setup Guide

← Previous Module Next Module →