Loung Attention

Building the attention mechanism introduced by Loung et al. (2015) that improves the approach Bahdanau et al. (2014) proposed in his paper.

import numpy as np
import torch
import torch.nn as nn

Data preperation

We will use the OPUS-100 "ar-en" data. Taking a max sentence lenght of 12, a train size of 3000, and test size of 300. The below code cell collects the data in pairs, each pair is an eng-arb snetances.

from data_pipeline import (
    train_loader, test_pairs,
    english_word_to_index, english_index_to_word,
    arabic_word_to_index, arabic_index_to_word,
    tokenize, encode_source, PAD_IDX,
)

src_batch, dec_in_batch, dec_tgt_batch, src_lengths = next(iter(train_loader))

Encoder

The encoder takes input tokens then create hidden states for each using RNN.

import torch.nn.utils.rnn as rnn_utils


class Encoder(nn.Module):
    def __init__(self, input_size, hidden_size):
        super(Encoder, self).__init__()
        self.hidden_size = hidden_size
        self.embedding = nn.Embedding(
            input_size, hidden_size, padding_idx=PAD_IDX)
        self.gru = nn.GRU(hidden_size, hidden_size, batch_first=True)

    def forward(self, input_seq, lengths):
        # input_seq: (batch, seq_len) — token IDs
        embedded = self.embedding(input_seq)   # (batch, seq_len, hidden_size)

        # telling PyTorch's RNN internals the true length of each sequence
        packed = rnn_utils.pack_padded_sequence(
            embedded, lengths.cpu(), batch_first=True, enforce_sorted=False
        )
        packed_outputs, hidden = self.gru(packed)

        # unpack the output
        outputs, _ = rnn_utils.pad_packed_sequence(
            packed_outputs, batch_first=True, total_length=input_seq.size(1)
        )
        # outputs: (batch, seq_len, hidden), hidden: (1, batch, hidden)
        return outputs, hidden


print("src_batch:", src_batch.shape)
print("src_lengths:", src_lengths)

encoder = Encoder(len(english_word_to_index), hidden_size=64)
outputs, hidden = encoder(src_batch, src_lengths)

print("outputs:", outputs.shape)
print("hidden :", hidden.shape)

src_batch: torch.Size([32, 13])
src_lengths: tensor([ 6,  5, 13,  6,  6,  9,  8,  6,  6,  5,  5,  8,  6,  7,  6,  8,  5,  9,
         8,  9,  6,  7,  4,  7, 11, 11,  8,  7,  6,  4,  9,  9])
outputs: torch.Size([32, 13, 64])
hidden : torch.Size([1, 32, 64])

outputs: torch.Size([32, 16, 64])

hidden : torch.Size([1, 32, 64])

• The encoder returns the batch examples outputs 32 examples, with a seq_lenth of the longest word in the batch 16, and each token in the sequence is represented by a 64 vector lentgh.

• The second output is a (,64) size vector for each example in the 32, representing the final hidden state.

Decoder

Dot attention (Global)

As we discussed on the paper discussion blog, instead of applying a nural network between decoder hidden states and the encoder states (Bahdanau approach), Loung used a simple dot product.

The class takes the following: - Single decoder_hidden at each timestep. —> (batch, hidden) - Full encoder_outputs –> (batch, seq_len, hidden)

class DotAttention(nn.Module):
    def __init__(self):
        super(DotAttention, self).__init__()

    def forward(self, encoder_outputs, decoder_hidden, src_lengths=None):
        
        ## Step 1
        # compute alignment scores: dot product of decoder_hidden with every encoder timestep
        alignment_scores = torch.bmm(
            decoder_hidden.unsqueeze(1),          # (batch, 1, hidden)
            encoder_outputs.transpose(1, 2)       # (batch, hidden, seq_len)
        ).squeeze(1)                              # (batch, seq_len)

        ## Step 1b: scale by sqrt(hidden_size) to prevent softmax from being too peaky
        div = np.sqrt(encoder_outputs.size(2))   # scale factor
        alignment_scores = alignment_scores / div
        
        # Step 2
        # softmax
        attention_weights = torch.softmax(alignment_scores, dim=1)   # (batch, seq_len)

        # Step 3
        # weighted sum to get the context vector
        context_vector = torch.bmm(
            attention_weights.unsqueeze(1),       # (batch, 1, seq_len)
            encoder_outputs                        # (batch, seq_len, hidden)
        ).squeeze(1)                              # (batch, hidden)
        
        return attention_weights, context_vector

