Implementing the architecture from the paper “Neural Machine Translation by Jointly Learning to Align and Translate” (Bahdanau et al., 2015) involves three main components: a Bidirectional GRU Encoder, the Bahdanau Attention mechanism, and an Attentional Decoder.
The key innovation is that the decoder doesn’t use a single fixed vector; instead, it “looks” at all encoder hidden states and calculates a weighted sum (the context vector) based on their relevance to the current decoding step.
1. The Encoder (Bidirectional GRU)
The encoder processes the source sequence in both directions. The annotation for each word is the concatenation of the forward ($\overrightarrow{h}$) and backward ($\overleftarrow{h}$) hidden states.
Python
import torch
import torch.nn as nn
import torch.nn.functional as F
class Encoder(nn.Module):
def __init__(self, input_dim, emb_dim, enc_hid_dim, dec_hid_dim, dropout):
super().__init__()
self.embedding = nn.Embedding(input_dim, emb_dim)
self.rnn = nn.GRU(emb_dim, enc_hid_dim, bidirectional=True)
self.fc = nn.Linear(enc_hid_dim * 2, dec_hid_dim)
self.dropout = nn.Dropout(dropout)
def forward(self, src):
# src: [src_len, batch_size]
embedded = self.dropout(self.embedding(src))
# outputs: [src_len, batch_size, enc_hid_dim * 2]
# hidden: [n_layers * directions, batch_size, enc_hid_dim]
outputs, hidden = self.rnn(embedded)
# Initial decoder hidden state is a linear transformation of the
# concatenated final forward and backward encoder hidden states
hidden = torch.tanh(self.fc(torch.cat((hidden[-2,:,:], hidden[-1,:,:]), dim=1)))
return outputs, hidden
2. The Attention Mechanism (Bahdanau Attention)
This is a feed-forward layer that takes the previous decoder hidden state $s_{i-1}$ and all encoder hidden states $h_j$ to calculate the energy (alignment) scores.
$$e_{ij} = v_a^\top \tanh(W_a s_{i-1} + U_a h_j)$$
Python
class Attention(nn.Module):
def __init__(self, enc_hid_dim, dec_hid_dim):
super().__init__()
self.attn = nn.Linear((enc_hid_dim * 2) + dec_hid_dim, dec_hid_dim)
self.v = nn.Linear(dec_hid_dim, 1, bias=False)
def forward(self, hidden, encoder_outputs):
# hidden: [batch_size, dec_hid_dim]
# encoder_outputs: [src_len, batch_size, enc_hid_dim * 2]
batch_size = encoder_outputs.shape[1]
src_len = encoder_outputs.shape[0]
# Repeat decoder hidden state src_len times
hidden = hidden.unsqueeze(1).repeat(1, src_len, 1)
encoder_outputs = encoder_outputs.permute(1, 0, 2)
# Energy: [batch_size, src_len, dec_hid_dim]
energy = torch.tanh(self.attn(torch.cat((hidden, encoder_outputs), dim=2)))
# Attention: [batch_size, src_len]
attention = self.v(energy).squeeze(2)
return F.softmax(attention, dim=1)
3. The Decoder
The decoder takes the attention weights, creates a context vector, and uses it along with the current input to predict the next word.
Python
class Decoder(nn.Module):
def __init__(self, output_dim, emb_dim, enc_hid_dim, dec_hid_dim, dropout, attention):
super().__init__()
self.output_dim = output_dim
self.attention = attention
self.embedding = nn.Embedding(output_dim, emb_dim)
self.rnn = nn.GRU((enc_hid_dim * 2) + emb_dim, dec_hid_dim)
self.fc_out = nn.Linear((enc_hid_dim * 2) + dec_hid_dim + emb_dim, output_dim)
self.dropout = nn.Dropout(dropout)
def forward(self, input, hidden, encoder_outputs):
# input: [batch_size]
input = input.unsqueeze(0)
embedded = self.dropout(self.embedding(input))
# Calculate attention weights
a = self.attention(hidden, encoder_outputs).unsqueeze(1)
# Context vector: [batch_size, 1, enc_hid_dim * 2]
encoder_outputs = encoder_outputs.permute(1, 0, 2)
weighted = torch.bmm(a, encoder_outputs)
weighted = weighted.permute(1, 0, 2)
# RNN Input: concat(embedded, context_vector)
rnn_input = torch.cat((embedded, weighted), dim=2)
output, hidden = self.rnn(rnn_input, hidden.unsqueeze(0))
# Final Linear layer
prediction = self.fc_out(torch.cat((output.squeeze(0), weighted.squeeze(0), embedded.squeeze(0)), dim=1))
return prediction, hidden.squeeze(0)
4. Seq2Seq Wrapper
Finally, combine everything into a single model to manage the training loop and inference.
Python
class Seq2Seq(nn.Module):
def __init__(self, encoder, decoder, device):
super().__init__()
self.encoder = encoder
self.decoder = decoder
self.device = device
def forward(self, src, trg, teacher_forcing_ratio=0.5):
batch_size = src.shape[1]
trg_len = trg.shape[0]
trg_vocab_size = self.decoder.output_dim
outputs = torch.zeros(trg_len, batch_size, trg_vocab_size).to(self.device)
enc_outputs, hidden = self.encoder(src)
input = trg[0,:]
for t in range(1, trg_len):
output, hidden = self.decoder(input, hidden, enc_outputs)
outputs[t] = output
teacher_force = torch.rand(1).item() < teacher_forcing_ratio
top1 = output.argmax(1)
input = trg[t] if teacher_force else top1
return outputs
Key Implementation Details:
- Context Vector: A weighted sum of encoder hidden states that dynamically represents the “attended” parts of the source sentence for each target word.
- Teacher Forcing: A technique used during training where the actual ground-truth word is fed as the next input instead of the model’s own prediction.
- Tanh Activation: Used in the attention scoring function to introduce non-linearity, which allows the model to learn complex alignment patterns.