attention.py

# attention.py
# Bahdanau Attention Mechanism
#
# The attention mechanism was introduced in this paper:
# "Neural Machine Translation by Jointly Learning to Align and Translate"
# Bahdanau, Cho, Bengio - 2015
# https://arxiv.org/abs/1409.0473
#
# Problem with basic Seq2Seq: the entire input is compressed into a single
# fixed-size vector. For long sentences this loses a lot of information.
#
# Attention solution: at each decoding step, let the model "look back" at
# all encoder outputs and decide which words are most relevant right now.

import torch
import torch.nn as nn
import torch.nn.functional as F


class BahdanauAttention(nn.Module):
    """
    Additive (Bahdanau) Attention.
    
    At each decoder step:
    1. Compare current decoder hidden state with all encoder outputs
    2. Compute attention weights (which encoder positions matter most?)
    3. Return weighted sum of encoder outputs as "context"
    
    This context is then used to generate the next output word.
    """

    def __init__(self, hidden_dim):
        super(BahdanauAttention, self).__init__()

        # learned linear transformations for computing attention energy
        self.attn = nn.Linear(hidden_dim * 2, hidden_dim)
        self.v    = nn.Linear(hidden_dim, 1, bias=False)

    def forward(self, decoder_hidden, encoder_outputs):
        """
        decoder_hidden:   [batch, hidden_dim]         -- current decoder state
        encoder_outputs:  [batch, src_len, hidden_dim] -- all encoder states

        returns:
            context:      [batch, hidden_dim]  -- weighted encoder summary
            attn_weights: [batch, src_len]     -- attention distribution (sums to 1)
        """
        src_len = encoder_outputs.shape[1]

        # repeat decoder hidden state so we can compare it with each encoder output
        hidden_expanded = decoder_hidden.unsqueeze(1).repeat(1, src_len, 1)  # [batch, src_len, hidden]

        # compute attention energy using tanh activation
        combined = torch.cat([hidden_expanded, encoder_outputs], dim=2)   # [batch, src_len, hidden*2]
        energy   = torch.tanh(self.attn(combined))                         # [batch, src_len, hidden]
        scores   = self.v(energy).squeeze(2)                               # [batch, src_len]

        # softmax to get probabilities (attention weights)
        attn_weights = F.softmax(scores, dim=1)                           # [batch, src_len]

        # weighted sum of encoder outputs
        context = torch.bmm(attn_weights.unsqueeze(1), encoder_outputs)   # [batch, 1, hidden]
        context = context.squeeze(1)                                       # [batch, hidden]

        return context, attn_weights