Bahdanau Attention

The Encoder

A bidirectional RNN (typically a BiGRU) reads the source sequence $x = (x_1, \ldots, x_T)$. At each position $j$, the forward and backward hidden states are concatenated to form an annotation vector:

Explanation: Each annotation $h_j$ captures local context around position $j$ from both directions, giving a richer representation than a unidirectional encoder. Here $\overrightarrow{h}_j \in \mathbb{R}^n$ is the forward hidden state, $\overleftarrow{h}_j \in \mathbb{R}^n$ the backward, and $[\cdot\,;\cdot]$ denotes concatenation, so $h_j \in \mathbb{R}^{2n}$.

The Alignment Model

At decoder time step $i$, before running the decoder RNN, an alignment score $e_{ij}$ measures how well source position $j$ matches the decoder's current search query $s_{i-1}$ (the previous decoder hidden state):

Explanation: This is a single-hidden-layer MLP that scores compatibility between the decoder's query and each encoder annotation — the "additive" structure is why this is called additive attention. Symbols:

• $s_{i-1} \in \mathbb{R}^m$ — previous decoder hidden state (the query)
• $h_j \in \mathbb{R}^{2n}$ — encoder annotation at position $j$ (the key/value, conflated here)
• $W_a \in \mathbb{R}^{k \times m}$, U_a \in \mathbb{R}^^{k \times 2n} — learned projection matrices
• $v_a \in \mathbb{R}^k$ — learned scoring vector; $k$ is the attention dimensionality (hyperparameter)

Attention Weights

The raw scores are normalized across all $T$ source positions via softmax to produce a proper probability distribution:

Explanation: The weights $\alpha_{ij}$ sum to 1 over $j$ and represent how much "attention" the decoder pays to each source position when generating the $i$-th output token.

Context Vector

The context vector is the attention-weighted sum of all encoder annotations:

Explanation: This is a soft, differentiable lookup — instead of hard-selecting one encoder state, the decoder reads a convex combination of all of them, parameterized by the learned alignment.

Decoder Step

The decoder RNN receives the context vector concatenated with the previous output embedding $y_{i-1}$:

Explanation: Every decoder state update is conditioned on the dynamically computed context $c_i$, so each output token has access to a fresh, task-relevant summary of the entire source sequence. The output distribution over the vocabulary is then $p(y_i \mid y_{<i}, x) = \text{softmax}(W_o s_i)$.

PyTorch Implementation Sketch

```python import torch import torch.nn as nn import torch.nn.functional as F class BahdanauAttention(nn.Module): def __init__(self, encoder_dim: int, decoder_dim: int, attn_dim: int): super().__init__()

# W_a projects decoder query self.W_a = nn.Linear(decoder_dim, attn_dim, bias=False)

# U_a projects encoder annotations self.U_a = nn.Linear(encoder_dim, attn_dim, bias=False)

# v_a produces scalar scores self.v_a = nn.Linear(attn_dim, 1, bias=False) def forward( self, s_prev: torch.Tensor, # (batch, decoder_dim) h: torch.Tensor, # (batch, src_len, encoder_dim) ):

# Expand query to match src_len for broadcasting query = self.W_a(s_prev).unsqueeze(1) # (batch, 1, attn_dim) keys = self.U_a(h) # (batch, src_len, attn_dim)

# Additive (Bahdanau) score energy = self.v_a(torch.tanh(query + keys)) # (batch, src_len, 1) energy = energy.squeeze(-1) # (batch, src_len) alpha = F.softmax(energy, dim=-1) # (batch, src_len)

# Weighted sum of encoder hidden states context = torch.bmm(alpha.unsqueeze(1), h) # (batch, 1, encoder_dim) context = context.squeeze(1) # (batch, encoder_dim) return context, alpha ```