Luong Attention › Implementation
Language AI HandbookAll chapters Luong Attention › ImplementationOne section · Sections in this chapter6 / 11 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let's implement Luong attention in PyTorch. We'll build a complete attention module that supports all three score functions, then integrate it into a sequence-to-sequence decoder.
Attention Module
First, we define the core attention computation. The module takes encoder outputs and a decoder hidden state, computes attention weights, and returns the context vector.
import torch
import torch.nn as nn
import torch.nn.functional as F
class LuongAttention(nn.Module):
"""Luong attention with configurable score function."""
def __init__(self, hidden_dim, encoder_dim=None, method="dot"):
super().__init__()
self.method = method
self.hidden_dim = hidden_dim
self.encoder_dim = encoder_dim or hidden_dim
if method == "general":
# W_a for bilinear scoring
self.W_a = nn.Linear(self.encoder_dim, hidden_dim, bias=False)
elif method == "concat":
# Two-layer scoring network
self.W_a = nn.Linear(
hidden_dim + self.encoder_dim, hidden_dim, bias=False
)
self.v_a = nn.Linear(hidden_dim, 1, bias=False)
def score(self, decoder_hidden, encoder_outputs):
"""Compute attention scores for all encoder positions."""
if self.method == "dot":
scores = torch.bmm(
encoder_outputs, decoder_hidden.unsqueeze(2)
).squeeze(2)
elif self.method == "general":
transformed = self.W_a(encoder_outputs)
scores = torch.bmm(
transformed, decoder_hidden.unsqueeze(2)
).squeeze(2)
elif self.method == "concat":
src_len = encoder_outputs.size(1)
decoder_expanded = decoder_hidden.unsqueeze(1).expand(
-1, src_len, -1
)
concat = torch.cat([decoder_expanded, encoder_outputs], dim=2)
scores = self.v_a(torch.tanh(self.W_a(concat))).squeeze(2)
return scores
def forward(self, decoder_hidden, encoder_outputs, mask=None):
"""Compute attention weights and context vector."""
scores = self.score(decoder_hidden, encoder_outputs)
if mask is not None:
scores = scores.masked_fill(mask == 0, float("-inf"))
attention_weights = F.softmax(scores, dim=1)
context = torch.bmm(
attention_weights.unsqueeze(1), encoder_outputs
).squeeze(1)
return context, attention_weightsThe LuongAttention class implements all three score functions. The constructor accepts the hidden dimension, an optional encoder dimension, and the scoring method. The score method computes alignment scores using dot product, general (bilinear), or concat approaches. The forward method orchestrates the full attention computation: scores, optional masking, softmax normalization, and weighted sum.
Attentional Decoder
Now let's build a decoder that uses Luong attention. The key difference from a standard RNN decoder is the attention step after the RNN and the combination layer that produces the attentional hidden state.
class LuongDecoder(nn.Module):
"""Decoder with Luong attention."""
def __init__(
self,
vocab_size,
embed_dim,
hidden_dim,
encoder_dim=None,
attention_method="dot",
num_layers=1,
dropout=0.1,
):
super().__init__()
self.hidden_dim = hidden_dim
self.encoder_dim = encoder_dim or hidden_dim
self.embedding = nn.Embedding(vocab_size, embed_dim)
self.rnn = nn.GRU(
embed_dim,
hidden_dim,
num_layers=num_layers,
batch_first=True,
dropout=dropout if num_layers > 1 else 0,
)
self.attention = LuongAttention(
hidden_dim, self.encoder_dim, attention_method
)
self.W_c = nn.Linear(
self.encoder_dim + hidden_dim, hidden_dim, bias=False
)
self.output_projection = nn.Linear(hidden_dim, vocab_size)
self.dropout = nn.Dropout(dropout)
def forward_step(self, input_token, hidden, encoder_outputs, mask=None):
"""Single decoding step with attention."""
embedded = self.dropout(self.embedding(input_token.unsqueeze(1)))
rnn_output, hidden = self.rnn(embedded, hidden)
rnn_output = rnn_output.squeeze(1)
context, attention_weights = self.attention(
hidden[-1], encoder_outputs, mask
)
combined = torch.cat([context, rnn_output], dim=1)
attentional_hidden = torch.tanh(self.W_c(combined))
output = self.output_projection(self.dropout(attentional_hidden))
return output, hidden, attention_weightsThe LuongDecoder class combines an embedding layer, a GRU, the attention module, a combination layer for producing the attentional hidden state, and an output projection. The forward_step method implements a single decoding step: embed, RNN, then attention. This order is the defining characteristic of Luong attention.
Testing the Implementation
Let's verify our implementation works correctly with some example data.
# Create sample data
batch_size = 4
src_len = 10
hidden_dim = 64
vocab_size = 1000
# Simulated encoder outputs (normally from an encoder RNN)
encoder_outputs = torch.randn(batch_size, src_len, hidden_dim)
# Padding mask (last 2 positions are padding)
mask = torch.ones(batch_size, src_len)
mask[:, -2:] = 0
# Initialize decoder
decoder = LuongDecoder(
vocab_size=vocab_size,
embed_dim=32,
hidden_dim=hidden_dim,
attention_method="dot",
)
# Initial hidden state
hidden = torch.zeros(1, batch_size, hidden_dim)
# Input token (e.g., start token)
input_token = torch.zeros(batch_size, dtype=torch.long)
# Run one decoding step
output, new_hidden, attention_weights = decoder.forward_step(
input_token, hidden, encoder_outputs, mask
)Output shape: torch.Size([4, 1000]) Hidden shape: torch.Size([1, 4, 64]) Attention weights shape: torch.Size([4, 10]) Attention weights sum to 1: 1.0000 Attention on padded positions: 0.000000
The output has the expected shape for vocabulary logits, and attention weights sum to 1 (ignoring padded positions). The near-zero attention on padded positions confirms our masking works correctly.
Visualizing Attention Patterns
Let's compare the attention patterns produced by different score functions on the same input.
# Create decoders with different attention methods
methods = ["dot", "general", "concat"]
decoders = {
method: LuongDecoder(vocab_size, 32, hidden_dim, attention_method=method)
for method in methods
}
# Run decoding steps and collect attention weights
all_attention = {}
for method, dec in decoders.items():
hidden = torch.zeros(1, batch_size, hidden_dim)
_, _, attn = dec.forward_step(input_token, hidden, encoder_outputs, mask)
all_attention[method] = attn.detach()
With random weights, the attention patterns are similar across methods. After training, each method would develop distinct patterns based on what it learns about alignment. Dot product attention tends to produce sharper distributions when encoder and decoder representations align well, while concat attention can learn more complex compatibility functions.
To see what trained attention looks like, let's simulate a full translation with realistic attention patterns. The heatmap below shows attention weights across multiple decoding steps, revealing how the decoder systematically aligns with different source positions.