torch.manual_seed(0)

batch_size = 3
seq_len = 5
hidden_size = 8

decoder_hidden = torch.randn(batch_size, hidden_size)          # (3, 8)
encoder_outputs = torch.randn(batch_size, seq_len, hidden_size)  # (3, 5, 8)

dot_attention = DotAttention()
attention_weights, context_vector = dot_attention(
    encoder_outputs, decoder_hidden)

print("attention_weights:", attention_weights.shape)   # (3, 5)
print(attention_weights)

print("\ncontext_vector:", context_vector.shape)        # (3, 8)
print(context_vector)

# sanity check: weights should sum to 1 across seq_len for every example
print("\nweights sum to 1:", torch.allclose(
    attention_weights.sum(dim=1), torch.ones(batch_size)))

attention_weights: torch.Size([3, 5])
tensor([[1.1674e-02, 2.5890e-04, 9.8594e-01, 2.0530e-03, 7.0815e-05],
        [5.4518e-03, 8.8998e-01, 4.2941e-02, 2.5005e-02, 3.6623e-02],
        [3.8082e-02, 3.6649e-01, 9.7240e-04, 5.8052e-03, 5.8865e-01]])

context_vector: torch.Size([3, 8])
tensor([[ 0.1978, -0.4521, -0.5588, -0.5454,  0.5807,  1.5118,  0.5116, -0.5739],
        [-0.0577, -0.4819,  1.0642,  0.4404,  1.7172,  0.6103, -0.0427, -0.4438],
        [ 0.6688,  0.4010,  0.0835, -0.2114, -0.1094, -0.6222, -0.1047, -0.7556]])

weights sum to 1: True

Monotonic Attention (local)

The authors proposed two types of loacal attention that they introduced to address the expensive computational resources needed for gloabl attention where at each decoder step, the alignment score is computed by attending to all encoder hidden states.

The two types are: - Monotonic Attention (local-m): decoder at step 1 looks at encoder hidden state 1, and so on.

• Predictive Alignment (local-p): the model learns where to look in the source sentence by predicting a position \(p_t\) using a small neural network.

Below we implement the first type.

class LocalMonotonicAttention(nn.Module):
    def __init__(self, window_size):
        super(LocalMonotonicAttention, self).__init__()
        self.D = window_size

    def forward(self, encoder_outpus, decoder_hidden, t, D):
        '''
            encoder_outputs:
                (batch_size, seq_len, hidden_size)

            decoder_hidden:
                (batch_size, hidden_size)

            t: int
                current decoder timestep. For local-m, this IS the aligned source
                position (p_t = t) — no prediction needed, unlike local-p.

            D: int
                half-width of the attention window. Only source positions in
                [t - D, t + D] are attended to (window size = 2D + 1).

            Returns:
            ---------
                attention_weights: (batch_size, 2D + 1)

                context_vector: (batch_size, hidden_size)
        '''
        #Step 1: get the window of encoder states
        start = max(0, t - self.D)
        end = min(seq_len, t + self.D + 1)
        encoder_window = encoder_outpus[:, start:end, :]  # (batch_size, window_size, hidden_size)

        #Step 2: compute alignment scores (dot product)
        alignment_scores = torch.bmm(
            decoder_hidden.unsqueeze(1),          # (batch, 1, hidden)
            encoder_window.transpose(1, 2)        # (batch, hidden, window_size)
        ).squeeze(1)                              # (batch, window_size)

        #Step 3: softmax to get attention weights
        attention_weights = torch.softmax(alignment_scores, dim=1)   # (batch, window_size)

        #Step 4: compute context vector (weighted sum)
        context_vector = torch.bmm(
            attention_weights.unsqueeze(1),     # (batch, 1, window_size)
            encoder_window                       # (batch, window_size, hidden_size)
        ).squeeze(1)                              # (batch, hidden_size)


        return attention_weights, context_vector

Predictive Alignment

Below is the class for the second type that uses predictive alignment as the second approach to find the alignment score then the context vector.

