Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Bahdanau Attention › From Theory to Code: Implementing Bahdanau Attention

Michael BrenndoerferUpdated May 22, 20255 min read

Language AI HandbookAll chapters

Part 12 of 67Sequence-to-Sequence

Chapter 87Bahdanau Attention › From Theory to Code: Implementing Bahdanau Attention

Reading section 5 of 11 · From Theory to Code: Implementing Bahdanau Attention · 5 minSection 5 of 11 · 5 min

Section 5 of 11 · 5 min read

Bahdanau Attention › From Theory to Code: Implementing Bahdanau AttentionOne section · Sections in this chapter5 / 11 ↓

Sections in this chapter11

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Having traced through the mathematics by hand, let's now implement Bahdanau attention as a PyTorch module. The code will mirror the formulas exactly, making it easy to see how each mathematical operation translates to tensor operations.

In[22]:

Code

class BahdanauAttention(nn.Module):
    """
    Bahdanau (additive) attention mechanism.

    Args:
        encoder_dim: Dimension of encoder hidden states
        decoder_dim: Dimension of decoder hidden states
        attention_dim: Dimension of the alignment space
    """

    def __init__(self, encoder_dim, decoder_dim, attention_dim):
        super().__init__()

        # Projection matrices
        self.W_a = nn.Linear(decoder_dim, attention_dim, bias=False)
        self.U_a = nn.Linear(encoder_dim, attention_dim, bias=False)
        self.v_a = nn.Linear(attention_dim, 1, bias=False)

    def forward(self, decoder_state, encoder_outputs):
        """
        Compute attention weights and context vector.

        Args:
            decoder_state: Previous decoder hidden state (batch, decoder_dim)
            encoder_outputs: All encoder hidden states (batch, seq_len, encoder_dim)

        Returns:
            context: Context vector (batch, encoder_dim)
            attention_weights: Attention distribution (batch, seq_len)
        """
        # Project decoder state: (batch, attention_dim)
        decoder_proj = self.W_a(decoder_state)

        # Project encoder outputs: (batch, seq_len, attention_dim)
        encoder_proj = self.U_a(encoder_outputs)

        # Add projections (broadcast decoder over sequence)
        # decoder_proj: (batch, 1, attention_dim) after unsqueeze
        # encoder_proj: (batch, seq_len, attention_dim)
        combined = torch.tanh(decoder_proj.unsqueeze(1) + encoder_proj)

        # Compute scores: (batch, seq_len, 1) -> (batch, seq_len)
        scores = self.v_a(combined).squeeze(-1)

        # Normalize to attention weights
        attention_weights = F.softmax(scores, dim=-1)

        # Compute context vector as weighted sum
        # attention_weights: (batch, seq_len) -> (batch, 1, seq_len)
        # encoder_outputs: (batch, seq_len, encoder_dim)
        # Result: (batch, encoder_dim)
        context = torch.bmm(
            attention_weights.unsqueeze(1), encoder_outputs
        ).squeeze(1)

        return context, attention_weights

The implementation follows our mathematical formulation exactly. Let's trace through the key operations:

Projection layers: W_a, U_a, and v_a are implemented as nn.Linear layers without bias terms, matching our weight matrices
Broadcasting: When we add decoder_proj.unsqueeze(1) to encoder_proj, PyTorch broadcasts the decoder projection across all sequence positions
Batch processing: The code handles batches of sequences simultaneously, which is essential for efficient training
Efficient aggregation: torch.bmm (batch matrix multiplication) computes the weighted sum $\sum_j \alpha_{ij} h_j$ efficiently as a matrix operation

Let's verify that our implementation produces the expected outputs:

In[23]:

Code

# Test the attention module
batch_size = 2
seq_len = 5
encoder_dim = 64
decoder_dim = 64
attention_dim = 32

attention = BahdanauAttention(encoder_dim, decoder_dim, attention_dim)

# Create dummy inputs
encoder_outputs = torch.randn(batch_size, seq_len, encoder_dim)
decoder_state = torch.randn(batch_size, decoder_dim)

# Compute attention
context, weights = attention(decoder_state, encoder_outputs)

Out[24]:

Console

Encoder outputs shape: torch.Size([2, 5, 64])
Decoder state shape: torch.Size([2, 64])
Context vector shape: torch.Size([2, 64])
Attention weights shape: torch.Size([2, 5])

