Bahdanau Attention: Dynamic Context for Neural Machine Translation

In the previous chapter, we developed an intuition for attention as a mechanism that allows models to selectively focus on different parts of the input when generating each output. We saw how attention acts as a soft lookup, assigning weights to encoder hidden states based on their relevance to the current decoding step. Now it's time to formalize this intuition into a concrete algorithm.

Bahdanau attention, introduced in the 2014 paper "Neural Machine Translation by Jointly Learning to Align and Translate," was the first attention mechanism to achieve broad success in sequence-to-sequence models. The key innovation was replacing the fixed context vector bottleneck with a dynamic, position-dependent context that changes at each decoder timestep. This change led to significant improvements in translation quality, especially for longer sentences.

Before diving into the mathematics, let's understand the problem that Bahdanau attention solves. In a standard encoder-decoder architecture, the encoder processes the entire input sequence and compresses it into a single fixed-length vector. The decoder then generates the output sequence using only this compressed representation.

This bottleneck creates a fundamental problem: the context vector must somehow encode everything about the input that the decoder might need. For short sentences, this works reasonably well. But as input length increases, the fixed-size context vector becomes increasingly inadequate. Information gets lost, compressed, or muddled together.

The insight behind attention is that different parts of the output depend on different parts of the input. When translating "The cat sat on the mat" to French, generating "Le" depends mostly on "The," while generating "chat" depends mostly on "cat." Instead of forcing all this information through a single bottleneck, why not let the decoder directly access the relevant encoder states at each step?

With the alignment problem clearly defined, we can now develop the mathematical machinery that Bahdanau attention uses to solve it. Our journey will take us through three interconnected concepts: computing relevance scores, converting those scores to probabilities, and aggregating information based on those probabilities. Each piece builds naturally on the previous one, culminating in a mechanism that allows the decoder to "look back" at exactly the right parts of the input at each generation step.

Attention addresses a simple question: when the decoder is about to generate output token $i$, how relevant is each encoder position $j$ to that decision? If we could answer this question with a single number for each encoder position, we'd have a way to rank and weight the encoder states.

But what should this "relevance score" capture? Consider translating "The cat sat on the mat" to French. When generating "chat" (cat), the decoder needs to know:

• What the decoder is currently looking for (encoded in its hidden state $s_{i-1}$)
• What information each encoder position offers (encoded in hidden states $h_1, h_2, \ldots, h_{T_x}$)

The relevance of position $j$ depends on both of these: it's not just about what $h_j$ contains, but whether that content matches what the decoder currently needs. This suggests we need a function that takes both the decoder state and an encoder state as inputs and produces a scalar score.

Bahdanau and colleagues designed a solution called the additive (or concatenative) alignment model. Rather than directly comparing the decoder and encoder states, which might have different dimensions and represent different types of information, they project both into a shared "alignment space" where comparison becomes meaningful.

The alignment score between decoder position $i$ and encoder position $j$ is computed as:

where:

• $e_{ij}$: the alignment score, a scalar indicating how relevant encoder position $j$ is for generating output $i$
• $s_{i-1} \in \mathbb{R}^{d_s}$: the decoder hidden state at the previous timestep, encoding what the decoder "knows" and "needs"
• $h_j \in \mathbb{R}^{d_h}$: the encoder hidden state at position $j$, encoding information about input token $j$ and its context
• $W_a \in \mathbb{R}^{d_a \times d_s}$: a learnable weight matrix that projects the decoder state into the alignment space
• $U_a \in \mathbb{R}^{d_a \times d_h}$: a learnable weight matrix that projects encoder states into the same alignment space
• $v_a \in \mathbb{R}^{d_a}$: a learnable vector that reduces the combined representation to a scalar
• $d_a$: the dimension of the alignment space (a hyperparameter, typically matching the hidden dimensions)

This formula might look complex at first glance, but it follows a clear logic. Let's trace through it step by step to understand why each component is necessary:

Step 1: Project into alignment space. The decoder state $s_{i-1}$ and encoder state $h_j$ live in potentially different vector spaces with different dimensions. Before we can meaningfully compare them, we need to transform them into a common representation. The matrices $W_a$ and $U_a$ learn these transformations during training:

• $W_a s_{i-1}$ produces a $d_a$-dimensional vector representing "what the decoder is looking for"
• $U_a h_j$ produces a $d_a$-dimensional vector representing "what encoder position $j$ offers"

Step 2: Combine additively. With both vectors in the same space, we add them element-wise: $W_a s_{i-1} + U_a h_j$. This additive combination allows the model to detect when the decoder's query and the encoder's content are compatible. If certain dimensions align well (both positive or both negative), they reinforce each other; if they conflict, they cancel out.