class PredictiveAlignmentAttention(nn.Module):
    def __init__(self, hidden_size, window_size):
        super(PredictiveAlignmentAttention, self).__init__()
        self.D = window_size
        self.position_predictor = nn.Linear(hidden_size, 1)  # Predicts p_t
        self.position_weights = nn.Linear(hidden_size, hidden_size)

    def forward(self, encoder_outputs, decoder_hidden, src_lengths=None):
        '''
            encoder_outputs:
                (batch_size, seq_len, hidden_size)

            decoder_hidden:
                (batch_size, hidden_size)

            src_lengths:
                (batch_size,) - actual lengths of source sequences

            Returns:
            ---------
                attention_weights: (batch_size, 2D + 1)

                context_vector: (batch_size, hidden_size)
        '''
        batch_size, seq_len, hidden_size = encoder_outputs.shape
        
        
        if src_lengths is None:   # fallback for standalone/mock testing
            src_lengths = torch.full((batch_size,), seq_len, dtype=torch.long, device=encoder_outputs.device)

        # Step1: predict a position within each example's REAL length
        pt = self.position_predictor(torch.tanh(
            self.position_weights(decoder_hidden)))
        pt = torch.sigmoid(pt).squeeze(-1) * \
            (src_lengths.float() - 1)   # (batch,)

        centers = pt.round().long()

        # clamp per example using ITS real length — never lets a center point into padding
        lower = torch.full_like(centers, self.D)
        upper = torch.clamp(src_lengths - 1 - self.D, min=self.D)
        centers = torch.max(torch.min(centers, upper), lower)
        self.last_centers = centers.detach()

        # Step2: get the window of encoder states
        context_vectors = []
        attention_weights = []
        sigma = self.D / 2.0
        for b in range(batch_size):
            center = centers[b].item()
            
            start = max(0, center - self.D)
            end = min(seq_len, center + self.D + 1)
            
            
            # encoder states inside the window
            encoder_window = encoder_outputs[b, start:end, :]          # (window, hidden)

            # decoder hidden for this sample
            decoder = decoder_hidden[b]
            
            # ---------- Dot scores ----------
            scores = torch.matmul(encoder_window, decoder)
            
            # ---------- Alignment ----------
            align = torch.softmax(scores, dim=0)

            # ---------- Gaussian ----------
            positions = torch.arange(
                start,
                end,
                device=encoder_outputs.device,
                dtype=torch.float
            )

            gaussian = torch.exp(-((positions - pt[b]) ** 2) / (2 * sigma ** 2))

            # ---------- Eq. (10) ----------
            weights = align * gaussian

            # Renormalize
            weights = weights / weights.sum()

            # ---------- Context ----------
            context = torch.matmul(
                weights,
                encoder_window
            )                                                          # (hidden)

            context_vectors.append(context)
            attention_weights.append(weights)
            
            
        context_vectors = torch.stack(context_vectors)
        attention_weights = torch.stack(attention_weights)

        return attention_weights, context_vectors

predictive_attention = PredictiveAlignmentAttention(hidden_size=hidden_size, window_size=2)
attention_weights, context_vectors = predictive_attention(encoder_outputs, decoder_hidden)

print("Attention weights:", attention_weights)
print("Context vectors:", context_vectors)

Attention weights: tensor([[2.8365e-03, 2.1148e-04, 9.9600e-01, 9.4357e-04, 5.4474e-06],
        [1.2321e-03, 8.9519e-01, 7.0720e-02, 2.4806e-02, 8.0506e-03],
        [2.2533e-02, 8.0329e-01, 2.9046e-03, 8.6933e-03, 1.6258e-01]],
       grad_fn=<StackBackward0>)
Context vectors: tensor([[ 0.2031, -0.4504, -0.5695, -0.5529,  0.5899,  1.5336,  0.5080, -0.5857],
        [-0.1077, -0.4836,  1.1161,  0.4807,  1.6977,  0.5830, -0.0319, -0.4016],
        [ 1.1515,  0.5790,  0.3544, -0.7929,  0.1045, -0.1037, -0.5392, -0.0390]],
       grad_fn=<StackBackward0>)

The above two outptus:

• Attention weights: It says where the decoder is paying attention to, to which source word. Each source word gets a weight, we have seq_len=5 thus 5 attention weights.

