The Attention Mechanism

I finally learned to pay attention, and now I can't stop staring at every token in the sequence. My therapist calls it hypervigilance; Bahdanau calls it alignment.

Attn, Hypervigilant AI Agent

2.2.1 The "Where to Look" Intuition

Consider a human translator converting an English sentence to French. They do not read the entire English sentence, memorize it as a single compressed thought, and then produce the French sentence from memory. Instead, they glance back at specific parts of the source sentence as they write each word of the translation. When writing the French verb, they look at the English verb. When writing the French adjective, they look at the English adjective (which might be in a different position due to word order differences between the two languages).

Attention gives neural networks this same ability. At each decoder step, the model computes a set of attention weights over all encoder positions. These weights determine how much "focus" to place on each source token. The model then takes a weighted sum of the encoder hidden states to produce a context vector that is specific to the current decoding step.

Crucially, these weights are not hardcoded. They are computed dynamically based on the decoder's current state and each encoder state. Different decoder steps attend to different parts of the source, creating a soft, differentiable alignment between source and target. Understanding how attention weights are interpreted remains an active research area, as discussed in Chapter 10: Interpretability.

2.2.2 Bahdanau Additive Attention

The first attention mechanism for seq2seq, proposed by Bahdanau, Cho, and Bengio (2014), uses a small feedforward network (often called an alignment model) to compute the compatibility between the decoder state and each encoder state.

The Computation

Let $s_{i}$ be the decoder hidden state at step $i$, and let $h_{j}$ be the encoder hidden state at position $j$. Bahdanau attention computes:

These energy scores are normalized via softmax to produce attention weights that sum to one:

The context vector is then the weighted sum of encoder hidden states:

To make this concrete, consider a decoder attending to four source tokens with energy scores [0.2, 2.8, 0.1, 1.5]. The softmax converts these to weights, and the weighted sum produces a context vector dominated by the highest-scoring position:

# Numeric example: Bahdanau attention on 4 source tokens
import torch, torch.nn.functional as F
scores = torch.tensor([0.2, 2.8, 0.1, 1.5])
weights = F.softmax(scores, dim=0)
print(f"Scores: {scores.tolist()}")
print(f"Weights: {weights.tolist()}")
# Weights: [0.049, 0.660, 0.045, 0.246] (sum = 1.0)
# Weighted sum of 4 encoder vectors (dim=3 for brevity)
h = torch.tensor([[1.0, 0.0, 0.0], # h1
    [0.0, 1.0, 0.0], # h2
    [0.0, 0.0, 1.0], # h3
    [0.5, 0.5, 0.0]]) # h4
context = weights @ h
print(f"Context: {context.tolist()}")
# Context: [0.172, 0.783, 0.045] (dominated by h2, weight=0.66)

Here is what each piece does:

1. Score function $e_{ij}$: Projects both the decoder state and encoder state into a common space, adds them, applies tanh, then reduces to a scalar with $v$. This is called "additive" attention because the two projections are added together.
2. Softmax normalization: Converts raw scores into a probability distribution over source positions. The weights $\alpha _{ij}$ sum to 1 and are all non-negative.
3. Weighted sum: Combines encoder states according to the attention weights, producing a context vector tailored to the current decoding step.

Figure 2.2.2: Bahdanau additive attention. At each decoder step, scores are computed between the decoder state $s_{i}$ and every encoder state $h_{j}$. After softmax normalization, the attention weights form a probability distribution that determines how much each source position contributes to the context vector.

Attention as Soft Alignment

In traditional machine translation, an alignment specifies which source word(s) correspond to each target word. Classic statistical MT systems learned hard alignments (each target word aligns to exactly one or a few source words). Attention produces a soft alignment: every target word is connected to every source word, but with varying strength.

This soft alignment has several advantages. It can handle one-to-many and many-to-one relationships naturally. It is fully differentiable, allowing end-to-end training with backpropagation. And it provides interpretability: by visualizing the attention weights as a matrix (source positions on one axis, target positions on the other), we get an alignment map that shows which source words the model "looked at" when generating each target word.

