Welcome back! It's amazing seeing you again in our course "Sequence Models & The Dawn of Attention". As we advance to our second lesson, we're about to explore one of the most revolutionary concepts in modern deep learning: the attention mechanism.

In our previous lesson, we witnessed firsthand how LSTMs struggle with long-range dependencies, observing performance degradation as sequence lengths increased. This limitation wasn't merely a technical curiosity; it revealed a fundamental bottleneck that restricted the potential of sequence models. The core issue lies in forcing models to compress all relevant information into a single fixed-size hidden state, creating an information bottleneck that becomes more severe as sequences grow longer.

Today, we'll discover how attention mechanisms elegantly solve this problem by allowing models to selectively focus on different parts of the input sequence. Rather than relying on a single summary vector, attention enables direct access to any position in the sequence, fundamentally changing how we approach sequence modeling. We'll implement two foundational attention variants: Luong (multiplicative) and Bahdanau (additive) attention, understanding their mathematical foundations and practical differences. Let's dive in!

The limitation we observed with LSTMs stems from a deeper architectural constraint: the fixed-size context bottleneck. Traditional sequence models, like RNNs and LSTMs, process information sequentially and attempt to compress all information encountered so far into a hidden state of a predetermined size. This hidden state must then serve as the sole basis for future predictions or for understanding the sequence as a whole.

Imagine trying to summarize an entire book in a single, short sentence while ensuring all crucial plot points, character developments, and themes are perfectly preserved. It's an incredibly difficult, if not impossible, task. As the book (or sequence) gets longer, more information needs to be crammed into that one sentence (the fixed-size hidden state), inevitably leading to information loss.

Consider a machine translation task. If we want to translate a long, complex sentence like "The cat, which had been lazily napping under the old oak tree in the sprawling garden all afternoon, suddenly awoke with a start," an LSTM would process this word by word. By the time it needs to translate "awoke," its hidden state must somehow retain the fact that "cat" is the subject, along with all the intervening descriptive clauses. This reliance on a compressed summary becomes increasingly problematic with longer sequences. Attention mechanisms offer a way out by allowing the model to look back at the entire input sequence at each step, rather than relying solely on a compressed summary.

At its core, attention operates through three fundamental components: Queries (Q), Keys (K), and Values (V). This QKV formulation provides an elegant and powerful framework for computing relevance and retrieving information. The concept is inspired by information retrieval systems, where you use a query to search a database (composed of key-value pairs) to find relevant information.

Let's break down these components in the context of sequence models:

• A Query represents the current point of interest or what the model is trying to figure out at a specific step. For example, in machine translation, a query might be related to the word currently being generated in the target language. It essentially asks, "Given my current context, what information from the input sequence is most relevant right now?"
• Keys are associated with different parts of the input sequence. Each key corresponds to a specific element or position in the input. They act like "labels" or "indices" for the information contained in the input.
• Values also correspond to the elements of the input sequence. They contain the actual information or content that we want to retrieve. Typically, for each key, there's an associated value.

The attention mechanism works by:

1. Comparing the Query with all the Keys to calculate a set of scores. These scores determine how relevant each input part (represented by its Key) is to the current Query.
2. Using these scores to compute a weighted sum of the Values.

This process allows the model to selectively focus on the most pertinent parts of the input sequence (those whose Keys best match the Query) and retrieve their corresponding Values.

Let's set up some sample data to see what these tensors look like. We'll use PyTorch for our implementation:

import torch

def create_sample_data(batch_size=2, seq_len=6, hidden_size=4):
    """Create sample query, keys, and values for demonstration"""
    torch.manual_seed(42) # For reproducibility
    keys = torch.randn(batch_size, seq_len, hidden_size)
    values = torch.randn(batch_size, seq_len, hidden_size)
    query = torch.randn(batch_size, hidden_size)
    
    return query, keys, values

In this function, create_sample_data, we generate sample random tensors for our query, keys, and values:

• query has a shape of (batch_size, hidden_size). This represents a single query vector for each item in our batch.
• keys has a shape of (batch_size, seq_len, hidden_size). For each item in the batch, there's a sequence of seq_len key vectors.
• values has a shape of (batch_size, seq_len, hidden_size), mirroring the shape of keys. Each key has a corresponding value.

These tensor shapes are typical for attention mechanisms. The hidden_size represents the dimensionality of our embeddings or feature vectors.

The heart of any attention mechanism lies in computing attention scores. These scores quantify the relevance between a Query (representing what we're looking for) and each Key in the input sequence (representing different pieces of information). Think of it like searching for a specific topic (the query) in a library catalog (the keys); some entries will be highly relevant, others not at all.

Once we have these raw scores, they need to be transformed into something more usable. This is where the softmax function comes in. Applying softmax to the scores converts them into a set of attention weights. These weights have two important properties:

1. Each weight is between 0 and 1.
2. All weights for a given query (across all keys in the sequence) sum up to 1.

This means the attention weights form a probability distribution, indicating how much "attention" or importance the model should pay to each part of the input sequence when constructing its output. A higher weight for a particular key-value pair means that value will contribute more to the final result, known as the context vector.

Different attention mechanisms use different scoring functions to calculate the initial raw scores. A common and intuitive one is the dot product, which measures the similarity in orientation between the query and key vectors. We'll see this in action with Luong attention.

Luong attention, also known as multiplicative attention, is a popular and efficient attention mechanism. It gets its "multiplicative" name because it primarily uses dot products (a form of multiplication) to calculate attention scores. Let's see how to implement it.

import torch.nn.functional as F

def luong_attention(query, keys, values):
    """
    Implement Luong (multiplicative) attention mechanism
    
    Args:
        query: (batch_size, hidden_size)
        keys: (batch_size, seq_len, hidden_size)
        values: (batch_size, seq_len, hidden_size)
    
    Returns:
        context: (batch_size, hidden_size)
        attention_weights: (batch_size, seq_len)
    """
    # Ensure query is ready for batch matrix multiplication
    query = query.unsqueeze(1)  # Shape: (batch_size, 1, hidden_size)
    
    # Compute attention scores using dot product
    # keys.transpose(1, 2) changes keys to (batch_size, hidden_size, seq_len)
    # scores will be (batch_size, 1, seq_len)
    scores = torch.bmm(query, keys.transpose(1, 2))
    scores = scores.squeeze(1)  # Shape: (batch_size, seq_len)
    
    # Apply softmax to get attention weights
    attention_weights = F.softmax(scores, dim=1) # Shape: (batch_size, seq_len)
    
    # Compute context vector as weighted sum of values
    # attention_weights.unsqueeze(1) gives (batch_size, 1, seq_len)
    # values are (batch_size, seq_len, hidden_size)
    # context will be (batch_size, 1, hidden_size)
    context = torch.bmm(attention_weights.unsqueeze(1), values)
    context = context.squeeze(1)  # Shape: (batch_size, hidden_size)
    
    return context, attention_weights

Let's break this down:

1. Prepare Query: We first unsqueeze the query tensor at dimension 1. This changes its shape from (batch_size, hidden_size) to (batch_size, 1, hidden_size), making it compatible for batch matrix multiplication (torch.bmm) with the keys.
2. Calculate Scores: The core of Luong attention is the dot product between the query and each key. We transpose the keys tensor's last two dimensions (from (batch_size, seq_len, hidden_size) to (batch_size, hidden_size, seq_len)). Then, torch.bmm(query_expanded, keys.transpose(1, 2)) computes the dot product of the query with every key in the sequence for each item in the batch. The result scores (after squeezing out the dimension of size 1) has a shape of (batch_size, seq_len), indicating how much the query aligns with each key.
3. Normalize with Softmax: The raw scores are passed through F.softmax(scores, dim=1). This operation is applied along dim=1 (the sequence length dimension), converting the scores into attention_weights. These weights now sum to 1 for each item in the batch, representing a probability distribution.
4. Compute Context Vector: Finally, the context vector is computed as a weighted sum of the values. We unsqueeze the attention_weights to make them (batch_size, 1, seq_len) and then perform another batch matrix multiplication with values (shape (batch_size, seq_len, hidden_size)). This effectively multiplies each value vector by its corresponding attention weight and sums them up. The resulting context vector (after squeezing) has shape (batch_size, hidden_size) and summarizes the relevant information from the input sequence based on the query.

Notice how this mechanism enables parallel processing: all dot products and weighted sums are computed simultaneously across the sequence and batch, a key advantage over sequential RNN processing.

Bahdanau attention, also called additive attention, takes a slightly more complex approach to computing query-key similarity. Instead of a simple dot product, it uses a small feed-forward neural network (often a single linear layer after combining query and key) to calculate the scores. This allows for potentially more expressive relationships to be learned.

import torch.nn as nn

def bahdanau_attention(query, keys, values, hidden_size):
    """
    Implement Bahdanau (additive) attention mechanism
    
    Args:
        query: (batch_size, hidden_size)
        keys: (batch_size, seq_len, hidden_size)
        values: (batch_size, seq_len, hidden_size)
        hidden_size: dimension of hidden state (and query/key/value vectors)
    
    Returns:
        context: (batch_size, hidden_size)
        attention_weights: (batch_size, seq_len)
    """
    # Define linear layers for Bahdanau attention (these would be part of a model's __init__)
    W_q = nn.Linear(hidden_size, hidden_size, bias=False)
    W_k = nn.Linear(hidden_size, hidden_size, bias=False)
    v = nn.Linear(hidden_size, 1, bias=False)
    
    # Transform query and keys
    query_transformed = W_q(query)  # Shape: (batch_size, hidden_size)
    keys_transformed = W_k(keys)    # Shape: (batch_size, seq_len, hidden_size)
    
    # Expand query to match keys dimensions for addition
    # query_transformed.unsqueeze(1) -> (batch_size, 1, hidden_size)
    # .expand(-1, keys.shape[1], -1) -> (batch_size, seq_len, hidden_size)
    query_expanded = query_transformed.unsqueeze(1).expand_as(keys_transformed)
    
    # Add transformed query and keys, apply tanh activation
    # This is the "additive" part
    combined_energy = torch.tanh(query_expanded + keys_transformed)  # Shape: (batch_size, seq_len, hidden_size)
    
    # Compute scores using the 'v' linear layer
    scores = v(combined_energy).squeeze(-1)  # Shape: (batch_size, seq_len)
    
    # Apply softmax to get attention weights
    attention_weights = F.softmax(scores, dim=1) # Shape: (batch_size, seq_len)
    
    # Compute context vector as weighted sum of values
    context = torch.bmm(attention_weights.unsqueeze(1), values) # Shape: (batch_size, 1, hidden_size)
    context = context.squeeze(1)  # Shape: (batch_size, hidden_size)
    
    return context, attention_weights

In Bahdanau attention:

1. Learnable Transformations: We define three linear layers: W_q to transform the query, W_k to transform the keys, and v to compute the final score from their combined representation. In a real model, these layers would be initialized once and their weights learned during training.
2. Transform Query and Keys: The query and keys are passed through their respective linear layers (W_q and W_k).
3. Combine and Activate: The transformed query is expanded to match the dimensions of the transformed keys. Then, they are added together (query_expanded + keys_transformed). The result is passed through a tanh activation function. This non-linear combination is a key feature of additive attention.
4. Score Calculation: The output of the tanh function is passed through the v linear layer, which projects it down to a single score per sequence position.
5. Softmax and Context Vector: The rest of the process is similar to Luong attention: the scores are normalized using softmax to get attention_weights, and these weights are used to compute a weighted sum of the values to get the context vector.

The main difference lies in how the scores are computed: Luong uses a direct dot product, while Bahdanau uses a small neural network with learned weights, allowing for more complex scoring functions.

Now let's examine how these two attention mechanisms behave in practice by running them on our sample data and analyzing their outputs:

import numpy as np

def main():   
    # Create sample data
    batch_size = 2
    seq_len = 6
    hidden_size = 4
    query, keys, values = create_sample_data(batch_size, seq_len, hidden_size)
    
    print(f"Query shape: {query.shape}")
    print(f"Keys shape: {keys.shape}")
    print(f"Values shape: {values.shape}")
    
    # Apply Luong attention
    context_luong, attention_weights_luong = luong_attention(query, keys, values)
    
    print(f"Context vector shape: {context_luong.shape}") # Luong context
    print(f"Attention weights shape: {attention_weights_luong.shape}") # Luong weights
    
    # Display Luong attention weights
    print("\nLuong Attention weights:")
    for i in range(attention_weights_luong.shape[0]): # Iterate through batch
        weights = attention_weights_luong[i].detach().numpy()
        print(f"Batch {i}: {np.round(weights, decimals=8)}") # Using numpy for cleaner print   
    
    # Apply Bahdanau attention
    context_bahdanau, attention_weights_bahdanau = bahdanau_attention(query, keys, values, hidden_size=hidden_size)
    
    # Display Bahdanau attention weights
    print("\nBahdanau Attention weights:")
    for i in range(attention_weights_bahdanau.shape[0]): # Iterate through batch
        weights = attention_weights_bahdanau[i].detach().numpy()
        print(f"Batch {i}: {np.round(weights, decimals=8)}")

When we run the main() function, we'll see the shapes of our tensors and the attention weights produced by each mechanism. The expected output is:

Query shape: torch.Size([2, 4])
Keys shape: torch.Size([2, 6, 4])
Values shape: torch.Size([2, 6, 4])
Context vector shape: torch.Size([2, 4])
Attention weights shape: torch.Size([2, 6])

Luong Attention weights:
Batch 0: [0.02072945 0.7967359  0.07351931 0.09071727 0.01487425 0.00342377]
Batch 1: [0.01174377 0.01791002 0.00222191 0.00307018 0.07629474 0.8887594 ]

Bahdanau Attention weights:
Batch 0: [0.14973381 0.21063526 0.1710429  0.18545005 0.14640713 0.1367309 ]
Batch 1: [0.11768297 0.15604942 0.18164434 0.24428545 0.1681909  0.13214697]

Looking at the output:

• The tensor shapes confirm our understanding. The context vector has the same hidden_size as the query, and attention weights are distributed over the seq_len.
• Luong Attention tends to produce more "peaked" or concentrated attention weights. For Batch 0, the second element (index 1) gets almost 80% of the attention. For Batch 1, the last element (index 5) gets about 89%. This is characteristic of dot-product attention, which can sharply focus if a key is very similar to the query.
• Bahdanau Attention, with its learnable layers and tanh activation, often produces more "diffuse" or spread-out attention weights, especially before training. In our example (with randomly initialized weights for W_q, W_k, v), the weights are more evenly distributed. For Batch 0, the highest weight is around 21%, and for Batch 1, it's around 24%. This doesn't mean it can't focus; during training, it can learn to produce sharper distributions if needed.

The choice between Luong and Bahdanau (or other attention variants) can depend on the specific task, dataset, and computational budget. Luong is generally faster due to its simpler computation, while Bahdanau offers more flexibility through its learnable components.

Today, we've taken a crucial step in our journey from traditional sequence models to the Transformer architecture. We explored how attention mechanisms solve the fixed-context bottleneck that limited RNNs and LSTMs, introducing the elegant Query-Key-Value paradigm that enables selective information retrieval from any sequence position.

Through implementing both Luong and Bahdanau attention, we discovered how different similarity functions produce distinct attention behaviors: multiplicative attention can create focused patterns, while additive attention with learnable transformations allows for more nuanced distributions. These mechanisms are foundational building blocks for more advanced attention systems. In our next lesson, we'll build upon these concepts to explore multi-head attention, discovering how parallel attention mechanisms can capture different types of relationships simultaneously, bringing us closer to the full Transformer. Until then, let's practice what we've learned today!