• Context vectors: Is a weighted average of the encoder hidden states, it is blending the encoder representations according to the attention weights. It comes from multiplying weights by encoder hidden states.

Step 1. What is the decoder trying to do?

If we are trying to translate this I love machine learning.

The decoder job is to produce the following:

<START>           
    
↓    
    
أنا    
    
↓    
    
أحب    
    
↓    
    
تعلم    
    
↓

الآلة

↓

<END>

The input to the encoder: - Previous generated word. — Convert this to embedding

• Previous hidden state. (This is decoder memory of what generated so far)

• Context vector

What is the output? - Next hidden state

• Vocabulary probabilities

We should concatenate both the embedding and the context before entering the GRU:

Embedding (64)

+

Context (64)

↓

128

What about the output of the GRU??

It is a 64 number, and this cannot tell us the next word because the vocaulary size is 15,000 words. So we need a nural net to map the 64 output to the vocab size (fc layer).

class Decoder(nn.Module):
    def __init__(self, output_size, embed_size, hidden_size, attention_module):
        super(Decoder, self).__init__()
        self.embedding = nn.Embedding(output_size, embed_size, padding_idx=PAD_IDX)
        self.gru = nn.GRU(embed_size + hidden_size, hidden_size, batch_first=True)
        self.attention = attention_module
        self.fc = nn.Linear(hidden_size, output_size)
        
    def forward(self, previous_word, previous_hidden, encoder_outputs, src_lengths=None):
        embedded = self.embedding(previous_word).unsqueeze(1)  # (batch, 1, embed_size)
        
        # Get the current hidden to compute attention
        current_hidden = previous_hidden.squeeze(0)  # (batch, hidden_size)
        attention_weights, context_vector = self.attention(encoder_outputs, current_hidden, src_lengths)  # (batch, seq_len), (batch, hidden_size)
        self.last_attention_weights = attention_weights.detach()
        
        rnn_input = torch.cat((embedded, context_vector.unsqueeze(1)), dim=2)  # (batch, 1, embed_size + hidden_size)
        output, hidden = self.gru(rnn_input, previous_hidden)  # output: (batch, 1, hidden_size), hidden: (1, batch, hidden_size)
        
        
        # Map the output to the vocabulary space
        logits  = self.fc(output.squeeze(1))  # (batch, output_size)
        
        return logits, hidden


## Example usage of the Decoder
output_size = len(arabic_word_to_index)
embed_size = 28
src_batch, dec_in_batch, dec_tgt_batch, src_lengths = next(iter(train_loader))

dot_attention = DotAttention()

# Construct the decoder class
decoder = Decoder(output_size, embed_size, hidden_size=64, attention_module=dot_attention)
decoder_input = dec_in_batch[:, 0]  # First token of the decoder input
encoder_final_hidden = hidden  # Use the encoder's final hidden state as the initial hidden state for
encoder_outputs = outputs  # Use the encoder's outputs for attention

print("decoder_input:", decoder_input.shape)
print("decoder_hidden:", encoder_final_hidden.shape)

output, decoder_hidden = decoder(decoder_input, encoder_final_hidden, encoder_outputs)
print("Decoder output (For logits prediction):", output.shape)  # (batch, output_size)
print("Decoder hidden:", decoder_hidden.shape)  # (1, batch, hidden_size)
    

decoder_input: torch.Size([32])
decoder_hidden: torch.Size([1, 32, 64])
Decoder output (For logits prediction): torch.Size([32, 1679])
Decoder hidden: torch.Size([1, 32, 64])

Below we will work on another decoder architecture that uses the Loung attention mechanism.