2.2.3 Luong Dot-Product Attention

Shortly after Bahdanau's work, Luong et al. (2015) proposed several alternative score functions. The most widely used is dot-product attention, which replaces the feedforward network with a simple dot product:

Luong also proposed a "general" variant that inserts a learnable matrix:

The dot-product score is computationally cheaper and can be computed as a single matrix multiplication across all source positions simultaneously. This efficiency advantage becomes critical as we scale attention to the Transformer architecture in Section 2.3. However, note that dot-product attention requires the decoder and encoder states to have the same dimensionality, while additive attention can handle different sizes through its projection matrices.

2.2.4 Attention as a Differentiable Dictionary Lookup

There is a powerful analogy that helps unify all forms of attention: think of it as a soft dictionary lookup.

In a traditional Python dictionary, you provide a key and retrieve the exact matching value. If the key is not present, you get nothing. Attention works similarly but in a "soft" way:

• Query: What you are looking for (the decoder state $s_{i}$)
• Keys: What you compare against (the encoder states $h_{j}$)
• Values: What you retrieve (also the encoder states $h_{j}$ in Bahdanau/Luong)

Instead of exact matching, the query is compared against all keys simultaneously using a similarity function. The result is not a single value but a weighted combination of all values, where the weights reflect how well each key matches the query.

2.2.5 Implementing Attention from Scratch

Let us build both Bahdanau and Luong attention in PyTorch (Section 0.3), step by step.

# Bahdanau (additive) attention: project decoder state and encoder outputs
# through a shared MLP, then softmax over positions to get alignment weights.
import torch
import torch.nn as nn
import torch.nn.functional as F
class BahdanauAttention(nn.Module):
    """Additive (Bahdanau) attention mechanism."""
    def __init__(self, enc_dim, dec_dim, attn_dim):
        super().__init__()
        self.W1 = nn.Linear(dec_dim, attn_dim, bias=False)
        self.W2 = nn.Linear(enc_dim, attn_dim, bias=False)
        self.v = nn.Linear(attn_dim, 1, bias=False)
    def forward(self, decoder_state, encoder_outputs):
        """
        decoder_state: (batch, dec_dim)
        encoder_outputs: (batch, src_len, enc_dim)
        Returns: context (batch, enc_dim), weights (batch, src_len)
        """
        # Expand decoder state to match encoder sequence length
        dec_proj = self.W1(decoder_state).unsqueeze(1) # (batch, 1, attn_dim)
        enc_proj = self.W2(encoder_outputs) # (batch, src_len, attn_dim)
        # Additive scoring: v^T tanh(W1*s + W2*h)
        scores = self.v(torch.tanh(dec_proj + enc_proj)) # (batch, src_len, 1)
        scores = scores.squeeze(-1) # (batch, src_len)
        # Normalize to get attention weights
        weights = F.softmax(scores, dim=-1) # (batch, src_len)
        # Weighted sum of encoder outputs
        context = torch.bmm(weights.unsqueeze(1), # (batch, 1, src_len)
            encoder_outputs) # x (batch, src_len, enc_dim)
        context = context.squeeze(1) # (batch, enc_dim)
        return context, weights
        # Test it
        attn = BahdanauAttention(enc_dim=128, dec_dim=128, attn_dim=64)
        enc_out = torch.randn(2, 6, 128) # batch=2, 6 source tokens
        dec_s = torch.randn(2, 128) # decoder state
        ctx, wts = attn(dec_s, enc_out)
        print(f"Context shape: {ctx.shape}")
        print(f"Weights shape: {wts.shape}")
        print(f"Weights sum: {wts[0].sum().item():.4f}")
        print(f"Weights[0]: {wts[0].detach().numpy().round(3)}")

The weights sum to 1 (as guaranteed by softmax) and the model has learned to concentrate most of its attention on positions 1 and 4 for this particular query. Now the Luong dot-product variant:

import torch.nn.functional as F
from torch import nn
import torch
# Luong (dot-product) attention: score = decoder_state @ encoder_outputs^T.
# Simpler and faster than Bahdanau because it skips the MLP projection.
class LuongDotAttention(nn.Module):
 """Dot-product (Luong) attention mechanism."""
# Forward pass: define computation graph
def forward(self, decoder_state, encoder_outputs):
 """
        decoder_state: (batch, dim)
        encoder_outputs: (batch, src_len, dim) , same dim required!
        """
# Dot product: s^T h for each encoder position
scores = torch.bmm(
encoder_outputs, # (batch, src_len, dim)
decoder_state.unsqueeze(-1) # (batch, dim, 1)
).squeeze(-1) # (batch, src_len)
# Convert scores to attention weights (probabilities summing to 1)
weights = F.softmax(scores, dim=-1)
context = torch.bmm(weights.unsqueeze(1), encoder_outputs).squeeze(1)
return context, weights
# Compare both mechanisms
luong_attn = LuongDotAttention()
ctx_l, wts_l = luong_attn(dec_s, enc_out)
print(f"Bahdanau weights: {wts[0].detach().numpy().round(3)}")
print(f"Luong weights: {wts_l[0].detach().numpy().round(3)}")
print(f"Both produce context of shape: {ctx_l.shape}")

Notice that Luong attention produces sharper weights (more concentrated on a single position). This is because the raw dot product can produce larger magnitude scores than the bounded tanh in Bahdanau attention, leading to more peaked softmax outputs.

Understanding how attention scores are computed is only half the story. To appreciate why attention works so well as a trainable mechanism, we need to see how gradients flow through the attention computation during learning.

2.2.6 Backpropagation Through Attention

Attention is fully differentiable, which means gradients flow through it naturally during backpropagation. The key takeaway: positions that receive high attention weights also receive strong gradient signals, creating a virtuous cycle where the model learns to attend to the right places.

Advanced Deep Dive (Optional): Matrix calculus of gradient flow through attention

The Jacobian of Softmax

The attention weights are computed as $\alpha = \operatorname{softmax}(e)$. The Jacobian of softmax with respect to its input is:

$$\partial \alpha _{i} / \partial e_{j} = \alpha _{i}( \delta _{ij} - \alpha _{j})$$

where $\delta _{ij}$ is the Kronecker delta (1 if i = j, 0 otherwise). This has an important consequence: increasing score $e_{j}$ increases weight $\alpha _{j}$ but decreases all other weights (since they must sum to 1). This competitive dynamic is what makes attention selective.

Gradient Flow Through the Context Vector

The context vector is $c = \sum _{j} \alpha _{j} h_{j}$. The gradient of the loss $L$ with respect to encoder state $h_{j}$ has two paths:

1. Direct path: $\partial L/ \partial c \cdot \alpha _{j}$. The gradient flows directly through the weighted sum, scaled by the attention weight. Positions with high attention receive more gradient.
2. Indirect path: Through the score function. Changing $h_{j}$ changes the scores, which changes the weights, which changes the context. This path allows the network to learn better scoring functions.

Note

The direct path is analogous to a Section 4.1: gradient flows directly from the output back to the relevant encoder states, scaled by the attention weight. This means that positions the model attends to receive strong gradient signals, making training much more effective than in a vanilla seq2seq model where gradients must traverse the entire encoder recurrence.

Let us verify this gradient flow empirically:

# Compare Bahdanau vs Luong attention: identical encoder outputs,
# different scoring functions. Check that gradients flow to all positions.
import torch
import torch.nn.functional as F
# Create encoder outputs with gradient tracking
enc = torch.randn(1, 5, 8, requires_grad=True) # 5 positions, dim 8
query = torch.randn(1, 8)
# Compute Luong dot-product attention
scores = torch.bmm(enc, query.unsqueeze(-1)).squeeze(-1)
weights = F.softmax(scores, dim=-1)
context = torch.bmm(weights.unsqueeze(1), enc).squeeze(1)
# Compute a scalar loss and backpropagate
loss = context.sum()
loss.backward()
print("Attention weights:", weights[0].detach().numpy().round(4))
print("Gradient norms per position:")
for j in range(5):
    grad_norm = enc.grad[0, j].norm().item()
    print(f" Position {j}: weight={weights[0,j]:.4f}, grad_norm={grad_norm:.4f}")

