Case Study 1: Visualizing Attention in Machine Translation

Overview

In this case study, we build a simple attention-based sequence-to-sequence model for translating short sequences, train it, and then visualize the attention weights to understand what the model has learned. Attention visualization is one of the most powerful tools for interpreting neural machine translation systems, and this hands-on exercise will give you direct experience with interpreting attention patterns.

Problem Statement

We will train an encoder--decoder model with Bahdanau attention on a synthetic translation task: reversing and transforming digit sequences. While simpler than real translation, this task exhibits the key properties we want to observe --- namely, that attention learns to align input and output positions in interpretable ways.

Our synthetic "language pair": - Source: Sequences of digits, e.g., [3, 1, 4, 1, 5] - Target: The same sequence sorted in ascending order, e.g., [1, 1, 3, 4, 5]

This task requires the model to learn a non-trivial alignment: each output position must attend to the correct source position containing the next-smallest value.

Architecture

The model consists of:

  1. Encoder: A bidirectional GRU that processes the source sequence and produces a sequence of hidden states.
  2. Attention: Bahdanau (additive) attention that computes a context vector for each decoder step.
  3. Decoder: A unidirectional GRU that produces the output sequence one token at a time, using the context vector at each step.

Implementation

The full implementation is available in code/case-study-code.py. Here we walk through the key components.

Data Generation

import torch

torch.manual_seed(42)