class LuongDecoder(nn.Module):
    def __init__(self, output_size, embed_size, hidden_size, attention_module):
        super().__init__()
        self.embedding = nn.Embedding(output_size, embed_size, padding_idx=PAD_IDX)
        self.gru = nn.GRU(embed_size + hidden_size,
                          hidden_size, batch_first=True)
        
        self.attention = attention_module
        self.Wc = nn.Linear(hidden_size * 2, hidden_size)
        self.fc = nn.Linear(hidden_size, output_size)
    
    def forward(self, previous_word, previous_hidden, prev_htilde, encoder_outputs, src_lengths=None):
        embedded = self.embedding(previous_word).unsqueeze(1)
        # concat the embedding with previous htilde (previous step's attentional vector) then feed into gru
        rnn_input = torch.cat((embedded, prev_htilde.unsqueeze(1)), dim=2)
        output, hidden = self.gru(rnn_input, previous_hidden)
        # We take the output for this step
        h_t = output.squeeze(1)
        
        attention_weights, context = self.attention(encoder_outputs, h_t, src_lengths)   # c_t

        # compute the attentional vector h_tilde
        h_tilde = torch.tanh(self.Wc(torch.cat((context, h_t), dim=1)))
        
        # compute the logits for the next token prediction
        logits = self.fc(h_tilde)
        self.last_attention_weights = attention_weights.detach()
        
        return logits, hidden, h_tilde
        
        
        

Training

Now we will train the whole model 3 times, each using a different attention mechanism:

• Dot Attention
• Local Monotonic Attention
• Predictive Alignment Attention

import torch.optim as optim

criterion = nn.CrossEntropyLoss(ignore_index=PAD_IDX)
optimizer = optim.Adam(list(encoder.parameters()) + list(decoder.parameters()), lr=1e-3)

def train_model(attention_module, num_epochs=5, hidden_size=64, embed_size=28, lr=1e-3):
    encoder = Encoder(len(english_word_to_index), hidden_size=hidden_size)
    decoder = Decoder(len(arabic_word_to_index), embed_size,
                      hidden_size, attention_module=attention_module)

    criterion = nn.CrossEntropyLoss(ignore_index=PAD_IDX)
    optimizer = optim.Adam(list(encoder.parameters()) +
                           list(decoder.parameters()), lr=lr)

    encoder.train()
    decoder.train()
    loss_history = []

    for epoch in range(1, num_epochs + 1):
        epoch_loss = 0.0

        for src_batch, dec_in_batch, dec_tgt_batch, src_lengths in train_loader:
            optimizer.zero_grad()

            encoder_outputs, encoder_hidden = encoder(src_batch, src_lengths)
            previous_word = dec_in_batch[:, 0]
            decoder_hidden = encoder_hidden
            loss = 0

            target_len = dec_tgt_batch.size(1)
            for t in range(target_len):
                output, decoder_hidden = decoder(
                    previous_word, decoder_hidden, encoder_outputs, src_lengths)
                loss += criterion(output, dec_tgt_batch[:, t])
                if t < dec_in_batch.size(1) - 1:
                    previous_word = dec_in_batch[:, t + 1]

            loss = loss / target_len
            loss.backward()
            torch.nn.utils.clip_grad_norm_(
                list(encoder.parameters()) + list(decoder.parameters()), max_norm=1.0
            )
            optimizer.step()
            epoch_loss += loss.item()

        avg_epoch_loss = epoch_loss / len(train_loader)
        loss_history.append(avg_epoch_loss)
        print(f"  epoch {epoch}/{num_epochs} | avg loss {avg_epoch_loss:.4f}")

    return encoder, decoder, loss_history

def train_model_1(attention_module, num_epochs=5, hidden_size=64, embed_size=28, lr=1e-3):
    encoder = Encoder(len(english_word_to_index), hidden_size=hidden_size)
    decoder = LuongDecoder(len(arabic_word_to_index), embed_size,        # <-- LuongDecoder
                           hidden_size, attention_module=attention_module)

    criterion = nn.CrossEntropyLoss(ignore_index=PAD_IDX)
    optimizer = optim.Adam(list(encoder.parameters()) +
                           list(decoder.parameters()), lr=lr)

    encoder.train()
    decoder.train()
    loss_history = []

    for epoch in range(1, num_epochs + 1):
        epoch_loss = 0.0

        for src_batch, dec_in_batch, dec_tgt_batch, src_lengths in train_loader:
            optimizer.zero_grad()

            encoder_outputs, encoder_hidden = encoder(src_batch, src_lengths)
            previous_word = dec_in_batch[:, 0]
            decoder_hidden = encoder_hidden

            # seed the attentional vector: h̃_0 = zeros (no past alignment yet)
            batch_size = src_batch.size(0)                                # not BATCH_SIZE: last batch is smaller
            prev_htilde = torch.zeros(src_batch.size(
                0), hidden_size)           # <-- new

            loss = 0
            target_len = dec_tgt_batch.size(1)
            for t in range(target_len):
                # capture h̃ and feed it back in next step  <-- 3-tuple return
                output, decoder_hidden, prev_htilde = decoder(
                    previous_word, decoder_hidden, prev_htilde, encoder_outputs, src_lengths)
                loss += criterion(output, dec_tgt_batch[:, t])
                if t < dec_in_batch.size(1) - 1:
                    previous_word = dec_in_batch[:, t + 1]

            loss = loss / target_len
            loss.backward()
            torch.nn.utils.clip_grad_norm_(
                list(encoder.parameters()) + list(decoder.parameters()), max_norm=1.0
            )
            optimizer.step()
            epoch_loss += loss.item()

        avg_epoch_loss = epoch_loss / len(train_loader)
        loss_history.append(avg_epoch_loss)
        print(f"  epoch {epoch}/{num_epochs} | avg loss {avg_epoch_loss:.4f}")

    return encoder, decoder, loss_history

