Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Bahdanau Attention › Computational Considerations

Michael BrenndoerferUpdated May 22, 20253 min read

Language AI HandbookAll chapters

Part 12 of 67Sequence-to-Sequence

Chapter 87Bahdanau Attention › Computational Considerations

Reading section 7 of 11 · Computational Considerations · 3 minSection 7 of 11 · 3 min

Section 7 of 11 · 3 min read

Bahdanau Attention › Computational ConsiderationsOne section · Sections in this chapter7 / 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

Bahdanau attention introduces additional computation compared to a standard encoder-decoder. Let's analyze the complexity by examining each operation required at every decoder step.

For each decoder step, we must perform the following operations:

Project the decoder state: Computing $W_a s_{i-1}$ requires $O(d_a \times d_s)$ operations (matrix-vector multiplication)
Project all encoder states: Computing $U_a h_j$ for all $j$ requires $O(T_x \times d_a \times d_h)$ operations
Compute scores: The addition, tanh, and dot product with $v_a$ require $O(T_x \times d_a)$ operations
Apply softmax: Normalizing scores requires $O(T_x)$ operations
Compute weighted sum: Computing $\sum_j \alpha_{ij} h_j$ requires $O(T_x \times d_h)$ operations

where:

$T_x$ : source sequence length
$T_y$ : target sequence length
$d_s$ : decoder hidden dimension
$d_h$ : encoder hidden dimension
$d_a$ : alignment dimension

The dominant term is the encoder projection, which is $O(T_x \times d_a \times d_h)$ per decoder step. Over an entire output sequence of length $T_y$ , the total complexity becomes $O(T_y \times T_x \times d_a \times d_h)$ . This quadratic dependence on sequence lengths ( $T_y \times T_x$ ) becomes significant for long sequences.

Out[32]:

Visualization

Line chart showing attention computation cost in MFLOPS growing quadratically with sequence length, for three hidden dimension sizes. — Computational cost of Bahdanau attention grows quadratically with sequence length. For very long sequences, this becomes a bottleneck, motivating more efficient attention variants.

For typical machine translation tasks with sequences of 20-50 tokens, this overhead is manageable. However, for very long sequences (documents, conversations), the quadratic scaling becomes problematic. This motivated later work on efficient attention variants, which we'll explore in subsequent chapters.

A practical optimization addresses the redundant computation in step 2. Notice that the encoder projections $U_a h_j$ depend only on the encoder outputs, which remain constant throughout decoding. By precomputing these projections once after encoding and caching them, we avoid repeating this $O(T_x \times d_a \times d_h)$ computation at every decoder step:

In[33]:

Code

class EfficientBahdanauAttention(nn.Module):
    """Bahdanau attention with precomputed encoder projections."""

    def __init__(self, encoder_dim, decoder_dim, attention_dim):
        super().__init__()
        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)

        # Cache for precomputed encoder projections
        self.encoder_proj_cache = None

    def precompute_encoder_proj(self, encoder_outputs):
        """Precompute encoder projections once per sequence."""
        self.encoder_proj_cache = self.U_a(encoder_outputs)

    def forward(self, decoder_state, encoder_outputs):
        """Compute attention using cached encoder projections."""
        decoder_proj = self.W_a(decoder_state)

        # Use cached projection if available
        if self.encoder_proj_cache is not None:
            encoder_proj = self.encoder_proj_cache
        else:
            encoder_proj = self.U_a(encoder_outputs)

        combined = torch.tanh(decoder_proj.unsqueeze(1) + encoder_proj)
        scores = self.v_a(combined).squeeze(-1)
        attention_weights = F.softmax(scores, dim=-1)
        context = torch.bmm(
            attention_weights.unsqueeze(1), encoder_outputs
        ).squeeze(1)

        return context, attention_weights