Attention weights sum to 1: tensor([1.0000, 1.0000], grad_fn=<SumBackward1>)

The attention weights sum to 1 for each batch element, confirming that softmax normalization is working correctly. The context vector has the same dimension as the encoder hidden states, ready to be fed into the decoder RNN.

Building the Complete DecoderLink Copied

The attention module computes context vectors, but we still need to integrate it into a full decoder that can generate output sequences. The decoder must coordinate several components: embedding the previous output token, computing attention, updating its hidden state, and predicting the next token.

In[25]:

Code

class AttentionDecoder(nn.Module):
    """
    Decoder with Bahdanau attention for sequence-to-sequence models.

    Args:
        vocab_size: Size of output vocabulary
        embed_dim: Dimension of word embeddings
        encoder_dim: Dimension of encoder hidden states
        decoder_dim: Dimension of decoder hidden states
        attention_dim: Dimension of attention alignment space
    """

    def __init__(
        self, vocab_size, embed_dim, encoder_dim, decoder_dim, attention_dim
    ):
        super().__init__()

        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.attention = BahdanauAttention(
            encoder_dim, decoder_dim, attention_dim
        )

        # GRU input: embedded token + context vector
        self.gru = nn.GRU(
            embed_dim + encoder_dim, decoder_dim, batch_first=True
        )

        # Output projection
        self.output_proj = nn.Linear(decoder_dim, vocab_size)

        self.decoder_dim = decoder_dim

    def forward_step(self, prev_token, prev_hidden, encoder_outputs):
        """
        Single decoding step with attention.

        Args:
            prev_token: Previous output token indices (batch,)
            prev_hidden: Previous decoder hidden state (1, batch, decoder_dim)
            encoder_outputs: Encoder hidden states (batch, seq_len, encoder_dim)

        Returns:
            output: Log probabilities over vocabulary (batch, vocab_size)
            hidden: New decoder hidden state (1, batch, decoder_dim)
            attention_weights: Attention distribution (batch, seq_len)
        """
        # Embed previous token: (batch, embed_dim)
        embedded = self.embedding(prev_token)

        # Compute attention using previous hidden state
        # prev_hidden shape: (1, batch, decoder_dim) -> (batch, decoder_dim)
        context, attention_weights = self.attention(
            prev_hidden.squeeze(0), encoder_outputs
        )

        # Concatenate embedding and context: (batch, 1, embed_dim + encoder_dim)
        gru_input = torch.cat([embedded, context], dim=-1).unsqueeze(1)

        # Update decoder state
        output, hidden = self.gru(gru_input, prev_hidden)

        # Project to vocabulary: (batch, vocab_size)
        logits = self.output_proj(output.squeeze(1))

        return F.log_softmax(logits, dim=-1), hidden, attention_weights

This decoder implements the complete Bahdanau architecture. The forward_step method executes one decoding step, following the exact sequence we outlined earlier:

Embed: Convert the previous token ID to a dense vector
Attend: Use the previous hidden state to compute attention over encoder outputs
Concatenate: Combine the embedding and context vector
Update: Pass through the GRU to get the new hidden state
Project: Transform the hidden state to vocabulary logits

Let's verify that the decoder produces outputs of the expected shapes:

In[26]:

Code

# Test the decoder
vocab_size = 1000
embed_dim = 128
encoder_dim = 256
decoder_dim = 256
attention_dim = 128

decoder = AttentionDecoder(
    vocab_size, embed_dim, encoder_dim, decoder_dim, attention_dim
)

# Simulate encoder outputs (from a bidirectional LSTM, for example)
batch_size = 4
src_len = 10
encoder_outputs = torch.randn(batch_size, src_len, encoder_dim)

# Initialize decoder
prev_token = torch.zeros(batch_size, dtype=torch.long)  # Start token
prev_hidden = torch.zeros(1, batch_size, decoder_dim)

# One decoding step
log_probs, hidden, attn_weights = decoder.forward_step(
    prev_token, prev_hidden, encoder_outputs
)

Out[27]:

Console

Output log probabilities shape: torch.Size([4, 1000])
New hidden state shape: torch.Size([1, 4, 256])
Attention weights shape: torch.Size([4, 10])

Sample attention distribution (first batch element):
  [0.098 0.176 0.09  0.087 0.079 0.092 0.124 0.096 0.078 0.079]