Output: Attention weights: [0.0312 0.7841 0.0028 0.1753 0.0066] Gradient norms per position: Position 0: weight=0.0312, grad_norm=0.2489 Position 1: weight=0.7841, grad_norm=1.3572 Position 2: weight=0.0028, grad_norm=0.0891 Position 3: weight=0.1753, grad_norm=0.6734 Position 4: weight=0.0066, grad_norm=0.1243

Code Fragment 2.2.4: Create encoder outputs with gradient tracking.

Observe the strong correlation: positions with higher attention weights receive proportionally larger gradients. Position 1, which has 78% of the attention weight, receives by far the largest gradient. This is the mechanism by which attention guides learning: the model receives the strongest training signal from the source positions it deems most relevant.

Figure 2.2.4a: Gradient flow through attention. The direct path scales gradients by α_j, giving highly-attended positions stronger learning signals. The indirect path flows through the score function, allowing the network to improve its attention distribution.

2.2.7 Integrating Attention into Seq2Seq

Now let us see how attention fits into a complete encoder-decoder model. The decoder uses the context vector (alongside its hidden state) to make predictions at each step:

# Full attention decoder: at each step, attend over encoder outputs,
# concatenate the context vector with the embedding, and predict the next token.
import torch
import torch.nn as nn
import torch.nn.functional as F
class AttentionDecoder(nn.Module):
    """Decoder with Bahdanau attention."""
    def __init__(self, vocab_size, emb_dim, hidden_dim, enc_dim):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, emb_dim)
        self.rnn = nn.GRU(emb_dim + enc_dim, hidden_dim, batch_first=True)
        self.attention = BahdanauAttention(enc_dim, hidden_dim, attn_dim=64)
        self.fc = nn.Linear(hidden_dim + enc_dim, vocab_size)
    def forward_step(self, token, hidden, encoder_outputs):
        """One decoding step with attention."""
        # Compute attention over encoder outputs
        context, weights = self.attention(
            hidden.squeeze(0), # (batch, hidden_dim)
            encoder_outputs # (batch, src_len, enc_dim)
            )
        # Embed current token and concatenate with context
        emb = self.embedding(token) # (batch, 1, emb_dim)
        rnn_input = torch.cat([emb, context.unsqueeze(1)], dim=-1)
        # RNN step
        output, hidden = self.rnn(rnn_input, hidden) # output: (batch, 1, hidden)
        # Combine RNN output with context for prediction
        combined = torch.cat([output.squeeze(1), context], dim=-1)
        logits = self.fc(combined)
        return logits, hidden, weights
        # Demo: decode 4 steps with attention
        dec = AttentionDecoder(vocab_size=6000, emb_dim=64, hidden_dim=128, enc_dim=128)
        enc_outputs = torch.randn(1, 8, 128) # 8 source tokens
        h = torch.randn(1, 1, 128) # initial decoder hidden
        tok = torch.tensor([[1]]) # <SOS>
        print("Attention patterns during decoding:")
        for step in range(4):
            logits, h, weights = dec.forward_step(tok, h, enc_outputs)
            tok = logits.argmax(dim=-1, keepdim=True)
            w = weights[0].detach().numpy()
            peak = w.argmax()
            print(f" Step {step}: peak attention at source pos {peak} ({w[peak]:.2%})")

Each decoding step focuses on a different source position, demonstrating the dynamic nature of attention. The decoder is no longer stuck with a single context vector; it constructs a fresh, step-specific context at every position.

Exercises