import matplotlib.pyplot as plt

results = {}

for name, attention_module in [
    ("Dot (Global)", DotAttention()),
    #("Predictive Alignment (local-p)",
     #PredictiveAlignmentAttention(hidden_size=64, window_size=4)),
]:
    print(f"Training with {name} attention...")
    enc, dec, loss_history = train_model_1(attention_module, num_epochs=50)
    results[name] = {"encoder": enc, "decoder": dec,
                     "loss_history": loss_history}

Training with Dot (Global) attention...
  epoch 1/50 | avg loss 5.2241

---------------------------------------------------------------------------
KeyboardInterrupt                         Traceback (most recent call last)
Cell In[42], line 11
      5 for name, attention_module in [
      6     ("Dot (Global)", DotAttention()),
      7     #("Predictive Alignment (local-p)",
      8      #PredictiveAlignmentAttention(hidden_size=64, window_size=4)),
      9 ]:
     10     print(f"Training with {name} attention...")
---> 11     enc, dec, loss_history = train_model_1(attention_module, num_epochs=50)
     12     results[name] = {"encoder": enc, "decoder": dec,
     13                      "loss_history": loss_history}

Cell In[38], line 39, in train_model_1(attention_module, num_epochs, hidden_size, embed_size, lr)
     36         previous_word = dec_in_batch[:, t + 1]
     38 loss = loss / target_len
---> 39 loss.backward()
     40 torch.nn.utils.clip_grad_norm_(
     41     list(encoder.parameters()) + list(decoder.parameters()), max_norm=1.0
     42 )
     43 optimizer.step()

File c:\Users\user\Documents\GitHub\simpe-AI\venv\Lib\site-packages\torch\_tensor.py:625, in Tensor.backward(self, gradient, retain_graph, create_graph, inputs)
    615 if has_torch_function_unary(self):
    616     return handle_torch_function(
    617         Tensor.backward,
    618         (self,),
   (...)    623         inputs=inputs,
    624     )
--> 625 torch.autograd.backward(
    626     self, gradient, retain_graph, create_graph, inputs=inputs
    627 )

File c:\Users\user\Documents\GitHub\simpe-AI\venv\Lib\site-packages\torch\autograd\__init__.py:354, in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)
    349     retain_graph = create_graph
    351 # The reason we repeat the same comment below is that
    352 # some Python versions print out the first line of a multi-line function
    353 # calls in the traceback and some print out the last line
--> 354 _engine_run_backward(
    355     tensors,
    356     grad_tensors_,
    357     retain_graph,
    358     create_graph,
    359     inputs_tuple,
    360     allow_unreachable=True,
    361     accumulate_grad=True,
    362 )

File c:\Users\user\Documents\GitHub\simpe-AI\venv\Lib\site-packages\torch\autograd\graph.py:841, in _engine_run_backward(t_outputs, *args, **kwargs)
    839     unregister_hooks = _register_logging_hooks_on_whole_graph(t_outputs)
    840 try:
--> 841     return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
    842         t_outputs, *args, **kwargs
    843     )  # Calls into the C++ engine to run the backward pass
    844 finally:
    845     if attach_logging_hooks:

KeyboardInterrupt: 