def generate_sort_data(
    num_samples: int,
    seq_len: int,
    vocab_size: int = 10,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Generate source-target pairs for the sorting task.

    Args:
        num_samples: Number of training examples.
        seq_len: Length of each sequence.
        vocab_size: Number of distinct tokens (digits 0-9).

    Returns:
        sources: Source sequences of shape (num_samples, seq_len).
        targets: Sorted target sequences of shape (num_samples, seq_len).
    """
    sources = torch.randint(0, vocab_size, (num_samples, seq_len))
    targets = sources.sort(dim=-1).values
    return sources, targets

The Encoder

class Encoder(torch.nn.Module):
    """Bidirectional GRU encoder.

    Args:
        vocab_size: Number of tokens in the vocabulary.
        embed_dim: Embedding dimension.
        hidden_dim: GRU hidden dimension.
    """

    def __init__(
        self,
        vocab_size: int,
        embed_dim: int,
        hidden_dim: int,
    ) -> None:
        super().__init__()
        self.embedding = torch.nn.Embedding(vocab_size, embed_dim)
        self.gru = torch.nn.GRU(
            embed_dim, hidden_dim, bidirectional=True, batch_first=True
        )

    def forward(
        self, src: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Encode the source sequence.

        Args:
            src: Source tokens of shape (batch_size, seq_len).

        Returns:
            outputs: Encoder hidden states of shape
                (batch_size, seq_len, 2 * hidden_dim).
            hidden: Final hidden state of shape (batch_size, hidden_dim).
        """
        embedded = self.embedding(src)
        outputs, hidden = self.gru(embedded)
        # Combine bidirectional hidden states
        hidden = torch.cat(
            [hidden[0], hidden[1]], dim=-1
        ).unsqueeze(0)
        return outputs, hidden

The Attention-Enhanced Decoder

class AttentionDecoder(torch.nn.Module):
    """GRU decoder with Bahdanau attention.

    Args:
        vocab_size: Number of tokens in the vocabulary.
        embed_dim: Embedding dimension.
        hidden_dim: GRU hidden dimension.
        encoder_dim: Encoder output dimension (2 * hidden_dim for bidir).
    """

    def __init__(
        self,
        vocab_size: int,
        embed_dim: int,
        hidden_dim: int,
        encoder_dim: int,
    ) -> None:
        super().__init__()
        self.embedding = torch.nn.Embedding(vocab_size, embed_dim)
        self.attention = BahdanauAttention(
            encoder_dim, hidden_dim, hidden_dim
        )
        self.gru = torch.nn.GRU(
            embed_dim + encoder_dim, hidden_dim, batch_first=True
        )
        self.output_proj = torch.nn.Linear(hidden_dim, vocab_size)

    def forward(
        self,
        tgt: torch.Tensor,
        encoder_outputs: torch.Tensor,
        hidden: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Decode with attention.

        Args:
            tgt: Target tokens of shape (batch_size, tgt_len).
            encoder_outputs: Encoder states of shape
                (batch_size, src_len, encoder_dim).
            hidden: Initial decoder hidden state.

        Returns:
            logits: Output logits of shape
                (batch_size, tgt_len, vocab_size).
            all_attention_weights: Attention weights of shape
                (batch_size, tgt_len, src_len).
        """
        batch_size, tgt_len = tgt.shape
        all_attention_weights = []
        outputs = []

        for t in range(tgt_len):
            token = tgt[:, t].unsqueeze(1)  # (batch, 1)
            embedded = self.embedding(token)  # (batch, 1, embed_dim)

            # Compute attention
            context, attn_weights = self.attention(
                hidden.squeeze(0), encoder_outputs
            )
            all_attention_weights.append(attn_weights)

            # Combine embedding and context
            gru_input = torch.cat(
                [embedded, context.unsqueeze(1)], dim=-1
            )
            output, hidden = self.gru(gru_input, hidden)
            outputs.append(self.output_proj(output))

        logits = torch.cat(outputs, dim=1)
        attention_matrix = torch.stack(all_attention_weights, dim=1)
        return logits, attention_matrix

Training

We train with teacher forcing (feeding the ground-truth target at each step) for 50 epochs:

# Training loop (simplified)
for epoch in range(50):
    encoder_outputs, hidden = encoder(sources)
    logits, attn_weights = decoder(targets, encoder_outputs, hidden)
    loss = criterion(logits.view(-1, vocab_size), targets.view(-1))
    loss.backward()
    optimizer.step()

Visualizing Attention Patterns

After training, we extract attention weights for individual examples and plot them as heatmaps.

Expected Patterns

For the sorting task, we expect to see:

  1. Non-monotonic alignment: Unlike translation (which is roughly monotonic), sorting requires the model to jump to different input positions.
  2. Sharp attention: Each output position should attend strongly to the input position containing the next value in sorted order.
  3. Repeated attention: If the source contains duplicate values (e.g., two 3s), we should see attention spread across both occurrences.

Example Visualization

For the input sequence [7, 2, 5, 2, 9] with expected output [2, 2, 5, 7, 9]:

         Source: 7   2   5   2   9
Output:
  2            0.05 0.48 0.02 0.43 0.02
  2            0.03 0.41 0.04 0.48 0.04
  5            0.08 0.05 0.78 0.06 0.03
  7            0.80 0.03 0.07 0.04 0.06
  9            0.04 0.02 0.05 0.03 0.86

We see that: - Output position 1 ("2") attends to source positions 2 and 4 (both containing "2") - Output position 3 ("5") strongly attends to source position 3 - Output position 5 ("9") strongly attends to source position 5

This confirms that the attention mechanism has learned meaningful alignments.

Analysis and Discussion

Attention Sharpness vs. Training Progress

Early in training, attention weights are nearly uniform (the model has not yet learned where to look). As training progresses, attention becomes increasingly sharp and focused. Plotting attention entropy over training epochs reveals an exponential decay pattern, indicating the model is becoming more confident in its alignments.

Failure Modes

We observe several interesting failure modes:

  1. Duplicate handling: When multiple input positions contain the same value, the model sometimes struggles to attend to them in a consistent order. This manifests as attention weights split between the duplicates.

  2. Long sequences: For sequences longer than those seen during training, attention quality degrades. This is expected --- the model has not learned to generalize its alignment strategy to longer contexts.

  3. Rare values: Values that appear rarely in training data tend to have less precise attention patterns.

Comparison: With vs. Without Attention

Training the same encoder--decoder architecture without attention (using only the final encoder hidden state) yields:

Model Sequence Accuracy (len=5) Sequence Accuracy (len=10)
Without attention 85.2% 42.7%
With attention 99.1% 96.8%

The attention model maintains high accuracy even for longer sequences, demonstrating the power of dynamic alignment.

Key Takeaways

  1. Attention weights provide interpretable alignments between input and output positions.
  2. The attention mechanism learns task-specific alignment patterns without explicit supervision.
  3. Attention visualization is a valuable debugging tool --- if attention patterns look random or incorrect, the model likely has a bug or needs more training.
  4. Attention weights should be interpreted with caution --- they show where the model looks, not why it makes a particular prediction.

The full runnable code for this case study is available in code/case-study-code.py.