Introducing the Attention Mechanism

Introduction

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 Problem with Fixed Context Windows

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.

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

Implementing 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.

Python

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

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:

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.
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.
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.
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.

Implementing Bahdanau Attention

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.

Python

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

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:

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.
Transform Query and Keys: The query and keys are passed through their respective linear layers (W_q and W_k).
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.
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.
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.

Comparing Attention Mechanisms

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

Python

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)}")

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:

text

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]

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.

Introduction

The Problem with Fixed Context Windows

Join the 1M+ learners on CodeSignal

The Query-Key-Value Paradigm

Setting Up Sample Data for Attention

Calculating Attention Scores

Implementing Luong Attention

Implementing Bahdanau Attention

Comparing Attention Mechanisms

Conclusion and Next Steps