plt.figure(figsize=(8, 5))
for name, run in results.items():
    loss_history = run["loss_history"]
    plt.plot(range(1, len(loss_history) + 1),
             loss_history, marker="o", label=name)

plt.xlabel("Epoch")
plt.ylabel("Average training loss")
plt.title("Luong Attention Variants — Training Loss")
plt.legend()
plt.grid(alpha=0.3)
plt.show()

Tracing Predition path

We will take a single example, then trace the prediction path and how tensors updates until making the predition.

import torch
from data_pipeline import (
    test_pairs, tokenize, encode_source, pairs,
    english_index_to_word, arabic_word_to_index, arabic_index_to_word,
)

# reuse the model trained earlier (no retraining)
trace_encoder = results["Dot (Global)"]["encoder"].eval()
trace_decoder = results["Dot (Global)"]["decoder"].eval()

SOS_IDX = arabic_word_to_index["<sos>"]
EOS_IDX = arabic_word_to_index["<eos>"]

# <-- change the index to trace another sentence
english, arabic_ref = pairs[16]
print("English (source)   :", english)
print("Arabic  (reference):", arabic_ref)

English (source)   : We have to do this by ourselves
Arabic  (reference): علينا عمل هذا بأنفسنا

Below we take the tokenized words then encode them to their index representation.

Below we feed the encoder the src tensor containing the source sentence indexes, and the src_len tensor containing the length of the source sentence. The encoder returns the encoder_outputs and the encoder_hidden state.

We take the first input word <sos> wich is the start of sentence token, having the index 1. Below we feed this to the embedder to get the embedding vector.

After reciveing the encoder_outputs and taking the encoder_hidden as the initial hidden state for the decoder. We can get the attention weights and the context vector.

We will see that the attention weights are a vector of size seq_len which is the number of tokens in the source sentence. Each weight corresponds to a source token, and the context vector is a weighted sum of the encoder outputs based on these attention weights.

Now that we have the context vector, we can concatenate it with the embedding of the previous word and feed it to the decoder. The decoder will then produce the next hidden state and the output probabilities for the next word in the sequence.

# ---- pick the example HERE; nothing else needs re-running ----
english, arabic_ref = test_pairs[23]        # change the index

toks = tokenize(english)
src = torch.tensor([encode_source(toks)], dtype=torch.long)
src_len = torch.tensor([src.size(1)], dtype=torch.long)
src_labels = toks + ["<eos>"]

with torch.no_grad():
    # re-encode THIS sentence every run -> no stale encoder_outputs
    encoder_outputs, hidden = trace_encoder(src, src_len)

    prev = torch.tensor([SOS_IDX], dtype=torch.long)
    generated = []

    for t in range(20):
        embedded = trace_decoder.embedding(prev).unsqueeze(1)
        attn_weights, context = trace_decoder.attention(
            encoder_outputs, hidden.squeeze(0))
        rnn_input = torch.cat((embedded, context.unsqueeze(1)), dim=2)
        gru_output, hidden = trace_decoder.gru(
            rnn_input, hidden)   # hidden updates each step
        logits = trace_decoder.fc(gru_output.squeeze(1))

        prev = logits.argmax(dim=1)
        if prev.item() == EOS_IDX:
            print(f"\nstep {t}: <eos> -> stop")
            break

        word = arabic_index_to_word[prev.item()]
        generated.append(word)

        print(f"\nstep {t}: produced '{word}'")
        for pos, (w, sw) in enumerate(zip(attn_weights.squeeze(0).tolist(), src_labels)):
            mark = " <== peak" if pos == attn_weights.argmax().item() else ""
            print(f"  {pos:>2} {sw:<10} {w:6.3f} {'#'*int(round(w*30))}{mark}")

print("\nEnglish  :", english)
print("Predicted:", " ".join(generated))
print("Reference:", arabic_ref)

step 0: produced 'لا'
   0 my          0.008 
   1 father      0.010 
   2 does        0.013 
   3 not         0.033 #
   4 eat         0.045 #
   5 much        0.089 ###
   6 fruit       0.265 ########
   7 <eos>       0.536 ################ <== peak