Step 3: Apply nonlinearity. The $\tanh$ function serves two purposes. First, it bounds the values to $[-1, 1]$, preventing any single dimension from dominating. Second, it introduces nonlinearity, enabling the model to learn complex, non-linear relationships between decoder queries and encoder contents. Without this nonlinearity, the entire alignment model would collapse to a simple linear function.

Step 4: Reduce to scalar. Finally, the dot product with $v_a$ compresses the $d_a$-dimensional result into a single scalar score. The vector $v_a$ learns which dimensions of the alignment space are most important for determining relevance. Some dimensions might capture semantic similarity, others syntactic compatibility; $v_a$ learns to weight these appropriately.

Why "additive" rather than some other approach? Additive combination with separate projection matrices gives the model maximum flexibility. The decoder and encoder might have different hidden dimensions, and even if they're the same size, they encode different types of information. By learning separate projections $W_a$ and $U_a$, the model can transform each representation into whatever form makes comparison most effective. Bahdanau and colleagues found this approach worked well in practice, and it became the foundation for subsequent attention mechanisms.

We now have a way to compute alignment scores $e_{ij}$ for each encoder position. But these raw scores can be any real number: positive, negative, or zero. To use them for weighting encoder states, we need to transform them into something more interpretable: a probability distribution over encoder positions.

This is where the softmax function enters the picture. Softmax takes a vector of arbitrary real numbers and converts it into a valid probability distribution where all values are positive and sum to 1:

where:

• $\alpha_{ij}$: the attention weight for encoder position $j$ when generating output $i$
• $e_{ij}$: the raw alignment score from our scoring function
• $\exp(\cdot)$: the exponential function ($e^x$)
• $T_x$: the length of the source sequence
• $\sum_{k=1}^{T_x} \exp(e_{ik})$: the normalizing constant that ensures weights sum to 1

The exponential function is the key to softmax's behavior. It has two essential properties that make it ideal for attention:

1. Positivity: $\exp(x) > 0$ for all real $x$, guaranteeing non-negative weights regardless of the input scores
2. Monotonicity with amplification: larger scores produce exponentially larger values, which amplifies differences between scores

This second property is particularly important. Consider two scores: $e_{i1} = 1.2$ and $e_{i2} = 3.5$. The difference looks modest, but after exponentiation, $\exp(3.5) \approx 33$ while $\exp(1.2) \approx 3.3$, a 10x ratio. Softmax creates a "winner-take-most" dynamic where the highest-scoring positions dominate, while still allowing lower-scoring positions to contribute.

The attention weight $\alpha_{ij}$ has a clear interpretation: it represents the probability that encoder position $j$ contains the information most relevant for generating output token $i$. When we compute attention, we're asking: "Given what the decoder currently needs, where in the input should it look?"

The visualization above shows this amplification in action. Position 4 has a raw score of 2.8, only about 2× higher than position 1's score of 0.8. But after softmax, position 4 receives about 4× the attention weight of position 1. This amplification helps the model make decisive choices about where to focus.

Beyond the mathematical properties, softmax has two important advantages for learning:

1. Differentiability: Gradients flow smoothly through softmax, enabling end-to-end training with backpropagation.
2. Soft selection: Unlike "hard" attention (which would pick exactly one position), soft attention maintains gradients to all positions, making optimization much easier.

The "softness" of attention can be controlled by scaling the scores before applying softmax. Dividing scores by a temperature parameter $\tau$ changes how sharply the attention focuses:

The entropy values in the figure quantify the "spread" of attention: lower entropy means more focused attention, higher entropy means more distributed. Standard Bahdanau attention uses $\tau = 1$, which provides a good balance between focusing on relevant positions and maintaining gradient flow to all positions during training.

The "sharpness" of the attention distribution depends on the magnitude of score differences. When scores are similar, attention spreads across many positions; when one score is much higher, attention concentrates sharply:

The entropy values quantify this difference: higher entropy indicates more spread-out attention, while lower entropy indicates more concentrated attention. A well-trained model learns to produce focused attention when one input position is clearly relevant (e.g., translating a content word) and spread attention when multiple positions contribute (e.g., translating a phrase that spans multiple source words).

We've now computed attention weights $\alpha_{ij}$ that tell us how much the decoder should attend to each encoder position. The final step is to use these weights to create a single vector that summarizes the relevant information from the entire input sequence. This is the context vector.

The context vector for decoder position $i$ is simply a weighted sum of the encoder hidden states:

where:

• $c_i \in \mathbb{R}^{d_h}$: the context vector, a $d_h$-dimensional summary of the relevant input information
• $\alpha_{ij}$: the attention weight for encoder position $j$ (how much to attend to that position)
• $h_j \in \mathbb{R}^{d_h}$: the encoder hidden state at position $j$
• $T_x$: the length of the source sequence

This formula is simple, but its implications are significant. Because the attention weights sum to 1, the context vector is a convex combination of the encoder states. It lies somewhere "between" them in the vector space, closer to the states with higher weights.

Consider what happens in different scenarios:

• Focused attention: If $\alpha_{i2} = 0.95$ and all other weights are near zero, then $c_i \approx h_2$. The context vector copies the encoder state at position 2.
• Distributed attention: If weights are spread across multiple positions, the context vector blends information from all of them. This is useful when the decoder needs information from multiple parts of the input.
• Uniform attention: If all weights are equal ($\alpha_{ij} = 1/T_x$), the context vector is the simple average of all encoder states, similar to a mean-pooling operation.

This geometric interpretation helps build intuition: the context vector "moves" through the encoder state space based on attention weights, getting pulled toward whichever states receive more attention. The model learns to position the context vector where it captures the most relevant information for the current generation step.

The key insight is that each decoder step gets its own context vector $c_i$, computed fresh based on the current decoder state $s_{i-1}$. This is the main improvement over standard encoder-decoder models: instead of forcing all information through a single fixed-size bottleneck, we dynamically select and aggregate the relevant parts of the input at each generation step.

We've now developed all three components of Bahdanau attention:

1. Scoring: $e_{ij} = v_a^\top \tanh(W_a s_{i-1} + U_a h_j)$ measures relevance
2. Normalization: $\alpha_{ij} = \text{softmax}(e_{ij})$ converts scores to probabilities
3. Aggregation: $c_i = \sum_j \alpha_{ij} h_j$ creates a weighted summary

But how does this attention mechanism integrate with the decoder RNN? The context vector $c_i$ contains a summary of the relevant input information, but the decoder also needs to know what it has generated so far and maintain its own internal state.

In Bahdanau's formulation, these three sources of information are combined by feeding them into the decoder RNN:

where:

• $s_i \in \mathbb{R}^{d_s}$: the new decoder hidden state after processing step $i$
• $s_{i-1} \in \mathbb{R}^{d_s}$: the previous decoder hidden state (the decoder's "memory")
• $y_{i-1} \in \mathbb{R}^{d_e}$: the embedding of the previous output token
• $c_i \in \mathbb{R}^{d_h}$: the context vector from attention
• $f$: a recurrent cell function (typically GRU or LSTM)

In practice, $y_{i-1}$ and $c_i$ are concatenated into a single vector of dimension $d_e + d_h$ before being fed to the RNN. This combines information about the previous prediction with the attention-weighted summary of the input.

The complete decoding process at step $i$ follows this sequence:

1. Compute alignment scores: For each encoder position $j$, calculate $e_{ij} = v_a^\top \tanh(W_a s_{i-1} + U_a h_j)$
2. Normalize to attention weights: Apply softmax to get $\alpha_{ij}$
3. Compute context vector: Calculate $c_i = \sum_j \alpha_{ij} h_j$
4. Update decoder state: Compute $s_i = f(s_{i-1}, y_{i-1}, c_i)$
5. Generate output: Predict the next token using $s_i$ (and often $c_i$ as well)

One subtle but important detail: attention is computed using the previous decoder state $s_{i-1}$, not the current state $s_i$. This makes sense because we need to know what the decoder is looking for before we can compute the context, and we need the context before we can update the decoder state.

The formulas we've developed might seem abstract, so let's ground them with a concrete numerical example. We'll trace through every step of the attention computation, watching the numbers flow through each transformation.

Imagine we're building a translation system for English to French. We're translating "I love cats" and have just generated "J'" (the French equivalent of "I"). Now we need to generate the next word, which should be "aime" (love). The question is: which part of the input should the decoder focus on?

For this example, we'll use deliberately small dimensions so we can see every number:

• Hidden dimension: 4 (real systems use 256-1024)
• Alignment dimension: 3
• Source sequence: 3 words ("I", "love", "cats")

Each encoder hidden state $h_j$ captures information about word $j$ and its surrounding context. The decoder state $s_{i-1}$ encodes what the decoder has generated so far and, implicitly, what it's looking for next. The weight matrices $W_a$, $U_a$, and vector $v_a$ would normally be learned during training; here we use fixed values for illustration.

The first step is to transform the decoder state into the alignment space. This projection happens once per decoder step and will be reused when comparing against each encoder position:

This 3-dimensional vector represents "what the decoder is looking for" in the alignment space. Think of it as a query that we'll compare against each encoder position.

Now we compute a score for each encoder position. For each position $j$, we project $h_j$ into the alignment space, add it to the projected decoder state, apply tanh, and compute the final score with $v_a$:

The scores reveal that position 2 ("love") has the highest alignment with what the decoder is looking for. This makes intuitive sense: to generate "aime" (the French word for love), the decoder should focus on the English word "love." Position 3 ("cats") has the lowest score, indicating it's least relevant for this particular generation step.

Raw scores can be any real number, but we need probabilities. Softmax transforms these scores into a proper probability distribution:

Position 2 receives about 45% of the attention, while positions 1 and 3 receive roughly 28% and 27% respectively. Notice that even though position 2 had the highest score, the attention isn't completely focused on it. This "soft" attention allows the model to hedge its bets, pulling in information from multiple positions when useful. The weights sum to exactly 1.0, confirming we have a valid probability distribution.

Finally, we use these attention weights to create a weighted combination of the encoder hidden states:

The context vector is a blend of all three encoder states, weighted by their relevance. It's closest to $h_2$ (the representation of "love") since that position received the highest attention weight, but it also incorporates information from the other positions. This context vector now gets fed into the decoder RNN along with the previous output embedding to generate the next token.

We can measure how the context vector relates to the encoder states using cosine similarity:

The similarity analysis confirms that the context vector is closest to the encoder state that received the highest attention weight. However, the context isn't identical to any single encoder state; it's a weighted blend that can capture information from multiple positions when needed.

This worked example illustrates the key insight of attention: instead of forcing the decoder to work with a single fixed representation of the entire input, we dynamically construct a context that's tailored to each generation step. When generating "aime," the context emphasizes "love"; when generating "chats" later, the context would shift to emphasize "cats."

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.

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

1. Projection layers: W_a, U_a, and v_a are implemented as nn.Linear layers without bias terms, matching our weight matrices
2. Broadcasting: When we add decoder_proj.unsqueeze(1) to encoder_proj, PyTorch broadcasts the decoder projection across all sequence positions
3. Batch processing: The code handles batches of sequences simultaneously, which is essential for efficient training
4. 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:

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.

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.

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

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

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

Attention mechanisms are interpretable. We can visualize the attention weights to see what the model is "looking at" when generating each output token.

The attention heatmap reveals several interesting patterns:

• Monotonic alignment: For this simple sentence, attention roughly follows the diagonal, reflecting the similar word order between English and French.
• Word-to-word correspondence: "Le" attends strongly to "The," "chat" to "cat," and "tapis" to "mat."
• Many-to-one mapping: Both "était" and "assis" attend primarily to "sat," since the single English word requires two French words.
• Soft attention: Even when one position dominates, other positions receive small but non-zero weights.

We can also visualize how the attention distribution evolves step by step during decoding. Each row in the heatmap above represents a single decoding step, but viewing them as individual distributions makes the shifting focus more apparent:

This interpretability is valuable for debugging and understanding model behavior. If a translation is wrong, we can inspect the attention weights to see if the model was looking at the right parts of the input.

A useful metric for understanding attention behavior is entropy, which quantifies how spread out or focused the attention distribution is. Low entropy indicates concentrated attention on few positions; high entropy indicates diffuse attention across many positions.

The entropy analysis reveals that most decoding steps in this translation have relatively focused attention (entropy below 1.0), indicating clear word-to-word correspondences. Steps with higher entropy might indicate cases where the model needs to aggregate information from multiple source positions, such as translating idiomatic expressions or handling word reordering.

While English-French translation often exhibits roughly monotonic alignment (words appear in similar order), attention can handle more complex patterns. Languages with different word orders, or sentences with long-distance dependencies, produce non-monotonic attention patterns.

This example demonstrates one of attention's key strengths: it can learn arbitrary alignment patterns without explicit supervision. The model discovers that German places the past participle ("gegessen") at the end of the clause, and learns to attend to the corresponding English words ("has eaten") regardless of their position.

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:

1. Project the decoder state: Computing $W_a s_{i-1}$ requires $O(d_a \times d_s)$ operations (matrix-vector multiplication)
2. Project all encoder states: Computing $U_a h_j$ for all $j$ requires $O(T_x \times d_a \times d_h)$ operations
3. Compute scores: The addition, tanh, and dot product with $v_a$ require $O(T_x \times d_a)$ operations
4. Apply softmax: Normalizing scores requires $O(T_x)$ operations
5. 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.

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:

Bahdanau attention changed how we think about sequence-to-sequence models. However, it has several limitations that motivated subsequent research.

The most significant limitation is the sequential nature of decoding. Because attention at step $i$ depends on the decoder state $s_{i-1}$, and $s_{i-1}$ depends on the attention at step $i-1$, we cannot parallelize across decoder timesteps. Each step must wait for the previous step to complete. This sequential bottleneck limits training throughput on modern parallel hardware like GPUs, where we'd prefer to process many positions simultaneously.

Another limitation is the additive score function's computational cost. Computing $\tanh(W_a s_{i-1} + U_a h_j)$ requires more operations than simpler alternatives like dot-product attention. While the difference is small for individual computations, it adds up over millions of training examples. Luong attention, which we'll cover in the next chapter, addresses this with more efficient score functions.

The attention mechanism also introduces additional hyperparameters (the alignment dimension $d_a$) and learnable parameters ($W_a$, $U_a$, $v_a$). While these provide flexibility, they also increase the risk of overfitting on small datasets and require careful tuning.

Despite these limitations, Bahdanau attention had a major impact on the field. It showed that attention mechanisms could improve performance on sequence-to-sequence tasks, reducing the BLEU score gap between neural and phrase-based machine translation systems. It also introduced the core ideas that would evolve into the self-attention mechanism at the heart of Transformers.

The key insights that carried forward include:

• Dynamic context: Computing a fresh context vector at each decoding step, rather than using a fixed representation
• Soft alignment: Learning to align input and output positions without explicit supervision
• Interpretability: Attention weights provide insight into model behavior
• Differentiable lookup: Treating attention as a soft, differentiable lookup table

These ideas laid the groundwork for the attention-based models that would reshape NLP over the following years.

A natural question is: how does attention learn to produce meaningful alignments? At initialization, the attention weights are random. Through training, the model gradually discovers which source positions are relevant for each target position.

The learning progression illustrates a key property of attention: it's trained end-to-end with the translation objective. The model receives no explicit supervision about which source words correspond to which target words. Instead, it discovers these alignments because they help minimize the translation loss. This emergent alignment is one of the most interesting aspects of attention mechanisms.

Bahdanau attention solves the bottleneck problem in encoder-decoder models by allowing the decoder to dynamically focus on relevant parts of the input at each generation step. The mechanism works through a learned alignment model that scores how well each encoder position matches the current decoder state.

The key components are:

• Additive score function: $e_{ij} = v_a^\top \tanh(W_a s_{i-1} + U_a h_j)$ computes alignment scores by projecting the decoder state $s_{i-1}$ and encoder states $h_j$ into a shared alignment space, combining them additively, and reducing to a scalar
• Softmax normalization: $\alpha_{ij} = \frac{\exp(e_{ij})}{\sum_k \exp(e_{ik})}$ converts raw scores to attention weights that form a probability distribution over encoder positions
• Context vector: $c_i = \sum_j \alpha_{ij} h_j$ computes a weighted sum of encoder states, where the weights $\alpha_{ij}$ determine how much each position contributes to the context for generating output $i$

The attention weights are interpretable, showing which input positions the model considers relevant for each output position. This interpretability, combined with strong empirical performance, made Bahdanau attention a foundational technique in neural machine translation and sequence-to-sequence modeling more broadly.

In the next chapter, we'll explore Luong attention, which offers alternative score functions with different computational trade-offs.

When implementing Bahdanau attention, several parameters significantly impact model performance:

• attention_dim (d_a): The dimension of the alignment space where encoder and decoder states are projected before comparison. Typical values range from 64 to 512. Larger values increase model capacity but also computation cost. A common heuristic is to set this equal to the encoder or decoder hidden dimension.

• encoder_dim (d_h): The dimension of encoder hidden states. For bidirectional encoders, this is typically twice the base RNN hidden size (forward + backward concatenated). Values of 256-1024 are common in practice.

• decoder_dim (d_s): The dimension of decoder hidden states. Often set equal to the encoder dimension for simplicity, though they can differ. The decoder dimension affects both the attention computation and the output projection.

• embed_dim: The dimension of word embeddings fed to the decoder. Typical values are 128-512. Smaller embeddings reduce parameters but may limit expressiveness for large vocabularies.

• vocab_size: The size of the output vocabulary. Larger vocabularies require more parameters in the output projection layer and can slow down training due to the softmax computation over all tokens.

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about Bahdanau attention and its role in neural machine translation.