step 1: produced '<unk>'
   0 my          0.016 
   1 father      0.028 #
   2 does        0.025 #
   3 not         0.109 ###
   4 eat         0.108 ###
   5 much        0.107 ###
   6 fruit       0.216 ######
   7 <eos>       0.392 ############ <== peak

step 2: produced '<unk>'
   0 my          0.038 #
   1 father      0.048 #
   2 does        0.045 #
   3 not         0.149 ####
   4 eat         0.177 #####
   5 much        0.123 ####
   6 fruit       0.175 #####
   7 <eos>       0.245 ####### <== peak

step 3: <eos> -> stop

English  : My father does not eat much fruit
Predicted: لا <unk> <unk>
Reference: أبي لا يأكل الكثير من الفواكه

The first question to ask is: why the attention weights all collapse to the <eos> token?

We have this finding: - “The weight vector is [0,0,…,1.0] — one-hot, every step.”

Now we ask: Where do these weights come from? - They come from alignment_scores. (See below the code example)

# One example, first decoder step — compute alignment scores explicitly
from torch import div


english, _ = test_pairs[19]
toks = tokenize(english)
src = torch.tensor([encode_source(toks)], dtype=torch.long)
src_len = torch.tensor([src.size(1)], dtype=torch.long)

# Print the attention weights for each decoder step
encoder_outputs, encoder_hidden = encoder(src, src_len)

# Attention weights and context vector for the first decoder step
attention_weights, context_vector = DotAttention()(encoder_outputs, encoder_hidden.squeeze(0))


# get the alignment scores for the first decoder step
decoder_hidden = encoder_hidden.squeeze(0)
scores = torch.bmm(
    decoder_hidden.unsqueeze(1),          # (batch, 1, hidden)
    encoder_outputs.transpose(1, 2)       # (batch, hidden, seq_len)
).squeeze(1)                              # (batch, seq_len)

print("Alignment scores:", scores)

print("alignment score per source token (raw dot product):\n")
for pos, (score, sw) in enumerate(zip(scores.squeeze(0).tolist(), src_labels)):
    print(f"  {pos:>2} {sw:<10} {score:6.3f} {'#'*int(round(score*10))}")

Alignment scores: tensor([[0.6264, 1.4047, 2.0138, 2.2400, 3.8166, 5.4974]],
       grad_fn=<SqueezeBackward1>)
alignment score per source token (raw dot product):

   0 i           0.626 ######
   1 didn't      1.405 ##############
   2 have        2.014 ####################
   3 to          2.240 ######################
   4 do          3.817 ######################################
   5 anything    5.497 #######################################################

What does the above output mean? - Try changing the input sentence to see how the weights focued heavely on the last token. Why??

Below we take the scoeres and pass them through a softmax function to get the attention weights. The softmax function converts the scores into probabilities that sum to 1. We see the last token is getting the highest weight, which means the decoder is focusing on the last token of the source sentence when generating the next word.

scores = scores.squeeze(0)   # to take 1 example
print("spread (max - min):", (scores.max() - scores.min()).item())
print("weights:", torch.softmax(scores, dim=0))     # <- watch this go one-hot

spread (max - min): 4.871016502380371
weights: tensor([0.0060, 0.0130, 0.0240, 0.0301, 0.1455, 0.7814],
       grad_fn=<SoftmaxBackward0>)

We can also experiment with softening the scores by dividing them by a factor before applying softmax. This can help to spread the attention weights more evenly across the source tokens, rather than focusing too heavily on the last token.

for div in [1, 2, 4, 8]:
    w = torch.softmax(scores / div, dim=0)
    print(
        f"divide scores by {div}: spread={(scores/div).max()-(scores/div).min():5.2f}  max weight={w.max():.3f}")

divide scores by 1: spread= 5.75  max weight=0.818
divide scores by 2: spread= 2.88  max weight=0.448
divide scores by 4: spread= 1.44  max weight=0.229
divide scores by 8: spread= 0.72  max weight=0.148

So the first fix will be: Adding a scale factor \(√d\) to the softmax function.

div = np.sqrt(hidden_size)   # scale factor
print("New weights:", torch.softmax(scores / div, dim=0))     # <- watch this go one-hot

New weights: tensor([0.0694, 0.0914, 0.1134, 0.1228, 0.2145, 0.3885],
       grad_fn=<SoftmaxBackward0>)